[None][perf] Kimi K3: overlap KDA/MLA attention decode onto the aux stream

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -845,7 +845,8 @@ class KimiKDARuntime(nn.Module):
     """
 
     def __init__(self, cfg, layer_idx: int, mapping=None,
-                 allreduce_strategy=AllReduceStrategy.AUTO):
+                 allreduce_strategy=AllReduceStrategy.AUTO,
+                 aux_stream: Optional[torch.cuda.Stream] = None):
         super().__init__()
         # Lazy import: pulls in fla/einops.
         from ..modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention
@@ -905,6 +906,15 @@ def __init__(self, cfg, layer_idx: int, mapping=None,
         # Persistent batch-row-dense staging for the fused decode kernel's
         # per-section conv windows (lazily sized to the pool slot count).
         self._cs_dense: Optional[torch.Tensor] = None
+        # Side stream (+ fork/join events) for the decode overlap in
+        # ``_forward_decode``: the in-projection GEMV chain against the
+        # conv-window/recurrent-state staging. Only engaged when multi-stream
+        # is active (i.e. under CUDA graphs) and a stream was supplied;
+        # otherwise both legs run in order on the default stream with an
+        # identical result.
+        self.aux_stream = aux_stream
+        self._kda_main_event = torch.cuda.Event()
+        self._kda_stage_event = torch.cuda.Event()
 
     def finalize_decode_weights(self) -> None:
         """Build the decode fast-path constants (once, after weight load).
@@ -1241,35 +1251,68 @@ def _forward_decode(self, x2d, conv_pool, ssm_pool, slot_indices,
                               device=x2d.device)
             self._cs_dense = buf
 
-        if self._in_proj_weight is not None:
-            # One GEMV over [q | k | v | g | f_a | b]; slices below are views.
-            proj = torch.nn.functional.linear(x2d, self._in_proj_weight)
-            x_qkv = proj[:, :3 * d]
-            onorm_g = proj[:, 3 * d:4 * d]
-            f_a = proj[:, 4 * d:4 * d + hd]
-            beta = proj[:, 4 * d + hd:4 * d + hd + H]
-        else:
-            # FP8 weight read (KIMI_K3_KDA_GLUE_FP8=1): the loader's fused
-            # FP8 [q | k | v | g] GEMM plus one BF16 GEMV over [f_a | b];
-            # slices below are views.
-            qkvg = mixer.qkvg_proj(x2d)
-            small = torch.nn.functional.linear(x2d,
-                                               self._in_proj_small_weight)
-            x_qkv = qkvg[:, :3 * d]
-            onorm_g = qkvg[:, 3 * d:4 * d]
-            f_a = small[:, :hd]
-            beta = small[:, hd:hd + H]
-        g = mixer.f_b_proj(f_a)  # [B, d]
-
-        # Gather the HF-layout conv windows once, then repack the
-        # historical W-1 columns into the kernel's dense per-section
-        # [B, d, W-1] layout (single strided copy kernel).
-        cs = conv_pool.index_select(0, slot_indices)  # [B, 3d, W]
-        cs_dense = buf[:, :B]
-        cs_dense.copy_(cs.view(B, 3, d, W)[:, :, :, 1:].permute(1, 0, 2, 3))
-
-        state = (ssm_pool if ssm_state_indices is not None else
-                 ssm_pool.index_select(0, slot_indices))
+        def _in_projection():
+            """In-projection GEMV chain: reads only ``x2d``."""
+            if self._in_proj_weight is not None:
+                # One GEMV over [q | k | v | g | f_a | b]; slices are views.
+                proj = torch.nn.functional.linear(x2d, self._in_proj_weight)
+                x_qkv = proj[:, :3 * d]
+                onorm_g = proj[:, 3 * d:4 * d]
+                f_a = proj[:, 4 * d:4 * d + hd]
+                beta = proj[:, 4 * d + hd:4 * d + hd + H]
+            else:
+                # FP8 weight read (KIMI_K3_KDA_GLUE_FP8=1): the loader's
+                # fused FP8 [q | k | v | g] GEMM plus one BF16 GEMV over
+                # [f_a | b]; slices below are views.
+                qkvg = mixer.qkvg_proj(x2d)
+                small = torch.nn.functional.linear(x2d,
+                                                   self._in_proj_small_weight)
+                x_qkv = qkvg[:, :3 * d]
+                onorm_g = qkvg[:, 3 * d:4 * d]
+                f_a = small[:, :hd]
+                beta = small[:, hd:hd + H]
+            return x_qkv, onorm_g, mixer.f_b_proj(f_a), beta
+
+        def _stage_conv_windows():
+            """Conv-window gather + repack, plus the recurrent-state gather.
+
+            Reads only the cache pools and ``slot_indices`` — nothing the
+            in-projection produces — so the two halves are data-independent
+            and may run concurrently. Gathers the HF-layout conv windows
+            once, then repacks the historical W-1 columns into the kernel's
+            dense per-section [B, d, W-1] layout (one strided copy).
+            """
+            cs = conv_pool.index_select(0, slot_indices)  # [B, 3d, W]
+            cs_dense = buf[:, :B]
+            cs_dense.copy_(cs.view(B, 3, d, W)[:, :, :, 1:].permute(
+                1, 0, 2, 3))
+            state = (ssm_pool if ssm_state_indices is not None else
+                     ssm_pool.index_select(0, slot_indices))
+            return cs, cs_dense, state
+
+        # The in-projection GEMV chain and the cache staging read disjoint
+        # inputs and write disjoint outputs, but the KDA decode kernel needs
+        # both — so under CUDA graphs issue them on two streams and let the
+        # staging hide behind the GEMV instead of following it. Sequential
+        # (identical result) when multi-stream is off; see the shared/routed
+        # expert overlap in ``KimiK3MoERuntime.forward`` for the same pattern.
+        ((x_qkv, onorm_g, g, beta),
+         (cs, cs_dense, state)) = maybe_execute_in_parallel(
+             _in_projection,
+             _stage_conv_windows,
+             self._kda_main_event,
+             self._kda_stage_event,
+             self.aux_stream,
+             disable_on_compile=True)
+        if self.aux_stream is not None:
+            # ``cs`` and the gathered ``state`` were allocated on the side
+            # stream but are consumed below on this one; tell the caching
+            # allocator they are still live here so it cannot recycle their
+            # storage into a later side-stream allocation mid-use.
+            cur_stream = torch.cuda.current_stream()
+            cs.record_stream(cur_stream)
+            if state is not ssm_pool:
+                state.record_stream(cur_stream)
 
         o = mixer._dispatch.decode_kda(
             x_q=x_qkv[:, :d].unflatten(-1, (H, hd)).unsqueeze(0),
@@ -1653,7 +1696,8 @@ class KimiMLARuntime(nn.Module):
     """Wraps ``KimiK3MLAAttention`` with the mixed-batch executor dispatch."""
 
     def __init__(self, cfg, layer_idx: int, mapping=None, quant_config=None,
-                 allreduce_strategy=AllReduceStrategy.AUTO):
+                 allreduce_strategy=AllReduceStrategy.AUTO,
+                 aux_stream: Optional[torch.cuda.Stream] = None):
         super().__init__()
         import os
 
@@ -1713,6 +1757,10 @@ def __init__(self, cfg, layer_idx: int, mapping=None, quant_config=None,
             use_output_gate=cfg.mla_use_output_gate,
             max_position_embeddings=max_positions,
             quant_config=quant_config,
+            # Shared side stream for the decode-path Q/KV-projection and
+            # bmm/rope overlaps; the MLA forks join before the MoE block
+            # reuses the same stream, so sharing it is safe.
+            aux_stream=aux_stream,
         )
 
     def forward(self, hidden_states: torch.Tensor,
@@ -1768,7 +1816,11 @@ def __init__(self,
                 cfg,
                 layer_idx,
                 mapping=model_config.mapping,
-                allreduce_strategy=model_config.allreduce_strategy)
+                allreduce_strategy=model_config.allreduce_strategy,
+                # Shared side stream for the decode in-projection / cache
+                # staging overlap. The attention fork always joins before the
+                # MoE block reuses the same stream, so sharing it is safe.
+                aux_stream=aux_stream)
         else:
             # Forward only the KV-cache quantization to the MLA attention
             # backends (enables FP8 KV cache). The attention projection
@@ -1786,7 +1838,8 @@ def __init__(self,
                 layer_idx,
                 mapping=model_config.mapping,
                 quant_config=mla_quant_config,
-                allreduce_strategy=model_config.allreduce_strategy)
+                allreduce_strategy=model_config.allreduce_strategy,
+                aux_stream=aux_stream)
 
         self.is_moe = (cfg.num_experts is not None
                        and layer_idx >= cfg.first_k_dense_replace
@@ -1923,9 +1976,11 @@ def __init__(self, model_config: ModelConfig):
         self._text_cfg = cfg
         dtype = torch.bfloat16
 
-        # One side stream shared across all layers, used by KimiK3MoERuntime
-        # to overlap the replicated shared-expert compute with the routed EP
-        # dispatch/combine collectives.
+        # One side stream shared across all layers: KimiK3MoERuntime overlaps
+        # the replicated shared-expert compute with the routed EP
+        # dispatch/combine collectives, and the attention runtimes overlap
+        # their decode-path independent halves on it. Each fork joins before
+        # the next one starts, so a single stream serves all of them.
         self.aux_stream = torch.cuda.Stream()
 
         self.embed_tokens = nn.Embedding(cfg.vocab_size,

tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py

@@ -89,6 +89,7 @@
     PredefinedAttentionMask,
     RopeParams,
 )
+from ..multi_stream_utils import maybe_execute_in_parallel
 from ..rms_norm import RMSNorm
 
 
@@ -328,6 +329,7 @@ def __init__(
         apply_rotary_mutation: bool = False,
         omit_output_gate_mutation: bool = False,
         quant_config: Optional[QuantConfig] = None,
+        aux_stream: Optional[torch.cuda.Stream] = None,
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -345,6 +347,15 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
         self.apply_rotary_mutation = apply_rotary_mutation
         self.omit_output_gate_mutation = omit_output_gate_mutation
+        # Side stream (+ fork/join events) for the decode-path overlaps in
+        # ``_forward_decode_fused``: the Q projection chain against the KV
+        # projection chain, and the absorb bmm against ``mla_rope_generation``
+        # (the DeepSeek MLA pattern, mla.py). Engaged only when multi-stream
+        # is active (i.e. under CUDA graphs) and a stream was supplied;
+        # otherwise both legs run in order with an identical result.
+        self.aux_stream = aux_stream
+        self._proj_events = (torch.cuda.Event(), torch.cuda.Event())
+        self._rope_events = (torch.cuda.Event(), torch.cuda.Event())
 
         self.q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False)
         self.q_a_layernorm = RMSNorm(hidden_size=q_lora_rank, eps=rms_norm_eps, dtype=dtype)
@@ -865,8 +876,26 @@ def _forward_decode_fused(
         """
         num_tokens = hidden_states.shape[0]
         metadata = rt.metadata
-        q_nope, q_rot = self._project_q_unabsorbed(hidden_states)
-        _, _, latent_cache = self._project_kv_and_latent(hidden_states)
+        # The Q chain (q_a_proj -> q_a_layernorm -> q_b_proj) and the KV chain
+        # (kv_a_proj_with_mqa -> kv_a_layernorm -> cat) both read only
+        # ``hidden_states`` and write disjoint tensors, so they are fully
+        # data-independent. DeepSeek MLA fuses q_a and kv_a into one GEMM and
+        # can therefore overlap only the two layernorms; K3 keeps them as
+        # separate projections, so the whole chains overlap.
+        (q_nope, q_rot), (_, _, latent_cache) = maybe_execute_in_parallel(
+            lambda: self._project_q_unabsorbed(hidden_states),
+            lambda: self._project_kv_and_latent(hidden_states),
+            self._proj_events[0],
+            self._proj_events[1],
+            self.aux_stream,
+            disable_on_compile=True,
+        )
+        if self.aux_stream is not None:
+            # ``latent_cache`` was allocated on the side stream but is consumed
+            # on this one (rope generation, then the FMHA); tell the caching
+            # allocator it is still live here so its storage cannot be recycled
+            # into a later side-stream allocation while still in use.
+            latent_cache.record_stream(torch.cuda.current_stream())
 
         k_absorb, v_absorb = self._kv_b_absorb_split()
 
@@ -877,24 +906,19 @@ def _forward_decode_fused(
         fused_q = hidden_states.new_empty(
             (num_tokens, self.num_heads, self.kv_lora_rank + self.qk_rope_head_dim)
         )
-        # Absorb: [H, T, m] × [H, m, kv] -> [H, T, kv] (m = qk_nope_head_dim),
-        # written directly into fused_q's latent slice.
-        torch.ops.trtllm.bmm_out(
-            q_nope.transpose(0, 1), k_absorb, fused_q[..., : self.kv_lora_rank].transpose(0, 1)
-        )
-
         num_seqs = metadata.num_seqs
         cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=hidden_states.device)
         cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=hidden_states.device)
         fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=hidden_states.device)
+        has_fp8_kv_cache = getattr(self._backend_gen, "has_fp8_kv_cache", False)
         # FP8 KV cache: the rope-generation op additionally quantizes the
         # appended latent row and fused_q, emitting the FMHA dequant scales
         # and the quantized-Q buffer (same contract as the DeepSeek MLA
         # decode path in mla.py).
         mla_bmm1_scale = None
         mla_bmm2_scale = None
         quant_q_buffer = None
-        if getattr(self._backend_gen, "has_fp8_kv_cache", False):
+        if has_fp8_kv_cache:
             mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=hidden_states.device)
             mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=hidden_states.device)
             quant_q_buffer = torch.empty(
@@ -904,21 +928,58 @@ def _forward_decode_fused(
                 dtype=torch.uint8,
                 device=hidden_states.device,
             )
-        # Identity rope (NoPE): writes q_rot into fused_q's rope tail,
-        # appends the latent row to the paged cache, and fills the
-        # scheduler buffers — one op, CUDA-graph-safe (the DeepSeek MLA
-        # generation path uses it under graphs in production).
-        self._backend_gen.mla_rope_generation(
-            fused_q,
-            q_rot_use,
-            latent_cache,
-            metadata,
-            cu_q_seqlens,
-            cu_kv_seqlens,
-            fmha_scheduler_counter,
-            mla_bmm1_scale,
-            mla_bmm2_scale,
-            quant_q_buffer,
+
+        def _absorb_bmm():
+            # Absorb: [H, T, m] × [H, m, kv] -> [H, T, kv]
+            # (m = qk_nope_head_dim), written directly into fused_q's latent
+            # slice.
+            torch.ops.trtllm.bmm_out(
+                q_nope.transpose(0, 1), k_absorb, fused_q[..., : self.kv_lora_rank].transpose(0, 1)
+            )
+
+        def _rope_generation():
+            # Identity rope (NoPE): writes q_rot into fused_q's rope tail,
+            # appends the latent row to the paged cache, and fills the
+            # scheduler buffers — one op, CUDA-graph-safe (the DeepSeek MLA
+            # generation path uses it under graphs in production).
+            self._backend_gen.mla_rope_generation(
+                fused_q,
+                q_rot_use,
+                latent_cache,
+                metadata,
+                cu_q_seqlens,
+                cu_kv_seqlens,
+                fmha_scheduler_counter,
+                mla_bmm1_scale,
+                mla_bmm2_scale,
+                quant_q_buffer,
+            )
+
+        # The absorb bmm writes only fused_q[..., :kv_lora_rank] while the rope
+        # kernel writes only the rope tail (mlaKernels.cu offsets its q store
+        # by +kv_lora_rank), the paged cache, and the scheduler buffers — so
+        # with a BF16 KV cache the two are data-independent and overlap on
+        # disjoint slices with no extra copy.
+        #
+        # NOT so under an FP8 KV cache: the extra CTAs that the grid gains for
+        # cache_type == FP8 read fused_q's latent slice back to quantize it
+        # into quant_q_buffer, i.e. they consume the bmm's output. Overlapping
+        # would race, so keep them serialized there — the same gate the
+        # DeepSeek MLA decode path applies in mla.py.
+        rope_stream = None if has_fp8_kv_cache else self.aux_stream
+        if rope_stream is not None:
+            # These were allocated on this stream but the rope leg writes them
+            # from the side stream, so mark them live there too before the fork.
+            for t in (fused_q, cu_q_seqlens, cu_kv_seqlens,
+                      fmha_scheduler_counter):
+                t.record_stream(rope_stream)
+        maybe_execute_in_parallel(
+            _absorb_bmm,
+            _rope_generation,
+            self._rope_events[0],
+            self._rope_events[1],
+            rope_stream,
+            disable_on_compile=True,
         )
 
         forward_args = AttentionForwardArgs(