diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 86291f2ea770..f83c490e03e8 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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, diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py deleted file mode 100644 index 9f2bfd441219..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MLA in-tree module for TensorRT-LLM's PyTorch backend. - -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. -* **Softmax scale.** ``(qk_nope + qk_rope) ** -0.5 = 192 ** -0.5`` for - real K3 dims — matches ``TrtllmAttention`` default MLA q_scaling. - -The module wraps the existing ``TrtllmAttention`` backend MLA path plus -``KVCacheManagerV2`` for both context and cached-decode. -""" - -from .kimi_k3_mla_attention import KimiK3MLAAttention - -__all__ = [ - "KimiK3MLAAttention", -] 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 deleted file mode 100644 index 4ab1d52f1c55..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ /dev/null @@ -1,237 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 specialization of the shared PyTorch-backend MLA module. - -The base ``MLA`` class owns context attention, cached/chunked prefill, -absorbed generation, and paged-cache handling. This module only supplies the -K3 projection topology, NoPE identity table, KV-B checkpoint layout, and -gated output projection. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from ....functional import PositionEmbeddingType -from ....mapping import Mapping -from ...attention.backends import TrtllmAttention -from ...attention.backends.interface import PositionalEmbeddingParams, RopeParams -from ...attention.mla import MLA -from ...model_config import ModelConfig -from ..linear import Linear, TensorParallelMode - - -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 - - -# --------------------------------------------------------------------------- -# KimiK3MLAAttention. -# --------------------------------------------------------------------------- - - -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 = Linear( - 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