diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh index a3b411575..0fda6b8a1 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -46,6 +46,8 @@ struct FusedNormRopeStoreParams { const float* __restrict__ freqs_cis; const int64_t* __restrict__ out_loc; uint8_t* __restrict__ kvcache; + // second pool for the fp8 two-pool store; the other layouts keep rope inline + uint8_t* __restrict__ kvcache_rope = nullptr; float eps; uint32_t compress_ratio; uint32_t num_tokens; @@ -375,12 +377,24 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor } } +// 448 B fp8 nope payload + 7 UE8M0 tile scales written twice, padded up to a power of +// two. Has to stay in step with DSV4_FP8_NOPE_ROW_BYTES in +// ops/attention/dsv4/unified_kv_kernels/layout.py; nothing checks that across the +// language boundary. +constexpr int64_t kFp8TwoPoolRowBytes = 512; + // ---------------------------------------------------------------------------- // FlashMLA variant: kHeadDim = 512, 1 token per *block* (256 threads). // Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems. // Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale. // ---------------------------------------------------------------------------- -template +template < + typename DType, + ForwardMode kMode, + int32_t kPageBits, + bool kUsePDL, + bool kBf16Store = false, + bool kFp8TwoPool = false> FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) { using namespace device; using enum ForwardMode; @@ -393,8 +407,12 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR constexpr uint32_t kRopeWarp = kNumWarps - 1; // kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a // [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc + // kFp8TwoPool: 512 B row holding the 448 fp8 nope + its UE8M0 scales, with rope + // split off into a second [num_slots, kRopeDim] bf16 pool at the same row + static_assert(!(kBf16Store && kFp8TwoPool)); + constexpr int64_t kRowBytes = kBf16Store ? (kHeadDim * 2ll) : (kFp8TwoPool ? kFp8TwoPoolRowBytes : 576ll); constexpr int64_t kPageBytes = - kBf16Store ? ((kHeadDim * 2ll) << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576; + (kBf16Store || kFp8TwoPool) ? (kRowBytes << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576; static_assert(kHeadDim == kBlockSize * kVecSize); static_assert(kRopeDim == kWarpThreads * kVecSize); static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize); @@ -465,7 +483,7 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR const int64_t page = out_loc >> kPageBits; const int64_t offset = out_loc & ((1 << kPageBits) - 1); const auto page_ptr = params.kvcache + page * kPageBytes; - const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576); + const auto value_ptr = page_ptr + offset * kRowBytes; PDLTriggerSecondary(); @@ -491,7 +509,9 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR data[0] = x_real * freq_real - x_imag * freq_imag; data[1] = x_real * freq_imag + x_imag * freq_real; const auto result = cast(fp32x2_t{data[0], data[1]}); - const auto rope_ptr = value_ptr + 448; + // out_loc indexes the rope pool directly: its rows are kRopeDim * 2 B wide no + // matter how the nope pool is paged + const auto rope_ptr = kFp8TwoPool ? (params.kvcache_rope + out_loc * kRopeDim * 2) : (value_ptr + 448); reinterpret_cast(rope_ptr)[lane_id] = result; } else { // Non-rope warp: per-warp UE8M0 group (64 elems -> 64 fp8 + 1 scale byte). @@ -504,10 +524,20 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR const auto scale_ue8m0 = cast_to_ue8m0(scale_raw); const auto inv_scale = inv_scale_ue8m0(scale_ue8m0); const auto result = pack_fp8(x * inv_scale, y * inv_scale); - const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8; reinterpret_cast(value_ptr)[tx] = result; // All lanes in this warp produce the same scale byte; let lane 0 publish. - if (lane_id == 0) static_cast(scale_ptr)[warp_id] = scale_ue8m0; + if (lane_id == 0) { + if constexpr (kFp8TwoPool) { + // scales sit in the same row behind the payload, and the decode reader + // loads each one twice + const auto scale_ptr = value_ptr + 448 + warp_id * 2; + scale_ptr[0] = scale_ue8m0; + scale_ptr[1] = scale_ue8m0; + } else { + const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8; + static_cast(scale_ptr)[warp_id] = scale_ue8m0; + } + } } } @@ -541,12 +571,61 @@ struct FusedNormRopeKernel { } } + template + static constexpr auto select_fp8_2buff_kernel() { + static_assert(!kIsIndexer, "fp8 two-pool store is only defined for the flashmla latent"); + static_assert(!kBf16Store, "fp8 two-pool store and bf16 store are separate layouts"); + return fused_norm_rope_flashmla; + } + template static constexpr auto select_fp4_kernel() { static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer"); return fused_norm_rope_indexer_fp4; } + // Everything except the cache tensors is the same whichever layout we store into. + // Each wrapper still matches its own cache in between these two, so the order a + // caller sees errors in does not change. + static void verify_operands( + const tvm::ffi::TensorView& input, + const tvm::ffi::TensorView& weight, + const tvm::ffi::TensorView& freqs_cis, + const tvm::ffi::TensorView& out_loc, + host::SymbolicSize& N, + host::SymbolicDevice& device_) { + using namespace host; + TensorMatcher({N, kHeadDim}).with_dtype().with_device(device_).verify(input); + TensorMatcher({kHeadDim}).with_dtype().with_device(device_).verify(weight); + TensorMatcher({-1, kRopeDim}).with_dtype().with_device(device_).verify(freqs_cis); + TensorMatcher({-1}).with_dtype().with_device(device_).verify(out_loc); + } + + // Careful with the extend bound: that arm addresses out_loc by the plan's ragged_id, + // i.e. by q token, so N (compressed tokens) is a floor and not a bound. Sizing + // out_loc to N passes this check and then reads off the end -- with page_size 1 and + // a c128 ratio the garbage row index faults outright. Nothing on the host side can + // see the ragged length, so the caller owns it. + static void verify_plan_for_mode( + const ForwardMode mode, + const tvm::ffi::TensorView& plan, + const tvm::ffi::TensorView& out_loc, + host::SymbolicSize& N, + host::SymbolicDevice& device_) { + using namespace host; + using enum ForwardMode; + switch (mode) { + case CompressExtend: + compress::verify_plan_c(plan, N, device_); + RuntimeCheck(out_loc.size(0) >= N.unwrap()); + break; + case CompressDecode: + compress::verify_plan_d(plan, N, device_); + RuntimeCheck(out_loc.size(0) == N.unwrap()); + break; + } + } + static void forward( const tvm::ffi::TensorView input, const tvm::ffi::TensorView plan, @@ -566,38 +645,13 @@ struct FusedNormRopeKernel { auto device_ = SymbolicDevice{}; device_.set_options(); - TensorMatcher({N, kHeadDim}) // input - .with_dtype() - .with_device(device_) - .verify(input); - TensorMatcher({kHeadDim}) // weight - .with_dtype() - .with_device(device_) - .verify(weight); - TensorMatcher({-1, kRopeDim}) // freqs_cis - .with_dtype() - .with_device(device_) - .verify(freqs_cis); - TensorMatcher({-1}) // out_loc - .with_dtype() - .with_device(device_) - .verify(out_loc); + verify_operands(input, weight, freqs_cis, out_loc, N, device_); TensorMatcher({-1, -1}) // cache .with_strides({kPageBytes, 1}) .with_dtype() .with_device(device_) .verify(kvcache); - - switch (mode) { - case CompressExtend: - compress::verify_plan_c(plan, N, device_); - RuntimeCheck(out_loc.size(0) >= N.unwrap()); - break; - case CompressDecode: - compress::verify_plan_d(plan, N, device_); - RuntimeCheck(out_loc.size(0) == N.unwrap()); - break; - } + verify_plan_for_mode(mode, plan, out_loc, N, device_); const auto num_tokens = static_cast(N.unwrap()); if (num_tokens == 0) return; @@ -620,6 +674,64 @@ struct FusedNormRopeKernel { LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params); } + // Same store as `forward` minus the packed 584 B layout: the fp8 nope row goes to + // `kvcache` (512 B rows) and rope to `kvcache_rope`, both indexed by out_loc. + // Callers pass byte views of the two unified_kv pools. + static void forward_fp8_2buff( + const tvm::ffi::TensorView input, + const tvm::ffi::TensorView plan, + const tvm::ffi::TensorView weight, + const float eps, + const tvm::ffi::TensorView freqs_cis, + const tvm::ffi::TensorView out_loc, + const tvm::ffi::TensorView kvcache, + const tvm::ffi::TensorView kvcache_rope, + const bool is_decode, + const uint32_t compress_ratio) { + using namespace host; + using enum ForwardMode; + + static_assert(!kIsIndexer, "fp8 two-pool store is only defined for the flashmla latent"); + constexpr int64_t kFp8PageBytes = kFp8TwoPoolRowBytes * kPageSize; + constexpr int64_t kRopeRowBytes = kRopeDim * 2; + const auto mode = static_cast(is_decode); + + auto N = SymbolicSize{"num_tokens"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + verify_operands(input, weight, freqs_cis, out_loc, N, device_); + TensorMatcher({-1, -1}).with_strides({kFp8PageBytes, 1}).with_dtype().with_device(device_).verify(kvcache); + TensorMatcher({-1, kRopeRowBytes}) + .with_strides({kRopeRowBytes, 1}) + .with_dtype() + .with_device(device_) + .verify(kvcache_rope); + // one row index addresses both pools, so a short rope pool would let the rope + // warp write past its end + RuntimeCheck(kvcache_rope.size(0) == kvcache.size(0) * static_cast(kPageSize)); + verify_plan_for_mode(mode, plan, out_loc, N, device_); + + const auto num_tokens = static_cast(N.unwrap()); + if (num_tokens == 0) return; + const auto params = FusedNormRopeStoreParams{ + .input = input.data_ptr(), + .handle = plan.data_ptr(), + .weight = weight.data_ptr(), + .freqs_cis = static_cast(freqs_cis.data_ptr()), + .out_loc = static_cast(out_loc.data_ptr()), + .kvcache = static_cast(kvcache.data_ptr()), + .kvcache_rope = static_cast(kvcache_rope.data_ptr()), + .eps = eps, + .compress_ratio = compress_ratio, + .num_tokens = num_tokens, + }; + const auto device = device_.unwrap(); + const auto kernel = + mode == CompressExtend ? select_fp8_2buff_kernel() : select_fp8_2buff_kernel(); + LaunchKernel(num_tokens, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params); + } + static void forward_fp4( const tvm::ffi::TensorView input, const tvm::ffi::TensorView plan, diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index 9650bb522..becf52f07 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -49,6 +49,7 @@ def _jit_compress_norm_rope_module( rope_dim: int, page_size: int, bf16_store: bool = False, + fp8_2buff: bool = False, ) -> Module: args = make_cpp_args( dtype, @@ -64,6 +65,13 @@ def _jit_compress_norm_rope_module( cuda_wrappers.append( ("forward_fp4", f"FusedNormRopeKernel<{args}>::forward_fp4") ) + # elif because forward_fp8_2buff cannot even instantiate at head_dim 128 -- the kernel + # static_asserts the two-pool store is latent-only. The default latent arm skips it as + # well, so it doesn't carry a symbol nothing calls. + elif fp8_2buff: + cuda_wrappers.append( + ("forward_fp8_2buff", f"FusedNormRopeKernel<{args}>::forward_fp8_2buff") + ) return load_jit( make_name(f"fused_norm_rope_v2"), *args, @@ -447,6 +455,8 @@ def compress_norm_rope_store( kvcache_scale: Optional[torch.Tensor] = None, rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None, fp4_k_write_metadata=None, + fp8_2buff: bool = False, + kvcache_rope: Optional[torch.Tensor] = None, ) -> None: if use_fp4: assert kv.shape[-1] == 128 @@ -470,6 +480,11 @@ def compress_norm_rope_store( ) return + if fp8_2buff: + assert not (use_fp4 or bf16_store), "fp8 two-pool store is its own layout" + assert kv.shape[-1] != 128, "fp8 two-pool store is the latent, not the indexer" + assert kvcache_rope is not None, "fp8 two-pool store needs the rope pool" + assert not _is_xpu, "fp8 two-pool store is only wired for the CUDA/HIP kernel" freq_cis = torch.view_as_real(freq_cis).flatten(-2) if _is_xpu: compress_norm_rope_store_xpu( @@ -487,9 +502,19 @@ def compress_norm_rope_store( ) else: module = _jit_compress_norm_rope_module( - kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store + kv.dtype, + kv.shape[-1], + freq_cis.shape[-1], + page_size, + bf16_store, + fp8_2buff, ) - fn = module.forward_fp4 if use_fp4 else module.forward + if use_fp4: + fn, extra = module.forward_fp4, () + elif fp8_2buff: + fn, extra = module.forward_fp8_2buff, (kvcache_rope,) + else: + fn, extra = module.forward, () if norm_weight.dtype != kv.dtype: norm_weight = norm_weight.to(dtype=kv.dtype) fn( @@ -500,6 +525,7 @@ def compress_norm_rope_store( freq_cis, out_loc, kvcache, + *extra, plan.is_decode, plan.compress_ratio, ) diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py index c55ba905e..6ef30f2ee 100644 --- a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py +++ b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py @@ -1,12 +1,35 @@ from __future__ import annotations import functools +import logging from sglang.srt.environ import envs -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_gfx95_supported, is_hip + +logger = logging.getLogger(__name__) @functools.lru_cache(maxsize=1) def is_unified_kv_triton() -> bool: # unified_kv_triton is only implemented on HIP (ROCm) return is_hip() and envs.SGLANG_HACK_FLASHMLA_BACKEND.get() == "unified_kv_triton" + + +@functools.lru_cache(maxsize=1) +def is_unified_kv_fp8() -> bool: + # fp8 is a layout variant of the unified pool, so it can never outlive the + # unified gate -- the sizing, the allocation and the writers all key off this + # one call, so an unsupported device has to be turned away here or the three + # will disagree. + if not (is_unified_kv_triton() and envs.SGLANG_DSV4_UNIFIED_KV_FP8.get()): + return False + # two-pool fp8 is OCP e4m3 plus E8M0 tile scales, so it only means anything + # where MX is native: on gfx94x sglang's fp8_dtype is e4m3fnuz (max 224, not + # 448) and the writers would feed values the pool's own dtype misreads. + if not is_gfx95_supported(): + logger.warning( + "SGLANG_DSV4_UNIFIED_KV_FP8=1 needs an AMD gfx95 GPU; falling back to " + "the bf16 unified_kv pool (see unified_fp8= in the DSV4 memory log)." + ) + return False + return True diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py new file mode 100644 index 000000000..472ce0add --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py @@ -0,0 +1,62 @@ +"""Row layout of the two-pool fp8 unified_kv cache, shared by its writers. + +The pools are separate allocations with the same row count and one row index +addresses both, so these numbers belong with the kernels that write the rows +rather than with the pool that allocates them. Neither writer bounds-checks that +index -- the Triton scatter walks off the end of the shorter pool, aiter's fused +store aborts the process with nothing on stderr -- so the pair has to be checked +before the launch. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +# The fp8 nope row is a fixed 512 B whatever the payload: 448 B latent, then +# 14 B of E8M0 tile scales (7 tiles, each written twice -- the asm reader reads +# every tile scale twice), then 50 B nobody touches. Keep in sync with aiter's +# pack_v4_nope_scale and with kFp8TwoPoolRowBytes in +# jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh; the 512 B stride is what the +# reader assumes and nothing checks it across the language boundary. +DSV4_FP8_NOPE_ROW_BYTES = 512 +DSV4_FP8_QUANT_TILE = 64 + + +def check_two_pool_pair( + nope_pool: torch.Tensor, + rope_pool: Optional[torch.Tensor], + *, + rope_width: int, + rope_dtype: torch.dtype, +) -> None: + """Reject two pools that aren't a pair, before anything is written. + + ``rope_width`` is what the caller believes the rope row is (rot_dim for the + fused store, the source row width for the scatter). Both writers take the rope + row stride off the tensor, so a wider row would still land in the right place; + a width that disagrees with the caller means the wrong pool was fetched. + """ + assert rope_pool is not None, ( + "the fp8 layout needs a rope pool next to the nope pool" + ) + assert nope_pool.shape[0] == rope_pool.shape[0], ( + f"pool rows differ: nope {nope_pool.shape[0]} vs rope {rope_pool.shape[0]}" + ) + assert ( + nope_pool.element_size() == 1 and nope_pool.shape[-1] == DSV4_FP8_NOPE_ROW_BYTES + ), ( + f"nope pool must be the packed {DSV4_FP8_NOPE_ROW_BYTES} B fp8 row, got " + f"{nope_pool.shape[-1]} x {nope_pool.dtype}" + ) + assert rope_pool.shape[-1] == rope_width and rope_pool.dtype == rope_dtype, ( + f"rope pool is {rope_pool.shape[-1]} x {rope_pool.dtype}, expected " + f"{rope_width} x {rope_dtype}" + ) + assert nope_pool.is_contiguous(), ( + f"nope pool must be contiguous, got strides {nope_pool.stride()}" + ) + assert rope_pool.is_contiguous(), ( + f"rope pool must be contiguous, got strides {rope_pool.stride()}" + ) diff --git a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py index 79ae94ab0..39f02a9d7 100644 --- a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py +++ b/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py @@ -1,16 +1,19 @@ """Runtime glue for the unified_kv backend. Builds unified_kv-style flat ``kv_indices`` / ``kv_indptr`` from SGLang's already-computed -DSV4 metadata, scatters SWA K into the bf16 ``unified_kv`` ring, and dispatches the +DSV4 metadata, scatters SWA K into the ``unified_kv`` ring, and dispatches the vendored paged decode/prefill kernels. -unified_kv[L] layout (page_size 1, bf16, row-major): +unified_kv[L] layout (page_size 1, row-major): - rows ``[0, swa_pages)`` = SWA ring (``state_slot * win + pos % win``); - rows ``[swa_pages, ...)`` = compressed K (``swa_pages + page_index``), where SGLang metadata already encodes the compressed slot id: HCA (ratio 128): ``c128_page_indices`` (== phys_block, k_per_block=1) CSA (ratio 4): ``c4_sparse_page_indices`` (== phys_block*32 + slot) +Under SGLANG_DSV4_UNIFIED_KV_FP8 each row is split over two pools (512 B packed fp8 +nope + 128 B bf16 rope); row indexing and every index builder below are unchanged. + Index layout: RAGGED-PACKED. Each token's segment is tightly packed (``kv_indptr`` is a true prefix sum of per-token valid lengths) so the attention K-loop scans only real entries. The backing buffer is still @@ -24,6 +27,7 @@ on via ``topk_length``); the per-token compressed count is recovered from the from __future__ import annotations +from functools import lru_cache from typing import Optional, Tuple import torch @@ -31,6 +35,9 @@ import torch.nn.functional as F import triton import triton.language as tl +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.layout import ( + check_two_pool_pair, +) from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.paged_decode import ( sparse_attn_v4_paged_decode, ) @@ -77,15 +84,32 @@ def _swa_scatter_kernel( def store_swa_into_unified( *, - kv: torch.Tensor, # [T, head_dim] bf16 + kv: torch.Tensor, # [T, head_dim] bf16, or [T, nope_row_bytes] packed fp8 state_slot: torch.Tensor, # [T] int positions: torch.Tensor, # [T] int - unified_kv: torch.Tensor, # [pages, head_dim] bf16 + unified_kv: torch.Tensor, # [pages, ...] same dtype and row width as kv win: int, # SWA attention window length ring_stride: int, # SWA ring stride final_pos: Optional[torch.Tensor] = None, # [T] req's last position + kv_rope: Optional[torch.Tensor] = None, # [T, rope_dim] bf16, fp8 layout only + unified_kv_rope: Optional[torch.Tensor] = None, # [pages, rope_dim] bf16 ) -> None: - n_rows, D = kv.shape + """Scatter SWA K into ring row ``state_slot * ring_stride + pos % ring_stride``. + + Under the fp8 layout the latent is split over two pools, so ``kv`` carries the + already-packed nope row (DSV4_FP8_NOPE_ROW_BYTES wide: values + E8M0 scales + + pad; nothing is quantized here) and ``kv_rope`` the bf16 rope half. That width + is a byte count that happens to equal the bf16 head_dim in elements -- the two + are not the same thing. The scatter itself takes the row width off the tensor; + only the pair check reads the constant. + + Both scatters recompute the row index from the same ``state_slot`` / + ``positions``, so the two pools stay in lockstep with each other and with the + bf16 layout. What a caller can still get wrong is passing two pools that aren't + a pair, so the pair goes through ``check_two_pool_pair`` before the first launch + -- shared with the fused store, which has the same coupling. + """ + n_rows = kv.shape[0] if n_rows == 0: return @@ -94,20 +118,52 @@ def store_swa_into_unified( assert kv.is_contiguous() and kv.dtype == unified_kv.dtype assert state_slot.is_contiguous() and positions.is_contiguous() assert fp_arg.is_contiguous() - _swa_scatter_kernel[(n_rows,)]( - kv, - state_slot, - positions, - fp_arg, - unified_kv, - n_rows, - ring_stride, - win=win, - D=D, - HAS_FINAL=has_final, - BLOCK_D=triton.next_power_of_2(D), - num_warps=8, + two_pool = kv_rope is not None + assert two_pool == (unified_kv_rope is not None), ( + "kv_rope and unified_kv_rope come together" ) + if two_pool: + assert kv_rope.is_contiguous(), ( + f"kv_rope must be contiguous, got strides {kv_rope.stride()}" + ) + assert kv_rope.shape[0] == n_rows, ( + f"kv_rope holds {kv_rope.shape[0]} rows, kv holds {n_rows}" + ) + check_two_pool_pair( + unified_kv, + unified_kv_rope, + rope_width=kv_rope.shape[1], + rope_dtype=kv_rope.dtype, + ) + + def _scatter(src: torch.Tensor, dst: torch.Tensor) -> None: + D = src.shape[1] + assert dst.shape[1] == D, f"row width {D} does not fit pool {dst.shape[1]}" + _swa_scatter_kernel[(n_rows,)]( + src, + state_slot, + positions, + fp_arg, + dst, + n_rows, + ring_stride, + win=win, + D=D, + HAS_FINAL=has_final, + BLOCK_D=triton.next_power_of_2(D), + num_warps=8, + ) + + if kv.element_size() == 1: + # single-byte rows (any fp8 variant) are a pure byte move, and the E8M0 + # scale bytes aren't floats -- uint8 avoids a triton convert for `other=` + assert unified_kv.is_contiguous() + _scatter(kv.view(torch.uint8), unified_kv.view(torch.uint8)) + else: + _scatter(kv, unified_kv) + + if two_pool: + _scatter(kv_rope, unified_kv_rope) @triton.jit @@ -183,6 +239,116 @@ def decode( ) +@lru_cache(maxsize=None) +def decode_qo_indptr(num_tokens: int, device: torch.device) -> torch.Tensor: + """``qo_indptr`` for the two-pool decode: one q token per sequence. + + Not ``cu_seqlens_q`` -- that one is per-request and differs once MTP puts + several draft tokens in a batch. Cached unbounded like _token_identity_map: + all 61 layers ask for the same answer each step, and a captured graph holds + the address it got back. + """ + return torch.arange(num_tokens + 1, dtype=torch.int32, device=device) + + +# aiter sizes the split count off CU occupancy and over-splits just past 40 +# tokens, where the stage-2 merge starts to dominate. 4 rather than each shape's +# own optimum -- neighbouring split counts swing ~1.5x either way. +_DECODE_SPLIT_TAIL_MIN_TOKENS = 40 +_DECODE_SPLIT_TAIL_VALUE = 4 + + +def decode_fp8_2buff( + *, + q: torch.Tensor, # [T, H, nope_row_bytes] fp8 packed nope + inline e8m0 scale + q_rope: torch.Tensor, # [T, H, rope_dim] bf16 + unified_kv: torch.Tensor, # [rows, nope_row_bytes] fp8 + unified_kv_rope: torch.Tensor, # [rows, rope_dim] bf16 + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + attn_sink: torch.Tensor, # [H] fp32 + v_head_dim: int, + qo_indptr: Optional[torch.Tensor] = None, + num_kv_splits: Optional[int] = None, +) -> torch.Tensor: + """Decode over the two-pool fp8 unified_kv, through aiter's v4 nm asm kernel. + + Q arrives in the same packed form as the pool rows (nope fp8 + duplicated + e8m0 tile scales) with its rope half beside it in bf16, which is why this + can't share ``decode``'s single bf16 tensor. The kernel takes the row stride + off ``kv_buffer.size(-1)`` and only requires Q to match it, so the 512 B row + is not baked into the reader. + + ``v_head_dim`` is an element count (448 nope + 64 rope) that happens to equal + the row's byte width; it comes from the caller so that nothing here reads one + as the other. + """ + from aiter.mla import mla_decode_fwd_v4_nm + + T, H, row_bytes = q.shape + check_two_pool_pair( + unified_kv, + unified_kv_rope, + rope_width=q_rope.shape[-1], + rope_dtype=q_rope.dtype, + ) + assert row_bytes == unified_kv.shape[-1], ( + f"aiter derives the row stride from the kv pool ({unified_kv.shape[-1]} B) " + f"and reads Q with that same stride, but the q row is {row_bytes} B" + ) + assert q_rope.shape[:2] == (T, H), ( + f"q pair disagrees: packed {tuple(q.shape)[:2]} vs rope " + f"{tuple(q_rope.shape)[:2]}" + ) + # the asm kernel walks all four as flat buffers, it has no stride arguments + assert q.is_contiguous(), f"q must be contiguous, strides {q.stride()}" + assert q_rope.is_contiguous(), ( + f"q_rope must be contiguous, strides {q_rope.stride()}" + ) + assert attn_sink.dtype == torch.float32 and attn_sink.numel() == H, ( + f"sink must be {H} fp32 values, got {attn_sink.numel()} x {attn_sink.dtype}" + ) + + if qo_indptr is None: + qo_indptr = decode_qo_indptr(T, q.device) + # num_seqs comes from qo_indptr.numel()-1 and the kernel writes + # num_seqs * max_seqlen_q rows into `out`, so both have to be sized off q's + # own T. A qo_indptr built from a padded token count writes past `out`. + assert qo_indptr.shape[0] >= T + 1, ( + f"qo_indptr holds {qo_indptr.shape[0]} entries, kernel reads {T + 1}" + ) + assert kv_indptr.shape[0] >= T + 1, ( + f"kv_indptr holds {kv_indptr.shape[0]} entries, kernel reads {T + 1}" + ) + qo_indptr = qo_indptr[: T + 1] + + rows = unified_kv.shape[0] + out = q_rope.new_empty((T, H, v_head_dim)) + # Left None, the wrapper's occupancy heuristic picks it, folds the cross-split + # merge back into `out`, and leaves the final bf16 there whether or not it + # split. Pinning it to 1 costs 6.9x at bs=1 kv=2048. + if num_kv_splits is None and T > _DECODE_SPLIT_TAIL_MIN_TOKENS: + num_kv_splits = _DECODE_SPLIT_TAIL_VALUE + mla_decode_fwd_v4_nm( + q, + q_rope, + unified_kv.view(rows, 1, 1, row_bytes), + unified_kv_rope.view(rows, 1, 1, unified_kv_rope.shape[-1]), + out, + qo_indptr, + kv_indptr, + kv_indices, + 1, # max_seqlen_q; qo_indptr is per-token so every sequence is one token + sink=attn_sink, + num_kv_splits=num_kv_splits, + ) + # No empty-segment mask: a CG-padded row gets seq_len 1 on the ring slot + # ReqToTokenPool reserves, so the builders can't emit a zero-length one, and + # the compare + masked_fill_ was costing a launch per layer for it. One would + # come back NaN now (all-sink denominator); the guard UT pins that. + return out + + @triton.jit def _fill_compress_tail_kernel( indices_ptr, # [*] int32 (out) @@ -466,6 +632,115 @@ def build_prefill_indices( return kv_indices_prefix, kv_indptr_prefix, kv_indices_extend, kv_indptr_extend +def prefill_fp8_2buff( + *, + q: torch.Tensor, # [T, H, nope_row_bytes] fp8 packed nope + inline e8m0 scale + q_rope: torch.Tensor, # [T, H, rope_dim] bf16 + unified_kv: torch.Tensor, # [rows, nope_row_bytes] fp8 prefix pool + unified_kv_rope: torch.Tensor, # [rows, rope_dim] bf16 prefix pool + kv_indices_prefix: torch.Tensor, + kv_indptr_prefix: torch.Tensor, + kv_extend: torch.Tensor, # [tokens, nope_row_bytes] fp8 packed current chunk + kv_extend_rope: torch.Tensor, # [tokens, rope_dim] bf16 + kv_indices_extend: torch.Tensor, + kv_indptr_extend: torch.Tensor, + attn_sink: torch.Tensor, # [H] fp32 + softmax_scale: float, + v_head_dim: int, +) -> torch.Tensor: + """Prefill over the two-pool fp8 unified_kv, through aiter's opus kernel. + + Same two regions as ``prefill`` -- paged prefix plus this chunk's flat extend + -- but every latent arrives as a pair, so there are four buffers instead of + two. The extend pair is the packed K the fused norm+rope store hands back; + the ring write after attention consumes that same pair, which is why the + caller materialises it rather than this function quantizing here. + + Unlike ``decode_fp8_2buff`` the scale is a real argument: this kernel takes + it, so nothing has to match a hardcoded 1/sqrt(512). + + ``v_head_dim`` is an element count (448 nope + 64 rope) that happens to equal + the packed row's byte width; it comes from the caller so that nothing here + reads one as the other. + + A token with neither region comes back zero rather than NaN, so unlike + ``decode_fp8_2buff`` there is nothing to mask off the result afterwards. An + empty prefix is the live case here, not a guard: chunk 0 has nothing + committed yet and every token's prefix segment is empty. + """ + from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_fp8_opus + + T, H, row_bytes = q.shape + check_two_pool_pair( + unified_kv, + unified_kv_rope, + rope_width=q_rope.shape[-1], + rope_dtype=q_rope.dtype, + ) + # The kernel walks the prefix pool and the extend buffer with the same row + # layout, so a narrower extend row would read the next token's bytes as this + # one's scales instead of failing. + assert kv_extend.shape[-1] == row_bytes and kv_extend.dtype == unified_kv.dtype, ( + f"extend nope row is {kv_extend.shape[-1]} x {kv_extend.dtype}, pool is " + f"{row_bytes} x {unified_kv.dtype}" + ) + assert ( + kv_extend_rope.shape[-1] == unified_kv_rope.shape[-1] + and kv_extend_rope.dtype == unified_kv_rope.dtype + ), ( + f"extend rope row is {kv_extend_rope.shape[-1]} x {kv_extend_rope.dtype}, " + f"pool is {unified_kv_rope.shape[-1]} x {unified_kv_rope.dtype}" + ) + assert kv_extend.shape[0] == kv_extend_rope.shape[0], ( + f"extend pair disagrees: nope {kv_extend.shape[0]} rows vs rope " + f"{kv_extend_rope.shape[0]}" + ) + assert row_bytes == unified_kv.shape[-1], ( + f"the kernel reads Q with the kv row stride ({unified_kv.shape[-1]} B), " + f"but the q row is {row_bytes} B" + ) + assert q_rope.shape[:2] == (T, H), ( + f"q pair disagrees: packed {tuple(q.shape)[:2]} vs rope " + f"{tuple(q_rope.shape)[:2]}" + ) + # no stride arguments anywhere in the op. The two pools got their + # is_contiguous() from check_two_pool_pair above; these are the rest. + for name, t in ( + ("q", q), + ("q_rope", q_rope), + ("kv_extend", kv_extend), + ("kv_extend_rope", kv_extend_rope), + ): + assert t.is_contiguous(), f"{name} must be contiguous, strides {t.stride()}" + assert attn_sink.dtype == torch.float32 and attn_sink.numel() == H, ( + f"sink must be {H} fp32 values, got {attn_sink.numel()} x {attn_sink.dtype}" + ) + for name, indptr in ( + ("prefix", kv_indptr_prefix), + ("extend", kv_indptr_extend), + ): + assert indptr.shape[0] >= T + 1, ( + f"{name} indptr holds {indptr.shape[0]} entries, kernel reads {T + 1}" + ) + + out = q_rope.new_empty((T, H, v_head_dim)) + return pa_sparse_prefill_fp8_opus( + q, + q_rope, + unified_kv, + unified_kv_rope, + kv_indices_prefix, + kv_indptr_prefix[: T + 1], + kv_extend, + kv_extend_rope, + kv_indices_extend, + kv_indptr_extend[: T + 1], + attn_sink, + softmax_scale, + out=out, + ) + + def prefill( *, q: torch.Tensor, # [T, H, D] diff --git a/python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py b/python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py index 1624aafeb..f9bb63ccc 100644 --- a/python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py +++ b/python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py @@ -9,16 +9,34 @@ Grid: (cdiv(M, BLOCK_SIZE_M), num_local_heads + 1). pid_h == num_local_heads: KV program (norm + RoPE + FP8 quant nope + paged scatter) """ +from functools import lru_cache from typing import Optional import torch import triton import triton.language as tl +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.layout import ( + check_two_pool_pair, +) from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.utils import is_gfx95_supported _fp8_fnuz = is_fp8_fnuz() +# The two-pool fp8 store defers to aiter: its kernel already emits the exact +# 512 B nope row (448 fp8 + 14 dup e8m0 + pad) that the v4 asm attention reader +# expects, so the Triton kernel below stays bf16/legacy-packed only. +try: + from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_norm_rope_group_quant + + _HAS_AITER_OP = True +except ImportError: + fused_qk_norm_rope_group_quant = None + _HAS_AITER_OP = False + +_HAS_GROUP_QUANT = _HAS_AITER_OP and is_gfx95_supported() + # --------------------------------------------------------------------------- # Triton JIT helpers @@ -298,6 +316,137 @@ def _fused_qk_norm_rope_store_kernel( # --------------------------------------------------------------------------- +@lru_cache(maxsize=None) +def _token_identity_map(num_tokens: int, device: torch.device) -> torch.Tensor: + # Cached because every layer asks for it on every decode step and the answer + # only depends on the token count. Never evict: a captured cuda graph holds + # this address, and the default capture list has ~36 distinct batch sizes, so + # a bounded cache would free a live graph's buffer back into the graph pool. + return torch.arange(num_tokens, dtype=torch.int32, device=device) + + +def _fp8_2buff_store( + q: torch.Tensor, + kv: torch.Tensor, + q_norm_weight: Optional[torch.Tensor], + kv_norm_weight: torch.Tensor, + rms_eps: float, + rope_head_dim: int, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + positions: torch.Tensor, + q_out: torch.Tensor, + swa_cache: Optional[torch.Tensor], + swa_rope_cache: Optional[torch.Tensor], + swa_loc: Optional[torch.Tensor], + k_nope_out: Optional[torch.Tensor], + k_rope_out: Optional[torch.Tensor], + q_rope_out: Optional[torch.Tensor], +) -> torch.Tensor: + if not _HAS_GROUP_QUANT: + # is_unified_kv_fp8() already falls back to bf16 off gfx95, so reaching here + # on gfx95 means the installed aiter predates the op + raise RuntimeError( + "fp8 two-pool unified_kv needs aiter's fused_qk_norm_rope_group_quant: " + f"aiter exports it={_HAS_AITER_OP}, gfx95={is_gfx95_supported()}" + ) + + assert q.dim() == 2, ( + f"aiter takes q as [T, H, D] with no split-K reduce, got {tuple(q.shape)}" + ) + assert cos_cache.shape[-1] * 2 == rope_head_dim, ( + f"rot_dim from cos_cache ({cos_cache.shape[-1] * 2}) != {rope_head_dim}" + ) + # int64 is what the kernel indexes with; a silent .to() here would copy every + # call and hide a caller that changed dtype. + assert positions.dtype == torch.int64, ( + f"positions must be int64, got {positions.dtype}" + ) + + batch_id = None + has_swa = swa_cache is not None + if has_swa: + # Same coupling as store_swa_into_unified: one row index addresses both + # pools, and aiter bounds-checks neither -- a short rope pool aborts the + # process with nothing on stderr, so the pair is checked there first. + assert swa_loc is not None, "fp8 SWA store needs swa_loc alongside the pools" + # int32 is the SWA loc contract across the tree (translate_loc_from_full_to_swa + # enforces it too), and aiter reads the dest-row array raw -- a wider dtype + # becomes garbage row ids and aborts with nothing on stderr + assert swa_loc.dtype == torch.int32, ( + f"swa_loc must be int32, got {swa_loc.dtype}" + ) + assert swa_loc.shape[0] == kv.shape[0], ( + f"swa_loc holds {swa_loc.shape[0]} rows, kv holds {kv.shape[0]}" + ) + check_two_pool_pair( + swa_cache, + swa_rope_cache, + rope_width=rope_head_dim, + rope_dtype=kv.dtype, + ) + # aiter rejects the SWA write without a token->seq map even in dest-row + # mode, where all it does with it is drop tokens whose id is negative + # (CG pad). The ring row itself comes from swa_loc, and stale tokens are + # dropped on positions < 0, so identity is the map decode wants: one + # token per sequence, nothing masked. A caller with several tokens per + # sequence would have to pass its own. + batch_id = _token_identity_map(kv.shape[0], kv.device) + + # aiter has no stride arguments, so it can only write a packed q_out. With + # attn_tp_size > 1 the caller hands us a slice of a head-padded buffer + # ([T, 64, D] sliced to [T, n_local_heads, D]), which is strided unless the + # padding happened to be zero -- stage through a packed buffer then. + # contiguous_format is explicit: empty_like's default would copy q_out's + # strides for any input that is dense, and a strided staging buffer would + # silently misplace the heads. + q_dst = ( + q_out + if q_out.is_contiguous() + else torch.empty_like(q_out, memory_format=torch.contiguous_format) + ) + + # Q mirrors the K pair when the reader is the v4 nm asm kernel: nope fp8 with + # the tile scales inline, rotated PE beside it in bf16. A bf16 q_out instead + # keeps the whole rotated Q in one tensor, which is what the Triton + # sparse_attn reader takes. aiter picks between the two on the buffer's dtype + # alone, so the rope buffer has to be present exactly when q_out is fp8 -- + # otherwise it silently allocates one and the PE half goes nowhere. + assert (q_dst.element_size() == 1) == (q_rope_out is not None), ( + f"q_out is {q_dst.dtype} but q_rope_out is " + f"{None if q_rope_out is None else tuple(q_rope_out.shape)}" + ) + + # The packed K pair lands in the caller's buffers when it passed them (verify + # and prefill hand those to store_swa_into_unified); decode only wants the + # fused ring write, so it lets aiter allocate them and drops them. + q_packed, _, _, _ = fused_qk_norm_rope_group_quant( + q.view(q_dst.shape), + kv, + kv_norm_weight, + positions, + cos_cache, + sin_cache, + rms_eps, + is_neox=False, + q_nope_scale_buff=q_dst, + q_rope_buff=q_rope_out, + k_nope_scale_buff=k_nope_out, + k_rope_buff=k_rope_out, + q_weight=q_norm_weight, + quant_group_size=64, + scale_dtype="e8m0", + swa_nope_scale_buff=swa_cache, + swa_rope_buff=swa_rope_cache, + swa_dest_row=swa_loc, + batch_id_per_token=batch_id, + ) + if q_dst is not q_out: + q_out.copy_(q_packed) + return q_out + return q_packed + + def fused_qk_norm_rope_swa_store( q: torch.Tensor, kv: torch.Tensor, @@ -315,6 +464,11 @@ def fused_qk_norm_rope_swa_store( q_out: Optional[torch.Tensor] = None, dtype: torch.dtype = torch.bfloat16, bf16_store: bool = False, + fp8_2buff: bool = False, + swa_rope_cache: Optional[torch.Tensor] = None, + k_nope_out: Optional[torch.Tensor] = None, + k_rope_out: Optional[torch.Tensor] = None, + q_rope_out: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fused Q norm + KV norm + RoPE + optional SWA store. @@ -325,6 +479,16 @@ def fused_qk_norm_rope_swa_store( swa_loc: [M] int32 pre-translated paged indices swa_page_size: tokens per SWA page (default 128) bf16_store: write the whole head_dim as plain bf16 at swa_cache[swa_loc] + fp8_2buff: two-pool fp8 unified_kv. Delegates to aiter; ``swa_cache`` is + the fp8 nope pool and ``swa_rope_cache`` the bf16 rope pool, both + addressed by ``swa_loc``. Unlike the Triton path this leaves ``kv`` + untouched -- the normed + RoPE'd K comes back packed in + ``k_nope_out`` / ``k_rope_out`` when the caller supplies them. + q_rope_out: [M, num_local_heads, rope_head_dim] bf16, fp8_2buff only. + Present means Q is packed like K (nope fp8 + inline scales in + ``q_out``, rotated PE here) for the v4 nm asm reader; absent means + ``q_out`` holds the whole rotated Q in bf16 for the Triton reader. + ``q_out``'s dtype has to agree. """ head_dim = kv.shape[1] @@ -347,6 +511,30 @@ def fused_qk_norm_rope_swa_store( (M, num_local_heads, head_dim), dtype=dtype, device=q.device ) + if fp8_2buff: + assert not bf16_store, "fp8_2buff and bf16_store are different stores" + assert q_rms_eps == kv_rms_eps, ( + f"aiter norms Q and K with one eps, got {q_rms_eps} / {kv_rms_eps}" + ) + return _fp8_2buff_store( + q, + kv, + q_norm_weight, + kv_norm_weight, + kv_rms_eps, + rope_head_dim, + cos_cache, + sin_cache, + positions, + q_out, + swa_cache, + swa_rope_cache, + swa_loc, + k_nope_out, + k_rope_out, + q_rope_out, + ) + HAS_SWA_STORE = swa_cache is not None and swa_loc is not None dim_nope = 448 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 222ade3b6..1ebba7547 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1455,6 +1455,12 @@ class Envs: # Quantize the SWA fp8 KV cache from bf16-rounded values (matches # trainer-side QAT and the DSA-CP path) instead of fp32 registers. SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False) + # unified_kv only: split the pool into an fp8 nope pool plus a parallel + # bf16 rope pool, 640 B/token instead of 1024. The unified pool takes no + # dtype, so --kv-cache-dtype has no effect there and this switch is the + # only way to ask; on separate-KV it is the reverse -- --kv-cache-dtype + # picks the buffer dtype and this switch is inert. + SGLANG_DSV4_UNIFIED_KV_FP8 = EnvBool(False) # Kernels and indexer SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 8e04117c8..b38113b55 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -1366,8 +1366,17 @@ class DeepseekV4HipRadixBackend( attn_sink: torch.Tensor, core_attn_metadata: DSV4AttnMetadata, save_kv_cache: bool = True, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """unified_kv paged-attention path over the bf16 unified_kv""" + """unified_kv paged-attention path over the unified_kv pool. + + ``q_rope`` is what tells the two layouts apart: present means ``q`` is a + packed fp8 row and the pool is the two-pool fp8 one, so decode goes to + the asm reader; absent means both are plain bf16 and it goes to Triton. + Prefill needs ``k_rope`` alongside it, because there the current chunk is + a KV source of its own and not just something to store. + """ from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime pool = self.token_to_kv_pool @@ -1401,6 +1410,10 @@ class DeepseekV4HipRadixBackend( else: state_slot = forward_batch.req_pool_indices[:T] if save_kv_cache: + # Only verify reaches this under fp8 -- plain decode's rows are + # written by the fused kernel itself, which leaves kv None. The + # pair arrives already packed, so this is the same scatter with + # a second pool hanging off it. runtime.store_swa_into_unified( kv=kv, state_slot=state_slot, @@ -1409,6 +1422,10 @@ class DeepseekV4HipRadixBackend( win=win, ring_stride=ring_stride, final_pos=positions, + kv_rope=k_rope, + unified_kv_rope=( + None if k_rope is None else pool.get_unified_kv_rope(layer_id) + ), ) unified_metadata = core_attn_metadata.unified if compress_ratio == 0: @@ -1430,6 +1447,25 @@ class DeepseekV4HipRadixBackend( ) else: raise ValueError(f"bad compress_ratio {compress_ratio}") + if q_rope is not None: + # softmax_scale is not passed on: the asm kernel hardcodes + # 1/sqrt(512), which is what self.softmax_scale already is for + # V4's head_dim=512. The other readers here take it explicitly, + # so a head_dim change would leave only this one mis-scaled. + assert self.softmax_scale == 512**-0.5, ( + "the v4 nm asm kernel hardcodes 1/sqrt(512), this backend is " + f"at {self.softmax_scale}" + ) + return runtime.decode_fp8_2buff( + q=q, + q_rope=q_rope, + unified_kv=unified, + unified_kv_rope=pool.get_unified_kv_rope(layer_id), + kv_indices=kv_indices, + kv_indptr=kv_indptr, + attn_sink=attn_sink, + v_head_dim=layer.v_head_dim, + ) return runtime.decode( q=q, unified_kv=unified, @@ -1505,17 +1541,42 @@ class DeepseekV4HipRadixBackend( pad = T + 1 - kpre_p.shape[0] kpre_p = torch.cat([kpre_p, kpre_p[-1:].expand(pad)]) kext_p = torch.cat([kext_p, kext_p[-1:].expand(pad)]) - o = runtime.prefill( - q=q, - unified_kv=unified, - kv_indices_prefix=kpre_i, - kv_indptr_prefix=kpre_p, - kv_extend=kv, - kv_indices_extend=kext_i, - kv_indptr_extend=kext_p, - attn_sink=attn_sink, - softmax_scale=self.softmax_scale, - ) + if q_rope is not None: + assert k_rope is not None, ( + "fp8 prefill needs the extend rope half beside the packed nope; " + "q_rope came through but k_rope did not" + ) + # No empty-segment mask on the result, unlike decode: this kernel + # returns zeros for a token with neither region where the asm decode + # reader leaves the row NaN. Chunk 0 tokens have an empty prefix and + # a non-empty extend, which both readers handle. + o = runtime.prefill_fp8_2buff( + q=q, + q_rope=q_rope, + unified_kv=unified, + unified_kv_rope=pool.get_unified_kv_rope(layer_id), + kv_indices_prefix=kpre_i, + kv_indptr_prefix=kpre_p, + kv_extend=kv, + kv_extend_rope=k_rope, + kv_indices_extend=kext_i, + kv_indptr_extend=kext_p, + attn_sink=attn_sink, + softmax_scale=self.softmax_scale, + v_head_dim=layer.v_head_dim, + ) + else: + o = runtime.prefill( + q=q, + unified_kv=unified, + kv_indices_prefix=kpre_i, + kv_indptr_prefix=kpre_p, + kv_extend=kv, + kv_indices_extend=kext_i, + kv_indptr_extend=kext_p, + attn_sink=attn_sink, + softmax_scale=self.softmax_scale, + ) # write this chunk's SWA K into the ring for future chunks / decode # only the final-window tokens per request @@ -1535,6 +1596,10 @@ class DeepseekV4HipRadixBackend( win=win, ring_stride=ring_stride, final_pos=_ring_final_pos, + kv_rope=None if k_rope is None else k_rope[:n_real], + unified_kv_rope=( + None if k_rope is None else pool.get_unified_kv_rope(layer_id) + ), ) return o @@ -1602,6 +1667,8 @@ class DeepseekV4HipRadixBackend( compress_ratio: Literal[0, 4, 128], save_kv_cache: bool = True, attn_sink: Optional[torch.Tensor] = None, + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, **_, ) -> torch.Tensor: if self.mtp_enabled and forward_batch.forward_mode.is_idle(): @@ -1630,6 +1697,8 @@ class DeepseekV4HipRadixBackend( attn_sink=attn_sink, core_attn_metadata=core_attn_metadata, save_kv_cache=save_kv_cache, + q_rope=q_rope, + k_rope=k_rope, ) if isinstance(core_attn_metadata, DSV4AttnMetadata): diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index b7e2040e9..f369cc50f 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -158,6 +158,8 @@ class CompressorBackendMixin: bf16_store: bool = False, kv_scale_cache: Optional[torch.Tensor] = None, rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + fp8_2buff: bool = False, + kv_cache_rope: Optional[torch.Tensor] = None, ) -> None: assert compress_ratio == 4 or compress_ratio == 128 assert rotate == is_indexer == (head_dim == 128) @@ -220,6 +222,8 @@ class CompressorBackendMixin: if _is_hip and use_fp4_indexer else None ), + fp8_2buff=fp8_2buff, + kvcache_rope=kv_cache_rope, ) def forward_unified( @@ -238,6 +242,7 @@ class CompressorBackendMixin: state_pool = compressor.get_state_pool(self) from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, is_unified_kv_triton, ) @@ -264,6 +269,8 @@ class CompressorBackendMixin: use_hip_fp4 = _is_hip and use_fp4_indexer bf16_store = False kv_scale_cache = None + fp8_2buff = False + kv_cache_rope = None if compressor.is_in_indexer: page_size = token_to_kv_pool.get_index_k_page_size(compressor.ratio) if use_hip_fp4: @@ -278,7 +285,11 @@ class CompressorBackendMixin: self.forward_metadata.core_metadata.unified, f"c{compressor.ratio}_out_loc", ) - bf16_store = True + if is_unified_kv_fp8(): + fp8_2buff = True + kv_cache_rope = token_to_kv_pool.get_unified_kv_rope(layer_id) + else: + bf16_store = True else: _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id] assert compress_kv_pool is not None @@ -305,6 +316,10 @@ class CompressorBackendMixin: rope_cache=( (compressor.fp4_cos, compressor.fp4_sin) if use_hip_fp4 else None ), + fp8_2buff=fp8_2buff, + kv_cache_rope=( + None if kv_cache_rope is None else kv_cache_rope.view(dtype=torch.uint8) + ), ) online_c128_mtp = getattr(self, "online_c128_mtp", None) if online_c128_mtp is not None: diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index c96085db2..deaeb52a9 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -16,6 +16,7 @@ from sglang.kernels.ops.attention.dsv4 import ( index_buf_accessor as dsv4_index_buf_accessor, ) from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack +from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import layout from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool @@ -511,11 +512,43 @@ class DeepSeekV4LayerItem(NamedTuple): compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None +# re-exported: the pool allocates the rows, but the kernels that write them own the +# layout (see unified_kv_kernels/layout.py) +DSV4_FP8_NOPE_ROW_BYTES = layout.DSV4_FP8_NOPE_ROW_BYTES +DSV4_FP8_QUANT_TILE = layout.DSV4_FP8_QUANT_TILE + + +def dsv4_unified_row_bytes( + qk_nope_head_dim: int, qk_rope_head_dim: int, fp8: bool +) -> int: + """Bytes one unified_kv token occupies, summed over both pools.""" + if not fp8: + return (qk_nope_head_dim + qk_rope_head_dim) * 2 + num_tiles = -(-qk_nope_head_dim // DSV4_FP8_QUANT_TILE) + scale_bytes = 2 * num_tiles + # not an assert: sizing runs under -O too, and a silently skipped check here + # overreports capacity + if qk_nope_head_dim + scale_bytes > DSV4_FP8_NOPE_ROW_BYTES: + raise ValueError( + f"fp8 nope row overflows: {qk_nope_head_dim} latent values at 1 B + " + f"{scale_bytes} B scales > {DSV4_FP8_NOPE_ROW_BYTES} B stride" + ) + return DSV4_FP8_NOPE_ROW_BYTES + qk_rope_head_dim * 2 + + # The following kv pool follows ATOM's unified_kv kernel layout. class DeepSeekV4UnifiedKVPool: """ - Layout: + Layout (bf16): unified_kv[L]: ``[swa_pages + padded_compress_rows, head_dim]`` bf16 + + Layout (fp8, ``SGLANG_DSV4_UNIFIED_KV_FP8``) -- two parallel pools with the + same row count, so a row index means the same thing in both. Named after the + accessors, which under fp8 each return one half -- ``get_unified_kv`` the + nope, ``get_unified_kv_rope`` the rope: + unified_kv[L] (nope): ``[rows, 512]`` fp8, see DSV4_FP8_NOPE_ROW_BYTES + unified_kv_rope[L] (rope): ``[rows, qk_rope_head_dim]`` bf16, never quantized + - rows ``[0, swa_pages)`` = SWA ring (``req_pool_indices * swa_window + pos % swa_window``) - rows ``[swa_pages, ...)`` = compressed (``swa_pages + page_index``) """ @@ -535,8 +568,11 @@ class DeepSeekV4UnifiedKVPool: memory_saver_adapter, custom_mem_pool, swa_ring_size: int, + fp8: bool = False, ): self.swa_ring_size = swa_ring_size + self.fp8 = fp8 + self.rope_head_dim = qk_rope_head_dim self.head_dim = qk_nope_head_dim + qk_rope_head_dim self.num_slots = num_slots self.swa_pages = num_slots * self.swa_ring_size @@ -545,6 +581,7 @@ class DeepSeekV4UnifiedKVPool: self.k_per_block = dict(self.K_PER_BLOCK) bufs = [] + rope_bufs = [] with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( torch.cuda.use_mem_pool(custom_mem_pool) @@ -557,20 +594,54 @@ class DeepSeekV4UnifiedKVPool: compress_rows = self.num_blocks * self.k_per_block[ratio] rows_per_page = self.page_size // ratio if ratio else 0 padded_compress_rows = compress_rows + rows_per_page - bufs.append( - torch.zeros( - self.swa_pages + padded_compress_rows, - self.head_dim, - dtype=torch.bfloat16, - device=device, + rows = self.swa_pages + padded_compress_rows + if self.fp8: + bufs.append( + torch.zeros( + rows, + DSV4_FP8_NOPE_ROW_BYTES, + dtype=torch.float8_e4m3fn, + device=device, + ) ) - ) + rope_bufs.append( + torch.zeros( + rows, + self.rope_head_dim, + dtype=torch.bfloat16, + device=device, + ) + ) + else: + bufs.append( + torch.zeros( + rows, + self.head_dim, + dtype=torch.bfloat16, + device=device, + ) + ) + rope_bufs.append(None) self.kv_buffer = bufs + self.kv_buffer_rope = rope_bufs def get_unified_kv(self, local_layer_id: int) -> torch.Tensor: return self.kv_buffer[local_layer_id] + def get_unified_kv_rope(self, local_layer_id: int) -> torch.Tensor: + assert self.fp8, "rope pool only exists under SGLANG_DSV4_UNIFIED_KV_FP8" + return self.kv_buffer_rope[local_layer_id] + def get_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: + if self.fp8: + # same single-pool assumption as the outer get_contiguous_buf_infos: + # one pointer and one row size per layer describes the nope pool only, + # so whoever picks this up next would move half a row and not notice. + # TODO(danli103): report both pools once a consumer needs them. + raise NotImplementedError( + "get_buf_infos describes one pool per layer; the fp8 rope pool " + "would be dropped (SGLANG_DSV4_UNIFIED_KV_FP8=1)." + ) data_ptrs = [b.data_ptr() for b in self.kv_buffer] data_lens = [b.nbytes for b in self.kv_buffer] item_lens = [b[0].nbytes for b in self.kv_buffer] @@ -578,6 +649,10 @@ class DeepSeekV4UnifiedKVPool: class DeepSeekV4TokenToKVPool(BaseSWAKVPool): + # object.__new__ stubs (disagg wire test) skip __init__; False is the env + # default, so the fp8 PD/HiCache refuses don't AttributeError on them. + _unified_kv_fp8 = False + def __init__( self, max_num_reqs: int, @@ -633,11 +708,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.c4_logical_size = c4_logical_size self.c128_size = c128_size from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, is_unified_kv_triton, ) # Resolve the unified-kv gate before any sizing so the two cannot drift. self._unified_kv = is_unified_kv_triton() + self._unified_kv_fp8 = is_unified_kv_fp8() # Uniform 512-dim e4m3 layout for the trtllm attention backend self.uniform_fp8 = ( not self._unified_kv @@ -721,6 +798,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): memory_saver_adapter=self.memory_saver_adapter, custom_mem_pool=self.custom_mem_pool, swa_ring_size=swa_ring_size, + fp8=self._unified_kv_fp8, ) self.unified_swa_window = self.sliding_window @@ -766,6 +844,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.wait_layer_transfer(layer_id) return self.unified_kv_pool.get_unified_kv(layer_id - self._stage_start) + def get_unified_kv_rope(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) + return self.unified_kv_pool.get_unified_kv_rope(layer_id - self._stage_start) + def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping @@ -782,6 +864,18 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): data_lens: List[int] = [] item_lens: List[int] = [] + if self._unified_kv_fp8: + # The page-block transfer below prices one row as buf[0].nbytes and + # ships a single pointer per layer. Under fp8 that covers the nope + # pool only -- the parallel bf16 rope pool would be dropped and the + # remote side would decode rows against stale rope. Refuse instead. + # TODO(danli103): ship the rope pool as a second per-layer entry. + raise NotImplementedError( + "PD disaggregation is not supported with " + "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the transfer assumes a single " + "unified pool; the rope pool would be silently dropped)." + ) + def append_page_buffer(buf: torch.Tensor) -> None: assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D" data_ptrs.append(buf.data_ptr()) @@ -827,6 +921,15 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): item_lens: List[int] = [] if not self._unified_kv: return data_ptrs, data_lens, item_lens + if self._unified_kv_fp8: + # Other half of the PD path -- get_contiguous_buf_infos ships the + # compressed region, this one the ring. Same single-pool assumption, + # same silently dropped rope, same fix -- land them together. + raise NotImplementedError( + "PD disaggregation is not supported with " + "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the SWA_RING component assumes a " + "single unified pool; the rope pool would be silently dropped)." + ) swa_pages = self.unified_kv_pool.swa_pages for buf in self.unified_kv_pool.kv_buffer: assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D" @@ -841,6 +944,17 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): # the unified pool stores individual token rows after its SWA region. assert self._unified_kv, "unified_region_buffers requires unified_kv layout" assert ratio in (4, 128), f"unsupported compression ratio: {ratio}" + if self._unified_kv_fp8: + # item_bytes below prices kv_buffer alone, so the rope pool would never + # be offloaded and a fetched page would carry stale rope -- wrong output, + # no crash. + # TODO(danli103): give rope its own host pool, the way C4_INDEXER + # already parallels C4. + raise NotImplementedError( + "HiCache offload is not supported with " + "SGLANG_DSV4_UNIFIED_KV_FP8=1 (the host pool assumes a single " + "unified pool; the rope pool would never be offloaded)." + ) swa_pages = self.unified_kv_pool.swa_pages head_dim = self.unified_kv_pool.head_dim diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 8ec857a6a..e3b08c0a3 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -986,16 +986,56 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4) self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128) + # Unified-KV uses a different physical layout than the non-unified V4 path: + # * one row carries the full latent -- 1024 B bf16, or 640 B under + # SGLANG_DSV4_UNIFIED_KV_FP8 (512 B fp8 nope + 128 B bf16 rope) -- not + # that path's 584-byte fp8(nope) + bf16(rope) + scales cell. + # * SWA is a fixed per-request ring (num_req_slots * ring_size), + # independent of full_token, so it is a fixed *bias* rather than a + # per-token term. Gate on the same switch the pool itself uses so the + # sizing and the allocation never drift apart. from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, is_unified_kv_triton, ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + dsv4_unified_row_bytes, + ) self._unified = is_unified_kv_triton() + self._unified_fp8 = is_unified_kv_fp8() self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + # Row width across both pools: 1024 B bf16, 640 B fp8. Read from the pool + # module so sizing can't drift from the allocation. + self._unified_row_bytes = dsv4_unified_row_bytes( + self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8 + ) # swa_page_size is the model's sliding window (cfg.window_size). self._swa_ring_size = get_swa_ring_size(self.swa_page_size, self.is_speculative) self._spec_infl = 1.0 + # The unified pool takes no dtype, so --kv-cache-dtype never reaches it. + # V4 defaults "auto" to fp8_e4m3 (overrides.py + # _deepseek_v4_kv_cache_dtype), so only a bfloat16 here tells us the user + # set it explicitly; warning on the fp8 side would fire on every run. + if self._unified_fp8 and self.kv_cache_dtype_str == "bfloat16": + logger.warning( + "--kv-cache-dtype=bfloat16 is ignored on the unified_kv path; " + "SGLANG_DSV4_UNIFIED_KV_FP8=1 stores the latent as fp8. Unset the " + "env switch to get a bf16 unified pool." + ) + + # get_contiguous_buf_infos ships one pointer per layer and prices a row as + # buf[0].nbytes, which under fp8 covers the nope pool only. Fail at startup + # rather than at the first transfer. + # TODO(danli103): drop this once the transfer ships the rope pool. + if self._unified_fp8 and self.disaggregation_mode != "null": + raise ValueError( + "SGLANG_DSV4_UNIFIED_KV_FP8=1 does not support PD disaggregation " + f"(disaggregation_mode={self.disaggregation_mode!r}). Unset the fp8 " + "switch or run without disaggregation." + ) + if self.is_speculative: # Ring is sized once here, so it must serve the largest adaptive tier. self._assert_ring_serves_draft_tokens( @@ -1065,8 +1105,10 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def _get_bytes_per_full_token(self) -> float: if self._unified: - # Unified_kv stores the whole latent in bf16. - kv_bytes = self.attn_head_dim * 2 + # Unified_kv stores the whole latent: one bf16 pool, or an fp8 nope + # pool plus a bf16 rope pool. kv_bytes also prices the compressed + # c4/c128 rows below, which live in the same pool(s). + kv_bytes = self._unified_row_bytes else: kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 @@ -1199,14 +1241,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): return min(estimated, full_token // 2) def _fixed_swa_bytes(self, max_running_requests: int) -> int: + """Unified_kv SWA is a fixed per-request ring, sized by concurrency + (num_req_slots) rather than by full_token. Return its byte footprint + across all full layers, inflated for the draft worker the same way as the + per-token coeff. Returns 0 on the non-unified path (where SWA is already + accounted per-token).""" if not self._unified: return 0 num_req_slots = self._get_num_req_slots(max_running_requests) ring_bytes = ( num_req_slots * self._swa_ring_size - * self.attn_head_dim - * 2 # bf16 + * self._unified_row_bytes * self.num_layers_total ) return int(ring_bytes * self._spec_infl) @@ -1277,6 +1323,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): sizes = self._compute_dsv4_sizes(full_token, page_size) logger.info( f"DSV4 memory calculation: unified={self._unified}, " + f"unified_fp8={self._unified_fp8}, " f"bytes_per_full_token={self.bytes_per_full_token:.2f}, " f"available_bytes={available_bytes / (1 << 30):.2f} GB, " f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, " diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9af352e22..af76c73f7 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1423,6 +1423,9 @@ class MQALayer(MqaAttentionBase): attn_backend, q_out: Optional[torch.Tensor] = None, x_quant=None, + q_rope_out: Optional[torch.Tensor] = None, + k_nope_out: Optional[torch.Tensor] = None, + k_rope_out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: x_linear = x_quant if x_quant is not None else x @@ -1437,22 +1440,52 @@ class MQALayer(MqaAttentionBase): kv: Optional[torch.Tensor] from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, is_unified_kv_triton, ) unified = is_unified_kv_triton() + fp8_2buff = is_unified_kv_fp8() is_decode = forward_batch.forward_mode.is_decode_or_idle() # The kernel is token-indexed (q, kv and positions are all length M), so # a verify batch carrying several draft tokens per request is a shape it # already handles. Only the cache store differs between decode and - # verify, and that half is left off below. + # verify, and under fp8 that store takes the packed pair instead of bf16. fuse_verify = ( envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get() and forward_batch.forward_mode.is_target_verify() ) - do_fused_qk_norm_rope = (unified and (is_decode or fuse_verify)) or ( - not unified and self.use_fused_qk_norm_rope + # fp8 verify packs like prefill but keeps verify's store timing: the pair + # lands in the caller's buffers and the backend writes the ring off the + # per-token slot map before attention. Keyed off those buffers the same + # way fuse_prefill is, so the two arms cannot disagree about the layout. + fuse_verify_fp8 = ( + fuse_verify + and unified + and fp8_2buff + and k_nope_out is not None + and k_rope_out is not None ) + # Prefill under fp8 goes through the same fused store: the 2-source + # kernel reads this chunk as its extend region in the pool's packed form, + # and the ring write after attention reuses those same rows, so they are + # materialised once here rather than quantized on both sides. Keyed off + # the caller's buffers the way q_rope_out keys the packed Q, so the two + # cannot disagree about the layout; both halves are required because the + # nope one leaves on the kv slot and a missing one would read as "the + # fused store did not run". Verify packs the same way but is its own arm + # above: it stores before attention, not after. + fuse_prefill = ( + unified + and fp8_2buff + and k_nope_out is not None + and k_rope_out is not None + and not is_decode + and not forward_batch.forward_mode.is_target_verify() + ) + do_fused_qk_norm_rope = ( + unified and (is_decode or fuse_verify or fuse_prefill) + ) or (not unified and self.use_fused_qk_norm_rope) if do_fused_qk_norm_rope: if _is_gfx95_supported or _is_gfx1250_supported: @@ -1473,6 +1506,7 @@ class MQALayer(MqaAttentionBase): ) token_to_kv_pool = get_token_to_kv_pool() + swa_rope_cache = None if unified and fuse_verify: # Target-verify runs through the unified_kv decode path. The # backend writes the current chunk's KV into the ring *before* @@ -1490,15 +1524,34 @@ class MQALayer(MqaAttentionBase): # contiguous buffer, so materialise it before the kernel norms # it in place. The unfused path pays the same copy inside # _compute_kv_bf16. + # + # Under fp8 the kernel writes the packed pair to the caller's + # buffers rather than norming kv in place, and the same backend + # store takes that pair -- only the row format changes. kv = kv.contiguous() swa_cache, swa_loc = None, None - swa_page_size, bf16_store = 1, True + swa_page_size, bf16_store = 1, not fuse_verify_fp8 + elif unified and fuse_prefill: + # No pools, so the kernel norms + RoPEs + packs and writes no + # ring row. It must not: those rows are this fwd's extend region + # and the prefix pool has to stay as attention expects to find + # it. The backend stores them after attention from the pair. + swa_cache, swa_loc = None, None + swa_page_size, bf16_store = 1, False + # kv stays the strided slice of qkv_a. Under fp8 the kernel only + # reads it -- the packed pair goes to k_nope_out/k_rope_out, it + # does not norm in place -- and it takes the row stride as an + # argument, so materialising it was a copy on every fp8 layer. elif unified: swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id) # swa_loc is layer-independent; computed once per forward by the # backend and cached on the metadata (read here by every layer). swa_loc = attn_backend.get_unified_swa_loc(forward_batch) - swa_page_size, bf16_store = 1, True + swa_page_size, bf16_store = 1, not fp8_2buff + if fp8_2buff: + swa_rope_cache = token_to_kv_pool.get_unified_kv_rope(self.layer_id) + # kv stays the strided slice of qkv_a -- the group-quant + # kernel takes the row stride as an argument. else: swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id) swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch) @@ -1528,13 +1581,25 @@ class MQALayer(MqaAttentionBase): q_out=q_out, dtype=x.dtype, bf16_store=bf16_store, + fp8_2buff=fp8_2buff, + swa_rope_cache=swa_rope_cache, + k_nope_out=k_nope_out if (fuse_prefill or fuse_verify_fp8) else None, + k_rope_out=k_rope_out if (fuse_prefill or fuse_verify_fp8) else None, + q_rope_out=q_rope_out, ) # On the verify path the kernel normed + RoPE'd kv in place and wrote # nothing, so hand it back: the caller feeds it to attention as the # current chunk (attn_k = kv) and save_kv_cache = kv is not None lets # the backend do its normal causally-indexed store into the ring # before the decode kernel runs -- exactly as the unfused path did. - if not (unified and fuse_verify): + if unified and (fuse_prefill or fuse_verify_fp8): + # The packed nope half rides out on the kv slot -- attention + # takes it as attn_k and save_kv_cache stays on so the backend + # does the ring write. Its rope half went to the caller's buffer, + # which has no second return slot here. Prefill's write lands + # after attention, verify's before it; both read this pair. + kv = k_nope_out + elif not (unified and fuse_verify): kv = None if not unified and use_cp: @@ -1657,21 +1722,97 @@ class MQALayer(MqaAttentionBase): and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() ) - tp_slice, q_padded, q_out = slice(None), None, None - kernel_num_heads = self._kernel_num_heads(x.shape[0]) - if kernel_num_heads != self.n_local_heads: - # Backends without an exact-head specialization retain the existing - # padded shape. attn_sink is sliced to this rank and padded to match. - # Only [0:n_local_heads] is written below. Uninitialized padded TP - # heads inject NaN into attention on gfx942 (fnuz), so zero-init - # there; other archs tolerate new_empty and skip the per-forward - # memset. - if _is_gfx942_supported: - q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim) - else: - q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim) - tp_slice = slice(0, self.n_local_heads) - q_out = q_padded[:, tp_slice, :] + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, + is_unified_kv_triton, + ) + + unified = is_unified_kv_triton() + unified_fp8_verify = ( + unified + and is_unified_kv_fp8() + and forward_batch.forward_mode.is_target_verify() + ) + # The v4 nm asm reader takes Q in the pool's own packed form, so fp8 + # decode wants a contiguous fp8 buffer of exactly the local heads -- + # q_padded below is a FlashMLA layout and buys nothing here. Verify runs + # that same reader over the ring, so it takes the same Q. + unified_fp8_decode = ( + unified + and is_unified_kv_fp8() + and (forward_batch.forward_mode.is_decode_or_idle() or unified_fp8_verify) + ) + # The 2-source prefill kernel wants the same packed Q plus this chunk's + # K in the pool's layout. Verify is not prefill here even though it takes + # the same branch below -- it reads rows the ring already holds, so it + # goes with decode above. Multi-stream picks a different prepare that has + # no unified arm at all, so it keeps the bf16 buffers it always had. + unified_fp8_prefill = ( + unified + and is_unified_kv_fp8() + and not enable_multi_stream + and not forward_batch.forward_mode.is_decode_or_idle() + and not forward_batch.forward_mode.is_target_verify() + ) + if unified_fp8_verify and not envs.SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY.get(): + # The packed pair is produced by the fused norm+RoPE store; with that + # off the unfused arm hands the backend bf16 kv and the ring scatter + # dies on a dtype assert that says nothing about MTP. + raise NotImplementedError( + "fp8 two-pool unified_kv needs the fused verify store for " + "speculative decoding: set " + "SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY=1, or run with " + "SGLANG_DSV4_UNIFIED_KV_FP8=0." + ) + if ( + unified + and is_unified_kv_fp8() + and self.dsa_enable_prefill_cp + and dsa_use_prefill_cp(forward_batch) + and not forward_batch.forward_mode.is_decode_or_idle() + ): + # The gather hands back bf16 kv in global token order *after* + # norm+RoPE, so packing would have to move ahead of it and re-derive + # RoPE from global-order positions. Whether the CP path has those + # ready is unverified, so refuse instead of packing the wrong order. + raise NotImplementedError( + "fp8 two-pool unified_kv does not support DSA prefill CP " + "(SGLANG_DSV4_UNIFIED_KV_FP8=1 with cp_size > 1)." + ) + + tp_slice, q_padded, q_out, q_rope = slice(None), None, None, None + k_nope, k_rope = None, None + if unified_fp8_decode or unified_fp8_prefill: + # width and dtype come off the pools themselves; the kernel reads Q + # with the kv row stride, so the two must not drift + kv_pool = get_token_to_kv_pool() + nope_pool = kv_pool.get_unified_kv(self.layer_id) + rope_pool = kv_pool.get_unified_kv_rope(self.layer_id) + q_out = nope_pool.new_empty( + (x.shape[0], self.n_local_heads, nope_pool.shape[-1]) + ) + q_rope = rope_pool.new_empty( + (x.shape[0], self.n_local_heads, rope_pool.shape[-1]) + ) + if unified_fp8_prefill or unified_fp8_verify: + k_nope = nope_pool.new_empty((x.shape[0], nope_pool.shape[-1])) + k_rope = rope_pool.new_empty((x.shape[0], rope_pool.shape[-1])) + kernel_num_heads = self.n_local_heads + else: + kernel_num_heads = self._kernel_num_heads(x.shape[0]) + if kernel_num_heads != self.n_local_heads: + # Backends without an exact-head specialization retain the existing + # padded shape. attn_sink is sliced to this rank and padded to match. + # Only [0:n_local_heads] is written below. Uninitialized padded TP + # heads inject NaN into attention on gfx942 (fnuz), so zero-init + # there; other archs tolerate new_empty and skip the per-forward + # memset. + if _is_gfx942_supported: + q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim) + else: + q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim) + tp_slice = slice(0, self.n_local_heads) + q_out = q_padded[:, tp_slice, :] attn_sink = self._local_attn_sink(kernel_num_heads) if enable_multi_stream: @@ -1713,6 +1854,9 @@ class MQALayer(MqaAttentionBase): attn_backend, q_out, x_quant=x_quant, + q_rope_out=q_rope, + k_nope_out=k_nope, + k_rope_out=k_rope, ) # save_kv_cache = kv is not None selects who writes the ring. When kv is @@ -1723,11 +1867,16 @@ class MQALayer(MqaAttentionBase): # _forward_prepare* deliberately left the store off and the backend does # its normal causally-indexed store from attn_k = kv. attn_k = kv if kv is not None else q - from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_triton, - ) - if is_unified_kv_triton(): + if unified: + # only the HIP radix backend takes these two; passing them always would + # leave non-ROCm depending on the **_ in its forward() to drop them, and + # no test on that side would notice if the **_ went away + rope_kwargs = {} + if q_rope is not None: + rope_kwargs["q_rope"] = q_rope + if k_rope is not None: + rope_kwargs["k_rope"] = k_rope o = attn_backend.forward( q=q_out if q_out is not None else q, k=attn_k, @@ -1737,6 +1886,7 @@ class MQALayer(MqaAttentionBase): compress_ratio=self.compress_ratio, attn_sink=attn_sink[: self.n_local_heads], save_kv_cache=kv is not None, + **rope_kwargs, ) else: attn_q = q_padded if q_padded is not None else q diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py new file mode 100644 index 000000000..eb294a389 --- /dev/null +++ b/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() diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py new file mode 100644 index 000000000..51c38390c --- /dev/null +++ b/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() diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py new file mode 100644 index 000000000..d47345d62 --- /dev/null +++ b/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() diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py new file mode 100644 index 000000000..0b33a1c83 --- /dev/null +++ b/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() diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py new file mode 100644 index 000000000..c74b1045f --- /dev/null +++ b/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() diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py new file mode 100644 index 000000000..7e82d66d3 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py b/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py new file mode 100644 index 000000000..f9ddfb49f --- /dev/null +++ b/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() diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 5a6b53292..2133fadaa 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/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 diff --git a/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py b/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py new file mode 100644 index 000000000..013763eb0 --- /dev/null +++ b/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()