From 1b9e4ba32ef7813f3ec485aed90e3844c39354ee Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Mon, 2 Mar 2026 14:47:55 -0500 Subject: [PATCH 1/4] Support different TP config for draft model Signed-off-by: Mike Iovine --- .../_torch/models/modeling_deepseekv3.py | 18 +- .../_torch/models/modeling_speculative.py | 46 +++- tensorrt_llm/_torch/modules/attention.py | 88 +++--- tensorrt_llm/_torch/pyexecutor/_util.py | 7 +- .../_torch/pyexecutor/model_loader.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 21 ++ tensorrt_llm/mapping.py | 256 +++++++++++++++++- 7 files changed, 382 insertions(+), 57 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index b6e62def933b..e0402fbeeafc 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -322,11 +322,15 @@ def split_kv_b_proj(kv_b_proj: torch.Tensor, v_head_dim = self.config.v_head_dim kv_lora_rank = self.config.kv_lora_rank - tp_rank = self.model_config.mapping.tp_rank - tp_size = self.model_config.mapping.tp_size + target_tp_rank = self.model_config.mapping.tp_rank + target_tp_size = self.model_config.mapping.tp_size + tp_rank = target_tp_rank + tp_size = target_tp_size cp_rank = self.model_config.mapping.cp_rank cp_size = self.model_config.mapping.cp_size + draft_mapping = getattr(self.model, '_draft_mapping', None) + params_map = {'gate_up_proj': ['gate_proj', 'up_proj']} all_named_modules = dict(self.model.named_modules()) @@ -342,6 +346,16 @@ def split_kv_b_proj(kv_b_proj: torch.Tensor, else: names = name.split('.') parent_module_name = '.'.join(names[:-1]) + is_mtp_layer = ("model.layers" in name and len(names) > 2 + and names[2].isdigit() + and int(names[2]) + >= self.config.num_hidden_layers) + if is_mtp_layer and draft_mapping is not None: + tp_rank = draft_mapping.tp_rank + tp_size = draft_mapping.tp_size + else: + tp_rank = target_tp_rank + tp_size = target_tp_size if "model.layers" in name and int( names[2]) >= self.config.num_hidden_layers: mtp_layer_idx = int( diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index f64febf02849..eabc3062d61a 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -7,6 +7,7 @@ from transformers import LlamaConfig, PretrainedConfig from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import create_draft_mapping from ...functional import PositionEmbeddingType from ..attention_backend import AttentionMetadata @@ -961,7 +962,11 @@ def forward(self, ) -def get_draft_model(model_config, draft_config, lm_head, model): +def get_draft_model(model_config, + draft_config, + lm_head, + model, + draft_model_config=None): assert getattr(model_config, 'spec_config', None) is not None spec_dec_mode = model_config.spec_config.spec_dec_mode if spec_dec_mode.is_eagle3_one_model(): @@ -979,7 +984,8 @@ def get_draft_model(model_config, draft_config, lm_head, model): ) elif spec_dec_mode.is_mtp_one_model(): - return MTPForCausalLM(model_config, + mtp_config = draft_model_config if draft_model_config is not None else model_config + return MTPForCausalLM(mtp_config, model_config.pretrained_config.num_hidden_layers, lm_head, model) elif spec_dec_mode.is_mtp_eagle(): @@ -1009,6 +1015,13 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): spec_config = getattr(model_config, 'spec_config', None) self.spec_config = spec_config if spec_config and spec_config.spec_dec_mode.use_one_engine(): + draft_tp = getattr(spec_config, 'draft_tp_size', None) + draft_mapping = create_draft_mapping(model_config.mapping, + draft_tp) + self._draft_mapping = (draft_mapping + if draft_mapping is not model_config.mapping + else None) + # Only create draft_model for modes MTP, Eagle3 (not SA) if not spec_config.spec_dec_mode.is_sa(): if spec_config.spec_dec_mode.is_eagle3_one_model(): @@ -1017,7 +1030,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): MistralConfigLoader self.draft_config = MistralConfigLoader().load( spec_config.speculative_model, - mapping=model_config.mapping, + mapping=draft_mapping, moe_backend=model_config.moe_backend, moe_max_num_tokens=model_config.moe_max_num_tokens, max_num_tokens=model_config.max_num_tokens, @@ -1031,7 +1044,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): trust_remote_code=True, attn_backend=model_config.attn_backend, moe_backend=model_config.moe_backend, - mapping=model_config.mapping, + mapping=draft_mapping, spec_config=model_config.spec_config, max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens) @@ -1048,7 +1061,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): trust_remote_code=True, attn_backend=model_config.attn_backend, moe_backend=model_config.moe_backend, - mapping=model_config.mapping, + mapping=draft_mapping, spec_config=None, # Avoid recursive spec-dec max_num_tokens=model_config.max_num_tokens, moe_max_num_tokens=model_config.moe_max_num_tokens) @@ -1058,23 +1071,30 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]): self.use_separate_draft_kv_cache = should_use_separate_draft_kv_cache( spec_config) + draft_model_config = None + if draft_mapping is not model_config.mapping: + draft_model_config = replace(model_config, + mapping=draft_mapping) + if self.draft_config is None: + self.draft_config = draft_model_config + self.draft_model = get_draft_model(model_config, self.draft_config, - self.lm_head, self.model) + self.lm_head, self.model, + draft_model_config) if self.draft_model is not None: self.epilogue.append(self.draft_model) if spec_config.spec_dec_mode.is_pard( ) and self.draft_model is not None: self.draft_model.logits_processor = self.logits_processor - # EAGLE3-specific logic: merge extra_attrs from draft model for Llama3 - if (self.draft_config is not None and model_config.spec_config. - spec_dec_mode.is_eagle3_one_model() - and model_config.spec_config.eagle3_model_arch == "llama3"): + if (self.draft_config is not None + and hasattr(self.draft_config, 'extra_attrs') + and self.draft_config is not model_config): for key, value in self.draft_config.extra_attrs.items(): - assert key in ('attn_layers', 'mla_layers') - assert key in model_config.extra_attrs - model_config.extra_attrs[key].update(value) + if key in ('attn_layers', 'mla_layers'): + model_config.extra_attrs.setdefault(key, {}).update( + value) # spec_worker is created for all one-engine modes (MTP, Eagle3, SA) self.spec_worker = get_spec_worker( diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 33aea6848765..74a867a7f582 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -438,16 +438,19 @@ def __init__( assert self.mapping.has_cp_helix( ), f"CP type must be HELIX for Attention, but got {self.mapping.cp_config['cp_type']}." - mapping = Mapping( - world_size=dp_size * tp_size * pp_size * cp_size, - tp_size=tp_size, - pp_size=pp_size * dp_size, - cp_size=cp_size, - cp_config=self.mapping.cp_config, - rank=self.mapping.rank, - gpus_per_node=self.mapping.gpus_per_node, - enable_attention_dp=self.mapping.enable_attention_dp, - ) + if dp_size == 1 and cp_size == 1: + mapping = self.mapping + else: + mapping = Mapping( + world_size=dp_size * tp_size * pp_size * cp_size, + tp_size=tp_size, + pp_size=pp_size * dp_size, + cp_size=cp_size, + cp_config=self.mapping.cp_config, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) self.tp_size = tp_size self.cp_size = cp_size self.tp_rank = mapping.tp_rank @@ -495,15 +498,18 @@ def __init__( # For Helix CP, combine TP and CP for the output projection so each # rank's o_proj input is num_heads_tp_cp * head_dim. - mapping_o = Mapping( - world_size=dp_size * tp_size * pp_size * cp_size, - tp_size=tp_size * cp_size, - pp_size=pp_size * dp_size, - cp_size=1, - rank=self.mapping.rank, - gpus_per_node=self.mapping.gpus_per_node, - enable_attention_dp=self.mapping.enable_attention_dp, - ) + if dp_size == 1 and cp_size == 1: + mapping_o = self.mapping + else: + mapping_o = Mapping( + world_size=dp_size * tp_size * pp_size * cp_size, + tp_size=tp_size * cp_size, + pp_size=pp_size * dp_size, + cp_size=1, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) self.mapping_o = mapping_o self.o_proj = Linear( @@ -1099,16 +1105,19 @@ def __init__( assert self.mapping.has_cp_helix( ), f"CP type must be HELIX for MLA, but got {self.mapping.cp_config['cp_type']}." - mapping = Mapping( - world_size=pp_size * dp_size * tp_size * cp_size, - tp_size=tp_size, - pp_size=pp_size * dp_size, - cp_size=cp_size, - cp_config=self.mapping.cp_config, - rank=self.mapping.rank, - gpus_per_node=self.mapping.gpus_per_node, - enable_attention_dp=self.mapping.enable_attention_dp, - ) + if dp_size == 1 and cp_size == 1: + mapping = self.mapping + else: + mapping = Mapping( + world_size=pp_size * dp_size * tp_size * cp_size, + tp_size=tp_size, + pp_size=pp_size * dp_size, + cp_size=cp_size, + cp_config=self.mapping.cp_config, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) assert self.num_heads % (tp_size * cp_size) == 0 self.num_heads_tp = self.num_heads // tp_size @@ -1204,15 +1213,18 @@ def __init__( requires_grad=False, ) - mapping_o = Mapping( - world_size=pp_size * dp_size * tp_size * cp_size, - tp_size=tp_size * cp_size, - pp_size=pp_size * dp_size, - cp_size=1, - rank=self.mapping.rank, - gpus_per_node=self.mapping.gpus_per_node, - enable_attention_dp=self.mapping.enable_attention_dp, - ) + if dp_size == 1 and cp_size == 1: + mapping_o = self.mapping + else: + mapping_o = Mapping( + world_size=pp_size * dp_size * tp_size * cp_size, + tp_size=tp_size * cp_size, + pp_size=pp_size * dp_size, + cp_size=1, + rank=self.mapping.rank, + gpus_per_node=self.mapping.gpus_per_node, + enable_attention_dp=self.mapping.enable_attention_dp, + ) self.mapping_o = mapping_o self.o_proj = Linear( self.num_key_value_heads * self.v_head_dim, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 544eea6daca9..0629d81a8921 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -141,9 +141,10 @@ def _get_kv_size_per_token(self): elif self._should_create_separate_draft_kv_cache(): # One-model draft with separate KV cache layout effective_draft_config = self._get_effective_draft_config() + draft_mapping = effective_draft_config.mapping kv_size_per_token += self._kv_cache_manager_cls.get_cache_size_per_token( effective_draft_config, - mapping, + draft_mapping, tokens_per_block=self._tokens_per_block) return kv_size_per_token @@ -637,11 +638,13 @@ def _create_one_model_draft_kv_cache_manager( "Falling back to KVCacheManager for draft model.") draft_kv_cache_manager_cls = KVCacheManager + draft_mapping = effective_draft_config.mapping + estimating_kv_cache = estimating_kv_cache and not self._skip_est 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=self._kv_cache_config, tokens_per_block=self._tokens_per_block, max_seq_len=self._max_seq_len, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 484556ac17cd..43e826969da0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -317,9 +317,10 @@ def init_meta_tensor(t: torch.Tensor): if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights( ): + draft_mapping = model.draft_config.mapping weights = checkpoint_loader.load_weights( self.spec_config.speculative_model, - mapping=self.mapping) + mapping=draft_mapping) draft_model_arch = model.draft_config.pretrained_config.architectures[ 0] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a446f22629ab..0333f49357a0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -703,6 +703,15 @@ class DecodingBaseConfig(StrictBaseModel): _allow_greedy_draft_tokens: bool = PrivateAttr(True) # Internal: record decoding_type alias used during parsing (for warnings). _decoding_type_alias: Optional[str] = PrivateAttr(default=None) + draft_tp_size: Optional[PositiveInt] = Field( + default=None, + description= + "Tensor parallelism size for the draft model in one-model speculative decoding. " + "When set, must divide the target model's TP size evenly. If None, the draft model " + "uses the same TP size as the target model. Smaller values reduce communication " + "overhead for the (typically small) draft model at the cost of redundant computation." + ) + # If set, drafting will use separate KV cache in one-model speculative decoding. _allow_separate_draft_kv_cache: bool = PrivateAttr(True) @@ -3294,6 +3303,18 @@ def validate_speculative_config(self): if self.backend == "_autodeploy": self.speculative_config._draft_target_one_model = False + draft_tp = self.speculative_config.draft_tp_size + if draft_tp is not None: + target_tp = self.parallel_config.tp_size + if draft_tp > target_tp: + raise ValueError( + f"draft_tp_size ({draft_tp}) must be <= target tp_size ({target_tp})." + ) + if target_tp % draft_tp != 0: + raise ValueError( + f"target tp_size ({target_tp}) must be divisible by " + f"draft_tp_size ({draft_tp}).") + else: self.decoding_config = None diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index 74d3e02ee7e7..4ae9185bc073 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -686,3 +686,257 @@ def __init__(self, *args, **kwargs): ), "DeviceMeshTopology is only available in Ray orchestrator mode." super().__init__(*args, **kwargs) + + +def _get_all_tp_groups(mapping: "Mapping") -> List[List[int]]: + """Return TP group rank-lists for every PP/CP stage in the cluster.""" + if hasattr(mapping, 'tp_groups') and mapping.tp_groups: + return mapping.tp_groups + tp_size = mapping.tp_size + pp_size = mapping.pp_size + cp_size = mapping.cp_size + groups = [] + for pp in range(pp_size): + for cp in range(cp_size): + start = pp * tp_size * cp_size + cp + end = (pp + 1) * tp_size * cp_size + cp + groups.append(list(range(start, end, cp_size))) + return groups + + +class DraftMapping: + """Mapping for draft models with smaller TP than the target model. + + In one-model speculative decoding, the draft model (typically much smaller) + can run with a smaller TP size than the target model. Ranks are partitioned + into groups of ``draft_tp_size``, each of which independently computes the + same draft tokens with reduced communication overhead. + + This object provides the subset of the Mapping interface that draft model + layers (Linear, Attention, etc.) and the weight loader need. + """ + + def __init__(self, target_mapping: "Mapping", draft_tp_size: int): + if draft_tp_size == target_mapping.tp_size: + raise ValueError( + "draft_tp_size equals target tp_size; use the target mapping directly." + ) + if target_mapping.tp_size % draft_tp_size != 0: + raise ValueError( + f"target tp_size ({target_mapping.tp_size}) must be divisible by " + f"draft_tp_size ({draft_tp_size}).") + + self._target = target_mapping + self._global_rank = target_mapping.rank + + self.tp_size = draft_tp_size + self.pp_size = target_mapping.pp_size + self.pp_rank = target_mapping.pp_rank + self.cp_size = 1 + self.cp_rank = 0 + self.gpus_per_node = target_mapping.gpus_per_node + + self.world_size = draft_tp_size * self.pp_size + self.rank = (self.pp_rank * draft_tp_size + + target_mapping.tp_rank % draft_tp_size) + self.enable_attention_dp = target_mapping.enable_attention_dp + self.enable_lm_head_tp_in_adp = target_mapping.enable_lm_head_tp_in_adp + + self.moe_tp_size = min(draft_tp_size, target_mapping.moe_tp_size) + self.moe_ep_size = 1 + self.moe_cluster_size = 1 + self.moe_tp_cluster_ep_size = self.moe_tp_size + + self.attn_tp_size = draft_tp_size + self.attn_cp_size = 1 + + target_tp_rank = target_mapping.tp_rank + target_tp_group = list(target_mapping.tp_group) + target_tp_size = target_mapping.tp_size + + self._tp_rank = target_tp_rank % draft_tp_size + + group_index = target_tp_rank // draft_tp_size + start = group_index * draft_tp_size + self._tp_group = target_tp_group[start:start + draft_tp_size] + + self._tp_group_pg = None + self._moe_tp_group = self._tp_group[:self.moe_tp_size] + + if mpi_disabled() and draft_tp_size > 1 and torch.distributed.is_initialized(): + all_tp_groups = _get_all_tp_groups(target_mapping) + num_sub = target_tp_size // draft_tp_size + for tp_grp in all_tp_groups: + for i in range(num_sub): + sub_ranks = tp_grp[i * draft_tp_size:(i + 1) * + draft_tp_size] + pg = torch.distributed.new_group(sub_ranks) + if target_mapping.rank in sub_ranks: + self._tp_group_pg = pg + elif mpi_disabled() and draft_tp_size == 1: + from tensorrt_llm._torch.device_mesh import SingleProcessGroup + self._tp_group_pg = SingleProcessGroup.get_group() + + self.pp_groups = target_mapping.pp_groups + self.pp_partition = getattr(target_mapping, 'pp_partition', None) + self.cp_config = {} + self.tp_groups = [] + self.cp_groups = [] + self.moe_cluster_groups = [] + self.moe_tp_groups = [] + self.moe_ep_groups = [] + + def __eq__(self, other): + if not isinstance(other, DraftMapping): + return NotImplemented + return (self.tp_size == other.tp_size and self.rank == other.rank + and self._tp_group == other._tp_group) + + def __hash__(self): + return hash((self.tp_size, self.rank, tuple(self._tp_group))) + + @property + def tp_rank(self): + return self._tp_rank + + @property + def tp_group(self): + return self._tp_group + + @property + def tp_group_pg(self): + if self._tp_group_pg is None: + from tensorrt_llm._torch.device_mesh import SingleProcessGroup + return SingleProcessGroup.get_group() + return self._tp_group_pg + + @property + def moe_tp_rank(self): + return self._tp_rank % self.moe_tp_size + + @property + def moe_ep_rank(self): + return 0 + + @property + def moe_cluster_rank(self): + return 0 + + @property + def moe_tp_group(self): + return self._moe_tp_group + + @property + def moe_tp_group_pg(self): + return self.tp_group_pg + + @property + def moe_ep_group(self): + return self._tp_group[:1] + + @property + def moe_ep_group_pg(self): + from tensorrt_llm._torch.device_mesh import SingleProcessGroup + return SingleProcessGroup.get_group() + + @property + def dp_size(self): + return self.tp_size if self.enable_attention_dp else 1 + + def has_tp(self): + return self.tp_size > 1 + + def has_pp(self): + return self.pp_size > 1 + + def has_cp(self): + return False + + def has_cp_ulysses(self): + return False + + def has_cp_helix(self): + return False + + def has_moe_tp(self): + return self.moe_tp_size > 1 + + def has_moe_ep(self): + return False + + def has_moe_cluster(self): + return False + + def is_first_pp_rank(self): + return self.pp_rank == 0 + + def is_last_pp_rank(self): + return self.pp_rank == self.pp_size - 1 + + @property + def node_rank(self): + return self._global_rank // self.gpus_per_node + + @property + def local_rank(self): + return self._global_rank % self.gpus_per_node + + def get_node_rank(self, rank: int): + return rank // self.gpus_per_node + + def get_local_rank(self, rank: int): + return rank % self.gpus_per_node + + def is_multi_node(self): + return self.world_size > self.gpus_per_node + + def is_second_last_pp_rank(self): + return self.pp_rank == self.pp_size - 2 + + def pp_layers(self, num_layers): + return self._target.pp_layers(num_layers) + + def pp_rank_of_layer(self, layer_idx: int, num_layers: int) -> int: + return self._target.pp_rank_of_layer(layer_idx, num_layers) + + def ep_experts(self, num_experts): + return list(range(num_experts)) + + def prev_pp_rank(self): + return self._target.prev_pp_rank() + + def next_pp_rank(self): + return self._target.next_pp_rank() + + def to_dict(self): + return { + 'world_size': self.world_size, + 'rank': self.rank, + 'gpus_per_node': self.gpus_per_node, + 'cp_size': self.cp_size, + 'tp_size': self.tp_size, + 'pp_size': self.pp_size, + 'moe_tp_size': self.moe_tp_size, + 'moe_cluster_size': self.moe_cluster_size, + 'moe_ep_size': self.moe_ep_size, + 'enable_attention_dp': self.enable_attention_dp, + 'enable_lm_head_tp_in_adp': self.enable_lm_head_tp_in_adp, + } + + +def create_draft_mapping(target_mapping: "Mapping", + draft_tp_size: int) -> "Mapping": + """Create a mapping for a draft model that uses smaller TP than the target. + + Args: + target_mapping: The target model's Mapping. + draft_tp_size: The TP size for the draft model. Must evenly divide + the target's TP size. + + Returns: + The target mapping unchanged if ``draft_tp_size`` equals the target + TP size, otherwise a ``DraftMapping`` with the appropriate sub-groups. + """ + if draft_tp_size is None or draft_tp_size == target_mapping.tp_size: + return target_mapping + return DraftMapping(target_mapping, draft_tp_size) From 2e9c0636ef0a3a8ca791c70e3fb7835b11d4ea91 Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Fri, 6 Mar 2026 15:06:51 -0500 Subject: [PATCH 2/4] Mapping refactor Signed-off-by: Mike Iovine --- tensorrt_llm/mapping.py | 70 +++++++---------------------------------- 1 file changed, 11 insertions(+), 59 deletions(-) diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index 4ae9185bc073..40836bf02690 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -704,7 +704,7 @@ def _get_all_tp_groups(mapping: "Mapping") -> List[List[int]]: return groups -class DraftMapping: +class DraftMapping(Mapping): """Mapping for draft models with smaller TP than the target model. In one-model speculative decoding, the draft model (typically much smaller) @@ -712,10 +712,14 @@ class DraftMapping: into groups of ``draft_tp_size``, each of which independently computes the same draft tokens with reduced communication overhead. - This object provides the subset of the Mapping interface that draft model - layers (Linear, Attention, etc.) and the weight loader need. + Inherits from Mapping so it passes ``isinstance(m, Mapping)`` checks, but + bypasses the normal ``__new__`` dispatch (MpiTopology / DeviceMeshTopology) + because it builds a *virtual* sub-topology over the target's physical ranks. """ + def __new__(cls, *args, **kwargs): + return object.__new__(cls) + def __init__(self, target_mapping: "Mapping", draft_tp_size: int): if draft_tp_size == target_mapping.tp_size: raise ValueError( @@ -737,10 +741,12 @@ def __init__(self, target_mapping: "Mapping", draft_tp_size: int): self.gpus_per_node = target_mapping.gpus_per_node self.world_size = draft_tp_size * self.pp_size - self.rank = (self.pp_rank * draft_tp_size + - target_mapping.tp_rank % draft_tp_size) self.enable_attention_dp = target_mapping.enable_attention_dp self.enable_lm_head_tp_in_adp = target_mapping.enable_lm_head_tp_in_adp + # rank setter (inherited from MappingBase) validates against + # world_size and enable_attention_dp, so those must be set first. + self.rank = (self.pp_rank * draft_tp_size + + target_mapping.tp_rank % draft_tp_size) self.moe_tp_size = min(draft_tp_size, target_mapping.moe_tp_size) self.moe_ep_size = 1 @@ -814,14 +820,6 @@ def tp_group_pg(self): def moe_tp_rank(self): return self._tp_rank % self.moe_tp_size - @property - def moe_ep_rank(self): - return 0 - - @property - def moe_cluster_rank(self): - return 0 - @property def moe_tp_group(self): return self._moe_tp_group @@ -839,16 +837,6 @@ def moe_ep_group_pg(self): from tensorrt_llm._torch.device_mesh import SingleProcessGroup return SingleProcessGroup.get_group() - @property - def dp_size(self): - return self.tp_size if self.enable_attention_dp else 1 - - def has_tp(self): - return self.tp_size > 1 - - def has_pp(self): - return self.pp_size > 1 - def has_cp(self): return False @@ -858,21 +846,12 @@ def has_cp_ulysses(self): def has_cp_helix(self): return False - def has_moe_tp(self): - return self.moe_tp_size > 1 - def has_moe_ep(self): return False def has_moe_cluster(self): return False - def is_first_pp_rank(self): - return self.pp_rank == 0 - - def is_last_pp_rank(self): - return self.pp_rank == self.pp_size - 1 - @property def node_rank(self): return self._global_rank // self.gpus_per_node @@ -881,18 +860,6 @@ def node_rank(self): def local_rank(self): return self._global_rank % self.gpus_per_node - def get_node_rank(self, rank: int): - return rank // self.gpus_per_node - - def get_local_rank(self, rank: int): - return rank % self.gpus_per_node - - def is_multi_node(self): - return self.world_size > self.gpus_per_node - - def is_second_last_pp_rank(self): - return self.pp_rank == self.pp_size - 2 - def pp_layers(self, num_layers): return self._target.pp_layers(num_layers) @@ -908,21 +875,6 @@ def prev_pp_rank(self): def next_pp_rank(self): return self._target.next_pp_rank() - def to_dict(self): - return { - 'world_size': self.world_size, - 'rank': self.rank, - 'gpus_per_node': self.gpus_per_node, - 'cp_size': self.cp_size, - 'tp_size': self.tp_size, - 'pp_size': self.pp_size, - 'moe_tp_size': self.moe_tp_size, - 'moe_cluster_size': self.moe_cluster_size, - 'moe_ep_size': self.moe_ep_size, - 'enable_attention_dp': self.enable_attention_dp, - 'enable_lm_head_tp_in_adp': self.enable_lm_head_tp_in_adp, - } - def create_draft_mapping(target_mapping: "Mapping", draft_tp_size: int) -> "Mapping": From 344acc723fab3b07b2b5fd01ebf234fbab00c6ec Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Fri, 6 Mar 2026 15:19:00 -0500 Subject: [PATCH 3/4] Add tests Signed-off-by: Mike Iovine --- .../defs/accuracy/test_llm_api_pytorch.py | 79 ++++++++++++++++++- .../test_lists/test-db/l0_dgx_b200.yml | 4 + 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 1168bdbe5a2b..ae1c5508055d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -55,10 +55,10 @@ def patched_start_mpi_pool(self): # isort: off from tensorrt_llm.llmapi import ( AutoDecodingConfig, CudaGraphConfig, DeepSeekSparseAttentionConfig, - Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig, - NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, - SamplingParams, SkipSoftmaxAttentionConfig, SADecodingConfig, - TorchCompileConfig) + DraftTargetDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, + MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, + RocketSparseAttentionConfig, SamplingParams, SkipSoftmaxAttentionConfig, + SADecodingConfig, TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -609,6 +609,60 @@ def test_fp8_beam_search(self, enable_cuda_graph, enable_padding, sampling_params=sampling_params, extra_acc_spec="beam_width=2") + @skip_pre_hopper + @pytest.mark.skip_less_device(2) + def test_eagle3_draft_tp_size(self): + pytorch_config = dict( + max_batch_size=1, + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig(max_batch_size=1, + enable_padding=True), + ) + kv_cache_config = KvCacheConfig(enable_block_reuse=True, + free_gpu_memory_fraction=0.8) + + eagle_model_dir = f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B" + spec_config = Eagle3DecodingConfig(max_draft_len=4, + speculative_model=eagle_model_dir, + eagle3_one_model=True, + draft_tp_size=1) + + with LLM(model=self.MODEL_PATH, + tensor_parallel_size=2, + **pytorch_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config) as llm: + task = CnnDailymail(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @skip_pre_hopper + @pytest.mark.skip_less_device(2) + def test_draft_target_draft_tp_size(self): + draft_model_dir = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-1B" + pytorch_config = dict( + max_batch_size=1, + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(max_batch_size=1), + ) + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.8) + + spec_config = DraftTargetDecodingConfig( + max_draft_len=4, + speculative_model=draft_model_dir, + draft_tp_size=1, + ) + + with LLM(model=self.MODEL_PATH, + tensor_parallel_size=2, + **pytorch_config, + kv_cache_config=kv_cache_config, + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + class TestLlama3_2_1B(LlmapiAccuracyTestHarness): MODEL_NAME = "meta-llama/Llama-3.2-1B" @@ -1512,6 +1566,23 @@ def test_bfloat16_mtp_sa(self): task = MMLU(self.MODEL_NAME) task.evaluate(llm, extra_acc_spec="use_sa_spec") + @pytest.mark.skip_less_device(4) + def test_mtp_draft_tp_size_4gpus(self): + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + pytorch_config = dict( + disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig(), + ) + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=2, + draft_tp_size=1) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + kv_cache_config=kv_cache_config, + **pytorch_config, + speculative_config=mtp_config) as llm: + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + @pytest.mark.skip_less_device(4) @parametrize_with_ids("torch_compile", [False, True]) @parametrize_with_ids("attention_dp,cuda_graph,overlap_scheduler", 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 1b548e55a051..68f8cfee594e 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -28,6 +28,10 @@ l0_dgx_b200: - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4] - accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] TIMEOUT (60) + # ------------- draft_tp_size tests --------------- + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_mtp_draft_tp_size_4gpus + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3_draft_tp_size + - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_draft_target_draft_tp_size # ------------- VisualGen multi-GPU tests --------------- - unittest/_torch/visual_gen/multi_gpu - unittest/_torch/visual_gen/test_wan.py::TestWanParallelism::test_cfg_2gpu_correctness From 588035c6377a1e7f9f5e00acf54968e3e15aff2a Mon Sep 17 00:00:00 2001 From: Mike Iovine Date: Fri, 6 Mar 2026 15:30:02 -0500 Subject: [PATCH 4/4] Add error messages Signed-off-by: Mike Iovine --- tensorrt_llm/_torch/pyexecutor/_util.py | 10 ++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 7 +++++++ 2 files changed, 17 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0629d81a8921..4768b144ba5e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -572,6 +572,16 @@ def _should_create_separate_draft_kv_cache(self) -> bool: back to the target model's config via _get_effective_draft_config(). """ if self._mapping.enable_attention_dp: + draft_tp = getattr(self._speculative_config, 'draft_tp_size', + None) + if (draft_tp is not None + and draft_tp < self._mapping.tp_size): + raise ValueError( + f"draft_tp_size ({draft_tp}) < tp_size " + f"({self._mapping.tp_size}) requires separate " + "draft/target KV caches, which are not supported " + "with attention data parallelism." + ) logger.info( "Attention DP is enabled, separate draft KV cache is not supported." ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 2d7614559e65..f296658d38f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -365,6 +365,13 @@ def create_py_executor( # Disable separate draft KV cache in disaggregated mode # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM if cache_transceiver_config is not None: + draft_tp = getattr(spec_config, 'draft_tp_size', None) + if draft_tp is not None and draft_tp < mapping.tp_size: + raise ValueError( + f"draft_tp_size ({draft_tp}) < tp_size ({mapping.tp_size}) " + "requires separate draft/target KV caches, which are not " + "supported in disaggregated serving mode." + ) spec_config._allow_separate_draft_kv_cache = False # chunk_unit_size may be changed to 64 when using flash mla