diff --git a/docs/source/features/feature-combination-matrix.md b/docs/source/features/feature-combination-matrix.md index b56c1b219af9..0f968729feac 100644 --- a/docs/source/features/feature-combination-matrix.md +++ b/docs/source/features/feature-combination-matrix.md @@ -11,7 +11,7 @@ | Attention Data Parallelism | Yes | Yes | Yes | Yes | Yes | Known issues | --- | | | | | | | | | | | | | | Disaggregated Serving | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | Chunked Prefill | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | -| Speculative Decoding — Linear | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | --- | | | | | | | | | | +| Speculative Decoding — Linear | Yes | Yes | Yes | No | Yes | Yes (Eagle3 one-model) | Yes | Yes | Yes | --- | | | | | | | | | | | Speculative Decoding — Dynamic Trees | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | --- | | | | | | | | | | Speculative Decoding — Legacy Path (NGram, user-provided) | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | No | No | --- | | | | | | | | | Torch Sampler | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 420fd14a1fe0..f4fdc2f6c368 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import functools import math import os @@ -443,6 +444,7 @@ def _post_init_with_buffers(self, buffers) -> None: device='cpu', pin_memory=prefer_pinned(), ) + # One inactive flag per sequence; Helix ownership is per verify group. self.helix_is_inactive_rank = self.get_empty( buffers, (self.max_num_sequences, ), @@ -456,6 +458,75 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) + # Static buffers for the flattened per-query-row generation read. + self._helix_gen_flatten_active = False + self._helix_flat_total_q = 0 + self._helix_flat_kv_lens_cuda = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="helix_flat_kv_lens_cuda", + dtype=torch.int, + capture_graph=capture_graph, + ) + self._helix_flat_kv_lens_cpu = torch.empty_like( + self._helix_flat_kv_lens_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Host-only prompt lengths for the flattened rows; the generation + # kernel reads only host_context_lengths, so no device copy is kept. + self._helix_flat_prompt_lens_cpu = torch.empty_like( + self._helix_flat_kv_lens_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Flattened rows are all generation requests (RequestType 1); the + # buffer is read as all-ones, so fill it once. + self._helix_flat_request_types_cpu = torch.empty_like( + self._helix_flat_prompt_lens_cpu) + self._helix_flat_request_types_cpu.fill_(1) + self._helix_flat_cu_q_seqlens = self.get_empty( + buffers, + (self.max_num_tokens + 1, ), + cache_name="helix_flat_cu_q_seqlens", + dtype=torch.int, + capture_graph=capture_graph, + ) + self._helix_flat_cu_kv_seqlens = self.get_empty( + buffers, + (self.max_num_tokens + 1, ), + cache_name="helix_flat_cu_kv_seqlens", + dtype=torch.int, + capture_graph=capture_graph, + ) + # Sequence index per flattened row (for overlap kv-length correction). + self._helix_flat_row_to_seq = self.get_empty( + buffers, + (self.max_num_tokens, ), + cache_name="helix_flat_row_to_seq", + dtype=torch.int, + capture_graph=capture_graph, + ) + self._helix_flat_host_total_kv_lens = torch.empty(2, + device='cpu', + dtype=torch.int) + if self.kv_cache_manager is not None: + num_attention_op_pools = getattr( + self.kv_cache_manager, "num_attention_op_pools", + self.kv_cache_manager.num_pools) + self._helix_flat_block_offsets = self.get_empty( + buffers, + [ + num_attention_op_pools, self.max_num_tokens, 2, + self.kv_cache_manager.max_blocks_per_seq + ], + cache_name="helix_flat_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + else: + self._helix_flat_block_offsets = None + def on_update_kv_lens(self): # After changing the kv_lens/kv_lens_cuda, we may need to update other metadata. # Especially for the changes in the _preprocess_inputs() of model_engine.py. @@ -469,6 +540,153 @@ def update_for_spec_dec(self) -> None: if self.enable_flash_mla: self._flash_mla_metadata_valid = False + def _maybe_prepare_helix_flatten(self, cached_token_lens: torch.Tensor, + kv_lens: torch.Tensor) -> None: + """Build the flattened per-query-row generation view for Helix MTP verify. + + Each query row is presented as q_len == 1 so the trtllm-gen causal slope + vanishes and each row attends its exact KV bound. Only runs for a pure- + generation multi-token verify (num_contexts == 0 and total_q > num_seqs). + """ + self._helix_gen_flatten_active = False + if not (self.enable_helix and self.num_contexts == 0 + and self.helix_is_inactive_rank_cpu is not None): + return + num_seqs = self.num_seqs + if num_seqs == 0: + return + seg_lens = self.seq_lens_kv[:num_seqs].to(torch.int64) + total_q = int(seg_lens.sum().item()) + if total_q <= num_seqs: + return + + # Active rank: inclusive row index (1..q_len). Inactive rank: 0. + active = (~self.helix_is_inactive_rank_cpu[:num_seqs]).to(torch.int64) + seg_starts = torch.cumsum(seg_lens, 0) - seg_lens + row_incl = (torch.arange(total_q, dtype=torch.int64) - + torch.repeat_interleave(seg_starts, seg_lens) + 1) + owned_incl = row_incl * torch.repeat_interleave(active, seg_lens) + + cached64 = cached_token_lens[:num_seqs].to(torch.int64) + # Per-row KV read bound (no extra tokens, matching kv_lens_cuda_runtime). + flat_kv = (torch.repeat_interleave(cached64, seg_lens) + owned_incl).to( + torch.int) + flat_prompt = torch.repeat_interleave(self.prompt_lens_cpu[:num_seqs], + seg_lens).to(torch.int) + + self._helix_flat_kv_lens_cpu[:total_q].copy_(flat_kv) + self._helix_flat_kv_lens_cuda[:total_q].copy_( + self._helix_flat_kv_lens_cpu[:total_q], non_blocking=True) + self._helix_flat_prompt_lens_cpu[:total_q].copy_(flat_prompt) + # _helix_flat_request_types_cpu is pre-filled with 1s at construction. + + # q_len == 1 per request -> cu_q_seqlens = [0, 1, 2, ..., total_q]. + q_offsets = torch.arange(total_q + 1, dtype=torch.int) + self._helix_flat_cu_q_seqlens[:total_q + 1].copy_(q_offsets, + non_blocking=True) + # Exclusive prefix sum of the per-row KV lengths (seqKVOffset layout). + kv_offsets = torch.zeros(total_q + 1, dtype=torch.int) + kv_offsets[1:] = torch.cumsum(flat_kv, 0) + self._helix_flat_cu_kv_seqlens[:total_q + 1].copy_(kv_offsets, + non_blocking=True) + + # Row to sequence index for on-device overlap kv-length correction. + row_to_seq = torch.repeat_interleave( + torch.arange(num_seqs, dtype=torch.int), seg_lens.to(torch.int)) + self._helix_flat_row_to_seq[:total_q].copy_(row_to_seq, + non_blocking=True) + + self._helix_flat_host_total_kv_lens[0] = 0 + self._helix_flat_host_total_kv_lens[1] = int(flat_kv.sum().item()) + + # Replicate each sequence's KV block table across its query rows; the + # rows share the same physical KV blocks. + if (self._helix_flat_block_offsets is not None + and self.kv_cache_block_offsets is not None): + src = self.kv_cache_block_offsets[:, :num_seqs] + rep = torch.repeat_interleave(src, seg_lens.to(src.device), dim=1) + self._helix_flat_block_offsets[:, :total_q].copy_(rep, + non_blocking=True) + + self._helix_flat_total_q = total_q + self._helix_gen_flatten_active = True + + def apply_helix_overlap_flatten_correction( + self, + per_request_kv_offset: torch.Tensor, + max_extra_per_row: int, + sign: int = 1) -> None: + """Correct flattened per-row KV bounds for the overlap scheduler. + + Broadcasts the per-request kv-length correction onto each flattened row + and rebuilds cumulative KV offsets. Applied with sign=+1 in + _preprocess_inputs and undone with sign=-1 in _postprocess_inputs. + """ + if not getattr(self, "_helix_gen_flatten_active", False): + return + n = self._helix_flat_total_q + if n == 0: + return + off_rows = per_request_kv_offset[self._helix_flat_row_to_seq[:n].to( + torch.long)].to(self._helix_flat_kv_lens_cuda.dtype) + self._helix_flat_kv_lens_cuda[:n] += sign * off_rows + # Device-only correction; host kv_lens stays stale (non-Helix overlap path). + # Leading zero via device memset (CUDA graph capture safe). + self._helix_flat_cu_kv_seqlens[:1].zero_() + torch.cumsum(self._helix_flat_kv_lens_cuda[:n], + 0, + out=self._helix_flat_cu_kv_seqlens[1:n + 1]) + # Host sizing hint only; exact total would require a device sync. + self._helix_flat_host_total_kv_lens[1] = ( + int(self._helix_flat_host_total_kv_lens[1]) + + sign * n * max_extra_per_row) + + @contextlib.contextmanager + def helix_flattened_generation(self): + """Swap generation attention metadata to the flattened per-row view. + + Yields (cu_q_seqlens, cu_kv_seqlens) for the flattened layout, or + (None, None) when flattening is inactive. Reshapes the read only; the KV + write already ran on the un-flattened layout. + """ + if not getattr(self, "_helix_gen_flatten_active", False): + yield None, None + return + n = self._helix_flat_total_q + saved = { + "kv_lens_cuda_runtime": self.kv_lens_cuda_runtime, + "kv_lens_runtime": self.kv_lens_runtime, + "prompt_lens_cuda_runtime": self.prompt_lens_cuda_runtime, + "prompt_lens_cpu_runtime": self.prompt_lens_cpu_runtime, + "host_request_types_runtime": self.host_request_types_runtime, + "kv_cache_block_offsets": self.kv_cache_block_offsets, + "max_num_requests": self.max_num_requests, + } + saved_total_kv = (int(self.host_total_kv_lens[0]), + int(self.host_total_kv_lens[1])) + + self.kv_lens_cuda_runtime = self._helix_flat_kv_lens_cuda[:n] + self.kv_lens_runtime = self._helix_flat_kv_lens_cpu[:n] + # Device prompt-lens view is only shape-checked, never read here; alias + # the kv-lens buffer to satisfy the assertion without a separate copy. + self.prompt_lens_cuda_runtime = self._helix_flat_kv_lens_cuda[:n] + self.prompt_lens_cpu_runtime = self._helix_flat_prompt_lens_cpu[:n] + self.host_request_types_runtime = self._helix_flat_request_types_cpu[:n] + self.host_total_kv_lens[0] = self._helix_flat_host_total_kv_lens[0] + self.host_total_kv_lens[1] = self._helix_flat_host_total_kv_lens[1] + if self._helix_flat_block_offsets is not None: + self.kv_cache_block_offsets = self._helix_flat_block_offsets[:, :n] + # Flattened batch has total_q requests; size C++ workspace accordingly. + self.max_num_requests = max(self.max_num_requests, n) + try: + yield (self._helix_flat_cu_q_seqlens[:n + 1], + self._helix_flat_cu_kv_seqlens[:n + 1]) + finally: + for key, value in saved.items(): + setattr(self, key, value) + self.host_total_kv_lens[0] = saved_total_kv[0] + self.host_total_kv_lens[1] = saved_total_kv[1] + def update_helix_param( self, helix_position_offsets: List[int], @@ -479,7 +697,9 @@ 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_is_inactive_rank: Whether the current rank is inactive, per request, + with shape (num_seqs,). Ownership is per verify group, so a single flag + covers all of a sequence's query tokens (golden + drafts). """ if helix_position_offsets is not None and self.helix_position_offsets is not None: num_tokens = len(helix_position_offsets) @@ -489,11 +709,11 @@ def update_helix_param( self.helix_position_offsets_cpu[:num_tokens], non_blocking=True) if helix_is_inactive_rank is not None and self.helix_is_inactive_rank is not None: - batch_size = len(helix_is_inactive_rank) - self.helix_is_inactive_rank_cpu[:batch_size].copy_( + num_flags = len(helix_is_inactive_rank) + self.helix_is_inactive_rank_cpu[:num_flags].copy_( torch.tensor(helix_is_inactive_rank, dtype=torch.bool)) - self.helix_is_inactive_rank[:batch_size].copy_( - self.helix_is_inactive_rank_cpu[:batch_size], non_blocking=True) + self.helix_is_inactive_rank[:num_flags].copy_( + self.helix_is_inactive_rank_cpu[:num_flags], non_blocking=True) def _bind_runtime_views( self, @@ -556,10 +776,10 @@ def prepare(self) -> None: # number of tokens needed in the kv cache for each sequence after the next pass. 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() + # Active sequences grow kv length by seq_lens_kv; inactive ones do not. + active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] kv_lens[active_rank] += self.seq_lens_kv[active_rank] else: kv_lens = cached_token_lens + \ @@ -572,6 +792,7 @@ def prepare(self) -> None: self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory( kv_lens[:self.num_seqs]), non_blocking=True) + # total kv lens for context requests and generation requests, without extra tokens self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self. @@ -605,6 +826,9 @@ def prepare(self) -> None: self.draft_kv_cache_block_offsets, self.request_ids, self.beam_width, self.num_contexts, self.num_seqs) + # Build flattened view after block offsets are ready. + self._maybe_prepare_helix_flatten(cached_token_lens, kv_lens) + # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for # kv_lens_runtime, which becomes host_past_key_value_lengths and diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 57627b4650e7..e4bfa5ad135f 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -1613,9 +1613,14 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], layer_idx: int, aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], - is_separate_draft_engine: bool = False): - super().__init__(model_config, layer_idx, aux_stream_dict, - is_separate_draft_engine) + is_separate_draft_engine: bool = False, + mapping_with_cp: Optional[Mapping] = None): + # Thread mapping_with_cp into attention so MLA keeps the full head count. + super().__init__(model_config, + layer_idx, + aux_stream_dict, + is_separate_draft_engine, + mapping_with_cp=mapping_with_cp) config = model_config.pretrained_config self.hidden_dim = config.hidden_size self.moe_intermediate_size = config.moe_intermediate_size @@ -1691,11 +1696,11 @@ def norm_hidden(): disable_on_compile=True, ) hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1) - # Split hidden_states columnwise based on TP - tp_size = self.model_config.mapping.tp_size - tp_rank = self.model_config.mapping.tp_rank + # Use self.mapping (captured at construction), not model_config.mapping. + tp_size = self.mapping.tp_size + tp_rank = self.mapping.tp_rank - if tp_size > 1 and not (self.model_config.mapping.enable_attention_dp): + if tp_size > 1 and not (self.mapping.enable_attention_dp): hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank] hidden_states = self.eh_proj(hidden_states) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 6e17258067a1..11641672edb9 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1439,6 +1439,8 @@ def __init__( # Import here to avoid circular import model_type = model_config.pretrained_config.model_type mtp_layer = None + # Only DeepseekV3MTP accepts mapping_with_cp for Helix CP attention. + supports_mapping_with_cp = False match model_type: case "glm4_moe": from .modeling_glm import Glm4MTP @@ -1446,6 +1448,7 @@ def __init__( case "deepseek_v3" | "deepseek_v32" | "glm_moe_dsa": from .modeling_deepseekv3 import DeepseekV3MTP mtp_layer = DeepseekV3MTP + supports_mapping_with_cp = True case "exaone_moe": from .modeling_exaone_moe import ExaoneMoeMTP mtp_layer = ExaoneMoeMTP @@ -1478,9 +1481,15 @@ def __init__( moe_load_balancer_set_repeated_for_next_layer(mtp_repeat_count) + # Thread mapping_with_cp into MTP attention for correct MLA head count. + mtp_layer_kwargs = {} + mapping_with_cp = getattr(model, 'mapping_with_cp', None) + if supports_mapping_with_cp and mapping_with_cp is not None: + mtp_layer_kwargs['mapping_with_cp'] = mapping_with_cp + self.mtp_layers = nn.ModuleList([ mtp_layer(model_config, layer_idx + start_layer_idx, - model.aux_stream_dict) + model.aux_stream_dict, **mtp_layer_kwargs) for layer_idx in range(mtp_num_layers) ]) self.lm_head = lm_head diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 75651c920acd..417b78cee805 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import functools import math import os @@ -931,13 +932,21 @@ def _attn_forward_gen( (q.shape[0], self.num_heads_tp, 2), device=q.device, dtype=torch.float32 ) kwargs["softmax_stats_tensor"] = softmax_stats - partial_o = attn_backend.forward( - q, - k, - v, - attn_metadata, - forward_args=AttentionForwardArgs(**kwargs), - ) + # Flatten generation read to one q_len==1 row per query token (Helix MTP verify). + flatten_ctx = getattr(attn_metadata, "helix_flattened_generation", None) + with ( + flatten_ctx() if flatten_ctx is not None else contextlib.nullcontext((None, None)) + ) as (flat_cu_q, flat_cu_kv): + if flat_cu_q is not None: + kwargs["cu_q_seqlens"] = flat_cu_q + kwargs["cu_kv_seqlens"] = flat_cu_kv + partial_o = attn_backend.forward( + q, + k, + v, + attn_metadata, + forward_args=AttentionForwardArgs(**kwargs), + ) kv_lora_rank = partial_o.shape[-1] // self.num_heads_tp assert self.kv_lora_rank == kv_lora_rank diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 10a635e01f61..c5993d0922cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -727,6 +727,17 @@ def __init__( self.seqlen_this_rank_cp = self.prompt_len self.total_input_len_cp = self.prompt_len self.py_helix_is_inactive_rank = False + # Committed decode length across CP ranks, advanced by 1 + accepted per + # verify group (unlike py_decoding_iter, which counts iterations). The + # per-rank cached KV length lives in seqlen_this_rank_cp. + self.py_helix_global_decode_len = 0 + # Reserve-side verify-group counter: owner = (index // tpb) % cp_size. + self.py_helix_decode_group_index = 0 + # Rewind-side verify-group counter; matches reserve in FIFO order and + # recovers per-group ownership deterministically under overlap. + self.py_helix_rewind_group_index = 0 + # Prior group's ownership; gates on-device kv-length correction under overlap. + self.py_helix_prev_group_owns = False self.py_batch_idx = None self.py_draft_pages_allocated = 0 self.py_rewind_len = 0 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5abc84022421..d1dfcc6f9acd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2425,6 +2425,15 @@ def _preprocess_inputs(self, inputs: Dict[str, Any]): ) inputs['attn_metadata'].on_update_kv_lens() + # On-device Helix overlap corrections (RoPE positions and flattened read). + if self.mapping.has_cp_helix(): + self._apply_helix_overlap_corrections( + inputs['attn_metadata'], + num_ctx_tokens, + previous_batch_tokens, + num_gen_requests, + sign=1) + if self.guided_decoder is not None: self.guided_decoder.token_event.record() @@ -2472,6 +2481,37 @@ def _postprocess_inputs(self, inputs: Dict[str, Any]): previous_kv_lens_offsets_cuda[:num_gen_requests] ) + # Undo Helix overlap corrections for CUDA graph replay consistency. + if self.mapping.has_cp_helix(): + self._apply_helix_overlap_corrections( + inputs['attn_metadata'], + num_ctx_tokens, + previous_batch_tokens, + num_gen_requests, + sign=-1) + + def _apply_helix_overlap_corrections(self, attn_metadata, num_ctx_tokens, + previous_batch_tokens, + num_gen_requests, sign): + """On-device Helix verify metadata corrections for overlap scheduling. + + Fixes accept-dependent RoPE positions and flattened per-row KV bounds. + kv_lens_cuda is corrected by the caller. No-op outside a Helix verify. + """ + helix_pos = getattr(attn_metadata, 'helix_position_offsets', None) + if helix_pos is not None and previous_batch_tokens > 0: + helix_pos[num_ctx_tokens:num_ctx_tokens + + previous_batch_tokens] += ( + sign * self. + previous_pos_id_offsets_cuda[:previous_batch_tokens]) + apply_flat = getattr(attn_metadata, + 'apply_helix_overlap_flatten_correction', None) + if apply_flat is not None and num_gen_requests > 0: + # Upper bound for host total-KV sizing hint (1 + runtime_draft_len). + apply_flat(self.previous_kv_lens_offsets_cuda[:num_gen_requests], + 1 + self.runtime_draft_len, + sign=sign) + def _get_all_rank_num_tokens(self, attn_metadata: AttentionMetadata): if self.enable_attention_dp: num_tokens = attn_metadata.num_tokens @@ -3187,6 +3227,23 @@ def _apply_incremental_update_target( return inputs, self.gather_ids_cuda[:num_generation_tokens] + def _helix_verify_token_params(self, request: LlmRequest, num_draft: int, + tokens_per_block: int): + """Compute Helix parameters for a speculative verify forward. + + Returns (global_positions, is_inactive, num_active). Ownership comes from + py_helix_is_inactive_rank, set at KV reserve time. Overlap corrections + run on-device in _preprocess_inputs. + """ + del tokens_per_block + g = request.py_helix_global_decode_len + total_input_len = request.total_input_len_cp + first_pos = total_input_len + g + global_positions = list(range(first_pos, first_pos + 1 + num_draft)) + is_inactive = request.py_helix_is_inactive_rank + num_active = 0 if is_inactive else (1 + num_draft) + return global_positions, is_inactive, num_active + def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, @@ -3488,6 +3545,16 @@ def append_cross_attention_state(request: LlmRequest, runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + + # helix_position_offsets: per query token. helix_is_inactive_rank: per request. + helix_is_inactive_rank, helix_position_offsets = [], [] + # Prior group's ownership for on-device kv-length correction under overlap. + helix_prev_group_owns = [] + # Cache invariant method result to avoid repeated calls per-request. + _has_cp_helix = self.mapping.has_cp_helix() + _helix_tokens_per_block = (kv_cache_manager.tokens_per_block + if _has_cp_helix + and kv_cache_manager is not None else None) for request in extend_requests: request_ids.append(request.py_request_id) request_accepted_path[ @@ -3523,10 +3590,19 @@ def append_cross_attention_state(request: LlmRequest, list( range(len(position_ids), len(position_ids) + 1 + num_draft_tokens))) - position_ids.extend( - list( - range(past_seen_token_num, - past_seen_token_num + 1 + num_draft_tokens))) + if _has_cp_helix: + # Global positions and per-rank cached length under Helix. + positions_h, is_inactive_h, _ = self._helix_verify_token_params( + request, num_draft_tokens, _helix_tokens_per_block) + past_seen_token_num = request.seqlen_this_rank_cp + position_ids.extend(positions_h) + helix_position_offsets.extend(positions_h) + helix_is_inactive_rank.append(is_inactive_h) + else: + position_ids.extend( + list( + range(past_seen_token_num, + past_seen_token_num + 1 + num_draft_tokens))) num_cached_tokens_per_seq.append(past_seen_token_num) request.cached_tokens = num_cached_tokens_per_seq[-1] # update batch index @@ -3546,17 +3622,34 @@ def append_cross_attention_state(request: LlmRequest, list( range(len(position_ids), len(position_ids) + runtime_tokens_per_gen_step))) - position_ids.extend( - list( - range(past_seen_token_num, past_seen_token_num + - runtime_tokens_per_gen_step))) + if _has_cp_helix: + # Overlap verify path under Helix. + positions_h, is_inactive_h, _ = self._helix_verify_token_params( + request, runtime_draft_token_buffer_width, + _helix_tokens_per_block) + position_ids.extend(positions_h) + helix_position_offsets.extend(positions_h) + helix_is_inactive_rank.append(is_inactive_h) + else: + position_ids.extend( + list( + range( + past_seen_token_num, past_seen_token_num + + runtime_tokens_per_gen_step))) # previous tensor previous_batch_indices.append(previous_batch_idx) previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) - num_cached_tokens_per_seq.append(past_seen_token_num + - runtime_tokens_per_gen_step) + if _has_cp_helix: + # Prior group's ownership for kv-length correction. + helix_prev_group_owns.append( + request.py_helix_prev_group_owns) + num_cached_tokens_per_seq.append( + request.seqlen_this_rank_cp) + else: + num_cached_tokens_per_seq.append( + past_seen_token_num + runtime_tokens_per_gen_step) request.cached_tokens = num_cached_tokens_per_seq[-1] if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: @@ -3620,9 +3713,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 @@ -3672,18 +3762,12 @@ def append_cross_attention_state(request: LlmRequest, position_id = past_seen_token_num if _has_cp_helix: - # We compute a global position_id because each helix rank has only a subset of - # tokens for a sequence. + # Global position and per-rank cached length under Helix. position_id = request.total_input_len_cp + request.py_decoding_iter - 1 - if request.py_helix_is_inactive_rank: - past_seen_token_num = request.seqlen_this_rank_cp - else: - # Discount the token added to active rank in resource manager as it hasn't - # been previously seen. - past_seen_token_num = request.seqlen_this_rank_cp - 1 + past_seen_token_num = request.seqlen_this_rank_cp for beam in range(beam_width): - # Update helix-specific parameters. + # Plain decode: one query token when active, none when inactive. helix_is_inactive_rank.append( request.py_helix_is_inactive_rank) helix_position_offsets.append(position_id) @@ -3968,11 +4052,25 @@ def previous_seq_slots_device(): 0:previous_batch_tokens]], non_blocking=True) - self.previous_kv_lens_offsets_cuda[ - num_extend_reqeust_wo_dummy - - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], - non_blocking=True) + if _has_cp_helix: + # Owner-gated kv-length correction for overlap scheduling. + helix_prev_owned_cuda = torch.tensor( + helix_prev_group_owns, + dtype=torch.int, + pin_memory=prefer_pinned()).to(device='cuda', + non_blocking=True) + self.previous_kv_lens_offsets_cuda[ + num_extend_reqeust_wo_dummy - + previous_batch_len:num_extend_reqeust_wo_dummy].copy_( + helix_prev_owned_cuda * + new_tokens_lens_device[previous_slots], + non_blocking=True) + else: + self.previous_kv_lens_offsets_cuda[ + num_extend_reqeust_wo_dummy - + previous_batch_len:num_extend_reqeust_wo_dummy].copy_( + kv_len_offsets_device[previous_slots], + non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 797f2fd48666..cfd3de07be5d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -23,7 +23,7 @@ _llguidance_tokenizer_info, _xgrammar_tokenizer_info) from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping +from tensorrt_llm.mapping import CpType, Mapping from tensorrt_llm.quantization import QuantAlgo from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator @@ -415,6 +415,31 @@ def create_py_executor( f"decoding requires multiple tokens per sequence. Please use 'TRTLLM' attention " f"backend instead by setting attn_backend='TRTLLM'.") + # Helix CP + spec decode: overlap only for MTP-Eagle/Eagle3 on TRTLLM backend. + # Ownership stays host-deterministic; accept-dependent metadata is fixed on-device. + if (not llm_args.disable_overlap_scheduler + and llm_args.context_parallel_size > 1 + and llm_args.cp_config is not None + and llm_args.cp_config.cp_type == CpType.HELIX): + spec_mode = spec_config.spec_dec_mode + helix_overlap_supported = (llm_args.attn_backend == "TRTLLM" + and (spec_mode.is_mtp_eagle_one_model() + or spec_mode.is_eagle3_one_model())) + if not helix_overlap_supported: + logger.warning( + "Disabling overlap scheduler: Helix context parallelism with " + f"speculative-decoding mode '{spec_mode.name}' on attention " + f"backend '{llm_args.attn_backend}' does not support the " + "overlap scheduler. Only one-model MTP-Eagle / Eagle3 on the " + "TRTLLM backend are supported. Running with " + "disable_overlap_scheduler=True.") + llm_args.disable_overlap_scheduler = True + else: + logger.info( + "Enabling overlap scheduler for Helix context parallelism " + f"with one-model speculative-decoding mode '{spec_mode.name}'." + ) + if mm_encoder_only: llm_args.mm_encoder_only = True llm_args.disable_overlap_scheduler = True diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index e7da86608153..61ecaf3b1564 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -370,6 +370,8 @@ def merge_helix_requests( position_ids=position_ids_this_rank, ) req.total_input_len_cp = input_len + # Per-rank cached KV length; seeded with this rank's context shard and + # advanced at rewind (see KVCacheManager._helix_rewind_generation_kv). req.seqlen_this_rank_cp = len(input_ids_this_rank) req_with_children.append(req) if req.child_requests: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 486c43e8e269..85ffd488f348 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -735,6 +735,65 @@ def _context_seq_len(self, req: LlmRequest, is_cross: bool, return None return req.prompt_len + def _helix_owns_decode_group(self, group_index: int) -> bool: + """Return whether this CP rank owns the verify group at group_index. + + Owner = (group_index // tokens_per_block) % cp_size. + """ + block_id = group_index // self.tokens_per_block + return block_id % self.mapping.cp_size == self.mapping.cp_rank + + def _helix_prepare_generation_kv(self, req: LlmRequest, + draft_len: int) -> None: + """Reserve KV cache slots for a Helix generation or verify forward. + + The owning rank reserves slots for the whole verify group plus slack. + Ownership is a deterministic function of the reserve-side group index, so + the matching rewind recomputes it from its own FIFO counter rather than a + stored per-group decision. + """ + group_index = req.py_helix_decode_group_index + owns_group = self._helix_owns_decode_group(group_index) + # Prior group's ownership (deterministic) gates the on-device kv-length + # correction applied for the overlapped previous group. + req.py_helix_prev_group_owns = ( + self._helix_owns_decode_group(group_index - + 1) if group_index > 0 else False) + req.py_helix_decode_group_index = group_index + 1 + + # Owner reserves the whole group plus slack; non-owner reserves none. + reserve = max(draft_len, self._kv_reserve_draft_tokens) + n_reserve = 1 + reserve + if owns_group: + for _ in range(n_reserve): + self.impl.add_token(req.py_request_id) + + # Per-request inactive flag for verify-path input prep. + req.py_helix_is_inactive_rank = not owns_group + + def _helix_rewind_generation_kv(self, req: LlmRequest) -> None: + """Rewind rejected and slack draft KV, then advance decode lengths.""" + g = req.py_helix_global_decode_len + accepted = req.py_num_accepted_draft_tokens + runtime_draft_len = req.py_rewind_len + accepted + reserve = max(runtime_draft_len, self._kv_reserve_draft_tokens) + # Ownership for this group in FIFO order, recomputed deterministically to + # match the reserve-time decision (correct for any overlap depth). + owns_group = self._helix_owns_decode_group( + req.py_helix_rewind_group_index) + req.py_helix_rewind_group_index += 1 + + # Only the owner rank rewinds and grows its per-rank KV length. + rewind_count = reserve - accepted + if owns_group: + if rewind_count > 0: + self.rewind_kv_cache(req, rewind_count) + # Golden plus accepted drafts (1 + accepted). + req.seqlen_this_rank_cp += 1 + accepted + + # Advance the committed (global) decode length. + req.py_helix_global_decode_len = g + 1 + accepted + def prepare_resources(self, scheduled_batch: ScheduledRequests): # Cross/encoder K/V is allocated once and never grows; handle it on a # dedicated path so the self-attention flow below stays unconditional. @@ -769,18 +828,10 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): req, block_ids) for req in scheduled_batch.generation_requests: - if self.mapping.has_cp_helix(): - # Distribute the decode blocks across CP ranks in a round-robin manner. - decode_block_id = (req.py_decoding_iter - - 1) // self.tokens_per_block - if decode_block_id % self.mapping.cp_size == self.mapping.cp_rank: - req.py_helix_is_inactive_rank = False - req.seqlen_this_rank_cp += 1 - else: - req.py_helix_is_inactive_rank = True - # Skip allocating KV cache at decode for inactive helix ranks. - continue draft_len = get_draft_token_length(req) + if self.mapping.has_cp_helix(): + self._helix_prepare_generation_kv(req, draft_len) + continue self.impl.add_token(req.py_request_id) for _ in range(max(draft_len, self._kv_reserve_draft_tokens)): self.impl.add_token(req.py_request_id) @@ -966,6 +1017,14 @@ def add_dummy_requests( req.seqlen_this_rank_cp = req.prompt_len req.total_input_len_cp = token_num * self.mapping.cp_size - 1 req.py_decoding_iter = 1 + # Initialize Helix speculative-decode bookkeeping; per-rank + # cached length lives in seqlen_this_rank_cp (set above). + # global_decode_len starts at 0, matching real requests. + req.py_helix_global_decode_len = 0 + # Fresh generation request: reset the reserve/rewind counters. + req.py_helix_decode_group_index = 0 + req.py_helix_rewind_group_index = 0 + req.py_helix_prev_group_owns = False req.py_draft_tokens = [1] * max_num_draft_tokens if prepare_resource: for _ in range(_kv_draft): @@ -1003,6 +1062,10 @@ def update_resources(self, if request.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT): continue + if self.mapping.has_cp_helix(): + # Rewind owned rejected and slack tokens; advance decode length. + self._helix_rewind_generation_kv(request) + continue if request.py_rewind_len > 0: self.rewind_kv_cache(request, request.py_rewind_len) # Symmetric companion to prepare_resources's reserve_slack diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 90979682943f..f8ab160ddd8e 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -623,8 +623,46 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): else: self._saved_generation_lengths = None + # Save Helix buffers overwritten by the draft loop (CUDA graph capture). + self._saved_helix_position_offsets = None + self._saved_helix_is_inactive_rank = None + self._saved_kv_lens_cuda = None + # Draft loop disables verify-flatten; restore after loop for CUDA graph capture. + self._saved_helix_gen_flatten_active = getattr( + attn_metadata, '_helix_gen_flatten_active', None) + if getattr(attn_metadata, 'helix_position_offsets', None) is not None: + self._saved_helix_position_offsets = attn_metadata.helix_position_offsets.clone( + ) + if getattr(attn_metadata, 'helix_is_inactive_rank', None) is not None: + self._saved_helix_is_inactive_rank = attn_metadata.helix_is_inactive_rank.clone( + ) + # Save kv_lens_cuda grown per draft step under Helix. + if getattr(attn_metadata, 'kv_lens_cuda', None) is not None: + self._saved_kv_lens_cuda = attn_metadata.kv_lens_cuda[: + batch_size].clone( + ) + def _restore_attn_metadata_from_spec_dec(self, attn_metadata): super()._restore_attn_metadata_from_spec_dec(attn_metadata) + if self._saved_helix_position_offsets is not None: + attn_metadata.helix_position_offsets.copy_( + self._saved_helix_position_offsets) + self._saved_helix_position_offsets = None + if self._saved_helix_is_inactive_rank is not None: + attn_metadata.helix_is_inactive_rank.copy_( + self._saved_helix_is_inactive_rank) + self._saved_helix_is_inactive_rank = None + if self._saved_kv_lens_cuda is not None: + batch_size = self._saved_kv_lens_cuda.shape[0] + attn_metadata.kv_lens_cuda[:batch_size].copy_( + self._saved_kv_lens_cuda) + self._saved_kv_lens_cuda = None + # Restore verify-flatten flag after the draft loop. + if self._saved_helix_gen_flatten_active is not None: + attn_metadata._helix_gen_flatten_active = ( + self._saved_helix_gen_flatten_active) + self._saved_helix_gen_flatten_active = None + if self._saved_packed_mask is not None: batch_size = self._saved_packed_mask.shape[0] attn_metadata.spec_decoding_packed_mask[:batch_size].copy_( @@ -643,6 +681,18 @@ def _restore_attn_metadata_from_spec_dec(self, attn_metadata): self._saved_generation_lengths) self._saved_generation_lengths = None + def _helix_draft_owner_mask(self, attn_metadata, position_ids, batch_size): + """Per-request owner mask for a draft-loop step under Helix. + + Reuses verify-time helix_is_inactive_rank; only RoPE positions refresh. + Returns shape [batch_size], or None when Helix is inactive. + """ + if self.mapping is None or not self.mapping.has_cp_helix(): + return None + pos = position_ids[:batch_size].to(torch.int32).reshape(-1) + attn_metadata.helix_position_offsets[:batch_size].copy_(pos) + return ~attn_metadata.helix_is_inactive_rank[:batch_size] + # Skip torch.compile for now since current Torch is not compatible with Triton 3.4 # @torch.compile(options={"max-autotune": True}) @@ -778,6 +828,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i in range(runtime_draft_len): + # Set per-token global positions and ownership before draft forward. + helix_owner_mask = self._helix_draft_owner_mask( + attn_metadata, inputs["position_ids"], batch_size) + # Run draft model (mode-specific via helper). The helper # passes ``all_rank_num_tokens`` as a kwarg so the draft model # handles save/restore internally (Eagle3DraftModel.forward @@ -878,6 +932,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, spec_metadata, batch_size, draft_step=i) + next_draft_tokens.append(new_draft_token) # Update hidden states for the next iteration. @@ -897,6 +952,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, attn_metadata._seq_lens[:batch_size].fill_(1) attn_metadata._seq_lens_cuda[:batch_size].fill_(1) attn_metadata.on_update() + # Verify-flatten applies only to draft step 0 (full verify input). + if getattr(attn_metadata, "_helix_gen_flatten_active", + False): + attn_metadata._helix_gen_flatten_active = False has_kv_cache = inputs[ "attn_metadata"].kv_cache_manager is not None if has_kv_cache: @@ -904,10 +963,17 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, num_contexts].fill_(1) attn_metadata.num_contexts = 0 if hasattr(attn_metadata, 'kv_lens_cuda'): - attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( - runtime_draft_len - - num_accepted_tokens[num_contexts:]) - attn_metadata.kv_lens_cuda[:num_contexts] += 1 + if helix_owner_mask is not None: + # Append new draft token on the owner rank only. + attn_metadata.kv_lens_cuda[:batch_size] += ( + helix_owner_mask.to( + attn_metadata.kv_lens_cuda.dtype)) + else: + attn_metadata.kv_lens_cuda[ + num_contexts:batch_size] -= ( + runtime_draft_len - + num_accepted_tokens[num_contexts:]) + attn_metadata.kv_lens_cuda[:num_contexts] += 1 if has_kv_cache: self._prepare_flash_mla_generation_layout( @@ -923,7 +989,13 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, attn_metadata.use_spec_decoding = False else: if hasattr(attn_metadata, 'kv_lens_cuda'): - attn_metadata.kv_lens_cuda[:batch_size] += 1 + if helix_owner_mask is not None: + # Only the owner rank grows kv_lens. + attn_metadata.kv_lens_cuda[:batch_size] += ( + helix_owner_mask.to( + attn_metadata.kv_lens_cuda.dtype)) + else: + attn_metadata.kv_lens_cuda[:batch_size] += 1 attn_metadata.update_for_spec_dec() inputs = { @@ -982,6 +1054,9 @@ def _run_draft_forward(self, draft_model, inputs, spec_metadata, spec_metadata, step_idx) if self.is_mtp_eagle: + assert len( + draft_model.mtp_layers + ) == 1, f"expect only one MTP layer, found {len(draft_model.mtp_layers)} instead." hidden_states = draft_model.mtp_layers[0]( embed_tokens=draft_model.embed_tokens, all_rank_num_tokens=all_rank_num_tokens, @@ -1014,7 +1089,7 @@ def _get_local_max_and_combined(self, logits, mapping_lm_tp=None): local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) vocab_per_rank = logits.shape[-1] mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else \ - self.model_config.mapping + self.sampler_mapping max_index_per_rank = local_argmax.type( torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) max_index_per_rank_float = max_index_per_rank.float() @@ -1042,18 +1117,19 @@ def draft_sampler( Falls back to simple argmax when no tensor parallelism is active or when only attention DP is enabled without LM-head TP. + + Under Helix CP the vocab is sharded over the repurposed CP-to-TP group, + so sampler_mapping (not model_config.mapping) is the group to reduce + over, exactly as plain TP would. """ - if (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): + mapping = self.sampler_mapping + if (mapping is not None and mapping.tp_size > 1 + and not mapping.enable_attention_dp): combined = self._get_local_max_and_combined(logits) - gathered = allgather(combined, self.model_config.mapping, dim=-1) + gathered = allgather(combined, mapping, dim=-1) return self._get_draft_tokens_from_gathered(gathered) - elif (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and self.model_config.mapping.enable_lm_head_tp_in_adp): + elif (mapping is not None and mapping.tp_size > 1 + and mapping.enable_lm_head_tp_in_adp): combined = self._get_local_max_and_combined(logits, mapping_lm_head_tp) gathered = allgather(combined, mapping_lm_head_tp, dim=-1) @@ -1221,10 +1297,11 @@ def draft_decoder( # d2t-aware argmax. (Routing ADP/LM-head-TP through draft_sampler # without its mapping_lm_head_tp arg hits the None-mapping branch # and crashes with 'NoneType has no attribute tp_group'.) - if (self.is_mtp_eagle and self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): + # Vocab-shard check after attention (CP ranks folded into TP). + mapping = self.sampler_mapping + if (self.is_mtp_eagle and mapping is not None + and mapping.tp_size > 1 + and not mapping.enable_attention_dp): return self.draft_sampler(logits) return self._draft_sampler_greedy(logits, d2t) # Non-greedy (advanced) draft sampling has the same TP hazard as the @@ -1237,11 +1314,10 @@ def draft_decoder( # shared seed. (Greedy uses draft_sampler()'s lighter max+index gather; # random sampling needs the full distribution. The LM-head-TP-in-ADP # case is handled upstream and must not be gathered again here.) - if (self.is_mtp_eagle and self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size > 1 - and not self.model_config.mapping.enable_attention_dp): - logits = allgather(logits, self.model_config.mapping, dim=-1) + mapping = self.sampler_mapping + if (self.is_mtp_eagle and mapping is not None and mapping.tp_size > 1 + and not mapping.enable_attention_dp): + logits = allgather(logits, mapping, dim=-1) if spec_metadata.use_rejection_sampling and draft_step is not None: return self._draft_sampler_advanced_for_rejection( logits, spec_metadata, batch_size, d2t, draft_step) @@ -1310,10 +1386,11 @@ class MTPEagleWorker(Eagle3OneModelWorker): def __init__(self, spec_config, model_config: Optional[ModelConfig] = None, + mapping: Optional[Mapping] = None, use_separate_draft_kv_cache: bool = False): super().__init__( spec_config, - mapping=None, + mapping=mapping, model_config=model_config, use_separate_draft_kv_cache=use_separate_draft_kv_cache) # Preserved for callers/tests that still expect this attribute. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 061b36b6fc98..557b406c215d 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -35,6 +35,8 @@ ResourceManagerType) if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..pyexecutor.llm_request import LlmRequest @@ -871,6 +873,25 @@ def __init__(self, use_separate_draft_kv_cache: bool = False): # seed/offset pattern in `_sample_tokens_for_batch`). self._force_accept_rng_pool: Optional[torch.Tensor] = None self._force_accept_rng_counter: Optional[torch.Tensor] = None + # Repurposed CP-to-TP mapping for draft-token sampling under Helix CP. + self._sampler_mapping_cache: Optional["Mapping"] = None + + @property + def sampler_mapping(self) -> Optional["Mapping"]: + """Mapping for vocab-parallel draft-token sampling. + + Under Helix CP, ranks are repurposed to TP past attention. Returns + model_config.mapping when Helix CP is inactive. + """ + model_config = getattr(self, "model_config", None) + if model_config is None or not hasattr(model_config, "mapping"): + return None + mapping = model_config.mapping + if mapping is None or not mapping.has_cp_helix(): + return mapping + if self._sampler_mapping_cache is None: + self._sampler_mapping_cache = mapping.repurpose_helix_cp_to_tp() + return self._sampler_mapping_cache @property @abstractmethod diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 6afed0fc02e9..eb48966b1b0a 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -1098,7 +1098,7 @@ def get_local_max_and_combined(self, logits, mapping_lm_tp=None): local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) # Adjust indices based on TP rank and size vocab_per_rank = logits.shape[-1] - mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.model_config.mapping + mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else self.sampler_mapping max_index_per_rank = local_argmax.type( torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) # Use torch.stack and flatten instead of view+cat to avoid torch.compile issues @@ -1142,18 +1142,19 @@ def draft_sampler( draft_tokens: torch.Tensor [batch_size * max_draft_len] Draft token ids. Flattened. + + Under Helix CP the vocab is sharded over the repurposed CP-to-TP group, + so sampler_mapping (not model_config.mapping) is the group to reduce + over, exactly as plain TP would. ''' - if (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size - > 1) and not (self.model_config.mapping.enable_attention_dp): + mapping = self.sampler_mapping + if (mapping is not None + and mapping.tp_size > 1) and not (mapping.enable_attention_dp): combined = self.get_local_max_and_combined(logits) - gathered = allgather(combined, self.model_config.mapping, dim=-1) + gathered = allgather(combined, mapping, dim=-1) draft_tokens = self.get_draft_tokens_from_gathered(gathered) - elif (self.model_config is not None - and hasattr(self.model_config, 'mapping') - and self.model_config.mapping.tp_size - > 1) and self.model_config.mapping.enable_lm_head_tp_in_adp: + elif (mapping is not None + and mapping.tp_size > 1) and mapping.enable_lm_head_tp_in_adp: # For ADP + LM head TP mode, we need to find the global argmax across all TP ranks combined = self.get_local_max_and_combined(logits, mapping_lm_head_tp) diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 164d6758ed4a..1a24514dbcb4 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -389,7 +389,7 @@ def get_spec_worker(spec_config, if spec_dec_mode.is_mtp_vanilla(): return MTPWorker(spec_config, model_config, use_separate_draft_kv_cache) if spec_dec_mode.is_mtp_eagle_one_model(): - return MTPEagleWorker(spec_config, model_config, + return MTPEagleWorker(spec_config, model_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_eagle3_one_model(): if _is_effective_dynamic_tree(spec_config): diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 7e5a21a29207..f2ea0805a615 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -216,7 +216,31 @@ def evaluate(self, is_integration_test = is_integration_test or os.getenv( 'INTEGRATION_TEST', '0') == '1' - if is_integration_test: + # Optional GSM8K debug env vars: TRTLLM_GSM8K_NUM_SAMPLES, TRTLLM_GSM8K_OUTPUT_DIR. + # A custom sample count skips the accuracy hypothesis test. + gsm8k_num_samples = None + gsm8k_output_dir = None + if self.DATASET == "gsm8k": + gsm8k_output_dir = os.getenv("TRTLLM_GSM8K_OUTPUT_DIR") or None + num_samples_env = os.getenv("TRTLLM_GSM8K_NUM_SAMPLES") + if num_samples_env is not None: + gsm8k_num_samples = int(num_samples_env) + if gsm8k_num_samples <= 0: + raise ValueError( + "TRTLLM_GSM8K_NUM_SAMPLES must be a positive integer, " + f"got {num_samples_env!r}.") + + if gsm8k_num_samples is not None: + logger.info( + f"Running GSM8K on a deterministic subset of {gsm8k_num_samples} " + "sample(s) (TRTLLM_GSM8K_NUM_SAMPLES) and skipping accuracy " + "verification.") + hypothesis_testing_params = HypothesisTestingParams( + ref_accuracy=0 if self.HIGHER_IS_BETTER else math.inf, + num_samples=gsm8k_num_samples, + metric_name=self.METRIC_NAME, + higher_is_better=self.HIGHER_IS_BETTER) + elif is_integration_test: logger.info( "Running in INTEGRATION_TEST mode: using only 1 sample and skipping accuracy verification" ) @@ -248,6 +272,12 @@ def evaluate(self, evaluator_kwargs.update(self.EVALUATOR_KWARGS) if extra_evaluator_kwargs is not None: evaluator_kwargs.update(extra_evaluator_kwargs) + if gsm8k_output_dir is not None: + evaluator_kwargs["output_dir"] = gsm8k_output_dir + logger.info( + "Dumping GSM8K inference inputs/outputs to " + f"{os.path.realpath(gsm8k_output_dir)} (TRTLLM_GSM8K_OUTPUT_DIR)." + ) evaluator = self.EVALUATOR_CLS( num_samples=hypothesis_testing_params.num_samples, **evaluator_kwargs) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 1f83df9610f5..cb3606d63e11 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -1195,6 +1195,7 @@ def test_auto_dtype(self, overlap_scheduler, mtp_nextn): @skip_pre_blackwell @pytest.mark.skip_less_device(8) + @parametrize_with_ids("mtp_nextn", [0, 3]) @pytest.mark.parametrize( "gen_pp,gen_tp,gen_cp,enable_attention_dp", [ (1, 1, 4, False), @@ -1219,8 +1220,9 @@ def test_auto_dtype(self, overlap_scheduler, mtp_nextn): "cudagraph:with_padding" ]) @pytest.mark.parametrize("comms_medium", ["fifo_v1", "fifo_v2", "nccl"]) - def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, - gen_pp, gen_tp, gen_cp, enable_attention_dp): + def test_auto_dtype_with_helix(self, mtp_nextn, comms_medium, + cuda_graph_config, gen_pp, gen_tp, gen_cp, + enable_attention_dp): # Parse comms_medium to get use_nccl_for_alltoall and fifo_version. if comms_medium == "nccl": use_nccl_for_alltoall = True @@ -1264,7 +1266,7 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, "use_nccl_for_alltoall": use_nccl_for_alltoall, "fifo_version": fifo_version, }, - "disable_overlap_scheduler": True, + "disable_overlap_scheduler": mtp_nextn == 0, "kv_cache_config": kv_cache_config, "enable_chunked_prefill": False, "cuda_graph_config": cuda_graph_config, @@ -1274,6 +1276,16 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, }, "enable_attention_dp": enable_attention_dp, } + # Enable MTP on both servers so ctx and gen agree on spec-decode layout. + if mtp_nextn > 0: + ctx_server_config["speculative_config"] = { + "decoding_type": "MTP", + "max_draft_len": mtp_nextn, + } + gen_server_config["speculative_config"] = { + "decoding_type": "MTP", + "max_draft_len": mtp_nextn, + } disaggregated_server_config = { "hostname": "localhost", "backend": "pytorch", @@ -1287,7 +1299,7 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, with launch_disaggregated_llm(disaggregated_server_config, ctx_server_config, gen_server_config, self.MODEL_PATH) as llm: - run_accuracy_test(llm, self.MODEL_NAME, ["MMLU", "GSM8K"]) + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) @pytest.mark.skip_less_device(2) @pytest.mark.skip_less_device_memory(60000) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 03a6f7a007e0..6c7a8a218e42 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -7,18 +7,18 @@ accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_ accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=0-overlap_scheduler=True] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1dp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp2tp1cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp2tp1cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1dp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] -accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1dp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp2tp1cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp2tp1cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1dp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] +accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2-mtp_nextn=0] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_sync accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[llguidance-mtp_nextn=0] accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_guided_decoding[llguidance-mtp_nextn=2] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 3997ab6b6197..ee8d91de40a2 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -142,8 +142,8 @@ l0_dgx_b200: backend: pytorch orchestrator: mpi tests: - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp2tp1cp2] TIMEOUT (60) - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp2tp1cp2-mtp_nextn=0] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1dp2cp2-mtp_nextn=0] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype_with_helix[fifo-cudagraph:with_padding-pp1tp1cp4] TIMEOUT (60) @@ -225,8 +225,9 @@ l0_dgx_b200: backend: pytorch orchestrator: mpi tests: - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp2cp2] TIMEOUT (60) - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] TIMEOUT (60) + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=3] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp2tp1cp2] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1dp2cp2] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestDeepSeekV32Exp::test_auto_dtype_with_helix[fifo-cudagraph:with_padding-pp1dp2cp2] TIMEOUT (60) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d6fddb1d0f83..570355bc6794 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -164,11 +164,11 @@ full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[in full:B200/perf/test_perf.py::test_perf[quant:int8_sq_per_tensor] SKIP (https://nvbugs/5161074) full:B200/perf/test_perf.py::test_perf[quant:int8_sq_per_token_channel] SKIP (https://nvbugs/5161074) full:B200/perf/test_perf.py::test_perf[quant:w4a8_awq] SKIP (https://nvbugs/5161074) -full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4] SKIP (https://nvbugs/6410881) -full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4] SKIP (https://nvbugs/6410881) -full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2] SKIP (https://nvbugs/6410881) -full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4] SKIP (https://nvbugs/6410881) -full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] SKIP (https://nvbugs/6410881) +full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] SKIP (https://nvbugs/6410881) +full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] SKIP (https://nvbugs/6410881) +full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] SKIP (https://nvbugs/6410881) +full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] SKIP (https://nvbugs/6410881) +full:B300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] SKIP (https://nvbugs/6410881) full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[baseline_pp4_mtp1] SKIP (https://nvbugs/6423845) full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus[latency] SKIP (https://nvbugs/6423866) full:B300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True-enable_chunked_prefill=True-v2_kv_cache=True] SKIP (https://nvbugs/6422343) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py b/tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py new file mode 100644 index 000000000000..23930e058ca4 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Hardware-agnostic tests for Helix + Eagle3 KV ownership and positions. + +Guards that each verify group is owned by exactly one CP rank, that global +positions and inactive flags are consistent, and that reserve/rewind FIFO +ownership matches under overlap scheduling. +""" + +import pytest + + +def owns_decode_index(decode_index: int, tpb: int, cp_size: int, cp_rank: int) -> bool: + return (decode_index // tpb) % cp_size == cp_rank + + +@pytest.mark.parametrize("tpb", [1, 4, 32]) +@pytest.mark.parametrize("cp_size", [2, 4, 8]) +@pytest.mark.parametrize("g", [0, 1, 5, 31, 32, 33, 64, 100, 257]) +def test_verify_group_owned_by_exactly_one_rank(tpb, cp_size, g): + # Exactly one rank owns decode-index g. + owners = [cp_rank for cp_rank in range(cp_size) if owns_decode_index(g, tpb, cp_size, cp_rank)] + assert owners == [(g // tpb) % cp_size] + + +def verify_token_params( + total_input_len: int, g: int, num_draft: int, tpb: int, cp_size: int, cp_rank: int +) -> tuple[list[int], list[bool], int]: + first_decode_index = g + first_pos = total_input_len + first_decode_index + global_positions = list(range(first_pos, first_pos + 1 + num_draft)) + owns_group = owns_decode_index(g, tpb, cp_size, cp_rank) + inactive_flags = [not owns_group] * (1 + num_draft) + num_active = (1 + num_draft) if owns_group else 0 + return global_positions, inactive_flags, num_active + + +@pytest.mark.parametrize("tpb", [4, 32]) +@pytest.mark.parametrize("cp_size", [2, 4]) +@pytest.mark.parametrize("num_draft", [0, 1, 3]) +@pytest.mark.parametrize("g", [0, 1, 5, 32]) +def test_verify_token_params_consistency(tpb, cp_size, num_draft, g): + total_input_len = 100 + group_size = 1 + num_draft + # Expected owner from anchor decode-index g. + expected_owner = (g // tpb) % cp_size + active_ranks = [] + for cp_rank in range(cp_size): + positions, inactive, num_active = verify_token_params( + total_input_len, g, num_draft, tpb, cp_size, cp_rank + ) + # Global positions start at total_input_len + g. + assert positions == list(range(positions[0], positions[0] + group_size)) + assert positions[0] == total_input_len + g + # One inactive flag per request, not per token. + assert len(set(inactive)) == 1 + # num_active counts owned query tokens. + assert num_active == sum(1 for f in inactive if not f) + if cp_rank == expected_owner: + assert all(not f for f in inactive) + assert num_active == group_size + active_ranks.append(cp_rank) + else: + assert all(f for f in inactive) + assert num_active == 0 + # No mixed ownership across ranks. + assert active_ranks == [expected_owner] + + +@pytest.mark.parametrize( + "tpb,cp_size,g,num_draft", + [ + # Per-token indices would straddle a block boundary; group stays on one rank. + (4, 2, 3, 3), + (32, 2, 30, 3), + (4, 4, 3, 3), + # Groups within one block. + (4, 2, 5, 3), + (32, 2, 1, 3), + ], +) +def test_verify_token_params_no_mixed_ownership(tpb, cp_size, g, num_draft): + group_size = 1 + num_draft + expected_owner = (g // tpb) % cp_size + per_rank_active = [] + for cp_rank in range(cp_size): + _, inactive, num_active = verify_token_params(100, g, num_draft, tpb, cp_size, cp_rank) + per_rank_active.append(num_active) + # Uniform inactive flag within the request. + assert len(set(inactive)) == 1 + if cp_rank == expected_owner: + assert num_active == group_size + else: + assert num_active == 0 + # Whole group on one rank, even across block boundaries. + ranks_with_writes = sum(1 for a in per_rank_active if a > 0) + assert ranks_with_writes == 1 + assert per_rank_active[expected_owner] == group_size + + +# --------------------------------------------------------------------------- +# Overlap scheduler: deterministic verify-group ownership + reserve/rewind FIFO +# --------------------------------------------------------------------------- + + +owns_decode_group = owns_decode_index + + +@pytest.mark.parametrize("tpb", [1, 4, 32]) +@pytest.mark.parametrize("cp_size", [2, 4, 8]) +@pytest.mark.parametrize("group_index", [0, 1, 5, 31, 32, 33, 64, 100, 257]) +def test_decode_group_owned_by_exactly_one_rank(tpb, cp_size, group_index): + owners = [ + cp_rank + for cp_rank in range(cp_size) + if owns_decode_group(group_index, tpb, cp_size, cp_rank) + ] + assert owners == [(group_index // tpb) % cp_size] + + +@pytest.mark.parametrize("tpb", [1, 4, 32]) +@pytest.mark.parametrize("cp_size", [2, 4]) +def test_decode_group_ownership_is_balanced(tpb, cp_size): + # Balanced ownership over cp_size * tpb groups. + counts = [0] * cp_size + for group_index in range(cp_size * tpb): + counts[(group_index // tpb) % cp_size] += 1 + assert counts == [tpb] * cp_size + + +class _FifoRankState: + """Per-rank reserve and rewind state.""" + + def __init__(self, tpb: int, cp_size: int, cp_rank: int) -> None: + self.tpb = tpb + self.cp_size = cp_size + self.cp_rank = cp_rank + self.group_index: int = 0 + self.pending: list[bool] = [] + self.reserve_log: list[bool] = [] + self.rewind_log: list[bool] = [] + + def reserve(self) -> bool: + owns = owns_decode_group(self.group_index, self.tpb, self.cp_size, self.cp_rank) + self.group_index += 1 + self.pending.append(owns) + self.reserve_log.append(owns) + return owns + + def rewind(self) -> bool: + owns = self.pending.pop(0) + self.rewind_log.append(owns) + return owns + + +@pytest.mark.parametrize("tpb", [1, 4]) +@pytest.mark.parametrize("cp_size", [2, 4]) +@pytest.mark.parametrize("pipeline_depth", [0, 1]) +@pytest.mark.parametrize("num_iters", [1, 2, 5, 40]) +def test_reserve_rewind_fifo_consistency(tpb, cp_size, pipeline_depth, num_iters): + """Reserve and rewind must agree on ownership for every group.""" + for cp_rank in range(cp_size): + state = _FifoRankState(tpb, cp_size, cp_rank) + # Prime overlap pipeline before first rewind. + for _ in range(pipeline_depth): + state.reserve() + for _ in range(num_iters): + state.reserve() + state.rewind() + # Drain in-flight reserves. + while state.pending: + state.rewind() + + # Rewind log must match reserve log in FIFO order. + assert state.rewind_log == state.reserve_log + # Sequence matches deterministic group-ownership formula. + expected = [ + owns_decode_group(i, tpb, cp_size, cp_rank) for i in range(len(state.reserve_log)) + ] + assert state.reserve_log == expected + + +@pytest.mark.parametrize("tpb", [1, 4]) +@pytest.mark.parametrize("cp_size", [2, 4]) +@pytest.mark.parametrize("num_iters", [8, 40]) +def test_each_group_owned_by_exactly_one_rank_across_ranks(tpb, cp_size, num_iters): + # Each group index is owned by exactly one rank. + for group_index in range(num_iters): + owners = [r for r in range(cp_size) if owns_decode_group(group_index, tpb, cp_size, r)] + assert len(owners) == 1