[None][feat] Kimi K3: DSpark speculative decoding under Helix CP

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);
 
                     {
@@ -1689,7 +1705,7 @@ 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.host_bmm1_scale, params.helix_position_offsets, params.helix_is_inactive_rank, params.helix_local_slots,
         params.precomputed_cu_seqlens, params.precomputed_fmha_scheduler);
 }
 

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>

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();

tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py

@@ -309,8 +309,15 @@ 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 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."
+            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 +520,17 @@ def _run_mla_decode(
             # Max batch size for the AutoTuner to profile.
             int(meta.max_num_requests),
             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
+            ),
         )
 
     def run_mla_generation(

tensorrt_llm/_torch/attention_backend/fmha/fallback.py

@@ -76,7 +76,20 @@ def is_supported(
         *,
         phase: Optional[FmhaPhase] = None,
     ) -> bool:
-        del q, k, v, phase
+        del k, v, phase
+        # 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
+            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
         )

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 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
+    _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,33 @@ def _post_init_with_buffers(self, buffers) -> None:
                 device='cpu',
                 pin_memory=prefer_pinned(),
             )
+            # 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, ),
+                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,
+            )
+            # 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',
+                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,13 +609,18 @@ 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.
 
         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 +636,67 @@ 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)
+        # 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))
+        # 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}")
+        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 _bind_runtime_views(
         self,
         *,
@@ -722,9 +824,16 @@ 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 self._helix_spec_tokens_valid:
+                # 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:
+                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 +2553,10 @@ def mla_rope_generation(
         helix_tensor_params = [
             metadata.helix_position_offsets, metadata.helix_is_inactive_rank
         ]
+        if metadata._helix_spec_tokens_valid:
+            # 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(
             fused_q,