[None][fix] DSv4 indexer: stable radix aux scratch for CUDA Graph safety

[longcheng-nv]

Description

The fp32 Radix path of indexer_topk_decode was not CUDA-Graph-safe: it allocated its split-work scratch buffers (aux_indices, aux_logits) via th::empty per call when blocks_per_row > 1. Under CUDA Graph capture + replay these per-call pointers become stale when the caching allocator is perturbed by chunked-prefill activations at high concurrency. The captured Radix part1 / part2 kernels then write to recycled memory; the resulting CUDA_ERROR_ILLEGAL_ADDRESS surfaces at the next sync (commonly inside DeepGEMM's smxx_layout.hpp:97 TMA descriptor pack), looking like a downstream FP8 GEMM crash but caused upstream by the indexer.

IndexerTopKOp.cpp:91 already documents the pitfall for the Heuristic path:

"Caller-owned scratch buffer for heuristic TopK output values. Must be pre-allocated with stable address for CUDA Graph compatibility."

and dsa.py:879 follows that contract for heuristic_scratch_values. The Radix path was the missing half of the same pattern — this PR adds it.

Observed failure (Flash MXFP4, ISL ≈ 100K, TP = EP = 4, BS = 32, MTP = 0)

cell	CONC	Mode	GVR	result
cell_013_C008_TEP_GVRF	8	TEP	OFF (Radix)	CUDA 700 at smxx_layout.hpp:97
cell_017_C016_TEP_GVRF	16	TEP	OFF (Radix)	CUDA 700 at smxx_layout.hpp:97
cell_021_C032_TEP_GVRF	32	TEP	OFF (Radix)	CUDA 700 at smxx_layout.hpp:97

All three crashes converge on the identical stack: fp8_swap_ab_gemm → _fp8_quantize_1x128_ue8m0 → deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor → smxx_layout.hpp:97. The DeepGEMM line is the next sync after the offending kernel, not the offender itself. Same-point GVR-ON cells (Heuristic path, uses persistent heuristic_scratch) run clean — confirming the asymmetry is in the indexer.

Change

• cpp/tensorrt_llm/thop/IndexerTopKOp.cpp — extend indexer_topk_decode with two optional Tensor? kwargs (radix_aux_indices, radix_aux_logits). When both are provided and blocks_per_row > 1, use the caller-owned buffers; otherwise fall back to per-call th::empty (back-compat for bench scripts / warmup helpers / callers not under CUDA Graph capture). Validates is_cuda, same device, contiguous, dtype (int32 / float32), numel ≥ num_rows × blocks_per_row × index_topk.
• tensorrt_llm/_torch/attention_backend/sparse/dsa.py — unconditionally allocate persistent radix_aux_indices / radix_aux_logits in DSAtrtllmAttentionMetadata sized to the worst case (max_gen_tokens × kMaxBlocksPerRowDecode = 10 × num_sparse_topk). Buffers are routed through self.get_empty(self.cuda_graph_buffers, ...) so they have stable addresses across graph replays and survive metadata-level CUDA-graph capture. Passed at the call site.
• tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py — extend the register_fake Python signature to match the new schema.

Side-effect surface

• Same Radix kernel invocation, same dispatch logic, same numerical output.
• DSv3.2 sparse path shares this op and metadata class — benefits transparently, no regression.
• API back-compat: new kwargs default to None; existing callers continue to use the th::empty fallback unchanged.
• Memory cost: ≈ 1.3 MB per DSAtrtllmAttentionMetadata instance on Flash defaults (max_num_seqs = 32 × kMaxBlocksPerRowDecode = 10 × index_topk = 512 × 2 buffers × 4 B/elt). Negligible.

Test Coverage

Unit tests added to tests/unittest/_torch/thop/parallel/test_indexer_topk.py:

• test_indexer_topk_decode_radix_aux_equivalence — Parametrised over (batch_size, num_tokens, compress_ratio) shapes that force blocks_per_row > 1 (including the exact batch_size=8, num_tokens=16384, compress_ratio=4 that mirrors the Flash CONC=8 failure point). Verifies caller-owned-aux output is bit-identical to the legacy th::empty fallback.
• test_indexer_topk_decode_radix_aux_cuda_graph_replay — Captures the op into a CUDA Graph with caller-owned aux, replays 8× and asserts each replay's output matches a non-graph reference. This is the direct regression test for the original stale-pointer bug.
• test_indexer_topk_decode_radix_aux_validation — Negative tests for the new validation paths: wrong dtype (indices not int32 / logits not float32), undersized buffer (< num_rows × blocks_per_row × index_topk), non-contiguous tensor.

Existing coverage retained:

• test_indexer_topk_decode / test_indexer_topk_decode_sm_saturation / test_indexer_topk_decode_launch_policy_transitions continue to exercise the th::empty fallback path (do not pass the new kwargs).
• bench_indexer_topk_kernel.py continues to work unchanged (does not pass new kwargs).

End-to-end validation:

• Re-run of pareto_v4_pro_swebench100k_b300 sweep against this binary on 8× B300 is in progress; CONC = 1 paired cells (P1 / P2) already completed cleanly with the new fix in place. High-CONC GVR-OFF cells (the ones that previously triggered G4) will be the live test.

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.

[coderabbitai]

Caution

Review failed

Failed to post review comments

📝 Walkthrough

Walkthrough

This PR introduces DeepSeek-V4 KV cache compression kernels, MHC hyper-connection fused GPU kernels, NIXL transfer agent virtual memory region mapping, TopK decode compression support, and FP8 block-scale MoE enhancements across the TensorRT-LLM kernel and executor stacks, along with supporting CMake wiring and infrastructure updates.

Changes

KV Cache Compression System

Layer / File(s)	Summary
DeepSeek-V4 KV Cache Compression Kernels cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt, compressorKernels.h, compressorKernels.cu	Three main CUDA kernels implement online-softmax-based KV cache compression: pagedKvCompressKernel for decode-step token compression, prefillReductionKernel for bulk prefill compression, and postProcessScatterKernel for RMSNorm + RoPE + Hadamard butterfly + paged scatter with FP8/MXFP4 quantization support.

MHC Hyper-Connection Fused Kernels

Layer / File(s)	Summary
MHC Kernel Infrastructure cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt, mhcKernels.h, mhcKernels.cu	Core MHC kernel launchers: mhcBigFuseKernel for token-wise split-K reduction + Sinkhorn normalization + RMSNorm fusion, mhcGemmSqrsumFmaKernel for split-N GEMM + sqr-sum fusion, and post-mapping/head-apply kernels for residual and gating operations.
MHC TF32 Fused Kernels cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh, mhcFusedHcKernel.cu	TF32 tcgen05 MMA + TMA pipelined kernels: two-step mhcFusedHcLaunch (GEMM → bigFuse postlog) and all-in-one mhcFusedHcAllInOneLaunch with done-counter-based last-CTA election for fused postlogue execution.
MHC FMA Fused Kernels cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh, mhcFusedHcKernel.cu	Register-path FMA variants: fused_pmap_gemm_fma_ksplit for two-step K-split execution and fused_pmap_gemm_fma_allinone for single-launch pmap + GEMM + bigFuse with atomicAdd reduction when multi-split.

NIXL Transfer Agent VMM Region Mapping

Layer / File(s)	Summary
VMM Region Utilities and Agent Metadata cpp/include/tensorrt_llm/executor/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp	Introduces VramRegionInfo, VramRegionMap, and VmmDescSplitter utilities for backend-agnostic chunk-aware descriptor splitting. VmmDescSplitter provides static methods: lookupChunkInfo(), splitDescsWithRegionMap(), splitTransferDescsWithRegionMaps(), splitVmmDescs(), and detectVramRegionMap(). Extends AgentDesc with serialization/deserialization of per-region VMM metadata for remote boundary computation.
NIXL Transfer Status and Agent Lifecycle cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h, transferAgent.cpp	Adds RAII cleanup for NixlTransferStatus with destructor, status tracking via mLastStatus atomic, and accessors getLastStatus()/getLastStatusStr(). Extends NixlTransferAgent with shutdown() method and updates registerMemory(), deregisterMemory(), loadRemoteAgent(), getLocalAgentDesc(), and submitTransferRequests() to use VmmDescSplitter region maps instead of registry-based splitting. NixlLoopbackAgent gains a destructor for safe resource cleanup.
Python Bindings cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp	Adds nanobind methods serialize() and deserialize() for AgentDesc, and extends BaseTransferAgent and NixlTransferAgent bindings with nb::call_guard<nb::gil_scoped_release>() to release Python GIL during long-running transfer operations.

TopK Decode Compression Support

Layer / File(s)	Summary
Heuristic TopK Decode Compression cpp/tensorrt_llm/kernels/heuristicTopKDecode.h, heuristicTopKDecode.cu	Updates heuristic TopK kernels to accept compressRatio parameter, compute actual_kv_len from token window, and derive search domain N as actual_kv_len / compressRatio. Adjusts preIdxOffset temporal-shift logic to depend on compressRatio mode (nonzero offset for ratio==1, zero otherwise).
Indexer TopK Multi-Block Scheduling cpp/tensorrt_llm/kernels/IndexerTopK.h, indexerTopK.cu	Adds computeIndexerTopKDecodeBlocksPerRow() helper for SM-aware multi-block split/merge sizing. Introduces file-scope tuning constants (kNumBins, split-work thresholds). Updates topKPerRowDecode kernel to compute rowEnd as actual_kv_len / compressRatio. Tightens heuristic eligibility gating to allowed compressRatio modes (1 or 4). Refactors dispatch logic around computed blocksPerRow and fixes histogram refinement with explicit counter resets and garbage index initialization.

FP8 Block-Scale MoE and Routing Enhancements

Layer / File(s)	Summary
SwiGLU Clamping for FP8 Activation cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h, DevKernel.cu, runner.cu	Extends activation::Data and activation::KernelParams with optional swigluLimit and hasSwigluLimit fields. Both standard and DeepSeek FP8 kernels conditionally clamp gate/up inputs before activation computation. Runner forwards scalar clamp value for DeepSeek FP8 path.
DeepSeekV4 Routing Mode cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h, runner.cu	Adds DeepSeekV4 = 7 to RoutingMethodType enum (renumbering Unspecified to 8), adds serialization support, and extends MoERunnerArgs with gemm1_clamp_limit_value and has_gemm1_clamp_limit_value for scalar clamp-limit FP8 path.
Top-K Sort Specializations cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh	Adds explicit Sort<N, RedType> template specializations for N=5,6,7,8 using hardcoded TOPK_SWAP sequences to enable up-to-8-way expert selection.
FP8 Block-Scale GEMM Quantization cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h, fp8_blockscale_quant_packed.cu	Fused FP8 1x128 quantization kernel: loads BF16, computes per-block amax, derives UE8M0 scale, quantizes to FP8 e4m3, and packs four scale bytes into uint32 MN-major layout for block-scale GEMM dispatch.
Custom MoE Gate Forward Kernel cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h, customMoeRoutingKernels.cu	New gate_forward kernel computes per-expert score transform (softplus), performs top-K selection in hash or topK mode, normalizes weights, and dispatches for n_experts 256 and 384.

Infrastructure and Configuration Updates

Layer / File(s)	Summary
Build System Wiring cpp/tensorrt_llm/CMakeLists.txt, cpp/tensorrt_llm/kernels/CMakeLists.txt	Adds mhcKernels_src and compressorKernels_src subdirectories and object libraries to CMake build tree and links them into the main shared target.
Attention Workspace Sizing cpp/tensorrt_llm/common/attentionOp.h, attentionOp.cpp	Extends getWorkspaceSizeForContext() with total_kv_len parameter for KV cache reuse scenarios where effective total KV length exceeds max_num_tokens. Updates FP8 Context MLA K/V buffer sizing accordingly.
KV Cache Manager Utils cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h	Adds IndexMapper::size() and IndexMapper::numFreeSlots() accessor methods for active sequence and free-slot counts.
NIXL CMake Update cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt	Narrows NIXL wrapper CUDA dependencies from cuda_driver + cudart to cudart only.
Configuration and Dependencies .pre-commit-config.yaml, 3rdparty/fetch_content.json, cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu	Excludes broader trtllm-gen generated artifacts in pre-commit config, updates deepgemm dependency pin, and adds DSV3 router GEMM template instantiations for kHiddenDim=4096.

Sequence Diagram(s)

sequenceDiagram
  participant App as Application
  participant Compressor as Compressor Kernel
  participant KVCache as Paged KV Cache
  participant Transfer as NIXL Transfer
  
  Note over App,Transfer: KV Cache Compression + Transfer Workflow
  
  App->>Compressor: Decode: pagedKvCompressLaunch(kv_score, ...)
  activate Compressor
  Compressor->>KVCache: Load prior compressor state
  Compressor->>Compressor: Online softmax (running max + sum)
  Compressor->>KVCache: Write updated state + compressed tokens
  deactivate Compressor
  
  App->>Compressor: Prefill: prefillReductionLaunch(kv_score, ...)
  activate Compressor
  Compressor->>KVCache: Load full sequence
  Compressor->>Compressor: Reduce COMPRESS_RATIO rows via softmax
  Compressor->>KVCache: Write compressed output
  deactivate Compressor
  
  App->>Compressor: PostProcess: postProcessScatterLaunch(kv_comp, ...)
  activate Compressor
  Compressor->>Compressor: RMSNorm + RoPE + Hadamard
  Compressor->>KVCache: Scatter with FP8/MXFP4 quantization
  deactivate Compressor
  
  App->>Transfer: submitTransferRequests(src, dst)
  activate Transfer
  Transfer->>Transfer: detectVramRegionMap(src)
  Transfer->>Transfer: splitTransferDescsWithRegionMaps
  Transfer->>Transfer: Create per-chunk transfer ops
  deactivate Transfer

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#13740: Updates DSV3 router GEMM template instantiation variants for bf16 with different hidden dimensions, directly related to the dsv3RouterGemm.cu changes in this PR.

Suggested reviewers

• 2ez4bz
• byshiue
• Wanli-Jiang
• omera-nv
• moraxu

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch fix/indexer-radix-aux-cuda-graph

[longcheng-nv]

Pushed commit ffbc5addce addressing @heyuhhh's second inline review comment on IndexerTopKOp.cpp:122.

The reviewer's point — the two th::empty({0}, ...) placeholder Tensors at lines 122-123 became vestigial after the prior commit 383d8c5 tightened the blocks_per_row > 1 contract. They were only reachable in the blocks_per_row == 1 case, where the kernel never dereferences them. Reviewer suggested using the caller-passed buffers directly to avoid the leftover per-call allocation.

The fix — replace both Tensor placeholders with raw nullable int32_t* / float* pointers, matching the existing heuristicScratchPtr convention at lines 98-111. Hoist all caller-buffer validation (dtype / device / contiguity) out of the bp > 1 branch; keep the numel check inside it since that bound only matters when the kernel actually uses the buffers. For bp == 1 without caller-owned scratch, pass nullptr through; for bp > 1 without scratch, fail early with the same TORCH_CHECK message as before.

Same-PR-introduced regressions also fixed — chased down two paths that the strict TORCH_CHECK in 383d8c5 would now trip but weren't caught locally (the previous run was filtered with -k radix_aux, and CI broke earlier at model_config.py so the indexer tests never executed):

1. warmup_heuristic_topk_decode (dsa.py:68) — default num_cols=4096 falls below kSeqSmall=12288 so canUseHeuristic is false and the call routes to Radix with blocks_per_row=2 (num_rows=1 sweeps bp ∈ [2, maxByCols=2]). Allocate worst-case kMaxBlocksPerRowDecode=10 caller-owned scratch in the warmup. Cost ~80 KB, one-shot.
2. _run_indexer_topk_decode_check (test_indexer_topk.py:229) — covers params like num_tokens=32768, compress_ratio=1 → numCols=32768 → bp=9–10. Same fix: pre-allocate via existing _build_radix_aux_buffers and pass via kwargs. Same change applied to test_indexer_topk_decode_dist (fp32 + num_tokens=8192 hits Radix bp>1) and _run_indexer_topk_decode_v4_gvr_check (defense-in-depth — GVR almost always succeeds for that param matrix, but the kwarg is now uniform).

The radix-aux equivalence test no longer made sense in its original form (compared caller-owned vs the removed th::empty fallback). Rewrote it to verify allocation-pointer-independence — two invocations with freshly-allocated aux at distinct backing addresses must yield the same selected top-K logit values.

pre-commit run is clean on all three changed files (clang-format, ruff-format, isort, codespell, yapf).

cc @heyuhhh @lfr-0531 @mingyangHao

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49446 [ run ] triggered by Bot. Commit: ffbc5ad Link to invocation

[tensorrt-cicd]

PR_Github #49397 [ run ] completed with state ABORTED. Commit: 383d8c5

Link to invocation

[tensorrt-cicd]

PR_Github #49446 [ run ] completed with state SUCCESS. Commit: ffbc5ad
/LLM/main/L0_MergeRequest_PR pipeline #39091 completed with status: 'SUCCESS'

CI Report

Link to invocation

[longcheng-nv]

@lfr-0531 Friendly ping — this PR is ready and would appreciate a merge when you have a moment.

Current status (head ffbc5ad):

• ✅ CI green — latest PR_Github #49446 on ffbc5ad → SUCCESS (L0_MergeRequest_PR pipeline #39091 SUCCESS)
• ✅ Approved by @heyuhhh on 2026-05-20: "LGTM. Leave little comments for refine the code. Thanks!"
• ✅ Both inline review comments resolved:
  • IndexerTopKOp.cpp:L132 — replaced the if (radix_aux_indices.has_value() && ...) runtime branch with a TORCH_CHECK(...) precondition and dropped the dead else fallback (commit 383d8c5).
  • IndexerTopKOp.cpp:L122 — removed the two vestigial th::empty({0}, ...) placeholder tensors; the caller-passed aux tensors are reused directly (commit ffbc5ad).
• ✅ Mergeable: clean, base = feat/deepseek_v4.

Scope recap: makes the fp32 Radix path of indexer_topk_decode CUDA-Graph-safe by moving aux_indices / aux_logits scratch to caller-owned stable allocations — mirroring the same contract heuristic_scratch_values already follows on the Heuristic path. Fixes the CUDA 700 at smxx_layout.hpp:97 seen at ISL ≈ 100K / BS ≥ 8 with GVR-OFF (Radix) on Flash MXFP4.

Thanks!