Add nvfp4 attention support for vLLM serving#1898
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
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:
📝 WalkthroughWalkthroughAdds split-K NVFP4 paged decode attention, V QDQ and dense-token-aware sparsity in Triton FlashAttention, quantized V-cache writes, vLLM sparse/NVFP4 runtime installation with FlashInfer support, and updated serving workers, documentation, and tests. ChangesNVFP4 quant-sparse attention kernels and vLLM integration
Estimated code review effort: 5 (Critical) | ~100 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1898 +/- ##
==========================================
- Coverage 77.87% 68.67% -9.21%
==========================================
Files 522 524 +2
Lines 58452 59132 +680
==========================================
- Hits 45522 40606 -4916
- Misses 12930 18526 +5596
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kaix-nv
force-pushed
the
kaix/sparse_attn_quant_compact
branch
5 times, most recently
from
62c1a25 to
5ef7dea
Compare
kaix-nv
requested review from
Edwardf0t1,
realAsma,
rohansjoshi and
sugunav14
and removed request for
Edwardf0t1,
realAsma,
rohansjoshi and
sugunav14
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Large NVFP4+sparse-attention vLLM serving PR (+3585/-617, 19 files). It extends the existing unified Triton FA kernel (adds symmetric V_QDQ next to the existing P_QDQ), adds a new split-K decode kernel, and adds a FlashInfer adapter alongside the existing FlashAttention adapter in the sparse-attn vLLM plugin. Design review: this is an extension of established in-repo subsystems (the unified triton_fa.py kernel and the attention_sparsity/plugins/vllm.py adapter) rather than a competing new system, and the new split-K decode kernel is justified by the split-local-P quantization contract in both the README and the module docstring. License headers on the new files match the canonical LICENSE_HEADER. Test coverage is extensive (GPU kernel tests, CPU wrapper tests, vLLM worker/adapter tests).
One correctness regression to fix, plus a couple of consistency questions.
Blocking / correctness:
fp4_kernel_hopper.py: removing theblock_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0)guard reintroduces a0/0 = NaNfor all-zero (or FP8-underflowing) blocks, sinceglobal_scale_safeonly guards the global scale, not the per-block scale. This kernel has its own inline quant (it does NOT callnvfp4_scalar_quant), so the paired change innvfp4_scalar_quant(which now correctly returns 0.0 for zero scale) does not cover it. This also diverges from thecuda_fp4CUDA kernel and the Python reference_py_fp4_fake_quant_ref, both of which still clamp tiny scales to 1.0 — thetest_vs_triton*cross-checks userandn*10inputs that rarely hit zero/underflow blocks, so CI likely won't catch the divergence.
I did not exhaustively verify the mixed decode/prefill FlashInfer split math or the split-K online-softmax combine on-device; those paths are GPU-only and the unit tests cover the routing/metadata contracts well, but numerical correctness relies on the B200/A6000 runs reported in the PR body.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🧹 Nitpick comments (1)
modelopt/torch/quantization/plugins/vllm.py (1)
34-48: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNarrow the caught exception to avoid masking real import failures.
Catching bare
ImportErroralso swallows failures caused by a genuinely brokenvllminstall (e.g. a transitive import error insidevllm.attention.layeritself), silently falling through to the next candidate and ultimately raising a generic "No supported vLLM Attention module was found" that hides the real root cause.♻️ Suggested narrowing
try: module = importlib.import_module(module_name) - except ImportError: + except ModuleNotFoundError: continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/quantization/plugins/vllm.py` around lines 34 - 48, The _import_attention_module helper is catching all ImportError cases, which can hide real failures from a broken vllm installation. Narrow the exception handling around importlib.import_module so only missing-module cases are treated as fallthrough, and let genuine import-time errors from modules like vllm.attention.layer surface instead of being replaced by the generic fallback ImportError.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/kernels/common/attention/triton_fa.py`:
- Line 284: Update the backward attention kernels in triton_fa.py so they mirror
the FP32 QK path used in forward for NVFP4. In `_attn_bwd_dq` and
`_attn_bwd_dkdv`, detect the same `Q_IS_FP32` condition and make the `tl.dot(q,
kT)` QK recomputation use the IEEE precision path, matching the forward
attention logic. Keep the change localized to the backward score recomputation
and preserve the existing `Q_IS_FP32` symbol as the selector.
In `@tests/unit/torch/kernels/common/attention/test_triton_fa.py`:
- Around line 16-23: Update the module docstring in the test_triton_fa module to
point to the correct GPU test locations. The existing reference in the top-level
docstring is wrong; change the path text near the descriptions of the Triton
kernels and wrappers so it names the GPU coverage under the test_triton_fa and
test_triton_fa_calibrate suites. Use the docstring in
tests/unit/torch/kernels/common/attention/test_triton_fa.py as the place to
edit.
---
Nitpick comments:
In `@modelopt/torch/quantization/plugins/vllm.py`:
- Around line 34-48: The _import_attention_module helper is catching all
ImportError cases, which can hide real failures from a broken vllm installation.
Narrow the exception handling around importlib.import_module so only
missing-module cases are treated as fallthrough, and let genuine import-time
errors from modules like vllm.attention.layer surface instead of being replaced
by the generic fallback ImportError.
🪄 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: Enterprise
Run ID: a948de07-7f82-46c0-af1f-be5e2d2c1642
📒 Files selected for processing (19)
examples/vllm_serve/README.mdexamples/vllm_serve/sparse_attn_worker.pyexamples/vllm_serve/vllm_serve_sparse_attn.pymodelopt/torch/kernels/common/attention/decode_attention.pymodelopt/torch/kernels/common/attention/triton_fa.pymodelopt/torch/kernels/quantization/attention/v_qdq.pymodelopt/torch/kernels/quantization/common/nvfp4_quant.pymodelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.pymodelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.pymodelopt/torch/quantization/plugins/vllm.pymodelopt/torch/sparsity/attention_sparsity/plugins/vllm.pytests/gpu/torch/kernels/common/attention/test_decode_attention.pytests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.pytests/gpu/torch/kernels/common/attention/test_triton_fa_paged.pytests/gpu/torch/quantization/test_tensor_quant_cuda.pytests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.pytests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.pytests/unit/torch/kernels/common/attention/test_triton_fa.py
💤 Files with no reviewable changes (2)
- modelopt/torch/kernels/quantization/gemm/fp4_kernel_hopper.py
- modelopt/torch/kernels/sparsity/attention/skip_softmax_helpers.py
| IS_CAUSAL: tl.constexpr, # Whether to apply causal mask | ||
| HEAD_DIM: tl.constexpr, # Actual head dimension (for d_mask) | ||
| STORE_LSE: tl.constexpr, # Whether to save LSE for backward pass | ||
| Q_IS_FP32: tl.constexpr, # Dynamic NVFP4 QDQ carrier uses FP32 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if any caller uses p_qdq="nvfp4" with fp32 q and requires_grad together.
rg -nP 'p_qdq\s*=\s*"nvfp4"' modelopt tests -C3
rg -nP 'requires_grad' tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py -C3Repository: NVIDIA/Model-Optimizer
Length of output: 3916
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant attention kernel sections.
sed -n '250,420p' modelopt/torch/kernels/common/attention/triton_fa.py
printf '\n---\n'
sed -n '540,780p' modelopt/torch/kernels/common/attention/triton_fa.py
printf '\n---\n'
sed -n '960,1105p' modelopt/torch/kernels/common/attention/triton_fa.py
# Inspect the tests that exercise p_qdq="nvfp4" and autograd behavior.
printf '\n=== test_triton_fa_p_qdq.py ===\n'
sed -n '1,220p' tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py
printf '\n---\n'
sed -n '220,420p' tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py
printf '\n=== test_decode_attention.py ===\n'
sed -n '160,250p' tests/gpu/torch/kernels/common/attention/test_decode_attention.pyRepository: NVIDIA/Model-Optimizer
Length of output: 42962
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the backward/autograd tests for P-QDQ and any explicit fp32-Q cases.
rg -n "requires_grad|backward\(|p_qdq=\"nvfp4\"|dtype=torch.float32" tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py -C 3
# Inspect the remainder of the P-QDQ test file around backward coverage.
sed -n '420,760p' tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py
# Check whether the wrapper blocks autograd for p_qdq/v_qdq modes anywhere in the implementation.
rg -n "NotImplementedError|autograd|v_qdq|p_qdq" modelopt/torch/kernels/common/attention/triton_fa.py -C 3Repository: NVIDIA/Model-Optimizer
Length of output: 19086
Mirror FP32 QK precision in backward
attention(..., p_qdq="nvfp4", q.dtype=torch.float32) is exercised in TestSoftmaxQdqBackward, but _attn_bwd_dq and _attn_bwd_dkdv still recompute scores = tl.dot(q, kT) * qk_scale without the input_precision="ieee" path used in forward. That leaves the forward/backward QK numerics inconsistent for the same FP32-Q NVFP4 case.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/kernels/common/attention/triton_fa.py` at line 284, Update the
backward attention kernels in triton_fa.py so they mirror the FP32 QK path used
in forward for NVFP4. In `_attn_bwd_dq` and `_attn_bwd_dkdv`, detect the same
`Q_IS_FP32` condition and make the `tl.dot(q, kT)` QK recomputation use the IEEE
precision path, matching the forward attention logic. Keep the change localized
to the backward score recomputation and preserve the existing `Q_IS_FP32` symbol
as the selector.
25ed74a to
37d1b17
Compare
|
I guess after you investigating the v_bmm_quantizer, you still decided to add a v_qdq into the kernel. Can you explain why? |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re-review of the NVFP4+sparse-attention vLLM serving PR (+3585/-616, 18 files, kaix-nv). Design review was applied previously: this extends established in-repo subsystems (unified triton_fa.py kernel + attention_sparsity/plugins/vllm.py adapter) rather than introducing a competing system, and the new split-K decode kernel is justified by the split-local-P quantization contract in the README/docstring. License headers on new files match the canonical NVIDIA Apache-2.0 header. No prompt-injection attempts in the untrusted content.
Critical comment status — RESOLVED:
- The blocking
fp4_kernel_hopper.pyregression is fixed: the per-block guardblock_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0)is present in the current file, so the0/0=NaNfor all-zero/underflowing blocks is guarded.nvfp4_scalar_quantnow correctly zeroes zero-scale blocks (return tl.where(zero_scale, 0.0, x_quant)), and both_p_qdq_nvfp4/_v_qdq_nvfp4route through it. New teststest_native_block_scale_underflows_to_zeroandtest_zero_block_scale_zeroes_blockcover this. Good.
Why nudge (remaining items for owner sign-off):
- 💬 Prior CodeRabbit comment (minor): asked to mirror the forward FP32-Q
input_precision="ieee"path into the backward kernels. Still unaddressed —_attn_bwd_dq/_attn_bwd_dkdvrecomputetl.dot(q, kT) * qk_scalewithout aQ_IS_FP32branch. Impact is limited: the serving path is inference-only (v_qdqrejects autograd), so this only affects forward/backward numeric consistency in thep_qdq="nvfp4"+ FP32-Q QAT training case. Low severity but worth an owner call. - 💬 sychen52 asked why a new
v_qdq.pywas added rather than reusing/renaming thep_qdqhelper. The current diff keeps them separate (_v_qdq_nvfp4usesabs()for signed V and adds thefake_quant_v_onwritecache-finalize kernel). Reasonable technical divergence, but the design question doesn't appear resolved in-code — human reviewer should confirm the author's rationale. - Minor doc nit (CodeRabbit):
tests/unit/torch/kernels/common/attention/test_triton_fa.pymodule docstring still points to the wrong GPU pathtests/gpu/torch/sparsity/attention_sparsity/test_triton_fa*.py; actual coverage lives undertests/gpu/torch/kernels/common/attention/. Not blocking. - Nitpick (CodeRabbit, unaddressed):
_import_attention_modulestill catchesImportError(could mask a genuinely broken vllm install);ModuleNotFoundErrorwould be tighter. Trivial. - Numerical correctness of the mixed decode/prefill FlashInfer split math and the split-K online-softmax combine is GPU-only; unit tests cover routing/metadata/cache-layout contracts well, but numerics rely on the B200/A6000 runs reported in the PR body.
Overall the blocking regression is fixed and test coverage is extensive; remaining items are minor/design-level plus GPU-only numerics, so a human with architectural context should give final sign-off.
In-kernel quant for V is necessary. V must be quantized in 16-token blocks along the sequence dimension. Because a block can span multiple decode steps or cache pages, the kernel quantizes completed blocks once after caching V and quantizes the remaining partial block on read. I agree that this should be consolidated. I'll move their thin wrappers into a neutrally named shared module (e.g., |
There was a problem hiding this comment.
We maintain our fake-quantized kernels in modelopt/torch/quantization/plugins/vllm. It looks to me that some of the contents of this file should be moved to either modelopt/torch/quantization/plugins/vllm or modelopt/torch/sparsity/plugins/vllm.
This way we can avoid workarounds like _install_quant_plans; We should use modelopt high-level APIs like mtq.quantize / mts.sparsify to automatically insert and initialize quantizers.
This is particularly important because our vLLM fakequant path is used in many places like Nemo-RL. In the current design, fake sparse kernels are not available outside of this example folder. Hence can we please move relevant contents from this folder to modelopt library?
There was a problem hiding this comment.
RB: Can you take a look at my above comment and propose a new design?
There was a problem hiding this comment.
🤖 Bot comment.
🐝 Looking at the current design and the surrounding vLLM plugin code now.
There was a problem hiding this comment.
🤖 Bot comment.
I agree with the direction in your comment. I would make the example worker a thin consumer of library APIs and move the reusable serving path into the vLLM plugins.
Proposed shape:
-
Move the vLLM sparse serving planner out of
examples/vllm_serve/sparse_attn_worker.pyintomodelopt/torch/sparsity/attention_sparsity/plugins/vllm.py. That module should own loadingsparse_attention_configfrom checkpoint metadata, matching layer names, validating supported vLLM backends, cloning the native impl, and attachingsparse_kw. The worker should call one helper such asmtsa.plugins.vllm.install_sparse_attention_from_checkpoint(model, hf_config, vllm_config=...)rather than open-coding_sparse_plans/_install_sparse_plans. -
Move the fixed NVFP4 attention quantizer setup into
modelopt/torch/quantization/plugins/vllm.py. In particular,_QuantVLLMAttentionshould own the Q/K/V/P BMM quantizer attributes, and the plugin should expose a reusable config/helper for dynamic block-16 NVFP4 attention. Then the worker can use the normal path, e.g.mtq.quantize(model, mtq.VLLM_NVFP4_ATTENTION_CFG, forward_loop=...), instead of manually callingQuantModuleRegistry.convert, addingp_bmm_quantizer, and runningset_quantizer_by_cfgin_install_quant_plans. -
Keep only runtime worker details in the example: unwrap the vLLM model, provide the dummy calibration loop if quantizing, disable compilation around
mtq.quantize, call the sparse install helper, and set vLLM-only runtime flags such ascascade_attn_enabled = Falseif the plugin helper reports that ModelOpt attention transforms are active. -
Add tests at the library level, not only under the example: one test that
mtq.quantizeinserts/enables Q/K/P/V attention quantizers for vLLM Attention, one test that the sparse plugin installs the backend-matched impl from exported checkpoint metadata, and one combined sparse+NVFP4 test proving the example worker no longer owns the conversion logic.
That gives Nemo-RL and other vLLM fakequant users an importable ModelOpt path, and the example becomes just a reference for wiring it into vLLM worker lifecycle.
There was a problem hiding this comment.
Agreed. I'll have a refactor for the orchestration (layer discovery, validation, and the module.impl swap) and it should live in the vLLM plugins, with a supported activation entry point so NeMo-RL and others don't have to copy the worker.
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
…prefill, share sparse layer validation, disable sparse cascade, require vLLM>=0.15, reject sm<8.9 Signed-off-by: Kai Xu <kaix@nvidia.com>
Signed-off-by: Kai Xu <kaix@nvidia.com>
kaix-nv
force-pushed
the
kaix/sparse_attn_quant_compact
branch
from
ab6dd3e to
f3021f9
Compare
Signed-off-by: Kai Xu <kaix@nvidia.com>
kaix-nv
force-pushed
the
kaix/sparse_attn_quant_compact
branch
from
f3021f9 to
c4e2146
Compare
kaix-nv
requested review from
a team,
ChenhanYu and
kevalmorabia97
and removed request for
a team
rohansjoshi
left a comment
rohansjoshi
left a comment
There was a problem hiding this comment.
LGTM for torch/sparsity
8 participants
What does this PR do?
Type of change: New feature
This PR adds an NVFP4+sparse attention serving path for vLLM and consolidates it with the existing checkpoint-driven sparse-attention integration.
NVFP4 attention
vLLM integration
SparseAttnWorkerfor checkpoint-driven sparse attention.QuantSparseAttnWorkerfor fixed NVFP4 Q/K/P/V plus optional checkpoint sparsity.Supported configurations are regular decoder self-attention with vLLM >= 0.15.0, FlashInfer or FlashAttention, fp16/bf16 model and KV cache, equal Q/K/V head dimensions divisible by 16, and DCP 1. The README documents unsupported features and CUDA-graph constraints.
Usage
Let vLLM select the backend for the model and platform. For example, Nemotron-H on Blackwell selects FlashInfer automatically.
If
<MODEL_PATH>/config.jsoncontainssparse_attention_config, the same worker also applies its N:M or skip-softmax settings. Otherwise, it runs NVFP4 attention only.Testing
Focused attention kernels:
PYTEST_VERSION=1 PYTHONPATH=$PWD python -m pytest -q \ tests/gpu/torch/kernels/common/attention/test_triton_fa_p_qdq.py \ tests/gpu/torch/kernels/common/attention/test_decode_attention.py \ tests/gpu/torch/kernels/common/attention/test_triton_fa_paged.pyFocused vLLM worker and adapter tests:
PYTHONPATH=$PWD python -m pytest -q \ tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py \ tests/gpu_vllm/torch/sparsity/attention_sparsity/test_quant_sparse_attn_worker.pyResult: 43 passed.
CPU import and launch-contract tests:
PYTHONPATH=$PWD python -m pytest -q \ tests/unit/torch/kernels/common/attention/test_triton_fa.pyResult: 9 passed.
GitHub Linux, Windows, multi-version, code-quality, documentation, and required unit checks pass.
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: N/A — no copied code or new dependency.Additional Information
The fixed attention recipe intentionally does not expose the integration branch's environment-variable matrix. Backend selection is automatic unless an explicit vLLM backend override is needed.
Summary by CodeRabbit