minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
amd-danli103
2026-09-14 02:49:11 -07:00
committed by GitHub
parent 95140a7b0c
commit 5aa9b8fb3e
21 changed files with 3594 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+145 -33
View File
@@ -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 <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
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<kUsePDL>();
@@ -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<bf16x2_t>(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<bf16x2_t*>(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<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
// All lanes in this warp produce the same scale byte; let lane 0 publish.
if (lane_id == 0) static_cast<uint8_t*>(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<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
}
}
}
}
@@ -541,12 +571,61 @@ struct FusedNormRopeKernel {
}
}
template <ForwardMode kMode>
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<DType, kMode, kLogPageSize, kUsePDL, false, true>;
}
template <ForwardMode kMode>
static constexpr auto select_fp4_kernel() {
static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer");
return fused_norm_rope_indexer_fp4<DType, kMode, kLogPageSize, kUsePDL>;
}
// 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<DType>().with_device(device_).verify(input);
TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight);
TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
TensorMatcher({-1}).with_dtype<int64_t>().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<kDLGPU>();
TensorMatcher({N, kHeadDim}) // input
.with_dtype<DType>()
.with_device(device_)
.verify(input);
TensorMatcher({kHeadDim}) // weight
.with_dtype<DType>()
.with_device(device_)
.verify(weight);
TensorMatcher({-1, kRopeDim}) // freqs_cis
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
TensorMatcher({-1}) // out_loc
.with_dtype<int64_t>()
.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<uint8_t>()
.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<uint32_t>(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<ForwardMode>(is_decode);
auto N = SymbolicSize{"num_tokens"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLGPU>();
verify_operands(input, weight, freqs_cis, out_loc, N, device_);
TensorMatcher({-1, -1}).with_strides({kFp8PageBytes, 1}).with_dtype<uint8_t>().with_device(device_).verify(kvcache);
TensorMatcher({-1, kRopeRowBytes})
.with_strides({kRopeRowBytes, 1})
.with_dtype<uint8_t>()
.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<int64_t>(kPageSize));
verify_plan_for_mode(mode, plan, out_loc, N, device_);
const auto num_tokens = static_cast<uint32_t>(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<const float*>(freqs_cis.data_ptr()),
.out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
.kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
.kvcache_rope = static_cast<uint8_t*>(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<CompressExtend>() : select_fp8_2buff_kernel<CompressDecode>();
LaunchKernel(num_tokens, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
}
static void forward_fp4(
const tvm::ffi::TensorView input,
const tvm::ffi::TensorView plan,
python/sglang/kernels/ops/attention/dsv4/compress.py
+28 -2
View File
@@ -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,
)
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/env_gate.py
+24 -1
View File
@@ -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
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/layout.py
+62
View File
@@ -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()}"
)
python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/runtime.py
+293 -18
View File
@@ -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]
python/sglang/kernels/ops/attention/fused_qk_norm_rope_store.py
+188
View File
@@ -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
python/sglang/srt/environ.py
+6
View File
@@ -1455,6 +1455,12 @@ class Envs:
# Quantize the SWA fp8 KV cache from bf16-rounded values (matches
# trainer-side QAT and the DSA-CP path) instead of fp32 registers.
SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False)
# unified_kv only: split the pool into an fp8 nope pool plus a parallel
# bf16 rope pool, 640 B/token instead of 1024. The unified pool takes no
# dtype, so --kv-cache-dtype has no effect there and this switch is the
# only way to ask; on separate-KV it is the reverse -- --kv-cache-dtype
# picks the buffer dtype and this switch is inert.
SGLANG_DSV4_UNIFIED_KV_FP8 = EnvBool(False)
# Kernels and indexer
SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True)
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+81 -12
View File
@@ -1366,8 +1366,17 @@ class DeepseekV4HipRadixBackend(
attn_sink: torch.Tensor,
core_attn_metadata: DSV4AttnMetadata,
save_kv_cache: bool = True,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""unified_kv paged-attention path over the bf16 unified_kv"""
"""unified_kv paged-attention path over the unified_kv pool.
``q_rope`` is what tells the two layouts apart: present means ``q`` is a
packed fp8 row and the pool is the two-pool fp8 one, so decode goes to
the asm reader; absent means both are plain bf16 and it goes to Triton.
Prefill needs ``k_rope`` alongside it, because there the current chunk is
a KV source of its own and not just something to store.
"""
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
pool = self.token_to_kv_pool
@@ -1401,6 +1410,10 @@ class DeepseekV4HipRadixBackend(
else:
state_slot = forward_batch.req_pool_indices[:T]
if save_kv_cache:
# Only verify reaches this under fp8 -- plain decode's rows are
# written by the fused kernel itself, which leaves kv None. The
# pair arrives already packed, so this is the same scatter with
# a second pool hanging off it.
runtime.store_swa_into_unified(
kv=kv,
state_slot=state_slot,
@@ -1409,6 +1422,10 @@ class DeepseekV4HipRadixBackend(
win=win,
ring_stride=ring_stride,
final_pos=positions,
kv_rope=k_rope,
unified_kv_rope=(
None if k_rope is None else pool.get_unified_kv_rope(layer_id)
),
)
unified_metadata = core_attn_metadata.unified
if compress_ratio == 0:
@@ -1430,6 +1447,25 @@ class DeepseekV4HipRadixBackend(
)
else:
raise ValueError(f"bad compress_ratio {compress_ratio}")
if q_rope is not None:
# softmax_scale is not passed on: the asm kernel hardcodes
# 1/sqrt(512), which is what self.softmax_scale already is for
# V4's head_dim=512. The other readers here take it explicitly,
# so a head_dim change would leave only this one mis-scaled.
assert self.softmax_scale == 512**-0.5, (
"the v4 nm asm kernel hardcodes 1/sqrt(512), this backend is "
f"at {self.softmax_scale}"
)
return runtime.decode_fp8_2buff(
q=q,
q_rope=q_rope,
unified_kv=unified,
unified_kv_rope=pool.get_unified_kv_rope(layer_id),
kv_indices=kv_indices,
kv_indptr=kv_indptr,
attn_sink=attn_sink,
v_head_dim=layer.v_head_dim,
)
return runtime.decode(
q=q,
unified_kv=unified,
@@ -1505,17 +1541,42 @@ class DeepseekV4HipRadixBackend(
pad = T + 1 - kpre_p.shape[0]
kpre_p = torch.cat([kpre_p, kpre_p[-1:].expand(pad)])
kext_p = torch.cat([kext_p, kext_p[-1:].expand(pad)])
o = runtime.prefill(
q=q,
unified_kv=unified,
kv_indices_prefix=kpre_i,
kv_indptr_prefix=kpre_p,
kv_extend=kv,
kv_indices_extend=kext_i,
kv_indptr_extend=kext_p,
attn_sink=attn_sink,
softmax_scale=self.softmax_scale,
)
if q_rope is not None:
assert k_rope is not None, (
"fp8 prefill needs the extend rope half beside the packed nope; "
"q_rope came through but k_rope did not"
)
# No empty-segment mask on the result, unlike decode: this kernel
# returns zeros for a token with neither region where the asm decode
# reader leaves the row NaN. Chunk 0 tokens have an empty prefix and
# a non-empty extend, which both readers handle.
o = runtime.prefill_fp8_2buff(
q=q,
q_rope=q_rope,
unified_kv=unified,
unified_kv_rope=pool.get_unified_kv_rope(layer_id),
kv_indices_prefix=kpre_i,
kv_indptr_prefix=kpre_p,
kv_extend=kv,
kv_extend_rope=k_rope,
kv_indices_extend=kext_i,
kv_indptr_extend=kext_p,
attn_sink=attn_sink,
softmax_scale=self.softmax_scale,
v_head_dim=layer.v_head_dim,
)
else:
o = runtime.prefill(
q=q,
unified_kv=unified,
kv_indices_prefix=kpre_i,
kv_indptr_prefix=kpre_p,
kv_extend=kv,
kv_indices_extend=kext_i,
kv_indptr_extend=kext_p,
attn_sink=attn_sink,
softmax_scale=self.softmax_scale,
)
# write this chunk's SWA K into the ring for future chunks / decode
# only the final-window tokens per request
@@ -1535,6 +1596,10 @@ class DeepseekV4HipRadixBackend(
win=win,
ring_stride=ring_stride,
final_pos=_ring_final_pos,
kv_rope=None if k_rope is None else k_rope[:n_real],
unified_kv_rope=(
None if k_rope is None else pool.get_unified_kv_rope(layer_id)
),
)
return o
@@ -1602,6 +1667,8 @@ class DeepseekV4HipRadixBackend(
compress_ratio: Literal[0, 4, 128],
save_kv_cache: bool = True,
attn_sink: Optional[torch.Tensor] = None,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
**_,
) -> torch.Tensor:
if self.mtp_enabled and forward_batch.forward_mode.is_idle():
@@ -1630,6 +1697,8 @@ class DeepseekV4HipRadixBackend(
attn_sink=attn_sink,
core_attn_metadata=core_attn_metadata,
save_kv_cache=save_kv_cache,
q_rope=q_rope,
k_rope=k_rope,
)
if isinstance(core_attn_metadata, DSV4AttnMetadata):
python/sglang/srt/layers/attention/dsv4/compressor_v2.py
+16 -1
View File
@@ -158,6 +158,8 @@ class CompressorBackendMixin:
bf16_store: bool = False,
kv_scale_cache: Optional[torch.Tensor] = None,
rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
fp8_2buff: bool = False,
kv_cache_rope: Optional[torch.Tensor] = None,
) -> None:
assert compress_ratio == 4 or compress_ratio == 128
assert rotate == is_indexer == (head_dim == 128)
@@ -220,6 +222,8 @@ class CompressorBackendMixin:
if _is_hip and use_fp4_indexer
else None
),
fp8_2buff=fp8_2buff,
kvcache_rope=kv_cache_rope,
)
def forward_unified(
@@ -238,6 +242,7 @@ class CompressorBackendMixin:
state_pool = compressor.get_state_pool(self)
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_fp8,
is_unified_kv_triton,
)
@@ -264,6 +269,8 @@ class CompressorBackendMixin:
use_hip_fp4 = _is_hip and use_fp4_indexer
bf16_store = False
kv_scale_cache = None
fp8_2buff = False
kv_cache_rope = None
if compressor.is_in_indexer:
page_size = token_to_kv_pool.get_index_k_page_size(compressor.ratio)
if use_hip_fp4:
@@ -278,7 +285,11 @@ class CompressorBackendMixin:
self.forward_metadata.core_metadata.unified,
f"c{compressor.ratio}_out_loc",
)
bf16_store = True
if is_unified_kv_fp8():
fp8_2buff = True
kv_cache_rope = token_to_kv_pool.get_unified_kv_rope(layer_id)
else:
bf16_store = True
else:
_, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
assert compress_kv_pool is not None
@@ -305,6 +316,10 @@ class CompressorBackendMixin:
rope_cache=(
(compressor.fp4_cos, compressor.fp4_sin) if use_hip_fp4 else None
),
fp8_2buff=fp8_2buff,
kv_cache_rope=(
None if kv_cache_rope is None else kv_cache_rope.view(dtype=torch.uint8)
),
)
online_c128_mtp = getattr(self, "online_c128_mtp", None)
if online_c128_mtp is not None:
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+122 -8
View File
@@ -16,6 +16,7 @@ from sglang.kernels.ops.attention.dsv4 import (
index_buf_accessor as dsv4_index_buf_accessor,
)
from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import layout
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
@@ -511,11 +512,43 @@ class DeepSeekV4LayerItem(NamedTuple):
compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None
# re-exported: the pool allocates the rows, but the kernels that write them own the
# layout (see unified_kv_kernels/layout.py)
DSV4_FP8_NOPE_ROW_BYTES = layout.DSV4_FP8_NOPE_ROW_BYTES
DSV4_FP8_QUANT_TILE = layout.DSV4_FP8_QUANT_TILE
def dsv4_unified_row_bytes(
qk_nope_head_dim: int, qk_rope_head_dim: int, fp8: bool
) -> int:
"""Bytes one unified_kv token occupies, summed over both pools."""
if not fp8:
return (qk_nope_head_dim + qk_rope_head_dim) * 2
num_tiles = -(-qk_nope_head_dim // DSV4_FP8_QUANT_TILE)
scale_bytes = 2 * num_tiles
# not an assert: sizing runs under -O too, and a silently skipped check here
# overreports capacity
if qk_nope_head_dim + scale_bytes > DSV4_FP8_NOPE_ROW_BYTES:
raise ValueError(
f"fp8 nope row overflows: {qk_nope_head_dim} latent values at 1 B + "
f"{scale_bytes} B scales > {DSV4_FP8_NOPE_ROW_BYTES} B stride"
)
return DSV4_FP8_NOPE_ROW_BYTES + qk_rope_head_dim * 2
# The following kv pool follows ATOM's unified_kv kernel layout.
class DeepSeekV4UnifiedKVPool:
"""
Layout:
Layout (bf16):
unified_kv[L]: ``[swa_pages + padded_compress_rows, head_dim]`` bf16
Layout (fp8, ``SGLANG_DSV4_UNIFIED_KV_FP8``) -- two parallel pools with the
same row count, so a row index means the same thing in both. Named after the
accessors, which under fp8 each return one half -- ``get_unified_kv`` the
nope, ``get_unified_kv_rope`` the rope:
unified_kv[L] (nope): ``[rows, 512]`` fp8, see DSV4_FP8_NOPE_ROW_BYTES
unified_kv_rope[L] (rope): ``[rows, qk_rope_head_dim]`` bf16, never quantized
- rows ``[0, swa_pages)`` = SWA ring (``req_pool_indices * swa_window + pos % swa_window``)
- rows ``[swa_pages, ...)`` = compressed (``swa_pages + page_index``)
"""
@@ -535,8 +568,11 @@ class DeepSeekV4UnifiedKVPool:
memory_saver_adapter,
custom_mem_pool,
swa_ring_size: int,
fp8: bool = False,
):
self.swa_ring_size = swa_ring_size
self.fp8 = fp8
self.rope_head_dim = qk_rope_head_dim
self.head_dim = qk_nope_head_dim + qk_rope_head_dim
self.num_slots = num_slots
self.swa_pages = num_slots * self.swa_ring_size
@@ -545,6 +581,7 @@ class DeepSeekV4UnifiedKVPool:
self.k_per_block = dict(self.K_PER_BLOCK)
bufs = []
rope_bufs = []
with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
with (
torch.cuda.use_mem_pool(custom_mem_pool)
@@ -557,20 +594,54 @@ class DeepSeekV4UnifiedKVPool:
compress_rows = self.num_blocks * self.k_per_block[ratio]
rows_per_page = self.page_size // ratio if ratio else 0
padded_compress_rows = compress_rows + rows_per_page
bufs.append(
torch.zeros(
self.swa_pages + padded_compress_rows,
self.head_dim,
dtype=torch.bfloat16,
device=device,
rows = self.swa_pages + padded_compress_rows
if self.fp8:
bufs.append(
torch.zeros(
rows,
DSV4_FP8_NOPE_ROW_BYTES,
dtype=torch.float8_e4m3fn,
device=device,
)
)
)
rope_bufs.append(
torch.zeros(
rows,
self.rope_head_dim,
dtype=torch.bfloat16,
device=device,
)
)
else:
bufs.append(
torch.zeros(
rows,
self.head_dim,
dtype=torch.bfloat16,
device=device,
)
)
rope_bufs.append(None)
self.kv_buffer = bufs
self.kv_buffer_rope = rope_bufs
def get_unified_kv(self, local_layer_id: int) -> torch.Tensor:
return self.kv_buffer[local_layer_id]
def get_unified_kv_rope(self, local_layer_id: int) -> torch.Tensor:
assert self.fp8, "rope pool only exists under SGLANG_DSV4_UNIFIED_KV_FP8"
return self.kv_buffer_rope[local_layer_id]
def get_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
if self.fp8:
# same single-pool assumption as the outer get_contiguous_buf_infos:
# one pointer and one row size per layer describes the nope pool only,
# so whoever picks this up next would move half a row and not notice.
# TODO(danli103): report both pools once a consumer needs them.
raise NotImplementedError(
"get_buf_infos describes one pool per layer; the fp8 rope pool "
"would be dropped (SGLANG_DSV4_UNIFIED_KV_FP8=1)."
)
data_ptrs = [b.data_ptr() for b in self.kv_buffer]
data_lens = [b.nbytes for b in self.kv_buffer]
item_lens = [b[0].nbytes for b in self.kv_buffer]
@@ -578,6 +649,10 @@ class DeepSeekV4UnifiedKVPool:
class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
# object.__new__ stubs (disagg wire test) skip __init__; False is the env
# default, so the fp8 PD/HiCache refuses don't AttributeError on them.
_unified_kv_fp8 = False
def __init__(
self,
max_num_reqs: int,
@@ -633,11 +708,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.c4_logical_size = c4_logical_size
self.c128_size = c128_size
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_fp8,
is_unified_kv_triton,
)
# Resolve the unified-kv gate before any sizing so the two cannot drift.
self._unified_kv = is_unified_kv_triton()
self._unified_kv_fp8 = is_unified_kv_fp8()
# Uniform 512-dim e4m3 layout for the trtllm attention backend
self.uniform_fp8 = (
not self._unified_kv
@@ -721,6 +798,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
memory_saver_adapter=self.memory_saver_adapter,
custom_mem_pool=self.custom_mem_pool,
swa_ring_size=swa_ring_size,
fp8=self._unified_kv_fp8,
)
self.unified_swa_window = self.sliding_window
@@ -766,6 +844,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.wait_layer_transfer(layer_id)
return self.unified_kv_pool.get_unified_kv(layer_id - self._stage_start)
def get_unified_kv_rope(self, layer_id: int) -> torch.Tensor:
self.wait_layer_transfer(layer_id)
return self.unified_kv_pool.get_unified_kv_rope(layer_id - self._stage_start)
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
self.full_to_swa_index_mapping = full_to_swa_index_mapping
@@ -782,6 +864,18 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
data_lens: List[int] = []
item_lens: List[int] = []
if self._unified_kv_fp8:
# The page-block transfer below prices one row as buf[0].nbytes and
# ships a single pointer per layer. Under fp8 that covers the nope
# pool only -- the parallel bf16 rope pool would be dropped and the
# remote side would decode rows against stale rope. Refuse instead.
# TODO(danli103): ship the rope pool as a second per-layer entry.
raise NotImplementedError(
"PD disaggregation is not supported with "
"SGLANG_DSV4_UNIFIED_KV_FP8=1 (the transfer assumes a single "
"unified pool; the rope pool would be silently dropped)."
)
def append_page_buffer(buf: torch.Tensor) -> None:
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
data_ptrs.append(buf.data_ptr())
@@ -827,6 +921,15 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
item_lens: List[int] = []
if not self._unified_kv:
return data_ptrs, data_lens, item_lens
if self._unified_kv_fp8:
# Other half of the PD path -- get_contiguous_buf_infos ships the
# compressed region, this one the ring. Same single-pool assumption,
# same silently dropped rope, same fix -- land them together.
raise NotImplementedError(
"PD disaggregation is not supported with "
"SGLANG_DSV4_UNIFIED_KV_FP8=1 (the SWA_RING component assumes a "
"single unified pool; the rope pool would be silently dropped)."
)
swa_pages = self.unified_kv_pool.swa_pages
for buf in self.unified_kv_pool.kv_buffer:
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
@@ -841,6 +944,17 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
# the unified pool stores individual token rows after its SWA region.
assert self._unified_kv, "unified_region_buffers requires unified_kv layout"
assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
if self._unified_kv_fp8:
# item_bytes below prices kv_buffer alone, so the rope pool would never
# be offloaded and a fetched page would carry stale rope -- wrong output,
# no crash.
# TODO(danli103): give rope its own host pool, the way C4_INDEXER
# already parallels C4.
raise NotImplementedError(
"HiCache offload is not supported with "
"SGLANG_DSV4_UNIFIED_KV_FP8=1 (the host pool assumes a single "
"unified pool; the rope pool would never be offloaded)."
)
swa_pages = self.unified_kv_pool.swa_pages
head_dim = self.unified_kv_pool.head_dim
python/sglang/srt/model_executor/pool_configurator.py
+51 -4
View File
@@ -986,16 +986,56 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4)
self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128)
# Unified-KV uses a different physical layout than the non-unified V4 path:
# * one row carries the full latent -- 1024 B bf16, or 640 B under
# SGLANG_DSV4_UNIFIED_KV_FP8 (512 B fp8 nope + 128 B bf16 rope) -- not
# that path's 584-byte fp8(nope) + bf16(rope) + scales cell.
# * SWA is a fixed per-request ring (num_req_slots * ring_size),
# independent of full_token, so it is a fixed *bias* rather than a
# per-token term. Gate on the same switch the pool itself uses so the
# sizing and the allocation never drift apart.
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_fp8,
is_unified_kv_triton,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
dsv4_unified_row_bytes,
)
self._unified = is_unified_kv_triton()
self._unified_fp8 = is_unified_kv_fp8()
self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
# Row width across both pools: 1024 B bf16, 640 B fp8. Read from the pool
# module so sizing can't drift from the allocation.
self._unified_row_bytes = dsv4_unified_row_bytes(
self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8
)
# swa_page_size is the model's sliding window (cfg.window_size).
self._swa_ring_size = get_swa_ring_size(self.swa_page_size, self.is_speculative)
self._spec_infl = 1.0
# The unified pool takes no dtype, so --kv-cache-dtype never reaches it.
# V4 defaults "auto" to fp8_e4m3 (overrides.py
# _deepseek_v4_kv_cache_dtype), so only a bfloat16 here tells us the user
# set it explicitly; warning on the fp8 side would fire on every run.
if self._unified_fp8 and self.kv_cache_dtype_str == "bfloat16":
logger.warning(
"--kv-cache-dtype=bfloat16 is ignored on the unified_kv path; "
"SGLANG_DSV4_UNIFIED_KV_FP8=1 stores the latent as fp8. Unset the "
"env switch to get a bf16 unified pool."
)
# get_contiguous_buf_infos ships one pointer per layer and prices a row as
# buf[0].nbytes, which under fp8 covers the nope pool only. Fail at startup
# rather than at the first transfer.
# TODO(danli103): drop this once the transfer ships the rope pool.
if self._unified_fp8 and self.disaggregation_mode != "null":
raise ValueError(
"SGLANG_DSV4_UNIFIED_KV_FP8=1 does not support PD disaggregation "
f"(disaggregation_mode={self.disaggregation_mode!r}). Unset the fp8 "
"switch or run without disaggregation."
)
if self.is_speculative:
# Ring is sized once here, so it must serve the largest adaptive tier.
self._assert_ring_serves_draft_tokens(
@@ -1065,8 +1105,10 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
def _get_bytes_per_full_token(self) -> float:
if self._unified:
# Unified_kv stores the whole latent in bf16.
kv_bytes = self.attn_head_dim * 2
# Unified_kv stores the whole latent: one bf16 pool, or an fp8 nope
# pool plus a bf16 rope pool. kv_bytes also prices the compressed
# c4/c128 rows below, which live in the same pool(s).
kv_bytes = self._unified_row_bytes
else:
kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8
@@ -1199,14 +1241,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
return min(estimated, full_token // 2)
def _fixed_swa_bytes(self, max_running_requests: int) -> int:
"""Unified_kv SWA is a fixed per-request ring, sized by concurrency
(num_req_slots) rather than by full_token. Return its byte footprint
across all full layers, inflated for the draft worker the same way as the
per-token coeff. Returns 0 on the non-unified path (where SWA is already
accounted per-token)."""
if not self._unified:
return 0
num_req_slots = self._get_num_req_slots(max_running_requests)
ring_bytes = (
num_req_slots
* self._swa_ring_size
* self.attn_head_dim
* 2 # bf16
* self._unified_row_bytes
* self.num_layers_total
)
return int(ring_bytes * self._spec_infl)
@@ -1277,6 +1323,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
sizes = self._compute_dsv4_sizes(full_token, page_size)
logger.info(
f"DSV4 memory calculation: unified={self._unified}, "
f"unified_fp8={self._unified_fp8}, "
f"bytes_per_full_token={self.bytes_per_full_token:.2f}, "
f"available_bytes={available_bytes / (1 << 30):.2f} GB, "
f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, "
python/sglang/srt/models/deepseek_v4.py
+175 -25
View File
@@ -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
test/registered/e2e/dsv4/test_dsv4_unified_fp8_backend_prefill.py
+247
View File
@@ -0,0 +1,247 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
"""The backend's own prefill wiring: what reaches the reader, and what lands in the ring.
The pieces on either side of this are covered elsewhere -- the scatter primitive by
test_dsv4_unified_fp8_scatter, the model->backend kwargs by the q_pair test -- but
the middle, where _forward_unified_kv picks the fp8 arm and hands the packed pair to
both attention and the ring write, had nothing running through it.
Losing the rope half of that write is silent: the nope pool gets this chunk's rows,
the rope pool keeps stale ones, and later chunks plus decode read a wrong RoPE with
no crash and no NaN. So these run the real store against real (small) pools and pin
that both pools got written, on the same ring row. The attention reader is stubbed:
it is covered by test_dsv4_unified_fp8_prefill, and the store is what is at stake.
"""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
import sglang.srt.layers.attention.deepseek_v4_backend_hip_radix as backend_mod
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# the store is a plain row move, but the two-pool layout it pins is gfx95-only, so
# run it where the feature lives rather than on the default mi300 runner
register_amd_ci(est_time=15, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
ROPE_DIM = 64
V_HEAD_DIM = 512
NUM_HEADS = 16
WIN = 8
RING_STRIDE = 8
SWA_PAGES = 24 # ring rows are state_slot * RING_STRIDE + pos % RING_STRIDE, so < 24
POOL_ROWS = 32
# distinctive fill, so "the store never ran here" and "the store wrote zeros" are
# different failures
NOPE_SENTINEL = 0xEE
ROPE_SENTINEL = -7.0
# two requests on ring slots 1 and 2, three tokens each at positions 0..2
STATE_SLOT = [1, 1, 1, 2, 2, 2]
POSITIONS = [0, 1, 2, 0, 1, 2]
CU_Q = [0, 0, 0, 3, 3, 3]
EXPECTED_ROWS = [8, 9, 10, 16, 17, 18]
_needs_gfx950 = unittest.skipUnless(
torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
"the two-pool fp8 layout is gfx95-only",
)
def _ints(values):
return torch.tensor(values, dtype=torch.int32, device=DEVICE).contiguous()
class _Pool:
"""Just the surface _forward_unified_kv touches."""
def __init__(self, fp8):
self.unified_swa_window = WIN
self.unified_swa_ring_size = RING_STRIDE
self.unified_swa_pages = SWA_PAGES
if fp8:
self.nope = torch.full(
(POOL_ROWS, NOPE_ROW_BYTES),
NOPE_SENTINEL,
dtype=torch.uint8,
device=DEVICE,
).view(torch.float8_e4m3fn)
else:
self.nope = torch.full(
(POOL_ROWS, V_HEAD_DIM),
ROPE_SENTINEL,
dtype=torch.bfloat16,
device=DEVICE,
)
self.rope = torch.full(
(POOL_ROWS, ROPE_DIM), ROPE_SENTINEL, dtype=torch.bfloat16, device=DEVICE
)
def get_unified_kv(self, layer_id):
return self.nope
def get_unified_kv_rope(self, layer_id):
return self.rope
def _chunk(fp8):
"""This fwd's K, one row per token, every row a different value."""
tokens = len(STATE_SLOT)
if fp8:
rows = torch.arange(1, tokens + 1, dtype=torch.uint8, device=DEVICE)
nope = rows[:, None].expand(tokens, NOPE_ROW_BYTES).contiguous()
nope = nope.view(torch.float8_e4m3fn)
else:
rows = torch.arange(1, tokens + 1, dtype=torch.bfloat16, device=DEVICE)
nope = rows[:, None].expand(tokens, V_HEAD_DIM).contiguous()
rope = (
torch.arange(1, tokens + 1, dtype=torch.bfloat16, device=DEVICE)[:, None]
.expand(tokens, ROPE_DIM)
.contiguous()
)
return nope, rope
class TestUnifiedFp8BackendPrefill(CustomTestCase):
def _run(self, fp8=True, save_kv_cache=True):
tokens = len(STATE_SLOT)
pool = _Pool(fp8)
k_nope, k_rope = _chunk(fp8)
if fp8:
q = torch.zeros(
tokens, NUM_HEADS, NOPE_ROW_BYTES, dtype=torch.uint8, device=DEVICE
).view(torch.float8_e4m3fn)
q_rope = torch.zeros(
tokens, NUM_HEADS, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
)
else:
q = torch.zeros(
tokens, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=DEVICE
)
q_rope, k_rope = None, None
unified_meta = SimpleNamespace(
pf_state_slot=_ints(STATE_SLOT),
pf_chunk_start=_ints([0] * tokens),
pf_cu_q=_ints(CU_Q),
pf_final_pos=_ints([max(POSITIONS)] * tokens),
)
core_meta = SimpleNamespace(
unified=unified_meta,
c128_page_indices=None,
c4_sparse_page_indices=None,
)
forward_batch = SimpleNamespace(
forward_mode=ForwardMode.EXTEND,
positions=torch.tensor(POSITIONS, dtype=torch.int64, device=DEVICE),
req_pool_indices=_ints(STATE_SLOT),
)
fake_self = SimpleNamespace(
token_to_kv_pool=pool, softmax_scale=V_HEAD_DIM**-0.5
)
reader_calls = []
def _fake_reader(**kwargs):
reader_calls.append(kwargs)
return torch.zeros(
tokens, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16, device=DEVICE
)
target = "prefill_fp8_2buff" if fp8 else "prefill"
with (
patch.object(runtime, target, _fake_reader),
patch.object(
backend_mod,
"get_parallel",
return_value=SimpleNamespace(attn_cp_size=1, attn_cp_rank=0),
),
):
backend_mod.DeepseekV4HipRadixBackend._forward_unified_kv(
fake_self,
q=q,
kv=k_nope,
layer=SimpleNamespace(layer_id=0, v_head_dim=V_HEAD_DIM),
forward_batch=forward_batch,
compress_ratio=0,
attn_sink=torch.zeros(NUM_HEADS, dtype=torch.float32, device=DEVICE),
core_attn_metadata=core_meta,
save_kv_cache=save_kv_cache,
q_rope=q_rope,
k_rope=k_rope,
)
self.assertEqual(len(reader_calls), 1)
return pool, k_nope, k_rope, reader_calls[0]
def _untouched(self):
return sorted(set(range(POOL_ROWS)) - set(EXPECTED_ROWS))
@_needs_gfx950
def test_both_pools_get_this_chunk_on_the_same_ring_row(self):
"""the regression this file exists for: a rope pool left holding stale rows"""
pool, k_nope, k_rope, _ = self._run()
for token, row in enumerate(EXPECTED_ROWS):
self.assertTrue(
torch.equal(
pool.nope[row].view(torch.uint8), k_nope[token].view(torch.uint8)
),
f"nope pool row {row} does not hold token {token}",
)
self.assertTrue(
torch.equal(pool.rope[row], k_rope[token]),
f"rope pool row {row} does not hold token {token} -- "
f"got {pool.rope[row][0].item()}, want {k_rope[token][0].item()}",
)
@_needs_gfx950
def test_rows_outside_the_window_are_left_alone(self):
"""both scatters take the same row, so neither may spray past it"""
pool, _, _, _ = self._run()
rest = self._untouched()
self.assertTrue(
bool((pool.nope[rest].view(torch.uint8) == NOPE_SENTINEL).all())
)
self.assertTrue(bool((pool.rope[rest] == ROPE_SENTINEL).all()))
@_needs_gfx950
def test_the_reader_gets_the_same_pair_the_ring_write_does(self):
_, k_nope, k_rope, call = self._run()
self.assertIs(call["kv_extend"], k_nope)
self.assertIs(call["kv_extend_rope"], k_rope)
self.assertIsNotNone(call["unified_kv_rope"])
@_needs_gfx950
def test_nothing_is_written_when_the_model_already_stored(self):
pool, _, _, _ = self._run(save_kv_cache=False)
self.assertTrue(bool((pool.nope.view(torch.uint8) == NOPE_SENTINEL).all()))
self.assertTrue(bool((pool.rope == ROPE_SENTINEL).all()))
@_needs_gfx950
def test_the_bf16_arm_never_touches_the_rope_pool(self):
"""one pool, one write -- the rope pool only exists under the fp8 layout"""
pool, k_nope, _, _ = self._run(fp8=False)
for token, row in enumerate(EXPECTED_ROWS):
self.assertTrue(torch.equal(pool.nope[row], k_nope[token]))
self.assertTrue(bool((pool.rope == ROPE_SENTINEL).all()))
if __name__ == "__main__":
unittest.main()
test/registered/e2e/dsv4/test_dsv4_unified_fp8_compress_store.py
+332
View File
@@ -0,0 +1,332 @@
"""Two-pool fp8 store tests for the compressor's norm+rope kernel.
Under SGLANG_DSV4_UNIFIED_KV_FP8 the c4/c128 compressor writes its compressed
latent through ``forward_fp8_2buff``: a 512 B fp8 nope row (448 B payload + 7
UE8M0 tile scales stored twice) in the unified_kv pool, plus a bf16 rope row in
the second pool, both at ``out_loc``. These tests pin that layout for both
compress ratios, against the bf16 store of the same kernel (which shares the
norm+rope math, so the comparison is byte-exact) and against a torch reference.
Both plans are covered. Most cases run the decode plan; the extend arm (what
prefill takes) gets the bf16 comparison only, since its plan check and its
out_loc bound are hand-copied from decode's and nothing else exercises them.
"""
import unittest
import torch
from sglang.kernels.ops.attention.deepseek_v4_rope import precompute_freqs_cis
from sglang.kernels.ops.attention.dsv4 import (
CompressorDecodePlan,
CompressorPrefillPlan,
compress_norm_rope_store,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DSV4_FP8_NOPE_ROW_BYTES,
DSV4_FP8_QUANT_TILE,
)
from sglang.srt.utils import is_gfx95_supported
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# the kernel takes E4M3FN vs E4M3FNUZ from the arch and the two-pool layout is only
# ever allocated on gfx95, so on the default mi300 runner every case here would skip
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
HEAD_DIM = 512
ROPE_DIM = 64
NOPE_DIM = HEAD_DIM - ROPE_DIM
NUM_TILES = NOPE_DIM // DSV4_FP8_QUANT_TILE
SCALE_OFF = NOPE_DIM
SCALE_BYTES = 2 * NUM_TILES
NUM_TOKENS = 6
POOL_ROWS = 32
EPS = 1e-6
FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
RATIOS = (4, 128)
def _inputs(compress_ratio, seq_lens=None):
torch.manual_seed(compress_ratio)
kv = torch.randn(NUM_TOKENS, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
if seq_lens is None:
seq_lens = (
torch.arange(1, NUM_TOKENS + 1, device=DEVICE, dtype=torch.int64)
* compress_ratio
)
plan = CompressorDecodePlan.generate_legacy(
compress_ratio,
torch.arange(NUM_TOKENS, device=DEVICE, dtype=torch.int64),
seq_lens,
)
# every other row, so a row that gets written always has an untouched neighbour
out_loc = torch.arange(1, 2 * NUM_TOKENS + 1, 2, device=DEVICE, dtype=torch.int64)
freqs_cis = precompute_freqs_cis(
ROPE_DIM, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
).to(DEVICE)
return kv, weight, seq_lens, plan, out_loc, freqs_cis
def _extend_inputs(compress_ratio):
"""one request whose extend spans several compress boundaries"""
torch.manual_seed(compress_ratio + 1)
total = compress_ratio * NUM_TOKENS
seq_lens = torch.tensor([total], dtype=torch.int64)
plan = CompressorPrefillPlan.generate_legacy(
compress_ratio,
torch.zeros(1, dtype=torch.int64, device=DEVICE),
seq_lens,
seq_lens.clone(), # the whole sequence is the extend
total,
DEVICE,
)
# the kernel binds its token count off the input and then requires the plan to
# have that many rows, so the fixture has to follow whatever the planner emitted
num_c = plan.plan_c.shape[0]
kv = torch.randn(num_c, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
# unlike decode, extend indexes out_loc by ragged_id -- one entry per q token, of
# which only the compress boundaries are ever read. Sizing this num_c long instead
# reads off the end and stores to whatever row index it finds there.
written = torch.arange(1, 2 * num_c + 1, 2, device=DEVICE, dtype=torch.int64)
out_loc = torch.zeros(total, dtype=torch.int64, device=DEVICE)
out_loc[compress_ratio - 1 :: compress_ratio] = written
freqs_cis = precompute_freqs_cis(ROPE_DIM, total + 1, 0, 10000, 1, 32, 1).to(DEVICE)
return kv, weight, plan, out_loc, written, freqs_cis
def _ref_norm_rope(kv, weight, freqs_cis, positions):
"""rmsnorm over the latent, then rope on the trailing 64, as the kernel does."""
x = kv.float()
x = x * torch.rsqrt(x.pow(2).sum(-1, keepdim=True) / HEAD_DIM + EPS)
x = x * weight.float()
nope, pe = x[:, :NOPE_DIM], x[:, NOPE_DIM:]
freqs = torch.view_as_real(freqs_cis).flatten(-2)[positions]
freqs = freqs.reshape(-1, ROPE_DIM // 2, 2).float()
pairs = pe.reshape(-1, ROPE_DIM // 2, 2)
out = torch.empty_like(pairs)
out[..., 0] = pairs[..., 0] * freqs[..., 0] - pairs[..., 1] * freqs[..., 1]
out[..., 1] = pairs[..., 0] * freqs[..., 1] + pairs[..., 1] * freqs[..., 0]
# the quant warps round through bf16 first, so the scales come off bf16 values
return nope.to(torch.bfloat16).float(), out.reshape(-1, ROPE_DIM)
def _tile_scale_bytes(nope):
"""cast_to_ue8m0(max(absmax, 1e-4) / fp8_max) per 1x64 tile."""
tiles = nope.reshape(nope.shape[0], NUM_TILES, DSV4_FP8_QUANT_TILE)
scale_raw = tiles.abs().amax(-1).clamp_min(1e-4) / FP8_MAX
bits = scale_raw.contiguous().view(torch.int32)
exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32)
return exp.to(torch.uint8)
@unittest.skipUnless(is_gfx95_supported(), "needs an AMD gfx95 GPU for e4m3fn")
class TestUnifiedFp8CompressStore(CustomTestCase):
def _store_fp8(self, compress_ratio, *, seq_lens=None, rope_rows=POOL_ROWS):
kv, weight, seq_lens, plan, out_loc, freqs_cis = _inputs(
compress_ratio, seq_lens
)
nope_pool = torch.zeros(
POOL_ROWS, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
)
rope_pool = torch.zeros(
rope_rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
)
compress_norm_rope_store(
kv.clone(),
plan,
norm_weight=weight,
norm_eps=EPS,
freq_cis=freqs_cis,
out_loc=out_loc,
kvcache=nope_pool.view(torch.uint8),
page_size=1,
fp8_2buff=True,
kvcache_rope=rope_pool.view(torch.uint8),
)
ref = _ref_norm_rope(kv, weight, freqs_cis, (seq_lens - compress_ratio).long())
return nope_pool, rope_pool, out_loc, ref
def _store_bf16(self, compress_ratio):
"""same inputs through the bf16 store, i.e. the values before quantization"""
kv, weight, _, plan, out_loc, freqs_cis = _inputs(compress_ratio)
cache = torch.zeros(POOL_ROWS, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE)
compress_norm_rope_store(
kv.clone(),
plan,
norm_weight=weight,
norm_eps=EPS,
freq_cis=freqs_cis,
out_loc=out_loc,
kvcache=cache.view(torch.uint8),
page_size=1,
bf16_store=True,
)
return cache[out_loc]
def _store_extend(self, compress_ratio, *, fp8):
kv, weight, plan, out_loc, written, freqs_cis = _extend_inputs(compress_ratio)
rows = int(written.max().item()) + 2
common = dict(
norm_weight=weight,
norm_eps=EPS,
freq_cis=freqs_cis,
out_loc=out_loc,
page_size=1,
)
if not fp8:
cache = torch.zeros(rows, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE)
compress_norm_rope_store(
kv.clone(),
plan,
kvcache=cache.view(torch.uint8),
bf16_store=True,
**common,
)
return cache[written]
nope_pool = torch.zeros(
rows, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
)
rope_pool = torch.zeros(rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
compress_norm_rope_store(
kv.clone(),
plan,
kvcache=nope_pool.view(torch.uint8),
fp8_2buff=True,
kvcache_rope=rope_pool.view(torch.uint8),
**common,
)
return nope_pool, rope_pool, written
def test_extend_plan_stores_the_same_rows(self):
for ratio in RATIOS:
with self.subTest(compress_ratio=ratio):
nope_pool, rope_pool, written = self._store_extend(ratio, fp8=True)
pre_quant = self._store_extend(ratio, fp8=False)
nope = pre_quant[:, :NOPE_DIM].float()
num_c = nope.shape[0]
scale_bytes = _tile_scale_bytes(nope)
scale = torch.exp2((scale_bytes.to(torch.int32) - 127).float())
want = (
nope.reshape(num_c, NUM_TILES, DSV4_FP8_QUANT_TILE)
/ scale[..., None]
).to(torch.float8_e4m3fn)
self.assertTrue(
torch.equal(
nope_pool[written][:, :NOPE_DIM].view(torch.uint8),
want.view(torch.uint8).reshape(num_c, NOPE_DIM),
)
)
self.assertTrue(
torch.equal(rope_pool[written], pre_quant[:, NOPE_DIM:])
)
def test_row_matches_the_bf16_store_byte_for_byte(self):
for ratio in RATIOS:
with self.subTest(compress_ratio=ratio):
nope_pool, rope_pool, out_loc, _ = self._store_fp8(ratio)
pre_quant = self._store_bf16(ratio)
nope = pre_quant[:, :NOPE_DIM].float()
scale_bytes = _tile_scale_bytes(nope)
scale = torch.exp2((scale_bytes.to(torch.int32) - 127).float())
want = (
nope.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
/ scale[..., None]
).to(torch.float8_e4m3fn)
# the fixture has to reach the top e4m3 exponent, otherwise it would
# not notice a cast that saturates everything above 256
self.assertTrue(bool((want.float().abs() >= 256).any()))
row = nope_pool[out_loc]
self.assertTrue(
torch.equal(
row[:, :NOPE_DIM].view(torch.uint8),
want.view(torch.uint8).reshape(NUM_TOKENS, NOPE_DIM),
)
)
got_scales = row.view(torch.uint8)[
:, SCALE_OFF : SCALE_OFF + SCALE_BYTES
].reshape(NUM_TOKENS, NUM_TILES, 2)
self.assertTrue(torch.equal(got_scales[..., 0], scale_bytes))
self.assertTrue(torch.equal(got_scales[..., 1], scale_bytes))
self.assertTrue(
torch.equal(rope_pool[out_loc], pre_quant[:, NOPE_DIM:])
)
def test_scale_bytes_track_the_torch_reference(self):
for ratio in RATIOS:
with self.subTest(compress_ratio=ratio):
nope_pool, _, out_loc, (ref_nope, _) = self._store_fp8(ratio)
got = nope_pool.view(torch.uint8)[
out_loc, SCALE_OFF : SCALE_OFF + SCALE_BYTES
].reshape(NUM_TOKENS, NUM_TILES, 2)
self.assertTrue(torch.equal(got[..., 0], _tile_scale_bytes(ref_nope)))
def test_dequantized_nope_tracks_the_reference(self):
for ratio in RATIOS:
with self.subTest(compress_ratio=ratio):
nope_pool, _, out_loc, (ref_nope, _) = self._store_fp8(ratio)
exps = _tile_scale_bytes(ref_nope).to(torch.int32) - 127
payload = nope_pool[out_loc, :NOPE_DIM].float()
deq = (
payload.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
* torch.exp2(exps.float())[..., None]
).reshape(NUM_TOKENS, NOPE_DIM)
# e4m3 carries 3 mantissa bits, so half a step is at most ~2^-4 of
# the tile's own absmax; beyond that the scale or the payload is off
tile_absmax = (
ref_nope.reshape(NUM_TOKENS, NUM_TILES, DSV4_FP8_QUANT_TILE)
.abs()
.amax(-1)
.repeat_interleave(DSV4_FP8_QUANT_TILE, dim=1)
)
self.assertTrue(torch.all((deq - ref_nope).abs() <= 0.07 * tile_absmax))
def test_rope_pool_matches_the_bf16_reference(self):
for ratio in RATIOS:
with self.subTest(compress_ratio=ratio):
_, rope_pool, out_loc, (_, ref_pe) = self._store_fp8(ratio)
torch.testing.assert_close(
rope_pool[out_loc].float(), ref_pe, rtol=2e-2, atol=2e-2
)
def test_pad_and_neighbour_rows_untouched(self):
nope_pool, rope_pool, out_loc, _ = self._store_fp8(4)
nope_bytes = nope_pool.view(torch.uint8)
self.assertTrue(torch.all(nope_bytes[out_loc, SCALE_OFF + SCALE_BYTES :] == 0))
untouched = torch.ones(POOL_ROWS, dtype=torch.bool, device=DEVICE)
untouched[out_loc] = False
self.assertTrue(torch.all(nope_bytes[untouched] == 0))
self.assertTrue(torch.all(rope_pool[untouched] == 0))
def test_non_boundary_decode_is_skipped(self):
# only sequences whose length is a multiple of the ratio produce a token
seq_lens = torch.full(
(NUM_TOKENS,), 4 * 128 + 1, device=DEVICE, dtype=torch.int64
)
nope_pool, rope_pool, _, _ = self._store_fp8(128, seq_lens=seq_lens)
self.assertTrue(torch.all(nope_pool.view(torch.uint8) == 0))
self.assertTrue(torch.all(rope_pool == 0))
def test_short_rope_pool_rejected(self):
# one row index addresses both pools, so a short rope pool has to be caught
# before either pool is written
with self.assertRaises(RuntimeError):
self._store_fp8(4, rope_rows=POOL_ROWS // 2)
if __name__ == "__main__":
unittest.main()
test/registered/e2e/dsv4/test_dsv4_unified_fp8_decode.py
+310
View File
@@ -0,0 +1,310 @@
"""Decode read path over the two-pool fp8 unified_kv (aiter's v4 nm asm kernel).
What these pin is the reader-side plumbing, not the kernel's arithmetic: the
packed 512 B nope row and the bf16 rope pool addressed by one shared row index,
a per-token ``qo_indptr``, and the ragged ``kv_indptr`` the existing index
builders emit -- including what they emit for a cuda-graph padded row. The
reference attends over the *dequantized* pools, so a mismatch is the wiring
rather than the fp8 round-trip.
The quantization helpers mirror aiter's own reference
(``op_tests/test_mla_v40_persistent.py``: ``quantize_v4_nope_bpad8`` /
``pack_v4_nope_scale``). They are duplicated rather than imported because that
file is a test, not part of the aiter package.
"""
import unittest
import torch
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# the asm shader is only shipped for gfx950
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
NOPE_DIM = 448 # fp8 values per row, in elements
ROPE_DIM = 64
QUANT_TILE = 64
NUM_TILES = NOPE_DIM // QUANT_TILE # 7
SCALE_OFF = NOPE_DIM # scales start where the values end
# latent element count; the same number as NOPE_ROW_BYTES, different unit
V_HEAD_DIM = NOPE_DIM + ROPE_DIM
SOFTMAX_SCALE = V_HEAD_DIM**-0.5 # what the kernel hardcodes
_needs_gfx950 = unittest.skipUnless(
torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
"two-pool fp8 decode runs on the gfx950 asm shader",
)
def _pow2_ceil_scale(amax: torch.Tensor) -> torch.Tensor:
"""amax/fp8_max -> the next power of two at or above it, as fp32"""
return torch.pow(2.0, torch.clamp_min(amax, 1e-4).log2().ceil()).to(torch.float32)
def _pow2_to_e8m0(pow2: torch.Tensor) -> torch.Tensor:
"""byte B encodes 2^(B-127); 0 means 0.0 and 255 means inf, so clamp to 254"""
biased = torch.log2(pow2).round().to(torch.int32) + 127
return torch.clamp(biased, 0, 254).to(torch.uint8)
def _e8m0_to_fp32(byte: torch.Tensor) -> torch.Tensor:
return torch.exp2((byte.to(torch.int32) - 127).to(torch.float32))
def _quantize_nope(nope_fp32: torch.Tensor):
"""[..., 448] fp32 -> (fp8 values, [..., 7] e8m0 bytes, bf16 round-trip)"""
fp8_max = float(torch.finfo(torch.float8_e4m3fn).max)
leading = nope_fp32.shape[:-1]
tiled = nope_fp32.reshape(*leading, NUM_TILES, QUANT_TILE)
scale = _pow2_ceil_scale(tiled.abs().amax(dim=-1) / fp8_max)
values = (tiled / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
dequant = (values.to(torch.float32) * scale.unsqueeze(-1)).reshape(
*leading, NOPE_DIM
)
return (
values.reshape(*leading, NOPE_DIM),
_pow2_to_e8m0(scale),
dequant.to(torch.bfloat16),
)
def _pack(values: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor:
"""448 values + each tile scale twice + pad, as one NOPE_ROW_BYTES fp8 row
The 50 pad bytes get garbage on purpose. Production allocates Q with
nope_pool.new_empty(), so a reader that ever starts looking past the scales
should fail here and not in a bf16-vs-fp8 accuracy chase.
"""
leading = values.shape[:-1]
row = torch.randint(
1, 256, (*leading, NOPE_ROW_BYTES), dtype=torch.uint8, device=values.device
)
row[..., :NOPE_DIM] = values.view(torch.uint8)
dup = scale_e8m0.unsqueeze(-1).expand(*scale_e8m0.shape, 2).reshape(*leading, -1)
row[..., SCALE_OFF : SCALE_OFF + 2 * NUM_TILES] = dup
return row.view(torch.float8_e4m3fn)
def _make_latent(*leading: int):
"""Return (packed fp8 rows, bf16 rope, bf16 latent the kernel effectively sees)."""
nope = torch.randn(*leading, NOPE_DIM, device=DEVICE, dtype=torch.float32)
rope = torch.randn(*leading, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
values, scale, nope_bf16 = _quantize_nope(nope)
silver = torch.cat([nope_bf16, rope], dim=-1)
return _pack(values, scale).contiguous(), rope.contiguous(), silver
def _ragged(lengths, rows, device=DEVICE):
"""per-token row lists -> (flat int32 indices, int32 indptr)"""
indptr = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=device)
indptr[1:] = torch.cumsum(
torch.tensor(lengths, dtype=torch.int32, device=device), dim=0
)
flat = torch.cat(
[
torch.randperm(rows, device=device)[:n].to(torch.int32)
for n in lengths
if n > 0
]
or [torch.empty(0, dtype=torch.int32, device=device)]
)
return flat.contiguous(), indptr
def _reference(q_silver, kv_silver, indices, indptr, sink):
"""Ragged sparse attention in fp32; V is the full latent, sink has zero V."""
T, H, _ = q_silver.shape
out = torch.zeros(T, H, V_HEAD_DIM, device=q_silver.device, dtype=torch.float32)
q = q_silver.float()
sink_f = sink.float()
for t in range(T):
lo, hi = int(indptr[t]), int(indptr[t + 1])
k = kv_silver[indices[lo:hi].long()].float() # [L, 512]
logits = q[t] @ k.transpose(0, 1) * SOFTMAX_SCALE # [H, L]
aug = torch.cat([logits, sink_f.unsqueeze(1)], dim=1)
m = aug.amax(dim=1, keepdim=True)
p = torch.exp(logits - m)
denom = p.sum(dim=1, keepdim=True) + torch.exp(sink_f.unsqueeze(1) - m)
out[t] = (p @ k) / denom
return out
class TestUnifiedFp8Decode(CustomTestCase):
def setUp(self):
torch.manual_seed(7)
self.rows = 256
def _run(self, lengths, num_heads):
T = len(lengths)
pool_nope, pool_rope, kv_silver = _make_latent(self.rows)
q_packed, q_rope, q_silver = _make_latent(T, num_heads)
indices, indptr = _ragged(lengths, self.rows)
sink = torch.randn(num_heads, device=DEVICE, dtype=torch.float32)
got = runtime.decode_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope,
kv_indices=indices,
kv_indptr=indptr,
attn_sink=sink,
v_head_dim=V_HEAD_DIM,
)
want = _reference(q_silver, kv_silver, indices, indptr, sink)
return got.float(), want
def _assert_close(self, got, want, atol=3e-2, rtol=3e-2):
"""torch-style combined bound.
A pure relative bound is useless here: the latent's outputs straddle
zero, so an absolute error of 4e-3 -- which is what bf16 accumulation
costs -- reads as 47% relative on the rows that land near zero.
"""
diff = (got - want).abs()
outside = diff > atol + rtol * want.abs()
self.assertEqual(
outside.sum().item(),
0,
f"{outside.sum().item()}/{outside.numel()} elements outside "
f"{atol}+{rtol}|ref|, max abs {diff.max().item():.4g}",
)
@_needs_gfx950
def test_matches_dequantized_reference(self):
for lengths in ([64] * 4, [17, 5, 128, 1], [200] * 8):
with self.subTest(lengths=lengths):
got, want = self._run(lengths, num_heads=16)
self._assert_close(got, want)
@_needs_gfx950
def test_head_count_64(self):
got, want = self._run([48, 96], num_heads=64)
self._assert_close(got, want)
@_needs_gfx950
def test_cuda_graph_pad_reads_only_the_reserved_ring_row(self):
"""What the real builder emits for a cuda-graph padded row.
Not an empty segment: both dsv4 backends fill padded ``seq_lens`` with 1,
so ``clamp(seq_lens, max=win)`` leaves the pad one row long. It lands on
ring row 0, the slot ReqToTokenPool reserves for exactly this
(``free_slots`` starts at 1), so a pad only ever reads and writes there.
"""
win = ring = 64
seq_lens = torch.tensor([37, 55, 1, 1], dtype=torch.int32, device=DEVICE)
state_slot = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device=DEVICE)
n = seq_lens.numel()
zero = torch.zeros(n, dtype=torch.int32, device=DEVICE)
indices, indptr = runtime.build_decode_streams(
state_slot=state_slot,
positions=seq_lens - 1, # raw_positions, as the backend derives it
swa_len=torch.clamp(seq_lens, max=win),
hca_len=zero,
csa_len=zero,
hca_page_indices=torch.zeros(n, 1, dtype=torch.int32, device=DEVICE),
csa_width=1,
win=win,
ring_stride=ring,
swa_pages=self.rows,
)[:2]
seg = (indptr[1 : n + 1] - indptr[:n]).tolist()
self.assertEqual(seg, [37, 55, 1, 1])
for pad in (2, 3):
self.assertEqual(indices[int(indptr[pad])].item(), 0)
live = indices[: int(indptr[2])]
self.assertGreaterEqual(int(live.min()), ring, "live rows hit slot 0's block")
pool_nope, pool_rope, _ = _make_latent(self.rows)
q_packed, q_rope, _ = _make_latent(n, 16)
out = runtime.decode_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope,
kv_indices=indices.contiguous(),
kv_indptr=indptr,
attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
v_head_dim=V_HEAD_DIM,
)
# the mask never fires here, so what matters is the reserved row keeping
# the pad finite rather than it coming back zeroed
self.assertTrue(bool(out.isfinite().all()))
@_needs_gfx950
def test_empty_segment_comes_back_nonfinite(self):
"""Guard for a shape the builders do not reach today.
Padded seq_lens are always filled with 1 (see
test_cuda_graph_pad_reads_only_the_reserved_ring_row), so an empty segment
can only come from a builder change -- and it comes back NaN, not zero,
since the asm kernel divides by an all-sink denominator.
"""
got, want = self._run([32, 0, 32], num_heads=16)
self.assertTrue(bool(torch.isnan(got[1]).any()))
for t in (0, 2):
self._assert_close(got[t], want[t])
@_needs_gfx950
def test_split_tail_override_matches_reference(self):
"""past 40 tokens runtime overrides the split count, moving the kernel onto
a different stage-2 merge partition -- must still match the reference
"""
lengths = [200, 64] * 24 # 48 tokens, both layer flavours' segment lengths
self.assertGreater(len(lengths), 40)
got, want = self._run(lengths, num_heads=16)
self._assert_close(got, want)
@_needs_gfx950
def test_rejects_pool_that_is_not_a_pair(self):
pool_nope, pool_rope, _ = _make_latent(self.rows)
q_packed, q_rope, _ = _make_latent(2, 16)
indices, indptr = _ragged([4, 4], self.rows)
sink = torch.zeros(16, device=DEVICE, dtype=torch.float32)
with self.assertRaises(AssertionError):
runtime.decode_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope[: self.rows // 2],
kv_indices=indices,
kv_indptr=indptr,
attn_sink=sink,
v_head_dim=V_HEAD_DIM,
)
@_needs_gfx950
def test_rejects_q_row_wider_than_the_pool_row(self):
pool_nope, pool_rope, _ = _make_latent(self.rows)
q_packed, q_rope, _ = _make_latent(2, 16)
indices, indptr = _ragged([4, 4], self.rows)
sink = torch.zeros(16, device=DEVICE, dtype=torch.float32)
wider = torch.zeros(
2, 16, NOPE_ROW_BYTES + 64, device=DEVICE, dtype=torch.float8_e4m3fn
)
wider[..., :NOPE_ROW_BYTES] = q_packed
with self.assertRaises(AssertionError):
runtime.decode_fp8_2buff(
q=wider,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope,
kv_indices=indices,
kv_indptr=indptr,
attn_sink=sink,
v_head_dim=V_HEAD_DIM,
)
if __name__ == "__main__":
unittest.main()
test/registered/e2e/dsv4/test_dsv4_unified_fp8_prefill.py
+298
View File
@@ -0,0 +1,298 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
"""Prefill read path over the two-pool fp8 unified_kv (aiter's opus kernel).
Two regions per token: the paged prefix pools and this chunk's flat extend pair.
What these pin is that both regions are addressed with the same row layout and
that the pair guards fire before the launch -- the reference attends over the
*dequantized* pools, so a mismatch is the wiring rather than the fp8 round-trip.
The quantization helpers are the ones from the decode test rather than a shared
module: files under test/registered/ are collected standalone (no __init__.py,
no conftest), so importing across them breaks in CI.
"""
import unittest
import torch
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
from sglang.srt.utils import is_gfx95_supported, is_hip
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
NOPE_DIM = 448 # fp8 values per row, in elements
ROPE_DIM = 64
QUANT_TILE = 64
NUM_TILES = NOPE_DIM // QUANT_TILE # 7
SCALE_OFF = NOPE_DIM
# latent element count; the same number as NOPE_ROW_BYTES, different unit
V_HEAD_DIM = NOPE_DIM + ROPE_DIM
SOFTMAX_SCALE = V_HEAD_DIM**-0.5
_needs_gfx950 = unittest.skipUnless(
torch.cuda.is_available() and is_hip() and is_gfx95_supported(),
"two-pool fp8 prefill runs on the gfx950 opus kernel",
)
def _pow2_ceil_scale(amax: torch.Tensor) -> torch.Tensor:
return torch.pow(2.0, torch.clamp_min(amax, 1e-4).log2().ceil()).to(torch.float32)
def _pow2_to_e8m0(pow2: torch.Tensor) -> torch.Tensor:
biased = torch.log2(pow2).round().to(torch.int32) + 127
return torch.clamp(biased, 0, 254).to(torch.uint8)
def _quantize_nope(nope_fp32: torch.Tensor):
"""[..., 448] fp32 -> (fp8 values, [..., 7] e8m0 bytes, bf16 round-trip)"""
fp8_max = float(torch.finfo(torch.float8_e4m3fn).max)
leading = nope_fp32.shape[:-1]
tiled = nope_fp32.reshape(*leading, NUM_TILES, QUANT_TILE)
scale = _pow2_ceil_scale(tiled.abs().amax(dim=-1) / fp8_max)
values = (tiled / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
dequant = (values.to(torch.float32) * scale.unsqueeze(-1)).reshape(
*leading, NOPE_DIM
)
return (
values.reshape(*leading, NOPE_DIM),
_pow2_to_e8m0(scale),
dequant.to(torch.bfloat16),
)
def _pack(values: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor:
"""448 values + each tile scale twice + pad, as one NOPE_ROW_BYTES fp8 row
Pad bytes get garbage on purpose, same reason as the decode test: production
allocates these with new_empty(), so a reader that walks past the scales
should fail here rather than as an accuracy drift.
"""
leading = values.shape[:-1]
row = torch.randint(
1, 256, (*leading, NOPE_ROW_BYTES), dtype=torch.uint8, device=values.device
)
row[..., :NOPE_DIM] = values.view(torch.uint8)
dup = scale_e8m0.unsqueeze(-1).expand(*scale_e8m0.shape, 2).reshape(*leading, -1)
row[..., SCALE_OFF : SCALE_OFF + 2 * NUM_TILES] = dup
return row.view(torch.float8_e4m3fn)
def _make_latent(*leading: int):
"""Return (packed fp8 rows, bf16 rope, bf16 latent the kernel effectively sees)."""
nope = torch.randn(*leading, NOPE_DIM, device=DEVICE, dtype=torch.float32)
rope = torch.randn(*leading, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
values, scale, nope_bf16 = _quantize_nope(nope)
silver = torch.cat([nope_bf16, rope], dim=-1)
return _pack(values, scale).contiguous(), rope.contiguous(), silver
def _ragged(lengths, rows):
"""per-token row lists -> (flat int32 indices, int32 indptr)"""
indptr = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=DEVICE)
indptr[1:] = torch.cumsum(
torch.tensor(lengths, dtype=torch.int32, device=DEVICE), dim=0
)
parts = [
torch.randperm(rows, device=DEVICE)[:n].to(torch.int32)
for n in lengths
if n > 0
]
flat = (
torch.cat(parts) if parts else torch.empty(0, dtype=torch.int32, device=DEVICE)
)
return flat.contiguous(), indptr
def _reference(q_silver, sources, sink, scale):
"""Ragged two-region attention in fp32; V is the full latent, sink V is zero.
``sources`` is [(silver, indices, indptr), ...]. The kernel shares one online
softmax across the regions, so order does not matter and this just
concatenates whatever each region selected.
"""
T, H, _ = q_silver.shape
out = torch.zeros(T, H, V_HEAD_DIM, device=q_silver.device, dtype=torch.float32)
q = q_silver.float()
sink_f = sink.float()
for t in range(T):
keys = []
for silver, indices, indptr in sources:
lo, hi = int(indptr[t]), int(indptr[t + 1])
if hi > lo:
keys.append(silver[indices[lo:hi].long()].float())
if not keys:
# only the sink is left: it contributes to the denominator and has
# V = 0, so the row is exactly zero
continue
k = torch.cat(keys, dim=0)
logits = q[t] @ k.transpose(0, 1) * scale
m = torch.cat([logits, sink_f.unsqueeze(1)], dim=1).amax(dim=1, keepdim=True)
p = torch.exp(logits - m)
denom = p.sum(dim=1, keepdim=True) + torch.exp(sink_f.unsqueeze(1) - m)
out[t] = (p @ k) / denom
return out
class TestUnifiedFp8Prefill(CustomTestCase):
def setUp(self):
torch.manual_seed(11)
self.rows = 256
def _run(self, prefix_lens, extend_lens, num_heads=16, scale=SOFTMAX_SCALE):
T = len(prefix_lens)
self.assertEqual(T, len(extend_lens))
extend_rows = max(max(extend_lens), 1)
pool_nope, pool_rope, pool_silver = _make_latent(self.rows)
ext_nope, ext_rope, ext_silver = _make_latent(extend_rows)
q_packed, q_rope, q_silver = _make_latent(T, num_heads)
pre_i, pre_p = _ragged(prefix_lens, self.rows)
ext_i, ext_p = _ragged(extend_lens, extend_rows)
sink = torch.randn(num_heads, device=DEVICE, dtype=torch.float32)
got = runtime.prefill_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope,
kv_indices_prefix=pre_i,
kv_indptr_prefix=pre_p,
kv_extend=ext_nope,
kv_extend_rope=ext_rope,
kv_indices_extend=ext_i,
kv_indptr_extend=ext_p,
attn_sink=sink,
softmax_scale=scale,
v_head_dim=V_HEAD_DIM,
)
want = _reference(
q_silver,
[(pool_silver, pre_i, pre_p), (ext_silver, ext_i, ext_p)],
sink,
scale,
)
return got.float(), want
def _assert_close(self, got, want, atol=3e-2, rtol=3e-2):
"""torch-style combined bound, same reasoning as the decode test.
A pure relative bound is useless here: the latent's outputs straddle
zero, so the absolute error bf16 accumulation costs reads as a huge
relative one on the rows that land near zero.
"""
diff = (got - want).abs()
outside = diff > atol + rtol * want.abs()
self.assertEqual(
outside.sum().item(),
0,
f"{outside.sum().item()}/{outside.numel()} elements outside "
f"the bound, max abs {diff.max().item():.4g}",
)
@_needs_gfx950
def test_matches_dequantized_reference(self):
cases = (
([64, 64, 64, 64], [1, 2, 3, 4]),
([17, 5, 128, 1], [4, 4, 4, 4]),
([200] * 6, [1, 3, 6, 2, 5, 4]),
)
for prefix_lens, extend_lens in cases:
with self.subTest(prefix=prefix_lens, extend=extend_lens):
got, want = self._run(prefix_lens, extend_lens)
self._assert_close(got, want)
@_needs_gfx950
def test_first_chunk_has_an_empty_prefix_for_every_token(self):
"""the real shape of chunk 0: nothing committed yet, extend is all there is"""
got, want = self._run([0, 0, 0, 0], [1, 2, 3, 4])
self.assertTrue(bool(got.isfinite().all()))
self._assert_close(got, want)
@_needs_gfx950
def test_a_token_with_neither_region_comes_back_zero(self):
"""Not NaN, which is where this differs from the asm decode reader.
decode_fp8_2buff has to mask that case itself; this kernel already
returns zeros, so there is deliberately no mask on this path. The
reference skips those rows for the same reason: with only the sink left
the numerator is zero.
"""
got, want = self._run([64, 0, 64], [2, 0, 2])
self.assertTrue(torch.equal(got[1], torch.zeros_like(got[1])))
self._assert_close(got, want)
@_needs_gfx950
def test_head_count_64(self):
got, want = self._run([48, 96], [3, 5], num_heads=64)
self._assert_close(got, want)
@_needs_gfx950
def test_scale_is_passed_through(self):
"""unlike the decode reader, this kernel takes the scale as an argument"""
got, want = self._run([32, 32], [2, 2], scale=0.5 * SOFTMAX_SCALE)
self._assert_close(got, want)
@_needs_gfx950
def test_extend_row_narrower_than_the_pool_is_rejected(self):
"""the two regions are walked with one row layout, so a short row would
read the next token's bytes as this one's scales"""
pool_nope, pool_rope, _ = _make_latent(self.rows)
q_packed, q_rope, _ = _make_latent(2, 16)
ext_nope, ext_rope, _ = _make_latent(4)
pre_i, pre_p = _ragged([8, 8], self.rows)
ext_i, ext_p = _ragged([1, 1], 4)
with self.assertRaisesRegex(AssertionError, "extend nope row"):
runtime.prefill_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=pool_rope,
kv_indices_prefix=pre_i,
kv_indptr_prefix=pre_p,
kv_extend=ext_nope[:, : NOPE_ROW_BYTES // 2].contiguous(),
kv_extend_rope=ext_rope,
kv_indices_extend=ext_i,
kv_indptr_extend=ext_p,
attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
softmax_scale=SOFTMAX_SCALE,
v_head_dim=V_HEAD_DIM,
)
@_needs_gfx950
def test_mismatched_pools_are_rejected_before_the_launch(self):
pool_nope, _, _ = _make_latent(self.rows)
short_rope = torch.zeros(
self.rows // 2, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
)
q_packed, q_rope, _ = _make_latent(2, 16)
ext_nope, ext_rope, _ = _make_latent(4)
pre_i, pre_p = _ragged([8, 8], self.rows)
ext_i, ext_p = _ragged([1, 1], 4)
with self.assertRaisesRegex(AssertionError, "pool rows differ"):
runtime.prefill_fp8_2buff(
q=q_packed,
q_rope=q_rope,
unified_kv=pool_nope,
unified_kv_rope=short_rope,
kv_indices_prefix=pre_i,
kv_indptr_prefix=pre_p,
kv_extend=ext_nope,
kv_extend_rope=ext_rope,
kv_indices_extend=ext_i,
kv_indptr_extend=ext_p,
attn_sink=torch.randn(16, device=DEVICE, dtype=torch.float32),
softmax_scale=SOFTMAX_SCALE,
v_head_dim=V_HEAD_DIM,
)
if __name__ == "__main__":
unittest.main()
test/registered/e2e/dsv4/test_dsv4_unified_fp8_qk_norm_rope.py
+539
View File
@@ -0,0 +1,539 @@
"""Two-pool fp8 store tests for the fused QK norm+RoPE kernel wrapper.
Under SGLANG_DSV4_UNIFIED_KV_FP8 ``fused_qk_norm_rope_swa_store`` delegates to
aiter, which packs K into a 512 B fp8 nope row (448 B payload + 14 B duplicated
E8M0 tile scales) plus a bf16 rope row and scatters both into the SWA ring.
These tests pin that layout, which the decode reader depends on, and which of the
two forms Q comes back in: the same packed pair when the caller supplies a rope
buffer (what the v4 nm asm reader takes), plain rotated bf16 when it does not
(what the Triton reader takes).
"""
import unittest
import torch
from sglang.kernels.ops.attention.fused_qk_norm_rope_store import (
_HAS_GROUP_QUANT,
fused_qk_norm_rope_swa_store,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DSV4_FP8_NOPE_ROW_BYTES,
DSV4_FP8_QUANT_TILE,
)
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# aiter's group-quant path is gfx95-only, so on the default mi300 runner every case
# here would skip
register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
NOPE_DIM = 448
ROPE_DIM = 64
HEAD_DIM = NOPE_DIM + ROPE_DIM
NUM_TILES = NOPE_DIM // DSV4_FP8_QUANT_TILE
SCALE_OFF = NOPE_DIM
SCALE_BYTES = 2 * NUM_TILES
NUM_HEADS = 4
EPS = 1e-6
MAX_POS = 256
RING_STRIDE = 16
def _cos_sin():
inv = 1.0 / (
10000 ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32) / ROPE_DIM)
)
ang = torch.arange(MAX_POS, dtype=torch.float32)[:, None] * inv[None, :]
return (
ang.cos().to(torch.bfloat16).to(DEVICE),
ang.sin().to(torch.bfloat16).to(DEVICE),
)
def _ref_norm_rope(kv, weight, cos, sin, positions):
"""rmsnorm over the whole latent, then GPT-J rope on the trailing pe half."""
x = kv.float()
scale = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS)
normed = x * scale * weight.float()
nope, pe = normed[:, :NOPE_DIM], normed[:, NOPE_DIM:]
c = cos.float()[positions]
s = sin.float()[positions]
even, odd = pe[:, 0::2], pe[:, 1::2]
out = torch.empty_like(pe)
out[:, 0::2] = even * c - odd * s
out[:, 1::2] = odd * c + even * s
return nope, out
def _ref_tile_scales(nope):
"""e8m0 exponent byte per 1x64 tile, from the fp32 reference nope."""
tiles = nope.reshape(nope.shape[0], NUM_TILES, DSV4_FP8_QUANT_TILE)
absmax = tiles.abs().amax(-1).clamp_min(1e-8)
fp8_max = torch.finfo(torch.float8_e4m3fn).max
return torch.ceil(torch.log2(absmax / fp8_max))
def _pools(n_rows):
nope_pool = torch.zeros(
n_rows, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
)
rope_pool = torch.zeros(n_rows, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
return nope_pool, rope_pool
class _StoreCase(CustomTestCase):
def setUp(self):
torch.manual_seed(7)
self.T = 6
self.cos, self.sin = _cos_sin()
self.weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
self.kv = torch.randn(self.T, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
self.q = torch.randn(
self.T, NUM_HEADS * HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
)
self.positions = torch.arange(self.T, device=DEVICE, dtype=torch.int64)
# distinct ring rows so each row has one unambiguous writer
self.swa_loc = (
self.positions.to(torch.int32) % RING_STRIDE + RING_STRIDE
).contiguous()
@unittest.skipUnless(_HAS_GROUP_QUANT, "needs aiter's group-quant kernel on gfx95x")
class TestUnifiedFp8QkNormRope(_StoreCase):
def _call(
self,
nope_pool=None,
rope_pool=None,
k_nope=None,
k_rope=None,
q_out=None,
q_rope_out=None,
):
return fused_qk_norm_rope_swa_store(
q=self.q,
kv=self.kv,
q_norm_weight=None,
kv_norm_weight=self.weight,
q_rms_eps=EPS,
kv_rms_eps=EPS,
rope_head_dim=ROPE_DIM,
cos_cache=self.cos,
sin_cache=self.sin,
positions=self.positions,
swa_cache=nope_pool,
swa_loc=None if nope_pool is None else self.swa_loc,
swa_page_size=1,
dtype=torch.bfloat16,
fp8_2buff=True,
swa_rope_cache=rope_pool,
k_nope_out=k_nope,
k_rope_out=k_rope,
q_out=q_out,
q_rope_out=q_rope_out,
)
def test_pool_rows_equal_the_dense_packed_output(self):
"""the ring write and the dense K buffers come from the same values"""
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
k_nope = torch.empty(
self.T, 1, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
)
k_rope = torch.empty(self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
self._call(nope_pool, rope_pool, k_nope, k_rope)
rows = self.swa_loc.long()
pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
dense_bytes = k_nope.view(torch.uint8)[:, 0, : SCALE_OFF + SCALE_BYTES]
self.assertTrue(torch.equal(pool_bytes, dense_bytes))
self.assertTrue(torch.equal(rope_pool[rows], k_rope[:, 0]))
def test_strided_kv_slice_matches_contiguous(self):
"""kv is a strided slice of qkv_a; aiter forwards kv.stride(0), so going
back to assuming a packed row would corrupt silently instead of erroring
Covers both callers: the fused ring write (pools) and the caller-buffer
pair (k_nope/k_rope) that prefill and target-verify pass instead.
"""
q_lora_rank = 1536 # DSV4-Pro; only its being != 0 matters here
wide = torch.randn(
self.T, q_lora_rank + HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
)
strided = wide[..., q_lora_rank:]
self.assertFalse(strided.is_contiguous())
self.assertEqual(strided.stride(-1), 1)
runs = []
for kv in (strided, strided.contiguous()):
self.kv = kv
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
k_nope = torch.zeros(
self.T,
1,
DSV4_FP8_NOPE_ROW_BYTES,
dtype=torch.float8_e4m3fn,
device=DEVICE,
)
k_rope = torch.zeros(
self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
)
packed = self._call(nope_pool, rope_pool, k_nope, k_rope)
runs.append((nope_pool, rope_pool, k_nope, k_rope, packed))
for got, want in zip(*runs):
self.assertTrue(torch.equal(got.view(torch.uint8), want.view(torch.uint8)))
def test_scale_bytes_are_duplicated_e8m0(self):
"""the asm reader reads each tile scale twice, so the 14 B must be 7 equal pairs"""
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
self._call(nope_pool, rope_pool)
rows = self.swa_loc.long()
scales = nope_pool.view(torch.uint8)[
rows, SCALE_OFF : SCALE_OFF + SCALE_BYTES
].reshape(self.T, NUM_TILES, 2)
self.assertTrue(torch.equal(scales[..., 0], scales[..., 1]))
ref_nope, _ = _ref_norm_rope(
self.kv, self.weight, self.cos, self.sin, self.positions
)
expected = (_ref_tile_scales(ref_nope) + 127).to(torch.uint8)
self.assertTrue(torch.equal(scales[..., 0], expected))
def test_dequantized_nope_tracks_the_reference(self):
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
self._call(nope_pool, rope_pool)
ref_nope, _ = _ref_norm_rope(
self.kv, self.weight, self.cos, self.sin, self.positions
)
exps = _ref_tile_scales(ref_nope)
payload = nope_pool[self.swa_loc.long(), :NOPE_DIM].float()
deq = (
payload.reshape(self.T, NUM_TILES, DSV4_FP8_QUANT_TILE)
* torch.exp2(exps)[..., None]
).reshape(self.T, NOPE_DIM)
# e4m3 carries 3 mantissa bits, so the worst case is ~2^-4 of the tile's
# own absmax. Anything beyond that means the scale or the payload is off,
# not rounding.
tile_absmax = (
ref_nope.reshape(self.T, NUM_TILES, DSV4_FP8_QUANT_TILE)
.abs()
.amax(-1)
.repeat_interleave(DSV4_FP8_QUANT_TILE, dim=1)
)
self.assertTrue(torch.all((deq - ref_nope).abs() <= 0.07 * tile_absmax))
def test_rope_pool_matches_the_bf16_reference(self):
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
self._call(nope_pool, rope_pool)
_, ref_pe = _ref_norm_rope(
self.kv, self.weight, self.cos, self.sin, self.positions
)
got = rope_pool[self.swa_loc.long()].float()
torch.testing.assert_close(got, ref_pe, rtol=2e-2, atol=2e-2)
def test_q_stays_bf16_and_rotated(self):
q_out = self._call()
self.assertEqual(q_out.dtype, torch.bfloat16)
self.assertEqual(tuple(q_out.shape), (self.T, NUM_HEADS, HEAD_DIM))
head = self.q.view(self.T, NUM_HEADS, HEAD_DIM)[:, 0]
ones = torch.ones(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
ref_nope, ref_pe = _ref_norm_rope(
head, ones, self.cos, self.sin, self.positions
)
got = q_out[:, 0].float()
torch.testing.assert_close(got[:, :NOPE_DIM], ref_nope, rtol=2e-2, atol=2e-2)
torch.testing.assert_close(got[:, NOPE_DIM:], ref_pe, rtol=2e-2, atol=2e-2)
def test_q_is_packed_like_k_when_a_rope_buffer_is_given(self):
"""the v4 nm asm reader wants Q in the same 512 B form as the pool rows
Pinned against the bf16 Q the same call produces without the rope buffer,
so this is the quantization of a known-good rotated Q rather than a
second reimplementation of norm+rope.
"""
q_packed = torch.empty(
self.T,
NUM_HEADS,
DSV4_FP8_NOPE_ROW_BYTES,
dtype=torch.float8_e4m3fn,
device=DEVICE,
)
q_rope = torch.empty(
self.T, NUM_HEADS, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE
)
got = self._call(q_out=q_packed, q_rope_out=q_rope)
self.assertIs(got, q_packed)
ref = self._call().float() # bf16 Q, same input
ref_nope, ref_pe = ref[..., :NOPE_DIM], ref[..., NOPE_DIM:]
raw = q_packed.view(torch.uint8)
exp = _ref_tile_scales(ref_nope.reshape(-1, NOPE_DIM)).reshape(
self.T, NUM_HEADS, NUM_TILES
)
scale_bytes = raw[..., SCALE_OFF : SCALE_OFF + SCALE_BYTES].reshape(
self.T, NUM_HEADS, NUM_TILES, 2
)
torch.testing.assert_close(
scale_bytes[..., 0].int() - 127, exp.int(), rtol=0, atol=0
)
self.assertTrue(torch.equal(scale_bytes[..., 0], scale_bytes[..., 1]))
scale = torch.exp2(scale_bytes[..., 0].float() - 127)
dq = (
raw[..., :NOPE_DIM]
.view(torch.float8_e4m3fn)
.float()
.reshape(self.T, NUM_HEADS, NUM_TILES, DSV4_FP8_QUANT_TILE)
* scale.unsqueeze(-1)
).reshape(self.T, NUM_HEADS, NOPE_DIM)
# atol is half an fp8 step at the *tile's* absmax, not a per-element
# relative error -- a small value sharing a tile with a large one carries
# the large one's step. Measured 0.125 worst case at absmax 3.9, and the
# parts that must be exact (scale bytes above, rope below) are pinned as
# such.
torch.testing.assert_close(dq, ref_nope, rtol=5e-2, atol=5e-2)
torch.testing.assert_close(q_rope.float(), ref_pe, rtol=0, atol=0)
def test_fp8_q_without_a_rope_buffer_is_rejected(self):
q_packed = torch.empty(
self.T,
NUM_HEADS,
DSV4_FP8_NOPE_ROW_BYTES,
dtype=torch.float8_e4m3fn,
device=DEVICE,
)
with self.assertRaises(AssertionError):
self._call(q_out=q_packed)
def test_strided_q_out_is_filled_without_touching_the_padding(self):
"""attn_tp_size > 1 hands us a slice of a head-padded [T, 64, D] buffer
The zero-init is this test's way of seeing whether the staging copy strays
outside the slice. gfx950 allocates that buffer with new_empty, so in
production the padding holds garbage, not zeros -- what matters is only that
nobody writes it.
"""
padded = torch.zeros(self.T, 64, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
q_out = padded[:, :NUM_HEADS, :]
self.assertFalse(q_out.is_contiguous())
packed = self._call()
got = fused_qk_norm_rope_swa_store(
q=self.q,
kv=self.kv,
q_norm_weight=None,
kv_norm_weight=self.weight,
q_rms_eps=EPS,
kv_rms_eps=EPS,
rope_head_dim=ROPE_DIM,
cos_cache=self.cos,
sin_cache=self.sin,
positions=self.positions,
q_out=q_out,
dtype=torch.bfloat16,
fp8_2buff=True,
)
self.assertIs(got, q_out)
self.assertTrue(torch.all(padded[:, NUM_HEADS:, :] == 0))
# staging must not reorder the heads, so the strided destination has to
# hold exactly what the contiguous call produced
self.assertTrue(torch.equal(q_out, packed))
def test_negative_position_skips_both_pools(self):
"""a stale/pad token must leave both pools alone, not half-write a row"""
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
self.positions[2] = -1
self._call(nope_pool, rope_pool)
row = self.swa_loc[2].item()
self.assertEqual(nope_pool.view(torch.uint8)[row].max().item(), 0)
self.assertEqual(rope_pool[row].abs().max().item(), 0)
def test_rope_pool_is_required_with_the_nope_pool(self):
nope_pool, _ = _pools(2 * RING_STRIDE)
with self.assertRaises(AssertionError):
self._call(nope_pool, None)
def test_mismatched_pools_are_rejected_before_the_launch(self):
"""aiter aborts the process on a short pool, so these must fail in python"""
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
short_rope = rope_pool[:RING_STRIDE].contiguous()
cases = {
"fewer rope rows": (nope_pool, short_rope),
"rope dtype": (nope_pool, rope_pool.to(torch.float16)),
"rope width": (nope_pool, rope_pool[:, : ROPE_DIM // 2].contiguous()),
"nope row bytes": (nope_pool[:, :NOPE_DIM].contiguous(), rope_pool),
}
for name, (nope, rope) in cases.items():
with self.subTest(name), self.assertRaises(AssertionError):
self._call(nope, rope)
def test_bf16_store_is_a_different_store(self):
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
with self.assertRaises(AssertionError):
fused_qk_norm_rope_swa_store(
q=self.q,
kv=self.kv,
q_norm_weight=None,
kv_norm_weight=self.weight,
q_rms_eps=EPS,
kv_rms_eps=EPS,
rope_head_dim=ROPE_DIM,
cos_cache=self.cos,
sin_cache=self.sin,
positions=self.positions,
swa_cache=nope_pool,
swa_loc=self.swa_loc,
swa_page_size=1,
dtype=torch.bfloat16,
bf16_store=True,
fp8_2buff=True,
swa_rope_cache=rope_pool,
)
class TestBf16StoreStillWorks(_StoreCase):
"""fp8_2buff returns before the Triton kernel, so pin the branch it skips"""
def test_bf16_store_writes_the_whole_row(self):
pool = torch.zeros(
2 * RING_STRIDE, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
)
ref_nope, ref_pe = _ref_norm_rope(
self.kv, self.weight, self.cos, self.sin, self.positions
)
q_out = fused_qk_norm_rope_swa_store(
q=self.q,
kv=self.kv,
q_norm_weight=None,
kv_norm_weight=self.weight,
q_rms_eps=EPS,
kv_rms_eps=EPS,
rope_head_dim=ROPE_DIM,
cos_cache=self.cos,
sin_cache=self.sin,
positions=self.positions,
swa_cache=pool,
swa_loc=self.swa_loc,
swa_page_size=1,
dtype=torch.bfloat16,
bf16_store=True,
)
self.assertEqual(q_out.dtype, torch.bfloat16)
self.assertEqual(tuple(q_out.shape), (self.T, NUM_HEADS, HEAD_DIM))
rows = self.swa_loc.long()
got = pool[rows].float()
torch.testing.assert_close(got[:, :NOPE_DIM], ref_nope, rtol=2e-2, atol=2e-2)
torch.testing.assert_close(got[:, NOPE_DIM:], ref_pe, rtol=2e-2, atol=2e-2)
untouched = torch.ones(pool.shape[0], dtype=torch.bool, device=DEVICE)
untouched[rows] = False
self.assertEqual(pool[untouched].abs().max().item(), 0)
@unittest.skipUnless(_HAS_GROUP_QUANT, "needs aiter's group-quant kernel on gfx95x")
class TestUnifiedFp8SwaRingWrap(CustomTestCase):
"""What the ring holds once a slot gets written a second time.
The two pools have to turn over together. A row whose nope came from the new
token but whose rope is still the old one decodes against the wrong angle,
and nothing downstream can notice -- both halves are individually
well-formed.
Wrap is driven across calls, not inside one. Within a launch the
out-of-window tokens carry loc -1 and get skipped, so every live row has a
single writer; two writers to one row in one launch would be a race with no
defined winner to assert on.
"""
def setUp(self):
torch.manual_seed(11)
self.T = 6
self.cos, self.sin = _cos_sin()
self.weight = torch.randn(HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
self.q = torch.randn(
self.T, NUM_HEADS * HEAD_DIM, device=DEVICE, dtype=torch.bfloat16
)
def _store(self, kv, positions, swa_loc, nope_pool, rope_pool):
"""one launch; hands back the dense K pair as the per-token truth"""
k_nope = torch.empty(
self.T, 1, DSV4_FP8_NOPE_ROW_BYTES, dtype=torch.float8_e4m3fn, device=DEVICE
)
k_rope = torch.empty(self.T, 1, ROPE_DIM, dtype=torch.bfloat16, device=DEVICE)
fused_qk_norm_rope_swa_store(
q=self.q,
kv=kv,
q_norm_weight=None,
kv_norm_weight=self.weight,
q_rms_eps=EPS,
kv_rms_eps=EPS,
rope_head_dim=ROPE_DIM,
cos_cache=self.cos,
sin_cache=self.sin,
positions=positions,
swa_cache=nope_pool,
swa_loc=swa_loc,
swa_page_size=1,
dtype=torch.bfloat16,
fp8_2buff=True,
swa_rope_cache=rope_pool,
k_nope_out=k_nope,
k_rope_out=k_rope,
)
return k_nope.view(torch.uint8)[:, 0, : SCALE_OFF + SCALE_BYTES], k_rope[:, 0]
def _pass(self, step, nope_pool, rope_pool, count=None):
"""step 0 fills the ring, step 1 comes back around onto the same slots"""
count = self.T if count is None else count
kv = torch.randn(self.T, HEAD_DIM, device=DEVICE, dtype=torch.bfloat16)
positions = (
torch.arange(self.T, device=DEVICE, dtype=torch.int64) + step * RING_STRIDE
)
swa_loc = (positions.to(torch.int32) % RING_STRIDE + RING_STRIDE).contiguous()
# tokens past `count` fall out of window on this pass, like a short step
if count < self.T:
positions = positions.clone()
positions[count:] = -1
nope, rope = self._store(kv, positions, swa_loc, nope_pool, rope_pool)
return swa_loc.long(), nope.clone(), rope.clone()
def test_wrap_turns_over_both_pools(self):
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
rows, old_nope, _ = self._pass(0, nope_pool, rope_pool)
rows2, new_nope, new_rope = self._pass(1, nope_pool, rope_pool)
self.assertTrue(torch.equal(rows, rows2), "the wrap must reuse the same slots")
# only meaningful if pass 1 actually changed the bytes
self.assertFalse(torch.equal(old_nope, new_nope))
pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
self.assertTrue(torch.equal(pool_bytes, new_nope))
self.assertTrue(torch.equal(rope_pool[rows], new_rope))
def test_a_slot_the_wrap_skipped_keeps_its_old_pair(self):
"""a short second pass must leave the rows it didn't address alone"""
keep = 2
nope_pool, rope_pool = _pools(2 * RING_STRIDE)
rows, old_nope, old_rope = self._pass(0, nope_pool, rope_pool)
_, new_nope, new_rope = self._pass(1, nope_pool, rope_pool, count=keep)
pool_bytes = nope_pool.view(torch.uint8)[rows, : SCALE_OFF + SCALE_BYTES]
got_rope = rope_pool[rows]
self.assertTrue(torch.equal(pool_bytes[:keep], new_nope[:keep]))
self.assertTrue(torch.equal(got_rope[:keep], new_rope[:keep]))
self.assertTrue(torch.equal(pool_bytes[keep:], old_nope[keep:]))
self.assertTrue(torch.equal(got_rope[keep:], old_rope[keep:]))
if __name__ == "__main__":
unittest.main()
test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py
+267
View File
@@ -0,0 +1,267 @@
"""SWA ring scatter tests for the two-pool fp8 unified_kv layout.
``store_swa_into_unified`` writes one latent row per token. Under
SGLANG_DSV4_UNIFIED_KV_FP8 that row is split over a packed fp8 nope pool and a
bf16 rope pool, so what these tests pin is that the ring row index -- derived
from state_slot/positions alone -- stays identical to the bf16 layout's and
identical between the two pools.
"""
import unittest
import torch
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import runtime
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DSV4_FP8_NOPE_ROW_BYTES
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# the scatter itself is a plain row move, but the layout it pins is gfx95-only, so
# run it where the feature lives rather than on the default mi300 runner
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
DEVICE = torch.device("cuda")
# 448 values + 14 E8M0 scales + 50 pad, in bytes
NOPE_ROW_BYTES = DSV4_FP8_NOPE_ROW_BYTES
ROPE_DIM = 64
# V4-Pro latent, in elements -- same number as NOPE_ROW_BYTES, different unit
BF16_LATENT = 448 + ROPE_DIM
RING_STRIDE = 16
WIN = 8
N_PAGES = 64
def _inputs(n_rows=12):
"""state_slot/positions whose ring rows are all distinct, so a row's writer is unambiguous"""
state_slot = torch.tensor(
[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3][:n_rows],
device=DEVICE,
dtype=torch.int32,
)
positions = torch.tensor(
[0, 1, 2, 16, 17, 18, 32, 33, 34, 48, 49, 50][:n_rows],
device=DEVICE,
dtype=torch.int32,
)
return state_slot, positions
def _expected_rows(state_slot, positions, final_pos=None):
loc = state_slot.long() * RING_STRIDE + positions.long() % RING_STRIDE
if final_pos is None:
keep = torch.ones_like(loc, dtype=torch.bool)
else:
keep = positions.long() > final_pos.long() - WIN
return loc, keep
def _packed_nope(n_rows):
"""random packed fp8 rows; byte 0 is forced nonzero so a written row is detectable"""
raw = torch.randint(
0, 256, (n_rows, NOPE_ROW_BYTES), device=DEVICE, dtype=torch.uint8
)
raw[:, 0] = torch.arange(1, n_rows + 1, device=DEVICE, dtype=torch.uint8)
return raw.view(torch.float8_e4m3fn), raw
def _bf16_rope(n_rows):
rope = torch.randn(n_rows, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
rope[:, 0] = torch.arange(1, n_rows + 1, device=DEVICE, dtype=torch.bfloat16)
return rope.contiguous()
def _store(kv, pool, state_slot, positions, final_pos=None, **kw):
runtime.store_swa_into_unified(
kv=kv,
state_slot=state_slot,
positions=positions,
unified_kv=pool,
win=WIN,
ring_stride=RING_STRIDE,
final_pos=final_pos,
**kw,
)
class TestUnifiedFp8SwaScatter(CustomTestCase):
def setUp(self):
torch.manual_seed(20)
self.state_slot, self.positions = _inputs()
self.n_rows = self.state_slot.shape[0]
def _run_two_pool(self, final_pos=None):
kv_nope, nope_bytes = _packed_nope(self.n_rows)
kv_rope = _bf16_rope(self.n_rows)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
_store(
kv_nope,
pool_nope,
self.state_slot,
self.positions,
final_pos=final_pos,
kv_rope=kv_rope,
unified_kv_rope=pool_rope,
)
return pool_nope, pool_rope, nope_bytes, kv_rope
def _run_bf16(self, final_pos=None):
kv = torch.randn(
self.n_rows, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16
).contiguous()
kv[:, 0] = torch.arange(1, self.n_rows + 1, device=DEVICE, dtype=torch.bfloat16)
pool = torch.zeros(N_PAGES, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16)
_store(kv, pool, self.state_slot, self.positions, final_pos=final_pos)
return pool, kv
def test_bf16_single_pool_unchanged(self):
"""the bf16 path still writes exactly the expected ring rows"""
pool, kv = self._run_bf16()
loc, keep = _expected_rows(self.state_slot, self.positions)
expected = torch.zeros_like(pool)
expected[loc[keep]] = kv[keep]
self.assertTrue(torch.equal(pool, expected))
def test_two_pool_bytes_exact(self):
"""each pool gets its half verbatim -- nope byte-for-byte, rope bit-for-bit"""
pool_nope, pool_rope, nope_bytes, kv_rope = self._run_two_pool()
loc, keep = _expected_rows(self.state_slot, self.positions)
exp_nope = torch.zeros_like(pool_nope).view(torch.uint8)
exp_nope[loc[keep]] = nope_bytes[keep]
self.assertTrue(torch.equal(pool_nope.view(torch.uint8), exp_nope))
exp_rope = torch.zeros_like(pool_rope)
exp_rope[loc[keep]] = kv_rope[keep]
self.assertTrue(torch.equal(pool_rope, exp_rope))
def test_two_pool_rows_match_bf16(self):
"""same state_slot/positions -> same ring rows as bf16, and the same in both pools"""
pool_nope, pool_rope, _, _ = self._run_two_pool()
pool_bf16, _ = self._run_bf16()
rows_nope = (pool_nope.view(torch.uint8) != 0).any(dim=1)
rows_rope = (pool_rope != 0).any(dim=1)
rows_bf16 = (pool_bf16 != 0).any(dim=1)
self.assertTrue(torch.equal(rows_nope, rows_bf16))
self.assertTrue(torch.equal(rows_rope, rows_bf16))
self.assertEqual(int(rows_bf16.sum()), self.n_rows)
def test_final_pos_skips_both_pools(self):
"""tokens already outside the window are skipped in nope and rope alike"""
# positions[t] <= final_pos[t] - WIN skips; give the first half a far
# final_pos and the second half its own position
final_pos = self.positions.clone()
final_pos[: self.n_rows // 2] = self.positions.max() + WIN
pool_nope, pool_rope, nope_bytes, kv_rope = self._run_two_pool(
final_pos=final_pos
)
loc, keep = _expected_rows(self.state_slot, self.positions, final_pos)
self.assertTrue(bool((~keep).any()), "test would be vacuous without a skip")
rows_nope = (pool_nope.view(torch.uint8) != 0).any(dim=1)
rows_rope = (pool_rope != 0).any(dim=1)
expected_rows = torch.zeros(N_PAGES, device=DEVICE, dtype=torch.bool)
expected_rows[loc[keep]] = True
self.assertTrue(torch.equal(rows_nope, expected_rows))
self.assertTrue(torch.equal(rows_rope, expected_rows))
def test_rope_tensor_and_pool_come_together(self):
kv_nope, _ = _packed_nope(self.n_rows)
kv_rope = _bf16_rope(self.n_rows)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
with self.assertRaises(AssertionError):
_store(
kv_nope,
pool_nope,
self.state_slot,
self.positions,
kv_rope=kv_rope,
)
with self.assertRaises(AssertionError):
_store(
kv_nope,
pool_nope,
self.state_slot,
self.positions,
unified_kv_rope=pool_rope,
)
def test_short_rope_pool_rejected(self):
"""the kernel doesn't bound-check the ring row, so a rope pool with fewer
rows than the nope pool writes into whatever tensor follows it"""
kv_nope, _ = _packed_nope(self.n_rows)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
# ring rows reach state_slot 3 -> row 48, well past this
pool_rope = torch.zeros(8, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
with self.assertRaises(AssertionError):
_store(
kv_nope,
pool_nope,
self.state_slot,
self.positions,
kv_rope=_bf16_rope(self.n_rows),
unified_kv_rope=pool_rope,
)
def test_rope_row_width_mismatch_rejected(self):
"""row width is read off src, so a wider pool would place row i at i * D"""
kv_nope, _ = _packed_nope(self.n_rows)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
pool_rope = torch.zeros(
N_PAGES, ROPE_DIM * 2, device=DEVICE, dtype=torch.bfloat16
)
with self.assertRaises(AssertionError):
_store(
kv_nope,
pool_nope,
self.state_slot,
self.positions,
kv_rope=_bf16_rope(self.n_rows),
unified_kv_rope=pool_rope,
)
def test_dtype_mismatch_rejected(self):
"""a bf16 row must not land in an fp8 pool (the DSpark-under-fp8 case)"""
kv = torch.randn(
self.n_rows, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.bfloat16
)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
with self.assertRaises(AssertionError):
_store(kv, pool_nope, self.state_slot, self.positions)
def test_empty_batch_is_a_noop(self):
empty_slot = torch.zeros(0, device=DEVICE, dtype=torch.int32)
kv_nope, _ = _packed_nope(0)
pool_nope = torch.zeros(
N_PAGES, NOPE_ROW_BYTES, device=DEVICE, dtype=torch.float8_e4m3fn
)
pool_rope = torch.zeros(N_PAGES, ROPE_DIM, device=DEVICE, dtype=torch.bfloat16)
_store(
kv_nope,
pool_nope,
empty_slot,
empty_slot,
kv_rope=_bf16_rope(0),
unified_kv_rope=pool_rope,
)
self.assertEqual(int((pool_nope.view(torch.uint8) != 0).sum()), 0)
self.assertEqual(int((pool_rope != 0).sum()), 0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py
+128
View File
@@ -0,0 +1,128 @@
import contextlib
import unittest
import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DSV4_FP8_NOPE_ROW_BYTES,
DSV4_FP8_QUANT_TILE,
DeepSeekV4UnifiedKVPool,
dsv4_unified_row_bytes,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
# DeepSeek-V4-Pro geometry.
NOPE_DIM = 448
ROPE_DIM = 64
class _StubMemorySaver:
def region(self, _tag):
return contextlib.nullcontext()
class TestDSV4UnifiedRowBytes(CustomTestCase):
"""Row width drives both `bytes_per_full_token` and `_fixed_swa_bytes`, so the
capacity claim for the fp8 pool is only as good as this arithmetic."""
def test_bf16_row_is_the_whole_latent(self):
self.assertEqual(
dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False),
(NOPE_DIM + ROPE_DIM) * 2,
)
def test_fp8_row_is_padded_nope_plus_bf16_rope(self):
self.assertEqual(
dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True),
DSV4_FP8_NOPE_ROW_BYTES + ROPE_DIM * 2,
)
def test_fp8_saves_exactly_three_eighths(self):
"""0.625x is where the >=1.40x capacity target comes from; the remaining
dilution is the fixed SWA/c4-state bias, not the row."""
bf16 = dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False)
fp8 = dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True)
self.assertEqual((bf16, fp8), (1024, 640))
self.assertAlmostEqual(fp8 / bf16, 0.625)
def test_scales_and_latent_fit_the_asm_stride(self):
"""7 tiles written twice = 14 B; 448 + 14 leaves 50 B the reader never
touches. If a future head_dim broke this the pack would silently overlap."""
num_tiles = NOPE_DIM // DSV4_FP8_QUANT_TILE
self.assertEqual(num_tiles, 7)
self.assertLessEqual(NOPE_DIM + 2 * num_tiles, DSV4_FP8_NOPE_ROW_BYTES)
def test_oversized_latent_is_rejected(self):
# ValueError, not assert: sizing has to keep checking under python -O
with self.assertRaises(ValueError):
dsv4_unified_row_bytes(DSV4_FP8_NOPE_ROW_BYTES, ROPE_DIM, fp8=True)
class TestDSV4UnifiedFp8PoolAllocation(CustomTestCase):
"""The sizing formula and the allocation are two separate code paths; this pins
them to the same row width so a change to one cannot silently outrun the other."""
STAGE_RATIOS = [4, 128]
NUM_SLOTS = 3
NUM_BLOCKS = 5
PAGE_SIZE = 256
SWA_RING = 8
def _pool(self, fp8):
return DeepSeekV4UnifiedKVPool(
stage_ratios=self.STAGE_RATIOS,
num_slots=self.NUM_SLOTS,
num_blocks=self.NUM_BLOCKS,
page_size=self.PAGE_SIZE,
qk_nope_head_dim=NOPE_DIM,
qk_rope_head_dim=ROPE_DIM,
device="cpu",
memory_saver_adapter=_StubMemorySaver(),
custom_mem_pool=None,
swa_ring_size=self.SWA_RING,
fp8=fp8,
)
def test_bf16_pool_is_unchanged(self):
"""fp8 defaults off, so the bf16 arm must keep one pool of bf16 latents."""
pool = self._pool(fp8=False)
for buf, rope in zip(pool.kv_buffer, pool.kv_buffer_rope):
self.assertEqual(buf.dtype, torch.bfloat16)
self.assertEqual(buf.shape[1], NOPE_DIM + ROPE_DIM)
self.assertIsNone(rope)
def test_fp8_pool_row_counts_match_across_both_pools(self):
"""A row index addresses the SWA ring and the compressed region in both
pools, so the two must have identical row counts."""
pool = self._pool(fp8=True)
for buf, rope in zip(pool.kv_buffer, pool.kv_buffer_rope):
self.assertEqual(buf.dtype, torch.float8_e4m3fn)
self.assertEqual(rope.dtype, torch.bfloat16)
self.assertEqual(buf.shape[0], rope.shape[0])
self.assertEqual(buf.shape[1], DSV4_FP8_NOPE_ROW_BYTES)
self.assertEqual(rope.shape[1], ROPE_DIM)
def test_fp8_pool_bytes_match_the_sizing_row_width(self):
bf16, fp8 = self._pool(fp8=False), self._pool(fp8=True)
for layer, buf in enumerate(bf16.kv_buffer):
rows = buf.shape[0]
self.assertEqual(fp8.kv_buffer[layer].shape[0], rows)
self.assertEqual(
buf.nbytes,
rows * dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=False),
)
self.assertEqual(
fp8.kv_buffer[layer].nbytes + fp8.kv_buffer_rope[layer].nbytes,
rows * dsv4_unified_row_bytes(NOPE_DIM, ROPE_DIM, fp8=True),
)
def test_rope_accessor_rejects_the_bf16_pool(self):
with self.assertRaises(AssertionError):
self._pool(fp8=False).get_unified_kv_rope(0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/test_pool_configurator.py
+3
View File
@@ -1118,6 +1118,9 @@ class TestSWAPoolFloor(CustomTestCase):
cfg.disaggregation_mode = None
cfg.disaggregation_decode_extra_slots = 0
cfg._unified = True
cfg._unified_fp8 = False
# object.__new__ skips __init__; bf16 unified row is 2B * latent
cfg._unified_row_bytes = cfg.attn_head_dim * 2
return cfg
# Token pool plus the three request-scoped fixed pools, sized from the
test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py
+279
View File
@@ -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()