[TRTLLM-15177][chore] Kimi K3: inline MLA module into modeling_kimi_linear

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -104,10 +104,13 @@
 from torch import nn
 
 from ..._utils import is_sm_100f
+from ...functional import PositionEmbeddingType
 from ...logger import logger
 from ...mapping import Mapping
 from ...models.modeling_utils import QuantAlgo, QuantConfig
-from ..attention.backends import AttentionMetadata
+from ..attention.backends import AttentionMetadata, TrtllmAttention
+from ..attention.backends.interface import PositionalEmbeddingParams, RopeParams
+from ..attention.mla import MLA
 from ..distributed import AllReduce, AllReduceParams
 from ..model_config import ModelConfig
 from ..modules.gated_mlp import GatedMLP
@@ -1426,6 +1429,238 @@ def _routed_output():
         return routed_out + shared_out
 
 
+# ---------------------------------------------------------------------------
+# MLA attention.
+#
+# K3 MLA is DeepSeek-V3-style multi-latent attention with three K3-specific
+# deltas that live at the module level (not the attention backend):
+#
+# * NoPE. ``mla_use_nope=True`` in K3 config disables the rotary embedding;
+#   both the query and key rope slots pass through the backend unchanged.
+# * Output gate before ``o_proj``. When ``mla_use_output_gate=True`` an extra
+#   ``g_proj`` computes ``sigmoid(g_proj(hidden_states)) * attn_output`` before
+#   the final projection (installed as a base-MLA output hook).
+# * Softmax scale. ``(qk_nope + qk_rope) ** -0.5 = 192 ** -0.5`` for real K3
+#   dims — matches ``TrtllmAttention`` default MLA q_scaling.
+#
+# The base ``MLA`` class owns context attention, cached/chunked prefill,
+# absorbed generation, and paged-cache handling. The code below only supplies
+# the K3 projection topology, NoPE identity table, KV-B checkpoint layout, and
+# gated output projection. It was inlined here (from the former
+# ``modules/kimi_k3_mla`` package) to match the per-model ``modeling_xxx.py``
+# convention; ``kimi_kda`` stays a standalone module.
+# ---------------------------------------------------------------------------
+
+
+def _meta_safe_cast_dtype(module, dtype):
+    """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``.
+
+    ``Module.to`` dispatches ``aten._to_copy``, which MetaInitMode rejects
+    (it would silently fall back to full CPU construction of the model —
+    ~70 GB of host RAM per rank for Kimi K3). Under meta init the values
+    are garbage anyway, so a dtype-only re-allocation via ``empty_like``
+    (an allowed init op) is equivalent; off meta this matches ``.to``.
+    """
+    import torch as _torch
+
+    def _cast(t):
+        if not t.is_floating_point():
+            return t
+        if t.is_meta:
+            return _torch.empty_like(t, dtype=dtype)
+        return t.to(dtype=dtype)
+
+    module._apply(_cast)
+
+
+def _make_pos_embd_params(
+    *,
+    qk_rope_head_dim: int,
+    max_position_embeddings: int,
+) -> PositionalEmbeddingParams:
+    """Build a valid rope config so the backend allocates a real cache.
+
+    We use rope_gpt_neox with default theta=10000 and ``duplicate_data
+    =True`` (the same convention DeepSeek-V3-style MLA uses when
+    ``qk_rope_head_dim`` is present). The resulting ``rotary_cos_sin``
+    has the exact shape the C++ MLA rope kernel indexes. Immediately
+    after backend construction we overwrite the tensor values with
+    ``(cos=1, sin=0)`` — an identity rotation, matching K3's NoPE.
+    """
+    rope_params = RopeParams(
+        dim=qk_rope_head_dim,
+        theta=10000.0,
+        max_positions=max_position_embeddings,
+        original_max_positions=max_position_embeddings,
+        duplicate_data=True,
+    )
+    return PositionalEmbeddingParams(
+        type=PositionEmbeddingType.rope_gpt_neox,
+        rope=rope_params,
+        # Match the working DeepSeek-V3-style MLA reference test
+        # (tests/unittest/_torch/attention/test_attention_mla.py) which
+        # sets ``is_neox=False``. The MLA fused rope kernel is GPT-J
+        # style regardless of this flag, but the C++ FMHA reads this bit
+        # elsewhere and stability under identity-cos-sin depends on the
+        # standard non-neox layout.
+        is_neox=False,
+    )
+
+
+def _write_identity_rope_values(cos_sin: torch.Tensor) -> None:
+    """Overwrite a rotary cos/sin table with identity values in place.
+
+    Interleaved (cos, sin) pairs: index [::2] = cos, [1::2] = sin.
+    Setting cos=1 and sin=0 per position makes the rotation the
+    identity — a mathematical no-op — which preserves K3's NoPE
+    semantics without patching the backend.
+    """
+    flat = cos_sin.reshape(-1)
+    with torch.no_grad():
+        flat[0::2] = 1.0
+        flat[1::2] = 0.0
+    # Ensure the identity write reaches CUDA memory before any kernel
+    # launched from a different stream can read the table.
+    if cos_sin.is_cuda:
+        torch.cuda.synchronize(cos_sin.device)
+
+
+def _install_identity_rope_table(backend: TrtllmAttention) -> None:
+    """Install an identity rotary cos/sin table on ``backend``.
+
+    The C++ MLA rope kernels (``mla_rope_generation`` and the context
+    preprocess) read this table and apply the rotation; identity values
+    make that a copy, preserving K3's NoPE.
+
+    The tensor SHAPE produced by ``create_rope_const_params`` is kept
+    intact so the C++ ``float2`` indexing stays valid. Only the values
+    are overwritten in place. ``_ensure_rope_table_size`` is replaced
+    with an identity-preserving resize: the table may GROW (so the
+    fused rope-generation op can never index out of bounds for long
+    sequences) but its values are always rewritten to identity right
+    after a regeneration, so the real sinusoids never leak in.
+    """
+    cos_sin = backend.rotary_cos_sin
+    if cos_sin is None:
+        raise RuntimeError(
+            "backend.rotary_cos_sin is None after construction; check "
+            "pos_embd_params has a valid RopeParams with dim > 0."
+        )
+    _write_identity_rope_values(cos_sin)
+
+    orig_resize = backend._ensure_rope_table_size  # bound method
+
+    def _identity_preserving_resize(required_max_positions: int) -> None:
+        if required_max_positions <= backend.rope_params.max_positions:
+            return
+        orig_resize(required_max_positions)
+        _write_identity_rope_values(backend.rotary_cos_sin)
+
+    backend._ensure_rope_table_size = _identity_preserving_resize
+
+
+class KimiK3MLAAttention(MLA):
+    """Kimi K3 MLA implemented as a thin specialization of :class:`MLA`.
+
+    K3 keeps the standard dense MLA attention/cache flow and only changes the
+    checkpoint projection topology, positional encoding, KV-B runtime layout,
+    and gated output projection.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_heads: int,
+        q_lora_rank: int,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        rms_norm_eps: Optional[float] = None,
+        dtype: Optional[torch.dtype] = None,
+        layer_idx: int = 0,
+        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,
+            max_position_embeddings=max_position_embeddings,
+        )
+        super().__init__(
+            hidden_size=hidden_size,
+            num_attention_heads=num_heads,
+            num_key_value_heads=num_heads,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            q_lora_rank=q_lora_rank,
+            kv_lora_rank=kv_lora_rank,
+            predicted_tokens_per_seq=1,
+            max_position_embeddings=max_position_embeddings,
+            bias=False,
+            pos_embd_params=pos_embd_params,
+            layer_idx=layer_idx,
+            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,
+        )
+        # Keep the base MLA registration enabled so breakable CUDA graphs use
+        # the shared custom op. The output gate is a base hook and runs on both
+        # the registered and eager paths before the row-parallel o_proj.
+
+        self.use_output_gate = use_output_gate
+
+        if use_output_gate:
+            # 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 = TrtllmLinear(
+                hidden_size,
+                num_heads * v_head_dim,
+                bias=False,
+                dtype=dtype,
+                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,
+                allreduce_strategy=model_config.allreduce_strategy,
+                force_dynamic_quantization=model_config.force_dynamic_quantization,
+                use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
+                use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
+            )
+
+        # K3 is NoPE. The base MLA backends still require real RoPE tables, so
+        # retain their expected shape and replace every rotation with identity.
+        assert isinstance(self.mha, TrtllmAttention)
+        assert isinstance(self.mqa, TrtllmAttention)
+        _install_identity_rope_table(self.mha)
+        _install_identity_rope_table(self.mqa)
+        self.rotary_emb = None
+        self.apply_rotary_emb = False
+
+        if dtype is not None:
+            _meta_safe_cast_dtype(self, dtype)
+
+    def _apply_output_gate(
+        self,
+        hidden_states: torch.Tensor,
+        attn_output: torch.Tensor,
+    ) -> torch.Tensor:
+        # Sigmoid gate on o_proj's input. g_proj matches o_proj's input
+        # sharding, so the multiply composes with the helix-CP output
+        # projection.
+        if self.use_output_gate:
+            return attn_output * self.g_proj(hidden_states).sigmoid()
+        return attn_output
+
+
 # ---------------------------------------------------------------------------
 # MLA runtime.
 # ---------------------------------------------------------------------------
@@ -1443,8 +1678,6 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        from ..modules.kimi_k3_mla import KimiK3MLAAttention
-
         max_positions = int(
             os.environ.get(
                 _KIMI_K3_MLA_MAX_POSITIONS_ENV,