chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support overlapping two batches (#4068)

Browse Source
This commit is contained in:
fzyzcjy
2025-05-25 08:39:07 +08:00
committed by GitHub
parent f456037396
commit 0d47788025
13 changed files with 1145 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -37,6 +37,7 @@ from sglang.srt.distributed import (
set_custom_all_reduce,
)
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
get_attention_tp_size,
@@ -198,6 +199,7 @@ class ModelRunner:
"disable_radix_cache": server_args.disable_radix_cache,
"enable_nan_detection": server_args.enable_nan_detection,
"enable_dp_attention": server_args.enable_dp_attention,
"enable_two_batch_overlap": server_args.enable_two_batch_overlap,
"enable_dp_lm_head": server_args.enable_dp_lm_head,
"enable_ep_moe": server_args.enable_ep_moe,
"enable_deepep_moe": server_args.enable_deepep_moe,
@@ -994,6 +996,13 @@ class ModelRunner:
def init_attention_backend(self):
"""Init attention kernel backend."""
if self.server_args.enable_two_batch_overlap:
self.attn_backend = TboAttnBackend.init_new(self._get_attention_backend)
else:
self.attn_backend = self._get_attention_backend()
# TODO unify with 6338
def _get_attention_backend(self):
if self.server_args.attention_backend == "flashinfer":
if not self.use_mla_backend:
from sglang.srt.layers.attention.flashinfer_backend import (
@@ -1003,17 +1012,17 @@ class ModelRunner:
# Init streams
if self.server_args.speculative_algorithm == "EAGLE":
self.plan_stream_for_flashinfer = torch.cuda.Stream()
self.attn_backend = FlashInferAttnBackend(self)
return FlashInferAttnBackend(self)
else:
from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAAttnBackend,
)
self.attn_backend = FlashInferMLAAttnBackend(self)
return FlashInferMLAAttnBackend(self)
elif self.server_args.attention_backend == "aiter":
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
self.attn_backend = AiterAttnBackend(self)
return AiterAttnBackend(self)
elif self.server_args.attention_backend == "triton":
assert self.sliding_window_size is None, (
"Window attention is not supported in the triton attention backend. "
@@ -1028,21 +1037,21 @@ class ModelRunner:
DoubleSparseAttnBackend,
)
self.attn_backend = DoubleSparseAttnBackend(self)
return DoubleSparseAttnBackend(self)
else:
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
self.attn_backend = TritonAttnBackend(self)
return TritonAttnBackend(self)
elif self.server_args.attention_backend == "torch_native":
from sglang.srt.layers.attention.torch_native_backend import (
TorchNativeAttnBackend,
)
self.attn_backend = TorchNativeAttnBackend(self)
return TorchNativeAttnBackend(self)
elif self.server_args.attention_backend == "flashmla":
from sglang.srt.layers.attention.flashmla_backend import FlashMLABackend
self.attn_backend = FlashMLABackend(self)
return FlashMLABackend(self)
elif self.server_args.attention_backend == "fa3":
assert (
torch.cuda.get_device_capability()[0] == 8 and not self.use_mla_backend
@@ -1054,13 +1063,13 @@ class ModelRunner:
FlashAttentionBackend,
)
self.attn_backend = FlashAttentionBackend(self)
return FlashAttentionBackend(self)
elif self.server_args.attention_backend == "cutlass_mla":
from sglang.srt.layers.attention.cutlass_mla_backend import (
CutlassMLABackend,
)
self.attn_backend = CutlassMLABackend(self)
return CutlassMLABackend(self)
else:
raise ValueError(
f"Invalid attention backend: {self.server_args.attention_backend}"