diff --git a/cpp/tensorrt_llm/kernels/mlaKernels.cu b/cpp/tensorrt_llm/kernels/mlaKernels.cu index e98768faea82..e2c1e2cae312 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::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(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(global_token_idx) * (c_k + ROPE_DIM); { @@ -1689,7 +1705,7 @@ void invokeMLARopeGeneration(MlaParams& 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); } 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 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 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 fused_q, // [tokens, num_hea = helix_position_offsets.has_value() ? helix_position_offsets->data_ptr() : nullptr; bool const* helix_is_inactive_rank_ptr = helix_is_inactive_rank.has_value() ? helix_is_inactive_rank->data_ptr() : 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(); + 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(cu_q_seqlens.data_ptr()); int* cu_kv_seqlens_ptr = reinterpret_cast(cu_kv_seqlens.data_ptr()); @@ -308,8 +320,8 @@ void MLARopeGeneration(std::optional 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(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(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..4bb6c364ff53 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,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( diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py index 4ce5f8e149f1..cb7db89d879f 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py +++ b/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 ) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 04b17e36702c..19192b8e1acd 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 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,6 +609,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 +617,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 +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, 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..7c4e1896dabf 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,33 @@ 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 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() + 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 +9864,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 +9887,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 +9918,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 +9946,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 +10014,7 @@ def _( output_scale: float, max_batch_size: int, softmax_stats: Optional[torch.Tensor], + kv_bounds: Optional[torch.Tensor], ) -> None: return None @@ -10005,8 +10040,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 +10081,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 +10119,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..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 @@ -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, @@ -328,6 +329,7 @@ def __call__( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -348,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, @@ -365,6 +368,7 @@ def run_with_softmax_stats( workspace, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale, output_scale, @@ -385,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, @@ -392,6 +397,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 +806,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -827,6 +842,7 @@ class SplitKVKernelSharedStorage: acc_lse, split_kv, cache_seqs, + kv_bounds, block_split_kvs, softmax_scale_log2, output_scale, @@ -859,6 +875,7 @@ class SplitKVKernelSharedStorage: split_kv, cache_seqs, block_split_kvs, + kv_bounds, ) else: reduction_kernel = self.reduction_kernel( @@ -938,6 +955,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 +1374,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 +1477,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 +1548,27 @@ 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): + # 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 + + q_tok] > 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 +2488,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 +2823,19 @@ 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): + # 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[ + 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 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 +2875,18 @@ 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): + # 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[ + 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 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..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 @@ -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, @@ -346,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, 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/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index dce32f886909..ec9cdca66f3d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2054,11 +2054,33 @@ 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 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 " + 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/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 86b960c29d93..9f10e8e31a45 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -769,6 +769,18 @@ 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 + ): + # 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( partial_o, softmax_stats, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b4422c084463..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( @@ -698,13 +712,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 +749,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 +771,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 +796,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, @@ -1562,10 +1599,18 @@ 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 + # 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() 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/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/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/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 5e188a1061da..4ad60b37c514 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -53,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 @@ -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): diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e0ab193c50fd..8c6c1f065dc4 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.enable_spec_decode and self.mapping.has_cp_helix(): + # 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 + and md._helix_spec_tokens_valid): + helix_gen_tokens = (inputs['input_ids'].shape[0] - + md.num_ctx_tokens) + if not self._disable_overlap_scheduler: + # 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( + 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() @@ -4005,6 +4027,18 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]): restore=True, ) + 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] + ) + def _get_all_rank_num_tokens(self, attn_metadata: AttentionMetadata): if self.enable_attention_dp: num_tokens = attn_metadata.num_tokens @@ -5595,6 +5629,36 @@ def append_cross_attention_state(request: LlmRequest, generation_requests.append(request) extend_requests += extend_dummy_requests + # 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() + 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) + + 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( @@ -5667,6 +5731,16 @@ 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: + # No in-flight predecessor, so every value is exact. + group = 1 + num_draft_tokens + 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) + 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 +5772,18 @@ 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: provisional values, per the + # 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 + 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] = ( + 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 +5851,6 @@ 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() _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 @@ -5880,10 +5963,12 @@ def append_cross_attention_state(request: LlmRequest, # 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 @@ -5897,6 +5982,9 @@ 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 list aligned in mixed batches. + 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 +6408,9 @@ def previous_seq_slots_device(): attn_metadata.update_helix_param( helix_position_offsets=helix_position_offsets, helix_is_inactive_rank=helix_is_inactive_rank, + # 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: 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."""