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

[PP][DeepSeek V4] Overlap communication and optimize SM120 prefill (#38792)

Browse Source 
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>
This commit is contained in:
jmydurant
2026-09-18 21:18:04 -07:00
committed by GitHub
co-authored by Yangmin Li YAMY
parent 929230a6f0
commit 5e9342d16f
17 changed files with 418 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/model_hook.py
+5 -3
View File
@@ -431,9 +431,11 @@ def handle_model_specific_adjustments(server_args: Any):
validate_deepseek_v4_mega_moe_token_budget(server_args)
if get_platform().is_sm120:
# SM120 lacks tcgen05/TMEM: disable features that depend on
# DeepGEMM or require >99KB SMEM (topk_v2).
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
# FP8 wo_a stays opt-in on SM120: only recent DeepGEMM builds ship
# the SM120 kernels, and deep_gemm_wrapper.configurer validates them.
if not envs.SGLANG_OPT_FP8_WO_A_GEMM.is_set():
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
# The default top-k v2 path still requires unsupported resources.
envs.SGLANG_OPT_USE_TOPK_V2.set(False)
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.is_set():
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
python/sglang/srt/arg_groups/serving_hook.py
+5 -5
View File
@@ -425,10 +425,10 @@ def handle_environment_variables(server_args: Any):
"--enable-deepseek-v4-fp4-indexer requires SM100, SM120, or gfx95 GPUs "
"with FP4 indexer support."
)
# FP8 W_o GEMM needs DeepGEMM JIT. Enable exactly where the runtime can run
# it, mirroring the forward scale split: the ue8m0 path
# (DEEPGEMM_SCALE_UE8M0, true sm100, default on) or an sm90 opt-in
# fp32-scale path (use FP4 expert ckpt). Disable in every other case.
# FP8 W_o GEMM needs DeepGEMM JIT. Enable exactly where the runtime can
# run it, mirroring the forward scale split: the default sm100 UE8M0
# path, or explicit opt-in on sm90 (FP32 scales) and sm120 (UE8M0).
# SM120 API compatibility is centralized in deep_gemm_wrapper.configurer.
if get_platform().is_cuda and envs.SGLANG_OPT_FP8_WO_A_GEMM.get():
from sglang.srt.layers import deep_gemm_wrapper
@@ -442,7 +442,7 @@ def handle_environment_variables(server_args: Any):
if not supported and explicit:
logger.warning(
"Disabling SGLANG_OPT_FP8_WO_A_GEMM: requires DeepGEMM JIT "
"and sm100+ (Blackwell), or explicit opt-in on sm90; "
"and a compatible sm100/sm120 build, or explicit opt-in on sm90; "
"detected sm%d.",
sm,
)
python/sglang/srt/environ.py
+6
View File
@@ -654,6 +654,9 @@ class Envs:
# PP: skip output send/recv when the entire batch consists of non-final chunked prefill requests,
# since process_batch_result_prefill discards next_token_ids for those anyway.
SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM = EnvBool(False)
# Run PP tensor communication on a dedicated stream so asynchronous sends
# do not fence the next forward through the scheduler stream.
SGLANG_PP_COMM_OVERLAP = EnvBool(False)
SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH = EnvBool(False)
# ===================================================================
@@ -1075,6 +1078,9 @@ class Envs:
SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS = EnvInt(1)
# SM120 FlashMLA decode backend: "flashinfer" (default), "triton", or "torch".
SGLANG_SM120_FLASHMLA_BACKEND = EnvStr("flashinfer")
# Store DeepSeek-V4 SWA KV directly in FlashInfer's 64-token SM120 page
# layout. The scheduler continues to allocate 256-token logical pages.
SGLANG_OPT_SM120_DIRECT_SWA_KV = EnvBool(False)
SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096)
SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048)
SGLANG_FLASHINFER_AUTOTUNE_CACHE = EnvBool(True)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+4 -4
View File
@@ -2518,7 +2518,7 @@ class DeepseekV4AttnBackend(
req_to_token=self.req_to_token,
full_to_swa=self.token_to_kv_pool.full_to_swa_index_mapping,
swa_window_size=SWA_WINDOW,
swa_page_size=self.token_to_kv_pool.swa_page_size,
swa_page_size=self.token_to_kv_pool.swa_kv_pool.page_size,
num_qo_tokens=num_qo_tokens,
max_seq_len=max(seq_lens_cpu_list),
total_swa=total_swa,
@@ -3749,13 +3749,13 @@ class DeepseekV4AttnBackend(
compress_ratio
)
swa_page_size = token_to_kv_pool.swa_page_size
swa_kv_page_size = token_to_kv_pool.swa_kv_pool.page_size
assert swa_k_cache.ndim == 2
# The kernel detects each cache's format from the last dim of this
# view: 584 (V4), 528 (V4.1 fp8) or 288 (V4.1 fp4, extra cache only).
k_cache_total_dim = token_to_kv_pool.get_swa_key_bytes_per_token()
swa_k_cache = swa_k_cache[:, : swa_page_size * k_cache_total_dim].view(
swa_k_cache.shape[0], swa_page_size, 1, k_cache_total_dim
swa_k_cache = swa_k_cache[:, : swa_kv_page_size * k_cache_total_dim].view(
swa_k_cache.shape[0], swa_kv_page_size, 1, k_cache_total_dim
)
if extra_k_cache is not None:
python/sglang/srt/layers/deep_gemm_wrapper/configurer.py
+26 -10
View File
@@ -14,26 +14,42 @@ _is_cuda = is_cuda()
_is_musa = is_musa()
def _sm120_deep_gemm_apis_available() -> bool:
try:
import deep_gemm
except (ImportError, OSError, RuntimeError):
return False
return all(
callable(getattr(deep_gemm, name, None))
for name in (
"fp8_einsum",
"m_grouped_fp8_fp4_gemm_nt_contiguous",
"transform_sf_into_required_layout",
)
)
def _compute_enable_deep_gemm():
if not (_is_cuda or _is_musa):
return False
if not envs.SGLANG_ENABLE_JIT_DEEPGEMM.get():
return False
sm_version = get_device_sm()
if (_is_cuda and sm_version < 90) or (_is_musa and sm_version < 31):
return False
# SM120/SM121 support (including GB10) landed in DeepGEMM#324;
# probe the entry point since installed builds may predate it.
if sm_version in (120, 121):
try:
from deep_gemm import m_grouped_fp8_fp4_gemm_nt_contiguous # noqa: F401
except (ImportError, AttributeError):
return False
if not (_is_cuda or _is_musa):
# SM120/SM121 support (mma.sync block-scale, no TMEM) landed in DeepGEMM#324;
# probe every API used by the SM120 DSV4 paths since installed builds may
# expose fp8_einsum but still predate the SM120 kernels.
if sm_version in (120, 121) and not _sm120_deep_gemm_apis_available():
return False
try:
import deep_gemm # noqa: F401
except ImportError:
except (ImportError, OSError, RuntimeError):
return False
return envs.SGLANG_ENABLE_JIT_DEEPGEMM.get()
return True
ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm()
python/sglang/srt/managers/scheduler.py
+5 -7
View File
@@ -279,6 +279,7 @@ from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin
from sglang.srt.managers.utils import (
EmbeddingBatchResult,
GenerationBatchResult,
allocate_distinct_stream,
is_health_check_generate_req,
validate_input_length,
)
@@ -1899,13 +1900,10 @@ class Scheduler(
# stream aliases forward_stream, which would eliminate scheduler
# overlap. Only CUDA/HIP streams expose a ``cuda_stream`` handle;
# other accelerators (e.g. NPU/XPU) skip the alias check.
_redraws = 0
while (
self.schedule_stream.cuda_stream == self.forward_stream.cuda_stream
and _redraws < 64
):
self.schedule_stream = self.device_module.Stream(priority=0)
_redraws += 1
if self.schedule_stream.cuda_stream == self.forward_stream.cuda_stream:
self.schedule_stream = allocate_distinct_stream(
self.device_module, (self.forward_stream,)
)
# The global WAR barrier fences the scheduler's next shared-buffer write
# on the previous forward's read of the unified memory pool.
self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get()
python/sglang/srt/managers/scheduler_pp_mixin.py
+116 -59
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import logging
from collections import defaultdict, deque
from contextlib import nullcontext
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
@@ -18,6 +19,7 @@ from sglang.srt.managers.overlap_utils import RelayPayload
from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ScheduleBatch
from sglang.srt.managers.utils import (
GenerationBatchResult,
allocate_distinct_stream,
get_logprob_dict_from_result,
get_logprob_from_pp_outputs,
)
@@ -144,7 +146,7 @@ class SchedulerPPMixin:
next_mb_id,
)
)
self._pp_commit_comm_work(self.send_proxy_work)
self._pp_commit_proxy_send_work()
if cur_batch:
result, self.launch_event = self._pp_launch_batch(
mb_id,
@@ -176,17 +178,7 @@ class SchedulerPPMixin:
self.last_mbs[next_mb_id] = self.mbs[next_mb_id]
if not self.pp_group.is_last_rank:
if cur_batch:
self.device_module.current_stream().wait_event(
self.launch_event
)
with torch.profiler.record_function(
"send_proxy_dict_to_next_stage"
):
self.send_proxy_work = self._pp_send_dict_to_next_stage(
result.pp_hidden_states_proxy_tensors.tensors,
async_send=True,
msg_type="proxy",
)
self._pp_send_proxy_to_next_stage(result)
self.pp_outputs = next_pp_outputs
@@ -296,7 +288,7 @@ class SchedulerPPMixin:
next_mb_id,
)
)
self._pp_commit_comm_work(self.send_proxy_work)
self._pp_commit_proxy_send_work()
if cur_batch:
if self.enable_staging:
self.maybe_prefetch_staging_for_batch(cur_batch)
@@ -361,14 +353,7 @@ class SchedulerPPMixin:
transferred_rids, async_send=True
)
if cur_batch:
self.device_module.current_stream().wait_event(
self.launch_event
)
self.send_proxy_work = self._pp_send_dict_to_next_stage(
result.pp_hidden_states_proxy_tensors.tensors,
async_send=True,
msg_type="proxy",
)
self._pp_send_proxy_to_next_stage(result)
self.pp_outputs = next_pp_outputs
release_rids = next_release_rids
@@ -456,7 +441,7 @@ class SchedulerPPMixin:
next_mb_id,
)
)
self._pp_commit_comm_work(self.send_proxy_work)
self._pp_commit_proxy_send_work()
if cur_batch:
result, self.launch_event = self._pp_launch_batch(
@@ -549,14 +534,7 @@ class SchedulerPPMixin:
transferred_rids, async_send=True
)
if cur_batch and not cur_batch.forward_mode.is_prebuilt():
self.device_module.current_stream().wait_event(
self.launch_event
)
self.send_proxy_work = self._pp_send_dict_to_next_stage(
result.pp_hidden_states_proxy_tensors.tensors,
async_send=True,
msg_type="proxy",
)
self._pp_send_proxy_to_next_stage(result)
self.pp_outputs = next_pp_outputs
release_rids = next_release_rids
@@ -597,10 +575,32 @@ class SchedulerPPMixin:
self.send_req_work = []
self.send_proxy_work = []
self.send_output_work = []
self.send_proxy_requires_forward_fence = False
self.launch_event = None
self._pp_tensor_dict_inbox: Dict[str, deque[Dict[str, torch.Tensor]]] = (
defaultdict(deque)
)
self.pp_proxy_recv_event = None
self.pp_send_done_event = None
# With SGLANG_PP_COMM_OVERLAP, PP tensor sends/recvs run on a dedicated
# stream so async sends do not fence the next forward through the
# scheduler stream. Otherwise they stay on the current stream.
self.pp_comm_stream = None
self.pp_comm_stream_ctx = nullcontext()
if (
envs.SGLANG_PP_COMM_OVERLAP.get()
and torch.cuda.is_available()
and str(self.device).startswith("cuda")
):
self.pp_comm_stream = allocate_distinct_stream(
self.device_module,
(self.schedule_stream, self.forward_stream, self.copy_stream),
)
self.pp_comm_stream_ctx = self.device_module.stream(self.pp_comm_stream)
logger.info(
"PP tensor communication overlap enabled on a dedicated CUDA stream"
)
self._pp_tensor_dict_inbox: Dict[
str, deque[Tuple[Dict[str, torch.Tensor], Optional[torch.Event]]]
] = defaultdict(deque)
def process_bootstrapped_queue(
self: Scheduler, bootstrapped_rids: Optional[List[str]]
@@ -734,11 +734,45 @@ class SchedulerPPMixin:
)
return send_release_work, release_rids
def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None:
for p2p_work in work:
p2p_work.work.wait()
def _pp_record_comm_event(self: Scheduler) -> Optional[torch.Event]:
"""Mark the tail of the PP comm stream; None when not overlapping."""
if self.pp_comm_stream is None:
return None
event = self.device_module.Event()
event.record(self.pp_comm_stream)
return event
def _pp_commit_comm_work(
self: Scheduler,
work: List[P2PWork],
fence_next_forward: bool = False,
) -> None:
with self.pp_comm_stream_ctx:
for p2p_work in work:
p2p_work.work.wait()
if fence_next_forward and work:
# CUDA Graph outputs are views of replay-owned static buffers. Do not
# replay the next forward until NCCL has stopped reading the previous
# proxy tensors. Eager prefill remains overlapped.
self.pp_send_done_event = self._pp_record_comm_event()
work.clear()
def _pp_commit_proxy_send_work(self: Scheduler) -> None:
self._pp_commit_comm_work(
self.send_proxy_work,
fence_next_forward=self.send_proxy_requires_forward_fence,
)
self.send_proxy_requires_forward_fence = False
def _pp_wait_forward_dependencies(self: Scheduler) -> None:
self.forward_stream.wait_stream(self.schedule_stream)
if self.pp_send_done_event is not None:
self.forward_stream.wait_event(self.pp_send_done_event)
self.pp_send_done_event = None
if self.pp_proxy_recv_event is not None:
self.forward_stream.wait_event(self.pp_proxy_recv_event)
self.pp_proxy_recv_event = None
def _pp_commit_send_output_work_and_preprocess_output_tensors(
self: Scheduler,
next_first_rank_mb_id: int,
@@ -859,6 +893,7 @@ class SchedulerPPMixin:
tensor_dict: Dict[str, torch.Tensor],
async_send: bool = True,
msg_type: str = "default",
ready_event: Optional[torch.Event] = None,
):
# Warn once if using default untyped messages
if msg_type == "default":
@@ -868,24 +903,40 @@ class SchedulerPPMixin:
)
tensor_dict["__msg_type__"] = msg_type
p2p_work = []
p2p_work.extend(
self.pp_group.send_tensor_dict(
tensor_dict=tensor_dict,
all_gather_group=(self.attn_tp_group),
async_send=async_send,
with self.pp_comm_stream_ctx:
if ready_event is not None:
self.device_module.current_stream().wait_event(ready_event)
p2p_work.extend(
self.pp_group.send_tensor_dict(
tensor_dict=tensor_dict,
all_gather_group=(self.attn_tp_group),
async_send=async_send,
)
)
)
return p2p_work
def _pp_send_proxy_to_next_stage(
self: Scheduler, result: GenerationBatchResult
) -> None:
with torch.profiler.record_function("send_proxy_dict_to_next_stage"):
self.send_proxy_work = self._pp_send_dict_to_next_stage(
result.pp_hidden_states_proxy_tensors.tensors,
async_send=True,
msg_type="proxy",
ready_event=self.launch_event,
)
self.send_proxy_requires_forward_fence = result.can_run_cuda_graph
def _pp_recv_typed_dict(
self: Scheduler,
expected_kind: str = "default",
all_gather_group: Optional = None,
) -> Dict[str, torch.Tensor]:
) -> Tuple[Dict[str, torch.Tensor], Optional[torch.Event]]:
"""Receive a typed tensor dict, demultiplexing by msg_type.
If a message of the wrong kind is received, it's stashed in the queue
and we continue receiving until we get the expected kind.
and we continue receiving until we get the expected kind. The returned
event marks receive completion on the PP comm stream (None otherwise).
"""
if expected_kind in self._pp_tensor_dict_inbox:
inbox_queue = self._pp_tensor_dict_inbox[expected_kind]
@@ -893,9 +944,11 @@ class SchedulerPPMixin:
return inbox_queue.popleft()
while True:
tensor_dict = self.pp_group.recv_tensor_dict(
all_gather_group=all_gather_group
)
with self.pp_comm_stream_ctx:
tensor_dict = self.pp_group.recv_tensor_dict(
all_gather_group=all_gather_group
)
recv_event = self._pp_record_comm_event()
received_kind = tensor_dict.get("__msg_type__", "default")
if received_kind == expected_kind:
if received_kind == "default":
@@ -903,27 +956,28 @@ class SchedulerPPMixin:
f"PP recv: got default untyped message. Content keys: {tensor_dict.keys()}"
"Consider adding msg_type='proxy' or 'output' to avoid recv conflicts."
)
return tensor_dict
return tensor_dict, recv_event
else:
logger.debug(
f"PP recv: expected {expected_kind}, got {received_kind}, stashing"
)
self._pp_tensor_dict_inbox[received_kind].append(tensor_dict)
self._pp_tensor_dict_inbox[received_kind].append(
(tensor_dict, recv_event)
)
def _pp_recv_proxy_tensors(self: Scheduler) -> Optional[PPProxyTensors]:
pp_proxy_tensors = None
if not self.pp_group.is_first_rank:
pp_proxy_tensors = PPProxyTensors(
self._pp_recv_typed_dict(
expected_kind="proxy",
all_gather_group=(self.attn_tp_group),
)
tensor_dict, self.pp_proxy_recv_event = self._pp_recv_typed_dict(
expected_kind="proxy",
all_gather_group=(self.attn_tp_group),
)
pp_proxy_tensors = PPProxyTensors(tensor_dict)
return pp_proxy_tensors
def _pp_recv_dict_from_prev_stage(
self: Scheduler,
) -> Dict[str, torch.Tensor]:
) -> Tuple[Dict[str, torch.Tensor], Optional[torch.Event]]:
return self._pp_recv_typed_dict(
expected_kind="output",
all_gather_group=(self.attn_tp_group),
@@ -1370,12 +1424,12 @@ class SchedulerPPMixin:
not target.forward_mode.is_prebuilt()
and not _pp_can_skip_output_comm(target)
):
self.device_module.current_stream().wait_event(q_event)
with torch.profiler.record_function("send_res_dict_to_next_stage"):
send_output_work = self._pp_send_dict_to_next_stage(
pp_outputs_to_send.tensors,
async_send=True,
msg_type="output",
ready_event=q_event,
)
# send the outputs from the last round to let the next stage worker run post processing
if not self.pp_group.is_last_rank:
@@ -1460,9 +1514,12 @@ class SchedulerPPMixin:
)
return
with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage())
tensor_dict, output_recv_event = self._pp_recv_dict_from_prev_stage()
next_pp_outputs = PPProxyTensors(tensor_dict)
with self.copy_stream_ctx:
self.copy_stream.wait_stream(self.schedule_stream)
if output_recv_event is not None:
self.copy_stream.wait_event(output_recv_event)
batch_result = self._pp_prep_batch_result(
target, mb_metadata[next_mb_id], next_pp_outputs
)
@@ -1593,7 +1650,7 @@ class SchedulerPPMixin:
elif should_recv:
# Recv only (no send needed)
with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
recv_dict = self._pp_recv_dict_from_prev_stage()
recv_dict, _ = self._pp_recv_dict_from_prev_stage()
_handle_recv_dict(recv_dict)
return next_pp_outputs, batch_result, d2h_event, send_output_work
@@ -1608,7 +1665,7 @@ class SchedulerPPMixin:
):
with torch.profiler.record_function("run_batch"):
with self.forward_stream_ctx:
self.forward_stream.wait_stream(self.schedule_stream)
self._pp_wait_forward_dependencies()
set_time_batch(
cur_batch.reqs,
"set_run_batch_cpu_start_time",
python/sglang/srt/managers/utils.py
+14
View File
@@ -31,6 +31,20 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def allocate_distinct_stream(device_module, avoid_streams):
"""Draw a stream that aliases none of ``avoid_streams``.
CUDA/HIP streams come from a fixed round-robin pool, so a fresh ``Stream()``
may hand back one that is already in use.
"""
avoid = {stream.cuda_stream for stream in avoid_streams}
for _ in range(65):
stream = device_module.Stream(priority=0)
if stream.cuda_stream not in avoid:
return stream
raise RuntimeError("Unable to allocate a distinct stream")
def _async_d2h(t: torch.Tensor) -> torch.Tensor:
"""Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H
to pageable host memory blocks the caller until done) and record_stream keeps
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+41 -7
View File
@@ -26,7 +26,7 @@ from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.runtime_context import get_exec, get_spec
from sglang.srt.runtime_context import get_exec, get_platform, get_spec
from sglang.srt.utils import ceil_div, is_hip
logger = logging.getLogger(__name__)
@@ -78,6 +78,14 @@ def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int:
return sliding_window + spec_extra
def _num_dsv4_physical_kv_pages(
size: int, physical_page_size: int, logical_page_size: int
) -> int:
# The paged allocator reserves one page at the logical page size, so the
# highest token index is size + logical_page_size - 1.
return ceil_div(size + logical_page_size, physical_page_size)
def resolve_compressed_kv_layout(
kv_layout: KVLayout, compress_ratio: int, option: Optional[str] = None
) -> KVLayout:
@@ -158,6 +166,7 @@ class DeepSeekV4SingleKVPool(KVCache):
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
kv_layout: Union[str, KVLayout] = KVLayout.V4,
global_page_size: Optional[int] = None,
):
super().__init__(
size,
@@ -171,6 +180,7 @@ class DeepSeekV4SingleKVPool(KVCache):
)
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.global_page_size = global_page_size or page_size
# Paged FlashMLA layout of this pool's pages; see KVLayout.
self.kv_layout = KVLayout.parse(kv_layout)
@@ -190,7 +200,9 @@ class DeepSeekV4SingleKVPool(KVCache):
):
self.kv_buffer = [
self.create_buffer(
num_pages=(self.size + self.page_size + 1) // self.page_size,
num_pages=_num_dsv4_physical_kv_pages(
self.size, self.page_size, self.global_page_size
),
)
for _ in range(self.layer_num)
]
@@ -356,6 +368,7 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool):
start_layer: int | None = None,
end_layer: int | None = None,
kv_layout: Union[str, KVLayout] = KVLayout.V4,
global_page_size: int | None = None,
):
super().__init__(
size,
@@ -368,6 +381,7 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool):
enable_memory_saver,
start_layer,
end_layer,
global_page_size=global_page_size,
kv_layout=kv_layout,
)
# The HiSparse transfer kernels hardcode the V4 token layout.
@@ -974,6 +988,23 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.swa_size = swa_size
self.swa_page_size = swa_page_size
# The allocator and compress state keep 256-token logical pages, while
# FlashInfer's SM120 DSV4 kernel consumes 64-token physical pages. Storing
# SWA KV in that layout removes the per-layer 256 -> 64 page split; the
# allocator's flat token indices stay valid.
swa_kv_page_size = (
64
if get_platform().is_sm120 and envs.SGLANG_OPT_SM120_DIRECT_SWA_KV.get()
else swa_page_size
)
assert swa_page_size % swa_kv_page_size == 0
if swa_kv_page_size != swa_page_size:
logger.info(
"DeepSeek-V4 SM120 direct SWA KV layout enabled: "
"logical_page_size=%d physical_page_size=%d",
swa_page_size,
swa_kv_page_size,
)
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
@@ -1059,7 +1090,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
kv_pool_cls = DeepSeekV4UniformFP8KVPool
self.swa_kv_pool = self._make_kv_pool(
size=swa_size,
page_size=swa_page_size,
page_size=swa_kv_page_size,
dtype=dtype,
layer_num=stage_layer_num,
device=device,
@@ -1283,11 +1314,14 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
item_lens: List[int] = []
if self.swa_kv_pool is not None:
physical_pages_per_logical_page = (
self.swa_page_size // self.swa_kv_pool.page_size
)
for buf in self.swa_kv_pool.kv_buffer:
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
data_ptrs.append(buf.data_ptr())
data_lens.append(buf.nbytes)
item_lens.append(buf[0].nbytes)
item_lens.append(buf[0].nbytes * physical_pages_per_logical_page)
for pools in [
self.compress_state_pools,
@@ -1435,11 +1469,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
kv_layout: KVLayout = KVLayout.V4,
) -> DeepSeekV4SingleKVPool:
"""Build a full / SWA / c4 / c128 single-KV pool. ``global_page_size``
is the model-wide page_size (== ``page_size`` for the SWA pool, larger
for the per-ratio c4/c128 pools); the default CUDA pool ignores it.
is the model-wide logical page size. CUDA pools use it to reserve enough
physical rows for the allocator's dummy logical page.
Overridden by :class:`DSV4NPUTokenToKVPool` to swap in the NPU bf16
PA_ND variant, which needs ``global_page_size`` for its kernel view."""
del global_page_size # CUDA pools key only off their own page_size
return cls(
size,
page_size,
@@ -1449,6 +1482,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
layer_num,
device,
enable_memory_saver,
global_page_size=global_page_size,
kv_layout=kv_layout,
)
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+23
View File
@@ -483,6 +483,19 @@ def _dsv4_compressed_region_buffers(kvcache: Any, ratio: int) -> tuple[list, int
return pool.kv_buffer, pool.bytes_per_page_padded
def _require_single_row_dsv4_swa_pages(
*, logical_page_size: int, physical_page_size: int, consumer: str
) -> None:
"""Reject consumers that cannot map one logical SWA page to many rows."""
if logical_page_size != physical_page_size:
raise ValueError(
f"{consumer} does not support the DeepSeek-V4 direct SWA KV layout "
f"({logical_page_size}-token logical pages stored as "
f"{physical_page_size}-token physical rows). Disable "
"SGLANG_OPT_SM120_DIRECT_SWA_KV for this configuration."
)
def _dsv4_page_aligned_only(pool: Any) -> bool:
"""Whether a pool may only move whole pages: the token-granular copy
(``transfer_cache_dsv4_mla``) hardcodes the V4 data/scale row split."""
@@ -743,6 +756,11 @@ def build_deepseek_v4_hicache_stack(
# Unified KV and encoder replay rebuild SWA state; keep it out of host cache.
swa_layer_mapping = {}
else:
_require_single_row_dsv4_swa_pages(
logical_page_size=kvcache.swa_page_size,
physical_page_size=kvcache.swa_kv_pool.page_size,
consumer="DeepSeek-V4 HiCache",
)
if len(kvcache.swa_kv_pool.kv_buffer) != transfer_layer_num:
raise ValueError(
"DeepSeek V4 SWA KV pool must be PP-stage-local: "
@@ -1431,6 +1449,11 @@ def build_swa_draft_pools(
target_swa_host_pool = host_pool_group.entry_map[PoolName.SWA].host_pool
if isinstance(target_swa_host_pool, DeepSeekV4PagedHostPool):
_require_single_row_dsv4_swa_pages(
logical_page_size=target_swa_host_pool.slot_page_size,
physical_page_size=draft_swa_pool.page_size,
consumer="DeepSeek-V4 MTP SWA HiCache sidecar",
)
host_pool = DeepSeekV4PagedHostPool(
pool_name=str(PoolName.DRAFT_SWA),
device_buffers=draft_swa_pool.kv_buffer,
python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py
+6
View File
@@ -263,6 +263,7 @@ def _build_deepseek_v4_device_pool_group(
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_dsv4_compressed_region_buffers,
_dsv4_indexer_regions,
_require_single_row_dsv4_swa_pages,
_resolve_deepseek_v4_layer_mappings,
)
@@ -273,6 +274,11 @@ def _build_deepseek_v4_device_pool_group(
is_unified_kv = getattr(kvcache, "_unified_kv", False)
entries = []
if not is_unified_kv:
_require_single_row_dsv4_swa_pages(
logical_page_size=kvcache.swa_page_size,
physical_page_size=kvcache.swa_kv_pool.page_size,
consumer="DeepSeek-V4 direct external linker",
)
if kvcache.swa_page_size != page_size:
raise ValueError(
"DeepSeek V4 SWA page size must match the tree page size: "
python/sglang/srt/models/deepseek_v4.py
+8 -13
View File
@@ -78,6 +78,7 @@ from sglang.srt.layers.cp.utils import (
cp_materialize_global_token_order,
is_cp_active,
)
from sglang.srt.layers.deep_gemm_wrapper.configurer import DEEPGEMM_SCALE_UE8M0
from sglang.srt.layers.dp_attention import (
_tbo_event,
attn_tp_all_gather,
@@ -259,6 +260,7 @@ def _get_mhc_ops() -> MhcOps:
logger = logging.getLogger(__name__)
_FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get()
_FP8_WO_A_UE8M0 = _FP8_WO_A_GEMM and DEEPGEMM_SCALE_UE8M0
def wo_a_fp8_gemm_enabled(quant_config: Optional[QuantizationConfig]) -> bool:
@@ -881,11 +883,7 @@ class MqaAttentionBase(nn.Module):
self.wo_a._dsv4_num_groups = self.n_local_groups
self.wo_a._dsv4_o_lora_rank = self.o_lora_rank
elif fp8:
from sglang.srt.layers import deep_gemm_wrapper
self.wo_a.weight_scale_inv.format_ue8m0 = (
deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
)
self.wo_a.weight_scale_inv.format_ue8m0 = _FP8_WO_A_UE8M0
# wo_a is quantized but never *applied* through its quant method:
# the absorb GEMM in forward() reads .weight / .weight_scale_inv and
# runs its own batched kernel (DeepGEMM fp8_einsum on CUDA, aiter
@@ -2408,12 +2406,11 @@ class MQALayer(MqaAttentionBase):
elif self.wo_a_fp8:
import deep_gemm
from sglang.srt.layers import deep_gemm_wrapper
T, G, D = o.shape
R = self.o_lora_rank
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
# sm100 (Blackwell): ue8m0 scales via the dedicated JIT kernel.
if _FP8_WO_A_UE8M0:
# Blackwell (including SM120): UE8M0 scales via the dedicated
# JIT kernel.
o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o)
recipe = (1, 1, 128)
else:
@@ -4896,9 +4893,7 @@ class DeepseekV4ForCausalLM(nn.Module):
return output
def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None:
from sglang.srt.layers import deep_gemm_wrapper
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
if _FP8_WO_A_UE8M0:
from deep_gemm import transform_sf_into_required_layout
if is_nextn:
@@ -4937,7 +4932,7 @@ class DeepseekV4ForCausalLM(nn.Module):
continue
raw_scale = attn.wo_a.weight_scale_inv.data.view(G, R // 128, D // 128)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
if _FP8_WO_A_UE8M0:
attn.wo_a.weight_scale_inv.data = transform_sf_into_required_layout(
raw_scale,
mn=R,
test/registered/unit/disaggregation/test_disaggregation_wire.py
+2 -1
View File
@@ -927,6 +927,7 @@ def _make_dsv4_draft(*, unified, mapping=None):
pool._unified_kv = unified
pool.compression_ratios = [0]
pool.page_size = 256
pool.swa_page_size = 256
pool.sliding_window = 128
pool.full_to_swa_index_mapping = mapping
pool.unified_swa_window = 128
@@ -941,7 +942,7 @@ def _make_dsv4_draft(*, unified, mapping=None):
)
else:
pool.swa_kv_pool = SimpleNamespace(
kv_buffer=[torch.empty((2, 16), dtype=torch.uint8)]
page_size=256, kv_buffer=[torch.empty((2, 16), dtype=torch.uint8)]
)
return pool
test/registered/unit/distributed/test_pp_comm_overlap.py
+102
View File
@@ -0,0 +1,102 @@
import unittest
from collections import defaultdict, deque
from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import Mock, call
import torch
from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class FakeStream:
def __init__(self, stream_id):
self.cuda_stream = stream_id
class FakeEvent:
def __init__(self):
self.recorded_stream = None
def record(self, stream):
self.recorded_stream = stream
def _make_scheduler(**attrs):
scheduler = object.__new__(SchedulerPPMixin)
scheduler.__dict__.update(attrs)
return scheduler
class TestPPCommOverlap(CustomTestCase):
def test_graph_proxy_send_records_forward_reuse_fence(self):
comm_stream = FakeStream(4)
work = Mock()
works = [SimpleNamespace(work=work)]
scheduler = _make_scheduler(
pp_comm_stream=comm_stream,
pp_comm_stream_ctx=nullcontext(),
pp_send_done_event=None,
device_module=SimpleNamespace(Event=FakeEvent),
)
scheduler._pp_commit_comm_work(works, fence_next_forward=True)
work.wait.assert_called_once_with()
self.assertEqual(works, [])
self.assertIs(scheduler.pp_send_done_event.recorded_stream, comm_stream)
def test_no_fence_event_without_comm_stream(self):
scheduler = _make_scheduler(
pp_comm_stream=None,
pp_comm_stream_ctx=nullcontext(),
pp_send_done_event=None,
)
scheduler._pp_commit_comm_work([SimpleNamespace(work=Mock())], True)
self.assertIsNone(scheduler.pp_send_done_event)
def test_forward_waits_for_graph_send_and_proxy_receive(self):
schedule_stream = FakeStream(1)
send_done_event = object()
recv_event = object()
forward_stream = Mock()
scheduler = _make_scheduler(
schedule_stream=schedule_stream,
forward_stream=forward_stream,
pp_send_done_event=send_done_event,
pp_proxy_recv_event=recv_event,
)
scheduler._pp_wait_forward_dependencies()
forward_stream.wait_stream.assert_called_once_with(schedule_stream)
self.assertEqual(
forward_stream.wait_event.call_args_list,
[call(send_done_event), call(recv_event)],
)
self.assertIsNone(scheduler.pp_send_done_event)
self.assertIsNone(scheduler.pp_proxy_recv_event)
def test_inbox_returns_original_receive_event(self):
recv_event = object()
tensor_dict = {"__msg_type__": "output", "value": torch.arange(2)}
scheduler = _make_scheduler(
_pp_tensor_dict_inbox=defaultdict(
deque, {"output": deque([(tensor_dict, recv_event)])}
),
)
received, event = scheduler._pp_recv_typed_dict("output")
self.assertIs(received, tensor_dict)
self.assertIs(event, recv_event)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_dsv4_compressed_pools.py
+35
View File
@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DeepSeekV4SingleKVPool,
DeepSeekV4TokenToKVPool,
_CompressedPoolConfig,
_num_dsv4_physical_kv_pages,
)
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
@@ -23,6 +24,40 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestDSV4CompressedPools(CustomTestCase):
def test_physical_kv_pages_cover_reserved_logical_page(self):
size = 8192
self.assertEqual(_num_dsv4_physical_kv_pages(size, 256, 256), 33)
self.assertEqual(_num_dsv4_physical_kv_pages(size, 64, 256), 132)
self.assertGreaterEqual(
_num_dsv4_physical_kv_pages(size, 64, 256) * 64,
size + 256,
)
def test_state_buf_item_covers_one_logical_swa_page(self):
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool._unified_kv = False
pool.swa_page_size = 256
pool.compress_state_pools = []
pool.indexer_compress_state_pools = []
for physical_page_size in (256, 64):
with self.subTest(physical_page_size=physical_page_size):
row_bytes = physical_page_size * 4
buf = torch.empty((8, row_bytes), dtype=torch.uint8)
pool.swa_kv_pool = SimpleNamespace(
page_size=physical_page_size, kv_buffer=[buf]
)
data_ptrs, data_lens, item_lens = pool.get_state_buf_infos()
self.assertEqual(data_ptrs, [buf.data_ptr()])
self.assertEqual(data_lens, [buf.nbytes])
self.assertEqual(item_lens, [256 * 4])
def test_state_buf_infos_without_paged_swa(self):
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool.swa_kv_pool = None
pool.compress_state_pools = []
pool.indexer_compress_state_pools = []
self.assertEqual(pool.get_state_buf_infos(), ([], [], []))
def test_pp_mapping_and_pd_buffer_order(self):
for unified, stage_ratios in product(
(False, True), ([4, 0, 128, 4], [128], [0])
test/registered/unit/mem_cache/test_hybrid_pool_assembler.py
+18
View File
@@ -11,6 +11,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_evict_swa_for_device_alloc,
_MambaStrategy,
_MambaSwaStrategy,
_require_single_row_dsv4_swa_pages,
_split_hicache_size,
_SwaStrategy,
build_full_draft_pools,
@@ -22,6 +23,23 @@ from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
class TestDeepSeekV4SWAPageLayout(CustomTestCase):
def test_split_physical_rows_are_rejected_for_hicache_consumers(self):
with self.assertRaisesRegex(ValueError, "direct SWA KV layout"):
_require_single_row_dsv4_swa_pages(
logical_page_size=256,
physical_page_size=64,
consumer="test consumer",
)
def test_matching_page_geometry_is_supported(self):
_require_single_row_dsv4_swa_pages(
logical_page_size=256,
physical_page_size=256,
consumer="test consumer",
)
class _Pool:
def __init__(self, kv_bytes):
self._kv_bytes = kv_bytes
test/registered/unit/mem_cache/test_linker_pool_assembler.py
+2 -1
View File
@@ -179,7 +179,8 @@ class TestHybridDevicePoolAssembler(CustomTestCase):
kvcache.end_layer = 4
kvcache.swa_page_size = 2
kvcache.swa_kv_pool = SimpleNamespace(
kv_buffer=[torch.zeros((8, 3), dtype=torch.uint8) for _ in range(3)]
page_size=2,
kv_buffer=[torch.zeros((8, 3), dtype=torch.uint8) for _ in range(3)],
)
kvcache.c4_kv_pool = SimpleNamespace(
kv_buffer=[torch.zeros((8, 5), dtype=torch.uint8) for _ in range(2)],