[TRTLLM-34871][feat] Add cute dsl FP8 paged MQA logits decode kernel#13219

limin2021 · 2026-04-20T11:51:22Z

@coderabbitai summary

Description

https://jirasw.nvidia.com/browse/CFK-34871

Summary

Add a CuTE DSL implementation of the FP8 paged MQA logits kernel for Blackwell (SM100+), replacing the C++ DeepGEMM-based fp8_paged_mqa_logits in the sparse attention indexer decode path
The kernel is gated by a new config flag use_cute_dsl_logits (default False) in DeepSeekSparseAttentionConfig, with automatic SM100+ detection
Register as PyTorch custom op trtllm::cute_dsl_fp8_paged_mqa_logits with fake tensor support for torch.compile

Key design details

Architecture: 384 threads (256 math in 2 warp groups + 2 TMA + 2 UMMA warps), persistent kernel with DeepGEMM-aligned scheduler
Full-K TMA: Single TMA load per KV block [128, 128], UMMA iterates 4x K=32
Fused KV layout: KV data and FP8 scales stored contiguously per physical block for L2 cache locality
Merged KV+Scale pipeline: KV and scales share a single TMA barrier, eliminating separate scale pipeline overhead
Dynamic shape support: Runtime-varying dims (batch_size, num_phys_blocks, etc.) marked dynamic via mark_compact_shape_dynamic for JIT cache reuse
Configurable epilogue dtypes: Supports FP32/FP16/BF16 accumulator, epilogue, and output dtypes

Test plan

Unit tests: pytest tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py on Blackwell GPU
Verify numerical correctness against C++ DeepGEMM baseline

# accuracy ut
python -m pytest tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py -v
# fp32 bench perf
python tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --batch_size 1 2 4 8 16 32 --next_n 1 2 3 4 --context_len 4096 8192 16384 32768 65536 
# fp16
python tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py --batch_size 1 2 4 8 16 32 --next_n 1 2 3 4 --context_len 4096 8192 16384 32768 65536  --output_dtype float16

TODO：
1.e2e accuracy test for fp32 and fp16 output dtype.

Benchmark results (B200, SM count=148)

output_dtype=float32 num_epi_subtiles=1

batch ctx next_n nblk | DSL(us) DG(fp32,us) DG/DSL

1    4096      1      32 |     3.7         3.7  1.015x
2    4096      1      64 |     3.7         3.8  1.026x
4    4096      1     128 |     3.7         3.8  1.025x
8    4096      1     256 |     3.8         3.8  1.022x
16    4096      1     512 |     4.3         4.3  1.004x
32    4096      1    1024 |     5.2         5.3  1.012x

1    8192      1      64 |     3.7         3.8  1.020x
2    8192      1     128 |     3.7         3.8  1.026x
4    8192      1     256 |     3.7         3.8  1.023x
8    8192      1     512 |     4.3         4.4  1.014x
16    8192      1    1024 |     5.2         5.2  1.015x
32    8192      1    2048 |     6.6         6.7  1.016x

1   16384      1     128 |     3.7         3.8  1.007x
2   16384      1     256 |     3.7         3.8  1.023x
4   16384      1     512 |     4.3         4.3  1.017x
8   16384      1    1024 |     5.1         5.4  1.048x
16   16384      1    2048 |     6.5         6.6  1.023x
32   16384      1    4096 |     9.5         9.5  0.995x

1   32768      1     256 |     3.7         3.8  1.027x
2   32768      1     512 |     4.3         4.3  1.019x
4   32768      1    1024 |     5.1         5.3  1.045x
8   32768      1    2048 |     6.5         6.6  1.005x
16   32768      1    4096 |     9.5         9.5  0.994x
32   32768      1    8192 |    25.5        25.7  1.008x

1   65536      1     512 |     4.3         4.3  1.006x
2   65536      1    1024 |     5.2         5.3  1.023x
4   65536      1    2048 |     6.5         6.6  1.010x
8   65536      1    4096 |     9.4         9.5  1.004x
16   65536      1    8192 |    25.7        26.1  1.013x
32   65536      1   16384 |    46.5        46.5  1.000x

1    4096      2      32 |     4.0         4.3  1.081x
2    4096      2      64 |     4.0         4.3  1.080x
4    4096      2     128 |     4.0         4.3  1.086x
8    4096      2     256 |     4.0         4.4  1.090x
16    4096      2     512 |     4.9         5.0  1.038x
32    4096      2    1024 |     6.2         6.5  1.045x

1    8192      2      64 |     4.0         4.3  1.061x
2    8192      2     128 |     4.0         4.3  1.081x
4    8192      2     256 |     4.0         4.4  1.083x
8    8192      2     512 |     4.8         5.0  1.043x
16    8192      2    1024 |     6.1         6.5  1.051x
32    8192      2    2048 |     8.3         8.7  1.050x

1   16384      2     128 |     4.0         4.3  1.079x
2   16384      2     256 |     4.0         4.4  1.090x
4   16384      2     512 |     4.8         5.0  1.045x
8   16384      2    1024 |     6.1         6.4  1.050x
16   16384      2    2048 |     8.3         8.7  1.056x
32   16384      2    4096 |    12.6        13.3  1.053x

1   32768      2     256 |     4.1         4.4  1.074x
2   32768      2     512 |     4.9         5.1  1.041x
4   32768      2    1024 |     6.1         6.4  1.053x
8   32768      2    2048 |     8.2         8.7  1.054x
16   32768      2    4096 |    12.5        13.2  1.063x
32   32768      2    8192 |    27.4        27.8  1.018x

1   65536      2     512 |     4.9         5.0  1.039x
2   65536      2    1024 |     6.2         6.5  1.051x
4   65536      2    2048 |     8.1         8.6  1.064x
8   65536      2    4096 |    12.5        13.2  1.053x
16   65536      2    8192 |    27.2        27.9  1.027x
32   65536      2   16384 |    48.9        49.4  1.012x

1    4096      3      32 |     4.2         3.6  0.858x
2    4096      3      64 |     4.2         3.7  0.895x
4    4096      3     128 |     4.3         4.1  0.972x
8    4096      3     256 |     4.3         4.8  1.116x
16    4096      3     512 |     5.3         5.9  1.128x
32    4096      3    1024 |     7.1         8.2  1.164x

1    8192      3      64 |     4.2         3.8  0.902x
2    8192      3     128 |     4.2         4.2  0.988x
4    8192      3     256 |     4.4         4.8  1.087x
8    8192      3     512 |     5.3         5.9  1.128x
16    8192      3    1024 |     7.0         8.1  1.157x
32    8192      3    2048 |    10.1        12.3  1.217x

1   16384      3     128 |     4.2         4.2  0.993x
2   16384      3     256 |     4.4         4.9  1.106x
4   16384      3     512 |     5.3         5.9  1.125x
8   16384      3    1024 |     7.0         8.1  1.153x
16   16384      3    2048 |    10.1        12.2  1.214x
32   16384      3    4096 |    16.5        20.3  1.233x

1   32768      3     256 |     4.4         4.8  1.110x
2   32768      3     512 |     5.3         6.0  1.129x
4   32768      3    1024 |     7.0         8.1  1.158x
8   32768      3    2048 |    10.0        12.1  1.214x
16   32768      3    4096 |    16.5        20.2  1.225x
32   32768      3    8192 |    31.0        49.5  1.594x

1   65536      3     512 |     5.2         5.8  1.123x
2   65536      3    1024 |     7.0         8.1  1.153x
4   65536      3    2048 |     9.9        12.2  1.229x
8   65536      3    4096 |    16.4        20.2  1.233x
16   65536      3    8192 |    31.1        41.0  1.320x
32   65536      3   16384 |    59.4       100.5  1.693x

1    4096      4      32 |     4.8         4.2  0.873x
2    4096      4      64 |     4.8         4.3  0.905x
4    4096      4     128 |     4.8         4.3  0.887x
8    4096      4     256 |     4.9         4.9  1.014x
16    4096      4     512 |     6.3         6.4  1.015x
32    4096      4    1024 |     8.8         8.5  0.972x

1    8192      4      64 |     4.9         4.2  0.861x
2    8192      4     128 |     4.7         4.3  0.905x
4    8192      4     256 |     4.9         4.9  1.001x
8    8192      4     512 |     6.3         6.3  1.006x
16    8192      4    1024 |     8.8         8.5  0.965x
32    8192      4    2048 |    12.7        13.1  1.031x

1   16384      4     128 |     4.8         4.3  0.897x
2   16384      4     256 |     4.9         4.9  1.015x
4   16384      4     512 |     6.3         6.3  1.004x
8   16384      4    1024 |     8.8         8.5  0.964x
16   16384      4    2048 |    12.6        13.1  1.036x
32   16384      4    4096 |    20.8        22.1  1.061x

1   32768      4     256 |     4.8         4.9  1.012x
2   32768      4     512 |     6.2         6.3  1.013x
4   32768      4    1024 |     8.7         8.5  0.969x
8   32768      4    2048 |    12.4        13.1  1.052x
16   32768      4    4096 |    20.7        21.8  1.056x
32   32768      4    8192 |    38.6        40.1  1.038x

1   65536      4     512 |     6.2         6.3  1.010x
2   65536      4    1024 |     8.7         8.4  0.959x
4   65536      4    2048 |    12.3        13.0  1.054x
8   65536      4    4096 |    21.1        21.8  1.033x
16   65536      4    8192 |    38.5        40.0  1.038x
32   65536      4   16384 |    69.9        76.2  1.089x

next_n=1: CuTE DSL matches DeepGEMM (~1.0-1.03x)
next_n=2: CuTE DSL consistently faster (1.02-1.08x)
next_n=4: CuTE DSL significantly faster (up to 1.42x at large batch×ctx), as DeepGEMM uses 2x multicast overhead
next_n=3: DeepGEMM N/A (unsupported); CuTE DSL handles arbitrary next_n

output_dtype=float16 num_epi_subtiles=1

batch ctx next_n nblk | DSL(us) DG(fp32,us) DG/DSL DSL(fp32,us) DSL(fp32)/DSL

1    4096      1      32 |     3.6         3.8  1.037x          3.6        1.003x
2    4096      1      64 |     3.6         3.8  1.039x          3.7        1.022x
4    4096      1     128 |     3.7         3.8  1.036x          3.7        1.007x
8    4096      1     256 |     3.6         3.8  1.055x          3.7        1.021x
16    4096      1     512 |     4.1         4.3  1.038x          4.3        1.027x
32    4096      1    1024 |     5.0         5.3  1.047x          5.1        1.022x

1    8192      1      64 |     3.7         3.8  1.031x          3.7        1.011x
2    8192      1     128 |     3.6         3.8  1.048x          3.7        1.029x
4    8192      1     256 |     3.6         3.8  1.046x          3.8        1.030x
8    8192      1     512 |     4.2         4.4  1.047x          4.3        1.030x
16    8192      1    1024 |     4.9         5.3  1.071x          5.2        1.047x
32    8192      1    2048 |     6.2         6.7  1.084x          6.6        1.061x

1   16384      1     128 |     3.7         3.8  1.022x          3.7        1.012x
2   16384      1     256 |     3.7         3.8  1.046x          3.7        1.025x
4   16384      1     512 |     4.2         4.3  1.040x          4.3        1.025x
8   16384      1    1024 |     5.0         5.3  1.079x          5.1        1.029x
16   16384      1    2048 |     6.2         6.7  1.081x          6.5        1.044x
32   16384      1    4096 |     8.9         9.5  1.063x          9.5        1.070x

1   32768      1     256 |     3.7         3.8  1.043x          3.7        1.017x
2   32768      1     512 |     4.1         4.3  1.050x          4.3        1.036x
4   32768      1    1024 |     4.9         5.3  1.075x          5.1        1.040x
8   32768      1    2048 |     6.1         6.6  1.081x          6.5        1.060x
16   32768      1    4096 |     8.8         9.4  1.075x          9.5        1.083x
32   32768      1    8192 |    24.9        25.7  1.030x         25.6        1.025x

1   65536      1     512 |     4.2         4.3  1.037x          4.3        1.030x
2   65536      1    1024 |     5.0         5.3  1.061x          5.1        1.032x
4   65536      1    2048 |     6.2         6.6  1.073x          6.5        1.060x
8   65536      1    4096 |     8.8         9.4  1.069x          9.5        1.081x
16   65536      1    8192 |    25.2        26.1  1.036x         25.7        1.021x
32   65536      1   16384 |    45.3        46.6  1.027x         46.5        1.025x

1    4096      2      32 |     3.9         4.3  1.120x          4.0        1.032x
2    4096      2      64 |     3.9         4.3  1.112x          4.0        1.028x
4    4096      2     128 |     3.9         4.3  1.118x          4.0        1.024x
8    4096      2     256 |     3.9         4.4  1.130x          4.0        1.041x
16    4096      2     512 |     4.7         5.1  1.083x          4.9        1.042x
32    4096      2    1024 |     6.0         6.5  1.080x          6.2        1.031x

1    8192      2      64 |     3.9         4.3  1.102x          4.0        1.031x
2    8192      2     128 |     3.8         4.3  1.128x          4.0        1.045x
4    8192      2     256 |     3.9         4.4  1.133x          4.0        1.043x
8    8192      2     512 |     4.6         5.0  1.080x          4.8        1.042x
16    8192      2    1024 |     5.9         6.5  1.098x          6.2        1.046x
32    8192      2    2048 |     8.2         8.7  1.063x          8.2        1.007x

1   16384      2     128 |     3.8         4.3  1.130x          4.0        1.048x
2   16384      2     256 |     3.9         4.4  1.132x          4.1        1.044x
4   16384      2     512 |     4.6         5.0  1.085x          4.8        1.031x
8   16384      2    1024 |     5.9         6.4  1.091x          6.1        1.031x
16   16384      2    2048 |     8.1         8.7  1.079x          8.2        1.012x
32   16384      2    4096 |    12.5        13.3  1.063x         12.6        1.007x

1   32768      2     256 |     3.9         4.4  1.121x          4.0        1.033x
2   32768      2     512 |     4.7         5.0  1.077x          4.9        1.042x
4   32768      2    1024 |     6.0         6.5  1.084x          6.1        1.027x
8   32768      2    2048 |     8.2         8.7  1.065x          8.1        0.995x
16   32768      2    4096 |    12.5        13.2  1.056x         12.5        0.997x
32   32768      2    8192 |    26.5        27.9  1.054x         27.4        1.033x

1   65536      2     512 |     4.7         5.0  1.082x          4.9        1.042x
2   65536      2    1024 |     5.9         6.5  1.087x          6.2        1.035x
4   65536      2    2048 |     8.1         8.7  1.065x          8.2        1.006x
8   65536      2    4096 |    12.5        13.1  1.052x         12.5        1.002x
16   65536      2    8192 |    26.5        28.0  1.055x         27.0        1.020x
32   65536      2   16384 |    47.5        49.5  1.041x         48.8        1.027x

1    4096      3      32 |     4.0         3.6  0.917x          4.3        1.071x
2    4096      3      64 |     3.9         3.7  0.953x          4.2        1.066x
4    4096      3     128 |     4.0         4.1  1.046x          4.3        1.073x
8    4096      3     256 |     4.0         4.8  1.197x          4.3        1.080x
16    4096      3     512 |     4.7         6.0  1.260x          5.3        1.112x
32    4096      3    1024 |     6.2         8.2  1.312x          7.0        1.118x

1    8192      3      64 |     3.9         3.8  0.964x          4.2        1.064x
2    8192      3     128 |     3.9         4.2  1.056x          4.2        1.072x
4    8192      3     256 |     4.0         4.8  1.189x          4.3        1.074x
8    8192      3     512 |     4.8         5.9  1.229x          5.3        1.094x
16    8192      3    1024 |     6.3         8.1  1.288x          7.1        1.118x
32    8192      3    2048 |     9.6        12.3  1.278x         10.1        1.055x

1   16384      3     128 |     3.9         4.2  1.068x          4.2        1.069x
2   16384      3     256 |     4.0         4.9  1.220x          4.3        1.080x
4   16384      3     512 |     4.8         6.0  1.249x          5.3        1.103x
8   16384      3    1024 |     6.3         8.2  1.298x          7.0        1.116x
16   16384      3    2048 |     9.4        12.2  1.304x         10.0        1.071x
32   16384      3    4096 |    14.6        20.3  1.390x         16.5        1.130x

1   32768      3     256 |     4.0         4.8  1.197x          4.3        1.072x
2   32768      3     512 |     4.8         6.0  1.248x          5.2        1.094x
4   32768      3    1024 |     6.3         8.1  1.288x          7.1        1.133x
8   32768      3    2048 |     9.3        12.1  1.309x         10.0        1.076x
16   32768      3    4096 |    14.5        20.3  1.396x         16.2        1.114x
32   32768      3    8192 |    32.0        49.2  1.540x         31.3        0.978x

1   65536      3     512 |     4.7         5.9  1.241x          5.2        1.100x
2   65536      3    1024 |     6.3         8.1  1.275x          7.0        1.098x
4   65536      3    2048 |     9.2        12.1  1.324x          9.9        1.077x
8   65536      3    4096 |    14.5        20.2  1.388x         16.3        1.119x
16   65536      3    8192 |    31.8        41.1  1.292x         31.2        0.983x
32   65536      3   16384 |    57.3       100.4  1.752x         59.7        1.041x

1    4096      4      32 |     4.4         4.2  0.942x          4.7        1.074x
2    4096      4      64 |     4.4         4.3  0.977x          4.8        1.077x
4    4096      4     128 |     4.5         4.3  0.955x          4.8        1.074x
8    4096      4     256 |     4.5         4.9  1.100x          4.9        1.083x
16    4096      4     512 |     5.7         6.3  1.115x          6.3        1.107x
32    4096      4    1024 |     7.7         8.5  1.107x          8.9        1.150x

1    8192      4      64 |     4.5         4.2  0.936x          4.9        1.090x
2    8192      4     128 |     4.4         4.3  0.963x          4.7        1.066x
4    8192      4     256 |     4.5         4.9  1.087x          4.9        1.076x
8    8192      4     512 |     5.7         6.3  1.099x          6.3        1.103x
16    8192      4    1024 |     7.9         8.5  1.075x          8.9        1.122x
32    8192      4    2048 |    11.1        13.1  1.187x         12.7        1.143x

1   16384      4     128 |     4.4         4.3  0.965x          4.8        1.076x
2   16384      4     256 |     4.5         4.9  1.098x          4.9        1.089x
4   16384      4     512 |     5.7         6.3  1.094x          6.3        1.105x
8   16384      4    1024 |     7.8         8.5  1.084x          8.9        1.135x
16   16384      4    2048 |    11.2        13.1  1.169x         12.6        1.132x
32   16384      4    4096 |    17.4        21.9  1.254x         20.7        1.188x

1   32768      4     256 |     4.5         4.9  1.101x          4.8        1.086x
2   32768      4     512 |     5.8         6.3  1.075x          6.3        1.091x
4   32768      4    1024 |     7.8         8.4  1.078x          8.7        1.118x
8   32768      4    2048 |    11.0        13.0  1.179x         12.5        1.137x
16   32768      4    4096 |    17.5        21.9  1.251x         20.8        1.188x
32   32768      4    8192 |    31.9        40.1  1.257x         38.7        1.213x

1   65536      4     512 |     5.7         6.3  1.103x          6.2        1.091x
2   65536      4    1024 |     7.8         8.4  1.083x          8.8        1.134x
4   65536      4    2048 |    11.0        13.0  1.186x         12.5        1.141x
8   65536      4    4096 |    17.4        21.8  1.257x         21.3        1.230x
16   65536      4    8192 |    31.8        39.9  1.253x         38.4        1.206x
32   65536      4   16384 |    56.6        76.5  1.353x         69.7        1.233x

FP16 epilogue provides additional 3-7% speedup over FP32 on top of the DSL baseline

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

…M100 Replace DeepGEMM-based indexer logits with a CuTE DSL kernel on SM100+, gated by `use_cute_dsl_logits` config flag. Includes kernel implementation, PyTorch custom op registration, config plumbing, and unit tests. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

Replace stale development script name with the actual module filename. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

coderabbitai · 2026-04-20T12:00:34Z

📝 Walkthrough

Walkthrough

This change introduces a new CuTE DSL-based FP8 paged MQA logits kernel for SM100+ GPUs as an alternative to the existing C++ implementation. Configuration, custom op registration, kernel integration in the sparse attention indexer, and comprehensive test coverage are added.

Changes

Cohort / File(s)	Summary
Configuration & Arguments `tensorrt_llm/llmapi/llm_args.py`, `tensorrt_llm/_torch/model_config.py`	Added `use_cute_dsl_logits` boolean field to `DeepSeekSparseAttentionConfig` with default `False`, and forwarded it through model configuration initialization paths.
Custom Operations `tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`	Implemented new `trtllm::cute_dsl_fp8_paged_mqa_logits` custom op via `CuteDSLPagedMQALogitsRunner` with kernel compilation/caching, tensor reshaping, stream management, and fake registration for shape inference.
Kernel Package `tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py`	Added package initializer re-exporting `FP8MQALogitsDGFullKKernel` from the `fp8_paged_mqa_logits` module.
Sparse Attention Integration `tensorrt_llm/_torch/attention_backend/sparse/dsa.py`	Added conditional decode-phase logits path: when `use_cute_dsl_logits` flag is enabled and SM≥100, calls the CuTE DSL custom op; otherwise retains the existing `fp8_paged_mqa_logits()` path.
Tests & Benchmarks `tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py`	Added comprehensive test/benchmark suite validating CuTE DSL logits kernel correctness against DeepGEMM or PyTorch reference, with parametrization over batch sizes, `next_n`, context lengths, output dtypes, and kernel profiling utilities.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 59.09% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description is well-structured and comprehensive, including a clear summary, key design details, test plan with commands, and detailed benchmark results on B200.
Title check	✅ Passed	The title accurately describes the main change: introducing a CuTE DSL FP8 paged MQA logits decode kernel for Blackwell SM100, which is the primary feature across all modified files.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/model_config.py (1)
526-548: ⚠️ Potential issue | 🟡 Minor

Add the required NVIDIA SPDX header to this modified file.

This file is modified in this PR but still has no NVIDIA copyright header at the top. Please add/update it in the same change so the file stays repo-compliant.

As per coding guidelines, **/*.{h,hpp,cpp,cc,cxx,cu,py}: All TensorRT-LLM source files must contain an NVIDIA copyright header with the year of latest meaningful modification.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/model_config.py` around lines 526 - 548, Add the required
NVIDIA SPDX copyright header to the top of tensorrt_llm/_torch/model_config.py
(the file modified in this PR) by inserting the standard NVIDIA SPDX header
block including the correct copyright year (use the year of the latest
meaningful modification) and SPDX identifier; ensure the header appears before
any imports or code so the module (including classes/functions like
DeepSeekSparseAttentionConfig and the surrounding kwargs handling) is
repo-compliant and save the file with that header included.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 1114-1117: The logits path can call
torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits before the custom op is
registered because cute_dsl_custom_ops is only imported under the independent
use_cute_dsl_topk branch; to fix this, ensure the module is explicitly imported
when use_cute_dsl_logits is true (before invoking
torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits) by adding an import of
cute_dsl_custom_ops from tensorrt_llm._torch.custom_ops in the branch guarded by
self.use_cute_dsl_logits (so use_cute_dsl_logits, cute_dsl_custom_ops, and the
call to torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits are all referenced).

In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 5290-5317: The cute_dsl_fp8_paged_mqa_logits wrapper must validate
the GPU SM/architecture before dispatch like the other Blackwell-only ops;
update the function (or immediately before calling
CuteDSLPagedMQALogitsRunner.forward) to check the current device's SM/compute
capability (reuse the same helper or check used by other *_blackwell wrappers)
and raise a clear, early exception if the GPU is not Blackwell (or not
supported), rather than proceeding to call CuteDSLPagedMQALogitsRunner.forward;
ensure the error message references this op name and suggests disabling
use_cute_dsl_logits or using a supported GPU.
- Around line 5267-5288: The code creates a fresh stream (torch_stream =
torch.cuda.Stream()) and synchronizes the device, causing data races and a
global barrier; instead use the caller's current CUDA stream: replace creating a
new torch.cuda.Stream() with torch.cuda.current_stream() (so torch_stream
variable refers to torch.cuda.current_stream()), keep constructing the
cuda.CUstream from that stream, pass that stream into the compiled kernel launch
(the compiled(*dl_args, num_phys_blocks, B, stream) call), and remove the
unconditional torch.cuda.synchronize() so we don't impose a device-wide barrier;
ensure the compile/cache logic (cls._compile, cls.kernel_cache,
cls._make_dlpacks, compiled) still works with the caller stream and let callers
manage any required synchronization.

In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py`:
- Line 16: Add an explicit module export list to document the public API: update
the package __init__ by defining __all__ = ["FP8MQALogitsDGFullKKernel"] so the
re-export of FP8MQALogitsDGFullKKernel from fp8_paged_mqa_logits becomes
explicit and avoids F401/implicit API issues; keep the existing import statement
(from .fp8_paged_mqa_logits import FP8MQALogitsDGFullKKernel) and add the
__all__ declaration in the same module.

In
`@tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py`:
- Around line 185-187: The current skip marker skip_if_unsupported gates the
whole test on has_deep_gemm(), which wrongly prevents running the CuTe DSL
kernel when DeepGEMM is absent; change the skip condition to only require
IS_CUTLASS_DSL_AVAILABLE so tests run when CuTe DSL/hardware is present (i.e.,
skip_if_unsupported = pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE,
reason="Requires CuTe DSL/hardware")); then update the test’s reference
selection to pick the best available reference at runtime (use has_deep_gemm()
to choose the DeepGEMM reference if available, otherwise fall back to the pure
PyTorch reference) so the test exercises the CuTe DSL kernel while still
comparing against a valid fallback.

---

Outside diff comments:
In `@tensorrt_llm/_torch/model_config.py`:
- Around line 526-548: Add the required NVIDIA SPDX copyright header to the top
of tensorrt_llm/_torch/model_config.py (the file modified in this PR) by
inserting the standard NVIDIA SPDX header block including the correct copyright
year (use the year of the latest meaningful modification) and SPDX identifier;
ensure the header appears before any imports or code so the module (including
classes/functions like DeepSeekSparseAttentionConfig and the surrounding kwargs
handling) is repo-compliant and save the file with that header included.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 12b7da53-887c-4a0e-b1e7-fe0d6e847a18

📥 Commits

Reviewing files that changed from the base of the PR and between 04915ad and 69d7170.

📒 Files selected for processing (7)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/__init__.py
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp8_paged_mqa_logits.py
tensorrt_llm/_torch/model_config.py
tensorrt_llm/llmapi/llm_args.py
tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py

- Use current_stream() instead of creating a new stream to avoid data races - Remove unnecessary torch.cuda.synchronize() - Add is_sm_100f() check in custom op to fail fast on unsupported GPUs - Rename FP8MQALogitsDGFullKKernel to FP8MQALogitsKernel - Add __all__ to paged_mqa_logits __init__.py - Fix test skip logic to precisely match SM 100/103 only Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…s enabled Broaden the cute_dsl_custom_ops import guard to also trigger when use_cute_dsl_logits is True, so the custom op is registered even when use_cute_dsl_topk is False. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-04-21T04:52:44Z

/bot run

…_logits Align config field name with the op name cute_dsl_fp8_paged_mqa_logits for clarity. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

tensorrt-cicd · 2026-04-21T04:58:47Z

PR_Github #44629 [ run ] triggered by Bot. Commit: 69d7170 Link to invocation

The custom op already validates SM 100/103 via is_sm_100f(). Also replace getattr with direct attribute access to match use_cute_dsl_topk style. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…ser-facing strings Remove DeepGEMM and DG-FullK references from module docstring, class docstring, function docstrings, print statements, and argparse descriptions. Inline implementation comments retained as design cross-references. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-04-21T05:13:42Z

/bot run

tensorrt-cicd · 2026-04-21T05:19:24Z

PR_Github #44637 [ run ] triggered by Bot. Commit: cf418ab Link to invocation

tensorrt-cicd · 2026-04-21T05:19:28Z

PR_Github #44629 [ run ] completed with state ABORTED. Commit: 69d7170

Link to invocation

limin2021 · 2026-04-21T05:37:30Z

/bot run

tensorrt-cicd · 2026-04-21T05:43:04Z

PR_Github #44645 [ run ] triggered by Bot. Commit: cf418ab Link to invocation

tensorrt-cicd · 2026-04-21T05:43:08Z

PR_Github #44637 [ run ] completed with state ABORTED. Commit: cf418ab

Link to invocation

limin2021 · 2026-05-06T00:45:43Z

/bot run

tensorrt-cicd · 2026-05-06T00:53:01Z

PR_Github #46879 [ run ] triggered by Bot. Commit: 9db90f0 Link to invocation

tensorrt-cicd · 2026-05-06T03:34:58Z

PR_Github #46879 [ run ] completed with state FAILURE. Commit: 9db90f0
/LLM/main/L0_MergeRequest_PR pipeline #36890 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

limin2021 · 2026-05-06T07:12:03Z

/bot run

limin2021 · 2026-05-06T07:47:06Z

/bot run

tensorrt-cicd · 2026-05-06T07:53:53Z

PR_Github #46943 [ run ] triggered by Bot. Commit: 9db90f0 Link to invocation

tensorrt-cicd · 2026-05-06T11:59:07Z

PR_Github #46943 [ run ] completed with state SUCCESS. Commit: 9db90f0
/LLM/main/L0_MergeRequest_PR pipeline #36945 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

limin2021 · 2026-05-07T02:14:09Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-07T02:20:43Z

PR_Github #47078 [ run ] triggered by Bot. Commit: 9db90f0 Link to invocation

tensorrt-cicd · 2026-05-07T10:23:31Z

PR_Github #47078 [ run ] completed with state SUCCESS. Commit: 9db90f0
/LLM/main/L0_MergeRequest_PR pipeline #37049 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…xer_gemm Resolve conflicts in DSA paged-MQA-logits dispatch and tests after DeepGEMM submodule bump (4ff3f54d -> c491439e via PR NVIDIA#13340 / DG NVIDIA#304): - dsa.py: take upstream's scheduler_metadata_buffer / _full_next_n selection (mtp3 buffer removed); add DSL early-branch using the existing scheduler_metadata_buffer (built with (num_gen, 1) shape, num_atoms=1, matching DSL's 1-atom-per-q design) and the 1D kv_lens_cuda_runtime slice for context_lens. - dsa.py: introduce module-level _DG_SCHEDULE_BLOCK_KV = 64, used by all 6 get_paged_mqa_logits_metadata calls (3 in on_update_kv_lens(), 3 in Indexer.prepare()) instead of cache tokens_per_block. Decouples schedule SPLIT_KV from cache page size and side-steps a SM100 + block_kv=32 latent regression in DG commit 7f2a703 (NVIDIA#304). - test_dsa_indexer.py: take upstream's scheduler buffer selection; DSL test branch reads scheduler_metadata_buffer + 1D kv_lens. - test_cute_dsl_fp8_paged_mqa_logits.py: 4 metadata calls now pass 2D context_lens via .unsqueeze(-1) and DG_METADATA_BLOCK_KV=64; DG bench drops cluster(2,1,1) for next_n=4 (SM100 always uses num_kv_multicast=1) and passes 2D context_lens to fp8_paged_mqa_logits. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-05-07T16:37:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-07T16:43:56Z

PR_Github #47224 [ run ] triggered by Bot. Commit: 0a84ecc Link to invocation

tensorrt-cicd · 2026-05-08T04:07:58Z

PR_Github #47224 [ run ] completed with state SUCCESS. Commit: 0a84ecc
/LLM/main/L0_MergeRequest_PR pipeline #37179 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

… MQA next_n coverage test_dsa_indexer: separate scheduler_metadata_buffer for DSL backend (kNumNextNAtoms=1) to avoid the next_n>1 alias used by DeepGEMM. Fixes test_indexer_decode_with_paged_kv_cache across {deepgemm, dsl} x {(4,1),(2,2),(4,3),(4,4)}. test_cute_dsl_fp8_paged_mqa_logits: extend multi_block next_n coverage to {1,2,3,4}; drop now-obsolete next_n==3 expansion in bench (newer DeepGEMM supports it natively); update --block_kv default from 128 to 64 to match the new {32,64} assertion in fp8_paged_mqa_logits. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-05-08T05:25:41Z

/bot run

tensorrt-cicd · 2026-05-08T05:32:15Z

PR_Github #47317 [ run ] triggered by Bot. Commit: 4bcbd79 Link to invocation

tensorrt-cicd · 2026-05-08T10:34:00Z

PR_Github #47317 [ run ] completed with state SUCCESS. Commit: 4bcbd79
/LLM/main/L0_MergeRequest_PR pipeline #37258 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

limin2021 · 2026-05-08T10:36:25Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-08T10:42:43Z

PR_Github #47386 [ run ] triggered by Bot. Commit: 4bcbd79 Link to invocation

tensorrt-cicd · 2026-05-08T17:03:04Z

PR_Github #47386 [ run ] completed with state SUCCESS. Commit: 4bcbd79
/LLM/main/L0_MergeRequest_PR pipeline #37316 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

limin2021 · 2026-05-09T00:10:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-09T00:15:35Z

PR_Github #47451 [ run ] triggered by Bot. Commit: 4bcbd79 Link to invocation

tensorrt-cicd · 2026-05-09T01:06:37Z

PR_Github #47451 [ run ] completed with state SUCCESS. Commit: 4bcbd79
/LLM/main/L0_MergeRequest_PR pipeline #37372 completed with status: 'SUCCESS'

CI Report

Link to invocation

…VIDIA#13219) Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 added 4 commits April 20, 2026 04:42

[None][fix] Fix docstring script filename in fp8_paged_mqa_logits.py

c32db8f

Replace stale development script name with the actual module filename. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

[None][fix] Update copyright year in paged_mqa_logits __init__.py

b235eb1

Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

[None][fix] Remove unused max_mem_gb argument from benchmark

69d7170

Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 requested review from a team as code owners April 20, 2026 11:51

limin2021 requested review from HuiGao-NV, Superjomn, brb-nv and hyukn April 20, 2026 11:51

github-actions Bot assigned limin2021 Apr 20, 2026

coderabbitai Bot reviewed Apr 20, 2026

View reviewed changes

limin2021 added 2 commits April 20, 2026 21:41

limin2021 requested review from Kefeng-Duan and benzh-2025 April 21, 2026 04:53

[None][refactor] Rename use_cute_dsl_logits to use_cute_dsl_paged_mqa…

fe637c8

…_logits Align config field name with the op name cute_dsl_fp8_paged_mqa_logits for clarity. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 added 2 commits April 20, 2026 22:00

[None][fix] Remove redundant SM version check from DSA logits config

8f9b0e1

The custom op already validates SM 100/103 via is_sm_100f(). Also replace getattr with direct attribute access to match use_cute_dsl_topk style. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 requested a review from lfr-0531 May 7, 2026 16:38

hyukn enabled auto-merge (squash) May 9, 2026 04:02

hyukn merged commit c0d86d0 into NVIDIA:main May 9, 2026
7 checks passed

ly865623 mentioned this pull request May 11, 2026

Attention daily report — 2026-05-10 ly865623/AI_report#2

Merged

3 tasks

yufeiwu-nv pushed a commit to yufeiwu-nv/TensorRT-LLM that referenced this pull request May 19, 2026

[TRTLLM-34871][feat] Add cute dsl FP8 paged MQA logits decode kernel (N…

8d2fc0e

…VIDIA#13219) Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

b8zhong mentioned this pull request Jun 27, 2026

Cute-DSL FP8 MQA logits sgl-project/sglang#25220

Merged

Uh oh!

Conversation

limin2021 commented Apr 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Apr 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

limin2021 commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

limin2021 commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

limin2021 commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

tensorrt-cicd commented Apr 21, 2026

Uh oh!

limin2021 commented May 6, 2026

Uh oh!

tensorrt-cicd commented May 6, 2026

Uh oh!

tensorrt-cicd commented May 6, 2026

Uh oh!

limin2021 commented May 6, 2026

Uh oh!

limin2021 commented May 6, 2026

Uh oh!

tensorrt-cicd commented May 6, 2026

Uh oh!

tensorrt-cicd commented May 6, 2026

Uh oh!

limin2021 commented May 7, 2026

Uh oh!

tensorrt-cicd commented May 7, 2026

Uh oh!

tensorrt-cicd commented May 7, 2026

Uh oh!

limin2021 commented May 7, 2026

Uh oh!

tensorrt-cicd commented May 7, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

limin2021 commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

limin2021 commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

tensorrt-cicd commented May 8, 2026

Uh oh!

limin2021 commented May 9, 2026

Uh oh!

limin2021 commented Apr 20, 2026 •

edited

Loading

coderabbitai Bot commented Apr 20, 2026 •

edited

Loading