feat: SM120 NVFP4 SVDQuant Gemm in CuteDSL - #4420
Merged
Merged
Conversation
rosenrodt
requested review from
Aneureka,
StudyingShao,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 8, 2026 13:06
Contributor
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThis change adds SM120/SM121 CuTe DSL fused and unfused NVFP4 SVDQuant paths, BF16 smooth quantization, backend selection, CUDA 12.9 support, expanded benchmarks, kernel fusion, tests, and trace definitions. ChangesNVFP4 SVDQuant expansion
Estimated code review effort: 5 (Critical) | ~120 minutes Merge Risk: 🟡 Moderate · up to The fused SVDQuant path aliases correction storage with released pipeline buffers, but the current checks may not prove that each aliased layout fits its destination buffer. An invalid layout could cause out-of-bounds writes or incorrect results, so this capacity validation should be addressed or explicitly accepted before merging. Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Collaborator
Author
|
/bot run tests/gemm |
Collaborator
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (12)
flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py (7)
474-477: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReorder the docstring arguments to match the signature.
The signature order is
svdquant_d,svdquant_l1,svdquant_bias, thenepilogue_op. The docstring listsepilogue_opfirst.🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 474 - 477, Reorder the parameter descriptions in the docstring for the affected function so `svdquant_d`, `svdquant_l1`, and `svdquant_bias` appear before `epilogue_op`, matching the function signature; do not change their descriptions or implementation.
216-218: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueComplete the documented tile constraints.
can_implementalso rejectstile_n > 128, and rejectstile_n < 64whenswap_abis False. The docstring lists only the divisibility rules. Add the two missing bounds so the documentation matches the check.📝 Proposed doc update
- Tile shape constraints: * tile_m must be divisible by 64 * tile_n must be divisible by 16, must be <= 128, and must be >= 64 unless swap_ab is set * tile_k must be divisible by 64 (sf_vec_size=16) or 128 (sf_vec_size=32)🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 216 - 218, Update the docstring constraints for can_implement to document that tile_n must not exceed 128 and, when swap_ab is False, must be at least 64. Keep the existing divisibility requirements unchanged and place the conditional lower-bound constraint alongside the other tile_n rules.
66-74: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider narrowing the caught exception type.
Ruff reports S110 and BLE001 here. The blind
except Exceptionwithpassis intentional, and the comment explains it. If the set of failures is known, catch(ImportError, AttributeError, NotImplementedError)instead, and add a# noqa: S110only where a broader catch stays necessary. This keeps the fallback behavior and removes the lint noise.🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 66 - 74, Narrow the exception handling in _load_iket to catch ImportError, AttributeError, and NotImplementedError while preserving the fallback loop and _IketShim return. If broader handling is required for environment-specific failures, retain it only with a targeted noqa: S110 annotation rather than a blanket lint suppression.Source: Linters/SAST tools
1913-1935: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winApply
alpha_valueoutside the bounds check.Inside the bias branch,
alpha_value * acc_valueruns only whenn_coord < directC_mnl.shape[1]. Out-of-range lanes keep an unscaled accumulator, and Line 2069 then skips the alpha multiply for the whole vector. The TMA store at Lines 2130-2134 discards those lanes, so the current output is correct.The behavior is still fragile: the alpha contract now depends on the store path clipping the padded columns. Scale unconditionally and add only the bias under the bounds check.
♻️ Proposed refactor
acc_value = tRS_rAcc_slice[elem_idx] if cutlass.const_expr( svdquant_bias is not None ): + acc_value = alpha_value * acc_value coord = tRS_cD_slice[elem_idx] n_coord = ( tile_coord_mnl[1] * Int32(self.tile_shape_mnk[1]) + coord[1] ) if n_coord < Int32(directC_mnl.shape[1]): - acc_value = ( - alpha_value * acc_value - + svdquant_bias[(n_coord,)].to( - Float32 - ) - ) + acc_value += svdquant_bias[ + (n_coord,) + ].to(Float32) tRS_rD_slice[elem_idx] = acc_value🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 1913 - 1935, Update the accumulator assignment in the bias branch around tRS_rD_slice to multiply acc_value by alpha_value unconditionally, then add svdquant_bias only when n_coord is within directC_mnl.shape[1]. Keep the bounds check solely around bias access and preserve the existing out-of-range lane handling.
1442-1455: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the duplicate copy atoms and document the hardcoded transpose flag.
svdquant_copy_atom_aandsvdquant_copy_atom_bare constructed with identical arguments. The residual path derives theLdMatrix8x8x16bOptranspose flag from the operand layout; the correction hardcodesFalse. That is correct only because bothsvdquant_d(M, rank)andsvdquant_l1(N, rank)are row-major, so both are k-major. Record that dependency in a comment.♻️ Proposed refactor
- svdquant_copy_atom_a = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), - self.svdquant_dtype, - ) - svdquant_copy_atom_b = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), - self.svdquant_dtype, - ) + # Both correction operands are row-major (M, rank) / (N, rank), + # so both are k-major and need no ldmatrix transpose. + svdquant_copy_atom = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), + self.svdquant_dtype, + ) svdquant_smem_copy_a = cute.make_tiled_copy_A( - svdquant_copy_atom_a, svdquant_tiled_mma + svdquant_copy_atom, svdquant_tiled_mma ) svdquant_smem_copy_b = cute.make_tiled_copy_B( - svdquant_copy_atom_b, svdquant_tiled_mma + svdquant_copy_atom, svdquant_tiled_mma )🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 1442 - 1455, In the setup around svdquant_copy_atom_a and svdquant_copy_atom_b, create one shared LdMatrix8x8x16bOp copy atom with the existing False transpose flag and reuse it for both operands. Add a concise comment documenting that False is valid because svdquant_d (M, rank) and svdquant_l1 (N, rank) are row-major and therefore k-major.
488-500: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate
svdquant_l1alongsidesvdquant_d.The check covers only
svdquant_d.element_type.svdquant_l1feeds the same BF16 correction MMA at Lines 1755-1761 and is dereferenced unconditionally at Lines 551-552. If a caller suppliessvdquant_dwithoutsvdquant_l1, or with a non-BF16svdquant_l1, the failure surfaces deep inside CuTe layout construction. Extend the guard so the error names the offending operand.🛡️ Proposed fix
if cutlass.const_expr( self.svdquant_enabled and self.svdquant_dtype != BFloat16 ): raise TypeError( f"SVDQuant rank operands must be BF16, got {self.svdquant_dtype}." ) + if cutlass.const_expr(self.svdquant_enabled and svdquant_l1 is None): + raise ValueError("svdquant_l1 is required when svdquant_d is supplied.") + if cutlass.const_expr( + self.svdquant_enabled and svdquant_l1.element_type != BFloat16 + ): + raise TypeError( + f"svdquant_l1 must be BF16, got {svdquant_l1.element_type}." + )🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 488 - 500, Extend the SVDQuant validation in the constructor around svdquant_enabled to require svdquant_l1 whenever svdquant_d is provided, and validate svdquant_l1.element_type is BFloat16 alongside svdquant_d.element_type. Ensure invalid or missing svdquant_l1 raises the explicit TypeError at this guard and identifies the offending operand.
1736-1739: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse
self.svdquant_rankinstead of recomputingrank.
self.svdquant_rankis already set at Line 491 from the same expression. The same recomputation appears again at Line 2464. Use the stored value in both places so the rank has one source.🤖 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 `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py` around lines 1736 - 1739, In the rank-tile loops near the shown code and the corresponding loop near line 2464, stop recomputing rank with cute.size and reuse self.svdquant_rank in the cutlass.range_constexpr calculation. Preserve the existing tiling behavior while making self.svdquant_rank the single source of rank.flashinfer/quantization/kernels/nvfp4_quantize.py (1)
2283-2283: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueName the unused placeholder argument.
This argument satisfies the new kernel signature. The silu specialization sets
smooth_quant=False, so the kernel never reads it. Add a short comment so a later reader does not mistake it for a real smoothing scale.📝 Proposed comment
global_scale_tensor, + # Unused placeholder: the SwiGLU specialization sets smooth_quant=False. input[0, :k],🤖 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 `@flashinfer/quantization/kernels/nvfp4_quantize.py` at line 2283, Add a short explanatory comment next to the unused placeholder argument in the kernel invocation around the input slice, noting that it is required by the new signature and ignored because the SiLU specialization uses smooth_quant=False; do not change the argument or behavior.tests/gemm/test_nvfp4_svdquant_gemm.py (1)
55-58: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the
enable_iketassertions into their own test.The test name is
test_sm120_svdquant_can_implement_rejects_ragged_rank. Lines 55-58 assert theenable_iketconstructor default and override. That behavior is unrelated to rank validation. A separate test keeps the failure signal precise.♻️ Proposed split
def test_sm120_svdquant_can_implement_rejects_ragged_rank(): _skip_unless_sm120() import cutlass from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel, ) - assert not Sm120B12xBlockScaledDenseGemmKernel(16, (64, 64), (1, 1)).enable_iket - assert Sm120B12xBlockScaledDenseGemmKernel( - 16, (64, 64), (1, 1), enable_iket=True - ).enable_iket - common_args = (Add the removed assertions as a new test:
def test_sm120_svdquant_kernel_iket_flag_defaults_off(): _skip_unless_sm120() from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel, ) assert not Sm120B12xBlockScaledDenseGemmKernel(16, (64, 64), (1, 1)).enable_iket assert Sm120B12xBlockScaledDenseGemmKernel( 16, (64, 64), (1, 1), enable_iket=True ).enable_iket🤖 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 `@tests/gemm/test_nvfp4_svdquant_gemm.py` around lines 55 - 58, Move the default and override assertions for enable_iket out of test_sm120_svdquant_can_implement_rejects_ragged_rank into a new focused test named test_sm120_svdquant_kernel_iket_flag_defaults_off. Keep the SM120 skip guard and Sm120B12xBlockScaledDenseGemmKernel setup, preserving both assertions unchanged in the new test.flashinfer/gemm/gemm_svdquant.py (3)
440-455: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe
workspace_buffersdict duplicates an existing cache.
_get_cache_bufinflashinfer/utils.py(Lines 229-240) already memoizes on(name, device)and reuses the same tensor. The localworkspace_buffersdict adds a second cache holding a strong reference to the same tensor. Call_get_cache_bufdirectly.♻️ Proposed simplification
- workspace_buffer = workspace_buffers.get(a.device) - if workspace_buffer is None: - workspace_buffer = _get_cache_buf( - "mm_fp4_workspace", - DEFAULT_WORKSPACE_SIZE, - a.device, - ) - workspace_buffers[a.device] = workspace_buffer + workspace_buffer = _get_cache_buf( + "mm_fp4_workspace", + DEFAULT_WORKSPACE_SIZE, + a.device, + )Then remove the
workspace_buffersdeclaration at Line 440.🤖 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 `@flashinfer/gemm/gemm_svdquant.py` around lines 440 - 455, Remove the local workspace_buffers declaration and simplify Sm120Nvfp4SvdquantUnfusedRunner._fp4_inputs to call _get_cache_buf directly for the "mm_fp4_workspace" name and device, relying on its existing memoization and preserving the current workspace size and device behavior.
483-490: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared correction epilogue.
Lines 485-489 duplicate the correction and bias math in
_mm_nvfp4_svdquant_sm120_unfusedat Lines 314-318. The two paths must stay numerically identical, because_mm_nvfp4_svdquant_sm120_unfusedserves as the differential oracle for the tuned unfused runner. A future change to one site can silently diverge from the other.Extract a small helper and call it from both sites.
♻️ Proposed refactor
+def _apply_svdquant_bf16_correction(d, l1, alpha, bias, out): + """Add the rank-r BF16 LoRA-up correction and optional bias to ``out``.""" + correction = torch.mm(d, l1.T) + correction.mul_(alpha) + out.add_(correction) + if bias is not None: + out.add_(bias) + return out_, _, _, _, alpha, d, l1, bias, out = inputs fp4_runner(inputs=self._fp4_inputs(inputs), tactic=tactic) - correction = torch.mm(d, l1.T) - correction.mul_(alpha) - out.add_(correction) - if bias is not None: - out.add_(bias) - return out + return _apply_svdquant_bf16_correction(d, l1, alpha, bias, out)🤖 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 `@flashinfer/gemm/gemm_svdquant.py` around lines 483 - 490, Extract the shared correction-and-bias epilogue from `_mm_nvfp4_svdquant_sm120_unfused` and the shown fused path into a small helper. The helper should compute `torch.mm(d, l1.T)`, scale it by `alpha`, add it to `out`, and conditionally add `bias`; replace both duplicated blocks with calls to this helper so both paths remain numerically identical.
322-323: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the runner factory.
_sm120_nvfp4_svdquant_runnerdefines a new class and creates a new instance on everymm_nvfp4_svdquantcall at Line 797. Class creation per call adds Python overhead on the inference path. The runner holds no per-call state, so one instance perenable_pdlvalue is sufficient.Apply
@functools.cacheto the factory. The same applies to_sm120_nvfp4_svdquant_unfused_runnerat Line 426, which additionally builds an_b12x_gemm_fp4_runnerand a workspace dict on each call.♻️ Proposed refactor
+@functools.cache def _sm120_nvfp4_svdquant_runner(enable_pdl: bool): class Sm120Nvfp4SvdquantRunner(TunableRunner):As per coding guidelines: "Public Python APIs should use
@functools.cachewhere applicable to provide Python-level module caching."🤖 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 `@flashinfer/gemm/gemm_svdquant.py` around lines 322 - 323, Apply `@functools.cache` to both _sm120_nvfp4_svdquant_runner and _sm120_nvfp4_svdquant_unfused_runner, ensuring functools is imported. Preserve enable_pdl as the cache key so each factory returns one reusable runner instance per value and avoids repeated class, nested runner, and workspace creation.Source: Coding guidelines
🤖 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 `@benchmarks/bench_nvfp4_svdquant_gemm.py`:
- Around line 111-114: In the setup immediately after `_to_float8(w.T)`, add a
cheap assertion that `w_fp8` has the expected column-major layout required by
`bmm_fp8`. Keep this validation outside the timed region and leave the
quantization flow unchanged.
- Around line 364-368: Update the BF16 and FP8 baseline descriptions near the
benchmark output to explicitly state that both omit the LoRA correction,
including the LoRA-down GEMM and LoRA-up correction. Keep the existing FP8
qualification about excluded scales and quantization, and ensure the published
table makes these omitted terms clear.
In `@flashinfer/gemm/gemm_svdquant.py`:
- Around line 629-634: The comment above _heuristic_func_nvfp4_svdquant
incorrectly claims backend requirements are architecture-disjoint. Update it to
state that automatic dispatch preserves the backend_checks dictionary order,
selecting “cute-dsl” before “cute-dsl-unfused” when both are supported on
SM120/SM121, and that this ordering affects the resulting implementation
configuration.
- Around line 696-697: Update the mm_nvfp4_svdquant docstring to describe alpha
as a float32 device scalar with exactly one element (numel == 1), matching the
validation and alpha.reshape(1) requirement. Leave the existing validation
unchanged.
In `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py`:
- Around line 1017-1032: Update the SVDQuant compatibility check around
svdquant_copy_bytes to reject use_m1_non_tma_a=True, matching the existing
load-path validation in can_implement. Extend the existing byte-budget assertion
with a message identifying the SVDQuant copy-size mismatch, while preserving the
current equality check.
In `@flashinfer/quantization/kernels/nvfp4_quantize.py`:
- Around line 2067-2081: Update process_nvfp4_block_bfloat_smooth to validate
that input and pre_quant_scale data_ptr() values are 16-byte aligned before the
contiguous conversions, matching the existing guard in
silu_and_mul_nvfp4_quantize_cute_dsl. Raise a clear ValueError when either
tensor is misaligned, while preserving the current validation and processing
flow for aligned tensors.
- Around line 2084-2116: Update nvfp4_quantize_smooth_cute_dsl to return the
already allocated empty fp4_output and scale_output tensors when m == 0, before
invoking kernel_fn. Match the early-return behavior of
silu_and_mul_nvfp4_quantize_cute_dsl and preserve the existing kernel path for
nonzero m.
In `@tests/gemm/test_nvfp4_svdquant_gemm.py`:
- Around line 845-856: Pin the svdquant_linear invocation in the affected test
to the fused backend by passing backend="triton" explicitly. Keep the existing
inputs and torch_mm_calls assertion unchanged so the test no longer depends on
the process-global AutoTuner state or execution order.
- Around line 97-105: Update the SM120 test gating around _skip_unless_sm120 to
also skip when is_cute_dsl_available() is false, using the existing CuTe DSL
availability marker or compatible wrapper with module-level skipping. Preserve
the current compute-capability check and its skip message for non-SM120 GPUs.
In `@tests/trace/example.py`:
- Line 32: Update the documented trace filename entry in the test fixture to
remove the SF_A24576 segment, matching the generator output and committed golden
filename while preserving the remaining suffix components.
---
Nitpick comments:
In `@flashinfer/gemm/gemm_svdquant.py`:
- Around line 440-455: Remove the local workspace_buffers declaration and
simplify Sm120Nvfp4SvdquantUnfusedRunner._fp4_inputs to call _get_cache_buf
directly for the "mm_fp4_workspace" name and device, relying on its existing
memoization and preserving the current workspace size and device behavior.
- Around line 483-490: Extract the shared correction-and-bias epilogue from
`_mm_nvfp4_svdquant_sm120_unfused` and the shown fused path into a small helper.
The helper should compute `torch.mm(d, l1.T)`, scale it by `alpha`, add it to
`out`, and conditionally add `bias`; replace both duplicated blocks with calls
to this helper so both paths remain numerically identical.
- Around line 322-323: Apply `@functools.cache` to both
_sm120_nvfp4_svdquant_runner and _sm120_nvfp4_svdquant_unfused_runner, ensuring
functools is imported. Preserve enable_pdl as the cache key so each factory
returns one reusable runner instance per value and avoids repeated class, nested
runner, and workspace creation.
In `@flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py`:
- Around line 474-477: Reorder the parameter descriptions in the docstring for
the affected function so `svdquant_d`, `svdquant_l1`, and `svdquant_bias` appear
before `epilogue_op`, matching the function signature; do not change their
descriptions or implementation.
- Around line 216-218: Update the docstring constraints for can_implement to
document that tile_n must not exceed 128 and, when swap_ab is False, must be at
least 64. Keep the existing divisibility requirements unchanged and place the
conditional lower-bound constraint alongside the other tile_n rules.
- Around line 66-74: Narrow the exception handling in _load_iket to catch
ImportError, AttributeError, and NotImplementedError while preserving the
fallback loop and _IketShim return. If broader handling is required for
environment-specific failures, retain it only with a targeted noqa: S110
annotation rather than a blanket lint suppression.
- Around line 1913-1935: Update the accumulator assignment in the bias branch
around tRS_rD_slice to multiply acc_value by alpha_value unconditionally, then
add svdquant_bias only when n_coord is within directC_mnl.shape[1]. Keep the
bounds check solely around bias access and preserve the existing out-of-range
lane handling.
- Around line 1442-1455: In the setup around svdquant_copy_atom_a and
svdquant_copy_atom_b, create one shared LdMatrix8x8x16bOp copy atom with the
existing False transpose flag and reuse it for both operands. Add a concise
comment documenting that False is valid because svdquant_d (M, rank) and
svdquant_l1 (N, rank) are row-major and therefore k-major.
- Around line 488-500: Extend the SVDQuant validation in the constructor around
svdquant_enabled to require svdquant_l1 whenever svdquant_d is provided, and
validate svdquant_l1.element_type is BFloat16 alongside svdquant_d.element_type.
Ensure invalid or missing svdquant_l1 raises the explicit TypeError at this
guard and identifies the offending operand.
- Around line 1736-1739: In the rank-tile loops near the shown code and the
corresponding loop near line 2464, stop recomputing rank with cute.size and
reuse self.svdquant_rank in the cutlass.range_constexpr calculation. Preserve
the existing tiling behavior while making self.svdquant_rank the single source
of rank.
In `@flashinfer/quantization/kernels/nvfp4_quantize.py`:
- Line 2283: Add a short explanatory comment next to the unused placeholder
argument in the kernel invocation around the input slice, noting that it is
required by the new signature and ignored because the SiLU specialization uses
smooth_quant=False; do not change the argument or behavior.
In `@tests/gemm/test_nvfp4_svdquant_gemm.py`:
- Around line 55-58: Move the default and override assertions for enable_iket
out of test_sm120_svdquant_can_implement_rejects_ragged_rank into a new focused
test named test_sm120_svdquant_kernel_iket_flag_defaults_off. Keep the SM120
skip guard and Sm120B12xBlockScaledDenseGemmKernel setup, preserving both
assertions unchanged in the new test.
🪄 Autofix
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fe056c03-5032-4bf5-818d-09902718ce33
📥 Commits
Reviewing files that changed from the base of the PR and between 29add4e and 55de7ad0d8af45e255c81b0d7a64721af9a16e85.
📒 Files selected for processing (13)
benchmarks/bench_nvfp4_svdquant_gemm.pyflashinfer/gemm/gemm_svdquant.pyflashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.pyflashinfer/quantization/kernels/nvfp4_quantize.pyflashinfer/quantization/quantization_cute_dsl_utils.pyflashinfer/trace/templates/gemm.pytests/gemm/test_nvfp4_svdquant_gemm.pytests/trace/example.pytests/trace/fi_trace_out/gemm_bf16_N32_K3072.jsontests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_SF_B589824_rank32.jsontests/trace/fi_trace_out/linear_nvfp4_svdquant_N3072_K3072_K_packed1536_SF_B589824_rank32.jsontests/trace/fi_trace_out/quantize_nvfp4_smooth_N3072.jsontests/trace/test_fi_trace_template_consistency.py
Collaborator
|
[FAILED] Pipeline #61715761 — 16/18 executed test jobs passed Compared with nightly #61544260. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsCould not compare
Timeouts, infrastructure, or incomplete jobs
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/gemm/gemm_base.py (1)
6605-6609: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the heuristic documentation.
The docstring still says that CUDA 12 selects
cutlassand that the heuristic routes only tocudnnorcutlass. This branch now returnsb12xfirst for SM120 with CUDA 12.9 or newer. Update the documented selection rules.As per coding guidelines, performance-critical Python paths must document the rationale for special algorithmic choices and relevant alternatives.
🤖 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 `@flashinfer/gemm/gemm_base.py` around lines 6605 - 6609, The heuristic documentation for the backend-selection function must reflect that SM120 with NVFP4 on CUDA 12.9+ prioritizes b12x, while SM121 remains excluded from this automatic preference and continues using cudnn or cutlass. Document the rationale and relevant backend alternatives near the existing selection rules.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@flashinfer/gemm/gemm_base.py`:
- Around line 6605-6609: The heuristic documentation for the backend-selection
function must reflect that SM120 with NVFP4 on CUDA 12.9+ prioritizes b12x,
while SM121 remains excluded from this automatic preference and continues using
cudnn or cutlass. Document the rationale and relevant backend alternatives near
the existing selection rules.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 7ff59a71-1e6b-4af1-96ee-41b6933eeec3
📥 Commits
Reviewing files that changed from the base of the PR and between 55de7ad0d8af45e255c81b0d7a64721af9a16e85 and a0260b046dba8b2b8b4b962c9829e7e9f0300822.
📒 Files selected for processing (3)
flashinfer/gemm/gemm_base.pyflashinfer/gemm/gemm_svdquant.pytests/gemm/test_mm_fp4.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/gemm/gemm_svdquant.py
rosenrodt
force-pushed
the
sm120-svdquant-cute-dsl
branch
from
August 9, 2026 13:33
a0260b0 to
38471b2
Compare
Collaborator
Author
|
/bot run tests/gemm |
Collaborator
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json (1)
105-106: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMatch the public BF16 LoRA-down construction.
Line 105 computes
dasx_hat @ L2ᵀ. The public SVDQuant path first builds the BF16-smoothedL2ᵀmatrix and then runstorch.mm(x, ...). These forms round different operands in BF16, so the embedded trace can produce a different LoRA correction than the public API path. (docs.flashinfer.ai)Use the canonical operand order, or verify that both forms meet the same target-device tolerance.
Proposed fix
- d = torch.mm(x_hat, l2.t().contiguous().to(torch.bfloat16)) + l2t_smoothed = ( + (pqs.float()[:, None] * l2.float().t()) + .to(torch.bfloat16) + .contiguous() + ) + d = torch.mm(x, l2t_smoothed)🤖 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 `@tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json` around lines 105 - 106, Update _mm_nvfp4_svdquant_init so the LoRA-down tensor d follows the public BF16 construction: first form the BF16-smoothed L2 transpose, then call torch.mm with x as the left operand. Avoid computing d from x_hat @ l2.t() unless you verify both operand orders meet the target-device tolerance.
🤖 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.
Nitpick comments:
In `@tests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json`:
- Around line 105-106: Update _mm_nvfp4_svdquant_init so the LoRA-down tensor d
follows the public BF16 construction: first form the BF16-smoothed L2 transpose,
then call torch.mm with x as the left operand. Avoid computing d from x_hat @
l2.t() unless you verify both operand orders meet the target-device tolerance.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: b7d47c17-88a3-4dc4-bda0-049888f68ca9
📥 Commits
Reviewing files that changed from the base of the PR and between a0260b046dba8b2b8b4b962c9829e7e9f0300822 and 38471b291eaa7da8021cac7c345bf8838955455d.
📒 Files selected for processing (5)
flashinfer/gemm/gemm_svdquant.pyflashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.pyflashinfer/trace/templates/gemm.pytests/trace/example.pytests/trace/fi_trace_out/gemm_nvfp4_svdquant_N3072_K_packed1536_rank32.json
🚧 Files skipped from review as they are similar to previous changes (4)
- tests/trace/example.py
- flashinfer/trace/templates/gemm.py
- flashinfer/gemm/gemm_svdquant.py
- flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py
Collaborator
Collaborator
|
[FAILED] Pipeline #63399423 — 7/16 executed test jobs passed Compared with nightly #63265553. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsPre-existing failures
Timeouts, infrastructure, or incomplete jobs
|
dhiraj113
approved these changes
Aug 20, 2026
Extend the existing NVFP4 SVDQuant API to SM120/SM121 with a fused CuTe DSL implementation while retaining the compositional backend as a differential oracle. Changes - Add fused residual, BF16 low-rank correction, bias, and automatic SM120 dispatch - Preserve explicit cute-dsl-unfused behavior and the existing SM100 CUTLASS contract - Add focused GEMM, linear, trace-template, and generated-trace coverage Validation - Preserve the original focused SM120, SM100, and trace acceptance changes - Run applicable repository commit hooks during history reconstruction Result - SM120/SM121 callers use the existing public SVDQuant APIs - Fused and unfused implementations are independently selectable for validation
Reduce end-to-end SVDQuant linear overhead and align SM120 behavior with the established SM100 alpha and epilogue contract. Changes - Fuse activation smoothing into NVFP4 quantization and vectorize BF16 correction staging - Add implementation and tactic autotuning with flattened storage and cache contracts - Extend dense BF16 baselines, benchmark reporting, trace integration, and focused coverage Validation - Preserve the original complete-linear, boundary-tactic, quantizer, and trace acceptance changes - Run applicable repository commit hooks during history reconstruction Result - Complete SM120 SVDQuant uses the fused quantizer and tensor-core low-rank path - Backend selection and alpha normalization match the public cross-architecture contract
Make the production LoRA-up path share the residual mainloop stage ring and tune its tile geometry without reserving a separate shared-memory pipeline. Changes - Match BF16 correction tiles to released NVFP4 A/B stage capacity and alias their storage - Reuse the mainloop producer-consumer barriers and autotune M/N/K tile choices - Use torch MM for LoRA-down and report the final cold-L2 benchmark comparisons Validation - Preserve the original deadlock, correctness, custom-shape, and RTX PRO 6000 performance acceptance changes - Run applicable repository commit hooks during history reconstruction Result - The correction pipeline avoids redundant barriers and shared-memory reservations - Fused execution remains faster than the unfused oracle across the recorded production shapes
Expose balanced warp-phase ranges for profiling while keeping diagnostic operations out of ordinary production specializations. Changes - Instrument main load, LoRA load, main MMA, LoRA MMA, and epilogue phases - Standardize the tile_k admission term - Guard all ranges with a default-off constexpr and separate _iket0/_iket1 cache identities Validation - Preserve the original five-phase trace, default-off correctness, and same-node performance A/B changes - Run applicable repository commit hooks during history reconstruction Result - IKET profiling remains explicitly available for diagnosis - Default SM120 kernels contain no active tracing operations and show no measurable regression
Permit the validated CuTe DSL SVDQuant and generic B12x NVFP4 paths on CUDA 12.9 while preserving the existing SM100 CUTLASS dispatch. Incorporate focused review corrections to backend ordering, scalar-alpha documentation, kernel compatibility checks, and SVDQuant trace axes. Changes - lower SM120 SVDQuant and B12x admission to CUDA 12.9 - prefer B12x for automatic SM120 FP4 dispatch on CUDA 12.9 - document fused-first SM120 auto-dispatch and exact scalar-alpha contract - reject incompatible M=1 non-TMA A loads for SVDQuant and diagnose copy-size mismatches - model SF_A and SF_B as derived trace variables while retaining K_packed as the K specialization axis - regenerate the SVDQuant trace golden and align its documented filename Validation - pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py (23 passed, 47 skipped on CUDA 12.9 RTX 5080) - focused generic B12x GPU cases (6 passed on CUDA 12.9 RTX 5080) - pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py -k 'sm120 or backend_arch_support' -x (22 passed, 48 deselected on RTX 5080, CUDA 13.2, CUTLASS DSL 4.5.2) - pytest -q tests/trace/test_fi_trace_template_consistency.py (465 passed) - pytest -q tests/trace/test_fi_trace.py (32 passed) - SVDQuant trace initializer executed on RTX 5080 with expected packed and scale shapes - pre-commit run -a - git diff --check Result - SM120 SVDQuant and B12x execute successfully with CUDA 12.9 - SM100 SVDQuant remains on the unchanged architecture-gated CUTLASS path - SVDQuant trace names contain only independent constant specialization axes - reviewed compatibility assumptions now fail early with actionable diagnostics
Preserve the historical alpha contract while normalizing every backend launch to a one-element scalar view. Remove the dead SM120 prefetch tactic dimension and harden cache, trace, quantization, graph, accuracy, and benchmark behavior. Changes - align and empty-guard smooth quantization inputs - separate SM100 prefetch from SM120 tactics and strengthen cache identities - add pooled-alpha, CUDA Graph, accuracy, trace, and benchmark coverage Validation - pre-commit run -a - python -m pytest -v tests/trace/ - python -m pytest -q tests/jit/test_cute_dsl_cache.py - python -m pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py Result - 1197 trace tests passed with 182 skipped - 34 SM120 SVDQuant tests passed with 48 skipped - corrected RTX PRO 6000 cold-L2 CUDA Graph benchmark completed 12 rows
Preserve the SVDQuant-tested NVFP4 SFA mode grouping while retaining upstream's trailing-K normalization for the new MXFP4 b12x path. Changes - select SFA rank normalization by scale-vector format - document why NVFP4 and MXFP4 collapse different modes Validation - pre-commit run -a - python -m pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py - focused b12x ragged-K, short-K, NVFP4, MXFP4, and MXFP4-alpha checks Result - 34 SVDQuant cases passed with 48 skipped - 6 upstream b12x regression cases and 6 representative format cases passed
rosenrodt
force-pushed
the
sm120-svdquant-cute-dsl
branch
from
August 20, 2026 03:33
f6051bb to
0373d0c
Compare
Collaborator
Author
|
/bot run tests/gemm |
Collaborator
Collaborator
Author
|
@flashinfer-bot run |
Collaborator
|
[SUCCESS] Pipeline #63609952: 16/16 executed test jobs passed |
Collaborator
Author
|
/bot run tests/moe |
Collaborator
Collaborator
|
[FAILED] Pipeline #63648648 — 10/16 executed test jobs passed Compared with nightly #63457917. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 4/6 passed
Failure detailsPre-existing failures
Timeouts, infrastructure, or incomplete jobs
|
Collaborator
|
@flashinfer-bot run |
Collaborator
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #63729891 — 8/16 executed test jobs passed Compared with nightly #63648836. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPre-existing failures
Could not compare
Timeouts, infrastructure, or incomplete jobs
|
georgeliu95
added a commit
to georgeliu95/flashinfer
that referenced
this pull request
Sep 9, 2026
Adds a CUTLASS C++ implementation of NVFP4 SVDQuant for SM120, registered as `backend="cutlass-sm120"`. It does not replace the CuTeDSL backend from flashinfer-ai#4420. That one remains the SM120 `auto` default; this one is reachable only by explicit name, and `_heuristic_func_nvfp4_svdquant` excludes it from `auto`. What it adds over the existing path is a fused prefix. Upstream's `svdquant_linear` computes the LoRA-down projection with a `torch.mm` after a separate quantize; this backend can fuse the smooth-quantize (K1) and the LoRA-down (K2) into one launch and hand the result straight to the residual GEMM with its LoRA-up epilogue (K3). Whether fusing wins is shape-dependent, so both arrangements are offered to the autotuner and it picks per shape. Measured against `backend="cute-dsl"` over 71 model shapes from Wan2.1/2.2, Qwen-Image and MiniMax-H3, under CUDA-graph replay with a cold L2: faster on essentially every shape, on both of the SM120 parts it was run on. Accuracy is checked as SQNR against `cute-dsl-unfused` rather than bitwise -- these are independent implementations, so equality is the wrong contract. Nothing here pairs a shape to a kernel. Which producer geometry a shape can run is computed from the launch constraints the kernels already assert, not read from a table, so a shape nobody measured reaches the autotuner as a candidate instead of being excluded by a list; and which candidate wins is decided by measurement on the device in front of it. That matters because the advantage is not uniform across parts: run on two SM120 SKUs, some shapes come out ahead on one and behind on the other. A table of where something was faster does not transport, so there is none -- not for producer geometry, not for the cuBLASLt prefix (it is a producer family the tactic selects), and not for admission. The LoRA rank is a build parameter of the module rather than a constant. Rank 32 leaves the define unset and keeps its own JIT cache identity; rank 64 is a separate build, served through the public `mm_nvfp4_svdquant` entry, which reads the rank from `d.shape[1]`. A tile that cannot stage the requested rank still compiles at 32 and is refused by `can_implement`, so the autotuner sees a smaller candidate plane rather than a build failure. Verified on two SM120 parts: independent cold JIT builds on each, and 701 tests passing with none failing. Four bitwise-equivalence tests for the diagnostic JIT variants are deselected in that run -- each builds an entire extra module, and the variants they cover are scaffolding that should come out before this merges. Co-Authored-By: Claude <noreply@anthropic.com>
georgeliu95
added a commit
to georgeliu95/flashinfer
that referenced
this pull request
Sep 15, 2026
Adds a CUTLASS C++ implementation of NVFP4 SVDQuant for SM120, registered as `backend="cutlass-sm120"`. It does not replace the CuTeDSL backend from flashinfer-ai#4420. That one remains the SM120 `auto` default; this one is reachable only by explicit name, and `_heuristic_func_nvfp4_svdquant` excludes it from `auto`. What it adds over the existing path is a fused prefix. Upstream's `svdquant_linear` computes the LoRA-down projection with a `torch.mm` after a separate quantize; this backend can fuse the smooth-quantize (K1) and the LoRA-down (K2) into one launch and hand the result straight to the residual GEMM with its LoRA-up epilogue (K3). Whether fusing wins is shape-dependent, so both arrangements are offered to the autotuner and it picks per shape. Measured against `backend="cute-dsl"` over 71 model shapes from Wan2.1/2.2, Qwen-Image and MiniMax-H3, under CUDA-graph replay with a cold L2: faster on essentially every shape, on both of the SM120 parts it was run on. Accuracy is checked as SQNR against `cute-dsl-unfused` rather than bitwise -- these are independent implementations, so equality is the wrong contract. Nothing here pairs a shape to a kernel. Which producer geometry a shape can run is computed from the launch constraints the kernels already assert, not read from a table, so a shape nobody measured reaches the autotuner as a candidate instead of being excluded by a list; and which candidate wins is decided by measurement on the device in front of it. That matters because the advantage is not uniform across parts: run on two SM120 SKUs, some shapes come out ahead on one and behind on the other. A table of where something was faster does not transport, so there is none -- not for producer geometry, not for the cuBLASLt prefix (it is a producer family the tactic selects), and not for admission. The LoRA rank is a build parameter of the module rather than a constant. Rank 32 leaves the define unset and keeps its own JIT cache identity; rank 64 is a separate build, served through the public `mm_nvfp4_svdquant` entry, which reads the rank from `d.shape[1]`. A tile that cannot stage the requested rank still compiles at 32 and is refused by `can_implement`, so the autotuner sees a smaller candidate plane rather than a build failure. Verified on two SM120 parts: independent cold JIT builds on each, and 701 tests passing with none failing. Four bitwise-equivalence tests for the diagnostic JIT variants are deselected in that run -- each builds an entire extra module, and the variants they cover are scaffolding that should come out before this merges. Co-Authored-By: Claude <noreply@anthropic.com>
11 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
This PR adds SM120/SM121 CuTe DSL support for NVFP4 SVDQuant GEMM under the
existing FlashInfer API. It provides a true fused residual-plus-BF16 LoRA-up
kernel and retains the unfused composition as an explicit correctness oracle.
The complete path also:
mm_nvfp4_svdquantandsvdquant_linearwithout changing API contracts;LoRA-down on BF16 tensor cores;
stages and uses the same producer/consumer pipeline;
tactics;
and generated goldens; and
ordinary and instrumented specializations have distinct cache identities.
Performance
Reportable performance is from RTX PRO 6000 only. The tables below use rank
32, CUDA Graph replay, and warm-L2 timing. They report complete
svdquant_linearlatency, gains over hidden BF16 and per-tensor FP8 baselines,and the fused/unfused correction comparison.
residual_fp4/fusedcompareslatency; TFLOPS and microseconds share one cell as
TFLOPS / us.Reproduce the benchmark script's complete default Qwen3-image sweep from the
repository root on an idle, exclusive RTX PRO 6000. The defaults cover three
(N,K)pairs, fourMvalues, and rank 32:svdquant_linearus🔍 Related Issues
N/A.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request,
please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).python -m pytest -q tests/gemm/test_nvfp4_svdquant_gemm.py -k "sm120"Summary by CodeRabbit
New Features
Bug Fixes