diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 86291f2ea770..f7c7f913c113 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1067,8 +1067,7 @@ def __init__( model_config: ModelConfig, cfg, layer_idx: int, - aux_stream: Optional[torch.cuda.Stream] = None, - moe_aux_stream_dict: Optional[Dict[AuxStreamType, torch.cuda.Stream]] = None, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], ): super().__init__() self.layer_idx = layer_idx @@ -1119,7 +1118,7 @@ def __init__( model_config=routed_moe_model_config, override_quant_config=routed_quant_config, layer_idx=layer_idx, - aux_stream_dict=moe_aux_stream_dict, + aux_stream_dict=aux_stream_dict, # Let CommunicationFactory select the best available strategy. communication_method=None, activation=SiTuActivation( @@ -1210,9 +1209,9 @@ def __init__( ) # Side stream (+ fork/join events) for overlapping shared-expert # compute with the routed chain. Only engaged when multi-stream is - # active (CUDA graphs on) and aux_stream is set; otherwise both run in - # order on the default stream. - self.aux_stream = aux_stream + # active (CUDA graphs on); otherwise both run in order on the default + # stream. + self.shared_expert_stream = aux_stream_dict[AuxStreamType.MoeShared] self.moe_main_event = torch.cuda.Event() self.moe_shared_event = torch.cuda.Event() self.routed_expert_down_proj = nn.Linear( @@ -1409,7 +1408,7 @@ def _routed_output(): lambda: self.shared_experts(identity), self.moe_main_event, self.moe_shared_event, - self.aux_stream, + self.shared_expert_stream, disable_on_compile=True, ) if self._use_combined_all_reduce: @@ -1439,6 +1438,7 @@ def __init__( cfg: "PretrainedConfig", layer_idx: int, model_config: ModelConfig, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], mapping_with_cp: Optional[Mapping] = None, ) -> None: super().__init__() @@ -1483,6 +1483,7 @@ def __init__( use_output_gate=cfg.mla_use_output_gate, max_position_embeddings=max_positions, model_config=model_config, + aux_stream_dict=aux_stream_dict, mapping_with_cp=mapping_with_cp, ) @@ -1509,8 +1510,7 @@ def __init__( model_config: ModelConfig, cfg, layer_idx: int, - aux_stream: Optional[torch.cuda.Stream] = None, - moe_aux_stream_dict: Optional[Dict[AuxStreamType, torch.cuda.Stream]] = None, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], ): super().__init__() self.layer_idx = layer_idx @@ -1528,7 +1528,7 @@ def __init__( layer_idx, mapping=model_config.mapping, allreduce_strategy=model_config.allreduce_strategy, - aux_stream=aux_stream, + aux_stream=aux_stream_dict[AuxStreamType.Attention], model_config=model_config, ) else: @@ -1550,6 +1550,7 @@ def __init__( cfg, layer_idx, model_config=mla_model_config, + aux_stream_dict=aux_stream_dict, # CP original stashed by _setup_helix_mappings; None outside helix. mapping_with_cp=getattr(model_config, "_helix_mapping_with_cp", None), ) @@ -1560,9 +1561,7 @@ def __init__( and layer_idx % getattr(cfg, "moe_layer_freq", 1) == 0 ) if self.is_moe: - self.block_sparse_moe = KimiK3MoERuntime( - model_config, cfg, layer_idx, aux_stream, moe_aux_stream_dict - ) + self.block_sparse_moe = KimiK3MoERuntime(model_config, cfg, layer_idx, aux_stream_dict) else: situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) @@ -1715,23 +1714,21 @@ def __init__(self, model_config: ModelConfig): self._text_cfg = cfg dtype = torch.bfloat16 - # Side streams shared across all layers. Keep MoE chunking separate - # from KDA/shared-expert overlap so both levels can run concurrently. - self.aux_stream = torch.cuda.Stream() - self.moe_aux_stream_dict = { - AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + # Attention and MoE phases are sequential, so their branch-overlap + # roles share one stream; MoE-internal overlap roles remain separate. + aux_stream_list = [torch.cuda.Stream() for _ in range(4)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], } self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) self.layers = nn.ModuleList( [ - KimiLinearDecoderLayer( - model_config, - cfg, - layer_idx, - self.aux_stream, - self.moe_aux_stream_dict, - ) + KimiLinearDecoderLayer(model_config, cfg, layer_idx, self.aux_stream_dict) for layer_idx in range(cfg.num_hidden_layers) ] ) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 4ab1d52f1c55..9970e71760af 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -20,6 +20,7 @@ from ...attention.backends.interface import PositionalEmbeddingParams, RopeParams from ...attention.mla import MLA from ...model_config import ModelConfig +from ...utils import AuxStreamType from ..linear import Linear, TensorParallelMode @@ -159,6 +160,7 @@ def __init__( use_output_gate: bool = True, max_position_embeddings: int = 8192, model_config: ModelConfig, + aux_stream_dict: dict[AuxStreamType, torch.cuda.Stream], mapping_with_cp: Optional[Mapping] = None, ) -> None: pos_embd_params = _make_pos_embd_params( @@ -182,6 +184,7 @@ def __init__( dtype=dtype, dense_bias=False, config=model_config, + aux_stream_dict=aux_stream_dict, mapping_with_cp=mapping_with_cp, reduce_output=False, fuse_qkv_a_proj=False, diff --git a/tests/unittest/_torch/modeling/test_kimi_linear_modeling.py b/tests/unittest/_torch/modeling/test_kimi_linear_modeling.py new file mode 100644 index 000000000000..62541347279e --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_linear_modeling.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models import modeling_kimi_linear +from tensorrt_llm._torch.utils import AuxStreamType + +pytestmark = pytest.mark.cpu_only + + +def test_kimi_linear_model_builds_shared_aux_stream_registry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + streams = [object() for _ in range(4)] + stream_iter = iter(streams) + layer_aux_stream_dicts = [] + + class _FakeDecoderLayer(nn.Module): + def __init__(self, _model_config, _config, _layer_idx, aux_stream_dict) -> None: + super().__init__() + layer_aux_stream_dicts.append(aux_stream_dict) + + monkeypatch.setattr(torch.cuda, "Stream", lambda: next(stream_iter)) + monkeypatch.setattr(modeling_kimi_linear, "KimiLinearDecoderLayer", _FakeDecoderLayer) + + config = KimiLinearConfig( + vocab_size=16, + hidden_size=8, + num_hidden_layers=3, + num_attention_heads=2, + rms_norm_eps=1e-5, + attn_res_block_size=1, + linear_attn_config={"kda_layers": [1, 3], "full_attn_layers": [2]}, + ) + model = modeling_kimi_linear.KimiLinearModel(ModelConfig(pretrained_config=config)) + + aux_stream_dict = model.aux_stream_dict + assert set(aux_stream_dict) == { + AuxStreamType.Attention, + AuxStreamType.MoeShared, + AuxStreamType.MoeChunkingOverlap, + AuxStreamType.MoeBalancer, + AuxStreamType.MoeOutputMemset, + } + assert aux_stream_dict[AuxStreamType.Attention] is streams[0] + assert aux_stream_dict[AuxStreamType.MoeShared] is streams[0] + assert aux_stream_dict[AuxStreamType.MoeChunkingOverlap] is streams[1] + assert aux_stream_dict[AuxStreamType.MoeBalancer] is streams[2] + assert aux_stream_dict[AuxStreamType.MoeOutputMemset] is streams[3] + assert len(layer_aux_stream_dicts) == config.num_hidden_layers + assert all(stream_dict is aux_stream_dict for stream_dict in layer_aux_stream_dicts) diff --git a/tests/unittest/_torch/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py index 330ee6a15039..f47a40f26ce2 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_mlp.py @@ -26,10 +26,21 @@ from tensorrt_llm._torch.modules.gated_mlp import GatedMLP from tensorrt_llm._torch.modules.situ import SituAndMul +from tensorrt_llm._torch.utils import AuxStreamType requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def _make_aux_stream_dict() -> dict[AuxStreamType, object]: + return { + AuxStreamType.Attention: object(), + AuxStreamType.MoeShared: object(), + AuxStreamType.MoeChunkingOverlap: object(), + AuxStreamType.MoeBalancer: object(), + AuxStreamType.MoeOutputMemset: object(), + } + + class _UnfusedKimiMLP(nn.Module): """HF ``KimiMLP`` layout: separate gate/up GEMMs + torch.cat.""" @@ -207,7 +218,6 @@ def test_kimi_k3_shared_expert_parallel_construction( from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_kimi_linear from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE - from tensorrt_llm._torch.utils import AuxStreamType from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -247,16 +257,18 @@ def __init__(self): moe_backend="TRTLLM", ) config = _runtime_config() - moe_aux_stream_dict = {AuxStreamType.MoeChunkingOverlap: object()} + aux_stream_dict = _make_aux_stream_dict() runtime = modeling_kimi_linear.KimiK3MoERuntime( model_config, config, layer_idx=1, - moe_aux_stream_dict=moe_aux_stream_dict, + aux_stream_dict=aux_stream_dict, ) shared = runtime.shared_experts - assert create_moe_kwargs["aux_stream_dict"] is moe_aux_stream_dict + assert create_moe_kwargs["aux_stream_dict"] is aux_stream_dict + assert AuxStreamType.MoeBalancer in create_moe_kwargs["aux_stream_dict"] + assert runtime.shared_expert_stream is aux_stream_dict[AuxStreamType.MoeShared] assert isinstance(shared, GatedMLP) assert shared.gate_up_proj.tp_size == expected_shared_tp assert shared.gate_up_proj.tp_rank == expected_shared_rank @@ -301,6 +313,7 @@ def test_kimi_k3_dense_layer_uses_gated_mlp( class _IdentityAttention(nn.Module): def __init__(self, *args, **kwargs): super().__init__() + self.aux_stream = kwargs["aux_stream"] def forward(self, hidden_states, attn_metadata): return hidden_states @@ -335,9 +348,16 @@ def forward(self, hidden_states, *args, **kwargs): activation_situ_beta=4.0, activation_situ_linear_beta=25.0, ) - layer = modeling_kimi_linear.KimiLinearDecoderLayer(model_config, config, layer_idx=0) + aux_stream_dict = _make_aux_stream_dict() + layer = modeling_kimi_linear.KimiLinearDecoderLayer( + model_config, + config, + layer_idx=0, + aux_stream_dict=aux_stream_dict, + ) assert not layer.is_moe + assert layer.linear_attn.aux_stream is aux_stream_dict[AuxStreamType.Attention] assert isinstance(layer.mlp, GatedMLP) assert layer.mlp_tp_size == expected_tp_size assert layer.mlp.gate_up_proj.tp_size == expected_tp_size