[AMD][DSV4] feat: enable fp8 two-pool unified_kv on gfx950 (#37413)

python/sglang/srt/environ.py

@@ -1455,6 +1455,12 @@ class Envs:
     # Quantize the SWA fp8 KV cache from bf16-rounded values (matches
     # trainer-side QAT and the DSA-CP path) instead of fp32 registers.
     SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False)
+    # unified_kv only: split the pool into an fp8 nope pool plus a parallel
+    # bf16 rope pool, 640 B/token instead of 1024. The unified pool takes no
+    # dtype, so --kv-cache-dtype has no effect there and this switch is the
+    # only way to ask; on separate-KV it is the reverse -- --kv-cache-dtype
+    # picks the buffer dtype and this switch is inert.
+    SGLANG_DSV4_UNIFIED_KV_FP8 = EnvBool(False)
 
     # Kernels and indexer
     SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True)

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -1366,8 +1366,17 @@ class DeepseekV4HipRadixBackend(
         attn_sink: torch.Tensor,
         core_attn_metadata: DSV4AttnMetadata,
         save_kv_cache: bool = True,
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        """unified_kv paged-attention path over the bf16 unified_kv"""
+        """unified_kv paged-attention path over the unified_kv pool.
+
+        ``q_rope`` is what tells the two layouts apart: present means ``q`` is a
+        packed fp8 row and the pool is the two-pool fp8 one, so decode goes to
+        the asm reader; absent means both are plain bf16 and it goes to Triton.
+        Prefill needs ``k_rope`` alongside it, because there the current chunk is
+        a KV source of its own and not just something to store.
+        """
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
 
         pool = self.token_to_kv_pool
@@ -1401,6 +1410,10 @@ class DeepseekV4HipRadixBackend(
             else:
                 state_slot = forward_batch.req_pool_indices[:T]
             if save_kv_cache:
+                # Only verify reaches this under fp8 -- plain decode's rows are
+                # written by the fused kernel itself, which leaves kv None. The
+                # pair arrives already packed, so this is the same scatter with
+                # a second pool hanging off it.
                 runtime.store_swa_into_unified(
                     kv=kv,
                     state_slot=state_slot,
@@ -1409,6 +1422,10 @@ class DeepseekV4HipRadixBackend(
                     win=win,
                     ring_stride=ring_stride,
                     final_pos=positions,
+                    kv_rope=k_rope,
+                    unified_kv_rope=(
+                        None if k_rope is None else pool.get_unified_kv_rope(layer_id)
+                    ),
                 )
             unified_metadata = core_attn_metadata.unified
             if compress_ratio == 0:
@@ -1430,6 +1447,25 @@ class DeepseekV4HipRadixBackend(
                 )
             else:
                 raise ValueError(f"bad compress_ratio {compress_ratio}")
+            if q_rope is not None:
+                # softmax_scale is not passed on: the asm kernel hardcodes
+                # 1/sqrt(512), which is what self.softmax_scale already is for
+                # V4's head_dim=512. The other readers here take it explicitly,
+                # so a head_dim change would leave only this one mis-scaled.
+                assert self.softmax_scale == 512**-0.5, (
+                    "the v4 nm asm kernel hardcodes 1/sqrt(512), this backend is "
+                    f"at {self.softmax_scale}"
+                )
+                return runtime.decode_fp8_2buff(
+                    q=q,
+                    q_rope=q_rope,
+                    unified_kv=unified,
+                    unified_kv_rope=pool.get_unified_kv_rope(layer_id),
+                    kv_indices=kv_indices,
+                    kv_indptr=kv_indptr,
+                    attn_sink=attn_sink,
+                    v_head_dim=layer.v_head_dim,
+                )
             return runtime.decode(
                 q=q,
                 unified_kv=unified,
@@ -1505,17 +1541,42 @@ class DeepseekV4HipRadixBackend(
             pad = T + 1 - kpre_p.shape[0]
             kpre_p = torch.cat([kpre_p, kpre_p[-1:].expand(pad)])
             kext_p = torch.cat([kext_p, kext_p[-1:].expand(pad)])
-        o = runtime.prefill(
-            q=q,
-            unified_kv=unified,
-            kv_indices_prefix=kpre_i,
-            kv_indptr_prefix=kpre_p,
-            kv_extend=kv,
-            kv_indices_extend=kext_i,
-            kv_indptr_extend=kext_p,
-            attn_sink=attn_sink,
-            softmax_scale=self.softmax_scale,
-        )
+        if q_rope is not None:
+            assert k_rope is not None, (
+                "fp8 prefill needs the extend rope half beside the packed nope; "
+                "q_rope came through but k_rope did not"
+            )
+            # No empty-segment mask on the result, unlike decode: this kernel
+            # returns zeros for a token with neither region where the asm decode
+            # reader leaves the row NaN. Chunk 0 tokens have an empty prefix and
+            # a non-empty extend, which both readers handle.
+            o = runtime.prefill_fp8_2buff(
+                q=q,
+                q_rope=q_rope,
+                unified_kv=unified,
+                unified_kv_rope=pool.get_unified_kv_rope(layer_id),
+                kv_indices_prefix=kpre_i,
+                kv_indptr_prefix=kpre_p,
+                kv_extend=kv,
+                kv_extend_rope=k_rope,
+                kv_indices_extend=kext_i,
+                kv_indptr_extend=kext_p,
+                attn_sink=attn_sink,
+                softmax_scale=self.softmax_scale,
+                v_head_dim=layer.v_head_dim,
+            )
+        else:
+            o = runtime.prefill(
+                q=q,
+                unified_kv=unified,
+                kv_indices_prefix=kpre_i,
+                kv_indptr_prefix=kpre_p,
+                kv_extend=kv,
+                kv_indices_extend=kext_i,
+                kv_indptr_extend=kext_p,
+                attn_sink=attn_sink,
+                softmax_scale=self.softmax_scale,
+            )
 
         # write this chunk's SWA K into the ring for future chunks / decode
         # only the final-window tokens per request
@@ -1535,6 +1596,10 @@ class DeepseekV4HipRadixBackend(
                 win=win,
                 ring_stride=ring_stride,
                 final_pos=_ring_final_pos,
+                kv_rope=None if k_rope is None else k_rope[:n_real],
+                unified_kv_rope=(
+                    None if k_rope is None else pool.get_unified_kv_rope(layer_id)
+                ),
             )
         return o
 
@@ -1602,6 +1667,8 @@ class DeepseekV4HipRadixBackend(
         compress_ratio: Literal[0, 4, 128],
         save_kv_cache: bool = True,
         attn_sink: Optional[torch.Tensor] = None,
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
         **_,
     ) -> torch.Tensor:
         if self.mtp_enabled and forward_batch.forward_mode.is_idle():
@@ -1630,6 +1697,8 @@ class DeepseekV4HipRadixBackend(
                 attn_sink=attn_sink,
                 core_attn_metadata=core_attn_metadata,
                 save_kv_cache=save_kv_cache,
+                q_rope=q_rope,
+                k_rope=k_rope,
             )
 
         if isinstance(core_attn_metadata, DSV4AttnMetadata):

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -158,6 +158,8 @@ class CompressorBackendMixin:
         bf16_store: bool = False,
         kv_scale_cache: Optional[torch.Tensor] = None,
         rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        fp8_2buff: bool = False,
+        kv_cache_rope: Optional[torch.Tensor] = None,
     ) -> None:
         assert compress_ratio == 4 or compress_ratio == 128
         assert rotate == is_indexer == (head_dim == 128)
@@ -220,6 +222,8 @@ class CompressorBackendMixin:
                 if _is_hip and use_fp4_indexer
                 else None
             ),
+            fp8_2buff=fp8_2buff,
+            kvcache_rope=kv_cache_rope,
         )
 
     def forward_unified(
@@ -238,6 +242,7 @@ class CompressorBackendMixin:
 
         state_pool = compressor.get_state_pool(self)
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
             is_unified_kv_triton,
         )
 
@@ -264,6 +269,8 @@ class CompressorBackendMixin:
         use_hip_fp4 = _is_hip and use_fp4_indexer
         bf16_store = False
         kv_scale_cache = None
+        fp8_2buff = False
+        kv_cache_rope = None
         if compressor.is_in_indexer:
             page_size = token_to_kv_pool.get_index_k_page_size(compressor.ratio)
             if use_hip_fp4:
@@ -278,7 +285,11 @@ class CompressorBackendMixin:
                 self.forward_metadata.core_metadata.unified,
                 f"c{compressor.ratio}_out_loc",
             )
-            bf16_store = True
+            if is_unified_kv_fp8():
+                fp8_2buff = True
+                kv_cache_rope = token_to_kv_pool.get_unified_kv_rope(layer_id)
+            else:
+                bf16_store = True
         else:
             _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
             assert compress_kv_pool is not None
@@ -305,6 +316,10 @@ class CompressorBackendMixin:
             rope_cache=(
                 (compressor.fp4_cos, compressor.fp4_sin) if use_hip_fp4 else None
             ),
+            fp8_2buff=fp8_2buff,
+            kv_cache_rope=(
+                None if kv_cache_rope is None else kv_cache_rope.view(dtype=torch.uint8)
+            ),
         )
         online_c128_mtp = getattr(self, "online_c128_mtp", None)
         if online_c128_mtp is not None:

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -16,6 +16,7 @@ from sglang.kernels.ops.attention.dsv4 import (
     index_buf_accessor as dsv4_index_buf_accessor,
 )
 from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import layout
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.environ import envs
 from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
@@ -511,11 +512,43 @@ class DeepSeekV4LayerItem(NamedTuple):
     compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None
 
 
+# re-exported: the pool allocates the rows, but the kernels that write them own the
+# layout (see unified_kv_kernels/layout.py)
+DSV4_FP8_NOPE_ROW_BYTES = layout.DSV4_FP8_NOPE_ROW_BYTES
+DSV4_FP8_QUANT_TILE = layout.DSV4_FP8_QUANT_TILE
+
+
+def dsv4_unified_row_bytes(
+    qk_nope_head_dim: int, qk_rope_head_dim: int, fp8: bool
+) -> int:
+    """Bytes one unified_kv token occupies, summed over both pools."""
+    if not fp8:
+        return (qk_nope_head_dim + qk_rope_head_dim) * 2
+    num_tiles = -(-qk_nope_head_dim // DSV4_FP8_QUANT_TILE)
+    scale_bytes = 2 * num_tiles
+    # not an assert: sizing runs under -O too, and a silently skipped check here
+    # overreports capacity
+    if qk_nope_head_dim + scale_bytes > DSV4_FP8_NOPE_ROW_BYTES:
+        raise ValueError(
+            f"fp8 nope row overflows: {qk_nope_head_dim} latent values at 1 B + "
+            f"{scale_bytes} B scales > {DSV4_FP8_NOPE_ROW_BYTES} B stride"
+        )
+    return DSV4_FP8_NOPE_ROW_BYTES + qk_rope_head_dim * 2
+
+
 # The following kv pool follows ATOM's unified_kv kernel layout.
 class DeepSeekV4UnifiedKVPool:
     """
-    Layout:
+    Layout (bf16):
     unified_kv[L]: ``[swa_pages + padded_compress_rows, head_dim]`` bf16
+
+    Layout (fp8, ``SGLANG_DSV4_UNIFIED_KV_FP8``) -- two parallel pools with the
+    same row count, so a row index means the same thing in both. Named after the
+    accessors, which under fp8 each return one half -- ``get_unified_kv`` the
+    nope, ``get_unified_kv_rope`` the rope:
+    unified_kv[L]      (nope): ``[rows, 512]`` fp8, see DSV4_FP8_NOPE_ROW_BYTES
+    unified_kv_rope[L] (rope): ``[rows, qk_rope_head_dim]`` bf16, never quantized
+
     - rows ``[0, swa_pages)``   = SWA ring (``req_pool_indices * swa_window + pos % swa_window``)
     - rows ``[swa_pages, ...)`` = compressed (``swa_pages + page_index``)
     """
@@ -535,8 +568,11 @@ class DeepSeekV4UnifiedKVPool:
         memory_saver_adapter,
         custom_mem_pool,
         swa_ring_size: int,
+        fp8: bool = False,
     ):
         self.swa_ring_size = swa_ring_size
+        self.fp8 = fp8
+        self.rope_head_dim = qk_rope_head_dim
         self.head_dim = qk_nope_head_dim + qk_rope_head_dim
         self.num_slots = num_slots
         self.swa_pages = num_slots * self.swa_ring_size
@@ -545,6 +581,7 @@ class DeepSeekV4UnifiedKVPool:
         self.k_per_block = dict(self.K_PER_BLOCK)
 
         bufs = []
+        rope_bufs = []
         with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
             with (
                 torch.cuda.use_mem_pool(custom_mem_pool)
@@ -557,20 +594,54 @@ class DeepSeekV4UnifiedKVPool:
                     compress_rows = self.num_blocks * self.k_per_block[ratio]
                     rows_per_page = self.page_size // ratio if ratio else 0
                     padded_compress_rows = compress_rows + rows_per_page
-                    bufs.append(
-                        torch.zeros(
-                            self.swa_pages + padded_compress_rows,
-                            self.head_dim,
-                            dtype=torch.bfloat16,
-                            device=device,
+                    rows = self.swa_pages + padded_compress_rows
+                    if self.fp8:
+                        bufs.append(
+                            torch.zeros(
+                                rows,
+                                DSV4_FP8_NOPE_ROW_BYTES,
+                                dtype=torch.float8_e4m3fn,
+                                device=device,
+                            )
                         )
-                    )
+                        rope_bufs.append(
+                            torch.zeros(
+                                rows,
+                                self.rope_head_dim,
+                                dtype=torch.bfloat16,
+                                device=device,
+                            )
+                        )
+                    else:
+                        bufs.append(
+                            torch.zeros(
+                                rows,
+                                self.head_dim,
+                                dtype=torch.bfloat16,
+                                device=device,
+                            )
+                        )
+                        rope_bufs.append(None)
         self.kv_buffer = bufs
+        self.kv_buffer_rope = rope_bufs
 
     def get_unified_kv(self, local_layer_id: int) -> torch.Tensor:
         return self.kv_buffer[local_layer_id]
 
+    def get_unified_kv_rope(self, local_layer_id: int) -> torch.Tensor:
+        assert self.fp8, "rope pool only exists under SGLANG_DSV4_UNIFIED_KV_FP8"
+        return self.kv_buffer_rope[local_layer_id]
+
     def get_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
+        if self.fp8:
+            # same single-pool assumption as the outer get_contiguous_buf_infos:
+            # one pointer and one row size per layer describes the nope pool only,
+            # so whoever picks this up next would move half a row and not notice.
+            # TODO(danli103): report both pools once a consumer needs them.
+            raise NotImplementedError(
+                "get_buf_infos describes one pool per layer; the fp8 rope pool "
+                "would be dropped (SGLANG_DSV4_UNIFIED_KV_FP8=1)."
+            )
         data_ptrs = [b.data_ptr() for b in self.kv_buffer]
         data_lens = [b.nbytes for b in self.kv_buffer]
         item_lens = [b[0].nbytes for b in self.kv_buffer]
@@ -578,6 +649,10 @@ class DeepSeekV4UnifiedKVPool:
 
 
 class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
+    # object.__new__ stubs (disagg wire test) skip __init__; False is the env
+    # default, so the fp8 PD/HiCache refuses don't AttributeError on them.
+    _unified_kv_fp8 = False
+
     def __init__(
         self,
         max_num_reqs: int,
@@ -633,11 +708,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         self.c4_logical_size = c4_logical_size
         self.c128_size = c128_size
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
             is_unified_kv_triton,
         )
 
         # Resolve the unified-kv gate before any sizing so the two cannot drift.
         self._unified_kv = is_unified_kv_triton()
+        self._unified_kv_fp8 = is_unified_kv_fp8()
         # Uniform 512-dim e4m3 layout for the trtllm attention backend
         self.uniform_fp8 = (
             not self._unified_kv
@@ -721,6 +798,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
                 memory_saver_adapter=self.memory_saver_adapter,
                 custom_mem_pool=self.custom_mem_pool,
                 swa_ring_size=swa_ring_size,
+                fp8=self._unified_kv_fp8,
             )
 
             self.unified_swa_window = self.sliding_window
@@ -766,6 +844,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         self.wait_layer_transfer(layer_id)
         return self.unified_kv_pool.get_unified_kv(layer_id - self._stage_start)
 
+    def get_unified_kv_rope(self, layer_id: int) -> torch.Tensor:
+        self.wait_layer_transfer(layer_id)
+        return self.unified_kv_pool.get_unified_kv_rope(layer_id - self._stage_start)
+
     def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
         self.full_to_swa_index_mapping = full_to_swa_index_mapping
 
@@ -782,6 +864,18 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         data_lens: List[int] = []
         item_lens: List[int] = []
 
+        if self._unified_kv_fp8:
+            # The page-block transfer below prices one row as buf[0].nbytes and
+            # ships a single pointer per layer. Under fp8 that covers the nope
+            # pool only -- the parallel bf16 rope pool would be dropped and the
+            # remote side would decode rows against stale rope. Refuse instead.
+            # TODO(danli103): ship the rope pool as a second per-layer entry.
+            raise NotImplementedError(
+                "PD disaggregation is not supported with "
+                "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the transfer assumes a single "
+                "unified pool; the rope pool would be silently dropped)."
+            )
+
         def append_page_buffer(buf: torch.Tensor) -> None:
             assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
             data_ptrs.append(buf.data_ptr())
@@ -827,6 +921,15 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         item_lens: List[int] = []
         if not self._unified_kv:
             return data_ptrs, data_lens, item_lens
+        if self._unified_kv_fp8:
+            # Other half of the PD path -- get_contiguous_buf_infos ships the
+            # compressed region, this one the ring. Same single-pool assumption,
+            # same silently dropped rope, same fix -- land them together.
+            raise NotImplementedError(
+                "PD disaggregation is not supported with "
+                "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the SWA_RING component assumes a "
+                "single unified pool; the rope pool would be silently dropped)."
+            )
         swa_pages = self.unified_kv_pool.swa_pages
         for buf in self.unified_kv_pool.kv_buffer:
             assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
@@ -841,6 +944,17 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         # the unified pool stores individual token rows after its SWA region.
         assert self._unified_kv, "unified_region_buffers requires unified_kv layout"
         assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
+        if self._unified_kv_fp8:
+            # item_bytes below prices kv_buffer alone, so the rope pool would never
+            # be offloaded and a fetched page would carry stale rope -- wrong output,
+            # no crash.
+            # TODO(danli103): give rope its own host pool, the way C4_INDEXER
+            # already parallels C4.
+            raise NotImplementedError(
+                "HiCache offload is not supported with "
+                "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the host pool assumes a single "
+                "unified pool; the rope pool would never be offloaded)."
+            )
 
         swa_pages = self.unified_kv_pool.swa_pages
         head_dim = self.unified_kv_pool.head_dim

python/sglang/srt/model_executor/pool_configurator.py

@@ -986,16 +986,56 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4)
         self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128)
 
+        # Unified-KV uses a different physical layout than the non-unified V4 path:
+        #  * one row carries the full latent -- 1024 B bf16, or 640 B under
+        #    SGLANG_DSV4_UNIFIED_KV_FP8 (512 B fp8 nope + 128 B bf16 rope) -- not
+        #    that path's 584-byte fp8(nope) + bf16(rope) + scales cell.
+        #  * SWA is a fixed per-request ring (num_req_slots * ring_size),
+        #    independent of full_token, so it is a fixed *bias* rather than a
+        #    per-token term. Gate on the same switch the pool itself uses so the
+        #    sizing and the allocation never drift apart.
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
             is_unified_kv_triton,
         )
+        from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+            dsv4_unified_row_bytes,
+        )
 
         self._unified = is_unified_kv_triton()
+        self._unified_fp8 = is_unified_kv_fp8()
         self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        # Row width across both pools: 1024 B bf16, 640 B fp8. Read from the pool
+        # module so sizing can't drift from the allocation.
+        self._unified_row_bytes = dsv4_unified_row_bytes(
+            self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8
+        )
         # swa_page_size is the model's sliding window (cfg.window_size).
         self._swa_ring_size = get_swa_ring_size(self.swa_page_size, self.is_speculative)
         self._spec_infl = 1.0
 
+        # The unified pool takes no dtype, so --kv-cache-dtype never reaches it.
+        # V4 defaults "auto" to fp8_e4m3 (overrides.py
+        # _deepseek_v4_kv_cache_dtype), so only a bfloat16 here tells us the user
+        # set it explicitly; warning on the fp8 side would fire on every run.
+        if self._unified_fp8 and self.kv_cache_dtype_str == "bfloat16":
+            logger.warning(
+                "--kv-cache-dtype=bfloat16 is ignored on the unified_kv path; "
+                "SGLANG_DSV4_UNIFIED_KV_FP8=1 stores the latent as fp8. Unset the "
+                "env switch to get a bf16 unified pool."
+            )
+
+        # get_contiguous_buf_infos ships one pointer per layer and prices a row as
+        # buf[0].nbytes, which under fp8 covers the nope pool only. Fail at startup
+        # rather than at the first transfer.
+        # TODO(danli103): drop this once the transfer ships the rope pool.
+        if self._unified_fp8 and self.disaggregation_mode != "null":
+            raise ValueError(
+                "SGLANG_DSV4_UNIFIED_KV_FP8=1 does not support PD disaggregation "
+                f"(disaggregation_mode={self.disaggregation_mode!r}). Unset the fp8 "
+                "switch or run without disaggregation."
+            )
+
         if self.is_speculative:
             # Ring is sized once here, so it must serve the largest adaptive tier.
             self._assert_ring_serves_draft_tokens(
@@ -1065,8 +1105,10 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
 
     def _get_bytes_per_full_token(self) -> float:
         if self._unified:
-            # Unified_kv stores the whole latent in bf16.
-            kv_bytes = self.attn_head_dim * 2
+            # Unified_kv stores the whole latent: one bf16 pool, or an fp8 nope
+            # pool plus a bf16 rope pool. kv_bytes also prices the compressed
+            # c4/c128 rows below, which live in the same pool(s).
+            kv_bytes = self._unified_row_bytes
         else:
             kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8
 
@@ -1199,14 +1241,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         return min(estimated, full_token // 2)
 
     def _fixed_swa_bytes(self, max_running_requests: int) -> int:
+        """Unified_kv SWA is a fixed per-request ring, sized by concurrency
+        (num_req_slots) rather than by full_token. Return its byte footprint
+        across all full layers, inflated for the draft worker the same way as the
+        per-token coeff. Returns 0 on the non-unified path (where SWA is already
+        accounted per-token)."""
         if not self._unified:
             return 0
         num_req_slots = self._get_num_req_slots(max_running_requests)
         ring_bytes = (
             num_req_slots
             * self._swa_ring_size
-            * self.attn_head_dim
-            * 2  # bf16
+            * self._unified_row_bytes
             * self.num_layers_total
         )
         return int(ring_bytes * self._spec_infl)
@@ -1277,6 +1323,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         sizes = self._compute_dsv4_sizes(full_token, page_size)
         logger.info(
             f"DSV4 memory calculation: unified={self._unified}, "
+            f"unified_fp8={self._unified_fp8}, "
             f"bytes_per_full_token={self.bytes_per_full_token:.2f}, "
             f"available_bytes={available_bytes / (1 << 30):.2f} GB, "
             f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, "

python/sglang/srt/models/deepseek_v4.py

@@ -1423,6 +1423,9 @@ class MQALayer(MqaAttentionBase):
         attn_backend,
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
+        q_rope_out: Optional[torch.Tensor] = None,
+        k_nope_out: Optional[torch.Tensor] = None,
+        k_rope_out: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         x_linear = x_quant if x_quant is not None else x
 
@@ -1437,22 +1440,52 @@ class MQALayer(MqaAttentionBase):
         kv: Optional[torch.Tensor]
 
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
             is_unified_kv_triton,
         )
 
         unified = is_unified_kv_triton()
+        fp8_2buff = is_unified_kv_fp8()
         is_decode = forward_batch.forward_mode.is_decode_or_idle()
         # The kernel is token-indexed (q, kv and positions are all length M), so
         # a verify batch carrying several draft tokens per request is a shape it
         # already handles. Only the cache store differs between decode and
-        # verify, and that half is left off below.
+        # verify, and under fp8 that store takes the packed pair instead of bf16.
         fuse_verify = (
             envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get()
             and forward_batch.forward_mode.is_target_verify()
         )
-        do_fused_qk_norm_rope = (unified and (is_decode or fuse_verify)) or (
-            not unified and self.use_fused_qk_norm_rope
+        # fp8 verify packs like prefill but keeps verify's store timing: the pair
+        # lands in the caller's buffers and the backend writes the ring off the
+        # per-token slot map before attention. Keyed off those buffers the same
+        # way fuse_prefill is, so the two arms cannot disagree about the layout.
+        fuse_verify_fp8 = (
+            fuse_verify
+            and unified
+            and fp8_2buff
+            and k_nope_out is not None
+            and k_rope_out is not None
         )
+        # Prefill under fp8 goes through the same fused store: the 2-source
+        # kernel reads this chunk as its extend region in the pool's packed form,
+        # and the ring write after attention reuses those same rows, so they are
+        # materialised once here rather than quantized on both sides. Keyed off
+        # the caller's buffers the way q_rope_out keys the packed Q, so the two
+        # cannot disagree about the layout; both halves are required because the
+        # nope one leaves on the kv slot and a missing one would read as "the
+        # fused store did not run". Verify packs the same way but is its own arm
+        # above: it stores before attention, not after.
+        fuse_prefill = (
+            unified
+            and fp8_2buff
+            and k_nope_out is not None
+            and k_rope_out is not None
+            and not is_decode
+            and not forward_batch.forward_mode.is_target_verify()
+        )
+        do_fused_qk_norm_rope = (
+            unified and (is_decode or fuse_verify or fuse_prefill)
+        ) or (not unified and self.use_fused_qk_norm_rope)
 
         if do_fused_qk_norm_rope:
             if _is_gfx95_supported or _is_gfx1250_supported:
@@ -1473,6 +1506,7 @@ class MQALayer(MqaAttentionBase):
             )
 
             token_to_kv_pool = get_token_to_kv_pool()
+            swa_rope_cache = None
             if unified and fuse_verify:
                 # Target-verify runs through the unified_kv decode path. The
                 # backend writes the current chunk's KV into the ring *before*
@@ -1490,15 +1524,34 @@ class MQALayer(MqaAttentionBase):
                 # contiguous buffer, so materialise it before the kernel norms
                 # it in place. The unfused path pays the same copy inside
                 # _compute_kv_bf16.
+                #
+                # Under fp8 the kernel writes the packed pair to the caller's
+                # buffers rather than norming kv in place, and the same backend
+                # store takes that pair -- only the row format changes.
                 kv = kv.contiguous()
                 swa_cache, swa_loc = None, None
-                swa_page_size, bf16_store = 1, True
+                swa_page_size, bf16_store = 1, not fuse_verify_fp8
+            elif unified and fuse_prefill:
+                # No pools, so the kernel norms + RoPEs + packs and writes no
+                # ring row. It must not: those rows are this fwd's extend region
+                # and the prefix pool has to stay as attention expects to find
+                # it. The backend stores them after attention from the pair.
+                swa_cache, swa_loc = None, None
+                swa_page_size, bf16_store = 1, False
+                # kv stays the strided slice of qkv_a. Under fp8 the kernel only
+                # reads it -- the packed pair goes to k_nope_out/k_rope_out, it
+                # does not norm in place -- and it takes the row stride as an
+                # argument, so materialising it was a copy on every fp8 layer.
             elif unified:
                 swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
                 # swa_loc is layer-independent; computed once per forward by the
                 # backend and cached on the metadata (read here by every layer).
                 swa_loc = attn_backend.get_unified_swa_loc(forward_batch)
-                swa_page_size, bf16_store = 1, True
+                swa_page_size, bf16_store = 1, not fp8_2buff
+                if fp8_2buff:
+                    swa_rope_cache = token_to_kv_pool.get_unified_kv_rope(self.layer_id)
+                    # kv stays the strided slice of qkv_a -- the group-quant
+                    # kernel takes the row stride as an argument.
             else:
                 swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id)
                 swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
@@ -1528,13 +1581,25 @@ class MQALayer(MqaAttentionBase):
                 q_out=q_out,
                 dtype=x.dtype,
                 bf16_store=bf16_store,
+                fp8_2buff=fp8_2buff,
+                swa_rope_cache=swa_rope_cache,
+                k_nope_out=k_nope_out if (fuse_prefill or fuse_verify_fp8) else None,
+                k_rope_out=k_rope_out if (fuse_prefill or fuse_verify_fp8) else None,
+                q_rope_out=q_rope_out,
             )
             # On the verify path the kernel normed + RoPE'd kv in place and wrote
             # nothing, so hand it back: the caller feeds it to attention as the
             # current chunk (attn_k = kv) and save_kv_cache = kv is not None lets
             # the backend do its normal causally-indexed store into the ring
             # before the decode kernel runs -- exactly as the unfused path did.
-            if not (unified and fuse_verify):
+            if unified and (fuse_prefill or fuse_verify_fp8):
+                # The packed nope half rides out on the kv slot -- attention
+                # takes it as attn_k and save_kv_cache stays on so the backend
+                # does the ring write. Its rope half went to the caller's buffer,
+                # which has no second return slot here. Prefill's write lands
+                # after attention, verify's before it; both read this pair.
+                kv = k_nope_out
+            elif not (unified and fuse_verify):
                 kv = None
 
             if not unified and use_cp:
@@ -1657,21 +1722,97 @@ class MQALayer(MqaAttentionBase):
             and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
         )
 
-        tp_slice, q_padded, q_out = slice(None), None, None
-        kernel_num_heads = self._kernel_num_heads(x.shape[0])
-        if kernel_num_heads != self.n_local_heads:
-            # Backends without an exact-head specialization retain the existing
-            # padded shape. attn_sink is sliced to this rank and padded to match.
-            # Only [0:n_local_heads] is written below. Uninitialized padded TP
-            # heads inject NaN into attention on gfx942 (fnuz), so zero-init
-            # there; other archs tolerate new_empty and skip the per-forward
-            # memset.
-            if _is_gfx942_supported:
-                q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim)
-            else:
-                q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim)
-            tp_slice = slice(0, self.n_local_heads)
-            q_out = q_padded[:, tp_slice, :]
+        from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
+            is_unified_kv_triton,
+        )
+
+        unified = is_unified_kv_triton()
+        unified_fp8_verify = (
+            unified
+            and is_unified_kv_fp8()
+            and forward_batch.forward_mode.is_target_verify()
+        )
+        # The v4 nm asm reader takes Q in the pool's own packed form, so fp8
+        # decode wants a contiguous fp8 buffer of exactly the local heads --
+        # q_padded below is a FlashMLA layout and buys nothing here. Verify runs
+        # that same reader over the ring, so it takes the same Q.
+        unified_fp8_decode = (
+            unified
+            and is_unified_kv_fp8()
+            and (forward_batch.forward_mode.is_decode_or_idle() or unified_fp8_verify)
+        )
+        # The 2-source prefill kernel wants the same packed Q plus this chunk's
+        # K in the pool's layout. Verify is not prefill here even though it takes
+        # the same branch below -- it reads rows the ring already holds, so it
+        # goes with decode above. Multi-stream picks a different prepare that has
+        # no unified arm at all, so it keeps the bf16 buffers it always had.
+        unified_fp8_prefill = (
+            unified
+            and is_unified_kv_fp8()
+            and not enable_multi_stream
+            and not forward_batch.forward_mode.is_decode_or_idle()
+            and not forward_batch.forward_mode.is_target_verify()
+        )
+        if unified_fp8_verify and not envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get():
+            # The packed pair is produced by the fused norm+RoPE store; with that
+            # off the unfused arm hands the backend bf16 kv and the ring scatter
+            # dies on a dtype assert that says nothing about MTP.
+            raise NotImplementedError(
+                "fp8 two-pool unified_kv needs the fused verify store for "
+                "speculative decoding: set "
+                "SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY=1, or run with "
+                "SGLANG_DSV4_UNIFIED_KV_FP8=0."
+            )
+        if (
+            unified
+            and is_unified_kv_fp8()
+            and self.dsa_enable_prefill_cp
+            and dsa_use_prefill_cp(forward_batch)
+            and not forward_batch.forward_mode.is_decode_or_idle()
+        ):
+            # The gather hands back bf16 kv in global token order *after*
+            # norm+RoPE, so packing would have to move ahead of it and re-derive
+            # RoPE from global-order positions. Whether the CP path has those
+            # ready is unverified, so refuse instead of packing the wrong order.
+            raise NotImplementedError(
+                "fp8 two-pool unified_kv does not support DSA prefill CP "
+                "(SGLANG_DSV4_UNIFIED_KV_FP8=1 with cp_size > 1)."
+            )
+
+        tp_slice, q_padded, q_out, q_rope = slice(None), None, None, None
+        k_nope, k_rope = None, None
+        if unified_fp8_decode or unified_fp8_prefill:
+            # width and dtype come off the pools themselves; the kernel reads Q
+            # with the kv row stride, so the two must not drift
+            kv_pool = get_token_to_kv_pool()
+            nope_pool = kv_pool.get_unified_kv(self.layer_id)
+            rope_pool = kv_pool.get_unified_kv_rope(self.layer_id)
+            q_out = nope_pool.new_empty(
+                (x.shape[0], self.n_local_heads, nope_pool.shape[-1])
+            )
+            q_rope = rope_pool.new_empty(
+                (x.shape[0], self.n_local_heads, rope_pool.shape[-1])
+            )
+            if unified_fp8_prefill or unified_fp8_verify:
+                k_nope = nope_pool.new_empty((x.shape[0], nope_pool.shape[-1]))
+                k_rope = rope_pool.new_empty((x.shape[0], rope_pool.shape[-1]))
+            kernel_num_heads = self.n_local_heads
+        else:
+            kernel_num_heads = self._kernel_num_heads(x.shape[0])
+            if kernel_num_heads != self.n_local_heads:
+                # Backends without an exact-head specialization retain the existing
+                # padded shape. attn_sink is sliced to this rank and padded to match.
+                # Only [0:n_local_heads] is written below. Uninitialized padded TP
+                # heads inject NaN into attention on gfx942 (fnuz), so zero-init
+                # there; other archs tolerate new_empty and skip the per-forward
+                # memset.
+                if _is_gfx942_supported:
+                    q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim)
+                else:
+                    q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim)
+                tp_slice = slice(0, self.n_local_heads)
+                q_out = q_padded[:, tp_slice, :]
         attn_sink = self._local_attn_sink(kernel_num_heads)
 
         if enable_multi_stream:
@@ -1713,6 +1854,9 @@ class MQALayer(MqaAttentionBase):
                 attn_backend,
                 q_out,
                 x_quant=x_quant,
+                q_rope_out=q_rope,
+                k_nope_out=k_nope,
+                k_rope_out=k_rope,
             )
 
         # save_kv_cache = kv is not None selects who writes the ring. When kv is
@@ -1723,11 +1867,16 @@ class MQALayer(MqaAttentionBase):
         # _forward_prepare* deliberately left the store off and the backend does
         # its normal causally-indexed store from attn_k = kv.
         attn_k = kv if kv is not None else q
-        from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
-            is_unified_kv_triton,
-        )
 
-        if is_unified_kv_triton():
+        if unified:
+            # only the HIP radix backend takes these two; passing them always would
+            # leave non-ROCm depending on the **_ in its forward() to drop them, and
+            # no test on that side would notice if the **_ went away
+            rope_kwargs = {}
+            if q_rope is not None:
+                rope_kwargs["q_rope"] = q_rope
+            if k_rope is not None:
+                rope_kwargs["k_rope"] = k_rope
             o = attn_backend.forward(
                 q=q_out if q_out is not None else q,
                 k=attn_k,
@@ -1737,6 +1886,7 @@ class MQALayer(MqaAttentionBase):
                 compress_ratio=self.compress_ratio,
                 attn_sink=attn_sink[: self.n_local_heads],
                 save_kv_cache=kv is not None,
+                **rope_kwargs,
             )
         else:
             attn_q = q_padded if q_padded is not None else q