diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 77da73e8f..5867fb299 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -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) diff --git a/python/sglang/srt/arg_groups/serving_hook.py b/python/sglang/srt/arg_groups/serving_hook.py index fc2ccfead..8993bd638 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -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, ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d13e1b4f6..cc9a1a59f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 96a1cc223..91fcd94f5 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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: diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index b467a153a..12c6963e2 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2fbc883bf..dc73c6e65 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 4a602b43a..2566554c4 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -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", diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index f739fb998..bfcb97124 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -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 diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 388bd1578..d05dd0819 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 78acf93e5..9cb0d6c1b 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -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, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py index f7973f0d6..0df7c4413 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py @@ -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: " diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a81cfa693..b1669f362 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 8649b95ba..af6a62fdc 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -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 diff --git a/test/registered/unit/distributed/test_pp_comm_overlap.py b/test/registered/unit/distributed/test_pp_comm_overlap.py new file mode 100644 index 000000000..2243ef754 --- /dev/null +++ b/test/registered/unit/distributed/test_pp_comm_overlap.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py index 0b345f3be..06e677a28 100644 --- a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py +++ b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py @@ -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]) diff --git a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py index 152aff6e7..2e958ffea 100644 --- a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py +++ b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_linker_pool_assembler.py b/test/registered/unit/mem_cache/test_linker_pool_assembler.py index 069cad2d6..de912c7a1 100644 --- a/test/registered/unit/mem_cache/test_linker_pool_assembler.py +++ b/test/registered/unit/mem_cache/test_linker_pool_assembler.py @@ -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)],