[None][perf] Skip DeepGEMM clean_logits in DSA indexer prefill on custom top-k path

[dc3671]

Description

Background

The DeepGEMM MQA-logits kernels used by the DeepSeek-V3.2/V4 DSA indexer (fp8_mqa_logits / fp8_fp4_mqa_logits) default clean_logits=true. After computing the logits they launch an extra smxx_clean_logits kernel that fills every element outside each row's causal window [ks, ke) with -inf. The logits buffer is torch.empty with a padded, 1024-byte-aligned row stride, so this fill exists to stop consumers that scan the full padded row from selecting uninitialized garbage into the top-k.

The default prefill path (use_custom_topk=True) does not scan the full row. It uses the custom indexer_topk_prefill kernel (cpp/tensorrt_llm/kernels/indexerTopK.cu), whose topKPerRowPrefill reads strictly [rowStart, rowEnd) == [ks, ke) per row (rowStart = rowStarts[rowIdx], rowEnd = rowEnds[rowIdx], and every load in topKPerRowJob / processHistogramStep is bounded by that window). The padding outside the window is never read, so the -inf fill is pure waste on this path.

Only the torch-fallback path (use_custom_topk=False), which runs logits.topk over the full padded row, actually needs the fill.

Solution

Plumb a clean_logits flag through Indexer._call_mqa_logits and pass clean_logits=not use_custom_topk at both prefill call sites (the chunked-prefill tile loop and the single-pass fallback). Production runs (custom top-k, the default; no caller overrides it) skip the clean kernel; the torch fallback keeps the fill and stays correct. Decode is unaffected — the paged variants already default clean_logits=false.

Impact

• Correctness: neutral. The custom top-k kernel never reads outside [ks, ke), so skipping the fill cannot change its output.
• Performance: removes one memory-bandwidth-bound kernel per indexer layer per prefill chunk. The win is small (the kernel is tiny and PDL-overlapped), so this is best viewed as dead-work removal rather than a headline speedup.

Validation

Rubin, DeepSeek-V4-Pro, DEP8 (TP8/EP8 attention-DP) MTP3, 8k-in/1k-out, aggregated ctx-only trtllm-bench, nsys on rank 0:

kernel (rank-0 GPU kern_sum)	baseline (clean on)	this PR (clean off)
smxx_clean_logits	5.96 ms × 210, 0.11% GPU time	absent
fp8_fp4_mqa_logits	19.45 ms × 210	19.48 ms × 210
topKPerRowPrefill	32.18 + 19.46 ms	32.08 + 19.50 ms

The MQA-logits and top-k kernels are byte-identical across the two runs, i.e. top-k output is unchanged. End-to-end throughput is flat within run-to-run noise.

Test Coverage

Existing indexer top-k tests in tests/unittest/_torch/attention/sparse/ compare against a reference implementation and pass on both flag states, so no new tests are added:

pytest tests/unittest/_torch/attention/sparse/ -k indexer

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.

[coderabbitai]

Walkthrough

Indexer now controls DeepGEMM logits cleaning through a clean_logits parameter, with chunked and non-chunked prefill paths deriving it from use_custom_topk.

Changes

DSA logits cleaning

Layer / File(s)	Summary
Expose logits-cleaning control tensorrt_llm/_torch/attention_backend/sparse/dsa.py	_call_mqa_logits accepts clean_logits and forwards it to FP4 and FP8 DeepGEMM dispatches.
Wire cleaning to prefill top-k selection tensorrt_llm/_torch/attention_backend/sparse/dsa.py	Chunked and non-chunked prefill calls set clean_logits=not use_custom_topk.

Estimated code review effort: 3 (Moderate) | ~15–30 minutes

Suggested reviewers: kris1025

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: skipping DeepGEMM clean_logits on the custom top-k DSA prefill path.
Description check	✅ Passed	The description follows the template and includes background, solution, impact, validation, test coverage, and checklist status.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[dc3671]

/bot run

[tensorrt-cicd]

PR_Github #61260 [ run ] triggered by Bot. Commit: 2c674e6 Link to invocation

[tensorrt-cicd]

PR_Github #61260 [ run ] completed with state FAILURE. Commit: 2c674e6
/LLM/main/L0_MergeRequest_PR pipeline #49496 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