Reorg the sparse/quant/common kernel dir - #1303
Conversation
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR reorganizes kernel code into a unified Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Common Attention<br/>(common/attention)
participant Sparsity as Sparsity Helpers<br/>(skip_softmax_helpers.py)
participant Calibrate as Calibrate Module<br/>(calibrate.py)
participant GPU as GPU Counters<br/>(device buffers)
Client->>Client: _load_sparsity_helpers() on first call
Client->>Sparsity: load _apply_sparse_nm, _skip_softmax_decision, _is_dense_region
alt Normal forward
Client->>Client: attention() compute qk scores
Client->>Sparsity: _apply_sparse_nm_to_qk_tile() if N:M enabled
Client->>Sparsity: _skip_softmax_decision() if skip-softmax enabled
Client->>Client: return attention output
else Calibration flow
Client->>Calibrate: attention_calibrate(thresholds)
Calibrate->>Calibrate: launch _attn_fwd_calibrate kernel
Calibrate->>GPU: write per-program totals/skipped counters
Calibrate->>Calibrate: reduce counters and return (output, sparsity_counters)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelopt/torch/kernels/__init__.py`:
- Line 16: Restore a temporary compatibility bridge in modelopt.torch.kernels by
re-exporting the legacy Triton-level symbols that were removed (reintroduce the
module-level names that downstream code imports) and emit a DeprecationWarning
when those legacy names are accessed; update the package __init__ (module
modelopt.torch.kernels) to import and re-export the new implementations (or
alias them) under the old names and add a short deprecation shim that warns
users and points to the new import path so existing runtime imports continue to
work during the migration window.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 0e39d97c-3957-4253-97cf-b9ad751b3c59
📒 Files selected for processing (49)
examples/deepseek/ptq.pyexamples/deepseek/quantize_to_nvfp4.pymodelopt/torch/kernels/__init__.pymodelopt/torch/kernels/common/__init__.pymodelopt/torch/kernels/common/hf_triton_attention.pymodelopt/torch/kernels/common/triton_fa.pymodelopt/torch/kernels/quantization/__init__.pymodelopt/torch/kernels/quantization/conv/README.mdmodelopt/torch/kernels/quantization/conv/__init__.pymodelopt/torch/kernels/quantization/conv/bench_implicit_gemm.pymodelopt/torch/kernels/quantization/conv/implicit_gemm_binding.cppmodelopt/torch/kernels/quantization/conv/implicit_gemm_cuda.pymodelopt/torch/kernels/quantization/conv/implicit_gemm_kernel.cumodelopt/torch/kernels/quantization/gemm/__init__.pymodelopt/torch/kernels/quantization/gemm/fp4_kernel.pymodelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.pymodelopt/torch/kernels/quantization/gemm/fp8_kernel.pymodelopt/torch/kernels/quantization/gemm/tensor_quant.cppmodelopt/torch/kernels/quantization/gemm/tensor_quant.hmodelopt/torch/kernels/quantization/gemm/tensor_quant_gpu.cumodelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cumodelopt/torch/kernels/quantization/gemm/tensor_quant_mx.cumodelopt/torch/kernels/quantization/gemm/tensor_quant_mx.hmodelopt/torch/kernels/sparsity/__init__.pymodelopt/torch/kernels/sparsity/attention/__init__.pymodelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.pymodelopt/torch/kernels/sparsity/attention/ltx_triton_attention.pymodelopt/torch/kernels/sparsity/gemm/__init__.pymodelopt/torch/quantization/extensions.pymodelopt/torch/quantization/nn/modules/quant_conv.pymodelopt/torch/quantization/plugins/huggingface.pymodelopt/torch/quantization/qtensor/nvfp4_tensor.pymodelopt/torch/quantization/tensor_quant.pymodelopt/torch/sparsity/attention_sparsity/conversion.pymodelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.pymodelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.pypyproject.tomltests/gpu/torch/quantization/kernels/test_implicit_gemm.pytests/gpu/torch/quantization/test_tensor_quant_cuda.pytests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_skip_softmax.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_sparse_nm.pytests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.pytests/unit/torch/kernels/test_triton_fa.pytests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.pytests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.pytests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1303 +/- ##
==========================================
+ Coverage 74.46% 75.68% +1.21%
==========================================
Files 464 467 +3
Lines 50089 50176 +87
==========================================
+ Hits 37300 37975 +675
+ Misses 12789 12201 -588
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@modelopt/torch/kernels/sparsity/attention/calibrate.py`:
- Around line 157-172: The kernel computes num_q_tiles using tl.load(b_seq_len +
0) (batch 0 length) which can under-estimate and cause prog_idx to exceed
allocated buffers (affecting Per_program_totals and Per_program_skipped); change
the kernel signature to accept max_input_len and use tl.cdiv(max_input_len,
BLOCK_M) when computing num_q_tiles/prog_idx, update the prog_idx/base
calculation sites (prog_idx, base) accordingly, and update the Python wrapper
that launches the kernel to pass the same max_input_len used to size the
Per_program_* buffers so the kernel and host agree on bounds.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: b59a3685-d843-4abe-8b17-41f6c2e12a9a
📒 Files selected for processing (17)
modelopt/torch/kernels/common/__init__.pymodelopt/torch/kernels/common/attention/__init__.pymodelopt/torch/kernels/common/attention/hf_triton_attention.pymodelopt/torch/kernels/common/attention/triton_fa.pymodelopt/torch/kernels/quantization/attention/__init__.pymodelopt/torch/kernels/sparsity/attention/__init__.pymodelopt/torch/kernels/sparsity/attention/calibrate.pymodelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.pymodelopt/torch/kernels/sparsity/attention/ltx_triton_attention.pymodelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.pytests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_skip_softmax.pytests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_sparse_nm.pytests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.pytests/unit/torch/kernels/test_triton_fa.py
✅ Files skipped from review due to trivial changes (6)
- modelopt/torch/kernels/quantization/attention/init.py
- modelopt/torch/kernels/common/attention/hf_triton_attention.py
- tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py
- modelopt/torch/kernels/sparsity/attention/init.py
- tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py
- modelopt/torch/kernels/common/init.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py
- modelopt/torch/kernels/sparsity/attention/diffusers_triton_attention.py
- tests/unit/torch/kernels/test_triton_fa.py
- modelopt/torch/kernels/sparsity/attention/ltx_triton_attention.py
- tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
This is a well-executed code reorganization moving kernels from scattered locations into a clean modelopt/torch/kernels/{common,quantization,sparsity} hierarchy. All internal import sites have been updated, tests are moved and adjusted, and the lazy-import pattern in triton_fa.py properly handles the circular dependency between common.attention and sparsity.attention.
Two items worth the owner's attention:
-
Backward compatibility claim is misleading: The PR checks "backward compatible ✅" but removes public symbols from
modelopt.torch.kernels(previously exportedIS_AVAILABLE,attention,attention_calibrate,register_triton_attention) and eliminates themodelopt.torch.quantization.tritonandmodelopt.torch.sparsity.attention_sparsity.kernelsimport paths. External consumers using these paths will break. Consider adding deprecation re-exports in the old__init__.pylocations, or at minimum update the backward-compatibility checkbox to ❌ with an explanation. -
Copyright year on new files:
calibrate.pyandskip_softmax_helpers.py(new files in this PR) haveCopyright (c) 2024— these should likely be2024-2025or2025since they're being created now (even though the code is moved from an existing file).
Thank you Chenjie, change the Backward compatibility to false since we changed the import path Copyright year is correct, all the file under this repo followed this pattern. |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
This is a well-organized directory restructuring, but two critical issues from previous reviews remain unaddressed:
-
Backward compatibility (critical, unresolved): The PR removes
modelopt.torch.quantization.triton,modelopt.torch.kernels.{IS_AVAILABLE, attention, ...}, andmodelopt.torch.sparsity.attention_sparsity.kernelsimport paths entirely with no deprecation shims. Any downstream code using these paths will break at import time. The author correctly marked backward compat as ❌, but no migration path is provided. At minimum, add__getattr__deprecation shims in the old__init__.pylocations. -
Buffer overrun bug in calibrate.py (critical, unresolved): Line 157 uses
tl.load(b_seq_len + 0)(batch 0's sequence length) to computenum_q_tilesfor the flat program index, but the Python wrapper (line ~243) sizesPer_program_*buffers usingmax_input_len. When batch 0 has a shorter sequence than other batches, the kernel computes a smallernum_q_tiles, causingprog_idxfor later batches to potentially exceed buffer bounds, leading to out-of-bounds writes. This was flagged in the previous review and is still present. Passmax_input_lenas a kernel parameter to fix the mismatch.
| acc = tl.dot(p.to(v.dtype), v, acc) | ||
| row_max = m_new | ||
|
|
||
| # --- Write per-program counters (no atomics, just stores) --- |
There was a problem hiding this comment.
Bug (unresolved from previous review): tl.load(b_seq_len + 0) loads batch 0's sequence length to compute num_q_tiles, but the Python wrapper sizes Per_program_totals/Per_program_skipped using max_input_len. If batch 0's seq_len < max_input_len (i.e., another batch is longer), the kernel computes a smaller num_q_tiles, causing prog_idx for later batches to exceed the allocated buffer — out-of-bounds write.
Fix: pass max_input_len as a kernel argument and use it here:
num_q_tiles = tl.cdiv(max_input_len, BLOCK_M)There was a problem hiding this comment.
Good catch, but it's not OOB, the buffer is sized by max_input_len, so the stride mismatch causes prog_idx collisions (silent wrong counts), not out-of-bounds. Hasn't fired because every current caller uses uniform-length batches.
Fixed: swap tl.load(b_seq_len + 0) -> tl.num_programs(2), plus a wrapper guard that raises on non-uniform b_seq_len until varlen is validated.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Two critical issues flagged in previous reviews remain unresolved in this revision:
-
Buffer overrun bug in
calibrate.py(UNRESOLVED): Line 157 still usestl.load(b_seq_len + 0)(batch 0's sequence length) to computenum_q_tilesfor the flat program index. The Python wrapper sizesPer_program_*buffers usingmax_input_len. The author acknowledged the bug and stated they'd fix it ("swap tl.load(b_seq_len + 0) -> tl.num_programs(2), plus a wrapper guard"), but the code is unchanged. -
No backward-compatibility deprecation shims (UNRESOLVED):
modelopt/torch/kernels/__init__.pyhad its public symbols (IS_AVAILABLE,attention,attention_calibrate,register_triton_attention) removed with no__getattr__deprecation bridge. The oldmodelopt.torch.quantization.tritonpackage is also gone entirely. The PR correctly marks backward compat as ❌ but provides no migration path for downstream consumers.
| "attention_calibrate", | ||
| "register_triton_attention", | ||
| ] | ||
| """ModelOpt kernel library: common, quantization (conv, gemm), sparsity (attention, gemm).""" |
There was a problem hiding this comment.
Critical (unresolved from two prior reviews): The old public symbols (IS_AVAILABLE, attention, attention_calibrate, register_triton_attention) were removed with no deprecation bridge. Any downstream code doing from modelopt.torch.kernels import attention will break.
Please add a __getattr__ shim that re-exports these names from modelopt.torch.kernels.common.attention with a DeprecationWarning, or at minimum document the breaking change with migration guidance in the CHANGELOG.
There was a problem hiding this comment.
Mentioned it in the changelog
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Re-review of the kernel directory reorganization PR. Of the two critical issues from previous reviews:
-
Buffer overrun in calibrate.py — RESOLVED ✅:
tl.load(b_seq_len + 0)has been replaced withtl.num_programs(2), and the Python wrapper now has guards that reject non-uniform-length batches withNotImplementedError. Good fix. -
Backward compatibility deprecation shims — STILL UNRESOLVED ❌: The CHANGELOG (line 8) claims "these four names continue to work at the old location via a
__getattr__shim that emits aDeprecationWarning" but no such shim exists in the code.modelopt/torch/kernels/__init__.pyis just a docstring with no__getattr__. The oldmodelopt.torch.quantization.tritonpackage is entirely deleted with no redirect. This is actually worse than before — the documentation now actively lies about backward compatibility that doesn't exist.
| "attention_calibrate", | ||
| "register_triton_attention", | ||
| ] | ||
| """ModelOpt kernel library: common, quantization (conv, gemm), sparsity (attention, gemm).""" |
There was a problem hiding this comment.
Critical (unresolved from three prior reviews): The CHANGELOG claims "these four names continue to work at the old location via a __getattr__ shim that emits a DeprecationWarning" — but this file contains no such shim. from modelopt.torch.kernels import attention will raise ImportError at runtime.
Please either:
- Add the
__getattr__shim as described in the CHANGELOG (see the suggested code from previous reviews), or - Remove the claim from the CHANGELOG that a shim exists.
The modelopt.torch.quantization.triton package is also entirely gone with no redirect — any downstream code using that path will break silently.
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Both critical issues from previous reviews are now resolved:
-
Buffer overrun in calibrate.py ✅:
tl.load(b_seq_len + 0)replaced withtl.num_programs(2), and the Python wrapper guards against non-uniform-length batches withNotImplementedError. -
Backward compatibility documentation ✅: The CHANGELOG no longer falsely claims a
__getattr__shim exists. It now honestly states "there is no backwards-compatibility shim" and provides a clear migration table mapping old import paths to new ones.
The reorganization into modelopt/torch/kernels/{common,quantization,sparsity} is clean, all internal import sites are updated, the lazy import pattern in triton_fa.py properly handles the circular dependency, and tests are moved/updated to match.
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
### What does this PR do? Type of change: re-org code <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> <!-- Details about the change. --> ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing <!-- Mention how have you tested your change if applicable. --> ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ❌ We changed the import path <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ❌ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Calibration support for skip-softmax multi-threshold measurement in sparse attention. * N:M sparse softmax masking and helpers for sparsity-aware attention. * **Chores** * Reorganized and consolidated kernel/backends for quantization and sparsity to a unified kernels layout, updating tests and examples to match. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com> Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
4 participants
What does this PR do?
Type of change: re-org code
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅Additional Information
Summary by CodeRabbit
New Features
Chores