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

test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""The backend's own prefill wiring: what reaches the reader, and what lands in the ring.
+
+The pieces on either side of this are covered elsewhere -- the scatter primitive by
+test_dsv4_unified_fp8_scatter, the model->backend kwargs by the q_pair test -- but
+the middle, where _forward_unified_kv picks the fp8 arm and hands the packed pair to
+both attention and the ring write, had nothing running through it.
+
+Losing the rope half of that write is silent: the nope pool gets this chunk's rows,
+the rope pool keeps stale ones, and later chunks plus decode read a wrong RoPE with
+no crash and no NaN. So these run the real store against real (small) pools and pin
+that both pools got written, on the same ring row. The attention reader is stubbed:
+it is covered by test_dsv4_unified_fp8_prefill, and the store is what is at stake.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+import sglang.srt.layers.attention.deepseek_v4_backend_hip_radix as backend_mod
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# the store is a plain row move, but the two-pool layout it pins is gfx95-only, so
+# run it where the feature lives rather than on the default mi300 runner
+register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
+ROPE_DIM = 64
+V_HEAD_DIM = 512
+NUM_HEADS = 16
+
+WIN = 8
+RING_STRIDE = 8
+SWA_PAGES = 24  # ring rows are state_slot * RING_STRIDE + pos % RING_STRIDE, so < 24
+POOL_ROWS = 32
+
+# distinctive fill, so "the store never ran here" and "the store wrote zeros" are
+# different failures
+NOPE_SENTINEL = 0xEE
+ROPE_SENTINEL = -7.0
+
+# two requests on ring slots 1 and 2, three tokens each at positions 0..2
+STATE_SLOT = [1, 1, 1, 2, 2, 2]
+POSITIONS = [0, 1, 2, 0, 1, 2]
+CU_Q = [0, 0, 0, 3, 3, 3]
+EXPECTED_ROWS = [8, 9, 10, 16, 17, 18]
+
+_needs_gfx950 = unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "the two-pool fp8 layout is gfx95-only",
+)
+
+
+def _ints(values):
+    return torch.tensor(values, dtype=torch.int32, device=DEVICE).contiguous()
+
+
+class _Pool:
+    """Just the surface _forward_unified_kv touches."""
+
+    def __init__(self, fp8):
+        self.unified_swa_window = WIN
+        self.unified_swa_ring_size = RING_STRIDE
+        self.unified_swa_pages = SWA_PAGES
+        if fp8:
+            self.nope = torch.full(
+                (POOL_ROWS, NOPE_ROW_BYTES),
+                NOPE_SENTINEL,
+                dtype=torch.uint8,
+                device=DEVICE,
+            ).view(torch.float8_e4m3fn)
+        else:
+            self.nope = torch.full(
+                (POOL_ROWS, V_HEAD_DIM),
+                ROPE_SENTINEL,
+                dtype=torch.bfloat16,
+                device=DEVICE,
+            )
+        self.rope = torch.full(
+            (POOL_ROWS, ROPE_DIM), ROPE_SENTINEL, dtype=torch.bfloat16, device=DEVICE
+        )
+
+    def get_unified_kv(self, layer_id):
+        return self.nope
+
+    def get_unified_kv_rope(self, layer_id):
+        return self.rope
+
+
+def _chunk(fp8):
+    """This fwd's K, one row per token, every row a different value."""
+    tokens = len(STATE_SLOT)
+    if fp8:
+        rows = torch.arange(1, tokens + 1, dtype=torch.uint8, device=DEVICE)
+        nope = rows[:, None].expand(tokens, NOPE_ROW_BYTES).contiguous()
+        nope = nope.view(torch.float8_e4m3fn)
+    else:
+        rows = torch.arange(1, tokens + 1, dtype=torch.bfloat16, device=DEVICE)
+        nope = rows[:, None].expand(tokens, V_HEAD_DIM).contiguous()
+    rope = (
+        torch.arange(1, tokens + 1, dtype=torch.bfloat16, device=DEVICE)[:, None]
+        .expand(tokens, ROPE_DIM)
+        .contiguous()
+    )
+    return nope, rope
+
+
+class TestUnifiedFp8BackendPrefill(CustomTestCase):
+    def _run(self, fp8=True, save_kv_cache=True):
+        tokens = len(STATE_SLOT)
+        pool = _Pool(fp8)
+        k_nope, k_rope = _chunk(fp8)
+        if fp8:
+            q = torch.zeros(
+                tokens, NUM_HEADS, NOPE_ROW_BYTES, dtype=torch.uint8, device=DEVICE
+            ).view(torch.float8_e4m3fn)
+            q_rope = torch.zeros(
+                tokens, NUM_HEADS, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
+            )
+        else:
+            q = torch.zeros(
+                tokens, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=DEVICE
+            )
+            q_rope, k_rope = None, None
+
+        unified_meta = SimpleNamespace(
+            pf_state_slot=_ints(STATE_SLOT),
+            pf_chunk_start=_ints([0] * tokens),
+            pf_cu_q=_ints(CU_Q),
+            pf_final_pos=_ints([max(POSITIONS)] * tokens),
+        )
+        core_meta = SimpleNamespace(
+            unified=unified_meta,
+            c128_page_indices=None,
+            c4_sparse_page_indices=None,
+        )
+        forward_batch = SimpleNamespace(
+            forward_mode=ForwardMode.EXTEND,
+            positions=torch.tensor(POSITIONS, dtype=torch.int64, device=DEVICE),
+            req_pool_indices=_ints(STATE_SLOT),
+        )
+        fake_self = SimpleNamespace(
+            token_to_kv_pool=pool, softmax_scale=V_HEAD_DIM**-0.5
+        )
+        reader_calls = []
+
+        def _fake_reader(**kwargs):
+            reader_calls.append(kwargs)
+            return torch.zeros(
+                tokens, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=DEVICE
+            )
+
+        target = "prefill_fp8_2buff" if fp8 else "prefill"
+        with (
+            patch.object(runtime, target, _fake_reader),
+            patch.object(
+                backend_mod,
+                "get_parallel",
+                return_value=SimpleNamespace(attn_cp_size=1, attn_cp_rank=0),
+            ),
+        ):
+            backend_mod.DeepseekV4HipRadixBackend._forward_unified_kv(
+                fake_self,
+                q=q,
+                kv=k_nope,
+                layer=SimpleNamespace(layer_id=0, v_head_dim=V_HEAD_DIM),
+                forward_batch=forward_batch,
+                compress_ratio=0,
+                attn_sink=torch.zeros(NUM_HEADS, dtype=torch.float32, device=DEVICE),
+                core_attn_metadata=core_meta,
+                save_kv_cache=save_kv_cache,
+                q_rope=q_rope,
+                k_rope=k_rope,
+            )
+        self.assertEqual(len(reader_calls), 1)
+        return pool, k_nope, k_rope, reader_calls[0]
+
+    def _untouched(self):
+        return sorted(set(range(POOL_ROWS)) - set(EXPECTED_ROWS))
+
+    @_needs_gfx950
+    def test_both_pools_get_this_chunk_on_the_same_ring_row(self):
+        """the regression this file exists for: a rope pool left holding stale rows"""
+        pool, k_nope, k_rope, _ = self._run()
+
+        for token, row in enumerate(EXPECTED_ROWS):
+            self.assertTrue(
+                torch.equal(
+                    pool.nope[row].view(torch.uint8), k_nope[token].view(torch.uint8)
+                ),
+                f"nope pool row {row} does not hold token {token}",
+            )
+            self.assertTrue(
+                torch.equal(pool.rope[row], k_rope[token]),
+                f"rope pool row {row} does not hold token {token} -- "
+                f"got {pool.rope[row][0].item()}, want {k_rope[token][0].item()}",
+            )
+
+    @_needs_gfx950
+    def test_rows_outside_the_window_are_left_alone(self):
+        """both scatters take the same row, so neither may spray past it"""
+        pool, _, _, _ = self._run()
+        rest = self._untouched()
+
+        self.assertTrue(
+            bool((pool.nope[rest].view(torch.uint8) == NOPE_SENTINEL).all())
+        )
+        self.assertTrue(bool((pool.rope[rest] == ROPE_SENTINEL).all()))
+
+    @_needs_gfx950
+    def test_the_reader_gets_the_same_pair_the_ring_write_does(self):
+        _, k_nope, k_rope, call = self._run()
+
+        self.assertIs(call["kv_extend"], k_nope)
+        self.assertIs(call["kv_extend_rope"], k_rope)
+        self.assertIsNotNone(call["unified_kv_rope"])
+
+    @_needs_gfx950
+    def test_nothing_is_written_when_the_model_already_stored(self):
+        pool, _, _, _ = self._run(save_kv_cache=False)
+
+        self.assertTrue(bool((pool.nope.view(torch.uint8) == NOPE_SENTINEL).all()))
+        self.assertTrue(bool((pool.rope == ROPE_SENTINEL).all()))
+
+    @_needs_gfx950
+    def test_the_bf16_arm_never_touches_the_rope_pool(self):
+        """one pool, one write -- the rope pool only exists under the fp8 layout"""
+        pool, k_nope, _, _ = self._run(fp8=False)
+
+        for token, row in enumerate(EXPECTED_ROWS):
+            self.assertTrue(torch.equal(pool.nope[row], k_nope[token]))
+        self.assertTrue(bool((pool.rope == ROPE_SENTINEL).all()))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py

@@ -0,0 +1,332 @@
+"""Two-pool fp8 store tests for the compressor's norm+rope kernel.
+
+Under SGLANG_DSV4_UNIFIED_KV_FP8 the c4/c128 compressor writes its compressed
+latent through ``forward_fp8_2buff``: a 512 B fp8 nope row (448 B payload + 7
+UE8M0 tile scales stored twice) in the unified_kv pool, plus a bf16 rope row in
+the second pool, both at ``out_loc``. These tests pin that layout for both
+compress ratios, against the bf16 store of the same kernel (which shares the
+norm+rope math, so the comparison is byte-exact) and against a torch reference.
+
+Both plans are covered. Most cases run the decode plan; the extend arm (what
+prefill takes) gets the bf16 comparison only, since its plan check and its
+out_loc bound are hand-copied from decode's and nothing else exercises them.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.deepseek_v4_rope import precompute_freqs_cis
+from sglang.kernels.ops.attention.dsv4 import (
+    CompressorDecodePlan,
+    CompressorPrefillPlan,
+    compress_norm_rope_store,
+)
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    DSV4_FP8_NOPE_ROW_BYTES,
+    DSV4_FP8_QUANT_TILE,
+)
+from sglang.srt.utils import is_gfx95_supported
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# the kernel takes E4M3FN vs E4M3FNUZ from the arch and the two-pool layout is only
+# ever allocated on gfx95, so on the default mi300 runner every case here would skip
+register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+HEAD_DIM = 512
+ROPE_DIM = 64
+NOPE_DIM = HEAD_DIM - ROPE_DIM
+NUM_TILES = NOPE_DIM // DSV4_FP8_QUANT_TILE
+SCALE_OFF = NOPE_DIM
+SCALE_BYTES = 2 * NUM_TILES
+
+NUM_TOKENS = 6
+POOL_ROWS = 32
+EPS = 1e-6
+FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
+RATIOS = (4, 128)
+
+
+def _inputs(compress_ratio, seq_lens=None):
+    torch.manual_seed(compress_ratio)
+    kv = torch.randn(NUM_TOKENS, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+    if seq_lens is None:
+        seq_lens = (
+            torch.arange(1, NUM_TOKENS + 1, device=DEVICE, dtype=torch.int64)
+            * compress_ratio
+        )
+    plan = CompressorDecodePlan.generate_legacy(
+        compress_ratio,
+        torch.arange(NUM_TOKENS, device=DEVICE, dtype=torch.int64),
+        seq_lens,
+    )
+    # every other row, so a row that gets written always has an untouched neighbour
+    out_loc = torch.arange(1, 2 * NUM_TOKENS + 1, 2, device=DEVICE, dtype=torch.int64)
+    freqs_cis = precompute_freqs_cis(
+        ROPE_DIM, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
+    ).to(DEVICE)
+    return kv, weight, seq_lens, plan, out_loc, freqs_cis
+
+
+def _extend_inputs(compress_ratio):
+    """one request whose extend spans several compress boundaries"""
+    torch.manual_seed(compress_ratio + 1)
+    total = compress_ratio * NUM_TOKENS
+    seq_lens = torch.tensor([total], dtype=torch.int64)
+    plan = CompressorPrefillPlan.generate_legacy(
+        compress_ratio,
+        torch.zeros(1, dtype=torch.int64, device=DEVICE),
+        seq_lens,
+        seq_lens.clone(),  # the whole sequence is the extend
+        total,
+        DEVICE,
+    )
+    # the kernel binds its token count off the input and then requires the plan to
+    # have that many rows, so the fixture has to follow whatever the planner emitted
+    num_c = plan.plan_c.shape[0]
+    kv = torch.randn(num_c, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+    weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+    # unlike decode, extend indexes out_loc by ragged_id -- one entry per q token, of
+    # which only the compress boundaries are ever read. Sizing this num_c long instead
+    # reads off the end and stores to whatever row index it finds there.
+    written = torch.arange(1, 2 * num_c + 1, 2, device=DEVICE, dtype=torch.int64)
+    out_loc = torch.zeros(total, dtype=torch.int64, device=DEVICE)
+    out_loc[compress_ratio - 1 :: compress_ratio] = written
+    freqs_cis = precompute_freqs_cis(ROPE_DIM, total + 1, 0, 10000, 1, 32, 1).to(DEVICE)
+    return kv, weight, plan, out_loc, written, freqs_cis
+
+
+def _ref_norm_rope(kv, weight, freqs_cis, positions):
+    """rmsnorm over the latent, then rope on the trailing 64, as the kernel does."""
+    x = kv.float()
+    x = x * torch.rsqrt(x.pow(2).sum(-1, keepdim=True) / HEAD_DIM + EPS)
+    x = x * weight.float()
+    nope, pe = x[:, :NOPE_DIM], x[:, NOPE_DIM:]
+
+    freqs = torch.view_as_real(freqs_cis).flatten(-2)[positions]
+    freqs = freqs.reshape(-1, ROPE_DIM // 2, 2).float()
+    pairs = pe.reshape(-1, ROPE_DIM // 2, 2)
+    out = torch.empty_like(pairs)
+    out[..., 0] = pairs[..., 0] * freqs[..., 0] - pairs[..., 1] * freqs[..., 1]
+    out[..., 1] = pairs[..., 0] * freqs[..., 1] + pairs[..., 1] * freqs[..., 0]
+    # the quant warps round through bf16 first, so the scales come off bf16 values
+    return nope.to(torch.bfloat16).float(), out.reshape(-1, ROPE_DIM)
+
+
+def _tile_scale_bytes(nope):
+    """cast_to_ue8m0(max(absmax, 1e-4) / fp8_max) per 1x64 tile."""
+    tiles = nope.reshape(nope.shape[0], NUM_TILES, DSV4_FP8_QUANT_TILE)
+    scale_raw = tiles.abs().amax(-1).clamp_min(1e-4) / FP8_MAX
+    bits = scale_raw.contiguous().view(torch.int32)
+    exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32)
+    return exp.to(torch.uint8)
+
+
+@unittest.skipUnless(is_gfx95_supported(), "needs an AMD gfx95 GPU for e4m3fn")
+class TestUnifiedFp8CompressStore(CustomTestCase):
+    def _store_fp8(self, compress_ratio, *, seq_lens=None, rope_rows=POOL_ROWS):
+        kv, weight, seq_lens, plan, out_loc, freqs_cis = _inputs(
+            compress_ratio, seq_lens
+        )
+        nope_pool = torch.zeros(
+            POOL_ROWS, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
+        )
+        rope_pool = torch.zeros(
+            rope_rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
+        )
+        compress_norm_rope_store(
+            kv.clone(),
+            plan,
+            norm_weight=weight,
+            norm_eps=EPS,
+            freq_cis=freqs_cis,
+            out_loc=out_loc,
+            kvcache=nope_pool.view(torch.uint8),
+            page_size=1,
+            fp8_2buff=True,
+            kvcache_rope=rope_pool.view(torch.uint8),
+        )
+        ref = _ref_norm_rope(kv, weight, freqs_cis, (seq_lens - compress_ratio).long())
+        return nope_pool, rope_pool, out_loc, ref
+
+    def _store_bf16(self, compress_ratio):
+        """same inputs through the bf16 store, i.e. the values before quantization"""
+        kv, weight, _, plan, out_loc, freqs_cis = _inputs(compress_ratio)
+        cache = torch.zeros(POOL_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE)
+        compress_norm_rope_store(
+            kv.clone(),
+            plan,
+            norm_weight=weight,
+            norm_eps=EPS,
+            freq_cis=freqs_cis,
+            out_loc=out_loc,
+            kvcache=cache.view(torch.uint8),
+            page_size=1,
+            bf16_store=True,
+        )
+        return cache[out_loc]
+
+    def _store_extend(self, compress_ratio, *, fp8):
+        kv, weight, plan, out_loc, written, freqs_cis = _extend_inputs(compress_ratio)
+        rows = int(written.max().item()) + 2
+        common = dict(
+            norm_weight=weight,
+            norm_eps=EPS,
+            freq_cis=freqs_cis,
+            out_loc=out_loc,
+            page_size=1,
+        )
+        if not fp8:
+            cache = torch.zeros(rows, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE)
+            compress_norm_rope_store(
+                kv.clone(),
+                plan,
+                kvcache=cache.view(torch.uint8),
+                bf16_store=True,
+                **common,
+            )
+            return cache[written]
+
+        nope_pool = torch.zeros(
+            rows, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
+        )
+        rope_pool = torch.zeros(rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
+        compress_norm_rope_store(
+            kv.clone(),
+            plan,
+            kvcache=nope_pool.view(torch.uint8),
+            fp8_2buff=True,
+            kvcache_rope=rope_pool.view(torch.uint8),
+            **common,
+        )
+        return nope_pool, rope_pool, written
+
+    def test_extend_plan_stores_the_same_rows(self):
+        for ratio in RATIOS:
+            with self.subTest(compress_ratio=ratio):
+                nope_pool, rope_pool, written = self._store_extend(ratio, fp8=True)
+                pre_quant = self._store_extend(ratio, fp8=False)
+
+                nope = pre_quant[:, :NOPE_DIM].float()
+                num_c = nope.shape[0]
+                scale_bytes = _tile_scale_bytes(nope)
+                scale = torch.exp2((scale_bytes.to(torch.int32) - 127).float())
+                want = (
+                    nope.reshape(num_c, NUM_TILES, DSV4_FP8_QUANT_TILE)
+                    / scale[..., None]
+                ).to(torch.float8_e4m3fn)
+
+                self.assertTrue(
+                    torch.equal(
+                        nope_pool[written][:, :NOPE_DIM].view(torch.uint8),
+                        want.view(torch.uint8).reshape(num_c, NOPE_DIM),
+                    )
+                )
+                self.assertTrue(
+                    torch.equal(rope_pool[written], pre_quant[:, NOPE_DIM:])
+                )
+
+    def test_row_matches_the_bf16_store_byte_for_byte(self):
+        for ratio in RATIOS:
+            with self.subTest(compress_ratio=ratio):
+                nope_pool, rope_pool, out_loc, _ = self._store_fp8(ratio)
+                pre_quant = self._store_bf16(ratio)
+
+                nope = pre_quant[:, :NOPE_DIM].float()
+                scale_bytes = _tile_scale_bytes(nope)
+                scale = torch.exp2((scale_bytes.to(torch.int32) - 127).float())
+                want = (
+                    nope.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
+                    / scale[..., None]
+                ).to(torch.float8_e4m3fn)
+
+                # the fixture has to reach the top e4m3 exponent, otherwise it would
+                # not notice a cast that saturates everything above 256
+                self.assertTrue(bool((want.float().abs() >= 256).any()))
+
+                row = nope_pool[out_loc]
+                self.assertTrue(
+                    torch.equal(
+                        row[:, :NOPE_DIM].view(torch.uint8),
+                        want.view(torch.uint8).reshape(NUM_TOKENS, NOPE_DIM),
+                    )
+                )
+                got_scales = row.view(torch.uint8)[
+                    :, SCALE_OFF : SCALE_OFF + SCALE_BYTES
+                ].reshape(NUM_TOKENS, NUM_TILES, 2)
+                self.assertTrue(torch.equal(got_scales[..., 0], scale_bytes))
+                self.assertTrue(torch.equal(got_scales[..., 1], scale_bytes))
+                self.assertTrue(
+                    torch.equal(rope_pool[out_loc], pre_quant[:, NOPE_DIM:])
+                )
+
+    def test_scale_bytes_track_the_torch_reference(self):
+        for ratio in RATIOS:
+            with self.subTest(compress_ratio=ratio):
+                nope_pool, _, out_loc, (ref_nope, _) = self._store_fp8(ratio)
+                got = nope_pool.view(torch.uint8)[
+                    out_loc, SCALE_OFF : SCALE_OFF + SCALE_BYTES
+                ].reshape(NUM_TOKENS, NUM_TILES, 2)
+                self.assertTrue(torch.equal(got[..., 0], _tile_scale_bytes(ref_nope)))
+
+    def test_dequantized_nope_tracks_the_reference(self):
+        for ratio in RATIOS:
+            with self.subTest(compress_ratio=ratio):
+                nope_pool, _, out_loc, (ref_nope, _) = self._store_fp8(ratio)
+                exps = _tile_scale_bytes(ref_nope).to(torch.int32) - 127
+                payload = nope_pool[out_loc, :NOPE_DIM].float()
+                deq = (
+                    payload.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
+                    * torch.exp2(exps.float())[..., None]
+                ).reshape(NUM_TOKENS, NOPE_DIM)
+
+                # e4m3 carries 3 mantissa bits, so half a step is at most ~2^-4 of
+                # the tile's own absmax; beyond that the scale or the payload is off
+                tile_absmax = (
+                    ref_nope.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
+                    .abs()
+                    .amax(-1)
+                    .repeat_interleave(DSV4_FP8_QUANT_TILE, dim=1)
+                )
+                self.assertTrue(torch.all((deq - ref_nope).abs() <= 0.07 * tile_absmax))
+
+    def test_rope_pool_matches_the_bf16_reference(self):
+        for ratio in RATIOS:
+            with self.subTest(compress_ratio=ratio):
+                _, rope_pool, out_loc, (_, ref_pe) = self._store_fp8(ratio)
+                torch.testing.assert_close(
+                    rope_pool[out_loc].float(), ref_pe, rtol=2e-2, atol=2e-2
+                )
+
+    def test_pad_and_neighbour_rows_untouched(self):
+        nope_pool, rope_pool, out_loc, _ = self._store_fp8(4)
+        nope_bytes = nope_pool.view(torch.uint8)
+        self.assertTrue(torch.all(nope_bytes[out_loc, SCALE_OFF + SCALE_BYTES :] == 0))
+
+        untouched = torch.ones(POOL_ROWS, dtype=torch.bool, device=DEVICE)
+        untouched[out_loc] = False
+        self.assertTrue(torch.all(nope_bytes[untouched] == 0))
+        self.assertTrue(torch.all(rope_pool[untouched] == 0))
+
+    def test_non_boundary_decode_is_skipped(self):
+        # only sequences whose length is a multiple of the ratio produce a token
+        seq_lens = torch.full(
+            (NUM_TOKENS,), 4 * 128 + 1, device=DEVICE, dtype=torch.int64
+        )
+        nope_pool, rope_pool, _, _ = self._store_fp8(128, seq_lens=seq_lens)
+        self.assertTrue(torch.all(nope_pool.view(torch.uint8) == 0))
+        self.assertTrue(torch.all(rope_pool == 0))
+
+    def test_short_rope_pool_rejected(self):
+        # one row index addresses both pools, so a short rope pool has to be caught
+        # before either pool is written
+        with self.assertRaises(RuntimeError):
+            self._store_fp8(4, rope_rows=POOL_ROWS // 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py

@@ -0,0 +1,310 @@
+"""Decode read path over the two-pool fp8 unified_kv (aiter's v4 nm asm kernel).
+
+What these pin is the reader-side plumbing, not the kernel's arithmetic: the
+packed 512 B nope row and the bf16 rope pool addressed by one shared row index,
+a per-token ``qo_indptr``, and the ragged ``kv_indptr`` the existing index
+builders emit -- including what they emit for a cuda-graph padded row. The
+reference attends over the *dequantized* pools, so a mismatch is the wiring
+rather than the fp8 round-trip.
+
+The quantization helpers mirror aiter's own reference
+(``op_tests/test_mla_v40_persistent.py``: ``quantize_v4_nope_bpad8`` /
+``pack_v4_nope_scale``). They are duplicated rather than imported because that
+file is a test, not part of the aiter package.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# the asm shader is only shipped for gfx950
+register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
+NOPE_DIM = 448  # fp8 values per row, in elements
+ROPE_DIM = 64
+QUANT_TILE = 64
+NUM_TILES = NOPE_DIM // QUANT_TILE  # 7
+SCALE_OFF = NOPE_DIM  # scales start where the values end
+# latent element count; the same number as NOPE_ROW_BYTES, different unit
+V_HEAD_DIM = NOPE_DIM + ROPE_DIM
+SOFTMAX_SCALE = V_HEAD_DIM**-0.5  # what the kernel hardcodes
+
+_needs_gfx950 = unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "two-pool fp8 decode runs on the gfx950 asm shader",
+)
+
+
+def _pow2_ceil_scale(amax: torch.Tensor) -> torch.Tensor:
+    """amax/fp8_max -> the next power of two at or above it, as fp32"""
+    return torch.pow(2.0, torch.clamp_min(amax, 1e-4).log2().ceil()).to(torch.float32)
+
+
+def _pow2_to_e8m0(pow2: torch.Tensor) -> torch.Tensor:
+    """byte B encodes 2^(B-127); 0 means 0.0 and 255 means inf, so clamp to 254"""
+    biased = torch.log2(pow2).round().to(torch.int32) + 127
+    return torch.clamp(biased, 0, 254).to(torch.uint8)
+
+
+def _e8m0_to_fp32(byte: torch.Tensor) -> torch.Tensor:
+    return torch.exp2((byte.to(torch.int32) - 127).to(torch.float32))
+
+
+def _quantize_nope(nope_fp32: torch.Tensor):
+    """[..., 448] fp32 -> (fp8 values, [..., 7] e8m0 bytes, bf16 round-trip)"""
+    fp8_max = float(torch.finfo(torch.float8_e4m3fn).max)
+    leading = nope_fp32.shape[:-1]
+    tiled = nope_fp32.reshape(*leading, NUM_TILES, QUANT_TILE)
+    scale = _pow2_ceil_scale(tiled.abs().amax(dim=-1) / fp8_max)
+    values = (tiled / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
+    dequant = (values.to(torch.float32) * scale.unsqueeze(-1)).reshape(
+        *leading, NOPE_DIM
+    )
+    return (
+        values.reshape(*leading, NOPE_DIM),
+        _pow2_to_e8m0(scale),
+        dequant.to(torch.bfloat16),
+    )
+
+
+def _pack(values: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor:
+    """448 values + each tile scale twice + pad, as one NOPE_ROW_BYTES fp8 row
+
+    The 50 pad bytes get garbage on purpose. Production allocates Q with
+    nope_pool.new_empty(), so a reader that ever starts looking past the scales
+    should fail here and not in a bf16-vs-fp8 accuracy chase.
+    """
+    leading = values.shape[:-1]
+    row = torch.randint(
+        1, 256, (*leading, NOPE_ROW_BYTES), dtype=torch.uint8, device=values.device
+    )
+    row[..., :NOPE_DIM] = values.view(torch.uint8)
+    dup = scale_e8m0.unsqueeze(-1).expand(*scale_e8m0.shape, 2).reshape(*leading, -1)
+    row[..., SCALE_OFF : SCALE_OFF + 2 * NUM_TILES] = dup
+    return row.view(torch.float8_e4m3fn)
+
+
+def _make_latent(*leading: int):
+    """Return (packed fp8 rows, bf16 rope, bf16 latent the kernel effectively sees)."""
+    nope = torch.randn(*leading, NOPE_DIM, device=DEVICE, dtype=torch.float32)
+    rope = torch.randn(*leading, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+    values, scale, nope_bf16 = _quantize_nope(nope)
+    silver = torch.cat([nope_bf16, rope], dim=-1)
+    return _pack(values, scale).contiguous(), rope.contiguous(), silver
+
+
+def _ragged(lengths, rows, device=DEVICE):
+    """per-token row lists -> (flat int32 indices, int32 indptr)"""
+    indptr = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=device)
+    indptr[1:] = torch.cumsum(
+        torch.tensor(lengths, dtype=torch.int32, device=device), dim=0
+    )
+    flat = torch.cat(
+        [
+            torch.randperm(rows, device=device)[:n].to(torch.int32)
+            for n in lengths
+            if n > 0
+        ]
+        or [torch.empty(0, dtype=torch.int32, device=device)]
+    )
+    return flat.contiguous(), indptr
+
+
+def _reference(q_silver, kv_silver, indices, indptr, sink):
+    """Ragged sparse attention in fp32; V is the full latent, sink has zero V."""
+    T, H, _ = q_silver.shape
+    out = torch.zeros(T, H, V_HEAD_DIM, device=q_silver.device, dtype=torch.float32)
+    q = q_silver.float()
+    sink_f = sink.float()
+    for t in range(T):
+        lo, hi = int(indptr[t]), int(indptr[t + 1])
+        k = kv_silver[indices[lo:hi].long()].float()  # [L, 512]
+        logits = q[t] @ k.transpose(0, 1) * SOFTMAX_SCALE  # [H, L]
+        aug = torch.cat([logits, sink_f.unsqueeze(1)], dim=1)
+        m = aug.amax(dim=1, keepdim=True)
+        p = torch.exp(logits - m)
+        denom = p.sum(dim=1, keepdim=True) + torch.exp(sink_f.unsqueeze(1) - m)
+        out[t] = (p @ k) / denom
+    return out
+
+
+class TestUnifiedFp8Decode(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(7)
+        self.rows = 256
+
+    def _run(self, lengths, num_heads):
+        T = len(lengths)
+        pool_nope, pool_rope, kv_silver = _make_latent(self.rows)
+        q_packed, q_rope, q_silver = _make_latent(T, num_heads)
+        indices, indptr = _ragged(lengths, self.rows)
+        sink = torch.randn(num_heads, device=DEVICE, dtype=torch.float32)
+
+        got = runtime.decode_fp8_2buff(
+            q=q_packed,
+            q_rope=q_rope,
+            unified_kv=pool_nope,
+            unified_kv_rope=pool_rope,
+            kv_indices=indices,
+            kv_indptr=indptr,
+            attn_sink=sink,
+            v_head_dim=V_HEAD_DIM,
+        )
+        want = _reference(q_silver, kv_silver, indices, indptr, sink)
+        return got.float(), want
+
+    def _assert_close(self, got, want, atol=3e-2, rtol=3e-2):
+        """torch-style combined bound.
+
+        A pure relative bound is useless here: the latent's outputs straddle
+        zero, so an absolute error of 4e-3 -- which is what bf16 accumulation
+        costs -- reads as 47% relative on the rows that land near zero.
+        """
+        diff = (got - want).abs()
+        outside = diff > atol + rtol * want.abs()
+        self.assertEqual(
+            outside.sum().item(),
+            0,
+            f"{outside.sum().item()}/{outside.numel()} elements outside "
+            f"{atol}+{rtol}|ref|, max abs {diff.max().item():.4g}",
+        )
+
+    @_needs_gfx950
+    def test_matches_dequantized_reference(self):
+        for lengths in ([64] * 4, [17, 5, 128, 1], [200] * 8):
+            with self.subTest(lengths=lengths):
+                got, want = self._run(lengths, num_heads=16)
+                self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_head_count_64(self):
+        got, want = self._run([48, 96], num_heads=64)
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_cuda_graph_pad_reads_only_the_reserved_ring_row(self):
+        """What the real builder emits for a cuda-graph padded row.
+
+        Not an empty segment: both dsv4 backends fill padded ``seq_lens`` with 1,
+        so ``clamp(seq_lens, max=win)`` leaves the pad one row long. It lands on
+        ring row 0, the slot ReqToTokenPool reserves for exactly this
+        (``free_slots`` starts at 1), so a pad only ever reads and writes there.
+        """
+        win = ring = 64
+        seq_lens = torch.tensor([37, 55, 1, 1], dtype=torch.int32, device=DEVICE)
+        state_slot = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device=DEVICE)
+        n = seq_lens.numel()
+        zero = torch.zeros(n, dtype=torch.int32, device=DEVICE)
+
+        indices, indptr = runtime.build_decode_streams(
+            state_slot=state_slot,
+            positions=seq_lens - 1,  # raw_positions, as the backend derives it
+            swa_len=torch.clamp(seq_lens, max=win),
+            hca_len=zero,
+            csa_len=zero,
+            hca_page_indices=torch.zeros(n, 1, dtype=torch.int32, device=DEVICE),
+            csa_width=1,
+            win=win,
+            ring_stride=ring,
+            swa_pages=self.rows,
+        )[:2]
+
+        seg = (indptr[1 : n + 1] - indptr[:n]).tolist()
+        self.assertEqual(seg, [37, 55, 1, 1])
+        for pad in (2, 3):
+            self.assertEqual(indices[int(indptr[pad])].item(), 0)
+        live = indices[: int(indptr[2])]
+        self.assertGreaterEqual(int(live.min()), ring, "live rows hit slot 0's block")
+
+        pool_nope, pool_rope, _ = _make_latent(self.rows)
+        q_packed, q_rope, _ = _make_latent(n, 16)
+        out = runtime.decode_fp8_2buff(
+            q=q_packed,
+            q_rope=q_rope,
+            unified_kv=pool_nope,
+            unified_kv_rope=pool_rope,
+            kv_indices=indices.contiguous(),
+            kv_indptr=indptr,
+            attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
+            v_head_dim=V_HEAD_DIM,
+        )
+        # the mask never fires here, so what matters is the reserved row keeping
+        # the pad finite rather than it coming back zeroed
+        self.assertTrue(bool(out.isfinite().all()))
+
+    @_needs_gfx950
+    def test_empty_segment_comes_back_nonfinite(self):
+        """Guard for a shape the builders do not reach today.
+
+        Padded seq_lens are always filled with 1 (see
+        test_cuda_graph_pad_reads_only_the_reserved_ring_row), so an empty segment
+        can only come from a builder change -- and it comes back NaN, not zero,
+        since the asm kernel divides by an all-sink denominator.
+        """
+        got, want = self._run([32, 0, 32], num_heads=16)
+        self.assertTrue(bool(torch.isnan(got[1]).any()))
+        for t in (0, 2):
+            self._assert_close(got[t], want[t])
+
+    @_needs_gfx950
+    def test_split_tail_override_matches_reference(self):
+        """past 40 tokens runtime overrides the split count, moving the kernel onto
+        a different stage-2 merge partition -- must still match the reference
+        """
+        lengths = [200, 64] * 24  # 48 tokens, both layer flavours' segment lengths
+        self.assertGreater(len(lengths), 40)
+        got, want = self._run(lengths, num_heads=16)
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_rejects_pool_that_is_not_a_pair(self):
+        pool_nope, pool_rope, _ = _make_latent(self.rows)
+        q_packed, q_rope, _ = _make_latent(2, 16)
+        indices, indptr = _ragged([4, 4], self.rows)
+        sink = torch.zeros(16, device=DEVICE, dtype=torch.float32)
+        with self.assertRaises(AssertionError):
+            runtime.decode_fp8_2buff(
+                q=q_packed,
+                q_rope=q_rope,
+                unified_kv=pool_nope,
+                unified_kv_rope=pool_rope[: self.rows // 2],
+                kv_indices=indices,
+                kv_indptr=indptr,
+                attn_sink=sink,
+                v_head_dim=V_HEAD_DIM,
+            )
+
+    @_needs_gfx950
+    def test_rejects_q_row_wider_than_the_pool_row(self):
+        pool_nope, pool_rope, _ = _make_latent(self.rows)
+        q_packed, q_rope, _ = _make_latent(2, 16)
+        indices, indptr = _ragged([4, 4], self.rows)
+        sink = torch.zeros(16, device=DEVICE, dtype=torch.float32)
+        wider = torch.zeros(
+            2, 16, NOPE_ROW_BYTES + 64, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        wider[..., :NOPE_ROW_BYTES] = q_packed
+        with self.assertRaises(AssertionError):
+            runtime.decode_fp8_2buff(
+                q=wider,
+                q_rope=q_rope,
+                unified_kv=pool_nope,
+                unified_kv_rope=pool_rope,
+                kv_indices=indices,
+                kv_indptr=indptr,
+                attn_sink=sink,
+                v_head_dim=V_HEAD_DIM,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py

@@ -0,0 +1,298 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Prefill read path over the two-pool fp8 unified_kv (aiter's opus kernel).
+
+Two regions per token: the paged prefix pools and this chunk's flat extend pair.
+What these pin is that both regions are addressed with the same row layout and
+that the pair guards fire before the launch -- the reference attends over the
+*dequantized* pools, so a mismatch is the wiring rather than the fp8 round-trip.
+
+The quantization helpers are the ones from the decode test rather than a shared
+module: files under test/registered/ are collected standalone (no __init__.py,
+no conftest), so importing across them breaks in CI.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
+from sglang.srt.utils import is_gfx95_supported, is_hip
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
+NOPE_DIM = 448  # fp8 values per row, in elements
+ROPE_DIM = 64
+QUANT_TILE = 64
+NUM_TILES = NOPE_DIM // QUANT_TILE  # 7
+SCALE_OFF = NOPE_DIM
+# latent element count; the same number as NOPE_ROW_BYTES, different unit
+V_HEAD_DIM = NOPE_DIM + ROPE_DIM
+SOFTMAX_SCALE = V_HEAD_DIM**-0.5
+
+_needs_gfx950 = unittest.skipUnless(
+    torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
+    "two-pool fp8 prefill runs on the gfx950 opus kernel",
+)
+
+
+def _pow2_ceil_scale(amax: torch.Tensor) -> torch.Tensor:
+    return torch.pow(2.0, torch.clamp_min(amax, 1e-4).log2().ceil()).to(torch.float32)
+
+
+def _pow2_to_e8m0(pow2: torch.Tensor) -> torch.Tensor:
+    biased = torch.log2(pow2).round().to(torch.int32) + 127
+    return torch.clamp(biased, 0, 254).to(torch.uint8)
+
+
+def _quantize_nope(nope_fp32: torch.Tensor):
+    """[..., 448] fp32 -> (fp8 values, [..., 7] e8m0 bytes, bf16 round-trip)"""
+    fp8_max = float(torch.finfo(torch.float8_e4m3fn).max)
+    leading = nope_fp32.shape[:-1]
+    tiled = nope_fp32.reshape(*leading, NUM_TILES, QUANT_TILE)
+    scale = _pow2_ceil_scale(tiled.abs().amax(dim=-1) / fp8_max)
+    values = (tiled / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
+    dequant = (values.to(torch.float32) * scale.unsqueeze(-1)).reshape(
+        *leading, NOPE_DIM
+    )
+    return (
+        values.reshape(*leading, NOPE_DIM),
+        _pow2_to_e8m0(scale),
+        dequant.to(torch.bfloat16),
+    )
+
+
+def _pack(values: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor:
+    """448 values + each tile scale twice + pad, as one NOPE_ROW_BYTES fp8 row
+
+    Pad bytes get garbage on purpose, same reason as the decode test: production
+    allocates these with new_empty(), so a reader that walks past the scales
+    should fail here rather than as an accuracy drift.
+    """
+    leading = values.shape[:-1]
+    row = torch.randint(
+        1, 256, (*leading, NOPE_ROW_BYTES), dtype=torch.uint8, device=values.device
+    )
+    row[..., :NOPE_DIM] = values.view(torch.uint8)
+    dup = scale_e8m0.unsqueeze(-1).expand(*scale_e8m0.shape, 2).reshape(*leading, -1)
+    row[..., SCALE_OFF : SCALE_OFF + 2 * NUM_TILES] = dup
+    return row.view(torch.float8_e4m3fn)
+
+
+def _make_latent(*leading: int):
+    """Return (packed fp8 rows, bf16 rope, bf16 latent the kernel effectively sees)."""
+    nope = torch.randn(*leading, NOPE_DIM, device=DEVICE, dtype=torch.float32)
+    rope = torch.randn(*leading, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+    values, scale, nope_bf16 = _quantize_nope(nope)
+    silver = torch.cat([nope_bf16, rope], dim=-1)
+    return _pack(values, scale).contiguous(), rope.contiguous(), silver
+
+
+def _ragged(lengths, rows):
+    """per-token row lists -> (flat int32 indices, int32 indptr)"""
+    indptr = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEVICE)
+    indptr[1:] = torch.cumsum(
+        torch.tensor(lengths, dtype=torch.int32, device=DEVICE), dim=0
+    )
+    parts = [
+        torch.randperm(rows, device=DEVICE)[:n].to(torch.int32)
+        for n in lengths
+        if n > 0
+    ]
+    flat = (
+        torch.cat(parts) if parts else torch.empty(0, dtype=torch.int32, device=DEVICE)
+    )
+    return flat.contiguous(), indptr
+
+
+def _reference(q_silver, sources, sink, scale):
+    """Ragged two-region attention in fp32; V is the full latent, sink V is zero.
+
+    ``sources`` is [(silver, indices, indptr), ...]. The kernel shares one online
+    softmax across the regions, so order does not matter and this just
+    concatenates whatever each region selected.
+    """
+    T, H, _ = q_silver.shape
+    out = torch.zeros(T, H, V_HEAD_DIM, device=q_silver.device, dtype=torch.float32)
+    q = q_silver.float()
+    sink_f = sink.float()
+    for t in range(T):
+        keys = []
+        for silver, indices, indptr in sources:
+            lo, hi = int(indptr[t]), int(indptr[t + 1])
+            if hi > lo:
+                keys.append(silver[indices[lo:hi].long()].float())
+        if not keys:
+            # only the sink is left: it contributes to the denominator and has
+            # V = 0, so the row is exactly zero
+            continue
+        k = torch.cat(keys, dim=0)
+        logits = q[t] @ k.transpose(0, 1) * scale
+        m = torch.cat([logits, sink_f.unsqueeze(1)], dim=1).amax(dim=1, keepdim=True)
+        p = torch.exp(logits - m)
+        denom = p.sum(dim=1, keepdim=True) + torch.exp(sink_f.unsqueeze(1) - m)
+        out[t] = (p @ k) / denom
+    return out
+
+
+class TestUnifiedFp8Prefill(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(11)
+        self.rows = 256
+
+    def _run(self, prefix_lens, extend_lens, num_heads=16, scale=SOFTMAX_SCALE):
+        T = len(prefix_lens)
+        self.assertEqual(T, len(extend_lens))
+        extend_rows = max(max(extend_lens), 1)
+        pool_nope, pool_rope, pool_silver = _make_latent(self.rows)
+        ext_nope, ext_rope, ext_silver = _make_latent(extend_rows)
+        q_packed, q_rope, q_silver = _make_latent(T, num_heads)
+        pre_i, pre_p = _ragged(prefix_lens, self.rows)
+        ext_i, ext_p = _ragged(extend_lens, extend_rows)
+        sink = torch.randn(num_heads, device=DEVICE, dtype=torch.float32)
+
+        got = runtime.prefill_fp8_2buff(
+            q=q_packed,
+            q_rope=q_rope,
+            unified_kv=pool_nope,
+            unified_kv_rope=pool_rope,
+            kv_indices_prefix=pre_i,
+            kv_indptr_prefix=pre_p,
+            kv_extend=ext_nope,
+            kv_extend_rope=ext_rope,
+            kv_indices_extend=ext_i,
+            kv_indptr_extend=ext_p,
+            attn_sink=sink,
+            softmax_scale=scale,
+            v_head_dim=V_HEAD_DIM,
+        )
+        want = _reference(
+            q_silver,
+            [(pool_silver, pre_i, pre_p), (ext_silver, ext_i, ext_p)],
+            sink,
+            scale,
+        )
+        return got.float(), want
+
+    def _assert_close(self, got, want, atol=3e-2, rtol=3e-2):
+        """torch-style combined bound, same reasoning as the decode test.
+
+        A pure relative bound is useless here: the latent's outputs straddle
+        zero, so the absolute error bf16 accumulation costs reads as a huge
+        relative one on the rows that land near zero.
+        """
+        diff = (got - want).abs()
+        outside = diff > atol + rtol * want.abs()
+        self.assertEqual(
+            outside.sum().item(),
+            0,
+            f"{outside.sum().item()}/{outside.numel()} elements outside "
+            f"the bound, max abs {diff.max().item():.4g}",
+        )
+
+    @_needs_gfx950
+    def test_matches_dequantized_reference(self):
+        cases = (
+            ([64, 64, 64, 64], [1, 2, 3, 4]),
+            ([17, 5, 128, 1], [4, 4, 4, 4]),
+            ([200] * 6, [1, 3, 6, 2, 5, 4]),
+        )
+        for prefix_lens, extend_lens in cases:
+            with self.subTest(prefix=prefix_lens, extend=extend_lens):
+                got, want = self._run(prefix_lens, extend_lens)
+                self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_first_chunk_has_an_empty_prefix_for_every_token(self):
+        """the real shape of chunk 0: nothing committed yet, extend is all there is"""
+        got, want = self._run([0, 0, 0, 0], [1, 2, 3, 4])
+        self.assertTrue(bool(got.isfinite().all()))
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_a_token_with_neither_region_comes_back_zero(self):
+        """Not NaN, which is where this differs from the asm decode reader.
+
+        decode_fp8_2buff has to mask that case itself; this kernel already
+        returns zeros, so there is deliberately no mask on this path. The
+        reference skips those rows for the same reason: with only the sink left
+        the numerator is zero.
+        """
+        got, want = self._run([64, 0, 64], [2, 0, 2])
+        self.assertTrue(torch.equal(got[1], torch.zeros_like(got[1])))
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_head_count_64(self):
+        got, want = self._run([48, 96], [3, 5], num_heads=64)
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_scale_is_passed_through(self):
+        """unlike the decode reader, this kernel takes the scale as an argument"""
+        got, want = self._run([32, 32], [2, 2], scale=0.5 * SOFTMAX_SCALE)
+        self._assert_close(got, want)
+
+    @_needs_gfx950
+    def test_extend_row_narrower_than_the_pool_is_rejected(self):
+        """the two regions are walked with one row layout, so a short row would
+        read the next token's bytes as this one's scales"""
+        pool_nope, pool_rope, _ = _make_latent(self.rows)
+        q_packed, q_rope, _ = _make_latent(2, 16)
+        ext_nope, ext_rope, _ = _make_latent(4)
+        pre_i, pre_p = _ragged([8, 8], self.rows)
+        ext_i, ext_p = _ragged([1, 1], 4)
+        with self.assertRaisesRegex(AssertionError, "extend nope row"):
+            runtime.prefill_fp8_2buff(
+                q=q_packed,
+                q_rope=q_rope,
+                unified_kv=pool_nope,
+                unified_kv_rope=pool_rope,
+                kv_indices_prefix=pre_i,
+                kv_indptr_prefix=pre_p,
+                kv_extend=ext_nope[:, : NOPE_ROW_BYTES // 2].contiguous(),
+                kv_extend_rope=ext_rope,
+                kv_indices_extend=ext_i,
+                kv_indptr_extend=ext_p,
+                attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
+                softmax_scale=SOFTMAX_SCALE,
+                v_head_dim=V_HEAD_DIM,
+            )
+
+    @_needs_gfx950
+    def test_mismatched_pools_are_rejected_before_the_launch(self):
+        pool_nope, _, _ = _make_latent(self.rows)
+        short_rope = torch.zeros(
+            self.rows // 2, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
+        )
+        q_packed, q_rope, _ = _make_latent(2, 16)
+        ext_nope, ext_rope, _ = _make_latent(4)
+        pre_i, pre_p = _ragged([8, 8], self.rows)
+        ext_i, ext_p = _ragged([1, 1], 4)
+        with self.assertRaisesRegex(AssertionError, "pool rows differ"):
+            runtime.prefill_fp8_2buff(
+                q=q_packed,
+                q_rope=q_rope,
+                unified_kv=pool_nope,
+                unified_kv_rope=short_rope,
+                kv_indices_prefix=pre_i,
+                kv_indptr_prefix=pre_p,
+                kv_extend=ext_nope,
+                kv_extend_rope=ext_rope,
+                kv_indices_extend=ext_i,
+                kv_indptr_extend=ext_p,
+                attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
+                softmax_scale=SOFTMAX_SCALE,
+                v_head_dim=V_HEAD_DIM,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py

@@ -0,0 +1,539 @@
+"""Two-pool fp8 store tests for the fused QK norm+RoPE kernel wrapper.
+
+Under SGLANG_DSV4_UNIFIED_KV_FP8 ``fused_qk_norm_rope_swa_store`` delegates to
+aiter, which packs K into a 512 B fp8 nope row (448 B payload + 14 B duplicated
+E8M0 tile scales) plus a bf16 rope row and scatters both into the SWA ring.
+These tests pin that layout, which the decode reader depends on, and which of the
+two forms Q comes back in: the same packed pair when the caller supplies a rope
+buffer (what the v4 nm asm reader takes), plain rotated bf16 when it does not
+(what the Triton reader takes).
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.fused_qk_norm_rope_store import (
+    _HAS_GROUP_QUANT,
+    fused_qk_norm_rope_swa_store,
+)
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    DSV4_FP8_NOPE_ROW_BYTES,
+    DSV4_FP8_QUANT_TILE,
+)
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# aiter's group-quant path is gfx95-only, so on the default mi300 runner every case
+# here would skip
+register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+NOPE_DIM = 448
+ROPE_DIM = 64
+HEAD_DIM = NOPE_DIM + ROPE_DIM
+NUM_TILES = NOPE_DIM // DSV4_FP8_QUANT_TILE
+SCALE_OFF = NOPE_DIM
+SCALE_BYTES = 2 * NUM_TILES
+
+NUM_HEADS = 4
+EPS = 1e-6
+MAX_POS = 256
+RING_STRIDE = 16
+
+
+def _cos_sin():
+    inv = 1.0 / (
+        10000 ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32) / ROPE_DIM)
+    )
+    ang = torch.arange(MAX_POS, dtype=torch.float32)[:, None] * inv[None, :]
+    return (
+        ang.cos().to(torch.bfloat16).to(DEVICE),
+        ang.sin().to(torch.bfloat16).to(DEVICE),
+    )
+
+
+def _ref_norm_rope(kv, weight, cos, sin, positions):
+    """rmsnorm over the whole latent, then GPT-J rope on the trailing pe half."""
+    x = kv.float()
+    scale = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)
+    normed = x * scale * weight.float()
+    nope, pe = normed[:, :NOPE_DIM], normed[:, NOPE_DIM:]
+    c = cos.float()[positions]
+    s = sin.float()[positions]
+    even, odd = pe[:, 0::2], pe[:, 1::2]
+    out = torch.empty_like(pe)
+    out[:, 0::2] = even * c - odd * s
+    out[:, 1::2] = odd * c + even * s
+    return nope, out
+
+
+def _ref_tile_scales(nope):
+    """e8m0 exponent byte per 1x64 tile, from the fp32 reference nope."""
+    tiles = nope.reshape(nope.shape[0], NUM_TILES, DSV4_FP8_QUANT_TILE)
+    absmax = tiles.abs().amax(-1).clamp_min(1e-8)
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    return torch.ceil(torch.log2(absmax / fp8_max))
+
+
+def _pools(n_rows):
+    nope_pool = torch.zeros(
+        n_rows, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
+    )
+    rope_pool = torch.zeros(n_rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
+    return nope_pool, rope_pool
+
+
+class _StoreCase(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(7)
+        self.T = 6
+        self.cos, self.sin = _cos_sin()
+        self.weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        self.kv = torch.randn(self.T, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        self.q = torch.randn(
+            self.T, NUM_HEADS * HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
+        )
+        self.positions = torch.arange(self.T, device=DEVICE, dtype=torch.int64)
+        # distinct ring rows so each row has one unambiguous writer
+        self.swa_loc = (
+            self.positions.to(torch.int32) % RING_STRIDE + RING_STRIDE
+        ).contiguous()
+
+
+@unittest.skipUnless(_HAS_GROUP_QUANT, "needs aiter's group-quant kernel on gfx95x")
+class TestUnifiedFp8QkNormRope(_StoreCase):
+    def _call(
+        self,
+        nope_pool=None,
+        rope_pool=None,
+        k_nope=None,
+        k_rope=None,
+        q_out=None,
+        q_rope_out=None,
+    ):
+        return fused_qk_norm_rope_swa_store(
+            q=self.q,
+            kv=self.kv,
+            q_norm_weight=None,
+            kv_norm_weight=self.weight,
+            q_rms_eps=EPS,
+            kv_rms_eps=EPS,
+            rope_head_dim=ROPE_DIM,
+            cos_cache=self.cos,
+            sin_cache=self.sin,
+            positions=self.positions,
+            swa_cache=nope_pool,
+            swa_loc=None if nope_pool is None else self.swa_loc,
+            swa_page_size=1,
+            dtype=torch.bfloat16,
+            fp8_2buff=True,
+            swa_rope_cache=rope_pool,
+            k_nope_out=k_nope,
+            k_rope_out=k_rope,
+            q_out=q_out,
+            q_rope_out=q_rope_out,
+        )
+
+    def test_pool_rows_equal_the_dense_packed_output(self):
+        """the ring write and the dense K buffers come from the same values"""
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        k_nope = torch.empty(
+            self.T, 1, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
+        )
+        k_rope = torch.empty(self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
+        self._call(nope_pool, rope_pool, k_nope, k_rope)
+
+        rows = self.swa_loc.long()
+        pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
+        dense_bytes = k_nope.view(torch.uint8)[:, 0, : SCALE_OFF + SCALE_BYTES]
+        self.assertTrue(torch.equal(pool_bytes, dense_bytes))
+        self.assertTrue(torch.equal(rope_pool[rows], k_rope[:, 0]))
+
+    def test_strided_kv_slice_matches_contiguous(self):
+        """kv is a strided slice of qkv_a; aiter forwards kv.stride(0), so going
+        back to assuming a packed row would corrupt silently instead of erroring
+
+        Covers both callers: the fused ring write (pools) and the caller-buffer
+        pair (k_nope/k_rope) that prefill and target-verify pass instead.
+        """
+        q_lora_rank = 1536  # DSV4-Pro; only its being != 0 matters here
+        wide = torch.randn(
+            self.T, q_lora_rank + HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
+        )
+        strided = wide[..., q_lora_rank:]
+        self.assertFalse(strided.is_contiguous())
+        self.assertEqual(strided.stride(-1), 1)
+
+        runs = []
+        for kv in (strided, strided.contiguous()):
+            self.kv = kv
+            nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+            k_nope = torch.zeros(
+                self.T,
+                1,
+                DSV4_FP8_NOPE_ROW_BYTES,
+                dtype=torch.float8_e4m3fn,
+                device=DEVICE,
+            )
+            k_rope = torch.zeros(
+                self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
+            )
+            packed = self._call(nope_pool, rope_pool, k_nope, k_rope)
+            runs.append((nope_pool, rope_pool, k_nope, k_rope, packed))
+
+        for got, want in zip(*runs):
+            self.assertTrue(torch.equal(got.view(torch.uint8), want.view(torch.uint8)))
+
+    def test_scale_bytes_are_duplicated_e8m0(self):
+        """the asm reader reads each tile scale twice, so the 14 B must be 7 equal pairs"""
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        self._call(nope_pool, rope_pool)
+
+        rows = self.swa_loc.long()
+        scales = nope_pool.view(torch.uint8)[
+            rows, SCALE_OFF : SCALE_OFF + SCALE_BYTES
+        ].reshape(self.T, NUM_TILES, 2)
+        self.assertTrue(torch.equal(scales[..., 0], scales[..., 1]))
+
+        ref_nope, _ = _ref_norm_rope(
+            self.kv, self.weight, self.cos, self.sin, self.positions
+        )
+        expected = (_ref_tile_scales(ref_nope) + 127).to(torch.uint8)
+        self.assertTrue(torch.equal(scales[..., 0], expected))
+
+    def test_dequantized_nope_tracks_the_reference(self):
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        self._call(nope_pool, rope_pool)
+
+        ref_nope, _ = _ref_norm_rope(
+            self.kv, self.weight, self.cos, self.sin, self.positions
+        )
+        exps = _ref_tile_scales(ref_nope)
+        payload = nope_pool[self.swa_loc.long(), :NOPE_DIM].float()
+        deq = (
+            payload.reshape(self.T, NUM_TILES, DSV4_FP8_QUANT_TILE)
+            * torch.exp2(exps)[..., None]
+        ).reshape(self.T, NOPE_DIM)
+
+        # e4m3 carries 3 mantissa bits, so the worst case is ~2^-4 of the tile's
+        # own absmax. Anything beyond that means the scale or the payload is off,
+        # not rounding.
+        tile_absmax = (
+            ref_nope.reshape(self.T, NUM_TILES, DSV4_FP8_QUANT_TILE)
+            .abs()
+            .amax(-1)
+            .repeat_interleave(DSV4_FP8_QUANT_TILE, dim=1)
+        )
+        self.assertTrue(torch.all((deq - ref_nope).abs() <= 0.07 * tile_absmax))
+
+    def test_rope_pool_matches_the_bf16_reference(self):
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        self._call(nope_pool, rope_pool)
+
+        _, ref_pe = _ref_norm_rope(
+            self.kv, self.weight, self.cos, self.sin, self.positions
+        )
+        got = rope_pool[self.swa_loc.long()].float()
+        torch.testing.assert_close(got, ref_pe, rtol=2e-2, atol=2e-2)
+
+    def test_q_stays_bf16_and_rotated(self):
+        q_out = self._call()
+        self.assertEqual(q_out.dtype, torch.bfloat16)
+        self.assertEqual(tuple(q_out.shape), (self.T, NUM_HEADS, HEAD_DIM))
+
+        head = self.q.view(self.T, NUM_HEADS, HEAD_DIM)[:, 0]
+        ones = torch.ones(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        ref_nope, ref_pe = _ref_norm_rope(
+            head, ones, self.cos, self.sin, self.positions
+        )
+        got = q_out[:, 0].float()
+        torch.testing.assert_close(got[:, :NOPE_DIM], ref_nope, rtol=2e-2, atol=2e-2)
+        torch.testing.assert_close(got[:, NOPE_DIM:], ref_pe, rtol=2e-2, atol=2e-2)
+
+    def test_q_is_packed_like_k_when_a_rope_buffer_is_given(self):
+        """the v4 nm asm reader wants Q in the same 512 B form as the pool rows
+
+        Pinned against the bf16 Q the same call produces without the rope buffer,
+        so this is the quantization of a known-good rotated Q rather than a
+        second reimplementation of norm+rope.
+        """
+        q_packed = torch.empty(
+            self.T,
+            NUM_HEADS,
+            DSV4_FP8_NOPE_ROW_BYTES,
+            dtype=torch.float8_e4m3fn,
+            device=DEVICE,
+        )
+        q_rope = torch.empty(
+            self.T, NUM_HEADS, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
+        )
+        got = self._call(q_out=q_packed, q_rope_out=q_rope)
+        self.assertIs(got, q_packed)
+
+        ref = self._call().float()  # bf16 Q, same input
+        ref_nope, ref_pe = ref[..., :NOPE_DIM], ref[..., NOPE_DIM:]
+
+        raw = q_packed.view(torch.uint8)
+        exp = _ref_tile_scales(ref_nope.reshape(-1, NOPE_DIM)).reshape(
+            self.T, NUM_HEADS, NUM_TILES
+        )
+        scale_bytes = raw[..., SCALE_OFF : SCALE_OFF + SCALE_BYTES].reshape(
+            self.T, NUM_HEADS, NUM_TILES, 2
+        )
+        torch.testing.assert_close(
+            scale_bytes[..., 0].int() - 127, exp.int(), rtol=0, atol=0
+        )
+        self.assertTrue(torch.equal(scale_bytes[..., 0], scale_bytes[..., 1]))
+
+        scale = torch.exp2(scale_bytes[..., 0].float() - 127)
+        dq = (
+            raw[..., :NOPE_DIM]
+            .view(torch.float8_e4m3fn)
+            .float()
+            .reshape(self.T, NUM_HEADS, NUM_TILES, DSV4_FP8_QUANT_TILE)
+            * scale.unsqueeze(-1)
+        ).reshape(self.T, NUM_HEADS, NOPE_DIM)
+        # atol is half an fp8 step at the *tile's* absmax, not a per-element
+        # relative error -- a small value sharing a tile with a large one carries
+        # the large one's step. Measured 0.125 worst case at absmax 3.9, and the
+        # parts that must be exact (scale bytes above, rope below) are pinned as
+        # such.
+        torch.testing.assert_close(dq, ref_nope, rtol=5e-2, atol=5e-2)
+        torch.testing.assert_close(q_rope.float(), ref_pe, rtol=0, atol=0)
+
+    def test_fp8_q_without_a_rope_buffer_is_rejected(self):
+        q_packed = torch.empty(
+            self.T,
+            NUM_HEADS,
+            DSV4_FP8_NOPE_ROW_BYTES,
+            dtype=torch.float8_e4m3fn,
+            device=DEVICE,
+        )
+        with self.assertRaises(AssertionError):
+            self._call(q_out=q_packed)
+
+    def test_strided_q_out_is_filled_without_touching_the_padding(self):
+        """attn_tp_size > 1 hands us a slice of a head-padded [T, 64, D] buffer
+
+        The zero-init is this test's way of seeing whether the staging copy strays
+        outside the slice. gfx950 allocates that buffer with new_empty, so in
+        production the padding holds garbage, not zeros -- what matters is only that
+        nobody writes it.
+        """
+        padded = torch.zeros(self.T, 64, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        q_out = padded[:, :NUM_HEADS, :]
+        self.assertFalse(q_out.is_contiguous())
+        packed = self._call()
+
+        got = fused_qk_norm_rope_swa_store(
+            q=self.q,
+            kv=self.kv,
+            q_norm_weight=None,
+            kv_norm_weight=self.weight,
+            q_rms_eps=EPS,
+            kv_rms_eps=EPS,
+            rope_head_dim=ROPE_DIM,
+            cos_cache=self.cos,
+            sin_cache=self.sin,
+            positions=self.positions,
+            q_out=q_out,
+            dtype=torch.bfloat16,
+            fp8_2buff=True,
+        )
+        self.assertIs(got, q_out)
+        self.assertTrue(torch.all(padded[:, NUM_HEADS:, :] == 0))
+        # staging must not reorder the heads, so the strided destination has to
+        # hold exactly what the contiguous call produced
+        self.assertTrue(torch.equal(q_out, packed))
+
+    def test_negative_position_skips_both_pools(self):
+        """a stale/pad token must leave both pools alone, not half-write a row"""
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        self.positions[2] = -1
+        self._call(nope_pool, rope_pool)
+
+        row = self.swa_loc[2].item()
+        self.assertEqual(nope_pool.view(torch.uint8)[row].max().item(), 0)
+        self.assertEqual(rope_pool[row].abs().max().item(), 0)
+
+    def test_rope_pool_is_required_with_the_nope_pool(self):
+        nope_pool, _ = _pools(2 * RING_STRIDE)
+        with self.assertRaises(AssertionError):
+            self._call(nope_pool, None)
+
+    def test_mismatched_pools_are_rejected_before_the_launch(self):
+        """aiter aborts the process on a short pool, so these must fail in python"""
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        short_rope = rope_pool[:RING_STRIDE].contiguous()
+        cases = {
+            "fewer rope rows": (nope_pool, short_rope),
+            "rope dtype": (nope_pool, rope_pool.to(torch.float16)),
+            "rope width": (nope_pool, rope_pool[:, : ROPE_DIM // 2].contiguous()),
+            "nope row bytes": (nope_pool[:, :NOPE_DIM].contiguous(), rope_pool),
+        }
+        for name, (nope, rope) in cases.items():
+            with self.subTest(name), self.assertRaises(AssertionError):
+                self._call(nope, rope)
+
+    def test_bf16_store_is_a_different_store(self):
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        with self.assertRaises(AssertionError):
+            fused_qk_norm_rope_swa_store(
+                q=self.q,
+                kv=self.kv,
+                q_norm_weight=None,
+                kv_norm_weight=self.weight,
+                q_rms_eps=EPS,
+                kv_rms_eps=EPS,
+                rope_head_dim=ROPE_DIM,
+                cos_cache=self.cos,
+                sin_cache=self.sin,
+                positions=self.positions,
+                swa_cache=nope_pool,
+                swa_loc=self.swa_loc,
+                swa_page_size=1,
+                dtype=torch.bfloat16,
+                bf16_store=True,
+                fp8_2buff=True,
+                swa_rope_cache=rope_pool,
+            )
+
+
+class TestBf16StoreStillWorks(_StoreCase):
+    """fp8_2buff returns before the Triton kernel, so pin the branch it skips"""
+
+    def test_bf16_store_writes_the_whole_row(self):
+        pool = torch.zeros(
+            2 * RING_STRIDE, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
+        )
+        ref_nope, ref_pe = _ref_norm_rope(
+            self.kv, self.weight, self.cos, self.sin, self.positions
+        )
+        q_out = fused_qk_norm_rope_swa_store(
+            q=self.q,
+            kv=self.kv,
+            q_norm_weight=None,
+            kv_norm_weight=self.weight,
+            q_rms_eps=EPS,
+            kv_rms_eps=EPS,
+            rope_head_dim=ROPE_DIM,
+            cos_cache=self.cos,
+            sin_cache=self.sin,
+            positions=self.positions,
+            swa_cache=pool,
+            swa_loc=self.swa_loc,
+            swa_page_size=1,
+            dtype=torch.bfloat16,
+            bf16_store=True,
+        )
+        self.assertEqual(q_out.dtype, torch.bfloat16)
+        self.assertEqual(tuple(q_out.shape), (self.T, NUM_HEADS, HEAD_DIM))
+
+        rows = self.swa_loc.long()
+        got = pool[rows].float()
+        torch.testing.assert_close(got[:, :NOPE_DIM], ref_nope, rtol=2e-2, atol=2e-2)
+        torch.testing.assert_close(got[:, NOPE_DIM:], ref_pe, rtol=2e-2, atol=2e-2)
+
+        untouched = torch.ones(pool.shape[0], dtype=torch.bool, device=DEVICE)
+        untouched[rows] = False
+        self.assertEqual(pool[untouched].abs().max().item(), 0)
+
+
+@unittest.skipUnless(_HAS_GROUP_QUANT, "needs aiter's group-quant kernel on gfx95x")
+class TestUnifiedFp8SwaRingWrap(CustomTestCase):
+    """What the ring holds once a slot gets written a second time.
+
+    The two pools have to turn over together. A row whose nope came from the new
+    token but whose rope is still the old one decodes against the wrong angle,
+    and nothing downstream can notice -- both halves are individually
+    well-formed.
+
+    Wrap is driven across calls, not inside one. Within a launch the
+    out-of-window tokens carry loc -1 and get skipped, so every live row has a
+    single writer; two writers to one row in one launch would be a race with no
+    defined winner to assert on.
+    """
+
+    def setUp(self):
+        torch.manual_seed(11)
+        self.T = 6
+        self.cos, self.sin = _cos_sin()
+        self.weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        self.q = torch.randn(
+            self.T, NUM_HEADS * HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
+        )
+
+    def _store(self, kv, positions, swa_loc, nope_pool, rope_pool):
+        """one launch; hands back the dense K pair as the per-token truth"""
+        k_nope = torch.empty(
+            self.T, 1, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
+        )
+        k_rope = torch.empty(self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
+        fused_qk_norm_rope_swa_store(
+            q=self.q,
+            kv=kv,
+            q_norm_weight=None,
+            kv_norm_weight=self.weight,
+            q_rms_eps=EPS,
+            kv_rms_eps=EPS,
+            rope_head_dim=ROPE_DIM,
+            cos_cache=self.cos,
+            sin_cache=self.sin,
+            positions=positions,
+            swa_cache=nope_pool,
+            swa_loc=swa_loc,
+            swa_page_size=1,
+            dtype=torch.bfloat16,
+            fp8_2buff=True,
+            swa_rope_cache=rope_pool,
+            k_nope_out=k_nope,
+            k_rope_out=k_rope,
+        )
+        return k_nope.view(torch.uint8)[:, 0, : SCALE_OFF + SCALE_BYTES], k_rope[:, 0]
+
+    def _pass(self, step, nope_pool, rope_pool, count=None):
+        """step 0 fills the ring, step 1 comes back around onto the same slots"""
+        count = self.T if count is None else count
+        kv = torch.randn(self.T, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
+        positions = (
+            torch.arange(self.T, device=DEVICE, dtype=torch.int64) + step * RING_STRIDE
+        )
+        swa_loc = (positions.to(torch.int32) % RING_STRIDE + RING_STRIDE).contiguous()
+        # tokens past `count` fall out of window on this pass, like a short step
+        if count < self.T:
+            positions = positions.clone()
+            positions[count:] = -1
+        nope, rope = self._store(kv, positions, swa_loc, nope_pool, rope_pool)
+        return swa_loc.long(), nope.clone(), rope.clone()
+
+    def test_wrap_turns_over_both_pools(self):
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        rows, old_nope, _ = self._pass(0, nope_pool, rope_pool)
+        rows2, new_nope, new_rope = self._pass(1, nope_pool, rope_pool)
+        self.assertTrue(torch.equal(rows, rows2), "the wrap must reuse the same slots")
+        # only meaningful if pass 1 actually changed the bytes
+        self.assertFalse(torch.equal(old_nope, new_nope))
+
+        pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
+        self.assertTrue(torch.equal(pool_bytes, new_nope))
+        self.assertTrue(torch.equal(rope_pool[rows], new_rope))
+
+    def test_a_slot_the_wrap_skipped_keeps_its_old_pair(self):
+        """a short second pass must leave the rows it didn't address alone"""
+        keep = 2
+        nope_pool, rope_pool = _pools(2 * RING_STRIDE)
+        rows, old_nope, old_rope = self._pass(0, nope_pool, rope_pool)
+        _, new_nope, new_rope = self._pass(1, nope_pool, rope_pool, count=keep)
+
+        pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
+        got_rope = rope_pool[rows]
+        self.assertTrue(torch.equal(pool_bytes[:keep], new_nope[:keep]))
+        self.assertTrue(torch.equal(got_rope[:keep], new_rope[:keep]))
+        self.assertTrue(torch.equal(pool_bytes[keep:], old_nope[keep:]))
+        self.assertTrue(torch.equal(got_rope[keep:], old_rope[keep:]))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py

@@ -0,0 +1,267 @@
+"""SWA ring scatter tests for the two-pool fp8 unified_kv layout.
+
+``store_swa_into_unified`` writes one latent row per token. Under
+SGLANG_DSV4_UNIFIED_KV_FP8 that row is split over a packed fp8 nope pool and a
+bf16 rope pool, so what these tests pin is that the ring row index -- derived
+from state_slot/positions alone -- stays identical to the bf16 layout's and
+identical between the two pools.
+"""
+
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+# the scatter itself is a plain row move, but the layout it pins is gfx95-only, so
+# run it where the feature lives rather than on the default mi300 runner
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+DEVICE = torch.device("cuda")
+
+# 448 values + 14 E8M0 scales + 50 pad, in bytes
+NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
+ROPE_DIM = 64
+# V4-Pro latent, in elements -- same number as NOPE_ROW_BYTES, different unit
+BF16_LATENT = 448 + ROPE_DIM
+
+RING_STRIDE = 16
+WIN = 8
+N_PAGES = 64
+
+
+def _inputs(n_rows=12):
+    """state_slot/positions whose ring rows are all distinct, so a row's writer is unambiguous"""
+    state_slot = torch.tensor(
+        [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3][:n_rows],
+        device=DEVICE,
+        dtype=torch.int32,
+    )
+    positions = torch.tensor(
+        [0, 1, 2, 16, 17, 18, 32, 33, 34, 48, 49, 50][:n_rows],
+        device=DEVICE,
+        dtype=torch.int32,
+    )
+    return state_slot, positions
+
+
+def _expected_rows(state_slot, positions, final_pos=None):
+    loc = state_slot.long() * RING_STRIDE + positions.long() % RING_STRIDE
+    if final_pos is None:
+        keep = torch.ones_like(loc, dtype=torch.bool)
+    else:
+        keep = positions.long() > final_pos.long() - WIN
+    return loc, keep
+
+
+def _packed_nope(n_rows):
+    """random packed fp8 rows; byte 0 is forced nonzero so a written row is detectable"""
+    raw = torch.randint(
+        0, 256, (n_rows, NOPE_ROW_BYTES), device=DEVICE, dtype=torch.uint8
+    )
+    raw[:, 0] = torch.arange(1, n_rows + 1, device=DEVICE, dtype=torch.uint8)
+    return raw.view(torch.float8_e4m3fn), raw
+
+
+def _bf16_rope(n_rows):
+    rope = torch.randn(n_rows, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+    rope[:, 0] = torch.arange(1, n_rows + 1, device=DEVICE, dtype=torch.bfloat16)
+    return rope.contiguous()
+
+
+def _store(kv, pool, state_slot, positions, final_pos=None, **kw):
+    runtime.store_swa_into_unified(
+        kv=kv,
+        state_slot=state_slot,
+        positions=positions,
+        unified_kv=pool,
+        win=WIN,
+        ring_stride=RING_STRIDE,
+        final_pos=final_pos,
+        **kw,
+    )
+
+
+class TestUnifiedFp8SwaScatter(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(20)
+        self.state_slot, self.positions = _inputs()
+        self.n_rows = self.state_slot.shape[0]
+
+    def _run_two_pool(self, final_pos=None):
+        kv_nope, nope_bytes = _packed_nope(self.n_rows)
+        kv_rope = _bf16_rope(self.n_rows)
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+        _store(
+            kv_nope,
+            pool_nope,
+            self.state_slot,
+            self.positions,
+            final_pos=final_pos,
+            kv_rope=kv_rope,
+            unified_kv_rope=pool_rope,
+        )
+        return pool_nope, pool_rope, nope_bytes, kv_rope
+
+    def _run_bf16(self, final_pos=None):
+        kv = torch.randn(
+            self.n_rows, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16
+        ).contiguous()
+        kv[:, 0] = torch.arange(1, self.n_rows + 1, device=DEVICE, dtype=torch.bfloat16)
+        pool = torch.zeros(N_PAGES, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16)
+        _store(kv, pool, self.state_slot, self.positions, final_pos=final_pos)
+        return pool, kv
+
+    def test_bf16_single_pool_unchanged(self):
+        """the bf16 path still writes exactly the expected ring rows"""
+        pool, kv = self._run_bf16()
+        loc, keep = _expected_rows(self.state_slot, self.positions)
+        expected = torch.zeros_like(pool)
+        expected[loc[keep]] = kv[keep]
+        self.assertTrue(torch.equal(pool, expected))
+
+    def test_two_pool_bytes_exact(self):
+        """each pool gets its half verbatim -- nope byte-for-byte, rope bit-for-bit"""
+        pool_nope, pool_rope, nope_bytes, kv_rope = self._run_two_pool()
+        loc, keep = _expected_rows(self.state_slot, self.positions)
+
+        exp_nope = torch.zeros_like(pool_nope).view(torch.uint8)
+        exp_nope[loc[keep]] = nope_bytes[keep]
+        self.assertTrue(torch.equal(pool_nope.view(torch.uint8), exp_nope))
+
+        exp_rope = torch.zeros_like(pool_rope)
+        exp_rope[loc[keep]] = kv_rope[keep]
+        self.assertTrue(torch.equal(pool_rope, exp_rope))
+
+    def test_two_pool_rows_match_bf16(self):
+        """same state_slot/positions -> same ring rows as bf16, and the same in both pools"""
+        pool_nope, pool_rope, _, _ = self._run_two_pool()
+        pool_bf16, _ = self._run_bf16()
+
+        rows_nope = (pool_nope.view(torch.uint8) != 0).any(dim=1)
+        rows_rope = (pool_rope != 0).any(dim=1)
+        rows_bf16 = (pool_bf16 != 0).any(dim=1)
+
+        self.assertTrue(torch.equal(rows_nope, rows_bf16))
+        self.assertTrue(torch.equal(rows_rope, rows_bf16))
+        self.assertEqual(int(rows_bf16.sum()), self.n_rows)
+
+    def test_final_pos_skips_both_pools(self):
+        """tokens already outside the window are skipped in nope and rope alike"""
+        # positions[t] <= final_pos[t] - WIN skips; give the first half a far
+        # final_pos and the second half its own position
+        final_pos = self.positions.clone()
+        final_pos[: self.n_rows // 2] = self.positions.max() + WIN
+        pool_nope, pool_rope, nope_bytes, kv_rope = self._run_two_pool(
+            final_pos=final_pos
+        )
+        loc, keep = _expected_rows(self.state_slot, self.positions, final_pos)
+        self.assertTrue(bool((~keep).any()), "test would be vacuous without a skip")
+
+        rows_nope = (pool_nope.view(torch.uint8) != 0).any(dim=1)
+        rows_rope = (pool_rope != 0).any(dim=1)
+        expected_rows = torch.zeros(N_PAGES, device=DEVICE, dtype=torch.bool)
+        expected_rows[loc[keep]] = True
+        self.assertTrue(torch.equal(rows_nope, expected_rows))
+        self.assertTrue(torch.equal(rows_rope, expected_rows))
+
+    def test_rope_tensor_and_pool_come_together(self):
+        kv_nope, _ = _packed_nope(self.n_rows)
+        kv_rope = _bf16_rope(self.n_rows)
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+        with self.assertRaises(AssertionError):
+            _store(
+                kv_nope,
+                pool_nope,
+                self.state_slot,
+                self.positions,
+                kv_rope=kv_rope,
+            )
+        with self.assertRaises(AssertionError):
+            _store(
+                kv_nope,
+                pool_nope,
+                self.state_slot,
+                self.positions,
+                unified_kv_rope=pool_rope,
+            )
+
+    def test_short_rope_pool_rejected(self):
+        """the kernel doesn't bound-check the ring row, so a rope pool with fewer
+        rows than the nope pool writes into whatever tensor follows it"""
+        kv_nope, _ = _packed_nope(self.n_rows)
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        # ring rows reach state_slot 3 -> row 48, well past this
+        pool_rope = torch.zeros(8, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+        with self.assertRaises(AssertionError):
+            _store(
+                kv_nope,
+                pool_nope,
+                self.state_slot,
+                self.positions,
+                kv_rope=_bf16_rope(self.n_rows),
+                unified_kv_rope=pool_rope,
+            )
+
+    def test_rope_row_width_mismatch_rejected(self):
+        """row width is read off src, so a wider pool would place row i at i * D"""
+        kv_nope, _ = _packed_nope(self.n_rows)
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        pool_rope = torch.zeros(
+            N_PAGES, ROPE_DIM * 2, device=DEVICE, dtype=torch.bfloat16
+        )
+        with self.assertRaises(AssertionError):
+            _store(
+                kv_nope,
+                pool_nope,
+                self.state_slot,
+                self.positions,
+                kv_rope=_bf16_rope(self.n_rows),
+                unified_kv_rope=pool_rope,
+            )
+
+    def test_dtype_mismatch_rejected(self):
+        """a bf16 row must not land in an fp8 pool (the DSpark-under-fp8 case)"""
+        kv = torch.randn(
+            self.n_rows, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.bfloat16
+        )
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        with self.assertRaises(AssertionError):
+            _store(kv, pool_nope, self.state_slot, self.positions)
+
+    def test_empty_batch_is_a_noop(self):
+        empty_slot = torch.zeros(0, device=DEVICE, dtype=torch.int32)
+        kv_nope, _ = _packed_nope(0)
+        pool_nope = torch.zeros(
+            N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
+        )
+        pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
+        _store(
+            kv_nope,
+            pool_nope,
+            empty_slot,
+            empty_slot,
+            kv_rope=_bf16_rope(0),
+            unified_kv_rope=pool_rope,
+        )
+        self.assertEqual(int((pool_nope.view(torch.uint8) != 0).sum()), 0)
+        self.assertEqual(int((pool_rope != 0).sum()), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py

@@ -0,0 +1,128 @@
+import contextlib
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    DSV4_FP8_NOPE_ROW_BYTES,
+    DSV4_FP8_QUANT_TILE,
+    DeepSeekV4UnifiedKVPool,
+    dsv4_unified_row_bytes,
+)
+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")
+
+# DeepSeek-V4-Pro geometry.
+NOPE_DIM = 448
+ROPE_DIM = 64
+
+
+class _StubMemorySaver:
+    def region(self, _tag):
+        return contextlib.nullcontext()
+
+
+class TestDSV4UnifiedRowBytes(CustomTestCase):
+    """Row width drives both `bytes_per_full_token` and `_fixed_swa_bytes`, so the
+    capacity claim for the fp8 pool is only as good as this arithmetic."""
+
+    def test_bf16_row_is_the_whole_latent(self):
+        self.assertEqual(
+            dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False),
+            (NOPE_DIM + ROPE_DIM) * 2,
+        )
+
+    def test_fp8_row_is_padded_nope_plus_bf16_rope(self):
+        self.assertEqual(
+            dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True),
+            DSV4_FP8_NOPE_ROW_BYTES + ROPE_DIM * 2,
+        )
+
+    def test_fp8_saves_exactly_three_eighths(self):
+        """0.625x is where the >=1.40x capacity target comes from; the remaining
+        dilution is the fixed SWA/c4-state bias, not the row."""
+        bf16 = dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False)
+        fp8 = dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True)
+        self.assertEqual((bf16, fp8), (1024, 640))
+        self.assertAlmostEqual(fp8 / bf16, 0.625)
+
+    def test_scales_and_latent_fit_the_asm_stride(self):
+        """7 tiles written twice = 14 B; 448 + 14 leaves 50 B the reader never
+        touches. If a future head_dim broke this the pack would silently overlap."""
+        num_tiles = NOPE_DIM // DSV4_FP8_QUANT_TILE
+        self.assertEqual(num_tiles, 7)
+        self.assertLessEqual(NOPE_DIM + 2 * num_tiles, DSV4_FP8_NOPE_ROW_BYTES)
+
+    def test_oversized_latent_is_rejected(self):
+        # ValueError, not assert: sizing has to keep checking under python -O
+        with self.assertRaises(ValueError):
+            dsv4_unified_row_bytes(DSV4_FP8_NOPE_ROW_BYTES, ROPE_DIM, fp8=True)
+
+
+class TestDSV4UnifiedFp8PoolAllocation(CustomTestCase):
+    """The sizing formula and the allocation are two separate code paths; this pins
+    them to the same row width so a change to one cannot silently outrun the other."""
+
+    STAGE_RATIOS = [4, 128]
+    NUM_SLOTS = 3
+    NUM_BLOCKS = 5
+    PAGE_SIZE = 256
+    SWA_RING = 8
+
+    def _pool(self, fp8):
+        return DeepSeekV4UnifiedKVPool(
+            stage_ratios=self.STAGE_RATIOS,
+            num_slots=self.NUM_SLOTS,
+            num_blocks=self.NUM_BLOCKS,
+            page_size=self.PAGE_SIZE,
+            qk_nope_head_dim=NOPE_DIM,
+            qk_rope_head_dim=ROPE_DIM,
+            device="cpu",
+            memory_saver_adapter=_StubMemorySaver(),
+            custom_mem_pool=None,
+            swa_ring_size=self.SWA_RING,
+            fp8=fp8,
+        )
+
+    def test_bf16_pool_is_unchanged(self):
+        """fp8 defaults off, so the bf16 arm must keep one pool of bf16 latents."""
+        pool = self._pool(fp8=False)
+        for buf, rope in zip(pool.kv_buffer, pool.kv_buffer_rope):
+            self.assertEqual(buf.dtype, torch.bfloat16)
+            self.assertEqual(buf.shape[1], NOPE_DIM + ROPE_DIM)
+            self.assertIsNone(rope)
+
+    def test_fp8_pool_row_counts_match_across_both_pools(self):
+        """A row index addresses the SWA ring and the compressed region in both
+        pools, so the two must have identical row counts."""
+        pool = self._pool(fp8=True)
+        for buf, rope in zip(pool.kv_buffer, pool.kv_buffer_rope):
+            self.assertEqual(buf.dtype, torch.float8_e4m3fn)
+            self.assertEqual(rope.dtype, torch.bfloat16)
+            self.assertEqual(buf.shape[0], rope.shape[0])
+            self.assertEqual(buf.shape[1], DSV4_FP8_NOPE_ROW_BYTES)
+            self.assertEqual(rope.shape[1], ROPE_DIM)
+
+    def test_fp8_pool_bytes_match_the_sizing_row_width(self):
+        bf16, fp8 = self._pool(fp8=False), self._pool(fp8=True)
+        for layer, buf in enumerate(bf16.kv_buffer):
+            rows = buf.shape[0]
+            self.assertEqual(fp8.kv_buffer[layer].shape[0], rows)
+            self.assertEqual(
+                buf.nbytes,
+                rows * dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False),
+            )
+            self.assertEqual(
+                fp8.kv_buffer[layer].nbytes + fp8.kv_buffer_rope[layer].nbytes,
+                rows * dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True),
+            )
+
+    def test_rope_accessor_rejects_the_bf16_pool(self):
+        with self.assertRaises(AssertionError):
+            self._pool(fp8=False).get_unified_kv_rope(0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_executor/test_pool_configurator.py

@@ -1118,6 +1118,9 @@ class TestSWAPoolFloor(CustomTestCase):
         cfg.disaggregation_mode = None
         cfg.disaggregation_decode_extra_slots = 0
         cfg._unified = True
+        cfg._unified_fp8 = False
+        # object.__new__ skips __init__; bf16 unified row is 2B * latent
+        cfg._unified_row_bytes = cfg.attn_head_dim * 2
         return cfg
 
     # Token pool plus the three request-scoped fixed pools, sized from the

test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py

@@ -0,0 +1,279 @@
+"""DeepSeek-V4 unified_kv fp8: the packed pairs handed to the two readers.
+
+Decode only needs Q packed -- its K is already in the ring. Prefill is a KV
+source of its own, so it gets a packed K pair beside the Q one, and the same
+buffers have to reach both attention and the ring write after it. Verify wants
+both halves: it reads the ring the way decode does and fills it the way prefill
+does, only the write lands before attention instead of after.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+import sglang.srt.models.deepseek_v4 as deepseek_v4
+from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import env_gate
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+# deliberately != head_dim below: the row width has to come off the pool, since
+# that is the stride the kernel reads Q with. Sharing head_dim's value would let
+# a regression that reads self.head_dim pass.
+NOPE_ROW_BYTES = 16
+ROPE_DIM = 2
+HEAD_DIM = 8
+N_LOCAL_HEADS = 16
+TOKENS = 3
+
+
+class _RecordingBackend:
+    def __init__(self):
+        self.calls = []
+
+    def forward(self, **kwargs):
+        self.calls.append(kwargs)
+        query = kwargs["q"]
+        # bf16 regardless of the q layout -- attention output is never fp8
+        return torch.zeros(
+            query.shape[0], query.shape[1], ROPE_DIM, dtype=torch.bfloat16
+        )
+
+
+class _Pool:
+    def __init__(self, fp8):
+        rows = 32
+        self.nope = torch.zeros(
+            rows, NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16
+        )
+        self.rope = torch.zeros(rows, ROPE_DIM, dtype=torch.bfloat16)
+
+    def get_unified_kv(self, layer_id):
+        return self.nope
+
+    def get_unified_kv_rope(self, layer_id):
+        return self.rope
+
+
+class _Harness(deepseek_v4.MQALayer):
+    def __init__(self, rank=3):
+        torch.nn.Module.__init__(self)
+        self.layer_id = 0
+        self.attn_tp_rank = rank
+        self.attn_tp_size = 8
+        self.n_heads = 128
+        self.n_local_heads = N_LOCAL_HEADS
+        self.head_dim = HEAD_DIM
+        self.n_local_groups = 1
+        self.o_lora_rank = 3
+        self.qk_rope_head_dim = ROPE_DIM
+        self.freqs_cis = torch.empty(0)
+        self.compress_ratio = 4
+        self.attn_mqa = SimpleNamespace(layer_id=0, v_head_dim=ROPE_DIM)
+        self.attn_sink = torch.nn.Parameter(torch.arange(128, dtype=torch.float32))
+        self._attn_sink_local = None
+        self.alt_streams = None
+        self.dsa_enable_prefill_cp = False
+        self.use_npu_arch35_mxfp8_wo_a = False
+        self.compressor = object()
+        self.wo_a = SimpleNamespace(
+            weight=torch.ones(
+                self.n_local_groups,
+                self.o_lora_rank,
+                self.n_local_heads * ROPE_DIM,
+                dtype=torch.bfloat16,
+            )
+        )
+        self.wo_b = lambda value: (value, None)
+        self.prepare_kwargs = None
+
+    def _forward_prepare(
+        self,
+        x,
+        positions,
+        forward_batch,
+        attn_backend,
+        q_out=None,
+        x_quant=None,
+        q_rope_out=None,
+        k_nope_out=None,
+        k_rope_out=None,
+    ):
+        self.prepare_kwargs = dict(
+            q_out=q_out,
+            q_rope_out=q_rope_out,
+            k_nope_out=k_nope_out,
+            k_rope_out=k_rope_out,
+        )
+        q_out.zero_()
+        # mirrors the prefill arm: the packed nope half leaves on the kv slot,
+        # which is what turns save_kv_cache on in the caller
+        return q_out, k_nope_out
+
+
+def _run(fp8, mode=ForwardMode.DECODE, cp=False, fused_verify=True):
+    layer = _Harness()
+    layer.dsa_enable_prefill_cp = cp
+    backend = _RecordingBackend()
+    forward_batch = SimpleNamespace(forward_mode=mode)
+
+    with (
+        envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.override(False),
+        envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.override(fused_verify),
+        patch.object(env_gate, "is_unified_kv_triton", return_value=True),
+        patch.object(env_gate, "is_unified_kv_fp8", return_value=fp8),
+        patch.object(deepseek_v4, "get_token_to_kv_pool", return_value=_Pool(fp8)),
+        patch.object(
+            deepseek_v4,
+            "get_attn_tp_context",
+            return_value=SimpleNamespace(input_scattered=True),
+        ),
+        patch.object(
+            deepseek_v4, "get_parallel", return_value=SimpleNamespace(tp_size=8)
+        ),
+        patch.object(deepseek_v4, "get_attn_backend", return_value=backend),
+        patch.object(deepseek_v4, "dsa_use_prefill_cp", return_value=cp),
+        patch.object(deepseek_v4, "fused_rope_inplace", return_value=None),
+        patch.object(deepseek_v4, "_FP8_WO_A_GEMM", False),
+        patch.object(deepseek_v4, "_is_gfx942_supported", False),
+        patch.object(deepseek_v4, "_is_hip", True),
+        patch.object(deepseek_v4, "_is_npu", False),
+    ):
+        layer.forward(
+            torch.zeros(TOKENS, 4, dtype=torch.bfloat16),
+            torch.arange(TOKENS),
+            forward_batch,
+        )
+
+    return layer, backend.calls[0]
+
+
+class TestUnifiedFp8QPair(unittest.TestCase):
+    def test_fp8_decode_hands_the_backend_a_packed_pair(self):
+        layer, call = _run(fp8=True)
+
+        q, q_rope = call["q"], call["q_rope"]
+        self.assertEqual(q.dtype, torch.float8_e4m3fn)
+        # width off the pool, not off head_dim
+        self.assertEqual(tuple(q.shape), (TOKENS, N_LOCAL_HEADS, NOPE_ROW_BYTES))
+        self.assertEqual(tuple(q_rope.shape), (TOKENS, N_LOCAL_HEADS, ROPE_DIM))
+        self.assertEqual(q_rope.dtype, torch.bfloat16)
+        # the asm kernel walks both as flat buffers, no stride arguments
+        self.assertTrue(q.is_contiguous())
+        self.assertTrue(q_rope.is_contiguous())
+        # same pair reached the store, or nothing would have written them
+        self.assertIs(layer.prepare_kwargs["q_out"], q)
+        self.assertIs(layer.prepare_kwargs["q_rope_out"], q_rope)
+
+    def test_bf16_decode_still_gets_one_plain_tensor(self):
+        layer, call = _run(fp8=False)
+
+        # q_rope absent is what routes the backend back to the Triton reader
+        self.assertNotIn("q_rope", call)
+        self.assertIsNone(layer.prepare_kwargs["q_rope_out"])
+        self.assertEqual(call["q"].dtype, torch.bfloat16)
+        self.assertEqual(tuple(call["q"].shape), (TOKENS, N_LOCAL_HEADS, HEAD_DIM))
+
+    def test_fp8_prefill_also_gets_a_packed_k_pair(self):
+        layer, call = _run(fp8=True, mode=ForwardMode.EXTEND)
+
+        k, k_rope = call["k"], call["k_rope"]
+        self.assertEqual(k.dtype, torch.float8_e4m3fn)
+        # one row per token, width off the pool like Q
+        self.assertEqual(tuple(k.shape), (TOKENS, NOPE_ROW_BYTES))
+        self.assertEqual(tuple(k_rope.shape), (TOKENS, ROPE_DIM))
+        self.assertEqual(k_rope.dtype, torch.bfloat16)
+        self.assertTrue(k.is_contiguous())
+        self.assertTrue(k_rope.is_contiguous())
+        # the buffers the fused store filled are the ones attention reads, and
+        # the ring write after it consumes the same rows
+        self.assertIs(layer.prepare_kwargs["k_nope_out"], k)
+        self.assertIs(layer.prepare_kwargs["k_rope_out"], k_rope)
+        self.assertTrue(call["save_kv_cache"])
+        # Q is packed here too, that is what picks the fp8 prefill kernel
+        self.assertEqual(call["q"].dtype, torch.float8_e4m3fn)
+        self.assertIsNotNone(call["q_rope"])
+
+    def test_fp8_decode_gets_no_k_pair(self):
+        """decode attends over rows the ring already holds, so it has no extend"""
+        layer, call = _run(fp8=True, mode=ForwardMode.DECODE)
+
+        self.assertNotIn("k_rope", call)
+        self.assertIsNone(layer.prepare_kwargs["k_nope_out"])
+        self.assertIsNone(layer.prepare_kwargs["k_rope_out"])
+
+    def test_bf16_prefill_keeps_one_plain_tensor(self):
+        layer, call = _run(fp8=False, mode=ForwardMode.EXTEND)
+
+        self.assertNotIn("q_rope", call)
+        self.assertNotIn("k_rope", call)
+        self.assertIsNone(layer.prepare_kwargs["k_nope_out"])
+        self.assertEqual(call["q"].dtype, torch.bfloat16)
+
+    def test_fp8_target_verify_gets_the_packed_pair(self):
+        """verify reads the ring like decode, but it also feeds it like prefill"""
+        layer, call = _run(fp8=True, mode=ForwardMode.TARGET_VERIFY)
+
+        # packed Q is what picks the decode reader over the Triton one
+        self.assertEqual(call["q"].dtype, torch.float8_e4m3fn)
+        self.assertIsNotNone(call["q_rope"])
+        k, k_rope = call["k"], call["k_rope"]
+        self.assertEqual(k.dtype, torch.float8_e4m3fn)
+        self.assertEqual(tuple(k.shape), (TOKENS, NOPE_ROW_BYTES))
+        self.assertEqual(tuple(k_rope.shape), (TOKENS, ROPE_DIM))
+        self.assertIs(layer.prepare_kwargs["k_nope_out"], k)
+        self.assertIs(layer.prepare_kwargs["k_rope_out"], k_rope)
+        # unlike prefill the ring write happens before attention, but it is the
+        # same flag and the same pair
+        self.assertTrue(call["save_kv_cache"])
+
+    def test_fp8_target_verify_needs_the_fused_store(self):
+        """nothing else packs the pair, so the unfused arm would hand over bf16"""
+        with self.assertRaisesRegex(
+            NotImplementedError, "SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY"
+        ):
+            _run(fp8=True, mode=ForwardMode.TARGET_VERIFY, fused_verify=False)
+
+    def test_bf16_target_verify_is_left_alone(self):
+        """the packing is fp8-only; bf16 verify keeps working as it always did"""
+        layer, call = _run(fp8=False, mode=ForwardMode.TARGET_VERIFY)
+
+        self.assertNotIn("q_rope", call)
+        self.assertNotIn("k_rope", call)
+        self.assertIsNone(layer.prepare_kwargs["k_nope_out"])
+
+    def test_fp8_prefill_cp_is_refused_with_a_reason(self):
+        """the gather hands kv back in global token order after norm+RoPE, so
+        packing would have to move ahead of it -- refuse rather than guess"""
+        with self.assertRaisesRegex(NotImplementedError, "cp_size"):
+            _run(fp8=True, mode=ForwardMode.EXTEND, cp=True)
+
+    def test_bf16_prefill_cp_is_left_alone(self):
+        """the refusal is fp8-only, CP prefill without it keeps working"""
+        _, call = _run(fp8=False, mode=ForwardMode.EXTEND, cp=True)
+
+        self.assertNotIn("q_rope", call)
+        self.assertNotIn("k_rope", call)
+
+    def test_fp8_decode_under_cp_is_not_refused(self):
+        """only prefill packs this chunk; decode reads rows the ring already has"""
+        _, call = _run(fp8=True, mode=ForwardMode.DECODE, cp=True)
+
+        self.assertEqual(call["q"].dtype, torch.float8_e4m3fn)
+
+    def test_sink_is_sliced_to_this_rank(self):
+        _, call = _run(fp8=True)
+
+        sink = call["attn_sink"]
+        self.assertEqual(tuple(sink.shape), (N_LOCAL_HEADS,))
+        torch.testing.assert_close(
+            sink, torch.arange(3 * N_LOCAL_HEADS, 4 * N_LOCAL_HEADS).float()
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()