From 295dfe7ad2b0a21f06689166760de358a2ee4857 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 30 Jul 2026 02:28:01 +0000 Subject: [PATCH] [None][perf] Kimi K3: overlap KDA/MLA attention decode onto the aux stream Route the model's existing side CUDA stream into KimiKDARuntime and KimiK3MLAAttention and fork three data-independent decode op pairs via maybe_execute_in_parallel. Before this change the attention layers ran 100% single-stream: 69 KDA segments/step held 4.667 ms of device-busy time with zero two-stream residency, and 24 MLA segments held 2.161 ms, likewise zero. The single side stream created for the layer loop was passed only to the MoE runtime, so the identical sm100_fp8_fp4_gemm_1d1d_impl kernel was 97% hidden inside MoE but 94% exposed inside KDA. The forks live inside forward() so the captured cuda graph records them, matching how MoE already does it. Pure scheduling: no dtype, precision, kv-cache or backend change. Measured on 16x GB300 TEP16, ISL 8192 / OSL 1024, c=1 median TPOT: 13.2106 ms -> 12.6789 ms (-4.025%, job 2720932, rack nvl72d127). Reproduced across both a rack change and a rebuild: -3.720% (job 2720192, rack d157, build 2719879) and -4.025% (job 2720932, rack d127, build 2720688). The two agree to 0.316%, against a measured c=1 noise band of 2.63%. A third observation and the full concurrency curve were taken afterwards on the same recipe (job 2721639), confirming the gain holds across the sweep and regresses nothing: c=1 13.5070 -> 12.6934 ms (-6.02%, band 2.63%) c=2 15.9358 -> 16.1229 ms (+1.17%, band 5.08% -- within noise) c=4 17.5480 -> 16.8362 ms (-4.06%, band 4.01%) c=8 22.7114 -> 21.7538 ms (-4.22%, band 3.63%) c=8 is the only genuinely independent corroborator: stream_interval=10 makes median ITL exactly 10x TPOT, so ITL, tps/user and output_throughput at c=1 are all restatements of the metric. Per-point noise bands are from three zero-change runs on the unmodified build. Post-change nsys on the same recipe confirms the forks execute on device: KDA two-stream residency 0.0000 -> 0.7016 ms/step (16.1% of its busy time), MLA 0.0000 -> 0.2829 ms/step (14.3%), with the aux stream appearing in both families' stream sets. Numerical parity: the KDA and MLA Q/KV forks were checked bitwise-identical against their serial equivalents under real cuda-graph capture on GB300. The third fork (absorb-bmm || mla_rope_generation) is argued disjoint from source -- mlaKernels.cu writes qkv_output[dst_q_idx] write-only, and both head_dim_idx and c_k=512 are multiples of ELTS_PER_VEC=8, so no vector store can straddle the bmm's slice -- but it does not yet have its own numerical probe, and no model-level greedy-token comparison has been run. Both are outstanding before this should be considered merge-ready. Candidates c-017 + c-018, campaign pa-tep-20260729. Signed-off-by: yukunh <23156053+hyukn@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 127 +++++++++++++----- .../kimi_k3_mla/kimi_k3_mla_attention.py | 109 +++++++++++---- 2 files changed, 176 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 3b260f6c3d7e..419ffce98474 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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, 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 284850306f09..df7d2bfceaf7 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 @@ -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,16 +906,11 @@ 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 @@ -894,7 +918,7 @@ def _forward_decode_fused( 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(