[AMD] Add unified kv attention support in dpsk-v4 (#27380)
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
This commit is contained in:
Thomas Wang
co-authored by
Xinyi Song
Xinyi Song
parent
95d8a75bc9
commit
f2bcdb0508
@@ -368,7 +368,7 @@ INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNor
|
||||
// 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>
|
||||
template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
|
||||
FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
|
||||
using namespace device;
|
||||
using enum ForwardMode;
|
||||
@@ -379,7 +379,10 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
|
||||
// Last warp owns the rope tail. The remaining 7 warps each emit one
|
||||
// 64-element fp8 group (own UE8M0 scale).
|
||||
constexpr uint32_t kRopeWarp = kNumWarps - 1;
|
||||
constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
|
||||
// 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
|
||||
constexpr int64_t kPageBytes =
|
||||
kBf16Store ? ((kHeadDim * 2ll) << 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);
|
||||
@@ -450,12 +453,23 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
|
||||
const int32_t page = out_loc >> kPageBits;
|
||||
const int32_t offset = out_loc & ((1 << kPageBits) - 1);
|
||||
const auto page_ptr = params.kvcache + page * kPageBytes;
|
||||
const auto value_ptr = page_ptr + offset * 576;
|
||||
const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
|
||||
if (warp_id == kRopeWarp) {
|
||||
if constexpr (kBf16Store) {
|
||||
Float2 d = data;
|
||||
if (warp_id == kRopeWarp) {
|
||||
const auto x_real = data[0];
|
||||
const auto x_imag = data[1];
|
||||
const auto freq_real = freq[0];
|
||||
const auto freq_imag = freq[1];
|
||||
d[0] = x_real * freq_real - x_imag * freq_imag;
|
||||
d[1] = x_real * freq_imag + x_imag * freq_real;
|
||||
}
|
||||
reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
|
||||
} else if (warp_id == kRopeWarp) {
|
||||
// Each rope-warp lane owns exactly one (real, imag) pair within the rope
|
||||
// tail. Apply rotation, downcast to BF16, write to the slot's rope region.
|
||||
const auto x_real = data[0];
|
||||
@@ -485,13 +499,15 @@ FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormR
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
|
||||
template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL, bool kBf16Store = false>
|
||||
struct FusedNormRopeKernel {
|
||||
static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
|
||||
static constexpr bool kIsIndexer = (kHeadDim == 128);
|
||||
static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
|
||||
static constexpr int64_t kIndexerBytes = 132 * kPageSize;
|
||||
static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
|
||||
static constexpr int64_t kPageBytes = kIsIndexer ? kIndexerBytes : kFlashMLABytes;
|
||||
static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize; // plain bf16 cache
|
||||
static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
|
||||
|
||||
/// TODO: Let's fix the config for now.
|
||||
static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
|
||||
@@ -502,7 +518,7 @@ struct FusedNormRopeKernel {
|
||||
if constexpr (kIsIndexer) {
|
||||
return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
|
||||
} else {
|
||||
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL>;
|
||||
return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -23,8 +23,11 @@ def _jit_compress_norm_rope_module(
|
||||
head_dim: int,
|
||||
rope_dim: int,
|
||||
page_size: int,
|
||||
bf16_store: bool = False,
|
||||
) -> Module:
|
||||
args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl())
|
||||
args = make_cpp_args(
|
||||
dtype, head_dim, rope_dim, page_size, is_arch_support_pdl(), bf16_store
|
||||
)
|
||||
cuda_wrappers = [("forward", f"FusedNormRopeKernel<{args}>::forward")]
|
||||
if head_dim == 128:
|
||||
cuda_wrappers.append(
|
||||
@@ -339,12 +342,13 @@ def compress_norm_rope_store(
|
||||
kvcache: torch.Tensor,
|
||||
page_size: int,
|
||||
use_fp4: bool = False,
|
||||
bf16_store: bool = False,
|
||||
) -> None:
|
||||
if use_fp4:
|
||||
assert kv.shape[-1] == 128
|
||||
freq_cis = torch.view_as_real(freq_cis).flatten(-2)
|
||||
module = _jit_compress_norm_rope_module(
|
||||
kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size
|
||||
kv.dtype, kv.shape[-1], freq_cis.shape[-1], page_size, bf16_store
|
||||
)
|
||||
fn = module.forward_fp4 if use_fp4 else module.forward
|
||||
fn(
|
||||
|
||||
@@ -238,6 +238,7 @@ class DSV4AttnMetadata:
|
||||
self.c4_topk_lengths_clamp1,
|
||||
self.c128_out_loc,
|
||||
_,
|
||||
_,
|
||||
self.c128_topk_lengths_clamp1,
|
||||
self.c128_page_indices,
|
||||
) = _init_compression_metadata_triton(
|
||||
|
||||
@@ -113,11 +113,28 @@ class DSV4AttnMetadata:
|
||||
c4_topk_lengths_raw: Optional[torch.Tensor] = None
|
||||
c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
|
||||
c4_sparse_topk_lengths: torch.Tensor = field(init=False)
|
||||
c4_sparse_topk_lengths_raw: torch.Tensor = field(init=False)
|
||||
c4_sparse_page_indices: torch.Tensor = field(init=False)
|
||||
c4_sparse_raw_indices: Optional[torch.Tensor] = field(init=False, default=None)
|
||||
|
||||
c128_out_loc: Optional[torch.Tensor] = None
|
||||
c128_page_indices: Optional[torch.Tensor] = None
|
||||
c128_topk_lengths_clamp1: Optional[torch.Tensor] = None
|
||||
c128_topk_lengths_raw: Optional[torch.Tensor] = None
|
||||
|
||||
# unified_kv: per-forward prebuilt ragged decode index
|
||||
unified_swa_indices: Optional[torch.Tensor] = None
|
||||
unified_swa_indptr: Optional[torch.Tensor] = None
|
||||
unified_hca_indices: Optional[torch.Tensor] = None
|
||||
unified_hca_indptr: Optional[torch.Tensor] = None
|
||||
unified_csa_indices: Optional[torch.Tensor] = None
|
||||
unified_csa_indptr: Optional[torch.Tensor] = None
|
||||
|
||||
# unified_kv: per-forward prefill/extend per-token mapping
|
||||
unified_pf_state_slot: Optional[torch.Tensor] = None
|
||||
unified_pf_chunk_start: Optional[torch.Tensor] = None
|
||||
unified_pf_cu_q: Optional[torch.Tensor] = None
|
||||
unified_pf_final_pos: Optional[torch.Tensor] = None
|
||||
|
||||
c1_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
|
||||
c4_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
|
||||
@@ -157,10 +174,23 @@ class DSV4AttnMetadata:
|
||||
"swa_topk_lengths",
|
||||
"c128_page_indices",
|
||||
"c128_topk_lengths_clamp1",
|
||||
"c128_topk_lengths_raw",
|
||||
"c4_topk_lengths_raw",
|
||||
"c4_topk_lengths_clamp1",
|
||||
"c4_sparse_topk_lengths",
|
||||
"c4_sparse_topk_lengths_raw",
|
||||
"c4_sparse_page_indices",
|
||||
"c4_sparse_raw_indices",
|
||||
"unified_swa_indices",
|
||||
"unified_swa_indptr",
|
||||
"unified_hca_indices",
|
||||
"unified_hca_indptr",
|
||||
"unified_csa_indices",
|
||||
"unified_csa_indptr",
|
||||
"unified_pf_state_slot",
|
||||
"unified_pf_chunk_start",
|
||||
"unified_pf_cu_q",
|
||||
"unified_pf_final_pos",
|
||||
],
|
||||
assign_fields=[
|
||||
# Recomputed by the recorded init_forward_metadata_in_graph op
|
||||
@@ -185,6 +215,7 @@ class DSV4AttnMetadata:
|
||||
self.c4_topk_lengths_clamp1,
|
||||
self.c128_out_loc,
|
||||
_,
|
||||
self.c128_topk_lengths_raw,
|
||||
self.c128_topk_lengths_clamp1,
|
||||
self.c128_page_indices,
|
||||
) = _init_compression_metadata_triton(
|
||||
@@ -209,6 +240,7 @@ class DSV4AttnMetadata:
|
||||
"c4_topk_lengths_clamp1",
|
||||
"c128_page_indices",
|
||||
"c128_topk_lengths_clamp1",
|
||||
"c128_topk_lengths_raw",
|
||||
]
|
||||
_CP_GLOBAL_FIELDS = [
|
||||
"raw_out_loc",
|
||||
@@ -260,6 +292,10 @@ class DSV4AttnMetadata:
|
||||
self.c4_sparse_topk_lengths = torch.clamp(
|
||||
self.c4_topk_lengths_clamp1, max=self.c4_sparse_topk
|
||||
)
|
||||
assert self.c4_topk_lengths_raw is not None
|
||||
self.c4_sparse_topk_lengths_raw = torch.clamp(
|
||||
self.c4_topk_lengths_raw, max=self.c4_sparse_topk
|
||||
)
|
||||
self.c4_sparse_page_indices = torch.full(
|
||||
(self.c4_topk_lengths_clamp1.size(0), self.c4_sparse_topk),
|
||||
-1,
|
||||
@@ -427,6 +463,7 @@ class DeepseekV4HipRadixBackend(
|
||||
out_loc=out_cache_loc,
|
||||
need_compress=True,
|
||||
)
|
||||
self._attach_unified_kv_decode_streams(core_attn_metadata, req_pool_indices)
|
||||
|
||||
indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
|
||||
|
||||
@@ -475,6 +512,9 @@ class DeepseekV4HipRadixBackend(
|
||||
need_compress=need_compress,
|
||||
is_prefill=True,
|
||||
)
|
||||
self._attach_unified_kv_prefill_meta(
|
||||
core_attn_metadata, req_pool_indices, seq_lens, extend_seq_lens
|
||||
)
|
||||
indexer_metadata = (
|
||||
self.init_forward_metadata_indexer(core_attn_metadata)
|
||||
if need_compress
|
||||
@@ -621,6 +661,7 @@ class DeepseekV4HipRadixBackend(
|
||||
out_loc=out_cache_loc,
|
||||
need_compress=True,
|
||||
)
|
||||
self._attach_unified_kv_decode_streams(core_attn_metadata, req_pool_indices)
|
||||
indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
|
||||
|
||||
create = functools.partial(
|
||||
@@ -954,6 +995,191 @@ class DeepseekV4HipRadixBackend(
|
||||
if current_raw is not None:
|
||||
self.forward_metadata = current_raw
|
||||
|
||||
def _attach_unified_kv_decode_streams(
|
||||
self, core: "DSV4AttnMetadata", req_pool_indices: torch.Tensor
|
||||
) -> None:
|
||||
"""build the ragged decode index streams once per forward"""
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
if not is_unified_kv_triton():
|
||||
return
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime
|
||||
|
||||
pool = self.token_to_kv_pool
|
||||
N = core.positions_casual.shape[0]
|
||||
(
|
||||
core.unified_swa_indices,
|
||||
core.unified_swa_indptr,
|
||||
core.unified_hca_indices,
|
||||
core.unified_hca_indptr,
|
||||
core.unified_csa_indices,
|
||||
core.unified_csa_indptr,
|
||||
) = runtime.build_decode_streams(
|
||||
state_slot=req_pool_indices[:N],
|
||||
positions=core.positions_casual,
|
||||
swa_len=core.swa_topk_lengths,
|
||||
hca_len=core.c128_topk_lengths_raw,
|
||||
csa_len=core.c4_sparse_topk_lengths_raw,
|
||||
hca_page_indices=core.c128_page_indices,
|
||||
csa_width=core.c4_sparse_page_indices.shape[1],
|
||||
win=pool.unified_swa_window,
|
||||
ring_stride=pool.unified_swa_ring_size,
|
||||
swa_pages=pool.unified_swa_pages,
|
||||
)
|
||||
|
||||
def _attach_unified_kv_prefill_meta(
|
||||
self,
|
||||
core: "DSV4AttnMetadata",
|
||||
req_pool_indices: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
extend_seq_lens: torch.Tensor,
|
||||
) -> None:
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
if not is_unified_kv_triton():
|
||||
return
|
||||
device = req_pool_indices.device
|
||||
bs = req_pool_indices.shape[0]
|
||||
seq_lens = seq_lens.to(torch.int64)
|
||||
extend_seq_lens = extend_seq_lens.to(torch.int64)
|
||||
# token -> req index (length L = sum(extend_seq_lens))
|
||||
bid = torch.repeat_interleave(
|
||||
torch.arange(bs, device=device, dtype=torch.int64), extend_seq_lens
|
||||
)
|
||||
core.unified_pf_state_slot = req_pool_indices[bid]
|
||||
core.unified_pf_chunk_start = (seq_lens - extend_seq_lens)[bid]
|
||||
cu_q_per_req = torch.cumsum(extend_seq_lens, dim=0) - extend_seq_lens
|
||||
core.unified_pf_cu_q = cu_q_per_req[bid]
|
||||
core.unified_pf_final_pos = (seq_lens - 1)[bid]
|
||||
|
||||
def _forward_unified_kv(
|
||||
self,
|
||||
*,
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
layer: RadixAttention,
|
||||
forward_batch: ForwardBatch,
|
||||
compress_ratio: Literal[0, 4, 128],
|
||||
attn_sink: torch.Tensor,
|
||||
core_attn_metadata: "DSV4AttnMetadata",
|
||||
save_kv_cache: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""unified_kv paged-attention path over the bf16 unified_kv"""
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels import runtime
|
||||
|
||||
pool = self.token_to_kv_pool
|
||||
layer_id = layer.layer_id
|
||||
unified = pool.get_unified_kv(layer_id)
|
||||
win = pool.unified_swa_window
|
||||
ring_stride = pool.unified_swa_ring_size
|
||||
swa_pages = pool.unified_swa_pages
|
||||
|
||||
if q.ndim == 4:
|
||||
q = q.squeeze(1)
|
||||
device = q.device
|
||||
positions = forward_batch.positions.to(torch.int64)
|
||||
T = q.shape[0]
|
||||
positions = positions[:T]
|
||||
|
||||
c128_pi = getattr(core_attn_metadata, "c128_page_indices", None)
|
||||
c4_pi = getattr(core_attn_metadata, "c4_sparse_page_indices", None)
|
||||
|
||||
# decode
|
||||
is_decode = forward_batch.forward_mode.is_decode_or_idle()
|
||||
if is_decode:
|
||||
state_slot = forward_batch.req_pool_indices[:T]
|
||||
if save_kv_cache:
|
||||
runtime.store_swa_into_unified(
|
||||
kv=kv,
|
||||
state_slot=state_slot,
|
||||
positions=positions,
|
||||
unified_kv=unified,
|
||||
win=win,
|
||||
ring_stride=ring_stride,
|
||||
final_pos=positions,
|
||||
)
|
||||
if compress_ratio == 0:
|
||||
kv_indices = core_attn_metadata.unified_swa_indices
|
||||
kv_indptr = core_attn_metadata.unified_swa_indptr
|
||||
elif compress_ratio == 128:
|
||||
kv_indices = core_attn_metadata.unified_hca_indices
|
||||
kv_indptr = core_attn_metadata.unified_hca_indptr
|
||||
elif compress_ratio == 4:
|
||||
kv_indices = core_attn_metadata.unified_csa_indices
|
||||
kv_indptr = core_attn_metadata.unified_csa_indptr
|
||||
runtime.fill_compress_tail(
|
||||
indices=kv_indices,
|
||||
indptr=kv_indptr,
|
||||
prefix_len=core_attn_metadata.swa_topk_lengths[:T],
|
||||
page_indices=c4_pi[:T],
|
||||
valid_len=core_attn_metadata.c4_sparse_topk_lengths_raw[:T],
|
||||
swa_pages=swa_pages,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"bad compress_ratio {compress_ratio}")
|
||||
return runtime.decode(
|
||||
q=q,
|
||||
unified_kv=unified,
|
||||
kv_indices=kv_indices,
|
||||
kv_indptr=kv_indptr,
|
||||
attn_sink=attn_sink,
|
||||
softmax_scale=self.softmax_scale,
|
||||
)
|
||||
|
||||
# prefill / extend
|
||||
state_slot = core_attn_metadata.unified_pf_state_slot
|
||||
chunk_start = core_attn_metadata.unified_pf_chunk_start
|
||||
cu_q = core_attn_metadata.unified_pf_cu_q
|
||||
final_pos = core_attn_metadata.unified_pf_final_pos
|
||||
|
||||
kpre_i, kpre_p, kext_i, kext_p = runtime.build_prefill_indices(
|
||||
compress_ratio=compress_ratio,
|
||||
state_slot=state_slot,
|
||||
positions=positions,
|
||||
chunk_start=chunk_start,
|
||||
cu_q=cu_q,
|
||||
win=win,
|
||||
ring_stride=ring_stride,
|
||||
swa_pages=swa_pages,
|
||||
c128_page_indices=c128_pi,
|
||||
c4_sparse_page_indices=c4_pi,
|
||||
)
|
||||
|
||||
if kpre_p.shape[0] < T + 1:
|
||||
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,
|
||||
)
|
||||
|
||||
# write this chunk's SWA K into the ring for future chunks / decode
|
||||
# only the final-window tokens per request
|
||||
if save_kv_cache:
|
||||
n_real = state_slot.shape[0]
|
||||
runtime.store_swa_into_unified(
|
||||
kv=kv[:n_real],
|
||||
state_slot=state_slot,
|
||||
positions=positions[:n_real],
|
||||
unified_kv=unified,
|
||||
win=win,
|
||||
ring_stride=ring_stride,
|
||||
final_pos=final_pos,
|
||||
)
|
||||
return o
|
||||
|
||||
def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
|
||||
"""Resolve the SWA KV-store write target for the current forward.
|
||||
|
||||
@@ -1022,6 +1248,22 @@ class DeepseekV4HipRadixBackend(
|
||||
token_to_kv_pool = self.token_to_kv_pool
|
||||
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
if is_unified_kv_triton():
|
||||
return self._forward_unified_kv(
|
||||
q=q,
|
||||
kv=swa_k,
|
||||
layer=layer,
|
||||
forward_batch=forward_batch,
|
||||
compress_ratio=compress_ratio,
|
||||
attn_sink=attn_sink,
|
||||
core_attn_metadata=core_attn_metadata,
|
||||
save_kv_cache=save_kv_cache,
|
||||
)
|
||||
|
||||
if isinstance(core_attn_metadata, DSV4AttnMetadata):
|
||||
if save_kv_cache:
|
||||
self.store_cache(layer_id, swa_k, forward_batch)
|
||||
@@ -1218,7 +1460,9 @@ class DeepseekV4HipRadixBackend(
|
||||
core_attn_metadata.init_flashmla_related()
|
||||
else:
|
||||
core_attn_metadata.c4_sparse_topk_lengths = None
|
||||
core_attn_metadata.c4_sparse_topk_lengths_raw = None
|
||||
core_attn_metadata.c4_sparse_page_indices = None
|
||||
core_attn_metadata.c4_sparse_raw_indices = None
|
||||
core_attn_metadata.c1_flashmla_metadata = _create_flashmla_metadata()
|
||||
core_attn_metadata.c4_flashmla_metadata = None
|
||||
core_attn_metadata.c128_flashmla_metadata = None
|
||||
|
||||
@@ -406,9 +406,9 @@ class Compressor(nn.Module):
|
||||
return ret
|
||||
|
||||
def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
|
||||
if _tgemm is not None:
|
||||
# linear_bf16_fp32 uses tgemm.mm + .float(); skip the .float() cast
|
||||
# because downstream Triton kernels promote bf16→fp32 internally.
|
||||
if _tgemm is not None and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
|
||||
# v1 compress goes through fused_compress_triton, which promotes
|
||||
# bf16->fp32 internally, so skip the .float() cast.
|
||||
kv_score = _tgemm.mm(x, self.wkv_gate.weight, otype=x.dtype)
|
||||
else:
|
||||
kv_score = linear_bf16_fp32(x, self.wkv_gate.weight)
|
||||
|
||||
@@ -428,6 +428,7 @@ class CompressorBackendMixin:
|
||||
page_size: int,
|
||||
out_loc: torch.Tensor,
|
||||
use_fp4_indexer: bool = False,
|
||||
bf16_store: bool = False,
|
||||
) -> None:
|
||||
assert compress_ratio == 4 or compress_ratio == 128
|
||||
assert rotate == is_indexer == (head_dim == 128)
|
||||
@@ -468,6 +469,7 @@ class CompressorBackendMixin:
|
||||
kvcache=kv_cache,
|
||||
page_size=page_size,
|
||||
use_fp4=use_fp4_indexer,
|
||||
bf16_store=bf16_store,
|
||||
)
|
||||
|
||||
def forward_unified(
|
||||
@@ -485,6 +487,10 @@ class CompressorBackendMixin:
|
||||
kv_score_input = compressor.compute_kv_score(x, forward_batch)
|
||||
|
||||
state_pool = compressor.get_state_pool(self)
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
if _is_hip and not envs.SGLANG_OPT_USE_JIT_NORM.get():
|
||||
self._forward_unified_hip(
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
@@ -498,9 +504,15 @@ class CompressorBackendMixin:
|
||||
use_fp4_indexer = (
|
||||
compressor.is_in_indexer and self.enable_deepseek_v4_fp4_indexer
|
||||
)
|
||||
bf16_store = False
|
||||
if compressor.is_in_indexer:
|
||||
kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_index_k_page_size()
|
||||
elif is_unified_kv_triton():
|
||||
kv_cache = token_to_kv_pool.get_unified_kv(layer_id)
|
||||
page_size = 1
|
||||
out_loc = out_loc + token_to_kv_pool.unified_swa_pages
|
||||
bf16_store = True
|
||||
else:
|
||||
_, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
|
||||
assert compress_kv_pool is not None
|
||||
@@ -527,6 +539,7 @@ class CompressorBackendMixin:
|
||||
page_size=page_size,
|
||||
out_loc=out_loc,
|
||||
use_fp4_indexer=use_fp4_indexer,
|
||||
bf16_store=bf16_store,
|
||||
)
|
||||
|
||||
def _forward_unified_hip(
|
||||
|
||||
@@ -17,6 +17,7 @@ def _init_compressed_attn_metadata_kernel(
|
||||
c4_seq_lens_clamp1_ptr,
|
||||
c128_out_loc_ptr,
|
||||
c128_positions_ptr,
|
||||
c128_seq_lens_raw_ptr,
|
||||
c128_seq_lens_clamp1_ptr,
|
||||
c128_page_indices_ptr,
|
||||
bs,
|
||||
@@ -54,6 +55,7 @@ def _init_compressed_attn_metadata_kernel(
|
||||
|
||||
tl.store(c128_out_loc_ptr + batch_id, c128_out_loc)
|
||||
tl.store(c128_positions_ptr + batch_id, c128_positions)
|
||||
tl.store(c128_seq_lens_raw_ptr + batch_id, c128_seq_lens_raw)
|
||||
tl.store(c128_seq_lens_clamp1_ptr + batch_id, c128_seq_lens_clamp1)
|
||||
|
||||
if COMPUTE_PAGE_INDICES:
|
||||
@@ -99,6 +101,7 @@ def _init_compressed_attn_metadata_triton(
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
Optional[torch.Tensor],
|
||||
]:
|
||||
bs = seq_lens.shape[0]
|
||||
@@ -111,6 +114,7 @@ def _init_compressed_attn_metadata_triton(
|
||||
|
||||
c128_out_loc = torch.empty(bs, dtype=torch.int32, device=device)
|
||||
c128_positions = torch.empty(bs, dtype=torch.int32, device=device)
|
||||
c128_seq_lens_raw = torch.empty(bs, dtype=torch.int32, device=device)
|
||||
c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
|
||||
|
||||
if compute_page_indices:
|
||||
@@ -146,6 +150,7 @@ def _init_compressed_attn_metadata_triton(
|
||||
c4_seq_lens_clamp1,
|
||||
c128_out_loc,
|
||||
c128_positions,
|
||||
c128_seq_lens_raw,
|
||||
c128_seq_lens_clamp1,
|
||||
(
|
||||
c128_page_indices
|
||||
@@ -168,6 +173,7 @@ def _init_compressed_attn_metadata_triton(
|
||||
c4_seq_lens_clamp1,
|
||||
c128_out_loc,
|
||||
c128_positions,
|
||||
c128_seq_lens_raw,
|
||||
c128_seq_lens_clamp1,
|
||||
c128_page_indices,
|
||||
)
|
||||
@@ -188,6 +194,7 @@ def init_compression_metadata(
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
torch.Tensor,
|
||||
Optional[torch.Tensor],
|
||||
]:
|
||||
return _init_compressed_attn_metadata_triton(
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import os
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def is_unified_kv_triton() -> bool:
|
||||
return os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "") == "unified_kv_triton"
|
||||
@@ -0,0 +1,881 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
"""Sparse decode attention over a unified KV pool with per-token paged indices.
|
||||
|
||||
Designed for V4 decode + CUDAGraph: replaces the per-fwd `kv_flat_sa`
|
||||
materialization (whose shape depends on `n_committed_per_seq` → varies per
|
||||
fwd → blocks CG capture) with a single unified KV pool indexed via paged
|
||||
indices, mirroring `aiter.mla.mla_decode_fwd`'s API style.
|
||||
|
||||
Caller contract:
|
||||
unified_kv: [total_pages, D] BF16 (page_size=1)
|
||||
Conceptually merges the SWA ring buffer and the compressor paged cache
|
||||
of a single V4 layer. Slots in `[0, swa_pages)` reference SWA entries
|
||||
(state_slot * win + ring); slots in `[swa_pages, ...)` reference
|
||||
compressed-K entries (block_id * K_PER_BLOCK + slot_in_block).
|
||||
kv_indices: [total_indices] int32 — per-token slot lists, flat.
|
||||
Per-token entries live in
|
||||
`kv_indices[kv_indptr[t] : kv_indptr[t+1]]`.
|
||||
**All entries MUST be valid slot ids in [0, unified_kv.shape[0]).**
|
||||
The production decode index builder (``write_v4_paged_decode_indices``)
|
||||
emits ragged-packed indices with no sentinels; CG-padded tokens get
|
||||
a zero-length slice via ``indptr[t+1] == indptr[t]``. The kernel no
|
||||
longer carries a per-iter ``slot >= 0`` sentinel check.
|
||||
kv_indptr: [N+1] int32 — true prefix sum (variable per-token len).
|
||||
attn_sink: [H] per-head learnable softmax-denom bias (V4 specific).
|
||||
softmax_scale: float.
|
||||
|
||||
Returns:
|
||||
out: [N, H, D] same dtype as q.
|
||||
|
||||
Numerics: online-softmax in log2 domain (exp2 with qk_scale = softmax_scale *
|
||||
LOG2E), with attention sink folded as a virtual K. Bit-close to the PyTorch
|
||||
reference (``_sparse_attn_ragged_torch``) within fp32-accumulation tolerance.
|
||||
|
||||
Architecture:
|
||||
- Small T (T * ceil(H/block_h) < ~1.7×CU): split-K + reduce-kernel path.
|
||||
Split kernel emits (m, l, acc) fp32 partials; reduce combines splits and
|
||||
folds attn_sink.
|
||||
- Large T: single-pass FUSED kernel that does softmax + sink + write in one
|
||||
shot (no partial-buffer alloc, no second kernel launch).
|
||||
|
||||
CUDAGraph-safe: kv_splits and tile config depend only on capture-time
|
||||
shapes (``T``, ``H``); the kernels' early-return / segment-mask logic is
|
||||
driven by ``kv_indptr`` values, which are runtime data but don't affect
|
||||
the captured launch sequence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from aiter.ops.triton.utils.device_info import get_num_sms
|
||||
|
||||
LOG2E = 1.4426950408889634 # log2(e); folded into qk_scale so softmax can use exp2.
|
||||
_MAX_KV_SPLITS = 64 # Hard cap on kv_splits (see _kv_splits_heuristic).
|
||||
|
||||
# FP8 KV cache (1xGROUP_SIZE block-scale quantization).
|
||||
#
|
||||
# Storage: unified_kv[total_pages, D] in e4m3fnuz + kv_scales[total_pages,
|
||||
# D // GROUP_SIZE] in fp32. Per-slot, D is split into NUM_GROUPS chunks of
|
||||
# GROUP_SIZE elements; each chunk shares one fp32 scale.
|
||||
# Dequant in-kernel: kv_bf16 = kv_fp8.to(fp32) * scale[d // GROUP_SIZE], cast
|
||||
# back to q.dtype before the second dot.
|
||||
#
|
||||
# GROUP_SIZE=64 matches the user's 1x64 quant spec. V4-Pro: D=512 → 8 scales
|
||||
# per slot, 4 bytes each → +6.25% storage on top of the fp8 pool (vs the
|
||||
# halving from bf16→fp8 = 2× saving — net ~46% read bandwidth reduction).
|
||||
_FP8_GROUP_SIZE = 64
|
||||
_FP8_DTYPE = torch.float8_e4m3fnuz
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _cu_count() -> int:
|
||||
"""Compute-unit count of the active GPU, queried once via aiter.
|
||||
|
||||
Wrapped in ``lru_cache`` so the first decode call pays the device-property
|
||||
lookup and all subsequent calls hit the cache — important inside a hot
|
||||
decode loop and CUDAGraph capture (no data-dependent host work).
|
||||
"""
|
||||
return get_num_sms()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Heuristics — pure-Python, capture-time deterministic.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _kernel_config(block_h: int) -> tuple[int, int, int]:
|
||||
"""Pick (BLOCK_K, num_warps, num_stages) without autotune.
|
||||
|
||||
Depends ONLY on ``block_h`` (a function of H, capture-time shape) so the
|
||||
config is identical between CUDAGraph capture and replay regardless of
|
||||
per-token K. In production ``kv_indices.shape[0]`` is a padded bucket
|
||||
whose value is unrelated to the true per-token kv_len, so any heuristic
|
||||
that reads it would mis-tune at capture time.
|
||||
|
||||
Derived from autotune statistics over ~150 shapes on MI355:
|
||||
- BLOCK_K=16 dominated (~78% of best configs for D=512); D=512 has
|
||||
enough load width that wider K tiles spill regs more than they buy.
|
||||
- num_warps:
|
||||
block_h ≤ 32 → 4 warps; block_h ≥ 64 → 8 warps.
|
||||
Pre-v11 the threshold was 16, which gave H=32 nw=8 — that was a
|
||||
regression worth +30% geomean (max +57%) on H=32 across all T,
|
||||
from a Phase-1 fine-tune sweep on top of v05+v09. With block_h=32
|
||||
the MFMA tile is 32×16×D; 4 waves (256 threads) is the sweet
|
||||
spot for AMD wave64 register budget. 8 warps over-distribute and
|
||||
leave each warp with too little MFMA to amortize pipeline fill.
|
||||
block_h=64 still wants 8 warps (one wave per row tile is too
|
||||
little ILP).
|
||||
- num_stages=2 is the safe default — deeper pipelining (3) helps only
|
||||
when the K loop has many iterations; we cannot know per-token K at
|
||||
capture time, so the conservative pick avoids regressing short-K.
|
||||
"""
|
||||
block_k = 16
|
||||
num_warps = 4 if block_h <= 32 else 8
|
||||
num_stages = 2
|
||||
return block_k, num_warps, num_stages
|
||||
|
||||
|
||||
def _prev_pow2(n: int) -> int:
|
||||
"""Largest power of two ≤ ``n``. For n < 1 returns 1."""
|
||||
if n < 1:
|
||||
return 1
|
||||
return 1 << (n.bit_length() - 1)
|
||||
|
||||
|
||||
def _kv_splits_heuristic(
|
||||
T: int,
|
||||
H: int,
|
||||
block_h: int,
|
||||
num_cu: int | None = None,
|
||||
target_wg_per_cu: float = 2.0,
|
||||
max_kv_splits: int = _MAX_KV_SPLITS,
|
||||
) -> int:
|
||||
"""Pick KV_SPLITS to fill the GPU. CUDAGraph-safe: depends ONLY on
|
||||
capture-time scalars (``T``, ``H``, ``block_h``). Never reads any
|
||||
tensor value or shape — production callers must not assume kv_indices
|
||||
layout encodes per-token kv_len.
|
||||
|
||||
Split-K trades a single-kernel pass for (split, reduce) two-kernel +
|
||||
partial-buffer allocation. The trade is worth it when the base grid
|
||||
``T * ceil(H/block_h)`` underfills the device.
|
||||
|
||||
base_ctas = T * ceil(H / block_h)
|
||||
target_wg = target_wg_per_cu * num_cu (≈ 1.7x to hide load-imbalance)
|
||||
if base_ctas >= target_wg: splits = 1 (grid already saturates GPU)
|
||||
else: splits = prev_pow2(min(target_wg/base_ctas,
|
||||
max_kv_splits))
|
||||
|
||||
``max_kv_splits`` (default 64) caps the number of split-kernel CTAs per
|
||||
token. Higher values would buy more parallelism for bs=1 long-ctx, but
|
||||
when per-token K is short most splits fall-through and the launch
|
||||
overhead dominates. 64 is the sweet spot for MI300/MI355.
|
||||
|
||||
Rounded DOWN to a power of two — rounding up over-splits when
|
||||
splits_to_fill isn't already pow2 (e.g. T=2 → 258 → 512 doubles the wg
|
||||
count past target, halves per-split work → 4× slowdown on bs=2 ctx=16384).
|
||||
"""
|
||||
if num_cu is None:
|
||||
num_cu = _cu_count()
|
||||
target_wg = max(1, int(target_wg_per_cu * num_cu))
|
||||
head_blocks = max(1, (H + block_h - 1) // block_h)
|
||||
base_ctas = max(1, T * head_blocks)
|
||||
if base_ctas >= target_wg:
|
||||
return 1
|
||||
|
||||
splits_to_fill = max(1, target_wg // base_ctas)
|
||||
return _prev_pow2(min(splits_to_fill, max_kv_splits))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kernels.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _paged_decode_fused_kernel(
|
||||
q_ptr, # [N, H, D]
|
||||
unified_kv_ptr, # [total_pages, D] bf16/fp16, or fp8 when QUANT_KV
|
||||
kv_scales_ptr, # [total_pages, NUM_GROUPS] fp32 when QUANT_KV (dummy otherwise)
|
||||
kv_indices_ptr, # [total_indices] int32
|
||||
kv_indptr_ptr, # [N+1] int32
|
||||
attn_sink_ptr, # [H]
|
||||
out_ptr, # [N, H, D]
|
||||
q_stride_t,
|
||||
q_stride_h,
|
||||
q_stride_d,
|
||||
kv_stride_n,
|
||||
kv_stride_d,
|
||||
ks_stride_n, # row stride of kv_scales (groups are contiguous, stride=1)
|
||||
out_stride_t,
|
||||
out_stride_h,
|
||||
out_stride_d,
|
||||
qk_scale, # = softmax_scale * LOG2E
|
||||
log2e, # = LOG2E, to lift natural-log sink into log2 domain
|
||||
H: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
QUANT_KV: tl.constexpr, # True → dequant fp8 KV via kv_scales
|
||||
GROUP_SIZE: tl.constexpr, # scale block width along D (e.g. 64)
|
||||
NUM_GROUPS: tl.constexpr, # D // GROUP_SIZE (constexpr; D % GROUP_SIZE == 0)
|
||||
):
|
||||
"""Single-pass online-softmax with sink folded inline — fast path for
|
||||
cases where ``kv_splits = 1`` (base grid already saturates the GPU). Skips
|
||||
the partial-buffer alloc + reduce-kernel launch that the 2-kernel
|
||||
split-K path needs.
|
||||
|
||||
Grid: ``(N, ceil(H / BLOCK_H))``. One CTA owns one token and one head-tile.
|
||||
"""
|
||||
t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
|
||||
h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
h_mask = h_offs < H
|
||||
d_mask = d_offs < D
|
||||
|
||||
q = tl.load(
|
||||
q_ptr
|
||||
+ t * q_stride_t
|
||||
+ h_offs[:, None] * q_stride_h
|
||||
+ d_offs[None, :] * q_stride_d,
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
kv_start = tl.load(kv_indptr_ptr + t)
|
||||
kv_end = tl.load(kv_indptr_ptr + t + 1)
|
||||
kv_len = kv_end - kv_start
|
||||
num_tiles = tl.cdiv(kv_len, BLOCK_K)
|
||||
|
||||
neg_large = -3.4028234663852886e38
|
||||
m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
|
||||
l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
|
||||
acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
|
||||
|
||||
k_offs = tl.arange(0, BLOCK_K)
|
||||
if QUANT_KV:
|
||||
# Compile-time per-D-element group index: d_offs // GROUP_SIZE has
|
||||
# NUM_GROUPS distinct values; redundant scale loads at the same
|
||||
# address are coalesced through L1, no per-element scalar issue.
|
||||
g_idx_per_d = d_offs // GROUP_SIZE
|
||||
# num_stages=3 on the inner K loop overrides the launch-time default
|
||||
# (2) for this loop only. Deeper SW pipeline keeps 2 in-flight KV
|
||||
# gathers (vs 1 with stages=2) while the current MFMA runs — better
|
||||
# hides AMD HBM gather latency on D=512. Cost: ~1 extra KV tile
|
||||
# (BLOCK_K*BLOCK_D*2 = 16KB at bf16) staged in regs/LDS per CTA.
|
||||
for j in tl.range(0, num_tiles, num_stages=3):
|
||||
k_start = j * BLOCK_K
|
||||
k_pos = k_start + k_offs
|
||||
valid = k_pos < kv_len # in_range; no sentinel check (see contract)
|
||||
slot = tl.load(
|
||||
kv_indices_ptr + kv_start + k_pos,
|
||||
mask=valid,
|
||||
other=0, # any in-bounds slot; the read is masked out below
|
||||
)
|
||||
|
||||
kv_raw = tl.load(
|
||||
unified_kv_ptr
|
||||
+ slot[:, None] * kv_stride_n
|
||||
+ d_offs[None, :] * kv_stride_d,
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
if QUANT_KV:
|
||||
# 1xGROUP_SIZE block-scale dequant via direct broadcast load —
|
||||
# avoids the explicit reshape + 3D intermediate that pinned
|
||||
# too many bf16 tiles in flight at once. The masked load with
|
||||
# d_offs // GROUP_SIZE as the column index produces a virtual
|
||||
# [BLOCK_K, BLOCK_D] scales tile but in IR is a coalesced
|
||||
# NUM_GROUPS-wide load per row.
|
||||
scales_full = tl.load(
|
||||
kv_scales_ptr + slot[:, None] * ks_stride_n + g_idx_per_d[None, :],
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
).to(q.dtype)
|
||||
kv = kv_raw.to(q.dtype) * scales_full
|
||||
else:
|
||||
kv = kv_raw
|
||||
|
||||
scores = tl.dot(q, tl.trans(kv)) * qk_scale
|
||||
# K: drop h_mask from the per-iter where. In V4-Pro every realistic
|
||||
# (H, BLOCK_H) pair has H % BLOCK_H == 0 → h_mask is statically
|
||||
# all-True, but the runtime ``h_offs < H`` compare prevents Triton
|
||||
# from constant-folding. Masking only on ``valid`` is sufficient:
|
||||
# ``neg_large`` on invalid k positions makes exp2(scores - m) ≈ 0
|
||||
# in the subsequent dot, and the masked store at the end gates
|
||||
# invalid h rows from polluting the output.
|
||||
scores = tl.where(valid[None, :], scores, neg_large)
|
||||
|
||||
m_block = tl.max(scores, axis=1)
|
||||
m_new = tl.maximum(m_i, m_block)
|
||||
alpha = tl.exp2(m_i - m_new)
|
||||
p = tl.exp2(scores - m_new[:, None])
|
||||
l_new = l_i * alpha + tl.sum(p, axis=1)
|
||||
|
||||
acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
# Fold attn_sink as a virtual K of weight 1. sink is a natural-log bias;
|
||||
# multiply by log2e so it lives in the same log2 domain as our online m_i.
|
||||
sink_raw = tl.load(attn_sink_ptr + h_offs, mask=h_mask, other=neg_large).to(
|
||||
tl.float32
|
||||
)
|
||||
sink = sink_raw * log2e
|
||||
m_final = tl.maximum(m_i, sink)
|
||||
alpha_kv = tl.exp2(m_i - m_final)
|
||||
alpha_sink = tl.exp2(sink - m_final)
|
||||
l_final = l_i * alpha_kv + alpha_sink
|
||||
|
||||
denom = tl.maximum(l_final, 1.0e-30)
|
||||
out = tl.where(
|
||||
l_final[:, None] > 0.0, (acc * alpha_kv[:, None]) / denom[:, None], 0.0
|
||||
)
|
||||
tl.store(
|
||||
out_ptr
|
||||
+ t * out_stride_t
|
||||
+ h_offs[:, None] * out_stride_h
|
||||
+ d_offs[None, :] * out_stride_d,
|
||||
out.to(out_ptr.dtype.element_ty),
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _paged_decode_split_kernel(
|
||||
q_ptr, # [N, H, D]
|
||||
unified_kv_ptr, # [total_pages, D] bf16/fp16, or fp8 when QUANT_KV
|
||||
kv_scales_ptr, # [total_pages, NUM_GROUPS] fp32 when QUANT_KV (dummy otherwise)
|
||||
kv_indices_ptr, # [total_indices] int32
|
||||
kv_indptr_ptr, # [N+1] int32
|
||||
m_partial_ptr, # [N, KV_SPLITS, H_padded] fp32
|
||||
l_partial_ptr, # [N, KV_SPLITS, H_padded] fp32
|
||||
acc_partial_ptr, # [N, KV_SPLITS, H_padded, D] fp32
|
||||
q_stride_t,
|
||||
q_stride_h,
|
||||
q_stride_d,
|
||||
kv_stride_n,
|
||||
kv_stride_d,
|
||||
ks_stride_n, # row stride of kv_scales (groups are contiguous, stride=1)
|
||||
mp_stride_t,
|
||||
mp_stride_k,
|
||||
mp_stride_h,
|
||||
lp_stride_t,
|
||||
lp_stride_k,
|
||||
lp_stride_h,
|
||||
ap_stride_t,
|
||||
ap_stride_k,
|
||||
ap_stride_h,
|
||||
ap_stride_d,
|
||||
H: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
KV_SPLITS: tl.constexpr,
|
||||
qk_scale, # = softmax_scale * LOG2E
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
QUANT_KV: tl.constexpr, # True → dequant fp8 KV via kv_scales
|
||||
GROUP_SIZE: tl.constexpr, # scale block width along D (e.g. 64)
|
||||
NUM_GROUPS: tl.constexpr, # D // GROUP_SIZE
|
||||
):
|
||||
"""3D split-K + exp2-softmax sparse paged-decode. Grid: (N, ceil(H/BLOCK_H), KV_SPLITS).
|
||||
|
||||
Emits pre-sink (m, l, acc) partials in fp32. The reduce kernel folds
|
||||
``attn_sink`` and combines splits.
|
||||
"""
|
||||
t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
pid_k = tl.program_id(2)
|
||||
|
||||
h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
h_mask = h_offs < H
|
||||
d_mask = d_offs < D
|
||||
|
||||
q = tl.load(
|
||||
q_ptr
|
||||
+ t * q_stride_t
|
||||
+ h_offs[:, None] * q_stride_h
|
||||
+ d_offs[None, :] * q_stride_d,
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
kv_start = tl.load(kv_indptr_ptr + t)
|
||||
kv_end = tl.load(kv_indptr_ptr + t + 1)
|
||||
kv_len = kv_end - kv_start
|
||||
|
||||
# tiles_per_segment pattern from aiter's kernel_unified_attention_3d:
|
||||
# KV_SPLITS is constexpr; tiles_per_segment derived at runtime. Splits
|
||||
# whose tile range is past kv_len early-return WITHOUT writing — the
|
||||
# reduce kernel uses the same constexpr BLOCK_K/KV_SPLITS to compute
|
||||
# ``act_num_segments = cdiv(kv_len, tiles_per_segment * BLOCK_K)`` and
|
||||
# masks unwritten slots out of its load. This saves a per-empty-split
|
||||
# write of BLOCK_H*BLOCK_D fp32 (16 KB at BLOCK_H=16, BLOCK_D=512) which
|
||||
# dominated short-K + many-splits cases.
|
||||
tiles_per_segment = tl.cdiv(kv_len, KV_SPLITS * BLOCK_K)
|
||||
if pid_k * tiles_per_segment * BLOCK_K >= kv_len:
|
||||
return
|
||||
num_tiles = tl.cdiv(kv_len, BLOCK_K)
|
||||
tile_start = pid_k * tiles_per_segment
|
||||
tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles)
|
||||
|
||||
neg_large = -3.4028234663852886e38
|
||||
m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
|
||||
l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
|
||||
acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
|
||||
|
||||
k_offs = tl.arange(0, BLOCK_K)
|
||||
if QUANT_KV:
|
||||
g_idx_per_d = d_offs // GROUP_SIZE
|
||||
# num_stages=3 (see fused kernel comment for rationale).
|
||||
for j in tl.range(tile_start, tile_end, num_stages=3):
|
||||
k_start = j * BLOCK_K
|
||||
k_pos = k_start + k_offs
|
||||
valid = k_pos < kv_len # in_range; no sentinel check (see contract)
|
||||
slot = tl.load(
|
||||
kv_indices_ptr + kv_start + k_pos,
|
||||
mask=valid,
|
||||
other=0, # any in-bounds slot; masked out below
|
||||
)
|
||||
|
||||
kv_raw = tl.load(
|
||||
unified_kv_ptr
|
||||
+ slot[:, None] * kv_stride_n
|
||||
+ d_offs[None, :] * kv_stride_d,
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
if QUANT_KV:
|
||||
scales_full = tl.load(
|
||||
kv_scales_ptr + slot[:, None] * ks_stride_n + g_idx_per_d[None, :],
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
).to(q.dtype)
|
||||
kv = kv_raw.to(q.dtype) * scales_full
|
||||
else:
|
||||
kv = kv_raw
|
||||
|
||||
scores = tl.dot(q, tl.trans(kv)) * qk_scale
|
||||
# K (same as fused kernel): drop h_mask from per-iter where.
|
||||
scores = tl.where(valid[None, :], scores, neg_large)
|
||||
|
||||
m_block = tl.max(scores, axis=1)
|
||||
m_new = tl.maximum(m_i, m_block)
|
||||
alpha = tl.exp2(m_i - m_new)
|
||||
p = tl.exp2(scores - m_new[:, None])
|
||||
l_new = l_i * alpha + tl.sum(p, axis=1)
|
||||
|
||||
acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
m_base = t * mp_stride_t + pid_k * mp_stride_k
|
||||
tl.store(m_partial_ptr + m_base + h_offs * mp_stride_h, m_i, mask=h_mask)
|
||||
l_base = t * lp_stride_t + pid_k * lp_stride_k
|
||||
tl.store(l_partial_ptr + l_base + h_offs * lp_stride_h, l_i, mask=h_mask)
|
||||
a_base = t * ap_stride_t + pid_k * ap_stride_k
|
||||
tl.store(
|
||||
acc_partial_ptr
|
||||
+ a_base
|
||||
+ h_offs[:, None] * ap_stride_h
|
||||
+ d_offs[None, :] * ap_stride_d,
|
||||
acc,
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _paged_decode_reduce_kernel(
|
||||
m_partial_ptr, # [N, KV_SPLITS, H_padded] fp32
|
||||
l_partial_ptr, # [N, KV_SPLITS, H_padded] fp32
|
||||
acc_partial_ptr, # [N, KV_SPLITS, H_padded, D] fp32
|
||||
attn_sink_ptr, # [H]
|
||||
kv_indptr_ptr, # [N+1] int32
|
||||
out_ptr, # [N, H, D]
|
||||
mp_stride_t,
|
||||
mp_stride_k,
|
||||
mp_stride_h,
|
||||
lp_stride_t,
|
||||
lp_stride_k,
|
||||
lp_stride_h,
|
||||
ap_stride_t,
|
||||
ap_stride_k,
|
||||
ap_stride_h,
|
||||
ap_stride_d,
|
||||
out_stride_t,
|
||||
out_stride_h,
|
||||
out_stride_d,
|
||||
log2e, # = LOG2E, used to convert natural-log sink → log2 domain
|
||||
H: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
KV_SPLITS: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
D_CHUNK: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
"""2D-tile reduce: combine KV_SPLITS partials, fold attn_sink, write
|
||||
final output. Grid: ``(T, H, ceil(D / D_CHUNK))`` — one CTA owns one
|
||||
(token, single-head, D-chunk) tuple.
|
||||
|
||||
Rewrite of the prior 3D-load reduce inspired by:
|
||||
- aiter ``_fwd_kernel_stage2`` (mla_decode_rope.py): scalar control
|
||||
flow + one D-tile per CTA, online accumulation across splits.
|
||||
- AKO4X ``hybrid_2d_reduce`` (B200 reference, +3.37x): merged 2D
|
||||
tile load ``[KV_SPLITS, D_CHUNK]`` replaces strided 3D access.
|
||||
|
||||
Why this wins on MI355 split path (T ≤ ~256, kv_splits > 1):
|
||||
|
||||
1. **2D tile fits registers**. The old 3D load tile
|
||||
``[KV_SPLITS=64, BLOCK_H=1, BLOCK_D=512]`` = 32K fp32 = 128 KB —
|
||||
didn't fit in registers, spilled to LDS, and used a strided 3D
|
||||
address compute. New ``[KV_SPLITS, D_CHUNK]`` = 64×64×4 = 16 KB
|
||||
fits one wave's VGPR with headroom.
|
||||
2. **D-chunked grid widens occupancy**. Old grid was ``(T, H)`` —
|
||||
e.g. T=1 H=16 → 16 CTAs into 256 CUs (6% occupancy). New grid
|
||||
``(T, H, D/D_CHUNK)`` → 16 × (512/64) = 128 CTAs at T=1 H=16
|
||||
(50% occupancy). Reduce was the latency bottleneck at small T.
|
||||
3. **Scalar control flow for sink fold**. Sink computation needs
|
||||
(m_max, l_combined) which only depend on (t, h) — same value
|
||||
across all dc programs for the same (t, h). They recompute it
|
||||
redundantly but it's a tiny scalar reduce; cheaper than passing
|
||||
through LDS.
|
||||
"""
|
||||
t = tl.program_id(0)
|
||||
h = tl.program_id(1)
|
||||
dc = tl.program_id(2)
|
||||
|
||||
d_offs = dc * D_CHUNK + tl.arange(0, D_CHUNK)
|
||||
k_offs = tl.arange(0, KV_SPLITS)
|
||||
d_mask = d_offs < D
|
||||
|
||||
neg_large = -3.4028234663852886e38
|
||||
|
||||
kv_start = tl.load(kv_indptr_ptr + t)
|
||||
kv_end = tl.load(kv_indptr_ptr + t + 1)
|
||||
kv_len = kv_end - kv_start
|
||||
# CTA-level early return for empty tokens (CUDAGraph padding, or any
|
||||
# caller-supplied zero-length slice). Split kernel skipped these without
|
||||
# writing partials → partial buffers hold garbage; the segm_mask path
|
||||
# below would still mask it correctly but consumes the masked-load BW
|
||||
# and the sink-fold arithmetic. Skipping the whole CTA also halves the
|
||||
# reduce-kernel cost on mixed-kv batches with many padded tokens.
|
||||
if kv_len == 0:
|
||||
out_off = t * out_stride_t + h * out_stride_h + d_offs * out_stride_d
|
||||
tl.store(
|
||||
out_ptr + out_off,
|
||||
tl.zeros([D_CHUNK], dtype=out_ptr.dtype.element_ty),
|
||||
mask=d_mask,
|
||||
)
|
||||
return
|
||||
tiles_per_segment = tl.cdiv(kv_len, KV_SPLITS * BLOCK_K)
|
||||
act_num_segments = tl.cdiv(kv_len, tl.maximum(tiles_per_segment, 1) * BLOCK_K)
|
||||
segm_mask = k_offs < act_num_segments
|
||||
|
||||
# 1D loads for (m, l) along splits — single head h.
|
||||
m_p = tl.load(
|
||||
m_partial_ptr + t * mp_stride_t + k_offs * mp_stride_k + h * mp_stride_h,
|
||||
mask=segm_mask,
|
||||
other=neg_large,
|
||||
) # [KV_SPLITS]
|
||||
l_p = tl.load(
|
||||
l_partial_ptr + t * lp_stride_t + k_offs * lp_stride_k + h * lp_stride_h,
|
||||
mask=segm_mask,
|
||||
other=0.0,
|
||||
) # [KV_SPLITS]
|
||||
|
||||
# 2D-tile load for acc partials — the key change vs old 3D-strided load.
|
||||
a_p = tl.load(
|
||||
acc_partial_ptr
|
||||
+ t * ap_stride_t
|
||||
+ k_offs[:, None] * ap_stride_k
|
||||
+ h * ap_stride_h
|
||||
+ d_offs[None, :] * ap_stride_d,
|
||||
mask=segm_mask[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
) # [KV_SPLITS, D_CHUNK]
|
||||
|
||||
# Combine across splits.
|
||||
m_max = tl.max(m_p, axis=0) # scalar
|
||||
alpha_split = tl.exp2(m_p - m_max) # [KV_SPLITS]
|
||||
l_combined = tl.sum(l_p * alpha_split, axis=0) # scalar
|
||||
acc_combined = tl.sum(a_p * alpha_split[:, None], axis=0) # [D_CHUNK]
|
||||
|
||||
# Fold attn_sink (recomputed across dc — scalar work, negligible).
|
||||
sink_raw = tl.load(attn_sink_ptr + h).to(tl.float32)
|
||||
sink = sink_raw * log2e
|
||||
m_final = tl.maximum(m_max, sink)
|
||||
alpha_kv = tl.exp2(m_max - m_final)
|
||||
alpha_sink = tl.exp2(sink - m_final)
|
||||
l_final = l_combined * alpha_kv + alpha_sink
|
||||
|
||||
denom = tl.maximum(l_final, 1.0e-30)
|
||||
# Direct divide (acc*alpha_kv)/denom, matching the single-CTA reference.
|
||||
# The prior `acc * (alpha_kv/denom)` precomputed a reciprocal-scaled scalar
|
||||
# (~1 extra ulp per element). Under split-K that per-kv_splits rounding
|
||||
# diverges across batch shapes (T) and, via MTP greedy spec-acceptance,
|
||||
# flips tokens — breaking MTP losslessness. The D-chunked multi-CTA layout
|
||||
# (perf) is untouched; only the final normalize arithmetic changes.
|
||||
acc_final = acc_combined * alpha_kv
|
||||
out = tl.where(l_final > 0.0, acc_final / denom, 0.0)
|
||||
|
||||
tl.store(
|
||||
out_ptr + t * out_stride_t + h * out_stride_h + d_offs * out_stride_d,
|
||||
out.to(out_ptr.dtype.element_ty),
|
||||
mask=d_mask,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wrapper.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sparse_attn_v4_paged_decode_triton(
|
||||
q: torch.Tensor,
|
||||
unified_kv: torch.Tensor,
|
||||
kv_indices: torch.Tensor,
|
||||
kv_indptr: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
kv_scales: torch.Tensor | None = None,
|
||||
block_h: int | None = None,
|
||||
kv_splits: int | None = None,
|
||||
block_k: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""V4 sparse decode Triton implementation: split-K with FUSED fast path,
|
||||
exp2 softmax, CG-safe heuristic. ``block_h`` and ``kv_splits`` are
|
||||
escape hatches for benchmarks; production callers pass neither.
|
||||
|
||||
When ``kv_scales`` is provided, ``unified_kv`` must be e4m3fnuz and
|
||||
``kv_scales`` must be ``[total_pages, D // GROUP_SIZE]`` fp32 — 1xGROUP_SIZE
|
||||
block-scale quantization. Dequant happens in-kernel; the dot still runs
|
||||
in q.dtype.
|
||||
"""
|
||||
if not q.is_cuda:
|
||||
raise RuntimeError(
|
||||
"Triton sparse_attn_v4_paged_decode requires CUDA/HIP tensors"
|
||||
)
|
||||
if q.dtype not in (torch.bfloat16, torch.float16):
|
||||
raise RuntimeError(
|
||||
f"sparse_attn_v4_paged_decode expects fp16/bf16 q, got {q.dtype}"
|
||||
)
|
||||
|
||||
quant_kv = kv_scales is not None
|
||||
if quant_kv:
|
||||
if unified_kv.dtype != _FP8_DTYPE:
|
||||
raise RuntimeError(
|
||||
f"kv_scales supplied but unified_kv is {unified_kv.dtype}, "
|
||||
f"expected {_FP8_DTYPE}"
|
||||
)
|
||||
if kv_scales.dtype != torch.float32:
|
||||
raise RuntimeError(f"kv_scales must be fp32, got {kv_scales.dtype}")
|
||||
D_check = unified_kv.shape[-1]
|
||||
if D_check % _FP8_GROUP_SIZE != 0:
|
||||
raise RuntimeError(
|
||||
f"D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE}"
|
||||
)
|
||||
expected_g = D_check // _FP8_GROUP_SIZE
|
||||
if kv_scales.shape != (unified_kv.shape[0], expected_g):
|
||||
raise RuntimeError(
|
||||
f"kv_scales shape {tuple(kv_scales.shape)} does not match "
|
||||
f"expected ({unified_kv.shape[0]}, {expected_g})"
|
||||
)
|
||||
if kv_scales.stride(-1) != 1:
|
||||
kv_scales = kv_scales.contiguous()
|
||||
else:
|
||||
if unified_kv.dtype != q.dtype:
|
||||
raise RuntimeError(
|
||||
f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
|
||||
)
|
||||
|
||||
T, H, D = q.shape
|
||||
out = torch.empty_like(q)
|
||||
|
||||
if block_h is None:
|
||||
block_h = triton.next_power_of_2(min(H, 64))
|
||||
else:
|
||||
block_h = triton.next_power_of_2(block_h)
|
||||
block_h = max(block_h, 16) # AMD MFMA min tile
|
||||
|
||||
n_head_blocks = (H + block_h - 1) // block_h
|
||||
h_padded = n_head_blocks * block_h
|
||||
block_d = triton.next_power_of_2(D)
|
||||
|
||||
if kv_splits is None:
|
||||
kv_splits = _kv_splits_heuristic(T, H, block_h)
|
||||
|
||||
qk_scale = float(softmax_scale) * LOG2E
|
||||
_bk, num_warps, num_stages = _kernel_config(block_h)
|
||||
if block_k is None:
|
||||
# fp8 dequant inflates per-tile ALU work ~4×; a wider K tile amortizes
|
||||
# the per-tile dequant cost (scale load + cast + multiply) over more
|
||||
# MFMA work. Empirically BLOCK_K=32 wins ~20% over BLOCK_K=16 on fp8
|
||||
# (bs=512 ctx=4096: 3000µs → 2300µs) without hurting bf16.
|
||||
block_k = 32 if quant_kv else _bk
|
||||
|
||||
# Kernel reads (kv_scales_ptr, ks_stride_n) only when QUANT_KV — supply a
|
||||
# dummy 1-element fp32 tensor on the bf16 path so the launch signature
|
||||
# stays uniform (avoids a separate JIT specialization per call).
|
||||
if quant_kv:
|
||||
kv_scales_arg = kv_scales
|
||||
ks_stride_n_arg = kv_scales.stride(0)
|
||||
num_groups_arg = D // _FP8_GROUP_SIZE
|
||||
else:
|
||||
kv_scales_arg = q.new_empty(1, dtype=torch.float32)
|
||||
ks_stride_n_arg = 1
|
||||
# NUM_GROUPS still needs a constexpr value (unused at compile time
|
||||
# because the QUANT_KV=False branch elides the dequant code).
|
||||
num_groups_arg = 1
|
||||
|
||||
# Fast path: when the base grid (T * n_head_blocks) already saturates the
|
||||
# GPU, kv_splits=1 and a single-pass fused kernel beats split+reduce by
|
||||
# skipping the partial-buffer alloc and the second kernel launch.
|
||||
if kv_splits == 1:
|
||||
grid_fused = (T, n_head_blocks)
|
||||
_paged_decode_fused_kernel[grid_fused](
|
||||
q,
|
||||
unified_kv,
|
||||
kv_scales_arg,
|
||||
kv_indices,
|
||||
kv_indptr,
|
||||
attn_sink,
|
||||
out,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
unified_kv.stride(0),
|
||||
unified_kv.stride(1),
|
||||
ks_stride_n_arg,
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
out.stride(2),
|
||||
qk_scale,
|
||||
LOG2E,
|
||||
H,
|
||||
D,
|
||||
BLOCK_H=block_h,
|
||||
BLOCK_D=block_d,
|
||||
BLOCK_K=block_k,
|
||||
QUANT_KV=quant_kv,
|
||||
GROUP_SIZE=_FP8_GROUP_SIZE,
|
||||
NUM_GROUPS=num_groups_arg,
|
||||
num_warps=num_warps,
|
||||
num_stages=num_stages,
|
||||
)
|
||||
return out
|
||||
|
||||
# Split-K path: split kernel writes (m, l, acc) partials in log2 domain;
|
||||
# reduce kernel combines them, folds attn_sink, writes final output.
|
||||
# Empty splits early-return without writing — reduce masks them out using
|
||||
# the same constexpr BLOCK_K to derive ``act_num_segments``.
|
||||
m_partial = torch.empty(
|
||||
(T, kv_splits, h_padded), dtype=torch.float32, device=q.device
|
||||
)
|
||||
l_partial = torch.empty_like(m_partial)
|
||||
acc_partial = torch.empty(
|
||||
(T, kv_splits, h_padded, D), dtype=torch.float32, device=q.device
|
||||
)
|
||||
|
||||
grid_split = (T, n_head_blocks, kv_splits)
|
||||
_paged_decode_split_kernel[grid_split](
|
||||
q,
|
||||
unified_kv,
|
||||
kv_scales_arg,
|
||||
kv_indices,
|
||||
kv_indptr,
|
||||
m_partial,
|
||||
l_partial,
|
||||
acc_partial,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
unified_kv.stride(0),
|
||||
unified_kv.stride(1),
|
||||
ks_stride_n_arg,
|
||||
m_partial.stride(0),
|
||||
m_partial.stride(1),
|
||||
m_partial.stride(2),
|
||||
l_partial.stride(0),
|
||||
l_partial.stride(1),
|
||||
l_partial.stride(2),
|
||||
acc_partial.stride(0),
|
||||
acc_partial.stride(1),
|
||||
acc_partial.stride(2),
|
||||
acc_partial.stride(3),
|
||||
H,
|
||||
D,
|
||||
kv_splits,
|
||||
qk_scale,
|
||||
BLOCK_H=block_h,
|
||||
BLOCK_D=block_d,
|
||||
BLOCK_K=block_k,
|
||||
QUANT_KV=quant_kv,
|
||||
GROUP_SIZE=_FP8_GROUP_SIZE,
|
||||
NUM_GROUPS=num_groups_arg,
|
||||
num_warps=num_warps,
|
||||
num_stages=num_stages,
|
||||
)
|
||||
|
||||
# 2D-tile reduce: grid = (T, H, ceil(D/D_CHUNK)). One CTA per
|
||||
# (token, single-head, D-chunk). Adaptive D_CHUNK based on whether the
|
||||
# base reduce grid (T*H) already saturates the GPU:
|
||||
# - large T*H (≥ 2*num_CU=512 on MI355): the grid is already at or
|
||||
# above the launch target; further D-splitting just over-fragments
|
||||
# (T=128 H=128 with D_CHUNK=64 → 131K CTAs → 2× slowdown). Use
|
||||
# D_CHUNK = block_d (1 d-block per CTA, grid = T*H).
|
||||
# - small T*H: D-split widens the grid to fill CUs at small-T
|
||||
# latency-critical decode. The 2D tile [KV_SPLITS, D_CHUNK] should
|
||||
# stay ≤ 16 KB fp32 to fit registers.
|
||||
base_grid_t_h = T * H
|
||||
target_reduce_wg = 2 * _cu_count()
|
||||
if base_grid_t_h >= target_reduce_wg:
|
||||
d_chunk = block_d
|
||||
else:
|
||||
d_chunks_needed = max(1, target_reduce_wg // base_grid_t_h)
|
||||
d_chunks_needed = min(d_chunks_needed, block_d // 32)
|
||||
d_chunk = max(32, triton.next_power_of_2(block_d // d_chunks_needed))
|
||||
grid_reduce = (T, H, (D + d_chunk - 1) // d_chunk)
|
||||
_paged_decode_reduce_kernel[grid_reduce](
|
||||
m_partial,
|
||||
l_partial,
|
||||
acc_partial,
|
||||
attn_sink,
|
||||
kv_indptr,
|
||||
out,
|
||||
m_partial.stride(0),
|
||||
m_partial.stride(1),
|
||||
m_partial.stride(2),
|
||||
l_partial.stride(0),
|
||||
l_partial.stride(1),
|
||||
l_partial.stride(2),
|
||||
acc_partial.stride(0),
|
||||
acc_partial.stride(1),
|
||||
acc_partial.stride(2),
|
||||
acc_partial.stride(3),
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
out.stride(2),
|
||||
LOG2E,
|
||||
H,
|
||||
D,
|
||||
kv_splits,
|
||||
BLOCK_D=block_d,
|
||||
D_CHUNK=d_chunk,
|
||||
BLOCK_K=block_k,
|
||||
num_warps=4,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def sparse_attn_v4_paged_decode(
|
||||
q: torch.Tensor,
|
||||
unified_kv: torch.Tensor,
|
||||
kv_indices: torch.Tensor,
|
||||
kv_indptr: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
kv_scales: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""V4 decode sparse attention over a unified KV pool with paged indices.
|
||||
|
||||
When ``kv_scales`` is provided, ``unified_kv`` must be fp8 (e4m3fnuz) and
|
||||
will be dequantized in-kernel using 1xGROUP_SIZE (default 64) block scales.
|
||||
"""
|
||||
return _sparse_attn_v4_paged_decode_triton(
|
||||
q,
|
||||
unified_kv,
|
||||
kv_indices,
|
||||
kv_indptr,
|
||||
attn_sink,
|
||||
softmax_scale,
|
||||
kv_scales=kv_scales,
|
||||
)
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
"""V4 paged-decode index scatter — single Triton kernel writes SWA window-
|
||||
prefix paged offsets into the three ragged-packed destination buffers
|
||||
(`kv_indices_swa` / `kv_indices_csa` / `kv_indices_hca`).
|
||||
|
||||
Replaces the prior chain (numpy `_build_window_topk_np` + `index_copy_`):
|
||||
|
||||
window_topk_np = _build_window_topk_np(positions, win, ring_stride) # [T, win]
|
||||
swa_paged_2d = torch.where(window_topk >= 0, slot * ring_stride + topk, -1)
|
||||
swa_paged_flat = swa_paged_2d.reshape(-1)
|
||||
swa_indices_gpu[:T*win].copy_(swa_paged_flat)
|
||||
csa_indices_gpu.index_copy_(0, csa_win_pos, swa_paged_flat)
|
||||
hca_indices_gpu.index_copy_(0, hca_win_pos, swa_paged_flat)
|
||||
|
||||
Two simplifications vs the prior implementation (see plan
|
||||
`sequential-noodling-turing.md` for details):
|
||||
|
||||
1. The ring-index formula `ring = (pos - win + 1 + w) % ring_stride` is computed
|
||||
inline inside the kernel from `positions[t]`. The `[T, win]`
|
||||
`window_topk` intermediate buffer (mnbt·win·4 = 4 MB at typical config)
|
||||
is gone; no separate CPU build + H2D copy.
|
||||
2. The destination layout is now ragged-packed (same as prefill): each
|
||||
token's SWA prefix segment has length `n = min(positions[t]+1, win)`
|
||||
(NOT a fixed `win` padded with `-1` sentinels). The caller's
|
||||
`swa_indptr` / `csa_indptr` / `hca_indptr` reflect this ragged sizing.
|
||||
|
||||
Bytewise correctness: for tokens with `position >= win-1` (all `n == win`),
|
||||
the output is identical to the prior implementation. For shorter
|
||||
positions, the prior layout wrote `(win - n)` leading `-1` entries that
|
||||
the sparse-attention kernel masked out; the new layout omits those slots
|
||||
entirely, saving sparse-attn loop iterations.
|
||||
|
||||
Caller contract:
|
||||
- Grid = T (one program per token).
|
||||
- `batch_id_per_token[:T]` may carry `-1` sentinels in the CG-padded tail —
|
||||
kernel checks and bails (matches `_attach_v4_per_fwd_meta` convention).
|
||||
- `swa_indptr` / `csa_indptr` / `hca_indptr` must reflect the ragged-packed
|
||||
sizing: per-token slot count = `min(positions[t]+1, win) + n_compress[t]`
|
||||
where `n_compress[t]` is 0 for SWA, `min(n_committed_csa, index_topk)`
|
||||
for CSA, `n_committed_hca` for HCA.
|
||||
- `swa_indices` / `csa_indices` / `hca_indices` capacity ≥ corresponding
|
||||
indptr[T]; this kernel only writes the SWA-prefix segment
|
||||
`[indptr[t], indptr[t] + n)` per token. The compress-tail is filled
|
||||
elsewhere (HCA: numpy fill in caller, CSA: `csa_translate_pack` per layer).
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _v4_paged_decode_indices_kernel(
|
||||
state_slot_per_seq_ptr, # [bs] int32
|
||||
batch_id_per_token_ptr, # [T+pad] int — sentinel -1 in pad tail
|
||||
positions_ptr, # [T+pad] int — global token position
|
||||
swa_indptr_ptr, # [T+1] int32 — ragged SWA-prefix cumsum
|
||||
csa_indptr_ptr, # [T+1] int32 — ragged (SWA + CSA topk)
|
||||
hca_indptr_ptr, # [T+1] int32 — ragged (SWA + HCA committed)
|
||||
swa_indices_ptr, # [swa_total] int32, output
|
||||
csa_indices_ptr, # [csa_total] int32, output (writes SWA-prefix segment only)
|
||||
hca_indices_ptr, # [hca_total] int32, output (writes SWA-prefix segment only)
|
||||
ring_stride, # win_with_spec — stride into unified_kv SWA region (paper §3.6.1)
|
||||
win: tl.constexpr, # window_size — max SWA prefix slots
|
||||
BLOCK_N: tl.constexpr, # next_pow2(win)
|
||||
):
|
||||
"""One program per token. Writes `n = min(positions[t]+1, win)` paged
|
||||
offsets to the SWA prefix segment of each of SWA/CSA/HCA index buffers.
|
||||
|
||||
For token `t`:
|
||||
bid = batch_id_per_token[t] # bail if -1 (CG pad)
|
||||
slot = state_slot_per_seq[bid]
|
||||
pos = positions[t]
|
||||
n = min(pos + 1, win)
|
||||
# Old -1 sentinels were at the leading `win - n` cols; reparameterize
|
||||
# to skip them: i in [0, n) → abs_pos = pos - n + 1 + i ∈ [0, pos].
|
||||
for i in range(n):
|
||||
abs_pos = pos - n + 1 + i
|
||||
ring = abs_pos % ring_stride
|
||||
paged = slot * ring_stride + ring
|
||||
swa_indices[swa_indptr[t] + i] = paged
|
||||
csa_indices[csa_indptr[t] + i] = paged
|
||||
hca_indices[hca_indptr[t] + i] = paged
|
||||
"""
|
||||
t = tl.program_id(0)
|
||||
bid = tl.load(batch_id_per_token_ptr + t)
|
||||
if bid < 0:
|
||||
return # CG-padded sentinel — leave outputs untouched
|
||||
|
||||
slot = tl.load(state_slot_per_seq_ptr + bid)
|
||||
pos = tl.load(positions_ptr + t)
|
||||
# `n` = actual valid SWA prefix count. Cast to match `win` (compile-time
|
||||
# int) — pos is i32/i64 from positions buffer.
|
||||
n = tl.minimum(pos + 1, win)
|
||||
swa_base = tl.load(swa_indptr_ptr + t)
|
||||
csa_base = tl.load(csa_indptr_ptr + t)
|
||||
hca_base = tl.load(hca_indptr_ptr + t)
|
||||
|
||||
i = tl.arange(0, BLOCK_N)
|
||||
mask = i < n
|
||||
abs_pos = pos - n + 1 + i # ∈ [0, pos] for valid i
|
||||
ring_idx = abs_pos % ring_stride
|
||||
paged = slot * ring_stride + ring_idx
|
||||
|
||||
tl.store(swa_indices_ptr + swa_base + i, paged, mask=mask)
|
||||
tl.store(csa_indices_ptr + csa_base + i, paged, mask=mask)
|
||||
tl.store(hca_indices_ptr + hca_base + i, paged, mask=mask)
|
||||
|
||||
|
||||
def write_v4_paged_decode_indices(
|
||||
*,
|
||||
state_slot_per_seq: torch.Tensor,
|
||||
batch_id_per_token: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
swa_indptr: torch.Tensor,
|
||||
csa_indptr: torch.Tensor,
|
||||
hca_indptr: torch.Tensor,
|
||||
swa_indices: torch.Tensor,
|
||||
csa_indices: torch.Tensor,
|
||||
hca_indices: torch.Tensor,
|
||||
T: int,
|
||||
win: int,
|
||||
ring_stride: int,
|
||||
) -> None:
|
||||
"""In-place fill SWA / CSA / HCA window-prefix offsets via a single
|
||||
Triton kernel. Replaces the prior `_build_window_topk_np` (CPU O(T·win))
|
||||
+ `index_copy_` chain. All inputs are persistent forward_vars buffers —
|
||||
no allocator churn.
|
||||
|
||||
Args (all GPU tensors except T/win/ring_stride):
|
||||
state_slot_per_seq: [bs] int32 — per-seq state cache slot.
|
||||
batch_id_per_token: [>=T] int — token→seq map; -1 sentinel skipped.
|
||||
positions: [>=T] int — global token position
|
||||
(forward_vars["positions"]); used to derive
|
||||
`n = min(pos+1, win)` per token + the ring
|
||||
index `(pos - n + 1 + i) % ring_stride`.
|
||||
swa_indptr: [>=T+1] int32 — ragged SWA-prefix cumsum, where
|
||||
`swa_indptr[t+1] - swa_indptr[t] =
|
||||
min(positions[t]+1, win)`.
|
||||
csa_indptr: [>=T+1] int32 — ragged CSA buffer indptr (SWA
|
||||
prefix + CSA topk per token).
|
||||
hca_indptr: [>=T+1] int32 — ragged HCA buffer indptr (SWA
|
||||
prefix + HCA committed per token).
|
||||
swa_indices: [>=swa_indptr[T]] int32 OUT — fully written by
|
||||
this kernel (no other source).
|
||||
csa_indices: [>=csa_indptr[T]] int32 OUT — window-prefix
|
||||
`[csa_indptr[t], +n)` written here; CSA
|
||||
topk tail filled per-layer by
|
||||
`csa_translate_pack`.
|
||||
hca_indices: [>=hca_indptr[T]] int32 OUT — same semantics; HCA
|
||||
compress tail filled in the caller via
|
||||
numpy fill.
|
||||
T: int — number of real tokens (grid size).
|
||||
win: int — SWA window size (typically 128 for V4-Pro).
|
||||
ring_stride: int — `win_with_spec = window_size + max_spec_steps`,
|
||||
stride into unified_kv SWA region per slot
|
||||
AND modulo for ring-index wrap.
|
||||
"""
|
||||
if T == 0:
|
||||
return
|
||||
assert state_slot_per_seq.dim() == 1
|
||||
assert batch_id_per_token.dim() == 1 and batch_id_per_token.shape[0] >= T
|
||||
assert positions.dim() == 1 and positions.shape[0] >= T
|
||||
assert swa_indptr.dim() == 1 and swa_indptr.shape[0] >= T + 1
|
||||
assert csa_indptr.dim() == 1 and csa_indptr.shape[0] >= T + 1
|
||||
assert hca_indptr.dim() == 1 and hca_indptr.shape[0] >= T + 1
|
||||
assert swa_indices.dim() == 1
|
||||
assert csa_indices.dim() == 1
|
||||
assert hca_indices.dim() == 1
|
||||
|
||||
BLOCK_N = triton.next_power_of_2(win)
|
||||
_v4_paged_decode_indices_kernel[(T,)](
|
||||
state_slot_per_seq,
|
||||
batch_id_per_token,
|
||||
positions,
|
||||
swa_indptr,
|
||||
csa_indptr,
|
||||
hca_indptr,
|
||||
swa_indices,
|
||||
csa_indices,
|
||||
hca_indices,
|
||||
ring_stride,
|
||||
win=win,
|
||||
BLOCK_N=BLOCK_N,
|
||||
)
|
||||
@@ -0,0 +1,354 @@
|
||||
# SPDX-License-Identifier: MIT
|
||||
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
"""Sparse prefill attention with two KV sources: paged `unified_kv` (history)
|
||||
and per-fwd flat `kv` (current chunk's input).
|
||||
|
||||
Designed for V4 prefill: indexes the two KV sources directly without
|
||||
materialising a per-fwd `kv_flat_sa` packed tensor.
|
||||
|
||||
Caller contract:
|
||||
unified_kv: [total_pages, D] BF16 — prefix source. Same buffer as
|
||||
decode kernel: SWA ring slots in `[0, swa_pages)`, compress pages in
|
||||
`[swa_pages, total_pages)`. For prefill, prefix indices select
|
||||
(a) prior-chunk SWA history, (b) CSA topk, (c) HCA all-committed.
|
||||
kv_indices_prefix: [total_prefix_indices] int32 — flat per-token slot
|
||||
lists. Per-token entries live in
|
||||
`kv_indices_prefix[kv_indptr_prefix[t] : kv_indptr_prefix[t+1]]`.
|
||||
`-1` entries are skipped (sentinel).
|
||||
kv_indptr_prefix: [N+1] int32 — true prefix sum (variable per-token len).
|
||||
|
||||
kv: [total_tokens, D] BF16 — extend source = current
|
||||
fwd's just-computed K (NOT yet written to swa_kv ring). Layout matches
|
||||
`swa_write` input.
|
||||
kv_indices_extend: [total_extend_indices] int32 — flat per-token row idx
|
||||
lists into `kv`. Per-token entries live in
|
||||
`kv_indices_extend[kv_indptr_extend[t] : kv_indptr_extend[t+1]]`.
|
||||
`-1` entries are skipped (rare for extend; usually all valid).
|
||||
kv_indptr_extend: [N+1] int32 — true prefix sum.
|
||||
|
||||
attn_sink: [H] per-head learnable softmax-denom bias (V4 specific).
|
||||
softmax_scale: float.
|
||||
|
||||
Per-token K loop iterates two regions sequentially, sharing the online
|
||||
softmax accumulator (m_i, l_i, acc) across regions. Order of regions does
|
||||
not affect correctness (online softmax is order-invariant).
|
||||
|
||||
Returns:
|
||||
out: [N, H, D] same dtype as q.
|
||||
|
||||
Numerics: identical online-softmax + sink finalization to
|
||||
`sparse_attn_v4_paged_decode` — bit-exact when the extend region is empty
|
||||
(then equivalent to a decode call with the same prefix indices).
|
||||
"""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
# OPUS gfx950 paged-prefill kernel is preferred when importable; otherwise fall
|
||||
# back to the Triton implementation below.
|
||||
try:
|
||||
from aiter.ops.pa_sparse_prefill_opus import pa_sparse_prefill_opus
|
||||
|
||||
_HAS_OPUS = True
|
||||
except ImportError:
|
||||
pa_sparse_prefill_opus = None
|
||||
_HAS_OPUS = False
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _sparse_attn_v4_paged_prefill_kernel(
|
||||
q_ptr, # [N, H, D]
|
||||
unified_kv_ptr, # [total_pages, D] — prefix source
|
||||
kv_indices_prefix_ptr, # [total_prefix_indices] int32
|
||||
kv_indptr_prefix_ptr, # [N+1] int32
|
||||
kv_ptr, # [total_tokens, D] — extend source
|
||||
kv_indices_extend_ptr, # [total_extend_indices] int32
|
||||
kv_indptr_extend_ptr, # [N+1] int32
|
||||
attn_sink_ptr, # [H]
|
||||
out_ptr, # [N, H, D]
|
||||
q_stride_t: tl.constexpr,
|
||||
q_stride_h: tl.constexpr,
|
||||
q_stride_d: tl.constexpr,
|
||||
pkv_stride_n: tl.constexpr, # unified_kv stride 0 (= D usually)
|
||||
pkv_stride_d: tl.constexpr, # unified_kv stride 1 (= 1 usually)
|
||||
ekv_stride_n: tl.constexpr, # kv stride 0
|
||||
ekv_stride_d: tl.constexpr, # kv stride 1
|
||||
out_stride_t: tl.constexpr,
|
||||
out_stride_h: tl.constexpr,
|
||||
out_stride_d: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
softmax_scale: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
|
||||
h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
d_offs = tl.arange(0, BLOCK_D)
|
||||
h_mask = h_offs < H
|
||||
d_mask = d_offs < D
|
||||
|
||||
q = tl.load(
|
||||
q_ptr
|
||||
+ t * q_stride_t
|
||||
+ h_offs[:, None] * q_stride_h
|
||||
+ d_offs[None, :] * q_stride_d,
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
neg_large = -3.4028234663852886e38
|
||||
m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
|
||||
l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
|
||||
acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
|
||||
|
||||
k_offs = tl.arange(0, BLOCK_K)
|
||||
|
||||
# ===== Region 1: prefix from unified_kv =====
|
||||
p_start = tl.load(kv_indptr_prefix_ptr + t)
|
||||
p_end = tl.load(kv_indptr_prefix_ptr + t + 1)
|
||||
p_len = p_end - p_start
|
||||
|
||||
for k_start in tl.range(0, p_len, BLOCK_K):
|
||||
k_pos = k_start + k_offs
|
||||
in_range = k_pos < p_len
|
||||
slot = tl.load(
|
||||
kv_indices_prefix_ptr + p_start + k_pos,
|
||||
mask=in_range,
|
||||
other=-1,
|
||||
)
|
||||
valid = in_range & (slot >= 0)
|
||||
slot_clamped = tl.maximum(slot, 0)
|
||||
|
||||
kv = tl.load(
|
||||
unified_kv_ptr
|
||||
+ slot_clamped[:, None] * pkv_stride_n
|
||||
+ d_offs[None, :] * pkv_stride_d,
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
scores = tl.dot(q, tl.trans(kv)) * softmax_scale
|
||||
scores = tl.where(h_mask[:, None] & valid[None, :], scores, neg_large)
|
||||
|
||||
m_block = tl.max(scores, axis=1)
|
||||
m_new = tl.maximum(m_i, m_block)
|
||||
alpha = tl.exp(m_i - m_new)
|
||||
p = tl.exp(scores - m_new[:, None])
|
||||
p = tl.where(h_mask[:, None] & valid[None, :], p, 0.0)
|
||||
l_new = l_i * alpha + tl.sum(p, axis=1)
|
||||
|
||||
acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
# ===== Region 2: extend from kv (per-fwd flat) =====
|
||||
e_start = tl.load(kv_indptr_extend_ptr + t)
|
||||
e_end = tl.load(kv_indptr_extend_ptr + t + 1)
|
||||
e_len = e_end - e_start
|
||||
|
||||
for k_start in tl.range(0, e_len, BLOCK_K):
|
||||
k_pos = k_start + k_offs
|
||||
in_range = k_pos < e_len
|
||||
slot = tl.load(
|
||||
kv_indices_extend_ptr + e_start + k_pos,
|
||||
mask=in_range,
|
||||
other=-1,
|
||||
)
|
||||
valid = in_range & (slot >= 0)
|
||||
slot_clamped = tl.maximum(slot, 0)
|
||||
|
||||
kv = tl.load(
|
||||
kv_ptr
|
||||
+ slot_clamped[:, None] * ekv_stride_n
|
||||
+ d_offs[None, :] * ekv_stride_d,
|
||||
mask=valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
scores = tl.dot(q, tl.trans(kv)) * softmax_scale
|
||||
scores = tl.where(h_mask[:, None] & valid[None, :], scores, neg_large)
|
||||
|
||||
m_block = tl.max(scores, axis=1)
|
||||
m_new = tl.maximum(m_i, m_block)
|
||||
alpha = tl.exp(m_i - m_new)
|
||||
p = tl.exp(scores - m_new[:, None])
|
||||
p = tl.where(h_mask[:, None] & valid[None, :], p, 0.0)
|
||||
l_new = l_i * alpha + tl.sum(p, axis=1)
|
||||
|
||||
acc = acc * alpha[:, None] + tl.dot(p.to(kv.dtype), kv)
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
# ===== Sink finalization =====
|
||||
# Online softmax + sink integration: sink is a virtual extra K with V=0,
|
||||
# contributing only to the denominator. After main loops, (m_i, l_i, acc)
|
||||
# are in m_i frame; sink may shift max to m_final = max(m_i, sink), so
|
||||
# rescale BOTH l_i (for denom) AND acc (for numerator) by alpha to switch
|
||||
# to m_final frame. The sink itself adds exp(sink - m_final) to l_final
|
||||
# but contributes 0 to acc since V_sink = 0.
|
||||
sink = tl.load(attn_sink_ptr + h_offs, mask=h_mask, other=neg_large).to(tl.float32)
|
||||
m_final = tl.maximum(m_i, sink)
|
||||
alpha = tl.exp(m_i - m_final)
|
||||
l_final = l_i * alpha + tl.exp(sink - m_final)
|
||||
|
||||
denom = tl.maximum(l_final, 1.0e-30)
|
||||
out = tl.where(l_final[:, None] > 0.0, (acc * alpha[:, None]) / denom[:, None], 0.0)
|
||||
tl.store(
|
||||
out_ptr
|
||||
+ t * out_stride_t
|
||||
+ h_offs[:, None] * out_stride_h
|
||||
+ d_offs[None, :] * out_stride_d,
|
||||
out,
|
||||
mask=h_mask[:, None] & d_mask[None, :],
|
||||
)
|
||||
|
||||
|
||||
def _sparse_attn_v4_paged_prefill_triton(
|
||||
q: torch.Tensor,
|
||||
unified_kv: torch.Tensor,
|
||||
kv_indices_prefix: torch.Tensor,
|
||||
kv_indptr_prefix: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
kv_indices_extend: torch.Tensor,
|
||||
kv_indptr_extend: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
if not q.is_cuda:
|
||||
raise RuntimeError(
|
||||
"Triton sparse_attn_v4_paged_prefill requires CUDA/HIP tensors"
|
||||
)
|
||||
if q.dtype not in (torch.bfloat16, torch.float16):
|
||||
raise RuntimeError(
|
||||
f"sparse_attn_v4_paged_prefill expects fp16/bf16 q, got {q.dtype}"
|
||||
)
|
||||
if unified_kv.dtype != q.dtype:
|
||||
raise RuntimeError(
|
||||
f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
|
||||
)
|
||||
if kv.dtype != q.dtype:
|
||||
raise RuntimeError(f"kv dtype mismatch: kv={kv.dtype}, q={q.dtype}")
|
||||
if unified_kv.size(-1) != kv.size(-1):
|
||||
raise RuntimeError(
|
||||
f"head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}"
|
||||
)
|
||||
|
||||
T, H, D = q.shape
|
||||
out = torch.empty_like(q)
|
||||
kv_indices_prefix = kv_indices_prefix.to(torch.int32).contiguous()
|
||||
kv_indptr_prefix = kv_indptr_prefix.to(torch.int32).contiguous()
|
||||
kv_indices_extend = kv_indices_extend.to(torch.int32).contiguous()
|
||||
kv_indptr_extend = kv_indptr_extend.to(torch.int32).contiguous()
|
||||
|
||||
block_h = 16 # AMD MFMA min tile
|
||||
block_d = triton.next_power_of_2(D)
|
||||
block_k = 16 if D >= 256 else 32
|
||||
_sparse_attn_v4_paged_prefill_kernel[(T, triton.cdiv(H, block_h))](
|
||||
q,
|
||||
unified_kv,
|
||||
kv_indices_prefix,
|
||||
kv_indptr_prefix,
|
||||
kv,
|
||||
kv_indices_extend,
|
||||
kv_indptr_extend,
|
||||
attn_sink,
|
||||
out,
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
q.stride(2),
|
||||
unified_kv.stride(0),
|
||||
unified_kv.stride(1),
|
||||
kv.stride(0),
|
||||
kv.stride(1),
|
||||
out.stride(0),
|
||||
out.stride(1),
|
||||
out.stride(2),
|
||||
H,
|
||||
D,
|
||||
float(softmax_scale),
|
||||
BLOCK_H=block_h,
|
||||
BLOCK_D=block_d,
|
||||
BLOCK_K=block_k,
|
||||
num_warps=8,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def sparse_attn_v4_paged_prefill(
|
||||
q: torch.Tensor,
|
||||
unified_kv: torch.Tensor,
|
||||
kv_indices_prefix: torch.Tensor,
|
||||
kv_indptr_prefix: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
kv_indices_extend: torch.Tensor,
|
||||
kv_indptr_extend: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
"""V4 prefill sparse attention over two KV sources (paged unified_kv +
|
||||
flat per-fwd kv).
|
||||
|
||||
Args:
|
||||
q: [T, H, D] BF16/FP16 — query.
|
||||
unified_kv: [total_pages, D] BF16/FP16 — prefix source (paged).
|
||||
kv_indices_prefix: [total_prefix] int32 — flat per-token slot lists into
|
||||
unified_kv. -1 sentinels skipped.
|
||||
kv_indptr_prefix: [T+1] int32 — true prefix sum.
|
||||
kv: [total_tokens, D] BF16/FP16 — extend source (this
|
||||
fwd's input K, NOT yet in swa_kv ring).
|
||||
kv_indices_extend: [total_extend] int32 — flat per-token row idx lists
|
||||
into kv. -1 sentinels skipped.
|
||||
kv_indptr_extend: [T+1] int32 — true prefix sum.
|
||||
attn_sink: [H] — per-head softmax-denom bias.
|
||||
softmax_scale: float.
|
||||
|
||||
Returns:
|
||||
out: [T, H, D] same dtype as q.
|
||||
"""
|
||||
if _HAS_OPUS:
|
||||
# OPUS contract differs from the Triton kernel in two ways the Triton
|
||||
# path tolerates implicitly:
|
||||
# - it requires a FULLY-contiguous q (it only asserts stride(2)==1 but
|
||||
# indexes assuming [T,H,D] contiguous); a non-contiguous head stride
|
||||
# silently reads wrong/out-of-bounds addresses. q from the model is
|
||||
# often a view, so force contiguity.
|
||||
# - it requires ``attn_sink.size(0) == H``; attn_sink is the full
|
||||
# per-head Parameter, so slice to the H query heads.
|
||||
q = q.contiguous()
|
||||
H = q.shape[1]
|
||||
if attn_sink.shape[0] != H:
|
||||
attn_sink = attn_sink[:H].contiguous()
|
||||
if (
|
||||
kv.stride(0) != unified_kv.stride(0)
|
||||
and kv.shape[0] == 1
|
||||
and kv.stride(1) == 1
|
||||
):
|
||||
kv = kv.as_strided(kv.shape, (kv.shape[1], 1))
|
||||
return pa_sparse_prefill_opus(
|
||||
q,
|
||||
unified_kv,
|
||||
kv_indices_prefix,
|
||||
kv_indptr_prefix,
|
||||
kv,
|
||||
kv_indices_extend,
|
||||
kv_indptr_extend,
|
||||
attn_sink,
|
||||
softmax_scale,
|
||||
)
|
||||
return _sparse_attn_v4_paged_prefill_triton(
|
||||
q,
|
||||
unified_kv,
|
||||
kv_indices_prefix,
|
||||
kv_indptr_prefix,
|
||||
kv,
|
||||
kv_indices_extend,
|
||||
kv_indptr_extend,
|
||||
attn_sink,
|
||||
softmax_scale,
|
||||
)
|
||||
@@ -0,0 +1,440 @@
|
||||
"""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
|
||||
vendored paged decode/prefill kernels.
|
||||
|
||||
unified_kv[L] layout (page_size 1, bf16, 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)
|
||||
|
||||
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
|
||||
allocated at the fixed worst-case capacity ``N * (win + Wc)`` so its shape is
|
||||
static across CUDA-graph replay; only ``kv_indptr`` values (and the written
|
||||
prefix) vary per forward. Compressed valid entries are front-packed in the
|
||||
``*_page_indices`` rows (the same contract the non-unified_kv flashmla path relies
|
||||
on via ``topk_length``); the per-token compressed count is recovered from the
|
||||
``kv_indptr`` delta inside the kernel, so no extra length tensor is threaded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode import (
|
||||
sparse_attn_v4_paged_decode,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_decode_indices import (
|
||||
write_v4_paged_decode_indices,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.paged_prefill import (
|
||||
sparse_attn_v4_paged_prefill,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SWA ring scatter
|
||||
# ---------------------------------------------------------------------------
|
||||
@triton.jit
|
||||
def _swa_scatter_kernel(
|
||||
kv_ptr, # [T, D] bf16
|
||||
state_slot_ptr, # [T] int
|
||||
positions_ptr, # [T] int
|
||||
final_pos_ptr, # [T] int
|
||||
unified_ptr, # [pages, D] bf16
|
||||
n_rows,
|
||||
ring_stride, # SWA ring per-slot stride
|
||||
win: tl.constexpr,
|
||||
D: tl.constexpr,
|
||||
HAS_FINAL: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
if row >= n_rows:
|
||||
return
|
||||
pos = tl.load(positions_ptr + row)
|
||||
if HAS_FINAL:
|
||||
fp = tl.load(final_pos_ptr + row)
|
||||
if pos <= fp - win:
|
||||
return
|
||||
s = tl.load(state_slot_ptr + row)
|
||||
loc = s * ring_stride + (pos % ring_stride)
|
||||
offs = tl.arange(0, BLOCK_D)
|
||||
mask = offs < D
|
||||
vals = tl.load(kv_ptr + row * D + offs, mask=mask, other=0.0)
|
||||
tl.store(unified_ptr + loc * D + offs, vals, mask=mask)
|
||||
|
||||
|
||||
def store_swa_into_unified(
|
||||
*,
|
||||
kv: torch.Tensor, # [T, head_dim] bf16
|
||||
state_slot: torch.Tensor, # [T] int
|
||||
positions: torch.Tensor, # [T] int
|
||||
unified_kv: torch.Tensor, # [pages, head_dim] bf16
|
||||
win: int, # SWA attention window length
|
||||
ring_stride: int, # SWA ring stride
|
||||
final_pos: Optional[torch.Tensor] = None, # [T] req's last position
|
||||
) -> None:
|
||||
n_rows, D = kv.shape
|
||||
if n_rows == 0:
|
||||
return
|
||||
|
||||
has_final = final_pos is not None
|
||||
fp_arg = final_pos if has_final else positions
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Ragged indptr helper (shared by the decode streams + prefill builders)
|
||||
# ---------------------------------------------------------------------------
|
||||
def _lengths_to_indptr(lengths: torch.Tensor) -> torch.Tensor:
|
||||
"""[N] int32 per-token lengths -> [N+1] int32 indptr"""
|
||||
return F.pad(torch.cumsum(lengths, dim=0, dtype=torch.int32), (1, 0))
|
||||
|
||||
|
||||
def decode(
|
||||
*,
|
||||
q: torch.Tensor, # [T, H, D] (local heads)
|
||||
unified_kv: torch.Tensor, # [pages, D] bf16
|
||||
kv_indices: torch.Tensor,
|
||||
kv_indptr: torch.Tensor,
|
||||
attn_sink: torch.Tensor, # [H] fp32
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
return sparse_attn_v4_paged_decode(
|
||||
q, unified_kv, kv_indices, kv_indptr, attn_sink, softmax_scale
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fill_compress_tail_kernel(
|
||||
indices_ptr, # [*] int32 (out)
|
||||
indptr_ptr, # [N+1] int32
|
||||
prefix_len_ptr, # [N] int
|
||||
page_idx_ptr, # [N, Wc] int
|
||||
valid_len_ptr, # [N] int
|
||||
swa_pages,
|
||||
Wc: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
"""Per token: write valid_len compressed slots (swa_pages+page_idx, -1 for empty) into the stream tail at indptr[t]+prefix_len[t]."""
|
||||
t = tl.program_id(0)
|
||||
cbase = tl.load(indptr_ptr + t) + tl.load(prefix_len_ptr + t).to(tl.int32)
|
||||
nc = tl.load(valid_len_ptr + t).to(tl.int32)
|
||||
for off in tl.range(0, Wc, BLOCK):
|
||||
j = off + tl.arange(0, BLOCK)
|
||||
m = j < nc
|
||||
j_clamped = tl.minimum(j, Wc - 1)
|
||||
pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=-1).to(tl.int32)
|
||||
slot = tl.where(pi >= 0, pi + swa_pages, -1)
|
||||
tl.store(indices_ptr + cbase + j, slot, mask=m)
|
||||
|
||||
|
||||
def fill_compress_tail(
|
||||
*,
|
||||
indices: torch.Tensor,
|
||||
indptr: torch.Tensor,
|
||||
prefix_len: torch.Tensor,
|
||||
page_indices: torch.Tensor, # [N, Wc] int32
|
||||
valid_len: torch.Tensor,
|
||||
swa_pages: int,
|
||||
) -> None:
|
||||
N, Wc = page_indices.shape
|
||||
if N == 0:
|
||||
return
|
||||
assert prefix_len.is_contiguous() and page_indices.is_contiguous()
|
||||
assert valid_len.is_contiguous()
|
||||
_fill_compress_tail_kernel[(N,)](
|
||||
indices,
|
||||
indptr,
|
||||
prefix_len,
|
||||
page_indices,
|
||||
valid_len,
|
||||
swa_pages,
|
||||
Wc=Wc,
|
||||
BLOCK=min(1024, triton.next_power_of_2(max(Wc, 1))),
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
|
||||
def build_decode_streams(
|
||||
*,
|
||||
state_slot: torch.Tensor, # [N] int
|
||||
positions: torch.Tensor, # [N] int
|
||||
swa_len: torch.Tensor, # [N] int
|
||||
hca_len: torch.Tensor, # [N] int
|
||||
csa_len: torch.Tensor, # [N] int
|
||||
hca_page_indices: torch.Tensor, # [N, hca_width] int32
|
||||
csa_width: int,
|
||||
win: int, # SWA attention window length
|
||||
ring_stride: int, # SWA ring per-slot stride
|
||||
swa_pages: int,
|
||||
) -> Tuple[
|
||||
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
|
||||
]:
|
||||
device = state_slot.device
|
||||
N = state_slot.shape[0]
|
||||
assert state_slot.is_contiguous() and positions.is_contiguous()
|
||||
state_slot = state_slot.to(torch.int32)
|
||||
positions = positions.to(torch.int32)
|
||||
hca_width = hca_page_indices.shape[1]
|
||||
|
||||
swa_p = _lengths_to_indptr(swa_len)
|
||||
hca_p = _lengths_to_indptr(swa_len + hca_len)
|
||||
csa_p = _lengths_to_indptr(swa_len + csa_len)
|
||||
|
||||
swa_i = torch.empty(N * win, dtype=torch.int32, device=device)
|
||||
hca_i = torch.empty(N * (win + hca_width), dtype=torch.int32, device=device)
|
||||
csa_i = torch.empty(N * (win + csa_width), dtype=torch.int32, device=device)
|
||||
|
||||
if N > 0:
|
||||
batch_id = torch.arange(N, dtype=torch.int32, device=device)
|
||||
write_v4_paged_decode_indices(
|
||||
state_slot_per_seq=state_slot,
|
||||
batch_id_per_token=batch_id,
|
||||
positions=positions,
|
||||
swa_indptr=swa_p,
|
||||
csa_indptr=csa_p,
|
||||
hca_indptr=hca_p,
|
||||
swa_indices=swa_i,
|
||||
csa_indices=csa_i,
|
||||
hca_indices=hca_i,
|
||||
T=N,
|
||||
win=win,
|
||||
ring_stride=ring_stride,
|
||||
)
|
||||
fill_compress_tail(
|
||||
indices=hca_i,
|
||||
indptr=hca_p,
|
||||
prefix_len=swa_len,
|
||||
page_indices=hca_page_indices[:N],
|
||||
valid_len=hca_len,
|
||||
swa_pages=swa_pages,
|
||||
)
|
||||
return swa_i, swa_p, hca_i, hca_p, csa_i, csa_p
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prefill index builder (ragged-packed: paged prefix + flat extend)
|
||||
# ---------------------------------------------------------------------------
|
||||
@triton.jit
|
||||
def _prefill_lengths_kernel(
|
||||
positions_ptr, # [T] int
|
||||
chunk_start_ptr, # [T] int
|
||||
page_idx_ptr, # [T, Wc] int (front-packed, -1 padded)
|
||||
prefix_len_ptr, # [T] int32 out
|
||||
extend_len_ptr, # [T] int32 out
|
||||
win: tl.constexpr,
|
||||
Wc: tl.constexpr,
|
||||
HAS_COMPRESS: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
"""Per token: write extend/prefix segment lengths"""
|
||||
t = tl.program_id(0)
|
||||
pos = tl.load(positions_ptr + t).to(tl.int32)
|
||||
cstart = tl.load(chunk_start_ptr + t).to(tl.int32)
|
||||
tpic = pos - cstart
|
||||
swa_low = tl.maximum(pos - win + 1, 0)
|
||||
extend_count = tl.minimum(tpic + 1, win)
|
||||
prefix_swa_count = tl.minimum(tl.maximum(cstart - swa_low, 0), win)
|
||||
tl.store(extend_len_ptr + t, extend_count)
|
||||
if HAS_COMPRESS:
|
||||
nc = 0
|
||||
for off in tl.range(0, Wc, BLOCK):
|
||||
j = off + tl.arange(0, BLOCK)
|
||||
m = j < Wc
|
||||
j_clamped = tl.minimum(j, Wc - 1)
|
||||
pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=-1)
|
||||
nc += tl.sum(tl.where(m & (pi >= 0), 1, 0))
|
||||
tl.store(prefix_len_ptr + t, prefix_swa_count + nc)
|
||||
else:
|
||||
tl.store(prefix_len_ptr + t, prefix_swa_count)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _build_prefill_indices_kernel(
|
||||
positions_ptr, # [T] int
|
||||
chunk_start_ptr, # [T] int
|
||||
cu_q_ptr, # [T] int
|
||||
state_slot_ptr, # [T] int
|
||||
page_idx_ptr, # [T, Wc] int (front-packed, -1 padded)
|
||||
pre_indptr_ptr, # [T+1] int32 (prefix stream ragged indptr)
|
||||
ext_indptr_ptr, # [T+1] int32 (extend stream ragged indptr)
|
||||
pre_out_ptr,
|
||||
ext_out_ptr,
|
||||
swa_pages,
|
||||
ring_stride, # SWA ring per-slot stride
|
||||
win: tl.constexpr,
|
||||
Wc: tl.constexpr,
|
||||
HAS_COMPRESS: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
"""Per token: write extend rows + prefix (SWA ring slots ++ swa_pages+compressed slots) as two ragged segments"""
|
||||
t = tl.program_id(0)
|
||||
pos = tl.load(positions_ptr + t).to(tl.int32)
|
||||
cstart = tl.load(chunk_start_ptr + t).to(tl.int32)
|
||||
cuq = tl.load(cu_q_ptr + t).to(tl.int32)
|
||||
s = tl.load(state_slot_ptr + t).to(tl.int32)
|
||||
|
||||
tpic = pos - cstart
|
||||
swa_low = tl.maximum(pos - win + 1, 0)
|
||||
extend_count = tl.minimum(tpic + 1, win)
|
||||
prefix_swa_count = tl.minimum(tl.maximum(cstart - swa_low, 0), win)
|
||||
|
||||
ebase = tl.load(ext_indptr_ptr + t)
|
||||
pbase = tl.load(pre_indptr_ptr + t)
|
||||
|
||||
# ---- extend: rows into the current-chunk kv tensor ----
|
||||
ext_start = cuq + tpic - extend_count + 1
|
||||
for off in tl.range(0, win, BLOCK):
|
||||
k = off + tl.arange(0, BLOCK)
|
||||
m = k < extend_count
|
||||
tl.store(ext_out_ptr + ebase + k, ext_start + k, mask=m)
|
||||
|
||||
# ---- prefix SWA: prior-chunk ring slots (stride = ring_stride) ----
|
||||
for off in tl.range(0, win, BLOCK):
|
||||
k = off + tl.arange(0, BLOCK)
|
||||
m = k < prefix_swa_count
|
||||
gp = swa_low + k
|
||||
tl.store(pre_out_ptr + pbase + k, s * ring_stride + (gp % ring_stride), mask=m)
|
||||
|
||||
# ---- prefix compressed: swa_pages + front-packed page index ----
|
||||
if HAS_COMPRESS:
|
||||
nc = tl.load(pre_indptr_ptr + t + 1) - pbase - prefix_swa_count
|
||||
cbase = pbase + prefix_swa_count
|
||||
for off in tl.range(0, Wc, BLOCK):
|
||||
j = off + tl.arange(0, BLOCK)
|
||||
m = j < nc
|
||||
j_clamped = tl.minimum(j, Wc - 1)
|
||||
pi = tl.load(page_idx_ptr + t * Wc + j_clamped, mask=m, other=0).to(
|
||||
tl.int32
|
||||
)
|
||||
tl.store(pre_out_ptr + cbase + j, pi + swa_pages, mask=m)
|
||||
|
||||
|
||||
def build_prefill_indices(
|
||||
*,
|
||||
compress_ratio: int,
|
||||
state_slot: torch.Tensor, # [T] int (per token)
|
||||
positions: torch.Tensor, # [T] int (per token absolute position)
|
||||
chunk_start: torch.Tensor, # [T] int (absolute start of this chunk for token's seq)
|
||||
cu_q: torch.Tensor, # [T] int (row in extend `kv` of the seq's first chunk token)
|
||||
win: int, # SWA attention window length
|
||||
ring_stride: int, # SWA ring per-slot stride / modulo (win_with_spec)
|
||||
swa_pages: int,
|
||||
c128_page_indices: Optional[torch.Tensor],
|
||||
c4_sparse_page_indices: Optional[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build ragged prefill indices: prefix (SWA ring + swa_pages + compressed) into unified_kv + extend into current-chunk kv; returns (prefix_indices, prefix_indptr, extend_indices, extend_indptr)."""
|
||||
device = state_slot.device
|
||||
T = state_slot.shape[0]
|
||||
assert positions.is_contiguous() and chunk_start.is_contiguous()
|
||||
assert cu_q.is_contiguous() and state_slot.is_contiguous()
|
||||
|
||||
if compress_ratio == 0:
|
||||
page_idx = None
|
||||
elif compress_ratio == 128:
|
||||
assert c128_page_indices is not None
|
||||
page_idx = c128_page_indices[:T]
|
||||
elif compress_ratio == 4:
|
||||
assert c4_sparse_page_indices is not None
|
||||
page_idx = c4_sparse_page_indices[:T]
|
||||
else:
|
||||
raise ValueError(f"bad compress_ratio {compress_ratio}")
|
||||
|
||||
has_compress = page_idx is not None
|
||||
if has_compress:
|
||||
assert page_idx.is_contiguous()
|
||||
Wc = page_idx.shape[1] if has_compress else 0
|
||||
|
||||
block = min(1024, triton.next_power_of_2(max(win, Wc, 1)))
|
||||
prefix_len = torch.empty(T, dtype=torch.int32, device=device)
|
||||
extend_len = torch.empty(T, dtype=torch.int32, device=device)
|
||||
_prefill_lengths_kernel[(T,)](
|
||||
positions,
|
||||
chunk_start,
|
||||
page_idx if has_compress else positions, # dummy ptr when no compress
|
||||
prefix_len,
|
||||
extend_len,
|
||||
win=win,
|
||||
Wc=Wc if has_compress else 1,
|
||||
HAS_COMPRESS=has_compress,
|
||||
BLOCK=block,
|
||||
num_warps=4,
|
||||
)
|
||||
kv_indptr_prefix = _lengths_to_indptr(prefix_len)
|
||||
kv_indptr_extend = _lengths_to_indptr(extend_len)
|
||||
|
||||
kv_indices_prefix = torch.empty(T * (win + Wc), dtype=torch.int32, device=device)
|
||||
kv_indices_extend = torch.empty(T * win, dtype=torch.int32, device=device)
|
||||
|
||||
_build_prefill_indices_kernel[(T,)](
|
||||
positions,
|
||||
chunk_start,
|
||||
cu_q,
|
||||
state_slot,
|
||||
page_idx if has_compress else state_slot, # dummy ptr when no compress
|
||||
kv_indptr_prefix,
|
||||
kv_indptr_extend,
|
||||
kv_indices_prefix,
|
||||
kv_indices_extend,
|
||||
swa_pages,
|
||||
ring_stride,
|
||||
win=win,
|
||||
Wc=Wc if has_compress else 1,
|
||||
HAS_COMPRESS=has_compress,
|
||||
BLOCK=block,
|
||||
num_warps=4,
|
||||
)
|
||||
return kv_indices_prefix, kv_indptr_prefix, kv_indices_extend, kv_indptr_extend
|
||||
|
||||
|
||||
def prefill(
|
||||
*,
|
||||
q: torch.Tensor, # [T, H, D]
|
||||
unified_kv: torch.Tensor, # [pages, D]
|
||||
kv_indices_prefix: torch.Tensor,
|
||||
kv_indptr_prefix: torch.Tensor,
|
||||
kv_extend: torch.Tensor, # [T, D] current-chunk K (bf16, norm+rope'd)
|
||||
kv_indices_extend: torch.Tensor,
|
||||
kv_indptr_extend: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
return sparse_attn_v4_paged_prefill(
|
||||
q,
|
||||
unified_kv,
|
||||
kv_indices_prefix,
|
||||
kv_indptr_prefix,
|
||||
kv_extend,
|
||||
kv_indices_extend,
|
||||
kv_indptr_extend,
|
||||
attn_sink,
|
||||
softmax_scale,
|
||||
)
|
||||
@@ -95,6 +95,7 @@ def _fused_qk_norm_rope_store_kernel(
|
||||
FP8_MAX: tl.constexpr,
|
||||
BYTES_PER_TOKEN: tl.constexpr,
|
||||
SWA_PAGE_SIZE: tl.constexpr,
|
||||
BF16_STORE: tl.constexpr,
|
||||
):
|
||||
pid_m = tl.program_id(0).to(tl.int64)
|
||||
pid_h = tl.program_id(1).to(tl.int64)
|
||||
@@ -206,7 +207,24 @@ def _fused_qk_norm_rope_store_kernel(
|
||||
VALUE_STRIDE: tl.constexpr = DIM_NOPE + ROPE_DIM * 2
|
||||
SCALE_BYTES: tl.constexpr = NUM_NOPE_TILES + 1
|
||||
|
||||
if HAS_SWA_STORE:
|
||||
if HAS_SWA_STORE and BF16_STORE:
|
||||
# unified_kv unified_kv: write the whole head_dim as plain bf16 into a
|
||||
# [num_slots, head_dim] bf16 cache at row=loc (no fp8 / no scale).
|
||||
loc = tl.load(swa_loc_ptr + src_id, mask=src_mask, other=0)
|
||||
row_base = loc.to(tl.int64)[:, None] * swa_cache_stride_page
|
||||
# nope
|
||||
tl.store(
|
||||
swa_cache_ptr + row_base + offs_d_full[None, :],
|
||||
kv_normed.to(swa_cache_ptr.dtype.element_ty),
|
||||
mask=src_mask[:, None] & nope_d_mask[None, :],
|
||||
)
|
||||
# pe
|
||||
tl.store(
|
||||
swa_cache_ptr + row_base + (NOPE_DIM + d_pe_offs[None, :]),
|
||||
kv_pe.to(swa_cache_ptr.dtype.element_ty),
|
||||
mask=src_mask[:, None],
|
||||
)
|
||||
elif HAS_SWA_STORE:
|
||||
loc = tl.load(swa_loc_ptr + src_id, mask=src_mask, other=0)
|
||||
page_id = loc // SWA_PAGE_SIZE
|
||||
page_off = loc % SWA_PAGE_SIZE
|
||||
@@ -288,15 +306,17 @@ def fused_qk_norm_rope_swa_store(
|
||||
swa_page_size: int = 128,
|
||||
q_out: Optional[torch.Tensor] = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
bf16_store: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Fused Q norm + KV norm + RoPE + optional FP8 paged SWA store.
|
||||
"""Fused Q norm + KV norm + RoPE + optional SWA store.
|
||||
|
||||
Args:
|
||||
q: [M, N] or [splitk, M, N] where N = num_local_heads * head_dim
|
||||
kv: [M, head_dim=512] mutated in-place (norm + RoPE)
|
||||
swa_cache: paged SWA KV pool buffer [num_pages, bytes_per_page] uint8
|
||||
swa_cache: paged SWA KV pool buffer [num_pages, bytes_per_page] uint8 OR a plain [num_slots, head_dim] bf16 cache
|
||||
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]
|
||||
"""
|
||||
head_dim = kv.shape[1]
|
||||
|
||||
@@ -375,6 +395,7 @@ def fused_qk_norm_rope_swa_store(
|
||||
FP8_MAX=fp8_info.max,
|
||||
BYTES_PER_TOKEN=bytes_per_token,
|
||||
SWA_PAGE_SIZE=swa_page_size,
|
||||
BF16_STORE=bf16_store,
|
||||
num_warps=num_warps,
|
||||
)
|
||||
return q_out
|
||||
|
||||
@@ -374,6 +374,65 @@ class DeepSeekV4LayerItem(NamedTuple):
|
||||
compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None
|
||||
|
||||
|
||||
class DeepSeekV4UnifiedKVPool:
|
||||
"""
|
||||
Layout:
|
||||
unified_kv[L]: ``[swa_pages + compress_pages, head_dim]`` bf16
|
||||
- rows ``[0, swa_pages)`` = SWA ring (``req_pool_indices * swa_window + pos % swa_window``)
|
||||
- rows ``[swa_pages, ...)`` = compressed (``swa_pages + page_index``)
|
||||
"""
|
||||
|
||||
K_PER_BLOCK = {0: 0, 4: 32, 128: 1}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
stage_ratios: List[int],
|
||||
num_slots: int,
|
||||
num_blocks: int,
|
||||
qk_nope_head_dim: int,
|
||||
qk_rope_head_dim: int,
|
||||
device: str,
|
||||
memory_saver_adapter,
|
||||
custom_mem_pool,
|
||||
swa_ring_size: int,
|
||||
):
|
||||
self.swa_ring_size = swa_ring_size
|
||||
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
|
||||
self.num_blocks = num_blocks
|
||||
self.k_per_block = dict(self.K_PER_BLOCK)
|
||||
|
||||
bufs = []
|
||||
with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
|
||||
with (
|
||||
torch.cuda.use_mem_pool(custom_mem_pool)
|
||||
if custom_mem_pool
|
||||
else nullcontext()
|
||||
):
|
||||
for ratio in stage_ratios:
|
||||
compress_pages = self.num_blocks * self.k_per_block[ratio]
|
||||
bufs.append(
|
||||
torch.zeros(
|
||||
self.swa_pages + compress_pages,
|
||||
self.head_dim,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
self.kv_buffer = bufs
|
||||
|
||||
def get_unified_kv(self, local_layer_id: int) -> torch.Tensor:
|
||||
return self.kv_buffer[local_layer_id]
|
||||
|
||||
def get_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
|
||||
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]
|
||||
return data_ptrs, data_lens, item_lens
|
||||
|
||||
|
||||
class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
|
||||
def __init__(
|
||||
@@ -395,6 +454,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
device: str,
|
||||
enable_memory_saver: bool,
|
||||
compression_ratios: List[int],
|
||||
sliding_window: int = 128,
|
||||
start_layer: Optional[int] = None,
|
||||
end_layer: Optional[int] = None,
|
||||
enable_hisparse: bool = False,
|
||||
@@ -441,6 +501,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
stage_ratios = compression_ratios[self._stage_start : self._stage_end]
|
||||
|
||||
assert page_size % swa_page_size == 0
|
||||
self.sliding_window = sliding_window
|
||||
|
||||
self.swa_size = swa_size
|
||||
self.swa_window_size = swa_page_size
|
||||
@@ -455,41 +516,75 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
c128_layer_num = sum(1 for r in stage_ratios if r == 128)
|
||||
c4_page_size = page_size // 4
|
||||
c128_page_size = page_size // 128
|
||||
self.swa_kv_pool = DeepSeekV4SingleKVPool(
|
||||
swa_size,
|
||||
swa_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
c4_kv_pool_type = DeepSeekV4SingleKVPool
|
||||
if enable_hisparse:
|
||||
c4_kv_pool_type = HiSparseC4DevicePool
|
||||
self.c4_kv_pool = c4_kv_pool_type(
|
||||
c4_size,
|
||||
c4_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
c4_layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
)
|
||||
self._unified_kv = is_unified_kv_triton()
|
||||
|
||||
self.c128_kv_pool = DeepSeekV4SingleKVPool(
|
||||
c128_size,
|
||||
c128_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
c128_layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
)
|
||||
if self._unified_kv:
|
||||
self.swa_kv_pool = None
|
||||
self.c4_kv_pool = None
|
||||
self.c128_kv_pool = None
|
||||
server_args = get_global_server_args()
|
||||
spec_extra = (
|
||||
(server_args.speculative_num_draft_tokens - 1)
|
||||
if server_args.speculative_algorithm is not None
|
||||
else 0
|
||||
)
|
||||
self.unified_kv_pool = DeepSeekV4UnifiedKVPool(
|
||||
stage_ratios=stage_ratios,
|
||||
num_slots=self.max_num_reqs + 1,
|
||||
num_blocks=self.c128_size,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
device=device,
|
||||
memory_saver_adapter=self.memory_saver_adapter,
|
||||
custom_mem_pool=self.custom_mem_pool,
|
||||
swa_ring_size=self.sliding_window + spec_extra,
|
||||
)
|
||||
|
||||
self.unified_swa_window = self.sliding_window
|
||||
self.unified_swa_ring_size = self.sliding_window + spec_extra
|
||||
self.unified_swa_pages = self.unified_kv_pool.swa_pages
|
||||
else:
|
||||
self.unified_kv_pool = None
|
||||
self.swa_kv_pool = DeepSeekV4SingleKVPool(
|
||||
swa_size,
|
||||
swa_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
)
|
||||
|
||||
c4_kv_pool_type = DeepSeekV4SingleKVPool
|
||||
if enable_hisparse:
|
||||
c4_kv_pool_type = HiSparseC4DevicePool
|
||||
self.c4_kv_pool = c4_kv_pool_type(
|
||||
c4_size,
|
||||
c4_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
c4_layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
)
|
||||
|
||||
self.c128_kv_pool = DeepSeekV4SingleKVPool(
|
||||
c128_size,
|
||||
c128_page_size,
|
||||
dtype,
|
||||
qk_nope_head_dim,
|
||||
qk_rope_head_dim,
|
||||
c128_layer_num,
|
||||
device,
|
||||
enable_memory_saver,
|
||||
)
|
||||
|
||||
indexer_size = (
|
||||
self.c4_logical_size
|
||||
@@ -513,6 +608,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
else:
|
||||
self._init_paged_compress_states(enable_memory_saver)
|
||||
|
||||
def get_unified_kv(self, layer_id: int) -> torch.Tensor:
|
||||
return self.unified_kv_pool.get_unified_kv(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
|
||||
|
||||
@@ -530,11 +628,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
data_lens: List[int] = []
|
||||
item_lens: List[int] = []
|
||||
|
||||
for bufs in [
|
||||
self.c4_kv_pool.kv_buffer,
|
||||
self.c4_indexer_kv_pool.index_k_with_scale_buffer,
|
||||
self.c128_kv_pool.kv_buffer,
|
||||
]:
|
||||
if self._unified_kv:
|
||||
buf_groups = [
|
||||
self.unified_kv_pool.kv_buffer,
|
||||
self.c4_indexer_kv_pool.index_k_with_scale_buffer,
|
||||
]
|
||||
else:
|
||||
buf_groups = [
|
||||
self.c4_kv_pool.kv_buffer,
|
||||
self.c4_indexer_kv_pool.index_k_with_scale_buffer,
|
||||
self.c128_kv_pool.kv_buffer,
|
||||
]
|
||||
|
||||
for bufs in buf_groups:
|
||||
for buf in bufs:
|
||||
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
|
||||
data_ptrs.append(buf.data_ptr())
|
||||
@@ -548,11 +654,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
data_lens: List[int] = []
|
||||
item_lens: List[int] = []
|
||||
|
||||
for buf in self.swa_kv_pool.kv_buffer:
|
||||
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
|
||||
data_ptrs.append(buf.data_ptr())
|
||||
data_lens.append(buf.nbytes)
|
||||
item_lens.append(buf[0].nbytes)
|
||||
if not self._unified_kv:
|
||||
for buf in self.swa_kv_pool.kv_buffer:
|
||||
assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D"
|
||||
data_ptrs.append(buf.data_ptr())
|
||||
data_lens.append(buf.nbytes)
|
||||
item_lens.append(buf[0].nbytes)
|
||||
|
||||
for pools in [
|
||||
self.compress_state_pools,
|
||||
|
||||
@@ -408,6 +408,7 @@ class ModelRunnerKVCacheMixin:
|
||||
c128_state_pool_size=self.c128_state_pool_size,
|
||||
page_size=self.page_size,
|
||||
swa_page_size=swa_page_size,
|
||||
sliding_window=self.model_config.window_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
state_dtype=self.state_dtype,
|
||||
qk_nope_head_dim=self.model_config.qk_nope_head_dim,
|
||||
|
||||
@@ -785,8 +785,17 @@ class MQALayer(nn.Module):
|
||||
use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
|
||||
kv: Optional[torch.Tensor]
|
||||
|
||||
if self.use_fused_qk_norm_rope:
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
unified = is_unified_kv_triton()
|
||||
is_decode = forward_batch.forward_mode.is_decode_or_idle()
|
||||
do_fused_store = (unified and is_decode) or (
|
||||
not unified and self.use_fused_qk_norm_rope
|
||||
)
|
||||
|
||||
if do_fused_store:
|
||||
if _is_gfx95_supported:
|
||||
q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant(
|
||||
q_lora,
|
||||
@@ -804,15 +813,33 @@ class MQALayer(nn.Module):
|
||||
else self.wkv(x_linear)[0]
|
||||
)
|
||||
|
||||
token_to_kv_pool = get_token_to_kv_pool()
|
||||
if unified:
|
||||
swa_ring_size = token_to_kv_pool.unified_swa_ring_size
|
||||
swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
|
||||
# ring slot = req_slot * ring + pos % ring, per token.
|
||||
# positions is per-token; req_pool_indices is per-req.
|
||||
req_slot = forward_batch.req_pool_indices.to(torch.int64)
|
||||
if req_slot.shape[0] != positions.shape[0]:
|
||||
req_slot = req_slot.repeat_interleave(
|
||||
positions.shape[0] // req_slot.shape[0]
|
||||
)
|
||||
swa_loc = (
|
||||
req_slot * swa_ring_size + positions.to(torch.int64) % swa_ring_size
|
||||
).to(torch.int32)
|
||||
swa_page_size, bf16_store = 1, True
|
||||
else:
|
||||
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
|
||||
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
|
||||
swa_page_size, bf16_store = (
|
||||
token_to_kv_pool.swa_kv_pool.page_size,
|
||||
False,
|
||||
)
|
||||
|
||||
from sglang.srt.layers.fused_qk_norm_rope_store import (
|
||||
fused_qk_norm_rope_swa_store,
|
||||
)
|
||||
|
||||
token_to_kv_pool = get_token_to_kv_pool()
|
||||
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
|
||||
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
|
||||
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
|
||||
|
||||
q = fused_qk_norm_rope_swa_store(
|
||||
q=q,
|
||||
kv=kv,
|
||||
@@ -829,9 +856,11 @@ class MQALayer(nn.Module):
|
||||
swa_page_size=swa_page_size,
|
||||
q_out=q_out,
|
||||
dtype=x.dtype,
|
||||
bf16_store=bf16_store,
|
||||
)
|
||||
kv = None
|
||||
|
||||
if use_cp:
|
||||
if not unified and use_cp:
|
||||
# DSA CP: keep bf16 kv around for the cross-rank all-gather, then
|
||||
# write to the FlashMLA cache after gather.
|
||||
kv = self._compute_kv_bf16(x, positions, qkv_a=qkv_a)
|
||||
@@ -844,7 +873,11 @@ class MQALayer(nn.Module):
|
||||
else:
|
||||
q_lora = self.q_norm(q_lora)
|
||||
q = self._compute_q_b(q_lora, positions, q_out)
|
||||
if use_cp:
|
||||
if unified:
|
||||
# unified_kv prefill: keep bf16 kv; the backend writes
|
||||
# the ring AFTER attention (2-source path).
|
||||
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
|
||||
elif use_cp:
|
||||
# NSA CP: keep bf16 kv around for the cross-rank all-gather, then
|
||||
# write to the FlashMLA cache after gather.
|
||||
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
|
||||
@@ -956,34 +989,50 @@ class MQALayer(nn.Module):
|
||||
# tell the backend to skip its own store_cache. When `kv is None`
|
||||
# (no DSA-CP), pass `q` as a sentinel for the `k is v` assert; the
|
||||
# attention path doesn't read it once `save_kv_cache=False`.
|
||||
attn_q = q_padded if q_padded is not None else q
|
||||
attn_k = kv if kv is not None else q
|
||||
save_kv_cache = False
|
||||
if forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph():
|
||||
o = attn_q.new_empty(
|
||||
(*attn_q.shape[:-1], self.attn_mqa.v_head_dim),
|
||||
)
|
||||
bcg_deepseek_v4_attention_with_output(
|
||||
attn_q,
|
||||
attn_k,
|
||||
o,
|
||||
self.attn_mqa.layer_id,
|
||||
self.compress_ratio,
|
||||
self.attn_sink,
|
||||
save_kv_cache,
|
||||
)
|
||||
else:
|
||||
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
|
||||
is_unified_kv_triton,
|
||||
)
|
||||
|
||||
if is_unified_kv_triton():
|
||||
o = attn_backend.forward(
|
||||
q=attn_q,
|
||||
q=q_out if q_out is not None else q,
|
||||
k=attn_k,
|
||||
v=attn_k,
|
||||
layer=self.attn_mqa,
|
||||
forward_batch=forward_batch,
|
||||
compress_ratio=self.compress_ratio,
|
||||
attn_sink=self.attn_sink,
|
||||
save_kv_cache=save_kv_cache,
|
||||
save_kv_cache=kv is not None,
|
||||
)
|
||||
o = o[:, tp_slice, :]
|
||||
else:
|
||||
attn_q = q_padded if q_padded is not None else q
|
||||
save_kv_cache = False
|
||||
if forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph():
|
||||
o = attn_q.new_empty(
|
||||
(*attn_q.shape[:-1], self.attn_mqa.v_head_dim),
|
||||
)
|
||||
bcg_deepseek_v4_attention_with_output(
|
||||
attn_q,
|
||||
attn_k,
|
||||
o,
|
||||
self.attn_mqa.layer_id,
|
||||
self.compress_ratio,
|
||||
self.attn_sink,
|
||||
save_kv_cache,
|
||||
)
|
||||
else:
|
||||
o = attn_backend.forward(
|
||||
q=attn_q,
|
||||
k=attn_k,
|
||||
v=attn_k,
|
||||
layer=self.attn_mqa,
|
||||
forward_batch=forward_batch,
|
||||
compress_ratio=self.compress_ratio,
|
||||
attn_sink=self.attn_sink,
|
||||
save_kv_cache=save_kv_cache,
|
||||
)
|
||||
o = o[:, tp_slice, :]
|
||||
fused_rope_inplace(
|
||||
o[..., -self.qk_rope_head_dim :],
|
||||
None,
|
||||
|
||||
Reference in New Issue
Block a user