[None][feat] Kimi K3: KDA-TP + MLA-DCP (helix) wiring

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

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,

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2201,6 +2201,11 @@ def __init__(
 
         # 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
+        if mapping.has_cp_helix():
+            # Helix stripes only the MLA KV; every other layer (incl. KDA)
+            # repurposes CP ranks as plain TP, so each rank holds a
+            # 1/(tp*cp) head slice of the mamba state.
+            tp_size = mapping.tp_size * mapping.cp_size
         d_inner = mamba_head_dim * mamba_num_heads
         conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state
         nheads = mamba_num_heads
@@ -2936,6 +2941,11 @@ def __init__(
 
         if self.local_num_mamba_layers > 0:
             tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1
+            if mapping.has_cp_helix():
+                # Helix stripes only the MLA KV; every other layer (incl.
+                # KDA) repurposes CP ranks as plain TP — 1/(tp*cp) head
+                # slice per rank.
+                tp_size = mapping.tp_size * mapping.cp_size
             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

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -1333,6 +1333,24 @@ def wrapper(self, resource_manager: ResourceManager, *args, **kwargs):
 
         return wrapper
 
+    def _helix_safe_warmup_configs(self, configs: List[Tuple[int, int]],
+                                   kv_cache_manager) -> List[Tuple[int, int]]:
+        """Floor context-shaped warmup lengths under helix.
+
+        The round-robin context partition gives rank r the blocks
+        [r::cp_size]; a context warmup shorter than cp_size*tokens_per_block
+        leaves some ranks with zero KV, which fails MLA TMA descriptor
+        creation. Skipping the warmup instead is not safe: the first real
+        request would drive Triton autotuning, whose repeated candidate
+        launches re-apply in-place recurrent-state updates.
+        """
+        if not self.mapping.has_cp_helix() or kv_cache_manager is None:
+            return configs
+        # +1 keeps the length non-aligned (the Kimi KDA prefill warmup needs
+        # a non-aligned length to enter the pure K123 path).
+        floor = self.mapping.cp_size * kv_cache_manager.tokens_per_block + 1
+        return [(t if g > 0 else max(t, floor), g) for (t, g) in configs]
+
     def _get_max_shape_warmup_requests(
             self, resource_manager: ResourceManager) -> List[Tuple[int, int]]:
         """
@@ -1355,6 +1373,9 @@ def _get_max_shape_warmup_requests(
             (max_batch_size, max_batch_size),  # max_batch_size, pure generation
         ]
 
+        warmup_requests_configs = self._helix_safe_warmup_configs(
+            warmup_requests_configs, kv_cache_manager)
+
         return warmup_requests_configs
 
     def _get_full_general_warmup_requests(
@@ -1375,6 +1396,10 @@ def _get_full_general_warmup_requests(
         # 1-token first to capture the 0→1 transition graph; max-shape next to seed
         # triton autotuning with the largest inputs; 2-token last for the small-ctx path.
         warmup_configs = one_token_configs + max_configs + small_ctx_configs
+        kv_cache_manager = resource_manager.get_resource_manager(
+            self.kv_cache_manager_key)
+        warmup_configs = self._helix_safe_warmup_configs(
+            warmup_configs, kv_cache_manager)
         # Deduplicate the warmup_configs while keeping the order.
         return list(dict.fromkeys(warmup_configs))
 
@@ -1823,7 +1848,10 @@ def trtllm_gen_fmha_jit_warmup():
 
         model_type = getattr(self.model.model_config.pretrained_config,
                              "model_type", None)
-        if can_run_general_warmup and model_type in ("kimi_k3", "kimi_linear"):
+        # Not gated on can_run_general_warmup: Kimi always uses a
+        # MambaHybridCacheManager, which forces that gate False, yet the KDA
+        # pure-K123 prefill warmup is still required.
+        if model_type in ("kimi_k3", "kimi_linear"):
             # Kimi's one-token context takes the NT < 4 FLA fallback and does
             # not compile the optimized single-sequence K123 variant. A
             # non-aligned five-chunk context enters the pure K123 path.
@@ -1843,6 +1871,11 @@ def trtllm_gen_fmha_jit_warmup():
                 "Skipped TRTLLM-Gen flashinfer_trtllm_gen FMHA lib JIT warmup When enable cute_dsl_mla FMHA lib"
             )
 
+        kv_cache_manager = resource_manager.get_resource_manager(
+            self.kv_cache_manager_key)
+        warmup_requests_configs = self._helix_safe_warmup_configs(
+            warmup_requests_configs, kv_cache_manager)
+
         for num_tokens, num_gen_requests in warmup_requests_configs:
             warmup_request = self._create_warmup_request(
                 resource_manager,
@@ -1851,6 +1884,15 @@ def trtllm_gen_fmha_jit_warmup():
 
             with self.no_cuda_graph(), self._release_batch_context(
                     warmup_request, resource_manager) as batch:
+                if batch is None and self.mapping.has_cp_helix():
+                    # The helix floor may exceed tiny token/KV budgets; a
+                    # skipped warmup leaves Triton autotuning to the first
+                    # real request.
+                    logger.warning(
+                        f"Helix: skipping attention warmup shape "
+                        f"(num_tokens={num_tokens}, "
+                        f"num_gen_requests={num_gen_requests}); it does not "
+                        f"fit the token/KV budgets.")
                 if batch is None and self.mapping.tp_size <= 1:
                     continue  # Not enough KV cache space (single rank, safe to skip)
                 self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens)