kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-12 18:36:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support piecewise cuda graph for MLA (#11812)

Browse Source
This commit is contained in:
Ke Bao
2025-11-10 09:13:48 +08:00
committed by GitHub
parent 885cfca273
commit db24d34603
10 changed files with 196 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/compilation/compile.py
View File

@@ -112,6 +112,9 @@ def _infer_dynamic_arg_dims_from_annotations(forward_fn):
for a in getattr(ann, "__args__", [])
):
dyn[name] = 0
elif ann == "torch.Tensor" or ann == "Optional[torch.Tensor]":
# For future import annotations (e.g. from __future__ import annotations), the annotation is a string
dyn[name] = 0
if not dyn:
raise ValueError("No dynamic dims inferred; pass dynamic_arg_dims explicitly.")
return dyn

3
python/sglang/srt/compilation/piecewise_context_manager.py
View File

@@ -30,11 +30,10 @@ def get_forward_context() -> Optional[ForwardContext]:
@contextmanager
def set_forward_context(forward_batch: ForwardBatch, attention_layers: List[Any]):
global _forward_context
prev_forward_context = _forward_context
_forward_context = ForwardContext()
_forward_context.set_forward_batch(forward_batch)
_forward_context.set_attention_layers(attention_layers)
try:
yield
finally:
_forward_context = prev_forward_context
_forward_context = None

5
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
View File

@@ -22,6 +22,9 @@ from sglang.srt.layers.attention.flashinfer_backend import (
)
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
is_in_piecewise_cuda_graph,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import (
@@ -322,6 +325,8 @@ class FlashInferMLAAttnBackend(AttentionBackend):
use_ragged = (
not get_global_server_args().flashinfer_mla_disable_ragged
and extend_no_prefix
# Piecewise cuda graph should use paged prefill to be compatible with prefix cache
and not is_in_piecewise_cuda_graph()
)
self.indices_updater_prefill.update(

26
python/sglang/srt/layers/radix_attention.py
View File

@@ -110,9 +110,12 @@ class RadixAttention(nn.Module):
k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
output = torch.empty_like(q)
if self.qk_head_dim != self.v_head_dim:
output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
else:
output = torch.empty_like(q)
torch.ops.sglang.unified_attention_with_output(
q, k, v, output, save_kv_cache, self.layer_id
q, k, v, output, save_kv_cache, self.layer_id, **kwargs
)
return output
else:
@@ -134,13 +137,26 @@ def unified_attention_with_output(
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
*,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
) -> None:
context = get_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
kwargs = {}
if q_rope is not None:
kwargs["q_rope"] = q_rope
if k_rope is not None:
kwargs["k_rope"] = k_rope
if sinks is not None:
kwargs["sinks"] = sinks
ret = forward_batch.attn_backend.forward(
query, key, value, attention_layer, forward_batch, save_kv_cache
query, key, value, attention_layer, forward_batch, save_kv_cache, **kwargs
)
assert (
output.numel() == ret.numel()
@@ -157,6 +173,10 @@ def unified_attention_with_output_fake(
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
*,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
) -> None:
return

1
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -823,7 +823,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
query_pass = query[..., self.rotary_dim :]
key_pass = key[..., self.rotary_dim :]
self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device)
cos_sin = self.cos_sin_cache[
torch.add(positions, offsets) if offsets is not None else positions
]

23
python/sglang/srt/model_executor/model_runner.py
View File

@@ -337,8 +337,13 @@ class ModelRunner:
):
self.attention_layers = []
for layer in self.model.model.layers:
if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "attn"):
self.attention_layers.append(layer.self_attn.attn)
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
self.attention_layers.append(layer.self_attn.attn)
elif hasattr(layer.self_attn, "attn_mqa"):
# For DeepSeek model
self.attention_layers.append(layer.self_attn.attn_mqa)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(
@@ -2075,9 +2080,6 @@ class ModelRunner:
skip_attn_backend_init: bool = False,
pp_proxy_tensors=None,
) -> LogitsProcessorOutput:
if not skip_attn_backend_init:
self.attn_backend.init_forward_metadata(forward_batch)
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
@@ -2086,9 +2088,14 @@ class ModelRunner:
if not self.is_generation:
kwargs["get_embedding"] = True
if self.piecewise_cuda_graph_runner is not None:
if self.piecewise_cuda_graph_runner.can_run(forward_batch):
return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
if (
self.piecewise_cuda_graph_runner is not None
and self.piecewise_cuda_graph_runner.can_run(forward_batch)
):
return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
if not skip_attn_backend_init:
self.attn_backend.init_forward_metadata(forward_batch)
return self.model.forward(
forward_batch.input_ids,

141
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -55,22 +55,21 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
# Detect whether the current forward pass is in capture mode
is_capture_mode = False
_in_piecewise_cuda_graph = False
def get_is_capture_mode():
return is_capture_mode
def is_in_piecewise_cuda_graph():
return _in_piecewise_cuda_graph
@contextmanager
def model_capture_mode():
global is_capture_mode
is_capture_mode = True
def enable_piecewise_cuda_graph():
global _in_piecewise_cuda_graph
_in_piecewise_cuda_graph = True
yield
is_capture_mode = False
_in_piecewise_cuda_graph = False
@contextmanager
@@ -125,6 +124,15 @@ def set_global_graph_memory_pool(val):
global_graph_memory_pool = val
def set_torch_compile_config():
import torch._dynamo.config
# Resolve torch._dynamo.exc.FailOnRecompileLimitHit
torch._dynamo.config.accumulated_cache_size_limit = 1024
if hasattr(torch._dynamo.config, "cache_size_limit"):
torch._dynamo.config.cache_size_limit = 1024
class PiecewiseCudaGraphRunner:
"""A PiecewiseCudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
@@ -142,6 +150,8 @@ class PiecewiseCudaGraphRunner:
self.attn_tp_size = get_attention_tp_size()
self.attn_tp_rank = get_attention_tp_rank()
set_torch_compile_config()
assert (
self.model_runner.server_args.piecewise_cuda_graph_tokens is not None
), "piecewise_cuda_graph_tokens is not set"
@@ -166,7 +176,6 @@ class PiecewiseCudaGraphRunner:
if model_runner.server_args.enable_return_hidden_states:
self.capture_hidden_mode = CaptureHiddenMode.FULL
# Attention backend
self.max_num_tokens = max(self.capture_num_tokens)
# Graph inputs
@@ -185,29 +194,29 @@ class PiecewiseCudaGraphRunner:
# Set graph pool id globally to be able to use symmetric memory
set_graph_pool_id(get_global_graph_memory_pool())
with patch_model(
self.model_runner.model.model, self.compile_config.compiler
) as patched_model:
install_torch_compiled(
patched_model,
fullgraph=True,
dynamic_arg_dims=None,
compile_config=self.compile_config,
graph_pool=get_global_graph_memory_pool(),
)
with set_compiled(True):
self.warmup_and_capture()
# Capture
try:
with model_capture_mode():
self.capture()
except RuntimeError as e:
raise Exception(
f"Capture cuda graph failed: {e}\n{PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
with enable_piecewise_cuda_graph():
with patch_model(
self.model_runner.model.model, self.compile_config.compiler
) as patched_model:
install_torch_compiled(
patched_model,
fullgraph=True,
dynamic_arg_dims=None,
compile_config=self.compile_config,
graph_pool=get_global_graph_memory_pool(),
)
with set_compiled(True):
self.warmup_and_capture()
# Capture
try:
self.capture()
except RuntimeError as e:
raise Exception(
f"Capture cuda graph failed: {e}\n{PIECEWISE_CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
self.raw_num_tokens = 0
def warmup_and_capture(self):
@@ -225,7 +234,9 @@ class PiecewiseCudaGraphRunner:
req_to_token_pool=self.model_runner.req_to_token_pool,
token_to_kv_pool=self.model_runner.token_to_kv_pool,
attn_backend=self.model_runner.attn_backend,
out_cache_loc=torch.randint(0, 100, (num_tokens,), device=self.device),
out_cache_loc=torch.zeros(
(num_tokens,), device=self.device, dtype=self._cache_loc_dtype()
),
seq_lens_sum=num_tokens,
encoder_lens=None,
return_logprob=False,
@@ -378,6 +389,8 @@ class PiecewiseCudaGraphRunner:
if lora_ids is not None:
self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
# Run and capture
def run_once():
# Clean intermediate result cache for DP attention
@@ -401,7 +414,7 @@ class PiecewiseCudaGraphRunner:
)
return
for _ in range(2):
for _ in range(3):
self.device_module.synchronize()
self.model_runner.tp_group.barrier()
run_once()
@@ -483,36 +496,40 @@ class PiecewiseCudaGraphRunner:
forward_batch: ForwardBatch,
**kwargs,
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
if self.model_runner.tp_group.ca_comm is not None:
old_ca_disable = self.model_runner.tp_group.ca_comm.disabled
self.model_runner.tp_group.ca_comm.disabled = True
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
# Replay
with set_forward_context(static_forward_batch, self.attention_layers):
with set_compiled(True):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
if isinstance(output, LogitsProcessorOutput):
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_tokens],
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
)
else:
assert isinstance(output, PPProxyTensors)
# TODO(Yuwei): support PP Support
raise NotImplementedError(
"PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet."
)
if self.model_runner.tp_group.ca_comm is not None:
self.model_runner.tp_group.ca_comm.disabled = old_ca_disable
with enable_piecewise_cuda_graph():
if self.model_runner.tp_group.ca_comm is not None:
old_ca_disable = self.model_runner.tp_group.ca_comm.disabled
self.model_runner.tp_group.ca_comm.disabled = True
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
# Replay
with set_forward_context(static_forward_batch, self.attention_layers):
with set_compiled(True):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,
static_forward_batch.positions,
static_forward_batch,
**kwargs,
)
if isinstance(output, LogitsProcessorOutput):
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[
: self.raw_num_tokens
],
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
)
else:
assert isinstance(output, PPProxyTensors)
# TODO(Yuwei): support PP Support
raise NotImplementedError(
"PPProxyTensors is not supported in PiecewiseCudaGraphRunner yet."
)
if self.model_runner.tp_group.ca_comm is not None:
self.model_runner.tp_group.ca_comm.disabled = old_ca_disable
def get_spec_info(self, num_tokens: int):
spec_info = None

41
python/sglang/srt/models/deepseek_v2.py
View File

@@ -112,6 +112,9 @@ from sglang.srt.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
is_in_piecewise_cuda_graph,
)
from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.server_args import get_global_server_args
@@ -320,6 +323,9 @@ def _support_mha_one_shot(attn: DeepseekV2AttentionMLA, forward_batch, backend_n
def _handle_attention_backend(
attn: DeepseekV2AttentionMLA, forward_batch, backend_name
):
if is_in_piecewise_cuda_graph():
return AttnForwardMethod.MLA
sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
disable_ragged = (
backend_name in ["flashinfer", "flashmla"]
@@ -427,6 +433,9 @@ def handle_attention_nsa(attn, forward_batch):
def handle_attention_triton(attn, forward_batch):
if is_in_piecewise_cuda_graph():
return AttnForwardMethod.MLA
# when deterministic inference is enabled, use MLA
if get_global_server_args().enable_deterministic_inference:
return _dispatch_mla_subtype(attn, forward_batch)
@@ -1841,18 +1850,26 @@ class DeepseekV2AttentionMLA(nn.Module):
)
attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
else:
attn_bmm_output = torch.empty(
(attn_output.shape[0], self.num_local_heads * self.v_head_dim),
dtype=attn_output.dtype,
device=attn_output.device,
)
torch.bmm(
attn_output.transpose(0, 1),
self.w_vc,
out=attn_bmm_output.view(
-1, self.num_local_heads, self.v_head_dim
).transpose(0, 1),
)
if is_in_piecewise_cuda_graph():
# torch dynamo requires out= op was called where output tensor was non-contiguous
attn_bmm_output = (
torch.bmm(attn_output.transpose(0, 1), self.w_vc)
.transpose(0, 1)
.flatten(1, 2)
)
else:
attn_bmm_output = torch.empty(
(attn_output.shape[0], self.num_local_heads * self.v_head_dim),
dtype=attn_output.dtype,
device=attn_output.device,
)
torch.bmm(
attn_output.transpose(0, 1),
self.w_vc,
out=attn_bmm_output.view(
-1, self.num_local_heads, self.v_head_dim
).transpose(0, 1),
)
output, _ = self.o_proj(attn_bmm_output)
return output