User/itabrizian/fix disagg gen init history length

[Tabrizian]

Title:

Body:

Description

The V2 scheduler's _try_schedule_disagg_gen_init was passing req.context_remaining_length as the history_length argument to kv_cache_manager.resize_context. That field is the prompt suffix still to be processed, not the historic prefix already in the cache.

When GEN-side block reuse hits a long matching prefix, prepare_context (resource_manager.py:2739) sets req.context_current [None][fix] disagg-gen-init: pass context_current_position as history_length _position to the number of reused tokens and the cache's _history_length to the same value. The subsequent resize_context(..., history_length=req.context_remaining_length) then violates the "history length cannot be decreased" invariant in _KVCache.resize() and aborts the GEN event loop with:

ValueError: History length cannot be decreased

The block comment right above the call already states the correct intent ("Prompt is already historic — pass history_length so SWA layers skip pre-window block allocation"); only the field name was wrong. Passing req.context_current_position matches the cache's established history
boundary in both the reuse-hit and reuse-miss paths.

Reproduction

DSV4-Pro disaggregated smoke on Lyris GB300, 1 × CTX TEP8 + 1 × GEN TEP8, concurrency 2, enable_block_reuse: true, enable_chunked_prefill: true. Without the fix, after 2 trajectories succeed (no prefix hit), the third request triggers reuse on GEN, the scheduler attempts to
decrease history length, and all 8 GEN ranks raise the ValueError simultaneously. With the fix, GEN keeps scheduling disagg-gen-init requests correctly.

Test Coverage

This path is exercised by existing disaggregated-serving tests when both enable_block_reuse: true and enable_chunked_prefill: true are set and a request reuses a prefix on the GEN side. I haven't added a new unit test because the failure requires a multi-step disagg setup; happy
to add one if requested.

PR Checklist

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

Summary by CodeRabbit

• New Features

  • Added KV cache compression support for sparse attention patterns to improve memory efficiency
  • Introduced optimized kernel implementations for attention operations and MoE routing

• Improvements

  • Enhanced memory management and workspace allocation for KV cache operations
  • Improved kernel dispatching and precision handling across various workloads

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces DeepSeek-V4 sparse-attention KV cache compression, MHC hyper-connection fused compute kernels, compressed KV-space indexing support, NIXL agent lifecycle improvements, and supporting infrastructure including FP8 quantization, custom MoE routing, and architectural enhancements.

Changes

KV Cache Compression & Workspace Sizing

Layer / File(s)	Summary
VMM Region Metadata & Descriptor Splitting cpp/include/tensorrt_llm/executor/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp	Introduces backend-agnostic VMM chunk-boundary detection and VmmDescSplitter utilities with static methods for splitting VRAM descriptors at chunk boundaries, detecting multi-chunk allocations via CUDA address ranges, and maintaining per-region chunk-size metadata maps independent of NIXL.
KV Compression Decode/Prefill/Scatter Kernels cpp/tensorrt_llm/kernels/compressorKernels/*	Implements three-stage DeepSeek-V4 KV cache compression: decode-stage pagedKvCompressKernel with online-softmax state persistence, prefill-stage prefillReductionKernel for bulk compression, and fused postProcessScatterKernel for RMSNorm + RoPE + Hadamard + paged scatter with multiple cache quantization formats (FP8 per-tensor, blockwise, MXFP4).
Attention Workspace Sizing for KV Reuse cpp/tensorrt_llm/common/attentionOp.h, cpp/tensorrt_llm/common/attentionOp.cpp	Adds total_kv_len parameter to AttentionOp::getWorkspaceSizeForContext to account for KV cache reuse scenarios, updating FP8 separate K/V buffer sizing via max(total_kv_len, mChunkPrefillBufferBatchSize * max_num_tokens).

NIXL Agent Lifecycle & Serialization

Layer / File(s)	Summary
AgentDesc Serialization & VMM Region Metadata cpp/include/tensorrt_llm/executor/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp, cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp	Extends AgentDesc with VMM region metadata vectors, implements serialize/deserialize for backend descriptors and regions, and adds Python bindings for serialization methods.
NixlTransferStatus Weak Pointer Lifecycle cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp	Replaces raw agent pointers with std::weak_ptr in NixlTransferStatus to defer cleanup, adds exception-safe destructor with conditional XferReq release, implements getLastStatus() / getLastStatusStr() accessors, and updates wait/isCompleted polling with agent-liveness checks.
NixlTransferAgent Shutdown & Concurrency cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h, cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp	Transitions agent ownership from unique_ptr to shared_ptr, adds std::shared_mutex for coordinating shutdown vs transfer operations, implements VmmDescSplitter-based VRAM registration/deregistration with region tracking, and adds noexcept shutdown() methods for idempotent draining and loopback agent lifecycle management.

MHC Hyper-Connection Fused Kernels

Layer / File(s)	Summary
MHC Kernel Launchers & Support Operations cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h, cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu	Defines public launcher APIs (mhcBigFuseLaunch, mhcGemmSqrsumFmaLaunch, mhcHcHeadApplyLaunch, mhcPostMappingLaunch) with parameter validation, heuristic block-size selection, and support for optional fused RMSNorm and split-K accumulation via atomics.
Fused TF32 GEMM Pipeline cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh	Implements TF32 (tcgen05) pipelined GEMM kernels with pmap-based post/comb mixing and RMSNorm + sigmoid + sinkhorn + matmul fusion, supporting K-split accumulation via atomics and optional inline bigFuse with configurable fused RMSNorm behavior.
Fused-HC Dispatch & All-in-One Launchers cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt, cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu	Implements launch paths for TF32 and FMA fused-HC variants, including all-in-one single-kernel pipelines, workspace zeroing for K-split accumulation, cached TMA descriptor management, and heuristic KS/tile selection fallbacks.
Fused FMA Kernel Variants cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh	Defines FMA K-split and all-in-one kernels with pmap + GEMM + bigFuse fusion, cross-split coordination via done_counter, optional fused RMSNorm, and configurable tile/hidden sizing with warp-level reductions and atomic epilogue.

Compressed KV Index Space Support

Layer / File(s)	Summary
IndexerTopK Decode CompressRatio Support cpp/tensorrt_llm/kernels/IndexerTopK.h, cpp/tensorrt_llm/kernels/indexerTopK.cu	Adds compressRatio parameter to indexerTopK kernel and host dispatch, introduces computeIndexerTopKDecodeBlocksPerRow helper for multi-block sizing, gates heuristic eligibility on supported compression ratios (1,4), and refactors split+merge dispatch with region-aware block-per-row selection.
Heuristic TopK Decode CompressRatio Support cpp/tensorrt_llm/kernels/heuristicTopKDecode.h, cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu	Adds compressRatio parameter to heuristic TopK kernels and host launchers, updates index-space computation to derive N = actual_kv_len / compressRatio, and adjusts preIdxOffset behavior based on compression mode (DSv3.2 vs DSv4).

MoE Enhancements & DeepSeek V4 Support

Layer / File(s)	Summary
Custom MoE Gate Forward Routing cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h, cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu	Implements gate_forward kernel for custom MoE routing with softplus + sqrt score transformation, top-K expert selection (hash or top-k mode), weight normalization, and host dispatch supporting n_experts 256/384.
DeepSeekV4 Routing & Activation Clamping cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h, cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h, cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu, cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu	Adds DeepSeekV4 routing mode to RoutingMethodType enum, extends MoE activation with optional SwiGLU clamp limit support in separate-activation FP8 path, and wires scalar clamp parameters through activation kernel dispatch.
Sort Specializations & Router GEMM Instantiations cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh, cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu	Adds Sort template specializations for top-K values 5-8 using unrolled comparator-swap sequences, and adds router GEMM template instantiations for kHiddenDim=4096 with kNumTokens 1-16.

FP8 Quantization & DeepSeek V4 QNorm

Layer / File(s)	Summary
FP8 Blockscale Quantization+Packing cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h, cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu	Implements fused FP8 (E4M3) quantization with UE8M0 scale packing in blockscale GEMM layout, computing per-block amax, converting scales, quantizing bf16 to fp8, and packing four scale bytes via warp shuffles.
DeepSeekV4 QNorm Kernel cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h, cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu	Implements per-row QNorm for DeepSeek V4 using vectorized bf16/half2 pairs, warp-level cross-lane reduction, and rsqrt-derived scaling with per-region epsilon, supporting both bf16 and half precision.

MLA RoPE & Build Configuration

Layer / File(s)	Summary
MLA RoPE Programmatic Launch Completion cpp/tensorrt_llm/kernels/mlaKernels.cu	Adds SM90+ architecture-gated cudaGridDependencySynchronize before bounds checks and cudaTriggerProgrammaticLaunchCompletion at exit points, replaces direct kernel launches with launchWithPdlWhenEnabled wrapper.
Build Configuration & Dependency Updates .github/workflows/blossom-ci.yml, .pre-commit-config.yaml, 3rdparty/fetch_content.json, ATTRIBUTIONS-Python.md, cpp/tensorrt_llm/CMakeLists.txt, cpp/tensorrt_llm/kernels/CMakeLists.txt, cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h, cpp/include/tensorrt_llm/batch_manager/llmRequest.h	Adds mhcKernels and compressorKernels subdirectories to CMake build with exclusions from glob sources, updates dependency versions (deepgemm, flashinfer-python, nvidia-cutlass-dsl), adds CI authorization entries, updates copyright years, and adds IndexMapper accessors (size(), numFreeSlots()) and updateKvCachePerfMetrics() helper.

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14304: Updates /bot command authorization in .github/workflows/blossom-ci.yml with overlapping GitHub username allowlist changes.
• NVIDIA/TensorRT-LLM#13992: Shares identical ATTRIBUTIONS-Python.md dependency version bumps for flashinfer-python 0.6.10→0.6.11 and nvidia-cutlass-dsl 4.4.2→4.5.0.
• NVIDIA/TensorRT-LLM#13740: Modifies the same dsv3RouterGemm.cu template instantiation set; both PRs add hidden-dimension variants to the explicit kernel instantiations.

Suggested reviewers

• qiaoxj07
• kaiyux
• PerkzZheng
• bo-nv

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch user/itabrizian/fix-disagg-gen-init-history-length

[coderabbitai]

Actionable comments posted: 2

Note

Due to the large number of review comments, Critical severity comments were prioritized as inline comments.

Caution

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

⚠️ Outside diff range comments (6)

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h (1)

2-2: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the copyright year to 2026.

The file is being modified (new methods added) but the copyright year remains 2022-2025. As per coding guidelines, the year should be updated to reflect the latest modification.

📅 Proposed fix

- * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

As per coding guidelines: "All TensorRT-LLM source files must contain an NVIDIA copyright header with year of latest modification."

🤖 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 `@cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h` at line 2, Update the
copyright header in cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h to
include the latest modification year (change "2022-2025" to "2022-2026"); locate
the top-of-file NVIDIA copyright comment block and replace the year range so it
reflects 2026, ensuring the header matches the project's required format.

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp (2)

1-16: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the NVIDIA copyright year.

This file is modified in this PR, but the header still says 2025. It should reflect the latest modification year.

Suggested fix

- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

As per coding guidelines, `**/*.{cpp,cc,h,hpp,py,cu,cuh}`: "Include NVIDIA copyright header on ALL new files; update year on modified files."

🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp`
around lines 1 - 16, Update the SPDX copyright header year from 2025 to 2026 in
this modified file by changing the "SPDX-FileCopyrightText" line (the line
containing "Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES") and ensure any
other occurrences of the year in the file's header (including
SPDX-License-Identifier block) are similarly updated to 2026.

---

261-283: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

NixlTransferAgent re-binds blocking methods without GIL guards, shadowing the protected base-class versions.

When a derived class re-binds a method with the same Python name, the derived binding shadows the base binding for instances of the derived class, and the derived binding's call_guard controls GIL release. Here, register_memory, deregister_memory, load_remote_agent, load_remote_agent_by_connection, get_local_agent_desc, get_local_connection_info, invalidate_remote_agent, notify_sync_message, get_notified_sync_messages, and check_remote_descs are all re-bound without nb::call_guard<nb::gil_scoped_release>(), causing these methods to hold the GIL when called on NixlTransferAgent instances, even though the base-class versions release it. Only submit_transfer_requests includes the guard.

Apply nb::call_guard<nb::gil_scoped_release>() to all blocking re-bound methods, or drop the duplicate bindings to use the guarded base-class versions:

Suggested fix

-        .def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs"))
-        .def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs"))
-        .def("load_remote_agent",
-            nb::overload_cast<std::string const&, kvc::AgentDesc const&>(&kvc::NixlTransferAgent::loadRemoteAgent),
-            nb::arg("name"), nb::arg("agent_desc"))
-        .def("load_remote_agent_by_connection",
-            nb::overload_cast<std::string const&, kvc::ConnectionInfoType const&>(
-                &kvc::NixlTransferAgent::loadRemoteAgent),
-            nb::arg("name"), nb::arg("connection_info"))
-        .def("get_local_agent_desc", &kvc::NixlTransferAgent::getLocalAgentDesc)
-        .def("get_local_connection_info", &kvc::NixlTransferAgent::getLocalConnectionInfo)
-        .def("invalidate_remote_agent", &kvc::NixlTransferAgent::invalidateRemoteAgent, nb::arg("name"))
+        .def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs"),
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs"),
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("load_remote_agent",
+            nb::overload_cast<std::string const&, kvc::AgentDesc const&>(&kvc::NixlTransferAgent::loadRemoteAgent),
+            nb::arg("name"), nb::arg("agent_desc"), nb::call_guard<nb::gil_scoped_release>())
+        .def("load_remote_agent_by_connection",
+            nb::overload_cast<std::string const&, kvc::ConnectionInfoType const&>(
+                &kvc::NixlTransferAgent::loadRemoteAgent),
+            nb::arg("name"), nb::arg("connection_info"), nb::call_guard<nb::gil_scoped_release>())
+        .def("get_local_agent_desc", &kvc::NixlTransferAgent::getLocalAgentDesc,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("get_local_connection_info", &kvc::NixlTransferAgent::getLocalConnectionInfo,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("invalidate_remote_agent", &kvc::NixlTransferAgent::invalidateRemoteAgent, nb::arg("name"),
+            nb::call_guard<nb::gil_scoped_release>())
-        .def(
-            "notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message"))
-        .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages)
-        .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"));
+        .def("notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message"),
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("check_remote_descs", &kvc::NixlTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"),
+            nb::call_guard<nb::gil_scoped_release>());

🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp`
around lines 261 - 283, The derived Python bindings for NixlTransferAgent
re-bind several blocking methods without releasing the GIL; update the bindings
for register_memory, deregister_memory, load_remote_agent,
load_remote_agent_by_connection, get_local_agent_desc,
get_local_connection_info, invalidate_remote_agent, notify_sync_message,
get_notified_sync_messages, and check_remote_descs to include
nb::call_guard<nb::gil_scoped_release>() (or remove these duplicate bindings so
the base-class guarded versions are used) so that they behave consistently with
the base class; locate the bindings for kvc::NixlTransferAgent in
agentBindings.cpp and add nb::call_guard<nb::gil_scoped_release>() to each
.def(...) for the listed methods (or delete the duplicate .def lines to fall
back to the base-class implementations).

cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu (1)

203-245: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Validate compressRatio in the public launcher.

These launchers are exposed in the header, but the new parameter is never checked before the kernels do actual_kv_len / compressRatio. Passing 0 is an immediate device-side divide-by-zero, and other values silently produce a mismatched index space. Please reject anything outside {1, 4} here.

Suggested fix

     TLLM_CHECK_WITH_INFO(
         topK == 512 || topK == 1024 || topK == 2048, "heuristicTopKDecode requires topK ∈ {512, 1024, 2048}");
+    TLLM_CHECK_WITH_INFO(
+        compressRatio == 1 || compressRatio == 4,
+        "heuristicTopKDecode only supports compressRatio ∈ {1, 4}. Got %d", compressRatio);
 
     constexpr int kAlign = std::is_same_v<InputT, float> ? 4 : 8;

🤖 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 `@cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu` around lines 203 - 245, The
launchHeuristicTopKDecodeImpl function must validate the new compressRatio
parameter before any kernel launch; add a check at the top of
launchHeuristicTopKDecodeImpl (near the other TLLM_CHECK_WITH_INFO calls) that
ensures compressRatio is either 1 or 4 and otherwise fails fast (e.g.,
TLLM_CHECK_WITH_INFO(compressRatio == 1 || compressRatio == 4,
"heuristicTopKDecode requires compressRatio ∈ {1,4}")). This prevents
device-side divide-by-zero and mismatched index space when kernels compute
actual_kv_len / compressRatio.

cpp/tensorrt_llm/kernels/IndexerTopK.h (1)

1-2: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the NVIDIA copyright year in this modified header.

This header now changes public declarations, but the NVIDIA header still says 2025. Please bump it to 2026 to match repository policy for modified source files. As per coding guidelines **/*.{cpp,cc,h,hpp,py,cu,cuh}: Include NVIDIA copyright header on ALL new files; update year on modified files.

🤖 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 `@cpp/tensorrt_llm/kernels/IndexerTopK.h` around lines 1 - 2, Update the
top-of-file NVIDIA copyright year in this modified header: find the header
comment containing the range "2019-2025" in IndexerTopK.h (the top comment
block) and change it to "2019-2026"; ensure the updated year appears in the
public-facing header comment for this modified file so it complies with the repo
policy for modified source files.

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp (1)

844-874: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Don't hold the exclusive agent lock through the remote-MD polling loop.

This keeps mLock in unique mode while checkRemoteMD() spins and sleeps. If the peer never publishes MD, every submit/read path and even shutdown() are blocked behind this one call indefinitely.

Take the lock only around each NIXL call, and re-check mShutdown between retries.

🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp`
around lines 844 - 874, In loadRemoteAgent, avoid holding the exclusive mLock
for the entire remote-MD polling; release the unique_lock before entering the
loop and only lock around each NIXL call (e.g., mRawAgent->checkRemoteMD and any
fetch/send calls) so other paths and shutdown can proceed; inside the retry loop
(where status and nixl_xfer_dlist_t descs are used) re-acquire the lock briefly
to call checkRemoteMD, release immediately, and re-check mShutdown between
retries to abort if shutdown was triggered.

🟡 Minor comments (1)

cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h-1-15 (1)

1-15: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the NVIDIA header year.

Line 2 still ends at 2025, but this new header is being added in a 2026 change set.

As per coding guidelines, **/*.{cpp,cc,h,hpp,py,cu,cuh}: Include NVIDIA copyright header on ALL new files; update year on modified files.

🤖 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 `@cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h` around lines 1 - 15, Update
the NVIDIA copyright header year range in the file header comment from
"2022-2025" to "2022-2026"; locate the top-of-file block comment (the file-level
license header) in mhcKernels.h and change the year range on the second line so
the header reflects the 2026 modification.

🧹 Nitpick comments (7)

cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (1)

344-345: ⚡ Quick win

Use braces for if body.

Please wrap this if body with braces to match the project’s required control-flow style.

As per coding guidelines, "If and else should always be followed by brace-delimited statements, even if empty or a single statement".

🤖 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 `@cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu` around lines 344 - 345,
The single-line if "if (global_warp_id >= batch_size) return;" must be changed
to use brace-delimited statements: replace it with an if block using braces so
the early-return is enclosed (i.e., if (global_warp_id >= batch_size) { return;
}) to conform to the project's control-flow style; locate the statement
referencing global_warp_id and batch_size in the customMoeRoutingKernels.cu
kernel and apply this minimal change.

cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h (1)

30-69: ⚡ Quick win

Update the public launch comments to match the implementation.

The declarations still describe pagedKvCompressLaunch as a cdiv(state_dim, block_size) grid and prefillReductionLaunch as a 3D grid with num_head_chunks, but compressorKernels.cu now launches (batch_size, head_blocks) and (batch_size, max_outputs) respectively. These comments are actively misleading for anyone debugging occupancy or call contracts.

cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h (1)

1-13: ⚡ Quick win

Add the required file banner and Doxygen docs to this new public header.

This header is missing the NVIDIA Apache-2.0 copyright block, and the exported declaration is undocumented from a Doxygen perspective. Please add the standard banner and convert the API comments to //! form before merging. As per coding guidelines, "Include NVIDIA copyright header on ALL new files; update year on modified files" and "Follow Doxygen rules for documenting class interfaces and function prototypes; use //! for C++-style single-line Doxygen comments and //!< for class member documentation."

🤖 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 `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h` around lines 1 - 13, Add
the standard NVIDIA Apache-2.0 file header banner to this new public header and
add Doxygen-style single-line comments using //! for the API and namespace so
the symbol is documented; specifically document the kernels namespace and the
invokeDeepseekV4QNorm function (parameters: input, output, totalRows, headDim,
isBfloat16, eps, stream) with //! comments describing purpose, parameter
meanings and units where applicable, and use //!< if you add any class/struct
members later—ensure the copyright year is correct and the file now contains the
required banner and Doxygen comments before merging.

cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h (1)

35-48: ⚡ Quick win

Use Doxygen comments for the exported launcher.

This is a public header, but the new API is documented with plain // comments, so it will not follow the repo's header-doc convention. Please convert this block to //! comments on the declaration. As per coding guidelines, "Follow Doxygen rules for documenting class interfaces and function prototypes; use //! for C++-style single-line Doxygen comments and //!< for class member documentation."

🤖 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
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h`
around lines 35 - 48, Replace the plain // comments above the public launcher
declaration with Doxygen single-line comments using //! so the API is documented
per project conventions; update the comment block for
launch_fp8_quantize_1x128_packed_bf16_e4m3 to use //! for each line (and use
//!< for any trailing member descriptions if applicable) and preserve the
existing explanatory text and parameter descriptions exactly but formatted as
Doxygen comment lines.

cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu (1)

1-7: ⚡ Quick win

Add the required NVIDIA copyright header.

This new .cu file is missing the standard NVIDIA Apache-2.0 header block that the repo requires on new source files. As per coding guidelines, "Include NVIDIA copyright header on ALL new files; update year on modified files."

🤖 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 `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu` around lines 1 - 7, This
new source file deepseekV4QNormKernel.cu is missing the required NVIDIA
Apache-2.0 copyright/header block; add the standard NVIDIA copyright header
comment block (with correct year) at the top of the file before any `#include`
lines so the file complies with repository licensing guidelines.

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h (1)

77-80: ⚡ Quick win

Use the header Doxygen form required by this tree.

These new docs use /// and plain //, but this repo's header rule requires //! for interface docs and //!< for member docs.

As per coding guidelines, "**/*.{h,hpp,hxx}: Follow Doxygen rules for documenting class interfaces and function prototypes; use //! for C++-style single-line Doxygen comments and //!< for class member documentation`."

Also applies to: 89-90, 125-143, 152-153, 170-171

🤖 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 `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h`
around lines 77 - 80, Update the header comments in transferAgent.h to use the
repository's Doxygen header style: replace any occurrences of /// or // used for
interface or member docs with //! for interface/class/function single-line docs
and //!< for trailing member documentation; specifically update the comments
above or beside members mWeakAgent, mHandle, and mLastStatus and the other
commented sections referenced (around the blocks at 89-90, 125-143, 152-153,
170-171) so all class/interface and member docs conform to the //! / //!<
convention.

cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h (1)

28-158: ⚡ Quick win

Switch these public API comments to Doxygen //!.

These declarations live in a public header, but the surrounding docs use plain //, so they will be skipped by the generated interface docs.

As per coding guidelines, **/*.{h,hpp,hxx}: Follow Doxygen rules for documenting class interfaces and function prototypes; use //! for C++-style single-line Doxygen comments and //!< for class member documentation.

🤖 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 `@cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h` around lines 28 - 158,
Replace the plain // comments that document the public API with Doxygen
single-line comments (//!), e.g. convert the descriptive blocks and the inline
comments immediately above/decorating the function prototypes such as
mhcBigFuseLaunch, mhcGemmSqrsumFmaLaunch, mhcHcHeadApplyLaunch,
mhcPostMappingLaunch, mhcFusedHcLaunch, mhcFusedHcFmaLaunch,
mhcFusedHcAllInOneLaunch, and mhcFusedHcFmaAllInOneLaunch to use //! so they are
picked up by the generated docs; keep the same text and punctuation, and where
any comment documents a parameter or member inline, use the appropriate Doxygen
form (//! or //!<) to preserve meaning.

🤖 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
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu`:
- Around line 80-104: The code currently does a full double4 vector load via
in_ptr[0] whenever k_in_range (k_base < k), which can read past the row end for
non-multiple-of-16 k; change the logic in the fp8_blockscale_quant_packed kernel
so you only perform the packed vector load (using in_ptr, double4 and LoadTrick)
when the whole tile fits (k_base + kLoadNumElems <= k). If the tile straddles
the boundary (k_base < k && k_base + kLoadNumElems > k) perform safe per-element
loads into load_trick.v using the scalar input pointer (indexed by k_base + i)
and zero-pad the remainder, instead of reading in_ptr[0] then masking. Apply the
same fix for the corresponding store path (the section around lines 146-161):
only do vector stores when the full tile fits, otherwise write elements
one-by-one bounded by k - k_base to avoid writing past output.

In `@cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh`:
- Around line 1170-1185: The hand-rolled grid barrier using done_counter and the
spin-wait loop (see done_counter, m_block_idx, kNumSplits) must be removed and
replaced with a safe synchronization strategy: either use CUDA cooperative
groups grid-wide sync (e.g., cooperative_groups::this_grid().sync()) if the
kernel will be launched with cooperative launch and kNumSplits>1, or eliminate
the in-kernel barrier by splitting Phase 4 into a separate kernel launch that
reads done_counter/state and proceeds; update the code paths in the same
function where the spin-wait occurs to call the cooperative sync API when
cooperative launch is used, otherwise refactor into a new Phase4 kernel and
ensure done_counter increments are still performed before returning so the host
can launch Phase4.

---

Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h`:
- Line 2: Update the copyright header in
cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h to include the latest
modification year (change "2022-2025" to "2022-2026"); locate the top-of-file
NVIDIA copyright comment block and replace the year range so it reflects 2026,
ensuring the header matches the project's required format.

In `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp`:
- Around line 1-16: Update the SPDX copyright header year from 2025 to 2026 in
this modified file by changing the "SPDX-FileCopyrightText" line (the line
containing "Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES") and ensure any
other occurrences of the year in the file's header (including
SPDX-License-Identifier block) are similarly updated to 2026.
- Around line 261-283: The derived Python bindings for NixlTransferAgent re-bind
several blocking methods without releasing the GIL; update the bindings for
register_memory, deregister_memory, load_remote_agent,
load_remote_agent_by_connection, get_local_agent_desc,
get_local_connection_info, invalidate_remote_agent, notify_sync_message,
get_notified_sync_messages, and check_remote_descs to include
nb::call_guard<nb::gil_scoped_release>() (or remove these duplicate bindings so
the base-class guarded versions are used) so that they behave consistently with
the base class; locate the bindings for kvc::NixlTransferAgent in
agentBindings.cpp and add nb::call_guard<nb::gil_scoped_release>() to each
.def(...) for the listed methods (or delete the duplicate .def lines to fall
back to the base-class implementations).

In `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp`:
- Around line 844-874: In loadRemoteAgent, avoid holding the exclusive mLock for
the entire remote-MD polling; release the unique_lock before entering the loop
and only lock around each NIXL call (e.g., mRawAgent->checkRemoteMD and any
fetch/send calls) so other paths and shutdown can proceed; inside the retry loop
(where status and nixl_xfer_dlist_t descs are used) re-acquire the lock briefly
to call checkRemoteMD, release immediately, and re-check mShutdown between
retries to abort if shutdown was triggered.

In `@cpp/tensorrt_llm/kernels/heuristicTopKDecode.cu`:
- Around line 203-245: The launchHeuristicTopKDecodeImpl function must validate
the new compressRatio parameter before any kernel launch; add a check at the top
of launchHeuristicTopKDecodeImpl (near the other TLLM_CHECK_WITH_INFO calls)
that ensures compressRatio is either 1 or 4 and otherwise fails fast (e.g.,
TLLM_CHECK_WITH_INFO(compressRatio == 1 || compressRatio == 4,
"heuristicTopKDecode requires compressRatio ∈ {1,4}")). This prevents
device-side divide-by-zero and mismatched index space when kernels compute
actual_kv_len / compressRatio.

In `@cpp/tensorrt_llm/kernels/IndexerTopK.h`:
- Around line 1-2: Update the top-of-file NVIDIA copyright year in this modified
header: find the header comment containing the range "2019-2025" in
IndexerTopK.h (the top comment block) and change it to "2019-2026"; ensure the
updated year appears in the public-facing header comment for this modified file
so it complies with the repo policy for modified source files.

---

Major comments:
In `@cpp/include/tensorrt_llm/executor/transferAgent.h`:
- Line 21: Update the copyright header in the modified header file
transferAgent.h to reflect the current modification year by changing the ending
year (currently 2025) to the latest year; locate the top-of-file NVIDIA
copyright block in transferAgent.h and bump the year range or final year
accordingly so it matches the modification date before merging.

In `@cpp/tensorrt_llm/common/attentionOp.cpp`:
- Around line 832-837: total_kv_len can be negative and casting it directly to
size_t causes wraparound and huge allocations; before computing kv_buf_tokens,
validate and clamp total_kv_len (e.g., compute a safe_total_kv_len from
total_kv_len by checking if total_kv_len > 0 then
static_cast<size_t>(total_kv_len) else 0 or assert/fail on negative), then use
that safe_total_kv_len in the std::max call that computes kv_buf_tokens and
subsequently fp8_k_buf_size and fp8_v_buf_size (referencing total_kv_len,
kv_buf_tokens, fp8_k_buf_size, fp8_v_buf_size, mChunkPrefillBufferBatchSize,
max_num_tokens).

In `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp`:
- Around line 712-720: The current logic appends entries into
mRemoteVramRegionInfo[name] when agentDesc.getVramRegions() is non-empty,
leaving stale entries if the new regions list is smaller or different; instead,
replace the map for that name entirely: when handling agentDesc.getVramRegions()
for a given name, build a fresh map from those regions (or clear
mRemoteVramRegionInfo[name] before inserting) so mRemoteVramRegionInfo[name]
exactly reflects the current regions and avoids stale boundaries that will
confuse splitTransferDescsWithRegionMaps().
- Around line 809-812: The code returns a NixlTransferStatus without checking
the result of mRawAgent->postXferReq(handle, &reqParams); — capture the return
in status and verify it equals NIXL_IN_PROG (use the same TLLM_CHECK(status ==
NIXL_IN_PROG) pattern used elsewhere) before constructing/returning the
NixlTransferStatus; if the check fails, handle the error path (log/abort/return
a failed indicator) instead of returning a live transfer object. Ensure you
reference and use the existing symbols postXferReq, status, NIXL_IN_PROG,
TLLM_CHECK, mRawAgent and handle in the fix.

In `@cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp`:
- Around line 84-87: The code calls su::deserialize<size_t>(is) into numRegions
and immediately does regions.reserve(numRegions), which can trigger huge
allocations from a malformed blob; before reserving or allocating, validate
numRegions against the remaining bytes in the input stream and against a sane
maximum. After reading numRegions from the stream (the same stream variable is),
compute the remaining byte count (or otherwise determine remaining serialized
length), ensure numRegions is non-negative and that numRegions *
estimated_per_region_size (or sizeof(VramRegionMeta) if appropriate) does not
exceed the remaining bytes, and clamp or reject values above a defined safe
limit; only then call regions.reserve(numRegions) and proceed to deserialize
VramRegionMeta entries. Ensure failure paths set stream error/fail state or
return an error rather than performing the reserve.
- Around line 20-24: Update the NVIDIA copyright header at the top of
transferAgent.cpp to reflect the latest modification year (bump 2025 to 2026) so
the file's header matches the new changes; open transferAgent.cpp, locate the
file-level comment/header block (above the includes shown: serializeUtils.h,
cuda.h, dlfcn.h, sstream) and change the year range to the current year.
- Around line 155-184: The fast-path and loop in
VmmDescSplitter::splitTransferDescsWithRegionMaps incorrectly only checks
srcDescs for VRAM and seeds remaining from src length, skipping destination-only
VRAM splitting and allowing mismatched lengths; change the initial fast-path to
return early only when both srcDescs.getType() and dstDescs.getType() are not
MemoryType::kVRAM, and inside the loop compute/validate the transfer length
per-pair by asserting or using a safe minLen between srcVec[i].getLen() and
dstVec[i].getLen(); then ensure splitting logic (using lookupChunkInfo for
src/dst chunkSize and base) handles cases where srcChunkSize==0 but
dstChunkSize!=0 (and vice versa) so destination-only chunk boundaries are
respected and no output piece extends past the declared dst length.

In `@cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt`:
- Around line 24-25: The target-wide CUDA compile flag --use_fast_math on
target_compile_options(compressorKernels_src ...) is too broad and can alter
expf/div/rsqrtf/denorm behavior used by compressorKernels.cu; remove this option
from the target and either leave fast-math off by default or apply it only to
specific validated kernel source(s) using file-scoped properties (e.g.,
set_source_files_properties on the validated .cu files) or per-file compile
flags so compressorKernels.cu remains built without --use_fast_math.

In `@cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu`:
- Around line 1-15: Update the copyright year range in the file header: locate
the header line that reads "Copyright (c) 2022-2025, NVIDIA CORPORATION." and
change the end year from 2025 to 2026 so it reads "Copyright (c) 2022-2026,
NVIDIA CORPORATION." before merging.
- Around line 1125-1137: The current dispatch block in compressorKernels.cu
aliases any non-512 head_dim to the 128 kernel (DISPATCH_PREFILL_DTYPE), which
will silently accept unsupported widths (e.g., 256) and corrupt memory; change
the logic so you explicitly handle only the supported head_dim values (512 and
128) and for any other head_dim throw or assert immediately (e.g., validate
head_dim at the top of the function or replace the final else with an explicit
if (head_dim == 128) branch and a default branch that raises an error), keeping
the existing compress_ratio branches (compress_ratio == 4 vs else) and the
DISPATCH_PREFILL_DTYPE calls intact for the supported cases.
- Around line 1658-1753: The dispatch macros map any non-128 head_dim to 512 and
any non-4 elem_bytes to 2, causing invalid kernel launches; add explicit runtime
validation before dispatch: use TLLM_CHECK_WITH_INFO to ensure head_dim is one
of the supported values (128 or 512) and elem_bytes is either 2 or 4, and
additionally enforce that when cst (CacheScaleType) indicates a quantized/BF16
path (used by DISPATCH_HD_BF16/postProcessScatterKernel) elem_bytes == 2; fail
fast with a clear message referencing head_dim, elem_bytes, and cst so
DISPATCH_HD_EB / DISPATCH_HD_BF16 only run for supported (head_dim, elem_bytes)
pairs.
- Around line 718-738: The code currently clamps next_n to 4 via next_n_dispatch
= std::min(next_n, 4), which can silently drop tokens when next_n > 4; instead
add an explicit runtime check before the TRY_LAUNCH/FOREACH_DECODE_CONFIG block
that detects next_n > 4 and calls TLLM_THROW with a clear message referencing
the offending value (e.g., use next_n in the message), so the code fails fast
rather than dispatching the NN=4 specialization; locate symbols next_n_dispatch,
TRY_LAUNCH, FOREACH_DECODE_CONFIG, pagedKvCompressKernel and replace the silent
clamp with this validation and throw.
- Around line 160-175: The kernel helper packE2M1x2() currently returns 0 on
pre-SM100 builds, so if cache_scale_type==kMXFP4Blockwise reaches the kernel on
older GPUs the compressed cache will be all zeros; update
postProcessScatterLaunch() to detect runtime GPU compute capability and
reject/disable MXFP4 mode (cache_scale_type=3) on devices with SM < 100 or
implement a software fallback path analogous to mxFp4BlockScaleMoe.cpp;
specifically, add a runtime SM check in postProcessScatterLaunch() and either
set cache_scale_type to a safe alternative or return an error, ensuring any code
paths that depend on packE2M1x2() won’t run on unsupported hardware.

In `@cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h`:
- Around line 1-16: Update the copyright header at the top of this new header
file by changing the year range to include 2026 (e.g., "2022-2026") so the
header reflects the file landing in 2026; edit the comment block above the
`#pragma` once in compressorKernels.h to update the final year value.

In `@cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu`:
- Around line 359-360: The current calculation uses sp = log1pf(expf(s)) which
overflows for large positive s; replace it with a numerically stable softplus:
compute sp as max(0, s) + log1p(exp(-abs(s))) (using s, sp, and then score =
sqrt(sp)) in the customMoeRoutingKernels.cu kernel where s, sp, and score are
set so that large positive or negative s won't produce infinities; update the
assignment to sp accordingly and leave the subsequent score = sqrtf(sp)
unchanged.
- Around line 418-419: Guard the normalization to avoid dividing by zero: in the
kernel where out_weights and out_indices are written (symbols: out_weights,
out_indices, weight_sum, my_topk_value, my_topk_index, route_scale), check if
weight_sum is zero (or <= a small epsilon) before doing (my_topk_value /
weight_sum) * route_scale; if weight_sum is zero, set out_weights[...] to 0.0f
(or another defined safe fallback) and still write out_indices[...] =
my_topk_index, otherwise perform the existing division; use an epsilon check
(e.g., weight_sum <= 1e-12) to avoid FP edge cases.

In `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu`:
- Around line 107-116: Replace the debug-only assert on headDim in
dispatchDeepseekV4QNorm with an unconditional runtime validation and safe
dispatch: check the incoming headDim argument (not just via assert) and either
(a) throw or return an error/log message if headDim is unsupported, or (b)
implement a runtime dispatch that calls the correct deepseekV4QNormKernel<T, N>
specialization for supported N values (e.g., 512) based on headDim; ensure you
reference dispatchDeepseekV4QNorm and deepseekV4QNormKernel and prevent
launching the 512 specialization when headDim != 512 to avoid OOB accesses.

In `@cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu`:
- Around line 245-292: The file's copyright header needs its year updated to the
latest modification year (replace the trailing "2023" with the current
year)—open the file that contains the explicit template instantiations for
invokeRouterGemm (e.g., the block with
tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16,
...>), locate the top-of-file NVIDIA copyright header, and update the end year
to the current modification year per the guideline.

In `@cpp/tensorrt_llm/kernels/indexerTopK.cu`:
- Around line 904-913: The code currently only checks compressRatioOk for the
heuristic path but allows unsupported compressRatio values to reach
topKPerRowDecode (which uses rowEnd = actual_kv_len / compressRatio), causing
divide-by-zero or incorrect indexing; add an explicit guard that hard-fails
(return error/throw/log and abort kernel dispatch) when compressRatio is not 1
or 4 before any kernel dispatch in both overloads that choose between the
heuristic and topKPerRowDecode paths (the code that computes compressRatioOk and
the code that calls topKPerRowDecode); specifically, replace the implicit
reliance on compressRatioOk with an early check on compressRatio (if
(compressRatio != 1 && compressRatio != 4) { return/report error; }) so neither
heuristic nor topKPerRowDecode is launched with an unsupported compressRatio.

In `@cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh`:
- Around line 400-402: The kernel signature currently defaults norm_weight to
nullptr which allows a null dereference when kFuseNorm is true; remove the
default nullptr from the parameter (make __nv_bfloat16 const* __restrict__
norm_weight required) for the kFuseNorm-enabled specialization or add an
immediate null-check that fails fast before the fused-norm pass; update the
function/constructor that declares norm_weight (the fused FMA kernel in
mhc_fused_fma.cuh referenced by kFuseNorm and the fused-norm path) so callers
must supply norm_weight when kFuseNorm is enabled and/or assert/processLogger an
explicit error if norm_weight is null before lines where norm_weight is
dereferenced.
- Around line 85-87: Ensure the per-split width used for vectorized bf16
loads/stores is validated: check that hidden_size divided by NUM_K_SPLITS
(hidden_per_split) is a multiple of BF16_VEC (and multiple of 8 for the
layer_input_out path) before launching/using the kernels (references:
hidden_size, NUM_K_SPLITS, hidden_per_split, h_start/h_end, BF16_VEC, KS,
layer_input_out). If the divisibility contract is not met either return a clear
launch-time error or implement a scalar-tail path that handles the remaining
elements safely (avoid walking past row end or dropping tail). Update kernel
launch checks and/or add conditional vector vs scalar loops where the vector
width assumption is used.

In `@cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu`:
- Around line 641-643: The code unconditionally sets ks = 1 when num_k_splits ==
0, preventing the public-API sentinel 0 from triggering internal heuristics;
change the ks initialization so that if num_k_splits == 0 you call the internal
heuristic (e.g., computeKSplitHeuristic or implement a small heuristic using M,
N, kernel/block sizes like FHC_BLOCK_M/FHC_BLOCK_N) and otherwise use
static_cast<uint32_t>(num_k_splits), or if you intentionally want ks==1 keep the
current behavior but update the public header/docs to state that 0 does not
trigger heuristics; update references to ks, num_k_splits, and the launcher that
relies on this behavior accordingly.

In `@cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu`:
- Around line 867-876: In mhcHcHeadApplyLaunch, validate inputs before
launching: if M <= 0 return immediately (no launch); ensure mult does not exceed
kernel's MAX_MULT (8) — return/error if mult > 8 to avoid corrupting s_pre used
by mhcHcHeadApplyKernel; and ensure hidden_size is a multiple of 8 (or
reject/handle the tail) because the kernel uses 8-wide uint4 vector paths with
no tail handling; reference mhcHcHeadApplyLaunch, mhcHcHeadApplyKernel,
MAX_MULT, s_pre, and the uint4 vectorized paths when adding these checks and
early exits/errors.
- Around line 807-837: mhcBigFuseLaunch must reject hidden sizes not divisible
by 8 before dispatching because mhcBigFuseKernel uses 8-element uint4 vector
paths; add an upfront validation in mhcBigFuseLaunch (the function shown) that
checks hidden_size % 8 == 0 and fails fast (e.g., return/error via
TLLM_CHECK_WITH_INFO or similar) with a clear message referencing hidden_size,
so we never call mhcBigFuseDispatch/mhcBigFuseKernel with an unaligned
hidden_size that would cause out-of-bounds loads/stores.
- Around line 878-885: The launcher mhcPostMappingLaunch must guard against
invalid launches and vector tail overruns: first, return early (no launch) if B
<= 0; second, ensure hidden_size is a multiple of 8 before calling
mhcPostMappingKernel (since the kernel processes 8 __nv_bfloat16 per iteration)
— either round hidden_size up to the next multiple of 8 and ensure buffers/out
are padded accordingly, or fail/LOG and return if hidden_size % 8 != 0; finally
keep the kernel invocation mhcPostMappingKernel<<<grid, block, 0, stream>>>(...)
but only after these checks so the kernel never runs with B <= 0 or a
non-8-aligned hidden_size.
- Around line 839-865: The switch falls back to launching
mhcGemmSqrsumFmaKernel<24> for any unsupported tileN, which can overrun buffers;
after computing tileN (from tile_n or selectFmaTileN) add a validation that
tileN is one of the supported values {1,2,3,4,6,8,12,24} and that N % tileN == 0
(use the existing TLLM_CHECK_WITH_INFO macro to fail with a clear message if
not), or explicitly map unsupported tileN to a supported divisor of N before the
switch; update mhcGemmSqrsumFmaLaunch to reject or remap unsupported tileN
instead of silently taking the default case so mhcGemmSqrsumFmaKernel<TN> is
only launched with safe TN values.

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu`:
- Around line 259-266: Update the file header year in DevKernel.cu to the latest
modification year: locate the copyright/header block at the top of
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu and change
the year range ending "2024" to the current year (e.g., "2026"); ensure the
NVIDIA copyright header format is preserved exactly as in other files and keep
no other edits to the file content (the rest of the code, including symbols like
hasSwigluLimit and swigluLimit, must remain unchanged).

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h`:
- Around line 264-273: Update the file header copyright year range in
DevKernel.h to the current modification year (change the end year from 2024 to
2026) so the NVIDIA copyright header complies with guidelines; locate the file
header at the top of
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h and update
the year range, and make the same year update for the related header block
referenced near the swigluLimit/hasSwigluLimit definitions.

In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh`:
- Around line 240-332: Update the file header year to reflect the latest
modification (change the ending year from 2025 to 2026) and ensure the NVIDIA
copyright header is present/complete; locate the header at the top of the same
file that contains templates like Sort<5, RedType>, Sort<6, RedType>, Sort<7,
RedType>, and Sort<8, RedType> and replace the old year range with the new 2026
year.

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu`:
- Around line 673-677: Update the file header in runner.cu to reflect the latest
modification year (change the trailing year range that currently ends in 2025 to
2026) so the NVIDIA copyright header matches the file changes; open the
top-of-file header block and update the year range while leaving the rest of the
NVIDIA copyright/header intact (the change affects runner.cu where symbols like
activationData and args.gemm1_clamp_limit_value are modified).
- Around line 676-677: When args.has_gemm1_clamp_limit_value is true, validate
that args.gemm1_clamp_limit_value is > 0 before assigning to
activationData.swigluLimit and activationData.hasSwigluLimit; if the value is
non-positive, either set activationData.hasSwigluLimit = false and skip the
assignment or fail fast (e.g., return an error/throw/log) so invalid clamp
behavior is avoided. Locate the assignment to activationData.swigluLimit and
activationData.hasSwigluLimit and add the conditional check on
args.has_gemm1_clamp_limit_value && args.gemm1_clamp_limit_value > 0, otherwise
handle the non-positive case as described.

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h`:
- Around line 88-91: RoutingMethodType::DeepSeekV4 was added to the enum but
Runner::run's switch/routing dispatch lacks a case for it, causing DeepSeekV4 to
hit the unimplemented/default path; add a case for RoutingMethodType::DeepSeekV4
inside Runner::run and implement its routing logic (or delegate to an existing
handler) consistent with other methods' patterns (e.g., mirror the behavior of a
similar method case or call the same helper function used by other routing
branches), and ensure any inputs/params used by DeepSeekV4 are validated and any
error path logs/returns match existing conventions.
- Around line 287-295: Update the file header's copyright year range in runner.h
to include the latest modification year (replace the trailing "2025" with the
current year); locate the file-level header at the top of
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h and modify the
copyright range string accordingly so the header matches the coding guideline
for modified files (no code changes to symbols like gemm1_clamp_limit,
gemm1_clamp_limit_value, or has_gemm1_clamp_limit_value are needed).

---

Minor comments:
In `@cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h`:
- Around line 1-15: Update the NVIDIA copyright header year range in the file
header comment from "2022-2025" to "2022-2026"; locate the top-of-file block
comment (the file-level license header) in mhcKernels.h and change the year
range on the second line so the header reflects the 2026 modification.

---

Nitpick comments:
In `@cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h`:
- Around line 77-80: Update the header comments in transferAgent.h to use the
repository's Doxygen header style: replace any occurrences of /// or // used for
interface or member docs with //! for interface/class/function single-line docs
and //!< for trailing member documentation; specifically update the comments
above or beside members mWeakAgent, mHandle, and mLastStatus and the other
commented sections referenced (around the blocks at 89-90, 125-143, 152-153,
170-171) so all class/interface and member docs conform to the //! / //!<
convention.

In `@cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu`:
- Around line 344-345: The single-line if "if (global_warp_id >= batch_size)
return;" must be changed to use brace-delimited statements: replace it with an
if block using braces so the early-return is enclosed (i.e., if (global_warp_id
>= batch_size) { return; }) to conform to the project's control-flow style;
locate the statement referencing global_warp_id and batch_size in the
customMoeRoutingKernels.cu kernel and apply this minimal change.

In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h`:
- Around line 35-48: Replace the plain // comments above the public launcher
declaration with Doxygen single-line comments using //! so the API is documented
per project conventions; update the comment block for
launch_fp8_quantize_1x128_packed_bf16_e4m3 to use //! for each line (and use
//!< for any trailing member descriptions if applicable) and preserve the
existing explanatory text and parameter descriptions exactly but formatted as
Doxygen comment lines.

In `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu`:
- Around line 1-7: This new source file deepseekV4QNormKernel.cu is missing the
required NVIDIA Apache-2.0 copyright/header block; add the standard NVIDIA
copyright header comment block (with correct year) at the top of the file before
any `#include` lines so the file complies with repository licensing guidelines.

In `@cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h`:
- Around line 1-13: Add the standard NVIDIA Apache-2.0 file header banner to
this new public header and add Doxygen-style single-line comments using //! for
the API and namespace so the symbol is documented; specifically document the
kernels namespace and the invokeDeepseekV4QNorm function (parameters: input,
output, totalRows, headDim, isBfloat16, eps, stream) with //! comments
describing purpose, parameter meanings and units where applicable, and use //!<
if you add any class/struct members later—ensure the copyright year is correct
and the file now contains the required banner and Doxygen comments before
merging.

In `@cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h`:
- Around line 28-158: Replace the plain // comments that document the public API
with Doxygen single-line comments (//!), e.g. convert the descriptive blocks and
the inline comments immediately above/decorating the function prototypes such as
mhcBigFuseLaunch, mhcGemmSqrsumFmaLaunch, mhcHcHeadApplyLaunch,
mhcPostMappingLaunch, mhcFusedHcLaunch, mhcFusedHcFmaLaunch,
mhcFusedHcAllInOneLaunch, and mhcFusedHcFmaAllInOneLaunch to use //! so they are
picked up by the generated docs; keep the same text and punctuation, and where
any comment documents a parameter or member inline, use the appropriate Doxygen
form (//! or //!<) to preserve meaning.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

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

[coderabbitai]

⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Guard partial tiles before the vector load/store.

Lines 91 and 161 still perform a full 16-element vector access whenever k_base < k, even if the tile straddles the row boundary. That means non-multiple-of-16 k can read past the end of the last row and write past the output allocation on the final row. Please split the tail path into bounded per-element loads/stores instead of masking after the vector access.

Suggested fix

-    load_trick.pack = k_in_range ? in_ptr[0] : double4{};
-
-    if (k_in_range && k_base + kLoadNumElems > k)
-    {
-        int const valid = k - k_base;
-#pragma unroll
-        for (int i = 0; i < kLoadNumElems; ++i)
-        {
-            if (i >= valid)
-            {
-                load_trick.v[i] = __nv_bfloat16(0.0f);
-            }
-        }
-    }
+    bool const fullTile = k_base + kLoadNumElems <= k;
+    if (fullTile)
+    {
+        load_trick.pack = in_ptr[0];
+    }
+    else
+    {
+#pragma unroll
+        for (int i = 0; i < kLoadNumElems; ++i)
+        {
+            int const col = k_base + i;
+            load_trick.v[i] = col < k ? input[static_cast<int64_t>(m_idx) * k + col] : __nv_bfloat16(0.0f);
+        }
+    }
@@
-        if (k_base + kStoreNumElems > k)
-        {
-            int const valid = k - k_base;
-#pragma unroll
-            for (int i = 0; i < kStoreNumElems; ++i)
-            {
-                if (i >= valid)
-                {
-                    store_trick.v[i] = __nv_fp8_e4m3(0.0f);
-                }
-            }
-        }
-        out_ptr[0] = store_trick.pack;
+        if (fullTile)
+        {
+            out_ptr[0] = store_trick.pack;
+        }
+        else
+        {
+#pragma unroll
+            for (int i = 0; i < kStoreNumElems; ++i)
+            {
+                int const col = k_base + i;
+                if (col < k)
+                {
+                    fp8_output[static_cast<int64_t>(m_idx) * k + col] = store_trick.v[i];
+                }
+            }
+        }
     }

Also applies to: 146-161

🤖 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
`@cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu`
around lines 80 - 104, The code currently does a full double4 vector load via
in_ptr[0] whenever k_in_range (k_base < k), which can read past the row end for
non-multiple-of-16 k; change the logic in the fp8_blockscale_quant_packed kernel
so you only perform the packed vector load (using in_ptr, double4 and LoadTrick)
when the whole tile fits (k_base + kLoadNumElems <= k). If the tile straddles
the boundary (k_base < k && k_base + kLoadNumElems > k) perform safe per-element
loads into load_trick.v using the scalar input pointer (indexed by k_base + i)
and zero-pad the remainder, instead of reading in_ptr[0] then masking. Apply the
same fix for the corresponding store path (the section around lines 146-161):
only do vector stores when the full tile fits, otherwise write elements
one-by-one bounded by k - k_base to avoid writing past output.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy lift

Replace this hand-rolled grid barrier before enabling kNumSplits > 1.

This done_counter spin-wait is a non-cooperative grid-wide barrier. With a normal launch, CTAs can park here while holding the kernel's full SMEM footprint, which means mixed m_block_idx residency can deadlock the launch once some splits are waiting for sibling CTAs that the scheduler can no longer issue. This needs either a real cooperative-grid sync or a separate Phase 4 kernel.

🤖 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 `@cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh` around lines
1170 - 1185, The hand-rolled grid barrier using done_counter and the spin-wait
loop (see done_counter, m_block_idx, kNumSplits) must be removed and replaced
with a safe synchronization strategy: either use CUDA cooperative groups
grid-wide sync (e.g., cooperative_groups::this_grid().sync()) if the kernel will
be launched with cooperative launch and kNumSplits>1, or eliminate the in-kernel
barrier by splitting Phase 4 into a separate kernel launch that reads
done_counter/state and proceeds; update the code paths in the same function
where the spin-wait occurs to call the cooperative sync API when cooperative
launch is used, otherwise refactor into a new Phase4 kernel and ensure
done_counter increments are still performed before returning so the host can
launch Phase4.