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

Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>

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
+            # FP8 wo_a stays opt-in on SM120: only recent DeepGEMM builds ship
-            # DeepGEMM or require >99KB SMEM (topk_v2).
+            # the SM120 kernels, and deep_gemm_wrapper.configurer validates them.
-            envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
+            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

@@ -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
+    # FP8 W_o GEMM needs DeepGEMM JIT. Enable exactly where the runtime can
-    # it, mirroring the forward scale split: the ue8m0 path
+    # run it, mirroring the forward scale split: the default sm100 UE8M0
-    # (DEEPGEMM_SCALE_UE8M0, true sm100, default on) or an sm90 opt-in
+    # path, or explicit opt-in on sm90 (FP32 scales) and sm120 (UE8M0).
-    # fp32-scale path (use FP4 expert ckpt). Disable in every other case.
+    # 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

@@ -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

@@ -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 = swa_k_cache[:, : swa_kv_page_size * k_cache_total_dim].view(
-                swa_k_cache.shape[0], swa_page_size, 1, k_cache_total_dim
+                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

@@ -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;
+    # SM120/SM121 support (mma.sync block-scale, no TMEM) landed in DeepGEMM#324;
-    # probe the entry point since installed builds may predate it.
+    # probe every API used by the SM120 DSV4 paths since installed builds may
-    if sm_version in (120, 121):
+    # expose fp8_einsum but still predate the SM120 kernels.
-        try:
+    if sm_version in (120, 121) and not _sm120_deep_gemm_apis_available():
-            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):
         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

@@ -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
+            if self.schedule_stream.cuda_stream == self.forward_stream.cuda_stream:
-            while (
+                self.schedule_stream = allocate_distinct_stream(
-                self.schedule_stream.cuda_stream == self.forward_stream.cuda_stream
+                    self.device_module, (self.forward_stream,)
-                and _redraws < 64
+                )
-            ):
-                self.schedule_stream = self.device_module.Stream(priority=0)
-                _redraws += 1
         # 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

@@ -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._pp_send_proxy_to_next_stage(result)
-                            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_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._pp_send_proxy_to_next_stage(result)
-                            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_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._pp_send_proxy_to_next_stage(result)
-                            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_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]]] = (
+        self.pp_proxy_recv_event = None
-            defaultdict(deque)
+        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:
+    def _pp_record_comm_event(self: Scheduler) -> Optional[torch.Event]:
-        for p2p_work in work:
+        """Mark the tail of the PP comm stream; None when not overlapping."""
-            p2p_work.work.wait()
+        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(
+        with self.pp_comm_stream_ctx:
-            self.pp_group.send_tensor_dict(
+            if ready_event is not None:
-                tensor_dict=tensor_dict,
+                self.device_module.current_stream().wait_event(ready_event)
-                all_gather_group=(self.attn_tp_group),
+            p2p_work.extend(
-                async_send=async_send,
+                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(
+            with self.pp_comm_stream_ctx:
-                all_gather_group=all_gather_group
+                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(
+            tensor_dict, self.pp_proxy_recv_event = self._pp_recv_typed_dict(
-                self._pp_recv_typed_dict(
+                expected_kind="proxy",
-                    expected_kind="proxy",
+                all_gather_group=(self.attn_tp_group),
-                    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

@@ -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

@@ -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
+        is the model-wide logical page size. CUDA pools use it to reserve enough
-        for the per-ratio c4/c128 pools); the default CUDA pool ignores it.
+        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

@@ -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

@@ -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

@@ -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 = _FP8_WO_A_UE8M0
-
-            self.wo_a.weight_scale_inv.format_ue8m0 = (
-                deep_gemm_wrapper.DEEPGEMM_SCALE_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:
+                if _FP8_WO_A_UE8M0:
-                    # sm100 (Blackwell): ue8m0 scales via the dedicated JIT kernel.
+                    # 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 _FP8_WO_A_UE8M0:
-
-        if deep_gemm_wrapper.DEEPGEMM_SCALE_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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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)],