[None][fix] share Kimi auxiliary streams

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)
             ]
         )

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,

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)

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