diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index f9c771797c59..7bce6227cf61 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -148,6 +148,8 @@ if TYPE_CHECKING: from transformers import PretrainedConfig + from ...llmapi.llm_args import DecodingBaseConfig + # Identity-RoPE table positions for the MLA backends. K3 is NoPE (the table # holds cos=1/sin=0), but the chunked-context path indexes the table by # absolute position, so it must cover max_position_embeddings (~512MB per @@ -2406,10 +2408,11 @@ class KimiMLARuntime(nn.Module): def __init__( self, - cfg, + cfg: "PretrainedConfig", layer_idx: int, model_config: ModelConfig, - ): + mapping_with_cp: Optional[Mapping] = None, + ) -> None: super().__init__() from ..modules.kimi_k3_mla import KimiK3MLAAttention @@ -2424,6 +2427,9 @@ def __init__( # KimiK3MLAAttention owns MLA projection/head sharding. Keep only the # final output reduction in this wrapper so the output gate remains # between attention and the row-parallel o_proj. + # Helix: mapping_with_cp (the CP original) activates the base MLA's + # helix machinery; this wrapper's allreduce over the repurposed + # mapping sums the base o_proj's tp*cp partials. mapping = model_config.mapping reduce_output = not mapping.enable_attention_dp and mapping.tp_size > 1 self._o_allreduce = ( @@ -2449,6 +2455,7 @@ def __init__( use_output_gate=cfg.mla_use_output_gate, max_position_embeddings=max_positions, model_config=model_config, + mapping_with_cp=mapping_with_cp, ) def forward( @@ -2512,6 +2519,8 @@ def __init__( cfg, layer_idx, model_config=mla_model_config, + # CP original stashed by _setup_helix_mappings; None outside helix. + mapping_with_cp=getattr(model_config, "_helix_mapping_with_cp", None), ) self.is_moe = ( @@ -2848,6 +2857,9 @@ def __init__(self, model_config: ModelConfig): cfg = _get_text_config(model_config.pretrained_config) assert model_config.mapping.pp_size == 1, "Kimi K3 does not support pipeline parallelism" spec_config = getattr(model_config, "spec_config", None) + + # Helix: swap in the repurposed mapping; restored after super().__init__. + self._setup_helix_mappings(model_config, cfg, spec_config) # Supported spec-dec modes: # - SA (suffix automaton): one-engine in-forward drafting, no draft # weights; the KDA/MLA verify paths below implement multi-token @@ -2873,6 +2885,91 @@ def __init__(self, model_config: ModelConfig): vocab_size=cfg.vocab_size, ) + # Restore the CP original: executor-side helix bookkeeping keys off + # has_cp_helix() at runtime. + if self.mapping_with_cp is not None: + model_config._frozen = False + model_config.mapping = self.mapping_with_cp + model_config._frozen = True + + def _setup_helix_mappings( + self, + model_config: ModelConfig, + cfg: "PretrainedConfig", + spec_config: Optional["DecodingBaseConfig"], + ) -> None: + """Validate helix preconditions and stage the dual-mapping swap. + + DeepseekV3 pattern: the MLA layers keep the CP original; everything + else is built against the repurposed mapping (CP ranks become TP + ranks). Sets ``mapping_with_cp`` (restored after construction) and + ``_repurposed_tp_mapping`` (load_weights shard selection); both stay + None outside helix. + """ + self.mapping_with_cp = None + self._repurposed_tp_mapping = None + if not model_config.mapping.has_cp_helix(): + return + if model_config.mapping.enable_attention_dp: + raise ValueError( + "Kimi K3 helix phase 1 requires enable_attention_dp=" + "False: the helix ADP token-scatter conflicts with the " + "per-request locality of KDA recurrent state." + ) + if spec_config is not None: + raise ValueError( + "Kimi K3 helix phase 1 does not support speculative " + "decoding (round-robin KV bookkeeping assumes one token " + "per decode step)." + ) + cp = model_config.mapping.cp_size + repurposed_tp = model_config.mapping.tp_size * cp + if cfg.num_attention_heads % repurposed_tp != 0: + raise ValueError( + f"Kimi K3 helix requires tp_size*cp_size ({repurposed_tp}) " + f"to divide the MLA head count ({cfg.num_attention_heads})." + ) + kda_heads = cfg.linear_attn_config["num_heads"] + if kda_heads % repurposed_tp != 0: + raise ValueError( + f"Kimi K3 helix requires tp_size*cp_size ({repurposed_tp}) to " + f"divide the KDA head count ({kda_heads})." + ) + # MoE splits apply to the repurposed tp*cp group (helix + # moe_world_size = tp*cp); default EP-only. The Mapping constructor + # skips its product check when both sizes are 1, so validate here. + moe_ep = repurposed_tp + if model_config.mapping.moe_tp_ep_user_specified: + moe_tp = model_config.mapping.moe_tp_size + moe_ep = model_config.mapping.moe_ep_size + if moe_tp * moe_ep != repurposed_tp: + raise ValueError( + f"Kimi K3 helix: moe_tensor_parallel_size ({moe_tp}) x " + f"moe_expert_parallel_size ({moe_ep}) must equal " + f"tp_size*cp_size ({repurposed_tp}): MoE runs on the " + "repurposed tp*cp group." + ) + if cfg.num_experts and cfg.num_experts % moe_ep != 0: + raise ValueError( + f"Kimi K3 helix requires the MoE EP size ({moe_ep}) to " + f"divide the routed expert count ({cfg.num_experts}): each " + "EP rank of the repurposed tp*cp group holds whole experts." + ) + self.mapping_with_cp = copy.deepcopy(model_config.mapping) + repurposed = model_config.mapping.repurpose_helix_cp_to_tp() + # repurpose passes resolved moe sizes, which the Mapping constructor + # mistakes for user-specified values; restore the flag. + repurposed.moe_tp_ep_user_specified = self.mapping_with_cp.moe_tp_ep_user_specified + # load_weights shard selection must use this tp_rank; the restored + # CP original's tp_rank is 0 on every rank. + self._repurposed_tp_mapping = repurposed + model_config._frozen = False + model_config.mapping = repurposed + # Side-channel for the MLA layers (avoids threading a kwarg through + # every intermediate signature). + model_config._helix_mapping_with_cp = self.mapping_with_cp + model_config._frozen = True + @classmethod def get_model_defaults(cls, llm_args) -> dict: # - enable_block_reuse defaults off: reuse is supported as an @@ -3061,7 +3158,13 @@ def _load_trunk_params( # MLP TP shard index. A dense MLP whose intermediate size does not # divide model TP uses a smaller repeated TP subgroup, so its local # shard rank is model tp_rank modulo the parameter's shard count. - model_tp_rank = self.model_config.mapping.tp_rank + # Under helix the modules were sharded against the repurposed + # mapping; the restored CP original's tp_rank is 0 on every rank. + model_tp_rank = ( + self._repurposed_tp_mapping.tp_rank + if self._repurposed_tp_mapping is not None + else self.model_config.mapping.tp_rank + ) # Keep each FP8_PB_WO checkpoint pair alongside the BF16 # parameter only when the later weight-read conversion consumes it. stash_ckpt_fp8 = _resolve_fp8_weight_read_gates()[0] @@ -3134,6 +3237,12 @@ def load_param(name: str, param: torch.nn.Parameter): ).to(param.dtype) ) mla_mixer.k_b_proj_trans.data.copy_(k_weight.transpose(1, 2)) + # Helix: v_b_proj holds this rank's 1/cp post-all-to-all + # head chunk; kv_b/k_b_proj_trans keep every tp-local head. + h_cp = mla_mixer.num_heads_tp_cp + if h_cp != h: + lo = mla_mixer.mapping.cp_rank * h_cp + v_weight = v_weight[lo : lo + h_cp] mla_mixer.v_b_proj.data.copy_(v_weight) return if name.endswith(".A_log") and src.numel() != param.numel(): diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index db0d816c42df..889ac18c7efa 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -18,6 +18,7 @@ from ....functional import PositionEmbeddingType from ....logger import logger +from ....mapping import Mapping from ....models.modeling_utils import QuantConfig from ...attention_backend import AttentionMetadata, TrtllmAttention, TrtllmAttentionMetadata from ...attention_backend.interface import PositionalEmbeddingParams, RopeParams @@ -231,6 +232,7 @@ def __init__( use_output_gate: bool = True, max_position_embeddings: int = 8192, model_config: ModelConfig, + mapping_with_cp: Optional[Mapping] = None, ) -> None: pos_embd_params = _make_pos_embd_params( qk_rope_head_dim=qk_rope_head_dim, @@ -253,6 +255,7 @@ def __init__( dtype=dtype, dense_bias=False, config=model_config, + mapping_with_cp=mapping_with_cp, reduce_output=False, fuse_qkv_a_proj=False, rms_norm_eps=rms_norm_eps, @@ -266,14 +269,15 @@ def __init__( self.use_output_gate = use_output_gate if use_output_gate: - # Follow q_b_proj's effective MLA mapping: replicated under - # attention-DP and column-sharded by head otherwise. + # The gate must match o_proj's input sharding (under helix the + # post-all-to-all 1/cp head chunk); outside helix this equals + # q_b_proj's head sharding, replicated under attention-DP. self.g_proj = Linear( hidden_size, num_heads * v_head_dim, bias=False, dtype=dtype, - mapping=self.q_b_proj.mapping, + mapping=self.o_proj.mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, quant_config=model_config.get_quant_config(), skip_create_weights_in_init=model_config.skip_create_weights_in_init, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 10a7dba55540..6660c0665ee6 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -126,6 +126,19 @@ class MambaRole: CONV_STATE = DataRole("conv_state") +def _mamba_effective_tp_size(mapping: Mapping) -> int: + """TP degree for sizing per-rank mamba/KDA state pools. + + Attention-DP replicates the state and takes precedence; helix + repurposes CP ranks as plain TP for recurrent-state layers. + """ + if mapping.enable_attention_dp: + return 1 + if mapping.has_cp_helix(): + return mapping.tp_size * mapping.cp_size + return mapping.tp_size + + def get_tensor_size_bytes(tensor): """Calculate tensor size in bytes.""" if isinstance(tensor, torch.Tensor): @@ -449,7 +462,7 @@ def __init__( self._seed_request_counter = 0 # get tp size - tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size + tp_size = _mamba_effective_tp_size(mapping) # derive mamba parameters for conv and ssm states d_inner = head_dim * num_heads @@ -2200,7 +2213,7 @@ def __init__( return # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) - tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + tp_size = _mamba_effective_tp_size(mapping) d_inner = mamba_head_dim * mamba_num_heads conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state nheads = mamba_num_heads @@ -2935,7 +2948,7 @@ def __init__( and self.local_num_mamba_layers > 0) if self.local_num_mamba_layers > 0: - tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + tp_size = _mamba_effective_tp_size(mapping) d_inner = mamba_head_dim * mamba_num_heads grouped_state_dim = mamba_n_groups * mamba_d_state conv_dim = d_inner + 2 * grouped_state_dim