diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 2bf00c929..9d43cfdf5 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -2037,12 +2037,12 @@ class KimiK3DeltaAttention(nn.Module): and get_is_capture_mode() and 0 < hidden_states.shape[0] <= self._bfa_bs_limit ): - # Issue the tiny [f_a|b] + f_b GEMVs on the side stream, - # then the wide [q,k,v,g] GEMM on the main stream (both - # read only hidden_states); join before the consumers. + # Fork before both branches; capture the main projection + # first to avoid CUDA graph replay stream expansion. alt = self._bfa_alt_stream cur = torch.cuda.current_stream() alt.wait_stream(cur) + fused_states, _ = self.fused_qkvg_proj(hidden_states) with torch.cuda.stream(alt): bfa = gemm(hidden_states, w) forget_gate = ( @@ -2051,7 +2051,6 @@ class KimiK3DeltaAttention(nn.Module): else gemm(bfa[..., :n_fa], self._bfa_f_b_w) ) beta = bfa[..., n_fa : n_fa + n_b] - fused_states, _ = self.fused_qkvg_proj(hidden_states) qkv, g_proj_states = torch.split( fused_states, self.split_sizes, dim=-1 ) @@ -2279,9 +2278,7 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): # DeepseekV2AttentionMLA forward cores, so wrap its forward at # the instance level (weights, reduce_results, loading untouched). self._gate_hidden_states = None - # (gate, producer stream) issued on the alt stream by forward(); - # None when the lazy path computes the gate here instead. - self._gate_precomputed = None + self._gate_pending_stream = None self._gate_alt_stream = gate_alt_stream # Above this token count the attention-core kernels fill the SMs # on their own and the overlap only adds sync overhead (same @@ -2296,19 +2293,8 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): def _gated_o_proj_forward(x, *args, **kwargs): gate_input = self._gate_hidden_states self._gate_hidden_states = None - precomputed = self._gate_precomputed - self._gate_precomputed = None - if precomputed is not None: - # Use wait_stream rather than an explicit event so the - # breakable-CUDA-graph runner can track the side-stream - # join across graph-segment boundaries. - torch.cuda.current_stream().wait_stream(precomputed[1]) if gate_input is not None and not isinstance(x, tuple): - gate = ( - precomputed[0] - if precomputed is not None - else self.g_proj(gate_input)[0] - ) + gate = self._compute_output_gate(gate_input) from sglang.kernels.ops.kimi_k3 import mla_output_gate if mla_output_gate.covered(x, gate): @@ -2317,6 +2303,10 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): x = mla_output_gate.kimi_k3_mla_output_gate(x, gate) else: x = x * torch.sigmoid(gate) + elif self._gate_pending_stream is not None: + # Even a skipped gate must close its capture branch. + torch.cuda.current_stream().wait_stream(self._gate_pending_stream) + self._gate_pending_stream = None return _orig_o_proj_forward(x, *args, **kwargs) self.o_proj.forward = _gated_o_proj_forward @@ -2339,27 +2329,29 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): return AttnForwardMethod.MLA return method - def _precompute_output_gate(self, hidden_states: torch.Tensor) -> None: - """Issue the output-gate GEMM on the alt stream so it overlaps the - attention core; the lazy path in the o_proj wrap otherwise computes - it on the critical path right before the gate multiply. The gate - tensor stays referenced via _gate_precomputed until the wrap joins, - so its memory cannot be reused while the alt stream still writes.""" - self._gate_precomputed = None + def _fork_output_gate(self, hidden_states: torch.Tensor) -> None: + """Fork early, but record the gate after attention to limit replay streams.""" + self._gate_pending_stream = None if ( self._gate_alt_stream is not None and get_is_capture_mode() - # The attention-core break ends the segment between the alt-stream - # event record and the o_proj-side wait, so under breakable capture - # the wait would cross graph segments; use the lazy path instead. + # Keep the fork and join within one capture segment. and not is_in_breakable_cuda_graph() and (0 < hidden_states.shape[0] <= self._gate_bs_limit) ): alt = self._gate_alt_stream alt.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(alt): - gate, _ = self.g_proj(hidden_states) - self._gate_precomputed = (gate, alt) + self._gate_pending_stream = alt + + def _compute_output_gate(self, hidden_states: torch.Tensor) -> torch.Tensor: + alt = self._gate_pending_stream + self._gate_pending_stream = None + if alt is None: + return self.g_proj(hidden_states)[0] + with torch.cuda.stream(alt): + gate, _ = self.g_proj(hidden_states) + torch.cuda.current_stream().wait_stream(alt) + return gate def forward( self, @@ -2371,7 +2363,7 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): ): if self.use_output_gate: self._gate_hidden_states = hidden_states - self._precompute_output_gate(hidden_states) + self._fork_output_gate(hidden_states) return super().forward( positions, hidden_states, forward_batch, zero_allocator, **kwargs ) diff --git a/test/registered/unit/models/test_kimi_k3_bfa_overlap.py b/test/registered/unit/models/test_kimi_k3_bfa_overlap.py index 1aa5520e8..ade7a3546 100644 --- a/test/registered/unit/models/test_kimi_k3_bfa_overlap.py +++ b/test/registered/unit/models/test_kimi_k3_bfa_overlap.py @@ -1,6 +1,4 @@ -"""KDA bfa side-stream overlap: forward_qkvbfg_fused must produce outputs -bit-identical to the serial path, both eager and under CUDA graph -capture/replay (the overlap only engages in capture mode).""" +"""K3 attention overlap parity under CUDA graph capture and changed-input replay.""" import unittest from types import SimpleNamespace @@ -10,6 +8,7 @@ import torch from sglang.srt.models.kimi_k3 import ( KimiK3DeltaAttention, + KimiK3MLAAttention, _get_k3_dense_weight, ) from sglang.test.ci.ci_register import register_cuda_ci @@ -67,31 +66,32 @@ class TestKimiK3BfaOverlap(CustomTestCase): def test_capture_replay_matches_serial(self): torch.manual_seed(0) for T in (1, 4, 12): - with self.subTest(T=T): - x = ( - torch.randn(T, _H, device="cuda", dtype=torch.float32) - .mul(0.05) - .to(torch.bfloat16) - ) - serial = _run(_make_owner(with_stream=False), x) - - owner = _make_owner(with_stream=True) - with patch( - "sglang.srt.models.kimi_k3.get_is_capture_mode", - return_value=True, - ): - # warm up allocations/JIT outside capture - _ = _run(owner, x) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = KimiK3DeltaAttention.forward_qkvbfg_fused(owner, x) - graph.replay() - torch.cuda.synchronize() - # note: owners share the same seeded weights - for got, ref, name in zip( - captured, serial, ("qkv", "beta", "forget_gate", "g") - ): - self.assertTrue(torch.equal(got, ref), f"T={T} {name} mismatch") + for defer_f_b in (False, True): + with self.subTest(T=T, defer_f_b=defer_f_b): + x = torch.empty(T, _H, device="cuda", dtype=torch.bfloat16) + x.normal_(std=0.05) + serial_owner = _make_owner(with_stream=False) + owner = _make_owner(with_stream=True) + forward = KimiK3DeltaAttention.forward_qkvbfg_fused + with patch( + "sglang.srt.models.kimi_k3.get_is_capture_mode", + return_value=True, + ): + # Warm up allocations/JIT outside capture. + forward(owner, x, defer_f_b=defer_f_b) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = forward(owner, x, defer_f_b=defer_f_b) + for _ in range(3): + # Changed inputs expose stale reads or missing dependencies. + x.normal_(std=0.05) + serial = forward(serial_owner, x, defer_f_b=defer_f_b) + graph.replay() + torch.cuda.synchronize() + for got, ref, name in zip( + captured, serial, ("qkv", "beta", "forget_gate", "g") + ): + self.assertTrue(torch.equal(got, ref), name) def test_eager_stream_branch_not_taken(self): x = torch.randn(3, _H, device="cuda", dtype=torch.bfloat16) @@ -100,6 +100,51 @@ class TestKimiK3BfaOverlap(CustomTestCase): for got, ref in zip(overlap, serial): self.assertTrue(torch.equal(got, ref)) + def test_mla_gate_capture_matches_serial(self): + x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) + qkv_weight = torch.randn(128, 64, device="cuda", dtype=torch.bfloat16) + gate_weight = torch.randn_like(qkv_weight) + project = torch.nn.functional.linear + owner = SimpleNamespace( + _gate_alt_stream=torch.cuda.Stream(), + _gate_bs_limit=128, + g_proj=lambda value: (project(value, gate_weight), None), + ) + + def forward(): + # Both branches must wait for this in-graph input producer. + hidden = x * 0.5 + KimiK3MLAAttention._fork_output_gate(owner, hidden) + qkv = project(hidden, qkv_weight) + gate = KimiK3MLAAttention._compute_output_gate(owner, hidden) + return qkv * torch.sigmoid(gate) + + for capture_mode, breakable in ((True, False), (True, True), (False, False)): + with ( + self.subTest(capture_mode=capture_mode, breakable=breakable), + patch( + "sglang.srt.models.kimi_k3.get_is_capture_mode", + return_value=capture_mode, + ), + patch( + "sglang.srt.models.kimi_k3.is_in_breakable_cuda_graph", + return_value=breakable, + ), + ): + forward() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = forward() + for _ in range(3): + x.normal_() + hidden = x * 0.5 + expected = project(hidden, qkv_weight) * torch.sigmoid( + project(hidden, gate_weight) + ) + graph.replay() + torch.cuda.synchronize() + self.assertTrue(torch.equal(captured, expected)) + def test_block_fp8_weight_is_dequantized_for_tiny_gemm(self): module = SimpleNamespace( weight=torch.nn.Parameter(