minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[Kimi K3] Fix CUDA graph stream explosion (#40640)

Browse Source
This commit is contained in:
Khoa Pham
2026-09-21 17:49:50 -07:00
committed by GitHub
parent 31b577bb08
commit c4d3770a68
2 changed files with 98 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/kimi_k3.py
+25 -33
View File
@@ -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
)
test/registered/unit/models/test_kimi_k3_bfa_overlap.py
+73 -28
View File
@@ -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(