cpp/tensorrt_llm/common/attentionOp.h (1)

416-416: Prefer fixed-width int over size_t for serialized config.

size_t varies by platform and can complicate stable serialization. Consider int32_t (or int64_t) for mChunkPrefillBufferBatchSize; if keeping size_t, always cast in serialization (as above).

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)

28-36: Duplicate imports of is_mla.

Two is_mla definitions are imported; one shadows the other and invites confusion. Keep one (or alias explicitly).

-from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
-                    create_py_executor_instance, instantiate_sampler, is_mla)
+from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
+                    create_py_executor_instance, instantiate_sampler)
-from .config_utils import is_mla
+from .config_utils import is_mla

tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)

1713-1713: Update stale comment contradicting new support.

Comment says FP8/NVFP4 chunked prefill isn’t supported, but you just enabled it. Remove or restate with SM gating.

Apply:

-    # currently, chunked prefill is not supported for fp8 and nvfp4
+    # Chunked prefill supports fp8/nvfp4 on SM90/SM100/SM120.

cpp/kernels/fmha_v2/src/fmha/fragment.h (2)

1741-1752: Avoid redeclaring traits/geometry already inherited from Base.

This specialization publicly derives from Tile_o_normalizer_fp32 but then redeclares Fragment_accu, Mma_tile, MMAS_M, MMAS_N, ROWS_PER_THREAD, REGS_PER_THREAD, WARPS_{M,N,K}, BYTES_PER_ELEMENT. It’s redundant and risks drift if Base changes.

Prefer bringing what you need into scope or referencing Base::… directly.

Example (minimal):

 struct Tile_o_normalizer<Ada_qmma_e4m3_fp32_traits, Cta_tile>
     : public Tile_o_normalizer_fp32<Ada_qmma_e4m3_fp32_traits, Cta_tile>
 {
-    using Fragment_accu = Fragment_accumulator<Traits>;
-    using Mma_tile = typename Traits::template Mma_tile<Cta_tile>;
-    enum { MMAS_M = Mma_tile::MMAS_M };
-    enum { MMAS_N = Mma_tile::VALID_MMAS_N };
-    enum { ROWS_PER_THREAD = 2 * MMAS_M };
-    enum { REGS_PER_THREAD = 8 };
-    enum { WARPS_M = Cta_tile::WARPS_M };
-    enum { WARPS_N = Cta_tile::WARPS_N };
-    enum { WARPS_K = Cta_tile::WARPS_K };
-    enum { BYTES_PER_ELEMENT = sizeof(float) };
+    using Base = Tile_o_normalizer_fp32<Ada_qmma_e4m3_fp32_traits, Cta_tile>;
+    using Fragment_accu = typename Base::Fragment_accu;
+    using Mma_tile     = typename Base::Mma_tile;
+    using Base::MMAS_M; using Base::MMAS_N; using Base::ROWS_PER_THREAD;
+    using Base::REGS_PER_THREAD; using Base::WARPS_M; using Base::WARPS_N;
+    using Base::WARPS_K; using Base::BYTES_PER_ELEMENT;

---

1822-1824: Typo: “diviser” → “divisor”.

cpp/kernels/fmha_v2/fmha_test.py (2)

164-169: Align epsilon for FP8 MLA context tests to avoid flakiness.

E4M3 typically needs a looser tolerance. Mirror the top-level tests by setting epsilon for FP8 here as well.

     epsilon = ''
-    if dtype == "-bf16" and s == 4096:
+    if dtype == "-bf16" and s == 4096:
         epsilon += ' -epsilon 0.03'
+    elif dtype in ["-e4m3", "-e4m3 -bf16-output"]:
+        epsilon += ' -epsilon 0.2'

---

179-191: Consider covering “-e4m3 -bf16-output” in save-softmax runs if supported.

The PR claims FP8 MLA chunked prefill support; if save-softmax is valid with BF16 outputs, include it to prevent regressions. If not supported, add an explicit skip.

-    if dtype in ["-bf16", "-e4m3"]:
+    if dtype in ["-bf16", "-e4m3", "-e4m3 -bf16-output"]:

If this combination is intentionally unsupported, please add a skip with a clear reason.

cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp (2)

489-497: Unify Hopper context-MLA detection; avoid duplicate, inconsistent predicates.

Here isHopperContextMLA also checks headSize == headSizeV + 64 and dtype in {BF16,E4M3}. Earlier (Line 456) a different predicate (sm90 && headSizeV == 128) is used. Merge into one helper so all decisions (tiling, WS, and softmax-stats) agree.

Minimal inline refactor within setupLaunchParams:

-    // Now we have SM90 context and FP8 generation MLA kernels
-    bool isHopperContextMLA = isSm90 && mFixedParams.headSizeV == 128;
+    auto const isContextMLAShape = (mFixedParams.headSize == mFixedParams.headSizeV + 64);
+    auto const isHopperContextMLA = isSm90 && isContextMLAShape
+        && (mFixedParams.dataType == DATA_TYPE_BF16 || mFixedParams.dataType == DATA_TYPE_E4M3)
+        && mFixedParams.headSizeV == 128;

Then reuse the same isHopperContextMLA at Lines 489–497 instead of recomputing it.

---

469-471: Comment typo.

“hooper style” → “Hopper style”.

cpp/tensorrt_llm/thop/attentionOp.h (1)

57-58: Prevent positional-arg break; keep param safe and keyword-only in bindings

Inserting chunk_prefill_buffer_batch_size before q_lora_rank can break existing positional callers. Ensure all Python bindings make args from attention_input_type onward keyword-only, and clamp the value to >= 1 when mapping into mChunkPrefillBufferBatchSize.

Would you like me to add kw-only in both pybind/nanobind bindings and add a clamp where the member is set?

tensorrt_llm/_torch/attention_backend/interface.py (2)

28-29: Validate and document runtime chunk buffer sizing

Add a quick validation that chunk_buffer_batch_size >= 1 and clarify units in the class docstring to avoid oversized/undersized allocations at runtime.

---

634-635: Keep MLA param default aligned and enforce bounds

Default chunk_prefill_buffer_batch_size=1 matches C++ expectations; add a bounds check (>=1) in the code path that forwards this into the ThOP call to prevent negative/zero from configs.

cpp/tensorrt_llm/thop/attentionOp.cpp (1)

631-633: SM gating uses magic numbers; clarify intent and ensure SM100 exclusion is intentional

Current condition enables FP8 Context MLA on SM90 and SM120 only. If SM100 is intentionally excluded, add a short comment. Also avoid magic numbers by capturing SM once.

Apply a small readability tweak:

-        op->mFP8ContextMLA = (tensorrt_llm::common::getSMVersion() == 120 || tensorrt_llm::common::getSMVersion() == 90)
-            && op->mKVCacheQuantMode.hasFp8KvCache();
+        // Enable FP8 Context MLA on Hopper (SM90) and Blackwell-next (SM120) only; SM100 intentionally excluded.
+        auto const sm = tensorrt_llm::common::getSMVersion();
+        op->mFP8ContextMLA = ((sm == 120) || (sm == 90)) && op->mKVCacheQuantMode.hasFp8KvCache();

tensorrt_llm/_torch/modules/attention.py (4)

1159-1165: Fix minor comment typo

“toal_token_q” -> “total_token_q”.

-        # [toal_token_q, num_heads, 2] -> [toal_token_q, num_heads] float2
+        # [total_token_q, num_heads, 2] -> [total_token_q, num_heads] float2

---

1186-1195: Assert device/dtype for chunked_ld_global_offset to match CUDA kernel expectations

Add lightweight asserts to prevent host/device mismatch or dtype errors at runtime.

Apply:

-            chunked_ld_global_offset = attn_metadata.chunked_ld_global_offset[loop_idx]
-            chunked_max_seq_len = attn_metadata.max_chunk_len_per_loop[loop_idx]
+            chunked_ld_global_offset = attn_metadata.chunked_ld_global_offset[loop_idx]
+            chunked_max_seq_len = attn_metadata.max_chunk_len_per_loop[loop_idx]
+            # Sanity checks for kernel inputs
+            assert temp_cu_chunked_seq_len.is_cuda, "cu_chunked_seq_len must be CUDA tensor"
+            assert chunked_ld_global_offset.is_cuda, "chunked_ld_global_offset must be CUDA tensor"
+            assert chunked_ld_global_offset.dtype == torch.int64, "chunked_ld_global_offset must be int64"

---

1243-1245: Graceful fallback for chunk_buffer_batch_size to keep backward compatibility

Use getattr to avoid attribute errors with older metadata.

Apply:

-                softmax_stats_tensor=self.temp_softmax_stats_tensor,
-                chunked_prefill_buffer_batch_size=attn_metadata.
-                runtime_features.chunk_buffer_batch_size,
+                softmax_stats_tensor=self.temp_softmax_stats_tensor,
+                chunked_prefill_buffer_batch_size=getattr(
+                    attn_metadata.runtime_features, "chunk_buffer_batch_size", None),

---

1295-1297: Repeat: fallback for final prefill call

Mirror the same defensive fallback here.

Apply:

-            softmax_stats_tensor=self.temp_softmax_stats_tensor,
-            chunked_prefill_buffer_batch_size=attn_metadata.runtime_features.
-            chunk_buffer_batch_size,
+            softmax_stats_tensor=self.temp_softmax_stats_tensor,
+            chunked_prefill_buffer_batch_size=getattr(
+                attn_metadata.runtime_features, "chunk_buffer_batch_size", None),

cpp/kernels/fmha_v2/setup.py (1)

3818-3821: Combination tuple order change is OK, but add a short comment for maintainers.

The new ordering (alibi, input_layout, enable_attn_logit_softcapping, return_softmax) differs from the hgmma path and can be easy to misread. A one-liner here would prevent future mistakes.

-    combinations = product([False, True], \
-        [InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV,
-         InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V],
-        [False, True], [False, True])
+    # combinations: (alibi, input_layout, enable_attn_logit_softcapping, return_softmax)
+    combinations = product([False, True],
+        [InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV,
+         InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V],
+        [False, True], [False, True])

cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu (3)

150-161: Remove duplicate early-exit check.

merge_op_val==0 is checked twice back-to-back; keep only one.

-    int64_t merge_op_val = merge_op[batch_idx];
-    if (merge_op_val == 0)
-    {
-        return; // skip this batch
-    }
+    int64_t merge_op_val = merge_op[batch_idx];
@@
-    if (merge_op_val == 0)
-    {
-        return; // skip this batch
-    }

---

274-286: K-PE write indexing is correct given head_dim gating; consider a comment.

head_dim_idx >= kLoraSize in this branch; subtracting kLoraSize is safe. A clarifying comment helps.

-            int const global_st_idx
-                = global_st_offset * KT::kRopeSize + local_token_idx * KT::kRopeSize + (head_dim_idx - KT::kLoraSize);
+            // head_dim_vec_idx >= kKVThreadPerHead => head_dim_idx >= kLoraSize; map to rope sub-dim.
+            int const global_st_idx = global_st_offset * KT::kRopeSize
+                + local_token_idx * KT::kRopeSize + (head_dim_idx - KT::kLoraSize);

---

337-349: Instantiation set expanded to bf16; good coverage.

Consider adding a unit test that mixes T=float with TCache=__nv_fp8_e4m3 to validate dequant path.

I can add a minimal test to cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu if helpful.

tensorrt_llm/_torch/attention_backend/trtllm.py (2)

191-192: Document and validate the new plan() parameter.

Add an Args entry for chunk_prefill_buffer_batch_size and (optionally) assert it’s >= 1 to prevent silent misconfigurations.

@@ def plan(
-        chunk_prefill_buffer_batch_size: int = 1,
+        chunk_prefill_buffer_batch_size: int = 1,
@@
-        """
+        """
@@
+            chunk_prefill_buffer_batch_size (int): Number of chunk buffers processed per loop (host-side staging/batching for chunked prefill). Default: 1.
@@
         self.is_spec_decoding_enabled = is_spec_decoding_enabled
@@
         self.spec_decoding_generation_lengths = spec_decoding_generation_lengths
         self.chunk_prefill_buffer_batch_size = chunk_prefill_buffer_batch_size
+        assert self.chunk_prefill_buffer_batch_size >= 1, "chunk_prefill_buffer_batch_size must be >= 1"
         self.kwargs.update(kwargs)

Also applies to: 276-277

---

896-959: Allocate per-context shapes to reduce memory.

chunked_seq_len is allocated with second dim = num_seqs, but only :num_contexts are used. Consider allocating [chunked_loop_num, num_contexts] for both device and host tensors to cut memory and H2D traffic.

cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu (2)

349-359: Debug helper overload is fine.

Named tensor printer improves diagnostics without impacting tests.

---

663-783: Metadata preparation mirrors runtime logic; consider replacing macros.

The algorithm matches the Python path. For readability, prefer inline lambdas over #define macros for 2D indexing.

-#define chunked_seq_len(loop,b) chunked_seq_len_vec[(loop) * (this->mBatchSize) + (b)]
-#define cu_chunked_seq_len(loop,b) h_cu_chunk_lens_ptr[(loop) * (this->mBatchSize + 1) + (b)]
-#define chunked_ld_global_offset(loop,b) h_chunked_ld_global_offset_ptr[(loop) * (this->mBatchSize) + (b)]
+auto chunked_seq_len = [&](int loop, int b){ return chunked_seq_len_vec[loop * this->mBatchSize + b]; };
+auto& cu_chunked_seq_len = [&](int loop, int b){ return h_cu_chunk_lens_ptr[loop * (this->mBatchSize + 1) + b]; };
+auto& chunked_ld_global_offset = [&](int loop, int b){ return h_chunked_ld_global_offset_ptr[loop * this->mBatchSize + b]; };
-#undef chunked_seq_len
-#undef cu_chunked_seq_len
-#undef chunked_ld_global_offset

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Learnt from: eopXD
PR: NVIDIA/TensorRT-LLM#6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.

Learnt from: eopXD
PR: NVIDIA/TensorRT-LLM#6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:2010-2045
Timestamp: 2025-08-21T09:41:49.347Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, updateSequenceCacheBlockOffsets is specifically for updating bookkeeping when blocks are added during the context phase, not for refreshing offsets after detach operations. During detach operations, GenerationRequest::removeFrontBlock handles the necessary cache block bookkeeping internally.

Learnt from: eopXD
PR: NVIDIA/TensorRT-LLM#6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.