[TRTLLM-35237][feat] Tune cute dsl paged MQA logits decode kernel by limin2021 · Pull Request #14133 · NVIDIA/TensorRT-LLM

limin2021 · 2026-05-14T09:26:17Z

Summary

Replace the legacy HBM-minimizing atom-split picker (_pick_fp4_dsl_expand)
with a wave-aware strategy that minimizes scheduler waves and tie-breaks by
max num_atoms (more SMs busy per wave). Generalize the picker to both FP4
and FP8 DSL paths via a kernel_atoms parameter (FP4: {1,2,3}; FP8:
{1,2,3,4}).
fix bug: https://nvbugspro.nvidia.com/bug/6186880

Motivation

At low batch / low ntask (typical small-B decode), the old picker chose
kernel-native next_n, leaving SMs idle (waves=1 but ntask < 148 on B200).
The new picker splits into more atoms so all SMs work in the single wave,
paying expand_factor x KV HBM re-read cost. On wave-saturated workloads,
both pickers converge — the change is a no-op.

FP8 production previously never atom-split (the legacy flag was gated on
use_fp4); the new generalized picker brings FP8 to parity with FP4.

Speedup vs legacy picker

performance ut

python tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
python tests/unittest/_torch/attention/sparse/test_cute_dsl_fp8_paged_mqa_logits.py

Path	next_n	`max_atom / max_na`
FP4	2	+2–9 %
FP4	3	+18–25 %
FP4	4	+2–10 %
FP8	2	+5–8 %
FP8	3	+15–23 %
FP8	4	+16–32 %

Full performance

FP4 indexer

epi_dtype=float32 output_dtype=bfloat16 num_epi_subtiles=1 mode=fix-len block_kv=64

next_n = 2

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	4	1/2	5.1	2/1	4.7	1.089x	5.2	1.034x	1.126x
2	1024	32	8	1/2	4.9	2/1	4.6	1.052x	5.4	1.099x	1.155x
4	1024	64	16	1/2	5.0	2/1	4.6	1.092x	5.4	1.073x	1.172x
8	1024	128	32	1/2	5.1	2/1	4.9	1.035x	5.3	1.056x	1.093x
16	1024	256	64	1/2	5.3	2/1	5.0	1.046x	5.4	1.035x	1.083x
1	2048	32	8	1/2	5.0	2/1	4.6	1.085x	5.4	1.079x	1.171x
2	2048	64	16	1/2	4.9	2/1	4.7	1.046x	5.4	1.094x	1.144x
4	2048	128	32	1/2	5.0	2/1	4.8	1.046x	5.4	1.074x	1.124x
8	2048	256	64	1/2	5.1	2/1	5.0	1.034x	5.4	1.046x	1.082x
1	4096	64	16	1/2	4.9	2/1	4.6	1.062x	5.2	1.071x	1.137x
2	4096	128	32	1/2	5.1	2/1	5.0	1.020x	5.4	1.053x	1.074x
4	4096	256	64	1/2	5.1	2/1	5.0	1.025x	5.5	1.068x	1.095x
1	8192	128	32	1/2	5.0	2/1	4.7	1.051x	5.4	1.081x	1.136x
1	16384	256	64	1/2	5.2	2/1	5.0	1.039x	5.5	1.063x	1.104x

next_n = 3

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	4	1/3	5.8	3/1	4.7	1.227x	4.9	0.858x	1.053x
2	1024	32	8	1/3	5.6	3/1	4.7	1.194x	4.8	0.858x	1.025x
4	1024	64	16	1/3	5.8	3/1	4.7	1.234x	5.1	0.868x	1.071x
8	1024	128	32	1/3	5.9	3/1	4.7	1.258x	5.1	0.858x	1.079x
1	2048	32	8	1/3	5.7	3/1	4.5	1.251x	4.8	0.842x	1.053x
2	2048	64	16	1/3	5.8	3/1	4.7	1.226x	4.7	0.814x	0.997x
4	2048	128	32	1/3	5.8	3/1	4.9	1.179x	4.9	0.851x	1.003x
1	4096	64	16	1/3	5.7	3/1	4.7	1.215x	4.9	0.868x	1.054x
2	4096	128	32	1/3	5.8	3/1	4.8	1.199x	4.9	0.854x	1.024x
1	8192	128	32	1/3	5.8	3/1	4.7	1.252x	4.9	0.842x	1.054x

next_n = 4

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	8	2/2	4.9	4/1	4.5	1.096x	4.9	1.012x	1.109x
2	1024	32	16	2/2	4.9	4/1	4.8	1.020x	5.1	1.046x	1.067x
4	1024	64	32	2/2	5.0	4/1	4.7	1.063x	5.4	1.072x	1.139x
8	1024	128	64	2/2	5.1	4/1	4.7	1.079x	5.4	1.056x	1.139x
1	2048	32	16	2/2	4.8	4/1	4.5	1.062x	5.4	1.112x	1.181x
2	2048	64	32	2/2	4.9	4/1	4.7	1.034x	5.2	1.070x	1.106x
4	2048	128	64	2/2	5.0	4/1	4.9	1.030x	5.3	1.049x	1.081x
1	4096	64	32	2/2	4.9	4/1	4.6	1.067x	5.1	1.036x	1.105x
2	4096	128	64	2/2	5.0	4/1	4.7	1.071x	5.2	1.049x	1.124x
1	8192	128	64	2/2	5.1	4/1	4.9	1.043x	5.3	1.047x	1.091x

FP8 indexer

output_dtype=float32 num_epi_subtiles=1 mode=fix-len block_kv=64

next_n = 2

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	4	1/2	4.1	2/1	3.8	1.074x	4.2	1.022x	1.097x
2	1024	32	8	1/2	4.1	2/1	3.8	1.073x	4.2	1.033x	1.109x
4	1024	64	16	1/2	4.1	2/1	3.8	1.078x	4.2	1.034x	1.116x
8	1024	128	32	1/2	4.2	2/1	3.9	1.080x	4.2	1.008x	1.089x
16	1024	256	64	1/2	4.2	2/1	3.9	1.075x	4.3	1.022x	1.099x
1	2048	32	8	1/2	4.0	2/1	3.8	1.062x	4.1	1.028x	1.092x
2	2048	64	16	1/2	4.1	2/1	3.8	1.068x	4.2	1.028x	1.098x
4	2048	128	32	1/2	4.1	2/1	3.8	1.063x	4.2	1.026x	1.090x
8	2048	256	64	1/2	4.2	2/1	3.9	1.081x	4.3	1.030x	1.113x
1	4096	64	16	1/2	4.1	2/1	3.8	1.067x	4.2	1.022x	1.091x
2	4096	128	32	1/2	4.1	2/1	3.8	1.071x	4.2	1.048x	1.122x
4	4096	256	64	1/2	4.1	2/1	3.9	1.062x	4.2	1.025x	1.089x
1	8192	128	32	1/2	4.1	2/1	3.8	1.066x	4.2	1.018x	1.086x
2	8192	256	64	1/2	4.1	2/1	3.9	1.062x	4.2	1.031x	1.095x
1	16384	256	64	1/2	4.1	2/1	3.9	1.057x	4.2	1.027x	1.086x

next_n = 3

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	4	1/3	4.4	3/1	3.8	1.164x	4.0	0.896x	1.043x
2	1024	32	8	1/3	4.4	3/1	3.8	1.155x	4.0	0.900x	1.039x
4	1024	64	16	1/3	4.7	3/1	3.8	1.228x	4.0	0.850x	1.044x
8	1024	128	32	1/3	4.5	3/1	3.8	1.176x	4.0	0.884x	1.039x
1	2048	32	8	1/3	4.5	3/1	3.8	1.184x	4.0	0.877x	1.039x
2	2048	64	16	1/3	4.5	3/1	3.8	1.170x	4.0	0.886x	1.037x
4	2048	128	32	1/3	4.5	3/1	3.8	1.171x	4.0	0.895x	1.048x
1	4096	64	16	1/3	4.5	3/1	3.8	1.179x	4.0	0.886x	1.045x
2	4096	128	32	1/3	4.7	3/1	3.9	1.200x	4.0	0.858x	1.029x
1	8192	128	32	1/3	4.5	3/1	3.8	1.166x	4.0	0.904x	1.054x

next_n = 4

batch	ctx	nblk	ntask	maxAtom	DSL (μs)	maxNa	DSL (μs)	max_atom / max_na	DG (μs)	DG / DSL_max_atom	DG / DSL_max_na
1	1024	16	4	1/4	4.8	4/1	3.8	1.257x	4.3	0.898x	1.128x
2	1024	32	8	1/4	4.8	4/1	3.8	1.282x	4.3	0.879x	1.127x
4	1024	64	16	1/4	4.9	4/1	3.8	1.275x	4.4	0.896x	1.143x
8	1024	128	32	1/4	4.9	4/1	3.9	1.263x	4.3	0.877x	1.107x
16	1024	256	64	1/4	4.9	2/2	4.2	1.160x	4.4	0.900x	1.044x
1	2048	32	8	1/4	4.9	4/1	3.8	1.267x	4.2	0.873x	1.106x
2	2048	64	16	1/4	5.0	4/1	3.8	1.315x	4.3	0.861x	1.132x
4	2048	128	32	1/4	4.9	4/1	3.9	1.272x	4.3	0.881x	1.120x
8	2048	256	64	1/4	4.9	2/2	4.2	1.177x	4.3	0.877x	1.033x
1	4096	64	16	1/4	4.9	4/1	3.8	1.305x	4.2	0.855x	1.116x
2	4096	128	32	1/4	4.9	4/1	3.8	1.274x	4.3	0.875x	1.115x
4	4096	256	64	1/4	4.9	2/2	4.2	1.180x	4.3	0.879x	1.037x
1	8192	128	32	1/4	4.9	4/1	3.9	1.270x	4.3	0.868x	1.102x
2	8192	256	64	1/4	4.9	2/2	4.2	1.179x	4.4	0.890x	1.050x
1	16384	256	64	1/4	5.1	2/2	4.2	1.215x	4.3	0.846x	1.027x

Changes

dsa.py
- _pick_fp4_dsl_expand → _pick_dsl_expand; adds batch_size,
  max_ctx, num_sms, kernel_atoms params; wave-aware decision with
  legacy max-atom fallback when batch/ctx unknown.
- Metadata fields renamed (drop fp4_ prefix): expand_for_dsl,
  dsl_expand_factor, dsl_atom. Decision is cached at prepare time;
  forward + Indexer.prepare read from cache (no per-call re-pick;
  prepare/forward guaranteed consistent).
- Trigger relaxed from max_draft_tokens >= 3 to >= 1, so next_n ∈ {2, 3} can also expand when the picker chooses factor > 1.
- FP4 forward: gate on metadata.dsl_expand_factor > 1 instead of
  next_n > 3.
- FP8 forward: add caller-side expand branch (previously absent).
test_cute_dsl_fp{4,8}_paged_mqa_logits.py: bench harness adds an
A/B comparison of maxAtom (baseline) vs maxNa (new) per row, plus
DG/DSL_max_atom, DG/DSL_max_na, and ntask columns.
test_dsa_indexer.py::test_indexer_decode_with_paged_kv_cache_fp4:
adapt to new picker signature and field names; mock expand setup runs
unconditionally for DSL backend so next_n ∈ {2, 3} paths are exercised.

Validation (4.4.2 PyPI wheel, B200)

FP4 paged MQA logits sweep: 810/810 PASS
FP8 paged MQA logits sweep: 176/176 PASS
test_dsa_indexer.py::test_indexer_decode_with_paged_kv_cache_fp4: 8/8 PASS

Description

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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>

- 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>

…_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>

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>

Replace from_dlpack + mark_compact_shape_dynamic compile pattern with make_fake_compact_tensor + TVM FFI in CuteDSLPagedMQALogitsRunner. This eliminates real tensor data at compile time, removes per-call dlpack wrapping and manual CUDA stream passing at runtime. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

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

…pper Add _check_fp8_paged_mqa_logits_dtypes helper that validates all tensor and compute dtypes upfront, collecting all errors into a single ValueError so callers see every mismatch at once rather than one at a time. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

- Add epi_dtype param to PyTorch ref so epilogue matches kernel precision - Add use_int_data mode for fp16 tests to isolate bugs from FP16 rounding - Tighten tolerances: fp32 atol 5e-3→5e-5, fp16 atol 1e-2→1e-3 with rtol - Use torch.testing.assert_close instead of cosine similarity - Add varlen benchmark support for mixed-length serving workloads - Clean up comments and remove stale parametrize line Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…e map Remove broken run_test/parse_args/__main__ block and its helper functions (make_fused_kv, fused_kv_views, _prepare_inputs, _make_dynamic_dlpacks, _get_or_compile_kernel, dsl_fp8_paged_mqa_logits_dg_fullk) from fp8_paged_mqa_logits.py — all had no callers after the missing paged_mqa_logits_helpers module made them non-functional. Accuracy validation is covered by the unit tests. Reuse existing module-level _TORCH_TO_CUTLASS_DTYPE in CuteDSLPagedMQALogitsRunner instead of duplicating it. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…ubtiles Move the FMA unroll granularity check outside the num_epi_subtiles > 1 guard so it also applies when num_epi_subtiles == 1. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

Remove is_umma_warp, num_q_stages (local copy), cons_group, num_tma_prod, num_mcast_a/b, and tCtAcc_fake (in __call__) which were assigned but never referenced. Remove unnecessary noqa F401 on pipeline_init_arrive/pipeline_init_wait which are actively used. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

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

…arbitrary next_n Fix cudaErrorMisalignedAddress when num_heads=32 with fp16 output by padding sW stage stride to 128-byte boundary for TMA bulk copy. Update qw_per_stage SMEM budget to use padded stride and actual epi_bytes. In dsa.py, skip MTP batch expansion for DSL kernel (supports arbitrary next_n natively) and use correct scheduler_metadata buffer selection. Expand test coverage with num_heads and fix_length parametrization, pure PyTorch reference, and DG expanded benchmark for next_n=3. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

Clamp start_q before reading mContextLens to prevent out-of-bounds access when the scheduler assigns batch_size as a sentinel for CTAs with no work. Move torch.manual_seed into _generate_test_data for reproducibility. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…aged MQA logits Decouple compute_block_kv (always 128) from phys_block_kv (physical page size). When phys_block_kv < 128, the kernel issues num_blocks_per_mma TMA copies per compute tile. Removes the tokens_per_block=128 constraint for the DSL kernel path. Signed-off-by: Mindy Li <lmin@nvidia.com> Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…eduling Hoist shuffle_sync calls before producer_acquire in both TMA warpgroups. The previous ordering placed shuffle after barrier acquire, causing a dependency chain that regressed next_n=3 by 5-7%. SASS now matches baseline exactly. Signed-off-by: Mindy Li <lmin@nvidia.com> Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

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

…and test Drop dead code flagged in review: - compute_schedule_metadata + cdiv (and now-unused torch import) from the kernel module — never called. - skip_if_unsupported decorator + has_deep_gemm helper + IS_CUTLASS_DSL_AVAILABLE import from the test — defined but never applied. SM100/103 gate already ensures DeepGEMM and CuTe DSL availability. Also refresh the test docstring: the reference is pure PyTorch, not DeepGEMM. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

The DSL variant of test_indexer_decode_with_paged_kv_cache calls torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits, which raises ValueError on SM != 100/103. Mark only the "dsl" parameter with a skipif so the "deepgemm" backend still runs on Hopper. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

…xer_gemm Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

…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>

… 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>

…ress) Snapshot of FP4 MXFP4 block-scaled paged_mqa_logits kernel based on the FP8 kernel: 5 TMA atoms (KV + B + W + SF_KV + SF_Q), per-WG SFA TMEM region, in-place SMEM transpose for UTCCP, per-k_block SFA/SFB slice for cute.gemm. Stage 0 still hangs at tmem.wait_for_alloc bar.sync 1, 320; committing state for further debugging. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

tensorrt-cicd · 2026-05-19T09:41:52Z

PR_Github #49045 [ run ] completed with state ABORTED. Commit: ac48ef7

Link to invocation

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

on_update_kv_lens() previously only refreshed kv_lens_expanded_cuda and scheduler_metadata_buffer_expanded under use_expanded_buffers_for_mtp. With DSL paged-mqa-logits and dsl_expand_factor > 1, overlap scheduler + spec-dec runtime corrections to kv_lens_cuda left both buffers stale, so the kernel scanned a pre-correction KV extent and the schedule referenced pre-correction lengths. Mirror the prepare()-time DSL split block next to the existing MTP branch. Reuse cached (dsl_expand_factor, dsl_atom) — re-running the picker here would let the split drift between prepare and forward and break CUDA graph capture. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-05-19T14:20:47Z

/bot run

tensorrt-cicd · 2026-05-19T14:29:22Z

PR_Github #49206 [ run ] triggered by Bot. Commit: aac6e9d Link to invocation

tensorrt-cicd · 2026-05-19T14:34:30Z

PR_Github #49168 [ run ] completed with state ABORTED. Commit: d87f661

Link to invocation

tensorrt-cicd · 2026-05-19T18:04:00Z

PR_Github #49206 [ run ] completed with state SUCCESS. Commit: aac6e9d
/LLM/main/L0_MergeRequest_PR pipeline #38881 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

mingyangHao

LGTM

lfr-0531

The changes in DSA part LGTM~

hyukn

LGTM.

- _pick_dsl_expand: make num_sms a required argument (drop magic default 148); both existing callers already pass it via keyword. - expand_for_dsl field comment: generalize from FP4-only to cover both FP4 ({1,2,3}) and FP8 ({1,2,3,4}) atom sets. - FP4 kernel: drop the print(...) fallback notice in remove_online_sf_transpose; silent fallback is sufficient. - FP4/FP8 kernels: remove commented-out phys_block_bytes = ... dead lines. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

limin2021 · 2026-05-21T05:39:54Z

/bot run

tensorrt-cicd · 2026-05-21T05:46:22Z

PR_Github #49597 [ run ] triggered by Bot. Commit: fed71fe Link to invocation

tensorrt-cicd · 2026-05-21T09:44:59Z

PR_Github #49597 [ run ] completed with state SUCCESS. Commit: fed71fe
/LLM/main/L0_MergeRequest_PR pipeline #39223 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

limin2021 · 2026-05-21T10:02:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T10:08:21Z

PR_Github #49676 [ run ] triggered by Bot. Commit: fed71fe Link to invocation

tensorrt-cicd · 2026-05-21T11:14:15Z

PR_Github #49676 [ run ] completed with state SUCCESS. Commit: fed71fe
/LLM/main/L0_MergeRequest_PR pipeline #39285 completed with status: 'SUCCESS'

CI Report

Link to invocation

PR NVIDIA#14133 updated the FP4 paged-MQA-logits CuTe DSL op to accept non-contiguous kv tensors, so the .contiguous() forcing in this custom op is no longer needed and was triggering an extra copy kernel. Revert to the original reshape-only path. Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>

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

limin2021 added 30 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>

[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>

[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>

[None][fix] Remove commented-out debug code from fp8_paged_mqa_logits

4ef3d13

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

[None][fix] Replace compile print with logger.debug in MQA logits

670cd16

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

[None][fix] Validate num_heads divisible by 4 regardless of num_epi_s…

d0fa1cb

…ubtiles Move the FMA unroll granularity check outside the num_epi_subtiles > 1 guard so it also applies when num_epi_subtiles == 1. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>

[None][fix] Add missing use_cute_dsl_paged_mqa_logits to test mock co…

719b333

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

[None][cleanup] Remove commented-out test_deepgemm_fp8_paged_mqa_logits

0331897

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

Merge branch 'main' into add_dsl_indexer_gemm

0ca0f27

Merge remote-tracking branch 'github-upstream/main' into add_dsl_inde…

9db90f0

…xer_gemm Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

limin2021 added 2 commits May 19, 2026 14:01

fix kv stride issue to not be compact

127a156

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

lfr-0531 requested a review from mingyangHao May 20, 2026 07:16

mingyangHao approved these changes May 20, 2026

View reviewed changes

lfr-0531 approved these changes May 20, 2026

View reviewed changes

hyukn approved these changes May 20, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/cute_dsl_kernels/blackwell/paged_mqa_logits/fp4_paged_mqa_logits.py Outdated

Comment thread tests/scripts/cute_dsl_kernels/paged_mqa_logits/run_fp4.py