From 92443feb8e3b991a4b2d783e88ba93b85ef4939d Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 24 Aug 2026 16:27:37 -0500 Subject: [PATCH 1/3] [TRTLLM-15177][chore] Kimi K3: inline MLA module, drop dead MoE comm plumbing Deferred cleanup from PR #17269, tracked in TRTLLM-15177. Redo of PR kimi_k3_moe module (via the shared modules/situ.py), so only the remaining items are carried over here. 1. Inline the K3-specific kimi_k3_mla module into modeling_kimi_linear.py, matching the per-model modeling_xxx.py convention (e.g. DeepSeek-V3). The moved code is unchanged apart from dropping a redundant local torch import and following the file's Linear-as-TrtllmLinear alias. kimi_kda stays a standalone module (general enough to warrant it). 2. Remove the unused communication_method parameter chain create_moe -> ConfigurableMoE -> CommunicationFactory.create_strategy. The only caller (modeling_kimi_linear.py) passed None, and TRTLLM_FORCE_COMM_METHOD already provides strategy forcing. The unit test covering the forwarding is deleted with it. 3. Tuple default for KimiLinearConfig.keys_to_ignore_at_inference, so the class-level default cannot be mutated in place. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/configs/kimi_linear.py | 4 +- .../_torch/models/modeling_kimi_linear.py | 356 ++++++++++++++++- .../_torch/modules/kimi_k3_mla/__init__.py | 25 -- .../kimi_k3_mla/kimi_k3_mla_attention.py | 369 ------------------ .../communication/communication_factory.py | 15 +- .../_torch/moe/fused_moe/configurable_moe.py | 3 - .../_torch/moe/fused_moe/create_moe.py | 5 - .../modules/test_kimi_k3_mla_backend.py | 4 +- .../_torch/moe/test_kimi_k3_situ_moe.py | 48 --- 9 files changed, 358 insertions(+), 471 deletions(-) delete mode 100644 tensorrt_llm/_torch/modules/kimi_k3_mla/__init__.py delete mode 100644 tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py diff --git a/tensorrt_llm/_torch/configs/kimi_linear.py b/tensorrt_llm/_torch/configs/kimi_linear.py index b00fa6627cd0..1513a3614a1d 100644 --- a/tensorrt_llm/_torch/configs/kimi_linear.py +++ b/tensorrt_llm/_torch/configs/kimi_linear.py @@ -17,7 +17,9 @@ class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" - keys_to_ignore_at_inference = ["past_key_values"] + # Tuple, not list: a class-level mutable default could be appended to + # in place, changing the default for every instance. + keys_to_ignore_at_inference = ("past_key_values",) def __init__( self, diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index dce32f886909..170bca1c0e6f 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -86,6 +86,7 @@ import os import threading from contextlib import ExitStack +from functools import partial from typing import ( TYPE_CHECKING, Any, @@ -104,16 +105,19 @@ 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_backend import AttentionMetadata +from ..attention_backend import AttentionMetadata, TrtllmAttention, TrtllmAttentionMetadata +from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..distributed import AllReduce, AllReduceParams from ..model_config import ModelConfig from ..modules.gated_mlp import GatedMLP from ..modules.kimi_kda import KimiKDALinearAttention from ..modules.linear import Linear as TrtllmLinear from ..modules.linear import TensorParallelMode, load_weight_shard +from ..modules.mla import MLA from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul @@ -1118,8 +1122,6 @@ def __init__( model_config=routed_moe_model_config, override_quant_config=routed_quant_config, layer_idx=layer_idx, - # Let CommunicationFactory select the best available strategy. - communication_method=None, ) # trtllm-gen ships SiTu cubins for exactly one dtype combination # (``Bmm_MxE4m3_MxE2m1MxE4m3`` = MXFP8 act x MXFP4 weight) and has no @@ -1485,6 +1487,352 @@ def _routed_output(): return routed_out + shared_out +# --------------------------------------------------------------------------- +# K3 MLA attention: specialization of the shared PyTorch-backend ``MLA``. +# +# The base ``MLA`` class owns context attention, cached/chunked prefill, +# absorbed generation, and paged-cache handling. This section only supplies +# the K3 projection topology, NoPE identity table, KV-B checkpoint layout, +# and gated output projection. +# --------------------------------------------------------------------------- + +_KIMI_K3_MLA_GEN_BACKEND_ENV = "TLLM_K3_MLA_GEN_BACKEND" +_KIMI_K3_MLA_GEN_BACKENDS = ("cute-dsl", "trtllm-gen") + + +def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: + """Select K3's absorbed-generation MLA backend. + + K3 was tuned with the FlashInfer CuTe-DSL backend for BF16 KV cache. + FP8 KV cache carries device scales that CuTe-DSL does not support, so + retain the TRTLLM-Gen fallback used by the pre-refactor implementation. + """ + backend = os.environ.get(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") + # Validate here, where the env var is read: an invalid value would + # otherwise surface only deep inside attention-backend construction, + # with an error that never names the knob that caused it. + if backend not in _KIMI_K3_MLA_GEN_BACKENDS: + raise ValueError( + f"{_KIMI_K3_MLA_GEN_BACKEND_ENV}={backend!r} is invalid; " + f"expected one of {list(_KIMI_K3_MLA_GEN_BACKENDS)}." + ) + has_fp8_kv_cache = bool( + quant_config is not None and quant_config.layer_quant_mode.has_fp8_kv_cache() + ) + if has_fp8_kv_cache and backend != "trtllm-gen": + # info_once: this runs once per MLA layer (~60x at startup for + # FP8-KV) and the decision is identical for every layer. + logger.info_once( + "Kimi K3 MLA: FP8 KV cache requires the trtllm-gen MLA " + f"generation backend; overriding '{backend}' -> 'trtllm-gen'.", + key="kimi_k3_mla_gen_backend_fp8_override", + ) + return "trtllm-gen" + return backend + + +def _validate_mla_generation_backend(backend: str, num_heads: int) -> None: + """Fail fast when `backend` can never run at this per-rank head count. + + FlashInfer's `trtllm_batch_decode_with_kv_cache_mla` rejects + `64 < num_heads_q < 128` for every batch shape, and the per-batch policy + never demotes away from an explicit `trtllm-gen` selection (the + FP8-KV-cache override or a `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen` request). + Without this check the conflict only surfaces as a FlashInfer error deep + in attention warmup. + + The bound mirrors FlashInfer's validation gate verbatim: it predicts + FlashInfer's rejection, it is not a verified support claim for the head + counts outside the range. For Kimi K3 the open side above 128 is + unreachable anyway — per-rank Q heads never exceed 96 (all heads + replicated under attention-DP, `96 / tp_size` under TEP head sharding). + """ + if backend == "trtllm-gen" and 64 < num_heads < 128: + raise ValueError( + "Kimi K3 MLA: the trtllm-gen generation backend cannot run with " + f"{num_heads} query heads per rank (trtllm-gen MLA decode rejects " + "64 < num_heads_q < 128; under attention-DP every rank keeps all " + "heads). trtllm-gen was selected explicitly — by the FP8-KV-cache " + f"override or by {_KIMI_K3_MLA_GEN_BACKEND_ENV}=trtllm-gen. Use " + "tensor-parallel head sharding (TEP) so each rank has <= 64 " + "heads, or a BF16 KV cache with the default cute-dsl backend." + ) + + +def _kimi_k3_mla_decode_backend_policy( + requested_backend: str, + metadata: TrtllmAttentionMetadata, + num_gen_tokens: int, + *, + num_heads: int, +) -> str: + """Per-batch MLA decode backend selection for Kimi K3. + + Installed as ``mla_backend_policy`` on K3's generation attention backend + (see :class:`KimiK3MLAAttention`); the general attention code applies no + such policy on its own. + + CuTe-DSL reuses one staged page table across MLA layers for a + generation-only, one-token-per-request batch. Other mixed batches repeat + the staging copies in every MLA layer and regress time to first token, so + they fall back to TRTLLM-Gen. The H=96 path is the correctness exception + and applies to EVERY fallback candidate: TRTLLM-Gen may select a 64-head + Q tile, which does not divide 96 (invalid after K3's head padding was + removed), and its decode gate rejects 64 < num_heads_q < 128 outright — + falling back would fail engine initialization. H=96 per rank is K3's + attention-DP shape, so this keeps attention-DP + speculative + verification (a generation-only multi-token batch) on CuTe-DSL, which + accepts multi-token queries; K3's decode tuning preference for + TRTLLM-Gen only applies where TRTLLM-Gen is valid at all. + """ + is_single_token_generation = num_gen_tokens == metadata.num_generations + requires_cute_dsl = num_heads == 96 + if ( + requested_backend == "cute-dsl" + and not requires_cute_dsl + and (metadata.num_contexts > 0 or not is_single_token_generation) + ): + return "trtllm-gen" + return requested_backend + + +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``. + """ + + 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, + flashinfer_mla_backend=_select_mla_generation_backend(model_config.get_quant_config()), + ) + # K3 calls forward_impl() directly to insert its output gate before + # the base row-parallel o_proj. The original executor metadata remains + # intact, so MLA performs its native mixed context/generation split. + self.register_to_config = False + + 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) + # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so + # only it needs K3's per-batch fallback policy; mha keeps the + # default trtllm-gen selection. + # Validate here rather than in _select_mla_generation_backend: the + # per-rank head count (replicated under attention-DP, sharded under + # TEP) is only authoritative once the base MLA module has built its + # generation backend. + _validate_mla_generation_backend(self.mqa.flashinfer_mla_backend, self.mqa.num_heads) + self.mqa.mla_backend_policy = partial( + _kimi_k3_mla_decode_backend_policy, + num_heads=self.mqa.num_heads, + ) + self.rotary_emb = None + self.apply_rotary_emb = False + + if dtype is not None: + _meta_safe_cast_dtype(self, dtype) + + def _apply_output_gate_and_o_proj( + self, + hidden_states: torch.Tensor, + attn_out: torch.Tensor, + ) -> torch.Tensor: + if self.use_output_gate: + attn_out = attn_out * self.g_proj(hidden_states).sigmoid() + return self.o_proj(attn_out) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + # _create_outputs() rather than create_output(): the base implementation + # takes a list so a sparse-attention backend can append its own buffers, + # and it routes through the sparse hooks when they are installed. The + # dense path this module uses is element 0. + attn_outputs = self._create_outputs(hidden_states, attn_metadata) + super().forward_impl( + None, + hidden_states, + attn_metadata, + attn_output=attn_outputs, + ) + return self._apply_output_gate_and_o_proj(hidden_states, attn_outputs[0]) + + # --------------------------------------------------------------------------- # MLA runtime. # --------------------------------------------------------------------------- @@ -1502,8 +1850,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 7e1682366863..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ /dev/null @@ -1,369 +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 - -import os -from functools import partial -from typing import Optional - -import torch - -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 -from ...model_config import ModelConfig -from ..linear import Linear, TensorParallelMode -from ..mla import MLA - -_KIMI_K3_MLA_GEN_BACKEND_ENV = "TLLM_K3_MLA_GEN_BACKEND" -_KIMI_K3_MLA_GEN_BACKENDS = ("cute-dsl", "trtllm-gen") - - -def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str: - """Select K3's absorbed-generation MLA backend. - - K3 was tuned with the FlashInfer CuTe-DSL backend for BF16 KV cache. - FP8 KV cache carries device scales that CuTe-DSL does not support, so - retain the TRTLLM-Gen fallback used by the pre-refactor implementation. - """ - backend = os.environ.get(_KIMI_K3_MLA_GEN_BACKEND_ENV, "cute-dsl") - # Validate here, where the env var is read: an invalid value would - # otherwise surface only deep inside attention-backend construction, - # with an error that never names the knob that caused it. - if backend not in _KIMI_K3_MLA_GEN_BACKENDS: - raise ValueError( - f"{_KIMI_K3_MLA_GEN_BACKEND_ENV}={backend!r} is invalid; " - f"expected one of {list(_KIMI_K3_MLA_GEN_BACKENDS)}." - ) - has_fp8_kv_cache = bool( - quant_config is not None and quant_config.layer_quant_mode.has_fp8_kv_cache() - ) - if has_fp8_kv_cache and backend != "trtllm-gen": - # info_once: this runs once per MLA layer (~60x at startup for - # FP8-KV) and the decision is identical for every layer. - logger.info_once( - "Kimi K3 MLA: FP8 KV cache requires the trtllm-gen MLA " - f"generation backend; overriding '{backend}' -> 'trtllm-gen'.", - key="kimi_k3_mla_gen_backend_fp8_override", - ) - return "trtllm-gen" - return backend - - -def _validate_mla_generation_backend(backend: str, num_heads: int) -> None: - """Fail fast when `backend` can never run at this per-rank head count. - - FlashInfer's `trtllm_batch_decode_with_kv_cache_mla` rejects - `64 < num_heads_q < 128` for every batch shape, and the per-batch policy - never demotes away from an explicit `trtllm-gen` selection (the - FP8-KV-cache override or a `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen` request). - Without this check the conflict only surfaces as a FlashInfer error deep - in attention warmup. - - The bound mirrors FlashInfer's validation gate verbatim: it predicts - FlashInfer's rejection, it is not a verified support claim for the head - counts outside the range. For Kimi K3 the open side above 128 is - unreachable anyway — per-rank Q heads never exceed 96 (all heads - replicated under attention-DP, `96 / tp_size` under TEP head sharding). - """ - if backend == "trtllm-gen" and 64 < num_heads < 128: - raise ValueError( - "Kimi K3 MLA: the trtllm-gen generation backend cannot run with " - f"{num_heads} query heads per rank (trtllm-gen MLA decode rejects " - "64 < num_heads_q < 128; under attention-DP every rank keeps all " - "heads). trtllm-gen was selected explicitly — by the FP8-KV-cache " - f"override or by {_KIMI_K3_MLA_GEN_BACKEND_ENV}=trtllm-gen. Use " - "tensor-parallel head sharding (TEP) so each rank has <= 64 " - "heads, or a BF16 KV cache with the default cute-dsl backend." - ) - - -def _kimi_k3_mla_decode_backend_policy( - requested_backend: str, - metadata: TrtllmAttentionMetadata, - num_gen_tokens: int, - *, - num_heads: int, -) -> str: - """Per-batch MLA decode backend selection for Kimi K3. - - Installed as ``mla_backend_policy`` on K3's generation attention backend - (see :class:`KimiK3MLAAttention`); the general attention code applies no - such policy on its own. - - CuTe-DSL reuses one staged page table across MLA layers for a - generation-only, one-token-per-request batch. Other mixed batches repeat - the staging copies in every MLA layer and regress time to first token, so - they fall back to TRTLLM-Gen. The H=96 path is the correctness exception - and applies to EVERY fallback candidate: TRTLLM-Gen may select a 64-head - Q tile, which does not divide 96 (invalid after K3's head padding was - removed), and its decode gate rejects 64 < num_heads_q < 128 outright — - falling back would fail engine initialization. H=96 per rank is K3's - attention-DP shape, so this keeps attention-DP + speculative - verification (a generation-only multi-token batch) on CuTe-DSL, which - accepts multi-token queries; K3's decode tuning preference for - TRTLLM-Gen only applies where TRTLLM-Gen is valid at all. - """ - is_single_token_generation = num_gen_tokens == metadata.num_generations - requires_cute_dsl = num_heads == 96 - if ( - requested_backend == "cute-dsl" - and not requires_cute_dsl - and (metadata.num_contexts > 0 or not is_single_token_generation) - ): - return "trtllm-gen" - return requested_backend - - -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, - flashinfer_mla_backend=_select_mla_generation_backend(model_config.get_quant_config()), - ) - # K3 calls forward_impl() directly to insert its output gate before - # the base row-parallel o_proj. The original executor metadata remains - # intact, so MLA performs its native mixed context/generation split. - self.register_to_config = False - - 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) - # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so - # only it needs K3's per-batch fallback policy; mha keeps the - # default trtllm-gen selection. - # Validate here rather than in _select_mla_generation_backend: the - # per-rank head count (replicated under attention-DP, sharded under - # TEP) is only authoritative once the base MLA module has built its - # generation backend. - _validate_mla_generation_backend(self.mqa.flashinfer_mla_backend, self.mqa.num_heads) - self.mqa.mla_backend_policy = partial( - _kimi_k3_mla_decode_backend_policy, - num_heads=self.mqa.num_heads, - ) - self.rotary_emb = None - self.apply_rotary_emb = False - - if dtype is not None: - _meta_safe_cast_dtype(self, dtype) - - def _apply_output_gate_and_o_proj( - self, - hidden_states: torch.Tensor, - attn_out: torch.Tensor, - ) -> torch.Tensor: - if self.use_output_gate: - attn_out = attn_out * self.g_proj(hidden_states).sigmoid() - return self.o_proj(attn_out) - - def forward( - self, - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - # _create_outputs() rather than create_output(): the base implementation - # takes a list so a sparse-attention backend can append its own buffers, - # and it routes through the sparse hooks when they are installed. The - # dense path this module uses is element 0. - attn_outputs = self._create_outputs(hidden_states, attn_metadata) - super().forward_impl( - None, - hidden_states, - attn_metadata, - attn_output=attn_outputs, - ) - return self._apply_output_gate_and_o_proj(hidden_states, attn_outputs[0]) diff --git a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py index c7628cc1c7cf..f03286a2d813 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/moe/fused_moe/communication/communication_factory.py @@ -88,7 +88,6 @@ def create_strategy( alltoall_result_do_sum: bool = True, use_flashinfer: bool = False, hidden_size: Optional[int] = None, - communication_method: Optional[str] = None, ) -> Optional[Communication]: """ Create the best communication method for the given configuration @@ -114,8 +113,6 @@ def create_strategy( hidden_size: Actual MoE activation dimension (the A2A payload width). For latent-MoE models this is moe_latent_size, not pretrained_config.hidden_size. Falls back to pretrained_config.hidden_size when not provided. - communication_method: Optional model-selected communication method. - ``TRTLLM_FORCE_COMM_METHOD`` takes precedence when set. # TODO: Need a way to indicate whether EPLB is enabled. Returns: @@ -151,15 +148,7 @@ def create_strategy( if mapping.moe_tp_size != 1: return AllGatherReduceScatter(mapping) - # A forced method comes either from the environment, which wins, or from the - # model-selected argument. Keep the source with the value so the log below can - # name the one the reader can actually go and change. - env_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD") - if env_method is not None: - force_method, force_source = env_method, "TRTLLM_FORCE_COMM_METHOD" - else: - force_method, force_source = communication_method, "communication_method" - + force_method = os.environ.get("TRTLLM_FORCE_COMM_METHOD") if force_method is not None: strategy = CommunicationFactory._create_forced_method( force_method, @@ -175,7 +164,7 @@ def create_strategy( ) logger.info( f"Selected communication strategy: {strategy.__class__.__name__} " - f"({force_source}={force_method})" + f"(TRTLLM_FORCE_COMM_METHOD={force_method})" ) return strategy diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index bd44358d74fa..9709e62e9a64 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -168,7 +168,6 @@ def __init__( trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, trtllm_gen_activation_alpha: Optional[float] = None, trtllm_gen_activation_beta: Optional[float] = None, - communication_method: Optional[str] = None, **kwargs, ): super().__init__( @@ -190,7 +189,6 @@ def __init__( # Store model_config and aux_stream_dict for later use (e.g., backend setter) self.model_config = model_config self.aux_stream_dict = aux_stream_dict - self.communication_method = communication_method # If True, the router weight will be multiplied on the input rather than at the end of FC2 self.apply_router_weight_on_input = apply_router_weight_on_input @@ -640,7 +638,6 @@ def _create_comm_strategy_auto(self) -> Optional[Communication]: alltoall_result_do_sum=True, use_flashinfer=self.use_flashinfer, hidden_size=self.hidden_size, - communication_method=self.communication_method, ) def forward_impl( diff --git a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py index 66d9871792d0..565ac4ecccda 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/create_moe.py @@ -383,7 +383,6 @@ def create_moe( trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None, trtllm_gen_activation_alpha: Optional[float] = None, trtllm_gen_activation_beta: Optional[float] = None, - communication_method: Optional[str] = None, ) -> MoE | VanillaMoE: """ Create MoE instance with automatic parameter inference from model_config. @@ -413,7 +412,6 @@ def create_moe( trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type trtllm_gen_activation_alpha: Optional backend-local activation alpha trtllm_gen_activation_beta: Optional backend-local activation beta - communication_method: Optional ConfigurableMoE communication method Returns: A complete MoE layer: a ``MoE`` (``ConfigurableMoE`` around an @@ -508,13 +506,10 @@ def create_moe( trtllm_gen_activation_type=trtllm_gen_activation_type, trtllm_gen_activation_alpha=trtllm_gen_activation_alpha, trtllm_gen_activation_beta=trtllm_gen_activation_beta, - communication_method=communication_method, ) # TritonFusedMoE and VanillaMoE are not wrapped by ConfigurableMoE # and own their communication and forward paths. - if communication_method is not None: - raise ValueError("communication_method requires ConfigurableMoE.") return create_moe_backend( moe_cls=moe_cls, routing_method=routing_method, diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 33f32cf693c5..3223d8d81679 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -8,9 +8,9 @@ import torch from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention -from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import ( +from tensorrt_llm._torch.models.modeling_kimi_linear import ( _KIMI_K3_MLA_GEN_BACKEND_ENV, + KimiK3MLAAttention, _kimi_k3_mla_decode_backend_policy, _select_mla_generation_backend, _validate_mla_generation_backend, diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 01bd5f81f6e9..563bc3f297e1 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -47,7 +47,6 @@ import tensorrt_llm._torch.models.modeling_kimi_linear as modeling_kimi_linear from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3MoERuntime -from tensorrt_llm._torch.moe.fused_moe.communication import CommunicationFactory from tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm import _MEGA_MOE_SYMM_BUFFER_CACHE from tensorrt_llm._torch.utils import ActType_TrtllmGen from tensorrt_llm._utils import get_free_port, get_sm_version @@ -202,52 +201,6 @@ def test_padded_fused_shapes(): assert padded_fused_shapes(2880, 96) == (3072, 2944, 128) -def test_communication_factory_accepts_model_selected_method(monkeypatch): - mapping = SimpleNamespace( - enable_attention_dp=True, - dp_size=16, - moe_tp_size=1, - moe_ep_size=16, - has_cp_helix=lambda: False, - ) - model_config = SimpleNamespace( - mapping=mapping, - pretrained_config=SimpleNamespace(hidden_size=3584), - torch_dtype=torch.bfloat16, - quant_config=None, - max_num_tokens=4096, - moe_max_num_tokens=65536, - use_cuda_graph=False, - use_low_precision_moe_combine=False, - ) - selected = object() - method = None - - def create_forced_method(force_method, *args, **kwargs): - nonlocal method - method = force_method - return selected - - monkeypatch.delenv("TRTLLM_FORCE_COMM_METHOD", raising=False) - monkeypatch.setattr( - CommunicationFactory, - "_create_forced_method", - staticmethod(create_forced_method), - ) - actual = CommunicationFactory.create_strategy( - model_config=model_config, - num_experts=896, - num_slots=896, - top_k=16, - expert_size_per_partition=56, - hidden_size=3584, - communication_method="ALLGATHER", - ) - - assert method == "ALLGATHER" - assert actual is selected - - @situ_supported def test_make_situ_alpha_beta_contract(): alpha, beta = make_situ_alpha_beta( @@ -712,7 +665,6 @@ def _make_routed_moe( else QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) ), layer_idx=0, - communication_method=None, ) if moe_backend == "TRTLLM": moe_kwargs.update( From 00426dc73e5058c22c95ed09053402491c3b2edb Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 25 Aug 2026 04:18:43 +0000 Subject: [PATCH 2/3] Address trivial review comments Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 170bca1c0e6f..c6c0b4aa8a8e 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1783,8 +1783,12 @@ def __init__( # 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) + if not isinstance(self.mha, TrtllmAttention) or not isinstance(self.mqa, TrtllmAttention): + raise ValueError( + "Kimi K3 MLA requires the TRTLLM attention backend; its NoPE " + "identity RoPE table is installed on TrtllmAttention only. Got " + f"mha={type(self.mha).__name__}, mqa={type(self.mqa).__name__}." + ) _install_identity_rope_table(self.mha) _install_identity_rope_table(self.mqa) # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so From d387e2865e52c8274e020f4ef245fb7b5e488782 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 28 Aug 2026 08:22:39 -0500 Subject: [PATCH 3/3] [TRTLLM-15177][chore] Kimi K3: address MLA review feedback - MLA decode policy: gate the H=96 CuTe-DSL requirement purely on head count so generation-only speculative batches (num_contexts == 0) can no longer fall through to TRTLLM-Gen's invalid 64-head Q tile. - _meta_safe_cast_dtype: preserve quantized weights (fp8/fp4), their float32 scale/dequant buffers, and nvfp4 alpha so the compute-dtype cast no longer breaks the quantized-kernel contract. - _install_identity_rope_table: clone the shared cached cos/sin table before writing identity values (the rope cache keys on RopeParams, so equal-param backends/layers share one tensor), on both install and resize paths. - kimi_linear config: revert keys_to_ignore_at_inference to a list to match the base PretrainedConfig type and every sibling config. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/configs/kimi_linear.py | 4 +- .../_torch/models/modeling_kimi_linear.py | 42 ++++++++++++++++++- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/configs/kimi_linear.py b/tensorrt_llm/_torch/configs/kimi_linear.py index 1513a3614a1d..b00fa6627cd0 100644 --- a/tensorrt_llm/_torch/configs/kimi_linear.py +++ b/tensorrt_llm/_torch/configs/kimi_linear.py @@ -17,9 +17,7 @@ class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" - # Tuple, not list: a class-level mutable default could be appended to - # in place, changing the default for every instance. - keys_to_ignore_at_inference = ("past_key_values",) + keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c6c0b4aa8a8e..c7fcc51fc165 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -81,6 +81,7 @@ import copy import gc +import itertools import json import math import os @@ -1597,18 +1598,48 @@ def _kimi_k3_mla_decode_backend_policy( def _meta_safe_cast_dtype(module, dtype): - """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. + """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``, + leaving quantized weights and their scales untouched. ``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``. + + ``skip_create_weights_in_init`` defaults to ``False``, so the MLA + projections may already hold quantized weights (fp8/fp4) plus float32 + scale/dequant buffers, and MLA's own ``k_b_proj_trans`` may be fp8 with + a float32 scale. Widening any of those to the compute dtype breaks the + quantized-kernel contract, so they are preserved: + + - every parameter/buffer of a quantized ``Linear`` (covers nvfp4 + ``alpha``, whose name carries no ``scale`` hint), + - any 1-byte float tensor (fp8 weights, incl. bare ``k_b_proj_trans``), + - any ``scale``/``dequant`` buffer (covers ``k_b_proj_trans_scale`` etc. + that live directly on the MLA module, not inside a ``Linear``). """ + preserve_ids = set() + for submod in module.modules(): + if ( + isinstance(submod, TrtllmLinear) + and getattr(submod, "_weights_created", False) + and submod.has_any_quant + ): + for t in itertools.chain( + submod.parameters(recurse=False), submod.buffers(recurse=False) + ): + preserve_ids.add(id(t)) + for name, t in itertools.chain(module.named_parameters(), module.named_buffers()): + if "scale" in name or "dequant" in name: + preserve_ids.add(id(t)) def _cast(t): if not t.is_floating_point(): return t + # 1-byte floats are quantized (fp8); never widen them. + if t.element_size() == 1 or id(t) in preserve_ids: + return t if t.is_meta: return torch.empty_like(t, dtype=dtype) return t.to(dtype=dtype) @@ -1689,6 +1720,12 @@ def _install_identity_rope_table(backend: TrtllmAttention) -> None: "backend.rotary_cos_sin is None after construction; check " "pos_embd_params has a valid RopeParams with dim > 0." ) + # The rope cache keys on (RopeParams, interleave), so mha, mqa, and every + # K3 MLA layer with equal params share ONE cached cos/sin tensor. Clone + # before overwriting so the in-place identity write never mutates the + # shared cached sinusoidal table (which other consumers may read). + cos_sin = cos_sin.clone() + backend.rotary_cos_sin = cos_sin _write_identity_rope_values(cos_sin) orig_resize = backend._ensure_rope_table_size # bound method @@ -1697,6 +1734,9 @@ def _identity_preserving_resize(required_max_positions: int) -> None: if required_max_positions <= backend.rope_params.max_positions: return orig_resize(required_max_positions) + # orig_resize regenerates from the shared cache again; clone the fresh + # table before rewriting identity so the cache stays pristine. + backend.rotary_cos_sin = backend.rotary_cos_sin.clone() _write_identity_rope_values(backend.rotary_cos_sin) backend._ensure_rope_table_size = _identity_preserving_resize