Add SM120 NVFP4 attention JIT path - #3640
Merged
Merged
Conversation
tiffany940107
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
June 15, 2026 11:46
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:
📝 WalkthroughWalkthroughAdds SM120-specific NVFP4 attention support across CUDA kernels, bindings, JIT wiring, Python APIs, validation tests, and benchmarking tooling. ChangesSM120 NVFP4 Attention
Sequence Diagram(s)sequenceDiagram
participant Bench as benchmark/test
participant API as flashinfer.nvfp4_attention_sm120
participant Quant as quantize backend
participant Fwd as fwd backend
Bench->>API: nvfp4_attention_sm120_quantize_qkv(q,k,v)
API->>Quant: quantize packed q/k/v + scales
Quant-->>API: q_fp4, k_fp4, v_fp4_t, scales, qk_correction
Bench->>API: nvfp4_attention_sm120_fwd(...)
API->>Fwd: launch forward kernel
Fwd-->>API: out, lse
API-->>Bench: out, lse
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
Suggested labels
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces SM120 NVFP4 attention support to FlashInfer, adding benchmarks, JIT compilation modules, quantization kernels, and various attention mainloop implementations (non-warp-specialized, split-Q, and cross-tile). The review identified several critical correctness and memory safety issues: an incomplete softmax row reduction causing accuracy loss, a potential out-of-bounds memory read when per_block_mean is false, unaligned shared memory access in the quantization kernel, and shadowed empty quantization functions in PVGemmComputer that silently disable quantization.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![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: 12
🧹 Nitpick comments (8)
benchmarks/bench_nvfp4_attention_sm120.py (1)
281-282: ⚡ Quick winAlign benchmark timer import with the repo benchmark API contract.
Please import/use
bench_gpu_timethroughflashinfer.testing.bench_gpu_time()to match the benchmark guideline contract.As per coding guidelines, files under
benchmarks/**/*.pyshould useflashinfer.testing.bench_gpu_time()for GPU kernel benchmarking.🤖 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 `@benchmarks/bench_nvfp4_attention_sm120.py` around lines 281 - 282, The current import statement imports bench_gpu_time directly from flashinfer.testing.utils, but the benchmark API contract requires using it through the flashinfer.testing module namespace. Change the import statement to import flashinfer.testing as a module (or import from flashinfer.testing directly), then update all calls to bench_gpu_time throughout the file to use the fully qualified name flashinfer.testing.bench_gpu_time() to comply with the benchmark guideline contract.Source: Coding guidelines
tests/attention/test_nvfp4_attention_sm120.py (1)
37-40: ⚡ Quick winUse
flashinfer.utilsarchitecture helpers for SM gating in tests.This file currently mixes manual capability checks and CUDA availability skip logic; please switch to the repository’s
flashinfer.utilsarchitecture helper pattern for consistent test gating.As per coding guidelines,
tests/**/*.pyshould useflashinfer.utilsarchitecture helper functions for CUDA-arch-based skips.Also applies to: 136-137
🤖 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/attention/test_nvfp4_attention_sm120.py` around lines 37 - 40, The function `_require_sm120()` and related code at lines 136-137 use manual CUDA capability checks with torch.cuda.get_device_capability() instead of the repository's standardized pattern. Replace the manual capability check logic in `_require_sm120()` with the appropriate `flashinfer.utils` architecture helper function, and apply the same pattern to the code at lines 136-137 to ensure consistent SM-based test gating across the file according to repository guidelines.Source: Coding guidelines
include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h (1)
1-1: ⚡ Quick winMissing copyright header.
This file lacks the standard copyright header present in other files in this PR (e.g., "Copyright (c) 2025 by SageAttention team. Licensed under the Apache License, Version 2.0").
Suggested fix
+/* + * Copyright (c) 2025 by SageAttention team. + * Licensed under the Apache License, Version 2.0 + */ + `#pragma` once🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h` at line 1, The file cute_extension.h is missing the standard copyright header that should appear at the top of the file. Add the copyright header before the `#pragma` once directive to match the format used in other files in this PR, including the copyright notice "Copyright (c) 2025 by SageAttention team" and the Apache License 2.0 reference.include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h (1)
64-84: 💤 Low valueUnused parameter
tiled_mmain layout deduction functions.The
tiled_mmaparameter is declared but unused indeduce_smem_layoutSFQ,deduce_smem_layoutSFKV, anddeduce_smem_layoutSFVt. Consider using[[maybe_unused]]attribute or removing if not needed for API consistency.Suggested fix
template <class TiledMma, class TileShape_MNK> - CUTE_HOST_DEVICE static constexpr auto deduce_smem_layoutSFQ(TiledMma tiled_mma, + CUTE_HOST_DEVICE static constexpr auto deduce_smem_layoutSFQ([[maybe_unused]] TiledMma tiled_mma, TileShape_MNK tileshape_mnk) {🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h` around lines 64 - 84, The parameter `tiled_mma` is declared but unused in the functions `deduce_smem_layoutSFQ`, `deduce_smem_layoutSFKV`, and `deduce_smem_layoutSFVt`. To suppress compiler warnings about unused parameters, add the `[[maybe_unused]]` attribute to the `tiled_mma` parameter in each of these three functions. This signals to the compiler that the parameter is intentionally unused, likely for API consistency purposes.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh (1)
226-239: 💤 Low valueMinor: Consider using
thread_idxparameter instead ofthreadIdx.x.The
add_delta_slambda usesthreadIdx.xdirectly at line 229, while the rest of the function uses thethread_idxparameter. This is functionally equivalent but inconsistent with the rest of the code.auto add_delta_s = [&](auto& acc) { auto tSsDS_stage = recast<float4>(sDS(_, _, smem_pipe_read_k.index())); auto acc_float4 = recast<float4>(acc); - int quad_id = (threadIdx.x % 4) * 2; + int quad_id = (thread_idx % 4) * 2; for (int i = 0; i < 4; i++) {🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh` around lines 226 - 239, The add_delta_s lambda function uses threadIdx.x directly to calculate quad_id instead of using the thread_idx parameter that is available and used consistently elsewhere in the function. Replace the threadIdx.x reference in the quad_id calculation line with the thread_idx parameter to maintain consistency with the rest of the codebase.include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h (1)
22-22: 💤 Low valueUnused
<vector>header.The
<vector>header is included but not used anywhere in this file. The structs only contain scalar fields and raw pointers.Suggested fix
`#include` <cuda.h> `#include` <cstdint> -#include <vector> `#include` "cutlass/fast_math.h"🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h` at line 22, Remove the unused `#include` <vector> header from the file. This header is included but not needed since all struct definitions in this file only contain scalar fields and raw pointers, with no use of std::vector or other vector-related functionality.include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h (1)
116-149: 💤 Low value
tile_count_semaphoreis stored but never used.
DynamicPersistentTileSchedulerstorestile_count_semaphoreinParamsbut neitherget_initial_work()norget_next_work()reference it. If dynamic scheduling is intended (e.g., atomically fetching tile indices), the semaphore should be used; otherwise, consider removing it or documenting why it exists for future use.🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h` around lines 116 - 149, The tile_count_semaphore field is stored in the Params struct within DynamicPersistentTileScheduler but is never referenced in either the get_initial_work() or get_next_work() methods. Resolve this by either removing the tile_count_semaphore parameter from both the Arguments struct and Params struct along with its assignment in to_underlying_arguments, or if dynamic scheduling using this semaphore is intended, implement the logic to use it atomically within get_initial_work() and get_next_work() to fetch tile indices dynamically, or add a clear comment explaining why it is stored for potential future use.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh (1)
11-14: ⚡ Quick win
using namespace cute;at namespace scope in headers pollutes all includer namespaces. This directive imports all CUTE symbols into every translation unit that includes these headers, potentially causing name collisions with user code or other libraries.
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh#L11-L14: Removeusing namespace cute;or move into function bodies; qualify CUTE names withcute::.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh#L11-L14: Same fix needed.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh#L11-L14: Same fix needed.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh#L11-L14: Same fix needed.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh#L15-L18: Same fix needed.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh#L11-L14: Same fix needed.🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh` around lines 11 - 14, Remove the `using namespace cute;` directive from all affected header files to prevent namespace pollution. In include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh (lines 11-14), remove the `using namespace cute;` statement. Apply the same fix to include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh (lines 11-14), include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh (lines 11-14), include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh (lines 11-14), include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh (lines 15-18), and include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh (lines 11-14). After removing these directives, qualify any CUTE symbols used in these files with the `cute::` prefix to maintain correctness while eliminating namespace scope pollution.
🤖 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 `@csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu`:
- Around line 266-269: The shape validation for qk_correction tensor in the
nvfp4_attention_sm120_binding.cu file is incomplete. Currently, the validation
checks dimensions 0, 1, and 3 using TVM_FFI_ICHECK_EQ, but dimension 2 is
missing. Add a validation check for qk_correction.size(2) which should
correspond to seqlen_s based on the expected tensor shape of [batch, heads,
seqlen_s, seqlen]. Determine the appropriate variable to validate dimension 2
against (likely related to per_block_mean or a derived seqlen_s value) and
insert the TVM_FFI_ICHECK_EQ call in the correct sequence with the other
dimension checks.
In `@flashinfer/jit/nvfp4_attention_sm120.py`:
- Around line 115-131: The function gen_nvfp4_attention_sm120_module() is
missing the required `@functools.cache` decorator. Add the `@functools.cache`
decorator from the functools module above the function definition to ensure the
JIT module generator result is cached, as required by the coding guidelines for
JIT module generator functions.
In `@flashinfer/nvfp4_attention_sm120.py`:
- Around line 123-124: Add the `trace=` argument to both `@flashinfer_api`
decorators on the public tensor I/O APIs. Specifically, add `trace=` to the
decorator for the function nvfp4_attention_sm120_quantize_qkv (around line 123)
and to the decorator for the other public API function at line 290. Both
decorators handle tensor inputs/outputs and per coding guidelines require the
`trace=` argument for proper TraceTemplate wiring.
- Around line 124-185: Add the supported_compute_capability([120, 121])
decorator to the nvfp4_attention_sm120_quantize_qkv function to gate this
SM120-specific API at the Python level. This decorator should be placed directly
above the function definition to ensure the function only runs on compatible
hardware. Apply the same decorator pattern to other SM120-specific public APIs
in the file (around lines 291-407) that similarly lack capability gating.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h`:
- Around line 198-202: The error message in the CUTE_INVALID_CONTROL_PATH call
contains an incorrect struct name. Update the string in the error message from
`SM120_16x8x64_TN_VS` to `SM120_16x32x64_TN_VS_NVFP4` to accurately reflect the
actual struct name where this code block is located.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h`:
- Around line 41-53: The HEADDIM_SWITCH macro is missing a fallback path that
handles cases where HEADDIM is not one of the supported values (64, 128, 256).
When HEADDIM matches none of these values, the lambda function will fall through
without returning, causing undefined behavior. Add a final else clause before
the closing brace of the macro that throws a runtime error or otherwise handles
unsupported HEADDIM values to ensure the lambda always returns a value for any
input.
- Around line 16-39: The PREC_SWITCH macro does not have a fallback path for
invalid PRECTYPE values, causing undefined behavior when PRECTYPE is not in {1,
2, 3, 4}. Add an else clause after the final else if (PRECTYPE == 4) branch that
either uses static_assert(false, "Invalid PRECTYPE") to catch invalid usage at
compile time, or uses __builtin_unreachable() to signal the compiler that this
path should never be reached. This ensures the macro always returns a value for
any input.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh`:
- Around line 147-154: The private static methods packed_float_to_ue4m3 and
packed_float_to_e2m1 in the pv_gemm.cuh struct have empty bodies that do not
initialize the out parameter, causing undefined behavior when these methods are
called in quantize_p() to populate SFP_uint32_view and tOrP_uint32_view which
are then used in the GEMM computation. Fix this by either including the
fp4_convert.cuh header file and delegating these method calls to the
corresponding global function implementations that use PTX inline assembly, or
by providing complete inline implementations of these conversion methods with
the same logic as the global versions in fp4_convert.cuh.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh`:
- Around line 142-149: The `copy_with_predicate` method has an empty body, which
means when `store_zero()` calls it to write data to global memory with
predicate-based bounds checking, nothing actually gets written. Implement the
actual predicated copy logic inside `copy_with_predicate` that uses the provided
tiled_copy, src, dst, coord, pred, and max_m parameters to perform the
bounds-checked copy operation. Alternatively, refactor to reuse the existing
`sage::copy_with_bounds_check` utility from `epilogue.cuh` to avoid duplicating
this logic.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.h`:
- Around line 1-4: Add the standard Apache 2.0 license header at the top of the
file attention_kernel_nonws.h, before the `#pragma` once directive. Copy the same
license header format used in other files in this PR such as attention_kernel.h
and launcher.h to maintain consistency across the codebase.
In `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h`:
- Around line 151-210: The member variables m_block_divmod and head_divmod in
the StaticPersistentTileSchedulerOld class are stored as const references, which
can become dangling if the FastDivmod objects passed to the constructor are
temporary or go out of scope. Store these members by value instead of by const
reference. Change the member variable declarations from const reference syntax
(const &m_block_divmod, &head_divmod) to value syntax (cutlass::FastDivmod
m_block_divmod, head_divmod), and update the constructor parameter list and
member initialization list accordingly to remove the reference syntax while
keeping the initialization logic the same.
In `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh`:
- Around line 35-55: The vectorized PTX instructions add.f32x2, mul.f32x2, and
fma.rn.f32x2 used in the add, mul, and fma functions decompose into scalar
operations on SM120 and do not provide actual vectorization benefits. Replace
these vector PTX instructions with scalar alternatives by removing the .f32x2
suffix (changing to add.f32, mul.f32, and fma.rn.f32), or implement scalar
fallback versions that perform component-wise float operations on the float2
structure elements (accessing c.x and c.y, a.x and a.y, b.x and b.y separately)
instead of using inline assembly for the vector instructions.
---
Nitpick comments:
In `@benchmarks/bench_nvfp4_attention_sm120.py`:
- Around line 281-282: The current import statement imports bench_gpu_time
directly from flashinfer.testing.utils, but the benchmark API contract requires
using it through the flashinfer.testing module namespace. Change the import
statement to import flashinfer.testing as a module (or import from
flashinfer.testing directly), then update all calls to bench_gpu_time throughout
the file to use the fully qualified name flashinfer.testing.bench_gpu_time() to
comply with the benchmark guideline contract.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h`:
- Line 1: The file cute_extension.h is missing the standard copyright header
that should appear at the top of the file. Add the copyright header before the
`#pragma` once directive to match the format used in other files in this PR,
including the copyright notice "Copyright (c) 2025 by SageAttention team" and
the Apache License 2.0 reference.
In `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h`:
- Line 22: Remove the unused `#include` <vector> header from the file. This header
is included but not needed since all struct definitions in this file only
contain scalar fields and raw pointers, with no use of std::vector or other
vector-related functionality.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh`:
- Around line 226-239: The add_delta_s lambda function uses threadIdx.x directly
to calculate quad_id instead of using the thread_idx parameter that is available
and used consistently elsewhere in the function. Replace the threadIdx.x
reference in the quad_id calculation line with the thread_idx parameter to
maintain consistency with the rest of the codebase.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh`:
- Around line 11-14: Remove the `using namespace cute;` directive from all
affected header files to prevent namespace pollution. In
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh
(lines 11-14), remove the `using namespace cute;` statement. Apply the same fix
to
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh
(lines 11-14),
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh
(lines 11-14),
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh
(lines 11-14),
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh
(lines 15-18), and
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh
(lines 11-14). After removing these directives, qualify any CUTE symbols used in
these files with the `cute::` prefix to maintain correctness while eliminating
namespace scope pollution.
In `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h`:
- Around line 116-149: The tile_count_semaphore field is stored in the Params
struct within DynamicPersistentTileScheduler but is never referenced in either
the get_initial_work() or get_next_work() methods. Resolve this by either
removing the tile_count_semaphore parameter from both the Arguments struct and
Params struct along with its assignment in to_underlying_arguments, or if
dynamic scheduling using this semaphore is intended, implement the logic to use
it atomically within get_initial_work() and get_next_work() to fetch tile
indices dynamically, or add a clear comment explaining why it is stored for
potential future use.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h`:
- Around line 64-84: The parameter `tiled_mma` is declared but unused in the
functions `deduce_smem_layoutSFQ`, `deduce_smem_layoutSFKV`, and
`deduce_smem_layoutSFVt`. To suppress compiler warnings about unused parameters,
add the `[[maybe_unused]]` attribute to the `tiled_mma` parameter in each of
these three functions. This signals to the compiler that the parameter is
intentionally unused, likely for API consistency purposes.
In `@tests/attention/test_nvfp4_attention_sm120.py`:
- Around line 37-40: The function `_require_sm120()` and related code at lines
136-137 use manual CUDA capability checks with
torch.cuda.get_device_capability() instead of the repository's standardized
pattern. Replace the manual capability check logic in `_require_sm120()` with
the appropriate `flashinfer.utils` architecture helper function, and apply the
same pattern to the code at lines 136-137 to ensure consistent SM-based test
gating across the file according to repository guidelines.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7aba14a9-0af1-4003-b3cb-ba90ab4edda8
📥 Commits
Reviewing files that changed from the base of the PR and between c15ac84 and c56b77dd6442adb350aed832f0c33c47bc0631f6.
📒 Files selected for processing (44)
benchmarks/bench_nvfp4_attention_sm120.pycsrc/nvfp4_attention_sm120/.clang-formatcsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cucsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cuflashinfer/__init__.pyflashinfer/jit/__init__.pyflashinfer/jit/nvfp4_attention_sm120.pyflashinfer/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/.clang-formatinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuhtests/attention/test_nvfp4_attention_sm120.pytests/conftest.py
tiffany940107
force-pushed
the
main
branch
2 times, most recently
from
June 15, 2026 12:37
a3603df to
0b2a9dd
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h (1)
61-70: 💤 Low valueUnused type alias
Traits.
Traitsis defined but never referenced. This appears to be dead code.♻️ Suggested removal
template <class TiledMma, class GapFn, class TA, class ALayout, class TB, class BLayout, class TC, class CLayout> CUTE_HOST_DEVICE void gemm_interleaved(TiledMma const& tiled_mma, Tensor<TC, CLayout>& C, Tensor<TA, ALayout> const& A, Tensor<TB, BLayout> const& B, GapFn&& gap_fn) { - using Traits = typename TiledMma::AtomThrID; - using MMAOp = typename TiledMma::MMA_Atom_Arch; mma_unpack_interleaved(MMA_Traits<MMAOp>{}, C, A, B, C, static_cast<GapFn&&>(gap_fn)); }🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h` around lines 61 - 70, In the gemm_interleaved function template, remove the unused type alias `using Traits = typename TiledMma::AtomThrID;` since it is declared but never referenced in the function body. This eliminates dead code and improves clarity.include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h (1)
75-140: 💤 Low valueUnused function parameters in
deduce_smem_layout*functions.The
tiled_mmaandtileshape_mnkparameters are declared but not used—the functions use the template typeTileShape_MNK{}directly. If these parameters are kept for future extensibility or interface consistency, this is fine; otherwise they could be removed.🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h` around lines 75 - 140, The three template functions deduce_smem_layoutSFQ, deduce_smem_layoutSFKV, and deduce_smem_layoutSFVt each declare a tileshape_mnk parameter that is not used in the function body. The tiled_mma parameter is already marked with [[maybe_unused]], but tileshape_mnk is not. Add the [[maybe_unused]] attribute to the tileshape_mnk parameter in all three function declarations to suppress compiler warnings about unused function parameters, indicating that these parameters are intentionally kept for interface consistency or future extensibility.include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh (1)
24-24: 💤 Low valueConsider using explicit
usingdeclarations for consistency.Other files in this PR (load_v.cuh, pv_gemm.cuh, softmax.cuh) use explicit
using cute::X;declarations rather thanusing namespace cute;. Usingusing namespacein a header can pollute the namespace of includers. While the effect is limited to thenvfp4_attentionnamespace here, switching to explicit declarations would maintain consistency across the codebase.🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh` at line 24, Replace the `using namespace cute;` declaration in output_writer.cuh with explicit using declarations for the specific cute namespace symbols actually used in this file. Identify all references to cute namespace members throughout the file and add explicit `using cute::SymbolName;` declarations for each one to maintain consistency with other files in the PR (load_v.cuh, pv_gemm.cuh, softmax.cuh) and avoid namespace pollution.
🤖 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
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h`:
- Around line 61-70: In the gemm_interleaved function template, remove the
unused type alias `using Traits = typename TiledMma::AtomThrID;` since it is
declared but never referenced in the function body. This eliminates dead code
and improves clarity.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh`:
- Line 24: Replace the `using namespace cute;` declaration in output_writer.cuh
with explicit using declarations for the specific cute namespace symbols
actually used in this file. Identify all references to cute namespace members
throughout the file and add explicit `using cute::SymbolName;` declarations for
each one to maintain consistency with other files in the PR (load_v.cuh,
pv_gemm.cuh, softmax.cuh) and avoid namespace pollution.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h`:
- Around line 75-140: The three template functions deduce_smem_layoutSFQ,
deduce_smem_layoutSFKV, and deduce_smem_layoutSFVt each declare a tileshape_mnk
parameter that is not used in the function body. The tiled_mma parameter is
already marked with [[maybe_unused]], but tileshape_mnk is not. Add the
[[maybe_unused]] attribute to the tileshape_mnk parameter in all three function
declarations to suppress compiler warnings about unused function parameters,
indicating that these parameters are intentionally kept for interface
consistency or future extensibility.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 41b86f57-f6ef-445d-b663-43a99213edff
📥 Commits
Reviewing files that changed from the base of the PR and between a3603df8ced5caf7bcf0ad248bab758747edb752 and 0b2a9dd63bc3486ac296357d1445b007059e308a.
📒 Files selected for processing (45)
benchmarks/bench_nvfp4_attention_sm120.pycsrc/nvfp4_attention_sm120/.clang-formatcsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cucsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cuflashinfer/__init__.pyflashinfer/jit/__init__.pyflashinfer/jit/nvfp4_attention_sm120.pyflashinfer/nvfp4_attention_sm120.pyflashinfer/trace/templates/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/.clang-formatinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuhtests/attention/test_nvfp4_attention_sm120.pytests/conftest.py
✅ Files skipped from review due to trivial changes (2)
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/.clang-format
- csrc/nvfp4_attention_sm120/.clang-format
🚧 Files skipped from review as they are similar to previous changes (28)
- flashinfer/init.py
- flashinfer/jit/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh
- tests/conftest.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h
- flashinfer/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu
- tests/attention/test_nvfp4_attention_sm120.py
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh
- benchmarks/bench_nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h
jiahanc
reviewed
Jun 15, 2026
jiahanc
left a comment
jiahanc
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the work! Left some comments
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/attention/test_nvfp4_attention_sm120.py (2)
22-29: 💤 Low valueConsider documenting the Cutlass patching rationale.
The patching logic works around a version-specific issue in the Cutlass DSL, but lacks a comment explaining which versions need the patch or why
OperandMajorModemight be missing fromcute.nvgpu.📝 Suggested documentation addition
def _patch_cutlass_dsl_operand_major_mode(): + """ + Compatibility shim for nvidia-cutlass-dsl. + Some versions place OperandMajorMode under tcgen05 but not in cute.nvgpu. + """ try:🤖 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/attention/test_nvfp4_attention_sm120.py` around lines 22 - 29, The _patch_cutlass_dsl_operand_major_mode() function lacks documentation explaining why the patching is necessary. Add a docstring or comments to the function that clearly document which Cutlass DSL versions require this patch and explain why OperandMajorMode might be missing from cute.nvgpu in certain versions. This will help maintainers understand the purpose and scope of the workaround.
72-84: ⚡ Quick winConsider validating LSE correctness, not just shape and dtype.
The reference implementation returns only the attention output, but the tested implementation also returns
lse(log-sum-exp). Currently, LSE is validated for shape, dtype, and finiteness (lines 127-133) but not for numerical correctness. Since LSE is used in chunked attention merging and other contexts, validating its correctness would strengthen test coverage.♻️ Proposed enhancement to validate LSE
def _reference_attention(q, k, v, causal): q, k, v, qk_correction = _preprocess_qkv_ref(q, k, v) sm_scale = q.shape[-1] ** -0.5 scores = torch.matmul(q.float(), k.float().transpose(-2, -1)) * sm_scale scores = scores + qk_correction * sm_scale if causal: seqlen = q.shape[2] mask = torch.triu( torch.ones(seqlen, seqlen, device=q.device, dtype=torch.bool), diagonal=1 ) scores.masked_fill_(mask, float("-inf")) + lse = torch.logsumexp(scores, dim=-1) probs = torch.softmax(scores, dim=-1) - return torch.matmul(probs, v.float()).to(q.dtype) + out = torch.matmul(probs, v.float()).to(q.dtype) + return out, lseThen update the caller at line 122:
- ref = _reference_attention(q, k, v, causal)[:, :, :seqlen, :] + ref, ref_lse = _reference_attention(q, k, v, causal) + ref = ref[:, :, :seqlen, :] + ref_lse = ref_lse[:, :, :seqlen]And add LSE validation after line 140:
assert cos_sim >= cos_threshold + + # Validate LSE correctness + lse_mae = (lse.float() - ref_lse.float()).abs().mean().item() + assert lse_mae <= 0.5, f"LSE MAE {lse_mae} exceeds threshold"🤖 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/attention/test_nvfp4_attention_sm120.py` around lines 72 - 84, The _reference_attention function currently only returns the attention output but the tested implementation also returns lse (log-sum-exp), which is only validated for shape and dtype. Modify _reference_attention to compute and return the LSE values alongside the attention output (LSE can be computed from the softmax scores or probabilities), update the caller at line 122 to capture both the attention output and LSE from the reference function, and add numerical validation after line 140 to compare the actual LSE values returned by the tested implementation against the reference LSE values to ensure correctness beyond just shape and dtype checks.
🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh`:
- Around line 55-57: The TMALoader class at line 55 and TMAStor class at line 77
both store const references to TMADesc, which creates a dangling reference risk
if the factory functions make_tma_loader and make_tma_storer (defined at lines
99-106) are ever called with temporary descriptor objects. To fix this, change
the member variable declarations in both TMALoader and TMAStor from storing
TMADesc const& to storing TMADesc by value, which eliminates the lifetime
dependency and prevents potential undefined behavior from dangling references.
---
Nitpick comments:
In `@tests/attention/test_nvfp4_attention_sm120.py`:
- Around line 22-29: The _patch_cutlass_dsl_operand_major_mode() function lacks
documentation explaining why the patching is necessary. Add a docstring or
comments to the function that clearly document which Cutlass DSL versions
require this patch and explain why OperandMajorMode might be missing from
cute.nvgpu in certain versions. This will help maintainers understand the
purpose and scope of the workaround.
- Around line 72-84: The _reference_attention function currently only returns
the attention output but the tested implementation also returns lse
(log-sum-exp), which is only validated for shape and dtype. Modify
_reference_attention to compute and return the LSE values alongside the
attention output (LSE can be computed from the softmax scores or probabilities),
update the caller at line 122 to capture both the attention output and LSE from
the reference function, and add numerical validation after line 140 to compare
the actual LSE values returned by the tested implementation against the
reference LSE values to ensure correctness beyond just shape and dtype checks.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3dfc81e0-d7cb-4452-b145-0f71b7ff15d4
📥 Commits
Reviewing files that changed from the base of the PR and between 0b2a9dd63bc3486ac296357d1445b007059e308a and 2206f69278843a8fda9e666356a55f632c3e474d.
📒 Files selected for processing (43)
benchmarks/bench_nvfp4_attention_sm120.pycsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cucsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cuflashinfer/__init__.pyflashinfer/jit/__init__.pyflashinfer/jit/nvfp4_attention_sm120.pyflashinfer/nvfp4_attention_sm120.pyflashinfer/trace/templates/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuhtests/attention/test_nvfp4_attention_sm120.pytests/conftest.py
🚧 Files skipped from review as they are similar to previous changes (39)
- flashinfer/init.py
- flashinfer/jit/init.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh
- tests/conftest.py
- flashinfer/trace/templates/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh
- flashinfer/jit/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu
- flashinfer/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h
- benchmarks/bench_nvfp4_attention_sm120.py
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh (1)
130-132: ⚡ Quick winDerive split-Q group geometry from one trait source. The mainloop and barrier setup both assume two 128-thread groups; if this templated path gains another group shape, thread partitioning and barrier membership diverge.
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh#L130-L132: useNumMmaThreadsandkBlockMPerWGforgroup_id,mma_thread_idx, andwg_m_offset.include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.h#L238-L245: derivemo_sizesandgroup_idfromsize(typename Ktraits::TiledMmaQK{})and add a matchingstatic_assert.🤖 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 `@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh` around lines 130 - 132, The split-Q group geometry is currently hardcoded in two places with a fixed 128-thread group size and 64-thread-per-group assumption, making the code brittle if group shapes change. In mainloop_nonws_splitq.cuh (lines 130-132), replace the hardcoded values in the group_id, mma_thread_idx, and wg_m_offset calculations with template traits: use NumMmaThreads instead of 128 and kBlockMPerWG instead of 64. In attention_kernel_nonws.h (lines 238-245), derive mo_sizes and group_id from the size of typename Ktraits::TiledMmaQK{} to match the mainloop's geometry, and add a static_assert to enforce consistency between the two locations.
🤖 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
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.h`:
- Around line 18-20: The file block_info.h uses the uint32_t type (referenced at
lines 48 and 54) but lacks the necessary standard header that defines it. Add
`#include` <cstdint> after the `#pragma` once directive and before the namespace
flash declaration to ensure uint32_t is properly defined and the file maintains
correct include order independence.
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh`:
- Around line 417-440: The PV accumulator fragments `tOrO_store` (used in the
first if branch) and `tOrO` (created with `make_fragment_like()` in the else
branch) must be explicitly cleared to zero before being used as accumulators in
the `cute::gemm` calls, as fragments created this way contain uninitialized
register values. Add a `clear(tOrO_store)` call before the first gemm loop in
the `is_first_compute` block, and add a `clear(tOrO)` call immediately after
`make_fragment_like(tOrO_store)` in the else block, before the second gemm loop
begins.
In `@include/flashinfer/math.cuh`:
- Around line 76-79: The fma function is using regular arithmetic operations (a
* b + c) instead of the CUDA fmaf intrinsic, which does not guarantee fused
multiply-add behavior. To fix this, replace each arithmetic operation in the
function body for both components (d.x and d.y assignments) with the fmaf
intrinsic, passing the corresponding components of a, b, and c as arguments
respectively. This ensures true fused multiply-add semantics with a single
rounding step as expected from a function named fma.
- Around line 87-90: Two locations perform left shifts on signed operands, which
causes undefined behavior. At include/flashinfer/math.cuh L87-90: Cast n_int to
uint32_t before performing the left shift operation so the expression becomes
((uint32_t)n_int) << 23. At
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh
L293-295: Replace the signed integer literal 0x00FF00FF with the unsigned
literal 0x00FF00FFu before performing the left shift operation to avoid shifting
on a signed operand.
- Around line 95-110: The fp32_vec_to_e4m3 and fp32_vec_to_e2m1 functions
silently return 0 for unsupported compute capabilities below SM100, which can
lead to corrupted outputs if called on incompatible devices. Replace the `return
0;` statements in the else branches (for both functions where the CUDA_ARCH
check fails) with an explicit fail-fast mechanism such as an assertion or
static_assert to immediately flag incorrect usage on unsupported architectures
instead of producing silently corrupted results.
---
Nitpick comments:
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuh`:
- Around line 130-132: The split-Q group geometry is currently hardcoded in two
places with a fixed 128-thread group size and 64-thread-per-group assumption,
making the code brittle if group shapes change. In mainloop_nonws_splitq.cuh
(lines 130-132), replace the hardcoded values in the group_id, mma_thread_idx,
and wg_m_offset calculations with template traits: use NumMmaThreads instead of
128 and kBlockMPerWG instead of 64. In attention_kernel_nonws.h (lines 238-245),
derive mo_sizes and group_id from the size of typename Ktraits::TiledMmaQK{} to
match the mainloop's geometry, and add a static_assert to enforce consistency
between the two locations.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 32d0e46c-1e5e-499a-8fb0-148171a7b2d5
📥 Commits
Reviewing files that changed from the base of the PR and between 2206f69278843a8fda9e666356a55f632c3e474d and c040cd8bceda5655130b36e186f3fbe7f3624bdf.
📒 Files selected for processing (46)
benchmarks/bench_nvfp4_attention_sm120.pycsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cucsrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cudocs/api/attention.rstflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/__init__.pyflashinfer/jit/nvfp4_attention_sm120.pyflashinfer/nvfp4_attention_sm120.pyflashinfer/trace/templates/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/block_info.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_splitq.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel_nonws.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_convert.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.hinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/math.cuhinclude/flashinfer/math.cuhtests/attention/test_nvfp4_attention_sm120.pytests/conftest.py
✅ Files skipped from review due to trivial changes (1)
- docs/api/attention.rst
🚧 Files skipped from review as they are similar to previous changes (37)
- flashinfer/jit/init.py
- tests/conftest.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/params.h
- flashinfer/init.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/tma.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_q.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_k.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/pipeline.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/layout.cuh
- flashinfer/trace/templates/nvfp4_attention_sm120.py
- flashinfer/jit/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/gemm_with_interleave.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/static_switch.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/utils/copy.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/warpgroup.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/delta_correction.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_gemm.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/quantization/fp4_layout.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/output_writer.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/qk_gemm.cuh
- tests/attention/test_nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/common/cute_extension.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/primitives/barrier.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/api/launcher.h
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh
- benchmarks/bench_nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop_nonws_crosstile.cuh
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/producer/load_v.cuh
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_quantize.cu
- flashinfer/nvfp4_attention_sm120.py
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/scheduler.h
- csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu
- include/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/traits.h
qsang-nv
reviewed
Jun 16, 2026
Collaborator
|
/bot run tests/attention |
Collaborator
tiffany940107
force-pushed
the
main
branch
2 times, most recently
from
June 16, 2026 16:01
3eb6288 to
31efc5b
Compare
Contributor
Author
|
/bot run tests/attention |
Collaborator
|
@tiffany940107 is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
tiffany940107
force-pushed
the
main
branch
2 times, most recently
from
June 16, 2026 16:31
b677340 to
d1fca27
Compare
Collaborator
|
[FAILED] Pipeline #54896756: 6/20 passed |
Collaborator
|
/bot run tests/attention |
Collaborator
Collaborator
|
/bot run tests/attention |
Collaborator
jiahanc
approved these changes
Jun 17, 2026
jiahanc
left a comment
jiahanc
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, thanks for the contribution!
jiahanc
enabled auto-merge (squash)
June 17, 2026 07:05
qsang-nv
approved these changes
Jun 17, 2026
This was referenced Jul 4, 2026
5 tasks
bkryu
pushed a commit
that referenced
this pull request
Jul 10, 2026
…uction, and lse (#3838) ## Description While benchmarking the SM120 NVFP4 attention path from #3640 on an RTX PRO 6000 for #3809, I found three defects versus the SageAttention3 kernel it ports. All are invisible to the current tests (iid inputs make the correction negligible, cosine/mean-abs-err thresholds absorb a uniform scale error, and lse is only checked for NaN), but they matter a lot on real inputs. **1. `qk_correction` is addressed as a compact tensor, but Python passes an expanded one.** The mainloop builds the DS TMA descriptor from `tile_to_shape(SmemLayoutAtomDS{}, ...)` where the atom is `Layout<Shape<kBlockM, kBlockN>, Stride<_0, _1>>` (`kernel/traits.h:165`, `compute/mainloop.cuh:193-200`). The resulting gmem layout has stride 0 within each 128-row block, i.e. the kernel addresses `ptr_ds` as a compact `[batch, heads, seq_len/128, seq_len]` tensor and never uses the strides passed from the binding. SageAttention3's `preprocess_qkv` produces exactly that compact `delta_s` (one `qm @ k^T` row per 128-token block). Our port added a `repeat_interleave(128, dim=2)` that materializes `[batch, heads, seq_len, seq_len]`, so at runtime the kernel read block 0's correction rows for every Q block and crossed head/batch boundaries for `bidh > 0` / `bidb > 0`. I verified the addressing empirically before changing anything: with B=1 H=1, scrambling every row `r >= seq_len/128` of the expanded tensor leaves the output bitwise identical (those rows are never read), and a marker placed in gmem row `r` moves exactly the output rows of Q block `r`, for `r = 0..seq_len/128-1`. The fix passes the compact tensor through (`quantize_qkv` output shape changes to `[B, H, S/128, S]`, or `[B, H, 1, S]` for `per_block_mean=False`) and updates the binding shape check. This also removes the O(seq_len^2) fp32 materialization, which was 99% of the preprocessing cost (7.3 ms of the 9.5 ms end-to-end at S=16K, and an 8.6 GB allocation). Per review feedback the correction matmul now also runs in fp32, since a float16 matmul output would overflow at 65504. **2. Row-sum reduction folds in the neighboring row.** One accumulator row spans 4 threads in the acc layout (2 columns per thread per 8-column group). `SoftmaxFused::RowReductionThr` was 8 -- SageAttention3 has 4 (`softmax_fused.h:35` upstream) -- so the `__shfl_xor(4)` step in `finalize()` added the neighboring row's sum into every `row_sum`, halving the output. A probe with uniform scores and V = all-ones (exact output must be 1.0 everywhere, independent of P quantization) returns 0.5156 = half the V-dequant value, exactly. **3. `lse` is never written.** `LSEWriter` (`compute/epilogue/lse_writer.cuh`) had no call site, so `fwd` returned an uninitialized buffer -- the existing non-NaN assertions pass or fail depending on allocator reuse (SageAttention3 has its LSE store commented out entirely; the writer here was half-adapted dead code with m16n8 fragment asserts that don't compile against this kernel's mma). It's now called from the consumer loop after the softmax finalize, using the per-warp-group PV mma for the row mapping plus an explicit row offset, and the value drops the `fp8_scalexfp4_scale_log2` factor that `row_sum` carries for the FP4 P quantization, so `lse` is the plain ln-sum-exp of the scaled scores. Uniform scores now give `lse = ln(seq_len)` exactly; iid/structured lse mean abs error vs exact `logsumexp` is 0.01-0.02. ### Accuracy A/B on RTX PRO 6000 (CUDA 13, torch 2.11, this repo at c53229e) Reference is exact fp32 SDPA on the same inputs; `alpha` is the least-squares scale `<out,ref>/<ref,ref>` (1.0 = unbiased), so it exposes the halving that cosine cannot see. | case | main: cos / relL2 / alpha | fixed: cos / relL2 / alpha | |---|---|---| | uniform scores, V=ones (exact 1.0) | out = 0.516 | out = 1.031 (= exact V-quant value) | | iid, B2 H4 S1024 D128 | 0.974 / 0.524 / 0.489 | 0.982 / 0.188 / 0.980 | | block-structured Q, per_block_mean=True | 0.408 / 0.914 / 0.186 | 0.984 / 0.181 / 0.983 | | per-(b,h)-shifted Q, per_block_mean=False | 0.455 / 0.891 / 0.212 | 0.983 / 0.182 / 0.982 | The structured-Q rows are the case the SageAttention smoothing machinery exists for; on main the correction is applied to the wrong blocks, so accuracy collapses exactly where the algorithm is supposed to help. ### End-to-end perf (quantize_qkv + fwd, in-repo benchmark, non-causal) | shape | main e2e | fixed e2e | BF16 SDPA-flash | |---|---|---|---| | B4 H8 S4096 D128 | 2.590 ms (106 TFLOPs/s) | 0.723 ms (380 TFLOPs/s) | 1.107 ms | | B1 H8 S16384 D128 | 9.528 ms (115 TFLOPs/s) | ~2.5 ms (~440 TFLOPs/s) | 5.594 ms | Kernel-only time is unchanged (~0.42 / ~1.9 ms, 660 / 570 TFLOPs/s); the fixed e2e includes the fp32 correction matmul and the lse store. The 16K shape shows some run-to-run spread on my box (2.5-3.1 ms) -- 3-4x faster than main either way. End-to-end goes from 2.3x slower than BF16 flash attention to roughly 1.5-2.2x faster, which is the regime diffusion workloads (fresh Q/K/V every step) actually run in. Note: `quantize_qkv`'s output shape for `qk_correction` changes; the pair of public APIs stays self-consistent, so code using them together is unaffected. ## Related Issues #3809 (this makes the existing SM120 NVFP4 path usable and fast end-to-end on diffusion shapes; kernel-level follow-ups tracked there), #3640. ## Pull Request Checklist ### Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). `pytest -q tests/attention/test_nvfp4_attention_sm120.py`: 13 passed (7 existing + 6 new) on RTX PRO 6000. The new regression tests fail on main: `test_nvfp4_attention_sm120_structured_q_correction` (cos 0.41/0.45 vs the 0.95 floor), `test_nvfp4_attention_sm120_output_magnitude` (0.516 vs 1.0 +/- 0.05), and `test_nvfp4_attention_sm120_lse` (uninitialized buffer). ## Reviewer Notes The decisive probes are easy to re-run: (a) scramble rows `>= seq_len/128` of the expanded correction on main and observe bitwise-identical output; (b) uniform scores with V = all-ones must return 1.0 and lse = ln(seq_len). Happy to split the three fixes into separate PRs if you prefer. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added LSE output support in the SM120 NVFP4 attention path, making log-sum-exp values available alongside attention results. * Updated FP32 correction handling to use a compact layout that matches block-based attention processing. * **Bug Fixes** * Corrected shape validation for correction data to reject outdated expanded layouts and accept the expected block-wise format. * Improved numerical handling for LSE and softmax computation, including more stable output scaling. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
saltyminty
pushed a commit
that referenced
this pull request
Jul 13, 2026
…3897) <!-- .github/pull_request_template.md --> ## 📌 Description Enables the SM120 NVFP4 attention kernel (#3640) on SM121 (GB10 / DGX Spark). The kernel code needs no changes — it compiles and runs correctly for `sm_121a`; this PR only widens the SM120-only gates - `pytest tests/attention/test_nvfp4_attention_sm120.py`: **7/7 pass** on this branch; **13/13 pass** with #3838's fixes and regression tests applied on top. - Kernel benchmarks vs FA2 BF16 ragged prefill at identical shapes (non-causal, attention-only kernel time): | Config (H8) | FA2 BF16 | NVFP4 | Speedup | |---|---|---|---| | B4 S4096 D128 | 2.97 ms (93 TF/s) | 1.28 ms (215 TF/s) | 2.32× | | B2 S8192 D128 | 5.81 ms (95 TF/s) | 2.48 ms (221 TF/s) | 2.34× | | B1 S32768 D128 | 46.09 ms (95 TF/s) | 19.51 ms (225 TF/s) | 2.36× | | B1 S16384 **D64** | 5.89 ms (93 TF/s) | 3.81 ms (144 TF/s) | 1.55× | Speedup is stable across head counts (2.2–2.6× for H=1…32 at D128); throughput saturates from B·H ≥ 16 at S4096. Speedup is around 1.5-1.6x for head dim 64. <details> <summary>Commands to reproduce the perf numbers</summary> NVFP4 (shape lists zip together; one row printed per config): ```bash # D=128 rows + head-dim sweep python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 4 2 2 1 1 --num-heads 8 --head-dim 128 \ --seq-len 4096 4096 8192 16384 32768 \ --no-causal --warmup 3 --repeat 10 # D=64 rows python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 4 2 2 1 1 --num-heads 8 --head-dim 64 \ --seq-len 4096 4096 8192 16384 32768 \ --no-causal --warmup 3 --repeat 10 # head-count sweep (B*H occupancy) python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 1 --num-heads 1 2 4 8 16 32 --head-dim 128 \ --seq-len 4096 --no-causal --warmup 3 --repeat 10 ``` The "NVFP4" column is the `attention_only` number (CUDA-graph replay, pure kernel time, quantization excluded); `end_to_end` additionally includes `quantize_qkv` each iteration. FA2 BF16 baseline (identical shapes; CUPTI kernel timing is on by default; uniform full-length sequences — do not pass `--random_actual_seq_len`): ```bash for cfg in "4 4096" "2 4096" "2 8192" "1 16384" "1 32768"; do set -- $cfg python benchmarks/flashinfer_benchmark.py \ --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 \ --batch_size $1 --s_qo $2 --s_kv $2 \ --num_qo_heads 8 --num_kv_heads 8 \ --head_dim_qk 128 --head_dim_vo 128 \ --q_dtype bfloat16 --kv_dtype bfloat16 --refcheck done # D=64 baseline: same loop with --head_dim_qk 64 --head_dim_vo 64 # head sweep baseline: --batch_size 1 --s_qo 4096 --s_kv 4096, loop --num_qo_heads/--num_kv_heads over 1 2 4 8 16 32 ``` Environment: NVIDIA GB10 (SM121), CUDA 13.0, torch 2.11. FA2 median of 30 iters (CUPTI); NVFP4 median of 10 iters (CUDA-graph). Comparing FA2 kernel time against NVFP4 `attention_only` is apples-to-apples — both exclude host launch overhead. </details> <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> - #3809 - #3838 - PR 3838 is an orthogonal correctness fix ## 🚀 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Expanded NVFP4 attention support to include GPU compute capability 12.1 in addition to 12.0. * Broadened which NVFP4 SM120 modules are generated and enabled, and improved JIT build flag selection across compatible CUDA environments. * **Bug Fixes** * Updated compute-capability validation and the related error/skip messaging to reflect 12.0 and 12.1 support. * Adjusted test gating to run when either supported compute capability is available, reducing unnecessary skips. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
68 tasks
5 tasks
saltyminty
pushed a commit
that referenced
this pull request
Aug 25, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This PR optimizes the CUTLASS/CUDA SM120 NVFP4 attention forward kernel, with a focus on head-dimension-128 workloads used by Cosmos inference. The largest performance contribution is **symmetric N64/N64 score-slot reuse** inside the existing N128 attention tile: 1. Compute the complete N128 QK score tile and establish the row maximum and online-softmax rescale across all 128 columns. 2. Softmax-quantize the first N64 score slot and consume it in the PV MMA. 3. Reuse that retired score-register slot immediately for the next tile's QK scores while the second N64 slot is consumed. 4. Repeat symmetrically for the second slot. This avoids keeping a second full score fragment live and exposes more QK/PV overlap while preserving the full-N128 softmax reduction. Additional kernel changes: - Split the main loop into a prologue, branch-free steady state, and final drain. - Remove the obsolete math-order barrier and reduce warp synchronization. - Compile masking out of the noncausal specialization. For causal traversal, only the first tile can intersect the diagonal; subsequent tiles are fully valid. - Derive mask row/column coordinates directly instead of retaining an identity tensor. - Tune consumer register allocation and persistent CTA scheduling/work distribution for both short and long causal/noncausal workloads. - Add a compile-time LSE specialization so the output-only path does not allocate or write LSE. ### LSE API behavior `return_lse=False` is now the default and returns only the attention output tensor: ~~~python out = flashinfer.nvfp4_attention_sm120_fwd(..., return_lse=False) ~~~ `return_lse=True` preserves the LSE-capable path and returns `(out, lse)`: ~~~python out, lse = flashinfer.nvfp4_attention_sm120_fwd(..., return_lse=True) ~~~ Disabling LSE only removes the LSE allocation/writeback; it does not remove any computation required for the attention output. The tests compare the output-only and LSE-enabled specializations, and the existing output/LSE reference checks continue to pass. ### Performance Local measurements used one NVIDIA RTX PRO 6000 Blackwell Server Edition GPU (SM120, 188 SMs), BF16 input/output, `D=128`, `per_block_mean=True`, and `return_lse=False`. The table reports CUDA-Graph attention-only kernel timing (median of 100 iterations after 10 warmups); QKV quantization is excluded. The maximum supported 2430 MHz graphics clock was requested and monitored during the runs. | Shape | Mode | [#3640](#3640) baseline (reported) | CuTe DSL reference (reported) | This PR (measured) | Latency vs #3640 | Latency vs CuTe DSL | |---|---:|---:|---:|---:|---:|---:| | B4, H8, S4096, D128 | Noncausal | 0.299 ms / 920.5 TFLOP/s | 0.262 ms / 1050.2 TFLOP/s | **0.262 ms / 1050.1 TFLOP/s** | **-12.4%** | same at reported precision | | B1, H8, S32768, D128 | Noncausal | 4.970 ms / 884.9 TFLOP/s | 4.398 ms / 1000.0 TFLOP/s | **4.079 ms / 1078.2 TFLOP/s** | **-17.9%** | **-7.3%** | | B4, H8, S4096, D128 | Causal | 0.223 ms / 616.6 TFLOP/s | 0.180 ms / 764.1 TFLOP/s | **0.165 ms / 830.6 TFLOP/s** | **-26.0%** | **-8.3%** | | B1, H8, S32768, D128 | Causal | 2.958 ms / 743.3 TFLOP/s | 2.516 ms / 874.0 TFLOP/s | **2.207 ms / 996.5 TFLOP/s** | **-25.4%** | **-12.3%** | The #3640 and CuTe DSL columns reproduce the values shared in the RTX Blackwell kernel performance design document; they were not remeasured in the same run as this PR. The relative percentages therefore use the published rounded latencies and should be treated as cross-run comparisons. ## Contributors - [Atharva Joshi (@atharvajoshi10)](https://github.com/atharvajoshi10) — contributed the CuTe DSL kernel optimization and design document, including the symmetric N64/N64 score-slot reuse strategy, together with the RTX PRO 6000 benchmark results that motivated this CUTLASS implementation. ## 🔍 Related Issues - Follow-up to #3640. - Builds on the SM120 NVFP4 attention/LSE implementation in #3838. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I installed `pre-commit` in an isolated environment. - [x] I ran `pre-commit run --all-files`; all hooks passed. ## 🧪 Tests - [x] Tests were added for both `return_lse=False` and `return_lse=True`. - [x] `python -m pytest -q tests/attention/test_nvfp4_attention_sm120.py` — **16 passed**. - [x] Output-only and LSE-enabled attention outputs match with `rtol=0, atol=5e-4`; LSE-enabled reference checks pass. ## Reviewer Notes Please pay particular attention to: - The intentional API default change: callers that need the previous `(out, lse)` result should pass `return_lse=True`. - The full-N128 max/rescale followed by per-N64 softmax quantization and score-slot reuse. - The persistent scheduler changes for triangular causal work distribution. The upstream `pre-commit` check is green. The full PR test matrix is currently skipped by the repository's unauthorized-PR gate and requires maintainer approval before it can run. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added grouped-query and multi-query attention support with separate query and key/value head counts. * Added optional LSE output and support for unpadded key/value sequence lengths. * Added configurable output buffers and output data types. * Enhanced tracing and benchmarking for the new attention options. * **Bug Fixes** * Improved validation for head ratios, tensor shapes, sequence lengths, and masking. * Improved scheduling and handling of empty or irregular sequence shapes. * **Tests** * Added coverage for GQA/MQA, LSE behavior, unpadded sequences, output dtypes, tracing, and invalid inputs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
5 tasks
jiahanc
pushed a commit
that referenced
this pull request
Sep 7, 2026
## 📌 Description This PR adds native NVFP4 sparse MLA support on SM120/SM121 for the DeepSeek V4 Flash attention path. The feature is opt-in through `kv_cache_format="nvfp4"` on the existing `flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4` API; the existing FP8 cache ABI and behavior remain unchanged and are still the default. ### Implementation - Adds a 384-byte/token paged NVFP4 cache ABI: - the 448-dimensional NoPE payload is quantized in groups of 16 to packed E2M1 values with one E4M3 scale per group; - the 64-dimensional RoPE payload remains BF16; - full-page packing and incremental slot-based append helpers support HND/NHD layouts and page-strided vLLM cache allocations. - Adds separate native attention kernels for both phases: - a single-launch streaming prefill kernel that gathers paged NVFP4 V directly, performs the candidate-tile transpose in CTA-local storage, and double-buffers the next source tile without a global transposed-V workspace; - a grouped split-K decode kernel with an autotuned chunks-per-block tactic and split reduction. - Reuses the existing sparse-MLA facade and planner persistence mechanism while maintaining an independent NVFP4 calibration namespace. The planner measures the streaming-prefill versus split-K-decode crossover and decode tactic for each supported serving shape; it does not reuse FP8 measurements. - Supports primary and optional extra sparse cache segments, dynamic top-k lengths, attention sinks, caller-owned workspace, CUDA Graph capture, AOT/JIT registration, and empty pipeline-parallel slices. - Adds correctness tests plus reproducible cache, prefill, decode, and planner benchmarks. The currently supported NVFP4 surface is 16/32/64/128 query heads, primary top-k 128 or 512, primary page size 64, and optional extra-cache page size 2 or 64. ### Performance Measurements below are medians on an NVIDIA RTX PRO 5000 Blackwell GPU (SM120), CUDA 13.0. FP8 and NVFP4 use the same generated inputs and independently selected tactics. | Phase and shape | FP8 | NVFP4 | Speedup | | --- | ---: | ---: | ---: | | Prefill: T=8192, H=64, K=128 | 3909.632 us | 2337.792 us | 1.6724x (+67.24%) | | Prefill: T=8192, H=64, K=128+512 | 10216.448 us | 7179.264 us | 1.4230x (+42.30%) | | Prefill: T=8192, H=128, K=128+512 | 20149.248 us | 14240.768 us | 1.4149x (+41.49%) | | Decode: T=8, H=64, K=128+512 | 43.008 us | 32.768 us | 1.3125x (+31.25%) | Reproduction: ```bash python benchmarks/bench_sparse_mla_nvfp4_prefill.py \ --num-tokens 8192 --num-heads 64 --topk 128 \ --extra-topk 512 --extra-page-size 64 python benchmarks/bench_sparse_mla_nvfp4_decode.py \ --num-tokens 8 --num-heads 64 --topk 128 \ --extra-topk 512 --extra-page-size 64 ``` A paired DeepSeek-V4-Flash serving run using `vllm/vllm-openai:v0.26.0`, PP=4, 256 requests, concurrency 32, ISL=8192, OSL=1, prefix cache disabled, and a 16K scheduled-token budget measured 80.466 s for FP8 and 73.808 s for NVFP4 over three runs: 1.0902x end-to-end speedup and +9.02% token throughput. ## 🔍 Related Issues Related NVFP4 SM120 work: #3640 and #4502. ## 🚀 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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). SM120 regression results: ```text pytest -q tests/attention/test_sparse_mla_nvfp4_sm120.py tests/attention/test_sparse_mla_nvfp4_sm120_plan.py tests/attention/test_sparse_mla_sm120_dispatch.py # 77 passed pytest -q tests/attention/test_sparse_mla_sm120.py -k dsv4_public_api # 5 passed, 483 deselected pre-commit run --all-files # all hooks passed (clang-format, mypy, ruff check, ruff format, and repository checks) ``` The GPU test matrix covers bit-level cache packing/append, HND/NHD and strided pages, numerical references for prefill and decode, single/dual cache layouts, supported head counts, ragged lengths, attention sinks, public API dispatch, empty PP slices, CUDA Graph replay, and planner policy. ## Reviewer Notes - FP8 remains the default. NVFP4 is explicitly selected with `kv_cache_format="nvfp4"`. - Operator gains are shape-dependent. Long-prompt production shapes show the largest prefill gains; short-K/small-head decode shapes can be neutral, so NVFP4 uses its own calibrated phase/CPB decisions rather than an FP8 tactic or a fixed query-token threshold. - In the strict ISL=8192/OSL=256 serving experiment, full-request throughput improved by 3.67% due primarily to lower TTFT, while steady-window decode throughput was 0.57% below FP8. The decode operator path is included and independently tuned, but additional serving-level decode optimization remains follow-up work. - Implementation, tests, benchmarking, and PR drafting were AI-assisted; the listed checks and measurements were run on hardware. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added NVFP4 cache packing and incremental append support for DeepSeek-V4 sparse MLA. * Added NVFP4 sparse MLA decode and prefill on compatible SM120/SM121 GPUs. * Added FP8 or NVFP4 cache format selection, optional extra KV data, attention sinks, and runtime top-k lengths. * Added automatic performance planning and calibration for decode versus prefill. * **Documentation** * Documented the new NVFP4 cache APIs. * **Benchmarks** * Added NVFP4 packing, append, decode, prefill, and planning benchmarks. * **Bug Fixes** * Added validation for invalid cache alignment and slot mappings. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: tiffany940107 <tiffany940107@users.noreply.github.com> Co-authored-by: Yan Wang <yanwa@smc521ge-0080.ipp2a2.colossus.nvidia.com> Co-authored-by: Yan Wang <yanwa@2u2g-spr-0094.ipp4a1.colossus.nvidia.com>
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 a dense SM120 pure NVFP4 attention JIT path.
Main changes:
csrc/nvfp4_attention_sm120/.include/flashinfer/attention/sm120/nvfp4_attention_sm120/.🔍 Related Issues
N/A
🚀 Pull Request Checklist
✅ 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
Validation:
pre-commit run --all-files: passeddockerhub.nvidia.com/flashinfer/flashinfer-ci-cu130:20260408-4cce866PYTHONPATH=/workspace FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer_nvfp4_sm120_rename_pytest python -m pytest -q tests/attention/ test_nvfp4_attention_sm120.py -s6 passed, 1 warning in 88.58sThe warning is from pytest cache write permission inside Docker and is unrelated to test correctness.
Reviewer Notes
This path is SM120-specific. Public module and header paths use explicit
sm120naming to avoid implying support on otherarchitectures.
Summary by CodeRabbit
Release Notes
causalandper_block_meanoptions).