minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 16 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