From 2eaa4eb99c88b0c18222a153cbd768ad40cd0029 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Mon, 24 Aug 2026 10:06:40 -0700
Subject: [PATCH 01/15] [None][feat] Kimi K3 helix x DSpark: verify groups on
 the V2 ledger

Speculative verify groups (1 target + k draft tokens) under helix
MLA-DCP + KDA-TP, overlap-scheduler-safe by construction:

- Per-token primitive: global positions -> owner rank, rank-local KV
  write slot (-1 = not owned) and per-token attention bound
  local_len(pos+1). A group may straddle a ledger-page boundary onto
  two CP ranks; every consumer below is per-token, so no group-affinity
  placement is needed and the stateless round-robin ledger stays intact.
- model_engine: extend-request packing emits provisional helix values
  (stale base / full-acceptance convention, mirroring the non-helix
  spec path); _preprocess_inputs applies the overlap accepted-count
  correction to helix_position_offsets and re-derives slots/bounds/
  rank-local kv_lens on device (recompute_helix_spec_buffers) -- CUDA
  graph safe, exact under overlap.
- trtllm attention metadata: helix_local_slots / helix_kv_bounds
  buffers, per-seq owned-new-token counts for the kv_lens math,
  vectorized helix_local_len_vec.
- mla_rope_generation (cpp): optional per-token slot table supersedes
  the per-sequence inactive-rank gate for KV appends (third
  helix_tensor_params entry; kernels index slots per token).
- CuTe DSL MLA decode (fp16/bf16): optional per-token kv_bounds
  replaces the implicit causal bound under helix; masked-phase span
  widened by one; stats epilogue emits the (-inf, 0) identity per
  token; gate now admits seq_len_q > 1 with helix when the spec
  buffers are live (fp8 KV stays rejected).
- Guards: helix spec allowlist = standalone DSpark linear chains only,
  loud rejection otherwise; the drafter's paged KV manager is built on
  the repurposed CP-free mapping (helix ledger applies only to the
  target KV).
- DFlash disagg gen-worker slot bootstrap (standalone twin of the
  embedded-DSpark #16767 fix): transferred requests get isolated
  context slots instead of aliasing the shared scratch row; acceptance
  is degraded until a ctx->gen window transfer exists, correctness is
  carried by verify.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
(cherry picked from commit fdb9594b21ff84bc45ed65bf94742646867d793d)
(cherry picked from commit dbd5dc613a1031bf48afc49b5b79a9393f697529)
---
 cpp/tensorrt_llm/kernels/mlaKernels.cu        |  32 +++-
 cpp/tensorrt_llm/kernels/mlaKernels.h         |   7 +
 cpp/tensorrt_llm/thop/dsv3RopeOp.cpp          |  20 ++-
 .../attention_backend/fmha/cute_dsl_mla.py    |  17 +-
 .../_torch/attention_backend/trtllm.py        | 170 +++++++++++++++++-
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  38 +++-
 .../attention/mla/mla_decode_fp16.py          |  56 +++++-
 .../blackwell/attention/mla/mla_decode_fp8.py |   4 +
 .../_torch/models/modeling_kimi_linear.py     |  23 ++-
 tensorrt_llm/_torch/pyexecutor/_util.py       |  11 +-
 .../_torch/pyexecutor/model_engine.py         |  85 ++++++++-
 11 files changed, 431 insertions(+), 32 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu
index e98768faea82..7fd6c12957e2 100644
--- a/cpp/tensorrt_llm/kernels/mlaKernels.cu
+++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu
@@ -447,7 +447,8 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe,
     int q_pe_stride, KvCacheDataType cache_type, float* bmm1_scale, float* bmm2_scale, float const* quant_scale_o,
     float const* quant_scale_q, float const* quant_scale_kv, float const* dequant_scale_q,
     float const* dequant_scale_kv, float host_bmm1_scale, int32_t const* helix_position_offsets,
-    bool const* helix_is_inactive_rank, bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false)
+    bool const* helix_is_inactive_rank, int32_t const* helix_local_slots = nullptr,
+    bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false)
 {
     // Constants.
     using VecT = typename VecType<T>::Type;
@@ -563,10 +564,18 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe,
             {
                 if (head_idx == head_num)
                 {
-                    // If helix parallelism is being used, only write to KV cache if current rank is active.
-                    if (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx])
+                    // If helix parallelism is being used, only write to KV cache if this rank
+                    // owns the token's position. With speculative verify groups the per-token
+                    // slot table decides (a group can straddle a page boundary onto two ranks);
+                    // otherwise the per-sequence flag does.
+                    bool const helix_write = helix_local_slots != nullptr
+                        ? helix_local_slots[global_token_idx] >= 0
+                        : (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]);
+                    if (helix_write)
                     {
-                        auto const token_kv_idx = kv_cache_lengths[batch_idx] - seq_len + local_token_idx;
+                        auto const token_kv_idx = helix_local_slots != nullptr
+                            ? helix_local_slots[global_token_idx]
+                            : kv_cache_lengths[batch_idx] - seq_len + local_token_idx;
 
                         {
                             auto kDst = reinterpret_cast<T*>(kv_cache.getKBlockPtr(batch_idx, token_kv_idx));
@@ -629,10 +638,17 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe,
                     }
                 }
 
-                // If helix parallelism is being used, only write to KV cache if current rank is active.
-                if (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx])
+                // If helix parallelism is being used, only write to KV cache if this rank
+                // owns the token's position (per-token slots for speculative verify
+                // groups, per-sequence flag otherwise; see the Q/K branch above).
+                bool const helix_write = helix_local_slots != nullptr
+                    ? helix_local_slots[global_token_idx] >= 0
+                    : (helix_is_inactive_rank == nullptr || !helix_is_inactive_rank[batch_idx]);
+                if (helix_write)
                 {
-                    auto const token_kv_idx = kv_cache_lengths[batch_idx] - seq_len + local_token_idx;
+                    auto const token_kv_idx = helix_local_slots != nullptr
+                        ? helix_local_slots[global_token_idx]
+                        : kv_cache_lengths[batch_idx] - seq_len + local_token_idx;
                     auto const src_kv_global_offset = static_cast<size_t>(global_token_idx) * (c_k + ROPE_DIM);
 
                     {
@@ -1690,7 +1706,7 @@ void invokeMLARopeGeneration(MlaParams<T>& params, KVCacheBuffer kv_cache_buffer
         params.q_pe_stride, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o,
         quant_scale_q_eff, params.quant_scale_kv, params.dequant_scale_q, params.dequant_scale_kv,
         params.host_bmm1_scale, params.helix_position_offsets, params.helix_is_inactive_rank,
-        params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler);
+        params.helix_local_slots, params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler);
 }
 
 template <typename T, typename KVCacheBuffer>
diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.h b/cpp/tensorrt_llm/kernels/mlaKernels.h
index 62bf1270290c..47978c1be50e 100644
--- a/cpp/tensorrt_llm/kernels/mlaKernels.h
+++ b/cpp/tensorrt_llm/kernels/mlaKernels.h
@@ -153,6 +153,13 @@ struct MlaParams
     // for Helix parallelism: whether the current rank is inactive, shape [b]
     // (the current query tokens are not appended to this rank's KV cache)
     bool const* helix_is_inactive_rank{nullptr};
+
+    // for Helix parallelism with speculative verify groups: per-token
+    // rank-local KV write slot, shape [num_tokens]; -1 means another CP rank
+    // owns the token's global position. Non-null supersedes the per-sequence
+    // helix_is_inactive_rank gate (a 1 + draft_len group can straddle a
+    // ledger-page boundary, splitting ownership between two ranks).
+    int32_t const* helix_local_slots{nullptr};
 };
 
 template <typename T, typename KVCacheBuffer>
diff --git a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp
index 191349aea3d2..f8e3d825e1e1 100644
--- a/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp
+++ b/cpp/tensorrt_llm/thop/dsv3RopeOp.cpp
@@ -74,6 +74,8 @@ struct MlaRopeGenArgs
     float host_bmm1_scale;
     int32_t const* helix_position_offsets_ptr;
     bool const* helix_is_inactive_rank_ptr;
+    // Per-token KV write slots for speculative verify groups (nullptr otherwise).
+    int32_t const* helix_local_slots_ptr;
     // `kv_norm_weight_ptr` set: `invokeMLAKvNormRopeQuantGeneration` produces the KV
     // half (norm + rope + fp8 + paged write) and the RoPE kernel runs Q-only.
     void const* kv_norm_weight_ptr;
@@ -130,6 +132,7 @@ void invokeMLARopeGenerationHelper(T const* latent_cache_ptr, T* q_pe_ptr, T* fu
     mla_params.host_bmm1_scale = args.host_bmm1_scale;
     mla_params.helix_position_offsets = args.helix_position_offsets_ptr;
     mla_params.helix_is_inactive_rank = args.helix_is_inactive_rank_ptr;
+    mla_params.helix_local_slots = args.helix_local_slots_ptr;
 
     mla_params.precomputed_cu_seqlens = args.precomputed_cu_seqlens;
     mla_params.precomputed_fmha_scheduler = args.precomputed_fmha_scheduler;
@@ -182,8 +185,9 @@ void MLARopeGeneration(std::optional<torch::Tensor> fused_q, // [tokens, num_hea
     TLLM_CHECK_WITH_INFO(
         head_size == kv_lora_rank + qk_rope_head_dim, "head_size must = kv_lora_rank + qk_rope_head_dim");
     TLLM_CHECK_WITH_INFO(num_kv_heads == 1, "num_kv_heads must = 1");
-    TORCH_CHECK(helix_tensor_params.size() == 2,
-        "Expecting 2 tensors for helix_tensor_params: helix_position_offsets and helix_is_inactive_rank.");
+    TORCH_CHECK(helix_tensor_params.size() == 2 || helix_tensor_params.size() == 3,
+        "Expecting 2 or 3 tensors for helix_tensor_params: helix_position_offsets, helix_is_inactive_rank "
+        "and optionally helix_local_slots (per-token KV write slots for speculative verify groups).");
 
     auto stream = at::cuda::getCurrentCUDAStream(latent_cache.get_device());
     auto const kv_cache_quant_mode = tc::QuantMode(uint32_t(quant_mode));
@@ -210,6 +214,14 @@ void MLARopeGeneration(std::optional<torch::Tensor> fused_q, // [tokens, num_hea
         = helix_position_offsets.has_value() ? helix_position_offsets->data_ptr<int32_t>() : nullptr;
     bool const* helix_is_inactive_rank_ptr
         = helix_is_inactive_rank.has_value() ? helix_is_inactive_rank->data_ptr<bool>() : nullptr;
+    int32_t const* helix_local_slots_ptr = nullptr;
+    if (helix_tensor_params.size() == 3 && helix_tensor_params[2].has_value())
+    {
+        helix_local_slots_ptr = helix_tensor_params[2]->data_ptr<int32_t>();
+        TORCH_CHECK(!kv_norm_weight.has_value(),
+            "helix_local_slots (speculative verify groups) is not supported on the fused "
+            "KV-norm RoPE path: its KV append kernel has no per-token helix gate.");
+    }
 
     int* cu_q_seqlens_ptr = reinterpret_cast<int*>(cu_q_seqlens.data_ptr());
     int* cu_kv_seqlens_ptr = reinterpret_cast<int*>(cu_kv_seqlens.data_ptr());
@@ -308,8 +320,8 @@ void MLARopeGeneration(std::optional<torch::Tensor> fused_q, // [tokens, num_hea
         block_ids_per_seq_ptr, cache_type, cu_q_seqlens_ptr, cu_kv_seqlens_ptr, fmha_tile_counter_ptr,
         mla_bmm1_scale_ptr, mla_bmm2_scale_ptr, quant_q_buffer_ptr, quant_scale_qkv_ptr, quant_scale_o_ptr,
         kv_scale_orig_quant_ptr, kv_scale_quant_orig_ptr, host_bmm1_scale, helix_position_offsets_ptr,
-        helix_is_inactive_rank_ptr, kv_norm_weight_ptr, static_cast<float>(kv_norm_eps), latent_row_stride,
-        precomputed_cu_seqlens, precomputed_fmha_scheduler, kv_only, kv_done_elsewhere};
+        helix_is_inactive_rank_ptr, helix_local_slots_ptr, kv_norm_weight_ptr, static_cast<float>(kv_norm_eps),
+        latent_row_stride, precomputed_cu_seqlens, precomputed_fmha_scheduler, kv_only, kv_done_elsewhere};
 
     void* q_pe_ptr = kv_only ? nullptr : q_pe->data_ptr();
     void* fused_q_ptr = kv_only ? nullptr : fused_q->data_ptr();
diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
index 65768b68d74f..a6ea42deb7d5 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
@@ -309,8 +309,16 @@ def _is_supported_with_reason(
         seq_len_q = q.shape[0] // meta.num_generations
         batch_size = meta.num_generations
         if meta.helix_position_offsets is not None:
-            if seq_len_q != 1:
+            if seq_len_q != 1 and not getattr(meta, "_helix_spec_tokens_valid",
+                                              False):
+                # Multi-token decode under helix needs the per-token bound /
+                # write-slot buffers of the speculative verify-group path.
                 return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix."
+            if seq_len_q != 1 and self._get_kernel_dtype(
+                    attn, q) == torch.float8_e4m3fn:
+                return False, (
+                    "CuTe DSL MLA FMHA helix verify groups require a bf16/fp16 "
+                    "KV cache (the fp8 kernel has no per-token bounds).")
             softmax_stats = fwd.softmax_stats_tensor
             if softmax_stats is None:
                 return False, "CuTe DSL MLA FMHA requires softmax_stats_tensor with Helix."
@@ -513,6 +521,13 @@ def _run_mla_decode(
             # Max batch size for the AutoTuner to profile.
             int(meta.max_num_requests),
             params.fwd.softmax_stats_tensor,
+            # Helix speculative verify groups: per-token rank-local bounds
+            # (filled by recompute_helix_spec_buffers). None on the
+            # single-token helix path and outside helix.
+            (meta.helix_kv_bounds[:num_tokens] if
+             (meta.helix_position_offsets is not None
+              and getattr(meta, "_helix_spec_tokens_valid", False)
+              and kernel_dtype != torch.float8_e4m3fn) else None),
         )
 
     def run_mla_generation(
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 04b17e36702c..12b1503f10a5 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -160,6 +160,15 @@ def effective_beam_width(self) -> int:
     helix_is_inactive_rank: Optional[torch.Tensor] = None
     helix_is_inactive_rank_cpu: Optional[torch.Tensor] = None
 
+    # Per-token helix state for speculative verify groups (a 1 + draft_len
+    # group may straddle a ledger-page boundary onto two CP ranks, so the
+    # per-sequence boolean above is insufficient there). See
+    # recompute_helix_spec_buffers for the derivation.
+    helix_local_slots: Optional[torch.Tensor] = None
+    helix_kv_bounds: Optional[torch.Tensor] = None
+    helix_owned_new_tokens_cpu: Optional[torch.Tensor] = None
+    _helix_spec_tokens_valid: bool = False
+
     # Block offsets for the target and draft KV caches
     kv_cache_block_offsets: Optional[torch.Tensor] = None
     host_kv_cache_block_offsets: Optional[torch.Tensor] = None
@@ -544,6 +553,41 @@ def _post_init_with_buffers(self, buffers) -> None:
                 device='cpu',
                 pin_memory=prefer_pinned(),
             )
+            # Per-token buffers for speculative verify groups under helix.
+            # A group of 1 + draft_len tokens can straddle a ledger-page
+            # boundary, splitting ownership between two CP ranks, so the
+            # per-sequence flag above is not expressive enough:
+            #   helix_local_slots[t]: rank-local KV write slot of gen token t
+            #     on this rank, or -1 when another rank owns its position
+            #     (consumed by the mla_rope_generation append kernel).
+            #   helix_kv_bounds[t]: number of rank-local KV entries token t
+            #     may attend to, i.e. local_len(pos_t + 1) (consumed by the
+            #     CuTe DSL MLA decode mask and the helix stats identity).
+            # Filled by recompute_helix_spec_buffers() on the spec path only.
+            self.helix_local_slots = self.get_empty(
+                buffers,
+                (self.max_num_tokens, ),
+                cache_name="helix_local_slots",
+                dtype=torch.int,
+                capture_graph=capture_graph,
+            )
+            self.helix_kv_bounds = self.get_empty(
+                buffers,
+                (self.max_num_tokens, ),
+                cache_name="helix_kv_bounds",
+                dtype=torch.int,
+                capture_graph=capture_graph,
+            )
+            # Host-side per-sequence count of this step's new tokens owned by
+            # this rank (spec path; single-token path derives it from the
+            # boolean flag). Consumed by prepare()'s helix kv_lens branch.
+            self.helix_owned_new_tokens_cpu = torch.zeros(
+                (self.max_num_sequences, ),
+                device='cpu',
+                dtype=torch.int,
+                pin_memory=prefer_pinned(),
+            )
+            self._helix_spec_tokens_valid = False
 
     def on_update_kv_lens(self):
         # After changing the kv_lens/kv_lens_cuda, we may need to update other metadata.
@@ -573,6 +617,7 @@ def update_helix_param(
         self,
         helix_position_offsets: List[int],
         helix_is_inactive_rank: List[bool],
+        helix_owned_new_tokens: Optional[List[int]] = None,
     ) -> None:
         """
         Update helix parameters by copying into static buffers for CUDA graph compatibility.
@@ -580,6 +625,10 @@ def update_helix_param(
         Args:
             helix_position_offsets: Position offsets for helix parallelism with shape (num_tokens,).
             helix_is_inactive_rank: Whether the current rank is inactive with shape (batch_size,).
+            helix_owned_new_tokens: Per-sequence count of this step's new
+                tokens owned by this rank (speculative verify groups; one
+                group may straddle a page boundary onto two ranks). None on
+                the single-token path, where the boolean flag carries it.
         """
         if helix_position_offsets is not None and self.helix_position_offsets is not None:
             num_tokens = len(helix_position_offsets)
@@ -595,6 +644,107 @@ def update_helix_param(
             self.helix_is_inactive_rank[:batch_size].copy_(
                 self.helix_is_inactive_rank_cpu[:batch_size], non_blocking=True)
 
+        self._helix_spec_tokens_valid = False
+        if helix_owned_new_tokens is not None:
+            batch_size = len(helix_owned_new_tokens)
+            self.helix_owned_new_tokens_cpu[:batch_size].copy_(
+                torch.tensor(helix_owned_new_tokens, dtype=torch.int))
+            self._helix_spec_tokens_valid = True
+
+    def helix_local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor:
+        """Vectorized rank-local prefix length for helix round-robin pages.
+
+        For each global sequence length g, returns the number of the first g
+        tokens whose ledger page lives on this CP rank (page b -> rank
+        b % cp_size). Mirrors KVCacheManagerV2._helix_local_len.
+        """
+        phys = self.kv_cache_manager.tokens_per_block
+        cp_size = self.mapping.cp_size
+        cp_rank = self.mapping.cp_rank
+        ledger = phys * cp_size
+        full = torch.div(global_lens, ledger, rounding_mode='floor')
+        rem = global_lens - full * ledger
+        return full * phys + (rem - cp_rank * phys).clamp_(0, phys)
+
+    def recompute_helix_spec_buffers(self, num_ctx_tokens: int,
+                                     num_gen_tokens: int,
+                                     tokens_per_gen_seq: int) -> None:
+        """Derive per-token helix buffers from (corrected) global positions.
+
+        Called after the overlap-scheduler device correction has been applied
+        to helix_position_offsets, so every derived quantity reflects the
+        real committed length even though the host packed provisional values.
+        Static shapes only; safe under CUDA graph capture.
+        """
+        pos = self.helix_position_offsets[num_ctx_tokens:num_ctx_tokens +
+                                          num_gen_tokens]
+        phys = self.kv_cache_manager.tokens_per_block
+        cp_rank = self.mapping.cp_rank
+        cp_size = self.mapping.cp_size
+        owner = torch.div(pos, phys, rounding_mode='floor') % cp_size
+        active = owner == cp_rank
+        local_before = self.helix_local_len_vec(pos)
+        self.helix_local_slots[num_ctx_tokens:num_ctx_tokens +
+                               num_gen_tokens].copy_(
+                                   torch.where(active, local_before,
+                                               local_before.new_full((), -1)))
+        self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens +
+                             num_gen_tokens].copy_(
+                                 self.helix_local_len_vec(pos + 1))
+        # Per-sequence rank-local kv length = bound of the sequence's last
+        # token (attention over committed + owned in-flight tokens).
+        assert num_gen_tokens % tokens_per_gen_seq == 0, (
+            f"helix spec expects uniform verify groups: {num_gen_tokens} gen "
+            f"tokens not divisible by group size {tokens_per_gen_seq}")
+        num_gen_seqs = num_gen_tokens // tokens_per_gen_seq
+        last_bounds = self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens +
+                                           num_gen_tokens].view(
+                                               num_gen_seqs,
+                                               tokens_per_gen_seq)[:, -1]
+        self.kv_lens_cuda[self.num_contexts:self.num_contexts +
+                          num_gen_seqs].copy_(last_bounds)
+
+    def build_ugpu_block_offsets(
+        self,
+        source_block_offsets: Optional[torch.Tensor],
+        buffer_attr: str,
+    ) -> Optional[List[torch.Tensor]]:
+        if (not self.ugpu_enabled or source_block_offsets is None
+                or self.kv_cache_manager is None):
+            return None
+
+        num_ugpus = self.kv_cache_manager.num_ugpus
+        buffers = getattr(self, buffer_attr, None)
+        if buffers is None:
+            buffers = [
+                torch.zeros_like(source_block_offsets) for _ in range(num_ugpus)
+            ]
+            setattr(self, buffer_attr, buffers)
+
+        packed_block_offsets = []
+        num_ctx = self.num_contexts
+        for ugpu_idx in range(num_ugpus):
+            ctx_start = self.ugpu_ctx_req_splits[ugpu_idx]
+            ctx_end = self.ugpu_ctx_req_splits[ugpu_idx + 1]
+            gen_start = num_ctx + self.ugpu_gen_req_splits[ugpu_idx]
+            gen_end = num_ctx + self.ugpu_gen_req_splits[ugpu_idx + 1]
+            num_ctx_reqs = ctx_end - ctx_start
+            num_gen_reqs = gen_end - gen_start
+            num_seqs = num_ctx_reqs + num_gen_reqs
+            block_offsets = buffers[ugpu_idx]
+            block_offsets.zero_()
+            if num_seqs > 0:
+                if num_ctx_reqs > 0:
+                    block_offsets[:, :num_ctx_reqs].copy_(
+                        source_block_offsets[:, ctx_start:ctx_end],
+                        non_blocking=True)
+                if num_gen_reqs > 0:
+                    block_offsets[:, num_ctx_reqs:num_seqs].copy_(
+                        source_block_offsets[:, gen_start:gen_end],
+                        non_blocking=True)
+            packed_block_offsets.append(block_offsets)
+        return packed_block_offsets
+
     def _bind_runtime_views(
         self,
         *,
@@ -722,9 +872,18 @@ def prepare(self) -> None:
         if self.enable_helix:
             # If helix is inactive, attend to the previously cached tokens only.
             assert cached_token_lens is not None, "cached_token_lens should be set for helix"
-            active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs]
-            kv_lens = cached_token_lens.clone()
-            kv_lens[active_rank] += self.seq_lens_kv[active_rank]
+            if getattr(self, '_helix_spec_tokens_valid', False):
+                # Speculative verify groups: a group may straddle a page
+                # boundary, so ownership of this step's new tokens is a
+                # per-sequence COUNT, not a boolean. Provisional host values;
+                # recompute_helix_spec_buffers overrides the device copy
+                # after the overlap correction.
+                kv_lens = cached_token_lens + \
+                    self.helix_owned_new_tokens_cpu[:self.num_seqs]
+            else:
+                active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs]
+                kv_lens = cached_token_lens.clone()
+                kv_lens[active_rank] += self.seq_lens_kv[active_rank]
         else:
             kv_lens = cached_token_lens + \
                 self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv
@@ -2444,6 +2603,11 @@ def mla_rope_generation(
         helix_tensor_params = [
             metadata.helix_position_offsets, metadata.helix_is_inactive_rank
         ]
+        if getattr(metadata, '_helix_spec_tokens_valid', False):
+            # Speculative verify groups: per-token KV write slots (-1 = this
+            # rank does not own the token's position). The append kernel then
+            # gates and addresses per token instead of per sequence.
+            helix_tensor_params.append(metadata.helix_local_slots)
 
         torch.ops.trtllm.mla_rope_generation(
             fused_q,
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index face56e566c9..d5c06ba5700a 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -9694,7 +9694,10 @@ def forward(
                     tensor of shape (H, S_q, B) remains in the workspace.
             """
             (q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o,
-             workspace, softmax_stats) = inputs
+             workspace, softmax_stats) = inputs[:9]
+            # inputs[9] (optional): helix per-token attention bounds of shape
+            # (B * S_q,), int32 — speculative verify groups only.
+            kv_bounds = inputs[9] if len(inputs) > 9 else None
             softmax_scale = float(kwargs.get("softmax_scale", 1.0))
             output_scale = float(kwargs.get("output_scale", 1.0))
 
@@ -9767,12 +9770,25 @@ def forward(
             split_workspace = workspace_bytes[split_kv_offset:split_kv_offset +
                                               split_kv_size]
 
+            if kv_bounds is not None:
+                expected_bounds_shape = (batch_size * seq_len_q, )
+                if (kv_bounds.shape != expected_bounds_shape
+                        or kv_bounds.dtype != torch.int32
+                        or kv_bounds.device != o.device
+                        or not kv_bounds.is_contiguous()):
+                    raise RuntimeError(
+                        "CuteDSLNVMlaDecodeBlackwellRunner requires contiguous "
+                        "int32 kv_bounds on the output device with shape "
+                        f"{expected_bounds_shape}, got shape="
+                        f"{tuple(kv_bounds.shape)}, dtype={kv_bounds.dtype}.")
+
             cache_key = self.unique_id() + (
                 out_dtype,
                 mma_qk_tiler_mn,
                 mma_pv_tiler_mn,
                 split_kv,
                 is_persistent,
+                kv_bounds is not None,
             )
             if cache_key not in CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache:
                 # A compile outside the tuning window stalls the serving loop
@@ -9840,6 +9856,9 @@ def forward(
                                 if use_workspace else None)
                 cache_seqs_ct = cute.runtime.from_dlpack(
                     cache_seqs, assumed_align=16).mark_layout_dynamic()
+                kv_bounds_ct = (cute.runtime.from_dlpack(
+                    kv_bounds, assumed_align=4).mark_layout_dynamic()
+                                if kv_bounds is not None else None)
                 # Variable split-KV (block_split_kvs) is not used on this path:
                 block_split_kvs_ct = None
 
@@ -9860,6 +9879,7 @@ def forward(
                     workspace_ct,
                     split_kv,
                     cache_seqs_ct,
+                    kv_bounds_ct,
                     block_split_kvs_ct,
                     cutlass.Float32(softmax_scale),
                     cutlass.Float32(output_scale),
@@ -9890,6 +9910,7 @@ def forward(
                 (split_kv > 1 and split_workspace.numel() > 0) else None,
                 split_kv,
                 cache_seqs,
+                kv_bounds,
                 None,  # block_split_kvs: var-split path unused (is_var_split_kv False)
                 softmax_scale,
                 output_scale,
@@ -9917,14 +9938,19 @@ def cute_dsl_mla_decode_fp8_blackwell(
         page_size: int,
         softmax_scale: float,
         output_scale: float,
-        # Keep the last two arguments required in the custom-op schema. PyTorch
+        # Keep the trailing arguments required in the custom-op schema. PyTorch
         # elides trailing default-valued arguments before its mutation fallback,
         # while mutates_args retains their positional indices.
         max_batch_size: int,
         softmax_stats: Optional[torch.Tensor],
+        kv_bounds: Optional[torch.Tensor],
     ) -> None:
         """CuTe DSL FP8 MLA decode (Blackwell SM100/SM103).
         """
+        if kv_bounds is not None:
+            raise ValueError(
+                "trtllm::cute_dsl_mla_decode_fp8_blackwell does not support "
+                "helix per-token kv_bounds (bf16/fp16 kernel only).")
         if (sm_version := get_sm_version()) not in (100, 103):
             raise ValueError(
                 f"trtllm::cute_dsl_mla_decode_fp8_blackwell requires SM 100 or "
@@ -9980,6 +10006,7 @@ def _(
         output_scale: float,
         max_batch_size: int,
         softmax_stats: Optional[torch.Tensor],
+        kv_bounds: Optional[torch.Tensor],
     ) -> None:
         return None
 
@@ -10005,8 +10032,12 @@ def cute_dsl_mla_decode_fp16_blackwell(
         # See the FP8 op above: these must remain required schema arguments.
         max_batch_size: int,
         softmax_stats: Optional[torch.Tensor],
+        kv_bounds: Optional[torch.Tensor],
     ) -> None:
         """CuTe DSL FP16/BF16 MLA decode (Blackwell SM100/SM103).
+
+        kv_bounds: helix speculative verify groups — per-token rank-local
+        attention bounds of shape (B * seq_len_q,), int32.
         """
         if (sm_version := get_sm_version()) not in (100, 103):
             raise ValueError(
@@ -10042,7 +10073,7 @@ def cute_dsl_mla_decode_fp16_blackwell(
         )
         inputs = [
             q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o,
-            workspace, softmax_stats
+            workspace, softmax_stats, kv_bounds
         ]
         tuner = AutoTuner.get()
         _, best_tactic = tuner.choose_one(
@@ -10080,5 +10111,6 @@ def _(
         output_scale: float,
         max_batch_size: int,
         softmax_stats: Optional[torch.Tensor],
+        kv_bounds: Optional[torch.Tensor],
     ) -> None:
         return None
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
index 7ed708faaacd..a19e74332835 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
@@ -311,6 +311,7 @@ def __call__(
         workspace: cute.Tensor,
         split_kv: cutlass.Int32,
         cache_seqs: Optional[cute.Tensor],
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         output_scale: cutlass.Float32,
@@ -392,6 +393,15 @@ def _run(
     ):
         """Execute the Multi-Head Latent Attention operation on the provided tensors.
 
+        kv_bounds (helix speculative verify groups): optional int32 tensor of
+        shape [batch_size * seq_len_q]; entry b*seq_len_q + q gives the number
+        of this rank's local KV entries query token q of sequence b may attend
+        to (committed prefix + owned in-flight group tokens up to and
+        including itself). When present it replaces the implicit causal bound
+        K - (seq_len_q - 1) + q_tok; values are guaranteed to lie in
+        [K - seq_len_q, K], i.e. inside the span the masked phase already
+        covers for the causal case.
+
         The method handles:
         1. Initialization of workspace for temporary split KV buffers
         2. Validation of tensor data types
@@ -792,6 +802,7 @@ class SplitKVKernelSharedStorage:
                 acc_lse,
                 split_kv,
                 cache_seqs,
+                kv_bounds,
                 block_split_kvs,
                 softmax_scale_log2,
                 output_scale,
@@ -827,6 +838,7 @@ class SplitKVKernelSharedStorage:
                 acc_lse,
                 split_kv,
                 cache_seqs,
+                kv_bounds,
                 block_split_kvs,
                 softmax_scale_log2,
                 output_scale,
@@ -859,6 +871,7 @@ class SplitKVKernelSharedStorage:
                     split_kv,
                     cache_seqs,
                     block_split_kvs,
+                    kv_bounds,
                 )
             else:
                 reduction_kernel = self.reduction_kernel(
@@ -938,6 +951,7 @@ def split_kv_kernel(
         mAccLSE: Optional[cute.Tensor],
         split_kv: cutlass.Int32,
         cache_seqs: cute.Tensor,
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: cute.Tensor,
         softmax_scale_log2: cutlass.Float32,
         output_scale: cutlass.Float32,
@@ -1356,6 +1370,7 @@ def split_kv_kernel(
                         mAccO=mAccO,
                         mO=mO,
                         K=cache_seqs[blk_coord[2]],
+                        kv_bounds=kv_bounds,
                         L=mCL.shape[1],
                         tmem_ptr=tmem_ptr,
                         tidx=tidx,
@@ -1458,6 +1473,7 @@ def reduction_kernel(
         split_kv: cutlass.Int32,
         cache_seqs: cute.Tensor,
         block_split_kvs: cute.Tensor,
+        kv_bounds: Optional[cute.Tensor] = None,
     ):
         """The reduction kernel for Multi-Head Latent Attention (MLA) that combines intermediate results
         from multiple split_kv blocks into final outputs.
@@ -1528,7 +1544,16 @@ def reduction_kernel(
             if tidx == 0:
                 mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse
                 if cutlass.const_expr(self.emit_softmax_stats):
-                    if cache_seqs[blk_coord[2]] > 0:
+                    # A rank joins the CP merge only for tokens with at least
+                    # one visible local KV entry. With verify groups that is
+                    # per-token: a straddling group leaves this rank zero
+                    # entries for its leading tokens while later ones have some.
+                    if cutlass.const_expr(kv_bounds is not None):
+                        has_local_kv = kv_bounds[blk_coord[2] * self.seq_len_q +
+                                                 blk_coord[1]] > 0
+                    else:
+                        has_local_kv = cache_seqs[blk_coord[2]] > 0
+                    if has_local_kv:
                         mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2],
                                       0] = global_lse / LOG2_E
                         mSoftmaxStats[blk_coord[0], blk_coord[1], blk_coord[2],
@@ -2448,8 +2473,14 @@ def compute(
         # positions. Min k_bound = K - (S_q-1), which can span up to
         # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For
         # S_q=1 this reduces to 1 tile -- identical to a plain K-bound check.
+        # With helix per-token bounds the minimum is K - S_q (a rank owning
+        # none of the group's tokens), one position deeper, so widen the span
+        # by one.
         tile_n = self.mma_qk_tiler[1]
-        mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1
+        if cutlass.const_expr(common_params.kv_bounds is not None):
+            mask_tile_count = (self.seq_len_q - 1 + tile_n - 1) // tile_n + 1
+        else:
+            mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1
 
         # first_mask_tile_idx is the global index of the first tile that may
         # need masking. Runtime because it depends on K (per-batch in
@@ -2777,7 +2808,16 @@ def softmax(
                              cta_m_rows) // self.num_heads)
                     else:
                         q_tok = common_params.blk_coord[1]
-                    k_bound = common_params.K - (self.seq_len_q - 1) + q_tok
+                    if cutlass.const_expr(common_params.kv_bounds is not None):
+                        # Helix verify groups: per-token rank-local bound
+                        # (committed prefix + owned group tokens <= q_tok);
+                        # subsumes the causal offset and non-owner ranks.
+                        k_bound = common_params.kv_bounds[
+                            common_params.blk_coord[2] * self.seq_len_q +
+                            q_tok]
+                    else:
+                        k_bound = common_params.K - (self.seq_len_q -
+                                                     1) + q_tok
                     tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less(
                         tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index,
                         k_bound,
@@ -2817,7 +2857,15 @@ def softmax(
                              cta_m_rows) // self.num_heads)
                     else:
                         q_tok = common_params.blk_coord[1]
-                    k_bound = common_params.K - (self.seq_len_q - 1) + q_tok
+                    if cutlass.const_expr(common_params.kv_bounds is not None):
+                        # Helix verify groups: per-token rank-local bound
+                        # (see the sm_100 branch above).
+                        k_bound = common_params.kv_bounds[
+                            common_params.blk_coord[2] * self.seq_len_q +
+                            q_tok]
+                    else:
+                        k_bound = common_params.K - (self.seq_len_q -
+                                                     1) + q_tok
                     tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less(
                         tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index,
                         k_bound,
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
index adfd5cc78726..4e3d14c17e37 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
@@ -309,6 +309,10 @@ def __call__(
         workspace: cute.Tensor,
         split_kv: cutlass.Int32,
         cache_seqs: Optional[cute.Tensor],
+        # Signature parity with the fp16 kernel (helix speculative verify
+        # groups); the fp8 kernel does not implement per-token bounds and the
+        # wrapper gate never selects fp8 KV under helix.
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         output_scale: cutlass.Float32,
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index dce32f886909..5675605da205 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2054,11 +2054,24 @@ def _setup_helix_mappings(
                 "per-request locality of KDA recurrent state."
             )
         if spec_config is not None:
-            raise ValueError(
-                "Kimi K3 helix phase 1 does not support speculative "
-                "decoding (round-robin KV bookkeeping assumes one token "
-                "per decode step)."
-            )
+            # Helix speculative decoding allowlist: standalone DSpark linear
+            # chains verified on the V2 superblock ledger (per-token ownership
+            # bookkeeping + CuTe DSL per-token bounds). Everything else stays
+            # loudly rejected — the #16003 review showed unsupported spec
+            # modes silently running wrong under helix.
+            decoding_type = getattr(spec_config, "decoding_type", None)
+            if decoding_type != "DSpark":
+                raise ValueError(
+                    "Kimi K3 helix supports speculative decoding only with "
+                    f"DSpark (standalone drafter); got {decoding_type!r}."
+                )
+            if getattr(spec_config, "draft_is_embedded_in_target", False):
+                raise ValueError(
+                    "Kimi K3 helix supports only the standalone DSpark "
+                    "drafter; the embedded (in-target) flavour shares the "
+                    "target KV bookkeeping in ways the helix ledger does "
+                    "not model."
+                )
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp
         if cfg.num_attention_heads % repurposed_tp != 0:
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index b4422c084463..9db0e6774389 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1562,10 +1562,19 @@ def _create_one_model_draft_kv_cache_manager(
         # the sparse_attention_config. Get it from effective_draft_config which
         # falls back to the target model's config for MTP mode.
         sparse_attn_config = effective_draft_config.sparse_attention_config
+        # The standalone drafter is a plain dense model built against the
+        # repurposed mapping under helix (CP ranks become TP ranks; each rank
+        # keeps its full drafter KV). Its paged KV manager must therefore use
+        # the repurposed, CP-free mapping — helix round-robin ledger semantics
+        # apply only to the TARGET KV (and KVCacheManagerV2 rejects
+        # is_draft x helix outright).
+        draft_mapping = self._mapping
+        if draft_mapping.has_cp_helix():
+            draft_mapping = draft_mapping.repurpose_helix_cp_to_tp()
         return _create_kv_cache_manager(
             model_engine=None,
             kv_cache_manager_cls=draft_kv_cache_manager_cls,
-            mapping=self._mapping,
+            mapping=draft_mapping,
             kv_cache_config=draft_kv_config,
             tokens_per_block=self._tokens_per_block,
             max_seq_len=max_seq_len,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index e0ab193c50fd..33c95a22cbd9 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3946,6 +3946,28 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                         num_chunked_contexts=num_chunked_ctx_requests,
                     )
 
+                if self.mapping.has_cp_helix() and getattr(
+                        inputs['attn_metadata'], '_helix_spec_tokens_valid',
+                        False):
+                    md = inputs['attn_metadata']
+                    # The helix position buffer holds the same provisional
+                    # global positions as position_ids (packed from the stale
+                    # base); apply the same accepted-count correction, then
+                    # re-derive every helix quantity (per-token write slots,
+                    # per-token attention bounds, per-seq rank-local kv lens)
+                    # from the corrected positions. The kv_lens override
+                    # below supersedes the generic previous_kv_lens_offsets
+                    # adjustment above, which is not ownership-aware.
+                    md.helix_position_offsets[:previous_batch_tokens] += (
+                        self.previous_pos_id_offsets_cuda[:
+                                                          previous_batch_tokens]
+                    )
+                    md.recompute_helix_spec_buffers(
+                        0, previous_batch_tokens,
+                        self.get_runtime_tokens_per_gen_step(
+                            self.runtime_draft_len))
+                    md.on_update_kv_lens()
+
         if self.guided_decoder is not None:
             self.guided_decoder.token_event.record()
 
@@ -5595,6 +5617,24 @@ def append_cross_attention_state(request: LlmRequest,
                 generation_requests.append(request)
         extend_requests += extend_dummy_requests
 
+        # Helix bookkeeping is needed by BOTH the extend (speculative verify
+        # group) and the plain generation packing loops below, so initialize
+        # it ahead of them. Positions are global; KV ownership follows the
+        # round-robin ledger (page b -> rank b % cp), mirrored host-side here
+        # (KVCacheManagerV2._helix_local_len) for provisional packing values.
+        helix_is_inactive_rank, helix_position_offsets = [], []
+        helix_owned_new_tokens = []
+        _has_cp_helix = self.mapping.has_cp_helix()
+        if _has_cp_helix and kv_cache_manager is not None:
+            _helix_phys = kv_cache_manager.tokens_per_block
+            _helix_ledger = _helix_phys * self.mapping.cp_size
+            _helix_rank_off = self.mapping.cp_rank * _helix_phys
+
+            def _helix_local_len_host(global_len: int) -> int:
+                full, rem = divmod(global_len, _helix_ledger)
+                return full * _helix_phys + min(
+                    max(rem - _helix_rank_off, 0), _helix_phys)
+
         spec_config = self.spec_config if self.enable_spec_decode else None
         if not self._disable_overlap_scheduler and spec_config is not None:
             assert spec_config.spec_dec_mode.support_overlap_scheduler(
@@ -5667,6 +5707,21 @@ def append_cross_attention_state(request: LlmRequest,
                 num_cached_tokens_per_seq.append(
                     past_seen_token_num - request.py_num_compressed_tokens)
                 request.cached_tokens = past_seen_token_num
+                if _has_cp_helix:
+                    # Verify group [base, base+group): the global positions
+                    # already sit in position_ids above; KV numbers are
+                    # rank-local. This branch has no in-flight predecessor,
+                    # so every value is exact (no device correction needed).
+                    group = 1 + num_draft_tokens
+                    base = past_seen_token_num
+                    helix_position_offsets.extend(range(base, base + group))
+                    helix_is_inactive_rank.append(False)
+                    local_cached = _helix_local_len_host(base)
+                    helix_owned_new_tokens.append(
+                        _helix_local_len_host(base + group) - local_cached)
+                    num_cached_tokens_per_seq[-1] = (
+                        local_cached - request.py_num_compressed_tokens)
+                    request.cached_tokens = local_cached
                 # update batch index
                 request.py_batch_idx = request.py_seq_slot
             else:
@@ -5698,6 +5753,21 @@ def append_cross_attention_state(request: LlmRequest,
                     request.py_num_compressed_tokens)
                 request.cached_tokens = (past_seen_token_num +
                                          runtime_tokens_per_gen_step)
+                if _has_cp_helix:
+                    # In-flight predecessor: mirror the non-helix convention
+                    # above — positions are packed from the stale base (the
+                    # overlap device correction adds the accepted count) and
+                    # KV numbers assume full acceptance (the device recompute
+                    # in recompute_helix_spec_buffers overrides them).
+                    group = runtime_tokens_per_gen_step
+                    base = past_seen_token_num
+                    helix_position_offsets.extend(range(base, base + group))
+                    helix_is_inactive_rank.append(False)
+                    local_full = _helix_local_len_host(base + group)
+                    helix_owned_new_tokens.append(0)
+                    num_cached_tokens_per_seq[-1] = (
+                        local_full - request.py_num_compressed_tokens)
+                    request.cached_tokens = local_full
                 if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend) and spec_config.is_linear_tree:
                     prompt_lengths.append(runtime_tokens_per_gen_step)
@@ -5765,9 +5835,9 @@ def append_cross_attention_state(request: LlmRequest,
             # update batch index
             request.py_batch_idx = request.py_seq_slot
 
-        helix_is_inactive_rank, helix_position_offsets = [], []
-        # Cache invariant method result to avoid repeated calls per-request
-        _has_cp_helix = self.mapping.has_cp_helix()
+        # (helix lists and _has_cp_helix are initialized ahead of the extend
+        # loop above, which also appends to them for verify groups.)
+
         _n_gen = len(generation_requests)
         # One-shot batch-level flag — True iff any generation request actually
         # carries multimodal payload. Lets the strip_mm_data branch below
@@ -5897,6 +5967,10 @@ def append_cross_attention_state(request: LlmRequest,
                         helix_is_inactive_rank.append(
                             request.py_helix_is_inactive_rank)
                         helix_position_offsets.append(position_id)
+                        # Keep the per-seq owned-count list aligned when the
+                        # spec path is active in the same batch.
+                        helix_owned_new_tokens.append(
+                            0 if request.py_helix_is_inactive_rank else 1)
 
                 request.cached_tokens = past_seen_token_num
                 for beam in range(beam_width):
@@ -6320,6 +6394,11 @@ def previous_seq_slots_device():
             attn_metadata.update_helix_param(
                 helix_position_offsets=helix_position_offsets,
                 helix_is_inactive_rank=helix_is_inactive_rank,
+                # Per-seq owned counts drive the kv_lens math only on the
+                # speculative path (verify groups); None keeps the
+                # single-token boolean convention.
+                helix_owned_new_tokens=(helix_owned_new_tokens
+                                        if self.enable_spec_decode else None),
             )
 
         if not attn_metadata.is_cuda_graph:

From 1f549dd93ea2421872989e760defb39e8f336adb Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Mon, 24 Aug 2026 14:02:33 -0700
Subject: [PATCH 02/15] [None][fix] helix x DSpark review fixes: global base,
 overlap-off recompute, loud fallback rejection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adversarial review of fdb9594b21 confirmed nine findings with four roots:
- CRITICAL: both extend packing branches used max_beam_num_tokens-1 as the
  global position base, but a helix gen worker's token list is the
  rank-LOCAL round-robin subset (merge_helix_requests). Reconstruct the
  global base as total_input_len_cp + (max_beam_num_tokens -
  py_prompt_len) - 1 and override the extend tokens' position_ids with the
  global values (same convention as the non-spec helix loop).
- CRITICAL: recompute_helix_spec_buffers only ran inside the
  overlap-enabled branch of _preprocess_inputs while every consumer arms on
  enable_spec_decode alone — overlap-off would read uninitialized slot/bound
  buffers. The helix recompute now runs on every spec step; the stale-base
  position correction stays overlap-only.
- FallbackFmha now rejects helix verify groups (per-token ownership is
  inexpressible in the fused thop path) so a CuTe DSL rejection surfaces as
  a loud no-library error instead of silently wrong attention.
- Autotuner profiling: re-derive a size-consistent kv_bounds dummy for
  bucketed batches (input 9 has no dynamic-dim spec).
- CuTe DSL fold_sq padding rows: clamp the per-token bounds index to stay
  in gmem bounds.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
(cherry picked from commit 37b561565da500c7f38dd9268a3ed3dd0b85fa3b)
(cherry picked from commit fb9b82c8f9ed9a16cda168622cf4fb1da172f0a8)
---
 .../_torch/attention_backend/fmha/fallback.py | 14 +++-
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  8 ++
 .../attention/mla/mla_decode_fp16.py          | 16 +++-
 .../_torch/pyexecutor/model_engine.py         | 74 ++++++++++++-------
 4 files changed, 81 insertions(+), 31 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
index 4ce5f8e149f1..4d9530f65a62 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
@@ -76,7 +76,19 @@ def is_supported(
         *,
         phase: Optional[FmhaPhase] = None,
     ) -> bool:
-        del q, k, v, phase
+        del k, v, phase
+        # Helix speculative verify groups carry per-token KV ownership (a
+        # group may straddle a page boundary onto two CP ranks) that the
+        # fused thop attention path cannot express: its spec-dec mask and the
+        # per-sequence helix_is_inactive_rank gate both assume the group's
+        # new KV entries are the trailing slots of one rank's kv_len. Running
+        # it would be silently wrong, so reject — with no remaining FMHA
+        # library, dispatch raises loudly instead.
+        if (metadata.helix_position_offsets is not None
+                and getattr(metadata, "_helix_spec_tokens_valid", False)
+                and metadata.num_generations > 0
+                and q.shape[0] > metadata.num_seqs):
+            return False
         return forward_args.attention_mask != CustomAttentionMask.CUSTOM and (
             forward_args.update_kv_cache or metadata.is_cross
         )
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index d5c06ba5700a..6e5843a5e76e 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -9770,6 +9770,14 @@ def forward(
             split_workspace = workspace_bytes[split_kv_offset:split_kv_offset +
                                               split_kv_size]
 
+            if kv_bounds is not None and AutoTuner.get().is_tuning_mode:
+                # Profiling rebuilds bucketed cache_seqs but carries
+                # kv_bounds through unchanged (input 9 has no dynamic-dim
+                # spec); re-derive a size-consistent dummy — bound values
+                # only affect masking depth, not the tactic space.
+                if kv_bounds.numel() != batch_size * seq_len_q:
+                    kv_bounds = cache_seqs.repeat_interleave(
+                        seq_len_q).contiguous()
             if kv_bounds is not None:
                 expected_bounds_shape = (batch_size * seq_len_q, )
                 if (kv_bounds.shape != expected_bounds_shape
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
index a19e74332835..7973c157284e 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
@@ -2812,9 +2812,15 @@ def softmax(
                         # Helix verify groups: per-token rank-local bound
                         # (committed prefix + owned group tokens <= q_tok);
                         # subsumes the causal offset and non-owner ranks.
+                        # Clamp for fold_sq M-tile padding rows (row beyond
+                        # num_heads * fold_sq_ratio derives q_tok >= S_q);
+                        # their results are discarded, but the gmem read must
+                        # stay in bounds.
+                        q_tok_c = (q_tok if cute.elem_less(
+                            q_tok, self.seq_len_q) else self.seq_len_q - 1)
                         k_bound = common_params.kv_bounds[
                             common_params.blk_coord[2] * self.seq_len_q +
-                            q_tok]
+                            q_tok_c]
                     else:
                         k_bound = common_params.K - (self.seq_len_q -
                                                      1) + q_tok
@@ -2860,9 +2866,15 @@ def softmax(
                     if cutlass.const_expr(common_params.kv_bounds is not None):
                         # Helix verify groups: per-token rank-local bound
                         # (see the sm_100 branch above).
+                        # Clamp for fold_sq M-tile padding rows (row beyond
+                        # num_heads * fold_sq_ratio derives q_tok >= S_q);
+                        # their results are discarded, but the gmem read must
+                        # stay in bounds.
+                        q_tok_c = (q_tok if cute.elem_less(
+                            q_tok, self.seq_len_q) else self.seq_len_q - 1)
                         k_bound = common_params.kv_bounds[
                             common_params.blk_coord[2] * self.seq_len_q +
-                            q_tok]
+                            q_tok_c]
                     else:
                         k_bound = common_params.K - (self.seq_len_q -
                                                      1) + q_tok
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 33c95a22cbd9..3607d08a1fbe 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3946,27 +3946,30 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                         num_chunked_contexts=num_chunked_ctx_requests,
                     )
 
-                if self.mapping.has_cp_helix() and getattr(
-                        inputs['attn_metadata'], '_helix_spec_tokens_valid',
-                        False):
-                    md = inputs['attn_metadata']
-                    # The helix position buffer holds the same provisional
-                    # global positions as position_ids (packed from the stale
-                    # base); apply the same accepted-count correction, then
-                    # re-derive every helix quantity (per-token write slots,
-                    # per-token attention bounds, per-seq rank-local kv lens)
-                    # from the corrected positions. The kv_lens override
-                    # below supersedes the generic previous_kv_lens_offsets
-                    # adjustment above, which is not ownership-aware.
-                    md.helix_position_offsets[:previous_batch_tokens] += (
-                        self.previous_pos_id_offsets_cuda[:
-                                                          previous_batch_tokens]
-                    )
-                    md.recompute_helix_spec_buffers(
-                        0, previous_batch_tokens,
-                        self.get_runtime_tokens_per_gen_step(
-                            self.runtime_draft_len))
-                    md.on_update_kv_lens()
+        if self.enable_spec_decode and self.mapping.has_cp_helix():
+            # Helix verify groups: the per-token device buffers (write slots,
+            # attention bounds, rank-local kv lens) must be derived on EVERY
+            # spec step, overlap or not — the append/mask kernels consume
+            # them whenever _helix_spec_tokens_valid is armed. Under overlap
+            # the host packed provisional positions from a stale base, so
+            # first apply the same accepted-count correction position_ids
+            # got above; without overlap the host values are already exact.
+            md = inputs.get('attn_metadata')
+            if (md is not None and md.kv_cache_manager is not None
+                    and getattr(md, '_helix_spec_tokens_valid', False)):
+                helix_gen_tokens = (inputs['input_ids'].shape[0] -
+                                    md.num_ctx_tokens)
+                if not self._disable_overlap_scheduler:
+                    # The kv_lens override in the recompute supersedes the
+                    # generic previous_kv_lens_offsets adjustment above,
+                    # which is not ownership-aware.
+                    md.helix_position_offsets[:helix_gen_tokens] += (
+                        self.previous_pos_id_offsets_cuda[:helix_gen_tokens])
+                md.recompute_helix_spec_buffers(
+                    0, helix_gen_tokens,
+                    self.get_runtime_tokens_per_gen_step(
+                        self.runtime_draft_len))
+                md.on_update_kv_lens()
 
         if self.guided_decoder is not None:
             self.guided_decoder.token_event.record()
@@ -5708,13 +5711,23 @@ def _helix_local_len_host(global_len: int) -> int:
                     past_seen_token_num - request.py_num_compressed_tokens)
                 request.cached_tokens = past_seen_token_num
                 if _has_cp_helix:
-                    # Verify group [base, base+group): the global positions
-                    # already sit in position_ids above; KV numbers are
-                    # rank-local. This branch has no in-flight predecessor,
-                    # so every value is exact (no device correction needed).
+                    # Verify group [base, base+group) in GLOBAL positions.
+                    # On a helix gen worker the request's token list is the
+                    # rank-LOCAL round-robin subset, so max_beam_num_tokens
+                    # (= local_prompt + generated) must NOT be used as a
+                    # global base; reconstruct it from the global prompt
+                    # length plus the (rank-invariant) generated count. This
+                    # branch has no in-flight predecessor, so every value is
+                    # exact (no device correction needed).
                     group = 1 + num_draft_tokens
-                    base = past_seen_token_num
+                    generated_len = (request.max_beam_num_tokens -
+                                     request.py_prompt_len)
+                    base = request.total_input_len_cp + generated_len - 1
                     helix_position_offsets.extend(range(base, base + group))
+                    # position_ids above were packed from the local base;
+                    # helix uses global position ids (same convention as the
+                    # non-spec helix generation loop below).
+                    position_ids[-group:] = range(base, base + group)
                     helix_is_inactive_rank.append(False)
                     local_cached = _helix_local_len_host(base)
                     helix_owned_new_tokens.append(
@@ -5758,10 +5771,15 @@ def _helix_local_len_host(global_len: int) -> int:
                     # above — positions are packed from the stale base (the
                     # overlap device correction adds the accepted count) and
                     # KV numbers assume full acceptance (the device recompute
-                    # in recompute_helix_spec_buffers overrides them).
+                    # in recompute_helix_spec_buffers overrides them). The
+                    # base is reconstructed GLOBALLY (see the no-previous
+                    # branch: the token list is rank-local under helix).
                     group = runtime_tokens_per_gen_step
-                    base = past_seen_token_num
+                    generated_len = (request.max_beam_num_tokens -
+                                     request.py_prompt_len)
+                    base = request.total_input_len_cp + generated_len - 1
                     helix_position_offsets.extend(range(base, base + group))
+                    position_ids[-group:] = range(base, base + group)
                     helix_is_inactive_rank.append(False)
                     local_full = _helix_local_len_host(base + group)
                     helix_owned_new_tokens.append(0)

From ab27100ccde5fda9d6fc668398688a59b26bcd3f Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Mon, 24 Aug 2026 22:17:53 -0700
Subject: [PATCH 03/15] [None][fix] helix x DSpark: sanitize zero-bound
 verify-group rows per token

External review finding: a CP rank owning only the tail page of a verify
group has zero visible KV for the group's leading tokens while its
per-sequence kv_len (bound of the LAST token) is nonzero, so the
per-sequence _helix_zero_kv_mask never sanitizes those partial_o rows.
The decode kernel fills fully-masked rows with a finite sentinel, making
them an average over arbitrary pool bytes; the stats identity gives them
corr = 0, but 0 * NaN = NaN would poison the token on every rank.
Sanitize by the per-token bound (helix_kv_bounds == 0) on the spec path.
Reachable only when a rank holds zero prompt KV (global prompt shorter
than cp * tokens_per_block), which GSM8K's ~1k prompts never hit.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
(cherry picked from commit 75565d3416de411310fa21370461f4e207daa11e)
(cherry picked from commit aa080622b38e1f8b8d5ad5e3dd07970dc9e47715)
---
 tensorrt_llm/_torch/modules/mla.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index 86b960c29d93..c4259b120930 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -769,6 +769,19 @@ def _attn_forward_gen(
                 seq_start=attn_metadata.num_contexts,
                 num_seqs=attn_metadata.num_generations,
             )
+            helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None)
+            if helix_kv_bounds is not None and getattr(
+                    attn_metadata, "_helix_spec_tokens_valid", False):
+                # Speculative verify groups: KV ownership is per-TOKEN. A rank
+                # owning only the tail page of a group has zero visible KV for
+                # the group's leading tokens while its per-sequence kv_len is
+                # nonzero, so the per-sequence mask above misses those rows.
+                # Their decode rows are fully masked with a finite sentinel,
+                # making partial_o an average over (possibly uninitialized)
+                # pool values; the combine multiplies by corr = 0 and
+                # 0 * NaN would poison the token on every CP rank — sanitize
+                # by the per-token bound instead.
+                zero_kv_mask = helix_kv_bounds[:partial_o.shape[0]] == 0
             return _helix_post_process(
                 partial_o,
                 softmax_stats,

From 421a85fd29eb1a8a4a2057e6c7fd9d03373411a7 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Mon, 24 Aug 2026 22:19:01 -0700
Subject: [PATCH 04/15] [None][fix] address review: enable DSpark in the K3
 spec-mode assert, sanitize zero-bound rows per token, drop leaked dead code

Three external-review findings on the draft:
- The class-level Kimi K3 spec-mode assert predates DSpark on this base
  and made the helix allowlist unreachable (helix+DSpark passed the
  allowlist then died on the assert); extend it with is_dspark(). The
  embedded-flavour rejection keeps a getattr that is a documented no-op
  until draft_is_embedded_in_target lands with the embedded DSv4 flavour.
- Sanitize zero-bound verify-group rows per token (cherry-picked from the
  integration branch): a rank owning only a group's tail page has
  nonzero per-seq kv_len but zero visible KV for leading tokens; their
  finite-sentinel decode rows average arbitrary pool bytes and 0 * NaN
  in the combine would poison the token. Reachable when the global
  prompt is shorter than cp * tokens_per_block.
- Remove build_ugpu_block_offsets: leaked from the integration branch by
  a conflict resolution; references attributes that do not exist on this
  base and has no callers.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
(cherry picked from commit 8d10f809eb5808f54079bc6086635adc456223e6)
---
 .../_torch/attention_backend/trtllm.py        | 41 -------------------
 .../_torch/models/modeling_kimi_linear.py     |  3 ++
 2 files changed, 3 insertions(+), 41 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 12b1503f10a5..a10a25180d62 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -704,47 +704,6 @@ def recompute_helix_spec_buffers(self, num_ctx_tokens: int,
         self.kv_lens_cuda[self.num_contexts:self.num_contexts +
                           num_gen_seqs].copy_(last_bounds)
 
-    def build_ugpu_block_offsets(
-        self,
-        source_block_offsets: Optional[torch.Tensor],
-        buffer_attr: str,
-    ) -> Optional[List[torch.Tensor]]:
-        if (not self.ugpu_enabled or source_block_offsets is None
-                or self.kv_cache_manager is None):
-            return None
-
-        num_ugpus = self.kv_cache_manager.num_ugpus
-        buffers = getattr(self, buffer_attr, None)
-        if buffers is None:
-            buffers = [
-                torch.zeros_like(source_block_offsets) for _ in range(num_ugpus)
-            ]
-            setattr(self, buffer_attr, buffers)
-
-        packed_block_offsets = []
-        num_ctx = self.num_contexts
-        for ugpu_idx in range(num_ugpus):
-            ctx_start = self.ugpu_ctx_req_splits[ugpu_idx]
-            ctx_end = self.ugpu_ctx_req_splits[ugpu_idx + 1]
-            gen_start = num_ctx + self.ugpu_gen_req_splits[ugpu_idx]
-            gen_end = num_ctx + self.ugpu_gen_req_splits[ugpu_idx + 1]
-            num_ctx_reqs = ctx_end - ctx_start
-            num_gen_reqs = gen_end - gen_start
-            num_seqs = num_ctx_reqs + num_gen_reqs
-            block_offsets = buffers[ugpu_idx]
-            block_offsets.zero_()
-            if num_seqs > 0:
-                if num_ctx_reqs > 0:
-                    block_offsets[:, :num_ctx_reqs].copy_(
-                        source_block_offsets[:, ctx_start:ctx_end],
-                        non_blocking=True)
-                if num_gen_reqs > 0:
-                    block_offsets[:, num_ctx_reqs:num_seqs].copy_(
-                        source_block_offsets[:, gen_start:gen_end],
-                        non_blocking=True)
-            packed_block_offsets.append(block_offsets)
-        return packed_block_offsets
-
     def _bind_runtime_views(
         self,
         *,
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 5675605da205..9ed340c001f6 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2065,6 +2065,9 @@ def _setup_helix_mappings(
                     "Kimi K3 helix supports speculative decoding only with "
                     f"DSpark (standalone drafter); got {decoding_type!r}."
                 )
+            # draft_is_embedded_in_target does not exist on this base yet (it
+            # arrives with the embedded DSv4 DSpark flavour); the getattr keeps
+            # this rejection forward-compatible and is a no-op until then.
             if getattr(spec_config, "draft_is_embedded_in_target", False):
                 raise ValueError(
                     "Kimi K3 helix supports only the standalone DSpark "

From 39a44a89fb2d3918b2d65c42410561346a7e097a Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Wed, 26 Aug 2026 00:58:32 -0700
Subject: [PATCH 05/15] [None][fix] helix x DSpark: CUDA-graph capture hygiene

Two capture-path defects in the helix speculative-decoding bookkeeping,
both of which only bite once CUDA graphs are enabled together with the
overlap scheduler.

1. `_preprocess_inputs()` corrects `helix_position_offsets` in place with
   `+= previous_pos_id_offsets_cuda`, but `_postprocess_inputs()` had no
   matching `-=`. Every other in-place correction on that path is
   symmetric (see `position_ids` a few lines above) precisely because the
   buffer is not rewritten between graph replays, so the missing reversal
   lets the offset accumulate step over step and the derived per-token
   owner/slot/bound values drift. Add the mirror, guarded by the same
   `_helix_spec_tokens_valid` condition. The recompute's OVERWRITES
   (slots/bounds/kv_lens) need no reversal: those buffers are rewritten
   from host state at the next step's prepare.

2. `recompute_helix_spec_buffers()` built the `-1` sentinel with
   `local_before.new_full((), -1)`, allocating a scalar tensor on every
   step of a routine that is captured into the graph. Use the scalar
   overload of `torch.where` instead so the captured region stays
   allocation-free.

Validated on Kimi-K3 disaggregated GSM8K with helix cp8 + DSpark
draft_len=7 and both CUDA graphs and the overlap scheduler enabled.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/trtllm.py |  5 +++--
 tensorrt_llm/_torch/pyexecutor/model_engine.py  | 11 +++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index a10a25180d62..b0f043072533 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -684,10 +684,11 @@ def recompute_helix_spec_buffers(self, num_ctx_tokens: int,
         owner = torch.div(pos, phys, rounding_mode='floor') % cp_size
         active = owner == cp_rank
         local_before = self.helix_local_len_vec(pos)
+        # Scalar overload: no per-step allocation (these ops are captured
+        # into the CUDA graph; keep them allocation-free).
         self.helix_local_slots[num_ctx_tokens:num_ctx_tokens +
                                num_gen_tokens].copy_(
-                                   torch.where(active, local_before,
-                                               local_before.new_full((), -1)))
+                                   torch.where(active, local_before, -1))
         self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens +
                              num_gen_tokens].copy_(
                                  self.helix_local_len_vec(pos + 1))
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 3607d08a1fbe..f7c54e3b4773 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -4030,6 +4030,17 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]):
                         restore=True,
                     )
 
+                if (self.mapping.has_cp_helix() and getattr(
+                        inputs['attn_metadata'], '_helix_spec_tokens_valid',
+                        False)):
+                    # Mirror of the helix position correction in
+                    # _preprocess_inputs (capture symmetry, like position_ids
+                    # above). The recompute's OVERWRITES (slots/bounds/
+                    # kv_lens) need no reversal: every consumer buffer is
+                    # rewritten from host state at the next step's prepare.
+                    inputs['attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= (
+                        self.previous_pos_id_offsets_cuda[:previous_batch_tokens])
+
     def _get_all_rank_num_tokens(self, attn_metadata: AttentionMetadata):
         if self.enable_attention_dp:
             num_tokens = attn_metadata.num_tokens

From addfde53c86bdef3288815fe0eed85ac2af93883 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Wed, 26 Aug 2026 01:08:47 -0700
Subject: [PATCH 06/15] [None][fix] helix x DSpark: make the embedded-drafter
 rejection fire, tighten metadata guards

Two review findings on the helix speculative-decoding guards.

1. The in-target-drafter rejection keyed on `draft_is_embedded_in_target`,
   an attribute that does not exist anywhere in the tree; the `getattr`
   default made the check a permanent no-op, and there is no guarantee the
   attribute would arrive under that name. Key it on
   `spec_config.speculative_model` instead, a `DecodingBaseConfig` field:
   a standalone drafter always carries its own checkpoint, so its absence
   identifies the embedded flavour. The check now actually fires.

   The rejection matters because the helix draft KV manager is built on the
   repurposed (CP-as-TP) mapping, where every rank keeps the FULL drafter
   KV and the round-robin ownership ledger applies to the target KV alone.
   An in-target drafter instead reads and writes the sharded target KV,
   which the per-token bookkeeping does not cover, and would run silently
   wrong.

2. Replace every `getattr(..., '_helix_spec_tokens_valid', False)` with a
   direct read. Where the object is statically `TrtllmAttentionMetadata`
   (`self` inside the metadata class, the annotated `metadata` parameter of
   `mla_rope_generation`, and the three FMHA entry points, all declared
   `TrtllmAttentionMetadata` in `fmha/interface.py`) the attribute is a
   dataclass field with a default and always exists. Where the object comes
   from generic code (`modules/mla.py`, `pyexecutor/model_engine.py`) guard
   with `isinstance(..., TrtllmAttentionMetadata)`, matching the existing
   idiom in those same functions, and read the fields directly after.

   Behaviour is unchanged: the helix fields and the `update_helix_param`
   override live only on `TrtllmAttentionMetadata` (the base-class hook is
   a no-op and neither FlashInfer nor Vanilla overrides it), so non-TRTLLM
   metadata could never have armed the flag. The guards are stricter,
   though: a `getattr` default silently degrades to False -- disabling the
   helix speculative path outright -- if a field is renamed, and would
   accept a same-named field on a backend that has no helix support.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 .../attention_backend/fmha/cute_dsl_mla.py    |  5 ++---
 .../_torch/attention_backend/fmha/fallback.py |  2 +-
 .../_torch/attention_backend/trtllm.py        |  4 ++--
 .../_torch/models/modeling_kimi_linear.py     | 14 ++++++-------
 tensorrt_llm/_torch/modules/mla.py            |  8 ++++----
 .../_torch/pyexecutor/model_engine.py         | 20 +++++++++----------
 6 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
index a6ea42deb7d5..ced37b77d8e3 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
@@ -309,8 +309,7 @@ def _is_supported_with_reason(
         seq_len_q = q.shape[0] // meta.num_generations
         batch_size = meta.num_generations
         if meta.helix_position_offsets is not None:
-            if seq_len_q != 1 and not getattr(meta, "_helix_spec_tokens_valid",
-                                              False):
+            if seq_len_q != 1 and not meta._helix_spec_tokens_valid:
                 # Multi-token decode under helix needs the per-token bound /
                 # write-slot buffers of the speculative verify-group path.
                 return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix."
@@ -526,7 +525,7 @@ def _run_mla_decode(
             # single-token helix path and outside helix.
             (meta.helix_kv_bounds[:num_tokens] if
              (meta.helix_position_offsets is not None
-              and getattr(meta, "_helix_spec_tokens_valid", False)
+              and meta._helix_spec_tokens_valid
               and kernel_dtype != torch.float8_e4m3fn) else None),
         )
 
diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
index 4d9530f65a62..2d251e396c75 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
@@ -85,7 +85,7 @@ def is_supported(
         # it would be silently wrong, so reject — with no remaining FMHA
         # library, dispatch raises loudly instead.
         if (metadata.helix_position_offsets is not None
-                and getattr(metadata, "_helix_spec_tokens_valid", False)
+                and metadata._helix_spec_tokens_valid
                 and metadata.num_generations > 0
                 and q.shape[0] > metadata.num_seqs):
             return False
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index b0f043072533..b65b7cbbb1fe 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -832,7 +832,7 @@ def prepare(self) -> None:
         if self.enable_helix:
             # If helix is inactive, attend to the previously cached tokens only.
             assert cached_token_lens is not None, "cached_token_lens should be set for helix"
-            if getattr(self, '_helix_spec_tokens_valid', False):
+            if self._helix_spec_tokens_valid:
                 # Speculative verify groups: a group may straddle a page
                 # boundary, so ownership of this step's new tokens is a
                 # per-sequence COUNT, not a boolean. Provisional host values;
@@ -2563,7 +2563,7 @@ def mla_rope_generation(
         helix_tensor_params = [
             metadata.helix_position_offsets, metadata.helix_is_inactive_rank
         ]
-        if getattr(metadata, '_helix_spec_tokens_valid', False):
+        if metadata._helix_spec_tokens_valid:
             # Speculative verify groups: per-token KV write slots (-1 = this
             # rank does not own the token's position). The append kernel then
             # gates and addresses per token instead of per sequence.
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 9ed340c001f6..e450216c187b 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2065,15 +2065,15 @@ def _setup_helix_mappings(
                     "Kimi K3 helix supports speculative decoding only with "
                     f"DSpark (standalone drafter); got {decoding_type!r}."
                 )
-            # draft_is_embedded_in_target does not exist on this base yet (it
-            # arrives with the embedded DSv4 DSpark flavour); the getattr keeps
-            # this rejection forward-compatible and is a no-op until then.
-            if getattr(spec_config, "draft_is_embedded_in_target", False):
+            # A standalone drafter always carries its own checkpoint, so
+            # its absence identifies the embedded flavour.
+            if spec_config.speculative_model is None:
                 raise ValueError(
                     "Kimi K3 helix supports only the standalone DSpark "
-                    "drafter; the embedded (in-target) flavour shares the "
-                    "target KV bookkeeping in ways the helix ledger does "
-                    "not model."
+                    "drafter (speculative_model must point at a drafter "
+                    "checkpoint); the embedded (in-target) flavour shares "
+                    "the target KV bookkeeping in ways the helix ledger "
+                    "does not model."
                 )
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index c4259b120930..f09c42c1b73d 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -769,9 +769,8 @@ def _attn_forward_gen(
                 seq_start=attn_metadata.num_contexts,
                 num_seqs=attn_metadata.num_generations,
             )
-            helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None)
-            if helix_kv_bounds is not None and getattr(
-                    attn_metadata, "_helix_spec_tokens_valid", False):
+            if (isinstance(attn_metadata, TrtllmAttentionMetadata)
+                    and attn_metadata._helix_spec_tokens_valid):
                 # Speculative verify groups: KV ownership is per-TOKEN. A rank
                 # owning only the tail page of a group has zero visible KV for
                 # the group's leading tokens while its per-sequence kv_len is
@@ -781,7 +780,8 @@ def _attn_forward_gen(
                 # pool values; the combine multiplies by corr = 0 and
                 # 0 * NaN would poison the token on every CP rank — sanitize
                 # by the per-token bound instead.
-                zero_kv_mask = helix_kv_bounds[:partial_o.shape[0]] == 0
+                zero_kv_mask = (
+                    attn_metadata.helix_kv_bounds[:partial_o.shape[0]] == 0)
             return _helix_post_process(
                 partial_o,
                 softmax_stats,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index f7c54e3b4773..ed5a1b672a2f 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3955,8 +3955,9 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
             # first apply the same accepted-count correction position_ids
             # got above; without overlap the host values are already exact.
             md = inputs.get('attn_metadata')
-            if (md is not None and md.kv_cache_manager is not None
-                    and getattr(md, '_helix_spec_tokens_valid', False)):
+            if (isinstance(md, TrtllmAttentionMetadata)
+                    and md.kv_cache_manager is not None
+                    and md._helix_spec_tokens_valid):
                 helix_gen_tokens = (inputs['input_ids'].shape[0] -
                                     md.num_ctx_tokens)
                 if not self._disable_overlap_scheduler:
@@ -4030,14 +4031,13 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]):
                         restore=True,
                     )
 
-                if (self.mapping.has_cp_helix() and getattr(
-                        inputs['attn_metadata'], '_helix_spec_tokens_valid',
-                        False)):
-                    # Mirror of the helix position correction in
-                    # _preprocess_inputs (capture symmetry, like position_ids
-                    # above). The recompute's OVERWRITES (slots/bounds/
-                    # kv_lens) need no reversal: every consumer buffer is
-                    # rewritten from host state at the next step's prepare.
+                if (self.mapping.has_cp_helix()
+                        and isinstance(inputs['attn_metadata'],
+                                       TrtllmAttentionMetadata)
+                        and inputs['attn_metadata']._helix_spec_tokens_valid):
+                    # Mirrors the correction in _preprocess_inputs. Only
+                    # the offsets need reversing: the recompute's other
+                    # outputs are rewritten from host state next prepare.
                     inputs['attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= (
                         self.previous_pos_id_offsets_cuda[:previous_batch_tokens])
 

From 7afd98ce86e8b1cfbf67c04d329dc592f943f5d9 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Wed, 26 Aug 2026 01:32:30 -0700
Subject: [PATCH 07/15] [None][chore] helix x DSpark: tighten the added
 comments

Trim the comments introduced by this series to the constraints the code
cannot state itself: drop restatements of adjacent code, of the error
messages right below them and of the field declarations they repeat, drop
one pure navigation note, and correct one claim (allocations during CUDA
graph capture come from the graph pool; the scalar overload saves an
allocation on the eager path, it is not a capture requirement). No
functional change.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 .../attention_backend/fmha/cute_dsl_mla.py    |  5 +-
 .../_torch/attention_backend/fmha/fallback.py | 13 ++--
 .../_torch/attention_backend/trtllm.py        | 49 ++++++---------
 .../_torch/custom_ops/cute_dsl_custom_ops.py  |  8 +--
 .../attention/mla/mla_decode_fp16.py          | 22 +++----
 tensorrt_llm/_torch/modules/mla.py            | 16 +++--
 tensorrt_llm/_torch/pyexecutor/_util.py       | 11 ++--
 .../_torch/pyexecutor/model_engine.py         | 60 +++++++------------
 8 files changed, 72 insertions(+), 112 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
index ced37b77d8e3..8c5f67b68755 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
@@ -520,9 +520,8 @@ def _run_mla_decode(
             # Max batch size for the AutoTuner to profile.
             int(meta.max_num_requests),
             params.fwd.softmax_stats_tensor,
-            # Helix speculative verify groups: per-token rank-local bounds
-            # (filled by recompute_helix_spec_buffers). None on the
-            # single-token helix path and outside helix.
+            # Per-token rank-local bounds, filled by
+            # recompute_helix_spec_buffers. None everywhere else.
             (meta.helix_kv_bounds[:num_tokens] if
              (meta.helix_position_offsets is not None
               and meta._helix_spec_tokens_valid
diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
index 2d251e396c75..566638672ec4 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
@@ -77,13 +77,12 @@ def is_supported(
         phase: Optional[FmhaPhase] = None,
     ) -> bool:
         del k, v, phase
-        # Helix speculative verify groups carry per-token KV ownership (a
-        # group may straddle a page boundary onto two CP ranks) that the
-        # fused thop attention path cannot express: its spec-dec mask and the
-        # per-sequence helix_is_inactive_rank gate both assume the group's
-        # new KV entries are the trailing slots of one rank's kv_len. Running
-        # it would be silently wrong, so reject — with no remaining FMHA
-        # library, dispatch raises loudly instead.
+        # A verify group may straddle a page boundary onto two CP ranks, so
+        # its KV ownership is per-token. The fused thop path cannot express
+        # that: its spec-dec mask and the per-sequence helix_is_inactive_rank
+        # gate both assume the new KV entries are the trailing slots of one
+        # rank's kv_len. Reject rather than run it silently wrong; being last
+        # in the library list, this makes dispatch raise.
         if (metadata.helix_position_offsets is not None
                 and metadata._helix_spec_tokens_valid
                 and metadata.num_generations > 0
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index b65b7cbbb1fe..19192b8e1acd 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -160,10 +160,10 @@ def effective_beam_width(self) -> int:
     helix_is_inactive_rank: Optional[torch.Tensor] = None
     helix_is_inactive_rank_cpu: Optional[torch.Tensor] = None
 
-    # Per-token helix state for speculative verify groups (a 1 + draft_len
-    # group may straddle a ledger-page boundary onto two CP ranks, so the
-    # per-sequence boolean above is insufficient there). See
-    # recompute_helix_spec_buffers for the derivation.
+    # Per-token helix state for speculative verify groups: a 1 + draft_len
+    # group may straddle a ledger page onto two CP ranks, which the
+    # per-sequence boolean above cannot express. See
+    # recompute_helix_spec_buffers.
     helix_local_slots: Optional[torch.Tensor] = None
     helix_kv_bounds: Optional[torch.Tensor] = None
     helix_owned_new_tokens_cpu: Optional[torch.Tensor] = None
@@ -553,17 +553,10 @@ def _post_init_with_buffers(self, buffers) -> None:
                 device='cpu',
                 pin_memory=prefer_pinned(),
             )
-            # Per-token buffers for speculative verify groups under helix.
-            # A group of 1 + draft_len tokens can straddle a ledger-page
-            # boundary, splitting ownership between two CP ranks, so the
-            # per-sequence flag above is not expressive enough:
-            #   helix_local_slots[t]: rank-local KV write slot of gen token t
-            #     on this rank, or -1 when another rank owns its position
-            #     (consumed by the mla_rope_generation append kernel).
-            #   helix_kv_bounds[t]: number of rank-local KV entries token t
-            #     may attend to, i.e. local_len(pos_t + 1) (consumed by the
-            #     CuTe DSL MLA decode mask and the helix stats identity).
-            # Filled by recompute_helix_spec_buffers() on the spec path only.
+            # helix_local_slots[t]: rank-local KV write slot of gen token t,
+            #   or -1 when another rank owns its position.
+            # helix_kv_bounds[t]: rank-local KV entries token t may attend
+            #   to, i.e. local_len(pos_t + 1).
             self.helix_local_slots = self.get_empty(
                 buffers,
                 (self.max_num_tokens, ),
@@ -578,9 +571,8 @@ def _post_init_with_buffers(self, buffers) -> None:
                 dtype=torch.int,
                 capture_graph=capture_graph,
             )
-            # Host-side per-sequence count of this step's new tokens owned by
-            # this rank (spec path; single-token path derives it from the
-            # boolean flag). Consumed by prepare()'s helix kv_lens branch.
+            # Per-sequence count of this step's new tokens owned by this
+            # rank. Provisional: the device recompute supersedes it.
             self.helix_owned_new_tokens_cpu = torch.zeros(
                 (self.max_num_sequences, ),
                 device='cpu',
@@ -684,16 +676,16 @@ def recompute_helix_spec_buffers(self, num_ctx_tokens: int,
         owner = torch.div(pos, phys, rounding_mode='floor') % cp_size
         active = owner == cp_rank
         local_before = self.helix_local_len_vec(pos)
-        # Scalar overload: no per-step allocation (these ops are captured
-        # into the CUDA graph; keep them allocation-free).
+        # Scalar overload: avoids a per-step tensor allocation on the
+        # eager path.
         self.helix_local_slots[num_ctx_tokens:num_ctx_tokens +
                                num_gen_tokens].copy_(
                                    torch.where(active, local_before, -1))
         self.helix_kv_bounds[num_ctx_tokens:num_ctx_tokens +
                              num_gen_tokens].copy_(
                                  self.helix_local_len_vec(pos + 1))
-        # Per-sequence rank-local kv length = bound of the sequence's last
-        # token (attention over committed + owned in-flight tokens).
+        # Positions rise within a group and local_len is monotonic, so the
+        # last token's bound is the group's maximum.
         assert num_gen_tokens % tokens_per_gen_seq == 0, (
             f"helix spec expects uniform verify groups: {num_gen_tokens} gen "
             f"tokens not divisible by group size {tokens_per_gen_seq}")
@@ -833,11 +825,9 @@ def prepare(self) -> None:
             # If helix is inactive, attend to the previously cached tokens only.
             assert cached_token_lens is not None, "cached_token_lens should be set for helix"
             if self._helix_spec_tokens_valid:
-                # Speculative verify groups: a group may straddle a page
-                # boundary, so ownership of this step's new tokens is a
-                # per-sequence COUNT, not a boolean. Provisional host values;
-                # recompute_helix_spec_buffers overrides the device copy
-                # after the overlap correction.
+                # A straddling group splits this step's new tokens between
+                # two ranks, so ownership is a COUNT, not a boolean.
+                # Provisional; the device recompute overrides it.
                 kv_lens = cached_token_lens + \
                     self.helix_owned_new_tokens_cpu[:self.num_seqs]
             else:
@@ -2564,9 +2554,8 @@ def mla_rope_generation(
             metadata.helix_position_offsets, metadata.helix_is_inactive_rank
         ]
         if metadata._helix_spec_tokens_valid:
-            # Speculative verify groups: per-token KV write slots (-1 = this
-            # rank does not own the token's position). The append kernel then
-            # gates and addresses per token instead of per sequence.
+            # With slots present the append kernel gates and addresses per
+            # token instead of per sequence.
             helix_tensor_params.append(metadata.helix_local_slots)
 
         torch.ops.trtllm.mla_rope_generation(
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 6e5843a5e76e..7c4e1896dabf 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -9771,10 +9771,10 @@ def forward(
                                               split_kv_size]
 
             if kv_bounds is not None and AutoTuner.get().is_tuning_mode:
-                # Profiling rebuilds bucketed cache_seqs but carries
-                # kv_bounds through unchanged (input 9 has no dynamic-dim
-                # spec); re-derive a size-consistent dummy — bound values
-                # only affect masking depth, not the tactic space.
+                # Profiling rebuilds cache_seqs at bucketed sizes but input 9
+                # has no dynamic-dim spec, so kv_bounds arrives at the old
+                # size. Bound values only affect masking depth, not the
+                # tactic space, so any size-consistent dummy will do.
                 if kv_bounds.numel() != batch_size * seq_len_q:
                     kv_bounds = cache_seqs.repeat_interleave(
                         seq_len_q).contiguous()
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
index 7973c157284e..69f216673e57 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
@@ -2809,13 +2809,11 @@ def softmax(
                     else:
                         q_tok = common_params.blk_coord[1]
                     if cutlass.const_expr(common_params.kv_bounds is not None):
-                        # Helix verify groups: per-token rank-local bound
-                        # (committed prefix + owned group tokens <= q_tok);
-                        # subsumes the causal offset and non-owner ranks.
-                        # Clamp for fold_sq M-tile padding rows (row beyond
-                        # num_heads * fold_sq_ratio derives q_tok >= S_q);
-                        # their results are discarded, but the gmem read must
-                        # stay in bounds.
+                        # Per-token rank-local bound; subsumes the causal
+                        # offset and non-owner ranks.
+                        # fold_sq M-tile padding rows derive q_tok >= S_q;
+                        # their results are discarded but the read must stay
+                        # in bounds.
                         q_tok_c = (q_tok if cute.elem_less(
                             q_tok, self.seq_len_q) else self.seq_len_q - 1)
                         k_bound = common_params.kv_bounds[
@@ -2864,12 +2862,10 @@ def softmax(
                     else:
                         q_tok = common_params.blk_coord[1]
                     if cutlass.const_expr(common_params.kv_bounds is not None):
-                        # Helix verify groups: per-token rank-local bound
-                        # (see the sm_100 branch above).
-                        # Clamp for fold_sq M-tile padding rows (row beyond
-                        # num_heads * fold_sq_ratio derives q_tok >= S_q);
-                        # their results are discarded, but the gmem read must
-                        # stay in bounds.
+                        # Per-token rank-local bound (see the sm_100 branch).
+                        # fold_sq M-tile padding rows derive q_tok >= S_q;
+                        # their results are discarded but the read must stay
+                        # in bounds.
                         q_tok_c = (q_tok if cute.elem_less(
                             q_tok, self.seq_len_q) else self.seq_len_q - 1)
                         k_bound = common_params.kv_bounds[
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index f09c42c1b73d..b99352411c15 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -771,15 +771,13 @@ def _attn_forward_gen(
             )
             if (isinstance(attn_metadata, TrtllmAttentionMetadata)
                     and attn_metadata._helix_spec_tokens_valid):
-                # Speculative verify groups: KV ownership is per-TOKEN. A rank
-                # owning only the tail page of a group has zero visible KV for
-                # the group's leading tokens while its per-sequence kv_len is
-                # nonzero, so the per-sequence mask above misses those rows.
-                # Their decode rows are fully masked with a finite sentinel,
-                # making partial_o an average over (possibly uninitialized)
-                # pool values; the combine multiplies by corr = 0 and
-                # 0 * NaN would poison the token on every CP rank — sanitize
-                # by the per-token bound instead.
+                # A rank owning only a group's tail page has zero visible KV
+                # for its leading tokens while the per-sequence kv_len is
+                # nonzero, so the mask above misses those rows. Their scores
+                # are fully masked with a finite sentinel, leaving partial_o
+                # an average over uninitialized pool values; the combine
+                # scales it by corr = 0, and 0 * NaN would poison the token
+                # on every CP rank.
                 zero_kv_mask = (
                     attn_metadata.helix_kv_bounds[:partial_o.shape[0]] == 0)
             return _helix_post_process(
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 9db0e6774389..806b297b9b89 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1562,12 +1562,11 @@ def _create_one_model_draft_kv_cache_manager(
         # the sparse_attention_config. Get it from effective_draft_config which
         # falls back to the target model's config for MTP mode.
         sparse_attn_config = effective_draft_config.sparse_attention_config
-        # The standalone drafter is a plain dense model built against the
-        # repurposed mapping under helix (CP ranks become TP ranks; each rank
-        # keeps its full drafter KV). Its paged KV manager must therefore use
-        # the repurposed, CP-free mapping — helix round-robin ledger semantics
-        # apply only to the TARGET KV (and KVCacheManagerV2 rejects
-        # is_draft x helix outright).
+        # Under helix the standalone drafter is built against the repurposed
+        # mapping (CP ranks become TP ranks) and every rank keeps its full
+        # drafter KV, so its paged manager needs the CP-free mapping: the
+        # round-robin ledger applies to the TARGET KV alone, and
+        # KVCacheManagerV2 rejects is_draft x helix outright.
         draft_mapping = self._mapping
         if draft_mapping.has_cp_helix():
             draft_mapping = draft_mapping.repurpose_helix_cp_to_tp()
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index ed5a1b672a2f..6198a9e8ce47 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3947,13 +3947,10 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                     )
 
         if self.enable_spec_decode and self.mapping.has_cp_helix():
-            # Helix verify groups: the per-token device buffers (write slots,
-            # attention bounds, rank-local kv lens) must be derived on EVERY
-            # spec step, overlap or not — the append/mask kernels consume
-            # them whenever _helix_spec_tokens_valid is armed. Under overlap
-            # the host packed provisional positions from a stale base, so
-            # first apply the same accepted-count correction position_ids
-            # got above; without overlap the host values are already exact.
+            # The per-token buffers must be derived on EVERY spec step, not
+            # just under overlap: the append/mask kernels read them whenever
+            # the flag is armed. Under overlap the host packed positions from
+            # a stale base, so correct them first (as position_ids was above).
             md = inputs.get('attn_metadata')
             if (isinstance(md, TrtllmAttentionMetadata)
                     and md.kv_cache_manager is not None
@@ -3961,9 +3958,8 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]):
                 helix_gen_tokens = (inputs['input_ids'].shape[0] -
                                     md.num_ctx_tokens)
                 if not self._disable_overlap_scheduler:
-                    # The kv_lens override in the recompute supersedes the
-                    # generic previous_kv_lens_offsets adjustment above,
-                    # which is not ownership-aware.
+                    # The recompute's kv_lens override supersedes the generic
+                    # adjustment above, which is not ownership-aware.
                     md.helix_position_offsets[:helix_gen_tokens] += (
                         self.previous_pos_id_offsets_cuda[:helix_gen_tokens])
                 md.recompute_helix_spec_buffers(
@@ -5631,11 +5627,9 @@ def append_cross_attention_state(request: LlmRequest,
                 generation_requests.append(request)
         extend_requests += extend_dummy_requests
 
-        # Helix bookkeeping is needed by BOTH the extend (speculative verify
-        # group) and the plain generation packing loops below, so initialize
-        # it ahead of them. Positions are global; KV ownership follows the
-        # round-robin ledger (page b -> rank b % cp), mirrored host-side here
-        # (KVCacheManagerV2._helix_local_len) for provisional packing values.
+        # Shared by the extend and plain generation loops below. Host mirror
+        # of KVCacheManagerV2._helix_local_len (page b -> rank b % cp), used
+        # for the provisional packing values.
         helix_is_inactive_rank, helix_position_offsets = [], []
         helix_owned_new_tokens = []
         _has_cp_helix = self.mapping.has_cp_helix()
@@ -5722,22 +5716,16 @@ def _helix_local_len_host(global_len: int) -> int:
                     past_seen_token_num - request.py_num_compressed_tokens)
                 request.cached_tokens = past_seen_token_num
                 if _has_cp_helix:
-                    # Verify group [base, base+group) in GLOBAL positions.
-                    # On a helix gen worker the request's token list is the
-                    # rank-LOCAL round-robin subset, so max_beam_num_tokens
-                    # (= local_prompt + generated) must NOT be used as a
-                    # global base; reconstruct it from the global prompt
-                    # length plus the (rank-invariant) generated count. This
-                    # branch has no in-flight predecessor, so every value is
-                    # exact (no device correction needed).
+                    # A helix gen worker's token list is the rank-LOCAL
+                    # round-robin subset, so max_beam_num_tokens is not a
+                    # global base; rebuild it from the global prompt length
+                    # plus the rank-invariant generated count.
                     group = 1 + num_draft_tokens
                     generated_len = (request.max_beam_num_tokens -
                                      request.py_prompt_len)
                     base = request.total_input_len_cp + generated_len - 1
                     helix_position_offsets.extend(range(base, base + group))
-                    # position_ids above were packed from the local base;
-                    # helix uses global position ids (same convention as the
-                    # non-spec helix generation loop below).
+                    # Repack: the loop above used the local base.
                     position_ids[-group:] = range(base, base + group)
                     helix_is_inactive_rank.append(False)
                     local_cached = _helix_local_len_host(base)
@@ -5778,13 +5766,10 @@ def _helix_local_len_host(global_len: int) -> int:
                 request.cached_tokens = (past_seen_token_num +
                                          runtime_tokens_per_gen_step)
                 if _has_cp_helix:
-                    # In-flight predecessor: mirror the non-helix convention
-                    # above — positions are packed from the stale base (the
-                    # overlap device correction adds the accepted count) and
-                    # KV numbers assume full acceptance (the device recompute
-                    # in recompute_helix_spec_buffers overrides them). The
-                    # base is reconstructed GLOBALLY (see the no-previous
-                    # branch: the token list is rank-local under helix).
+                    # In-flight predecessor: provisional values, per the
+                    # non-helix convention above. Positions use the stale base
+                    # and the KV numbers assume full acceptance; the device
+                    # recompute overrides both. Base is global (see above).
                     group = runtime_tokens_per_gen_step
                     generated_len = (request.max_beam_num_tokens -
                                      request.py_prompt_len)
@@ -5864,8 +5849,6 @@ def _helix_local_len_host(global_len: int) -> int:
             # update batch index
             request.py_batch_idx = request.py_seq_slot
 
-        # (helix lists and _has_cp_helix are initialized ahead of the extend
-        # loop above, which also appends to them for verify groups.)
 
         _n_gen = len(generation_requests)
         # One-shot batch-level flag — True iff any generation request actually
@@ -5996,8 +5979,7 @@ def _helix_local_len_host(global_len: int) -> int:
                         helix_is_inactive_rank.append(
                             request.py_helix_is_inactive_rank)
                         helix_position_offsets.append(position_id)
-                        # Keep the per-seq owned-count list aligned when the
-                        # spec path is active in the same batch.
+                        # Keep the per-seq list aligned in mixed batches.
                         helix_owned_new_tokens.append(
                             0 if request.py_helix_is_inactive_rank else 1)
 
@@ -6423,9 +6405,7 @@ def previous_seq_slots_device():
             attn_metadata.update_helix_param(
                 helix_position_offsets=helix_position_offsets,
                 helix_is_inactive_rank=helix_is_inactive_rank,
-                # Per-seq owned counts drive the kv_lens math only on the
-                # speculative path (verify groups); None keeps the
-                # single-token boolean convention.
+                # None keeps the single-token boolean convention.
                 helix_owned_new_tokens=(helix_owned_new_tokens
                                         if self.enable_spec_decode else None),
             )

From 7af05127383581f7a20e2c979f434e0f353ad179 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Wed, 26 Aug 2026 02:58:35 -0700
Subject: [PATCH 08/15] [None][chore] helix x DSpark: factor the shared
 extend-loop packing

The two helix blocks in the extend loop shared six of nine lines,
including the global-base reconstruction that is the easiest part to get
wrong (max_beam_num_tokens is rank-local on a helix gen worker, so it
cannot serve as a global base). Move that shared part into
_helix_pack_extend so the reconstruction exists once; each branch keeps
only its three real differences (group size, whether the cached length is
taken before or after the group, and whether the owned-token count is
exact or a placeholder). No functional change.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 40 +++++++++----------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 6198a9e8ce47..3104d9bd63a8 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -5643,6 +5643,20 @@ def _helix_local_len_host(global_len: int) -> int:
                 return full * _helix_phys + min(
                     max(rem - _helix_rank_off, 0), _helix_phys)
 
+            def _helix_pack_extend(request, group: int) -> int:
+                # A helix gen worker's token list is the rank-LOCAL
+                # round-robin subset, so max_beam_num_tokens is not a global
+                # base; rebuild it from the global prompt length plus the
+                # rank-invariant generated count. Also repacks position_ids,
+                # which the caller filled from the local base.
+                generated_len = (request.max_beam_num_tokens -
+                                 request.py_prompt_len)
+                base = request.total_input_len_cp + generated_len - 1
+                helix_position_offsets.extend(range(base, base + group))
+                position_ids[-group:] = range(base, base + group)
+                helix_is_inactive_rank.append(False)
+                return base
+
         spec_config = self.spec_config if self.enable_spec_decode else None
         if not self._disable_overlap_scheduler and spec_config is not None:
             assert spec_config.spec_dec_mode.support_overlap_scheduler(
@@ -5716,18 +5730,9 @@ def _helix_local_len_host(global_len: int) -> int:
                     past_seen_token_num - request.py_num_compressed_tokens)
                 request.cached_tokens = past_seen_token_num
                 if _has_cp_helix:
-                    # A helix gen worker's token list is the rank-LOCAL
-                    # round-robin subset, so max_beam_num_tokens is not a
-                    # global base; rebuild it from the global prompt length
-                    # plus the rank-invariant generated count.
+                    # No in-flight predecessor, so every value is exact.
                     group = 1 + num_draft_tokens
-                    generated_len = (request.max_beam_num_tokens -
-                                     request.py_prompt_len)
-                    base = request.total_input_len_cp + generated_len - 1
-                    helix_position_offsets.extend(range(base, base + group))
-                    # Repack: the loop above used the local base.
-                    position_ids[-group:] = range(base, base + group)
-                    helix_is_inactive_rank.append(False)
+                    base = _helix_pack_extend(request, group)
                     local_cached = _helix_local_len_host(base)
                     helix_owned_new_tokens.append(
                         _helix_local_len_host(base + group) - local_cached)
@@ -5767,16 +5772,11 @@ def _helix_local_len_host(global_len: int) -> int:
                                          runtime_tokens_per_gen_step)
                 if _has_cp_helix:
                     # In-flight predecessor: provisional values, per the
-                    # non-helix convention above. Positions use the stale base
-                    # and the KV numbers assume full acceptance; the device
-                    # recompute overrides both. Base is global (see above).
+                    # non-helix convention above. The base is stale and the KV
+                    # numbers assume full acceptance; the device recompute
+                    # overrides both.
                     group = runtime_tokens_per_gen_step
-                    generated_len = (request.max_beam_num_tokens -
-                                     request.py_prompt_len)
-                    base = request.total_input_len_cp + generated_len - 1
-                    helix_position_offsets.extend(range(base, base + group))
-                    position_ids[-group:] = range(base, base + group)
-                    helix_is_inactive_rank.append(False)
+                    base = _helix_pack_extend(request, group)
                     local_full = _helix_local_len_host(base + group)
                     helix_owned_new_tokens.append(0)
                     num_cached_tokens_per_seq[-1] = (

From b17afdd6e242622a8bb4176ca67ec8754c69e16d Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Thu, 27 Aug 2026 21:34:18 -0700
Subject: [PATCH 09/15] [None][fix] Drop the speculative_model-presence guard

Both DSpark flavours carry a speculative_model path (the embedded one
resolves it to the target checkpoint and probes the weight index), so a
missing path does not identify the embedded flavour. The embedded case
is already rejected downstream by draft_is_embedded_in_target.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_linear.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index e450216c187b..b1ec4e2f2751 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2065,16 +2065,6 @@ def _setup_helix_mappings(
                     "Kimi K3 helix supports speculative decoding only with "
                     f"DSpark (standalone drafter); got {decoding_type!r}."
                 )
-            # A standalone drafter always carries its own checkpoint, so
-            # its absence identifies the embedded flavour.
-            if spec_config.speculative_model is None:
-                raise ValueError(
-                    "Kimi K3 helix supports only the standalone DSpark "
-                    "drafter (speculative_model must point at a drafter "
-                    "checkpoint); the embedded (in-target) flavour shares "
-                    "the target KV bookkeeping in ways the helix ledger "
-                    "does not model."
-                )
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp
         if cfg.num_attention_heads % repurposed_tp != 0:

From 65bc5e8c65dc58d913b86d181c1391b08531c1d8 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Thu, 27 Aug 2026 22:32:28 -0700
Subject: [PATCH 10/15] [None][fix] Helix: account KDA state and draft KV with
 the repurposed mapping

The KV budget split for a separate DSpark draft cache computed both cost
terms with the raw helix mapping (tp_size=1):

- MambaKVCacheParams.get_states_bytes_per_layer counted the KDA conv/ssm
  state unsharded, withholding 29.2 GiB/rank from the budget on a helix16
  gen worker whose real pool is 1/16-sharded and costs 1.7 GiB. The
  estimator now shares the allocator's effective-TP rule
  (mamba_effective_tp_size, moved to config_utils so the two can never
  diverge again).
- The drafter's per-token cost used the unrepurposed mapping; the value
  (20480 B/token) was right only because this drafter's kv-head count
  happens to equal cp_size. It is now computed per GLOBAL token with the
  repurposed mapping runtime construction uses, then scaled by cp_size
  into the target's rank-local-token unit.

Non-helix paths are byte-identical (scale 1, same mapping). On helix16
mb64 the target KV capacity grows ~28% and the gen-only fill that
previously died with 'Insufficient KV cache' has room to complete.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 39 +++++++++++++++----
 .../_torch/pyexecutor/config_utils.py         | 18 ++++++++-
 .../_torch/pyexecutor/mamba_cache_manager.py  | 17 ++------
 3 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 806b297b9b89..de58c07b9daa 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -698,13 +698,14 @@ def _per_manager_cache_cost(self,
                                 manager_cls,
                                 model_config,
                                 kv_cache_config: Optional[KvCacheConfig] = None,
+                                mapping=None,
                                 **extra_kwargs) -> CacheCost:
         kv_cache_config = (kv_cache_config if kv_cache_config is not None else
                            self._kv_cache_config)
         return CacheCost.from_raw(
             manager_cls.get_cache_size_per_token(
                 model_config,
-                self._mapping,
+                mapping if mapping is not None else self._mapping,
                 tokens_per_block=self._tokens_per_block,
                 max_seq_len=self._max_seq_len,
                 max_batch_size=self._max_batch_size,
@@ -734,6 +735,17 @@ def _get_kv_size_per_token(self,
         model_config = self._model_engine.model.model_config
         use_separate_draft_kv_cache = (
             self._should_create_separate_draft_kv_cache())
+        # DRAFT cost unit convention under helix: compute per GLOBAL token with
+        # the same repurposed mapping runtime construction uses (the drafter is
+        # dense, not helix-sharded), then multiply by cp_size to express it per
+        # rank-LOCAL target token so the split's slopes share a unit (the
+        # target stores only every cp_size-th page per rank). Intercepts are
+        # per-request rank-local bytes and stay unscaled.
+        draft_mapping = self._mapping
+        helix_cp_scale = 1
+        if self._mapping.has_cp_helix():
+            draft_mapping = self._mapping.repurpose_helix_cp_to_tp()
+            helix_cp_scale = self._mapping.cp_size
         total = self._per_manager_cache_cost(
             self._kv_cache_manager_cls,
             model_config,
@@ -745,9 +757,13 @@ def _get_kv_size_per_token(self,
             draft_model_config = self._draft_model_engine.model.model_config
             draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls(
                 self._draft_model_engine, kv_cache_config)
-            total += self._per_manager_cache_cost(draft_kv_cache_manager_cls,
-                                                  draft_model_config,
-                                                  kv_cache_config)
+            draft_cost = self._per_manager_cache_cost(
+                draft_kv_cache_manager_cls,
+                draft_model_config,
+                kv_cache_config,
+                mapping=draft_mapping)
+            total += CacheCost(slope=draft_cost.slope * helix_cp_scale,
+                               intercept=draft_cost.intercept)
         elif use_separate_draft_kv_cache:
             # One-model draft with separate KV cache layout.
             # Pass num_layers explicitly since the HF config may report a
@@ -766,17 +782,24 @@ def _get_kv_size_per_token(self,
                     effective_draft_config,
                     draft_kv_cache_config,
                     is_disagg=self._is_disagg)
-                total += self._per_manager_cache_cost(
-                    draft_kv_cache_manager_cls, effective_draft_config,
-                    draft_kv_cache_config)
+                draft_cost = self._per_manager_cache_cost(
+                    draft_kv_cache_manager_cls,
+                    effective_draft_config,
+                    draft_kv_cache_config,
+                    mapping=draft_mapping)
+                total += CacheCost(slope=draft_cost.slope * helix_cp_scale,
+                                   intercept=draft_cost.intercept)
             elif self._mapping.is_last_pp_rank():
                 # EAGLE3/MTP: draft layers only on last PP rank
-                total += self._per_manager_cache_cost(
+                draft_cost = self._per_manager_cache_cost(
                     self._kv_cache_manager_cls,
                     effective_draft_config,
                     draft_kv_cache_config,
+                    mapping=draft_mapping,
                     num_layers=self._get_num_draft_layers(),
                     is_draft=True)
+                total += CacheCost(slope=draft_cost.slope * helix_cp_scale,
+                                   intercept=draft_cost.intercept)
         return total
 
     def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py
index ed02327e3d66..01aff79de2b1 100644
--- a/tensorrt_llm/_torch/pyexecutor/config_utils.py
+++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py
@@ -305,6 +305,22 @@ def get_qwen3_hybrid_num_attention_layers(config):
     return sum(layer_mask)
 
 
+def mamba_effective_tp_size(mapping) -> int:
+    """TP degree for sizing per-rank mamba/KDA state (budgeting AND allocation).
+
+    Attention-DP replicates the state and takes precedence; helix repurposes
+    CP ranks as plain TP for recurrent-state layers. Must match the runtime
+    pool construction (mamba_cache_manager) or the budget split withholds
+    unsharded-state bytes the allocator never uses (observed: 27.2 GiB/rank
+    mis-withheld on a helix16 gen worker whose real pool is 1/16-sharded).
+    """
+    if mapping.enable_attention_dp:
+        return 1
+    if mapping.has_cp_helix():
+        return mapping.tp_size * mapping.cp_size
+    return mapping.tp_size
+
+
 @dataclasses.dataclass
 class MambaKVCacheParams:
     """Normalized mamba-related inputs for kv_cache_manager_cls.
@@ -364,7 +380,7 @@ def get_layer_masks(
 
     def get_states_bytes_per_layer(self, mapping) -> int:
         """Return the total bytes of Mamba state per layer, used for budgeting."""
-        tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1
+        tp_size = mamba_effective_tp_size(mapping)
         d_inner = self.head_dim * self.num_heads
         conv_dim = (d_inner + 2 * self.n_groups * self.state_size) // tp_size
         nheads = self.num_heads // tp_size
diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
index 5e188a1061da..f2d575d823b7 100644
--- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
@@ -44,6 +44,10 @@
     LinearAttentionMetadata, LinearCacheType)
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig
 from tensorrt_llm.logger import logger
+
+# Shared with the KV budget estimator so allocator and budgeting can never
+# diverge on the sharding rule (config_utils is import-cycle-free).
+from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX,
                                                       BatchDesc, BufferConfig,
@@ -127,19 +131,6 @@ class MambaRole:
     CONV_STATE = DataRole("conv_state")
 
 
-def _mamba_effective_tp_size(mapping: Mapping) -> int:
-    """TP degree for sizing per-rank mamba/KDA state pools.
-
-    Attention-DP replicates the state and takes precedence; helix
-    repurposes CP ranks as plain TP for recurrent-state layers.
-    """
-    if mapping.enable_attention_dp:
-        return 1
-    if mapping.has_cp_helix():
-        return mapping.tp_size * mapping.cp_size
-    return mapping.tp_size
-
-
 def get_tensor_size_bytes(tensor):
     """Calculate tensor size in bytes."""
     if isinstance(tensor, torch.Tensor):

From 980c1db9e6a7b22630340bd255ffac6b6f64322a Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Fri, 28 Aug 2026 00:55:13 -0700
Subject: [PATCH 11/15] [None][chore] Use the canonical spec-mode check in the
 helix allowlist

Every concrete decoding config carries decoding_type as a required
Literal field, so the defensive getattr was dead weight; match the
file's existing spec_dec_mode.is_dspark() idiom instead.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_linear.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index b1ec4e2f2751..2fe19fc0cfe3 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2054,16 +2054,14 @@ def _setup_helix_mappings(
                 "per-request locality of KDA recurrent state."
             )
         if spec_config is not None:
-            # Helix speculative decoding allowlist: standalone DSpark linear
-            # chains verified on the V2 superblock ledger (per-token ownership
-            # bookkeeping + CuTe DSL per-token bounds). Everything else stays
-            # loudly rejected — the #16003 review showed unsupported spec
-            # modes silently running wrong under helix.
-            decoding_type = getattr(spec_config, "decoding_type", None)
-            if decoding_type != "DSpark":
+            # Helix supports only the standalone DSpark drafter (verified on
+            # the V2 superblock ledger); reject everything else loudly rather
+            # than let an unsupported spec mode run silently wrong.
+            if not spec_config.spec_dec_mode.is_dspark():
                 raise ValueError(
                     "Kimi K3 helix supports speculative decoding only with "
-                    f"DSpark (standalone drafter); got {decoding_type!r}."
+                    f"DSpark (standalone drafter); got "
+                    f"{spec_config.decoding_type!r}."
                 )
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp

From 2c4ed6a5f2c3e72035bb2a732bd1dbe082ff24d6 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Sun, 30 Aug 2026 19:20:04 -0700
Subject: [PATCH 12/15] [None][fix] Re-thread kv_bounds through the post-#18131
 kernel entry points

The #18131 merge split the kernel entry into __call__ /
run_with_softmax_stats / _run; rebasing the kv_bounds port over that
restructure landed the signature in __call__ but the body references in
_run, so every fp16 CuTe DSL MLA decode traced NameError and the stats
entry points took one fewer argument than the runner passes. Add the
parameter to _run and run_with_softmax_stats (fp8: signature parity,
value deliberately dropped) and forward it from __call__.

Also make the reduction kernel's per-token CP-merge gating fold-aware:
its grid is the folded (H*F, S_q/F, B) geometry, so the true token for
the kv_bounds lookup is chunk * F + row // num_heads; the previous
formula read the wrong token whenever fold_sq_ratio > 1 (latent for
K3 helix, whose H=96 folds at ratio 1).

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 .../blackwell/attention/mla/mla_decode_fp16.py  | 17 ++++++++++++++++-
 .../blackwell/attention/mla/mla_decode_fp8.py   |  4 ++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
index 69f216673e57..f7ce10b639ce 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
@@ -329,6 +329,7 @@ def __call__(
             workspace,
             split_kv,
             cache_seqs,
+            kv_bounds,
             block_split_kvs,
             softmax_scale,
             output_scale,
@@ -349,6 +350,7 @@ def run_with_softmax_stats(
         workspace: cute.Tensor,
         split_kv: cutlass.Int32,
         cache_seqs: Optional[cute.Tensor],
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         output_scale: cutlass.Float32,
@@ -366,6 +368,7 @@ def run_with_softmax_stats(
             workspace,
             split_kv,
             cache_seqs,
+            kv_bounds,
             block_split_kvs,
             softmax_scale,
             output_scale,
@@ -386,6 +389,7 @@ def _run(
         workspace: cute.Tensor,
         split_kv: cutlass.Int32,
         cache_seqs: Optional[cute.Tensor],
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         output_scale: cutlass.Float32,
@@ -1549,8 +1553,19 @@ def reduction_kernel(
                     # per-token: a straddling group leaves this rank zero
                     # entries for its leading tokens while later ones have some.
                     if cutlass.const_expr(kv_bounds is not None):
+                        # blk_coord runs over the folded reduction grid
+                        # (H*F, S_q/F, B) while kv_bounds is indexed by the
+                        # true token. A folded chunk packs its rows as
+                        # tok_in_chunk * H + head, so the true token is
+                        # chunk * F + row // H (self.num_heads and
+                        # self.seq_len_q stay pre-fold).
+                        if cutlass.const_expr(self.fold_sq):
+                            q_tok = (blk_coord[1] * self.fold_sq_ratio +
+                                     blk_coord[0] // self.num_heads)
+                        else:
+                            q_tok = blk_coord[1]
                         has_local_kv = kv_bounds[blk_coord[2] * self.seq_len_q +
-                                                 blk_coord[1]] > 0
+                                                 q_tok] > 0
                     else:
                         has_local_kv = cache_seqs[blk_coord[2]] > 0
                     if has_local_kv:
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
index 4e3d14c17e37..c995610c5bd1 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
@@ -350,6 +350,10 @@ def run_with_softmax_stats(
         workspace: cute.Tensor,
         split_kv: cutlass.Int32,
         cache_seqs: Optional[cute.Tensor],
+        # Signature parity with the fp16 kernel; the fp8 kernel does not
+        # implement per-token bounds, so the value is accepted and dropped
+        # (the custom op rejects a non-None kv_bounds for fp8 upstream).
+        kv_bounds: Optional[cute.Tensor],
         block_split_kvs: Optional[cute.Tensor],
         softmax_scale: cutlass.Float32,
         output_scale: cutlass.Float32,

From 94128e55fa4a67c430155c93f60025302b0ab232 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Sun, 30 Aug 2026 19:20:19 -0700
Subject: [PATCH 13/15] [None][fix] Guard helix x speculation against V1
 managers and the acceptance gate

Two review findings on the same theme (config combinations that run
silently wrong instead of failing loudly):

- Require KVCacheManagerV2 when helix runs with speculation: per-token
  verify-group bookkeeping exists only on the V2 superblock ledger,
  while the default K3 resolution (use_kv_cache_manager_v2 unset) picks
  the V1-family Mixed manager whose helix accounting is one token per
  iteration.

- Reject acceptance_rate_window_size / acceptance_rate_threshold in the
  K3 helix allowlist: the SpeculationGate trip permanently disables
  speculation mid-flight, dropping in-flight helix requests into the
  plain generation loop whose position formula is stale once any draft
  token was accepted.

Also refresh the two comments that still claimed draft-token modes are
rejected under helix.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 .../_torch/models/modeling_kimi_linear.py         | 15 +++++++++++++++
 tensorrt_llm/_torch/pyexecutor/_util.py           | 14 ++++++++++++++
 .../_torch/pyexecutor/kv_cache_manager_v2.py      |  5 +++--
 tensorrt_llm/_torch/pyexecutor/model_engine.py    | 10 ++++++----
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 2fe19fc0cfe3..b52325409f7f 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2063,6 +2063,21 @@ def _setup_helix_mappings(
                     f"DSpark (standalone drafter); got "
                     f"{spec_config.decoding_type!r}."
                 )
+            # The SpeculationGate acceptance-rate trip permanently disables
+            # speculation mid-flight while enable_spec_decode stays True;
+            # in-flight helix requests then fall into the plain generation
+            # loop whose position math (total_input_len_cp +
+            # py_decoding_iter - 1) is stale once any draft token was
+            # accepted -> silently wrong RoPE positions and KV slots. Reject
+            # the trip wires until that loop is helix-group aware.
+            if (spec_config.acceptance_rate_window_size is not None
+                    or spec_config.acceptance_rate_threshold is not None):
+                raise ValueError(
+                    "Kimi K3 helix does not support the speculation "
+                    "acceptance-rate gate (acceptance_rate_window_size / "
+                    "acceptance_rate_threshold): dynamically disabling "
+                    "speculation mid-flight leaves helix requests on a "
+                    "single-token position formula.")
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp
         if cfg.num_attention_heads % repurposed_tp != 0:
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index de58c07b9daa..a57ebfe577ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -184,6 +184,20 @@ def get_kv_cache_manager_cls(
         # the shared hybrid transceiver validation below: the Python NIXL
         # transceiver selects the Mixed manager, whose KDA recurrent/conv
         # states transfer through the bounce buffer.
+        # Helix x speculation bookkeeping (per-token verify groups on the
+        # superblock ledger, py_helix_decode_group_index advancement) exists
+        # only in KVCacheManagerV2. The V1-family hybrid managers account
+        # helix decode one token per iteration and have no helix-x-spec
+        # path, so a default (V1) resolution would run silently wrong.
+        if (model_config.mapping is not None
+                and model_config.mapping.has_cp_helix()
+                and model_config.spec_config is not None and not use_v2):
+            raise ValueError(
+                "Kimi K3 helix with speculative decoding requires "
+                "kv_cache_config.use_kv_cache_manager_v2=True; the V1-family "
+                "hybrid managers do not implement per-token verify-group "
+                "bookkeeping.")
+
         if is_kimi_linear(config) and not use_v2 and not is_disagg:
             if kv_cache_config.enable_block_reuse:
                 logger.info(
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index e6b5bc932a25..4d49c0c17343 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -2384,8 +2384,9 @@ def _set_helix_rank_fields(self, req: LlmRequest) -> None:
         from ``py_decoding_iter``: the sampler advances that counter after
         scheduling under the overlap loop, so a schedule-time read is one
         step behind and would repeat the first decode position, overwriting
-        the first generated token's KV. Assumes one new token per step
-        (draft-token modes are rejected under helix).
+        the first generated token's KV. Multi-token verify groups advance
+        py_helix_decode_group_index per committed group, so this formula
+        stays exact under DSpark speculation.
         """
         step = req.py_helix_decode_group_index + 1
         pos = req.total_input_len_cp + step - 1
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 3104d9bd63a8..292c31844a3b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -5962,10 +5962,12 @@ def _helix_pack_extend(request, group: int) -> int:
                         # once (L, L, L+1, ...) and the new token's K is roped
                         # at the wrong position before being written to the KV
                         # cache, corrupting every later step.
-                        # TODO: revisit for helix x speculative decoding -
-                        # the base formula and this +1 both assume exactly
-                        # one new token per step (draft-token modes are
-                        # currently rejected under helix).
+                        # The base formula and this +1 assume exactly one
+                        # new token per step. Helix x DSpark never reaches
+                        # this plain-generation branch: the K3 allowlist
+                        # rejects the acceptance-rate gate, which is the only
+                        # dynamic path that could strip speculation from an
+                        # in-flight helix request.
                         position_id += 1
                     if request.py_helix_is_inactive_rank:
                         past_seen_token_num = request.seqlen_this_rank_cp

From af11d518b1036cb8744fea8be9557274b4edb663 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Sun, 30 Aug 2026 23:08:32 -0700
Subject: [PATCH 14/15] [None][chore] Apply pre-commit formatting

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 cpp/tensorrt_llm/kernels/mlaKernels.cu        |  8 ++++----
 .../attention_backend/fmha/cute_dsl_mla.py    | 19 ++++++++++++-------
 .../_torch/attention_backend/fmha/fallback.py | 10 ++++++----
 .../attention/mla/mla_decode_fp16.py          |  6 ++----
 .../_torch/models/modeling_kimi_linear.py     |  9 ++++++---
 tensorrt_llm/_torch/modules/mla.py            |  9 +++++----
 .../_torch/pyexecutor/mamba_cache_manager.py  |  8 ++++----
 .../_torch/pyexecutor/model_engine.py         | 17 +++++++++--------
 8 files changed, 48 insertions(+), 38 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu
index 7fd6c12957e2..e2c1e2cae312 100644
--- a/cpp/tensorrt_llm/kernels/mlaKernels.cu
+++ b/cpp/tensorrt_llm/kernels/mlaKernels.cu
@@ -447,8 +447,8 @@ __global__ void applyMLARopeAndAssignQKVKernelGeneration(T* qkv_output, T* q_pe,
     int q_pe_stride, KvCacheDataType cache_type, float* bmm1_scale, float* bmm2_scale, float const* quant_scale_o,
     float const* quant_scale_q, float const* quant_scale_kv, float const* dequant_scale_q,
     float const* dequant_scale_kv, float host_bmm1_scale, int32_t const* helix_position_offsets,
-    bool const* helix_is_inactive_rank, int32_t const* helix_local_slots = nullptr,
-    bool precomputed_cu_seqlens = false, bool precomputed_fmha_scheduler = false)
+    bool const* helix_is_inactive_rank, int32_t const* helix_local_slots = nullptr, bool precomputed_cu_seqlens = false,
+    bool precomputed_fmha_scheduler = false)
 {
     // Constants.
     using VecT = typename VecType<T>::Type;
@@ -1705,8 +1705,8 @@ void invokeMLARopeGeneration(MlaParams<T>& params, KVCacheBuffer kv_cache_buffer
         params.seqQOffset, params.fmha_tile_counter, params.cache_seq_lens, params.cu_kv_seqlens, params.q_pe_ld,
         params.q_pe_stride, params.cache_type, params.bmm1_scale, params.bmm2_scale, params.quant_scale_o,
         quant_scale_q_eff, params.quant_scale_kv, params.dequant_scale_q, params.dequant_scale_kv,
-        params.host_bmm1_scale, params.helix_position_offsets, params.helix_is_inactive_rank,
-        params.helix_local_slots, params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler);
+        params.host_bmm1_scale, params.helix_position_offsets, params.helix_is_inactive_rank, params.helix_local_slots,
+        params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler);
 }
 
 template <typename T, typename KVCacheBuffer>
diff --git a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
index 8c5f67b68755..4bb6c364ff53 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
@@ -313,11 +313,11 @@ def _is_supported_with_reason(
                 # Multi-token decode under helix needs the per-token bound /
                 # write-slot buffers of the speculative verify-group path.
                 return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix."
-            if seq_len_q != 1 and self._get_kernel_dtype(
-                    attn, q) == torch.float8_e4m3fn:
+            if seq_len_q != 1 and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn:
                 return False, (
                     "CuTe DSL MLA FMHA helix verify groups require a bf16/fp16 "
-                    "KV cache (the fp8 kernel has no per-token bounds).")
+                    "KV cache (the fp8 kernel has no per-token bounds)."
+                )
             softmax_stats = fwd.softmax_stats_tensor
             if softmax_stats is None:
                 return False, "CuTe DSL MLA FMHA requires softmax_stats_tensor with Helix."
@@ -522,10 +522,15 @@ def _run_mla_decode(
             params.fwd.softmax_stats_tensor,
             # Per-token rank-local bounds, filled by
             # recompute_helix_spec_buffers. None everywhere else.
-            (meta.helix_kv_bounds[:num_tokens] if
-             (meta.helix_position_offsets is not None
-              and meta._helix_spec_tokens_valid
-              and kernel_dtype != torch.float8_e4m3fn) else None),
+            (
+                meta.helix_kv_bounds[:num_tokens]
+                if (
+                    meta.helix_position_offsets is not None
+                    and meta._helix_spec_tokens_valid
+                    and kernel_dtype != torch.float8_e4m3fn
+                )
+                else None
+            ),
         )
 
     def run_mla_generation(
diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
index 566638672ec4..cb7db89d879f 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
@@ -83,10 +83,12 @@ def is_supported(
         # gate both assume the new KV entries are the trailing slots of one
         # rank's kv_len. Reject rather than run it silently wrong; being last
         # in the library list, this makes dispatch raise.
-        if (metadata.helix_position_offsets is not None
-                and metadata._helix_spec_tokens_valid
-                and metadata.num_generations > 0
-                and q.shape[0] > metadata.num_seqs):
+        if (
+            metadata.helix_position_offsets is not None
+            and metadata._helix_spec_tokens_valid
+            and metadata.num_generations > 0
+            and q.shape[0] > metadata.num_seqs
+        ):
             return False
         return forward_args.attention_mask != CustomAttentionMask.CUSTOM and (
             forward_args.update_kv_cache or metadata.is_cross
diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
index f7ce10b639ce..73040326012e 100644
--- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
+++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.py
@@ -2835,8 +2835,7 @@ def softmax(
                             common_params.blk_coord[2] * self.seq_len_q +
                             q_tok_c]
                     else:
-                        k_bound = common_params.K - (self.seq_len_q -
-                                                     1) + q_tok
+                        k_bound = common_params.K - (self.seq_len_q - 1) + q_tok
                     tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less(
                         tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index,
                         k_bound,
@@ -2887,8 +2886,7 @@ def softmax(
                             common_params.blk_coord[2] * self.seq_len_q +
                             q_tok_c]
                     else:
-                        k_bound = common_params.K - (self.seq_len_q -
-                                                     1) + q_tok
+                        k_bound = common_params.K - (self.seq_len_q - 1) + q_tok
                     tTR_rAcc[i] = (tTR_rAcc[i] if cute.elem_less(
                         tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index,
                         k_bound,
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index b52325409f7f..ec9cdca66f3d 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -2070,14 +2070,17 @@ def _setup_helix_mappings(
             # py_decoding_iter - 1) is stale once any draft token was
             # accepted -> silently wrong RoPE positions and KV slots. Reject
             # the trip wires until that loop is helix-group aware.
-            if (spec_config.acceptance_rate_window_size is not None
-                    or spec_config.acceptance_rate_threshold is not None):
+            if (
+                spec_config.acceptance_rate_window_size is not None
+                or spec_config.acceptance_rate_threshold is not None
+            ):
                 raise ValueError(
                     "Kimi K3 helix does not support the speculation "
                     "acceptance-rate gate (acceptance_rate_window_size / "
                     "acceptance_rate_threshold): dynamically disabling "
                     "speculation mid-flight leaves helix requests on a "
-                    "single-token position formula.")
+                    "single-token position formula."
+                )
         cp = model_config.mapping.cp_size
         repurposed_tp = model_config.mapping.tp_size * cp
         if cfg.num_attention_heads % repurposed_tp != 0:
diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py
index b99352411c15..9f10e8e31a45 100644
--- a/tensorrt_llm/_torch/modules/mla.py
+++ b/tensorrt_llm/_torch/modules/mla.py
@@ -769,8 +769,10 @@ def _attn_forward_gen(
                 seq_start=attn_metadata.num_contexts,
                 num_seqs=attn_metadata.num_generations,
             )
-            if (isinstance(attn_metadata, TrtllmAttentionMetadata)
-                    and attn_metadata._helix_spec_tokens_valid):
+            if (
+                isinstance(attn_metadata, TrtllmAttentionMetadata)
+                and attn_metadata._helix_spec_tokens_valid
+            ):
                 # A rank owning only a group's tail page has zero visible KV
                 # for its leading tokens while the per-sequence kv_len is
                 # nonzero, so the mask above misses those rows. Their scores
@@ -778,8 +780,7 @@ def _attn_forward_gen(
                 # an average over uninitialized pool values; the combine
                 # scales it by corr = 0, and 0 * NaN would poison the token
                 # on every CP rank.
-                zero_kv_mask = (
-                    attn_metadata.helix_kv_bounds[:partial_o.shape[0]] == 0)
+                zero_kv_mask = attn_metadata.helix_kv_bounds[: partial_o.shape[0]] == 0
             return _helix_post_process(
                 partial_o,
                 softmax_stats,
diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
index f2d575d823b7..4ad60b37c514 100644
--- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
@@ -44,10 +44,6 @@
     LinearAttentionMetadata, LinearCacheType)
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig
 from tensorrt_llm.logger import logger
-
-# Shared with the KV budget estimator so allocator and budgeting can never
-# diverge on the sharding rule (config_utils is import-cycle-free).
-from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX,
                                                       BatchDesc, BufferConfig,
@@ -57,6 +53,10 @@
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode,
                                                       SsmLayerConfig)
 
+# Shared with the KV budget estimator so allocator and budgeting can never
+# diverge on the sharding rule (config_utils is import-cycle-free).
+from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size
+
 GB = 1 << 30
 
 # Replay kernels pad the token/window dimension to at least 16 for tensor-core
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 292c31844a3b..8c6c1f065dc4 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -4027,15 +4027,17 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]):
                         restore=True,
                     )
 
-                if (self.mapping.has_cp_helix()
-                        and isinstance(inputs['attn_metadata'],
-                                       TrtllmAttentionMetadata)
+                if (self.mapping.has_cp_helix() and isinstance(
+                        inputs['attn_metadata'], TrtllmAttentionMetadata)
                         and inputs['attn_metadata']._helix_spec_tokens_valid):
                     # Mirrors the correction in _preprocess_inputs. Only
                     # the offsets need reversing: the recompute's other
                     # outputs are rewritten from host state next prepare.
-                    inputs['attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= (
-                        self.previous_pos_id_offsets_cuda[:previous_batch_tokens])
+                    inputs[
+                        'attn_metadata'].helix_position_offsets[:previous_batch_tokens] -= (
+                            self.
+                            previous_pos_id_offsets_cuda[:previous_batch_tokens]
+                        )
 
     def _get_all_rank_num_tokens(self, attn_metadata: AttentionMetadata):
         if self.enable_attention_dp:
@@ -5640,8 +5642,8 @@ def append_cross_attention_state(request: LlmRequest,
 
             def _helix_local_len_host(global_len: int) -> int:
                 full, rem = divmod(global_len, _helix_ledger)
-                return full * _helix_phys + min(
-                    max(rem - _helix_rank_off, 0), _helix_phys)
+                return full * _helix_phys + min(max(rem - _helix_rank_off, 0),
+                                                _helix_phys)
 
             def _helix_pack_extend(request, group: int) -> int:
                 # A helix gen worker's token list is the rank-LOCAL
@@ -5849,7 +5851,6 @@ def _helix_pack_extend(request, group: int) -> int:
             # update batch index
             request.py_batch_idx = request.py_seq_slot
 
-
         _n_gen = len(generation_requests)
         # One-shot batch-level flag — True iff any generation request actually
         # carries multimodal payload. Lets the strip_mm_data branch below

From abf746cbcd5e5d56fe92bf4232cf7b654d604388 Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Mon, 31 Aug 2026 09:15:52 -0700
Subject: [PATCH 15/15] [None][fix] Make MoonViT replication helix-aware

_vision_requires_replication only tested num_heads % tp_size, but helix
places its parallelism on cp (tp_size=1), so the test never trips and
the tower inherits the cp mapping, failing the repurposed tp*cp
head-divisibility assertion (K3's 12 heads under tp1xcp8/16). The
replicated-mapping constructor had the same blind spot: world_size =
pp*tp collapses to 1 under helix while ranks span cp_size, which
Mapping rejects. Force replication for any cp > 1 and fold cp into the
replicated world size.

Latent on text-only disagg deployments (TLLM_MULTIMODAL_DISAGGREGATED=1
skips tower construction); live on multimodal helix serving.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_k25.py    | 14 ++++++++++++--
 .../_torch/modeling/test_kimi_k3_config_routing.py | 14 ++++++++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py
index d5ad4564a43f..99753c238ae7 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py
@@ -372,6 +372,11 @@ def _vision_requires_replication(model_config: ModelConfig, num_heads: int) -> b
     mapping = model_config.mapping
     if mapping.enable_attention_dp:
         return True
+    # Helix carries its parallelism in cp with tp_size=1, so a tp-only test
+    # never trips; the tower has no context-parallel form, so any cp > 1
+    # must replicate.
+    if mapping.cp_size > 1:
+        return True
     return (num_heads % mapping.tp_size) != 0
 
 
@@ -379,12 +384,17 @@ def _get_vision_tp_mapping(model_config: ModelConfig, num_heads: int) -> Mapping
     if not _vision_requires_replication(model_config, num_heads):
         return model_config.mapping
 
+    # Fold every parallel dimension (incl. helix cp) into pp so each rank
+    # runs the tower replicated; without cp the world size collapses below
+    # the rank range under helix.
+    attn_ranks = (model_config.mapping.pp_size * model_config.mapping.tp_size
+                  * model_config.mapping.cp_size)
     return Mapping(
-        world_size=model_config.mapping.pp_size * model_config.mapping.tp_size,
+        world_size=attn_ranks,
         rank=model_config.mapping.rank,
         gpus_per_node=model_config.mapping.gpus_per_node,
         tp_size=1,
-        pp_size=model_config.mapping.pp_size * model_config.mapping.tp_size,
+        pp_size=attn_ranks,
     )
 
 
diff --git a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py
index c8f23083dbb0..01cf6cfbe542 100644
--- a/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py
+++ b/tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py
@@ -84,10 +84,12 @@ def test_other_model_type_with_subconfigs(self):
         self.assertFalse(is_kimi_k3_multimodal_config(cfg))
 
 
-def _vision_model_config(tp_size, enable_attention_dp):
+def _vision_model_config(tp_size, enable_attention_dp, cp_size=1):
     """Minimal stand-in exposing the mapping fields the predicate reads."""
     return SimpleNamespace(
-        mapping=SimpleNamespace(tp_size=tp_size, enable_attention_dp=enable_attention_dp)
+        mapping=SimpleNamespace(
+            tp_size=tp_size, enable_attention_dp=enable_attention_dp, cp_size=cp_size
+        )
     )
 
 
@@ -108,6 +110,14 @@ def test_k25_16_heads_under_tp8_shards(self):
     def test_attention_dp_always_replicates(self):
         self.assertTrue(_vision_requires_replication(_vision_model_config(16, True), num_heads=16))
 
+    def test_helix_cp_requires_replication(self):
+        # Helix carries its parallelism in cp with tp_size=1; the tower has no
+        # context-parallel form, so any cp > 1 must replicate even when
+        # num_heads % tp_size == 0.
+        self.assertTrue(
+            _vision_requires_replication(_vision_model_config(1, False, cp_size=8), num_heads=12)
+        )
+
 
 def _load_config_from_dict(cfg: dict[str, Any]) -> PretrainedConfig:
     """Round-trip a raw config dict through load_pretrained_config."""