From f2bcdb05086a9e982fc4891e7eeb333da1967f70 Mon Sep 17 00:00:00 2001 From: Thomas Wang Date: Wed, 10 Jun 2026 14:13:37 +0800 Subject: [PATCH] [AMD] Add unified kv attention support in dpsk-v4 (#27380) Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> --- .../csrc/deepseek_v4/fused_norm_rope_v2.cuh | 30 +- python/sglang/jit_kernel/dsv4/compress.py | 8 +- .../layers/attention/deepseek_v4_backend.py | 1 + .../deepseek_v4_backend_hip_radix.py | 244 +++++ .../srt/layers/attention/dsv4/compressor.py | 6 +- .../layers/attention/dsv4/compressor_v2.py | 13 + .../layers/attention/dsv4/metadata_kernel.py | 7 + .../dsv4/unified_kv_kernels/env_gate.py | 9 + .../dsv4/unified_kv_kernels/paged_decode.py | 881 ++++++++++++++++++ .../paged_decode_indices.py | 187 ++++ .../dsv4/unified_kv_kernels/paged_prefill.py | 354 +++++++ .../dsv4/unified_kv_kernels/runtime.py | 440 +++++++++ .../srt/layers/fused_qk_norm_rope_store.py | 27 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 191 +++- .../model_runner_kv_cache_mixin.py | 1 + python/sglang/srt/models/deepseek_v4.py | 103 +- 16 files changed, 2418 insertions(+), 84 deletions(-) create mode 100644 python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py create mode 100644 python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py create mode 100644 python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py create mode 100644 python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py create mode 100644 python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh index 5ffe8fdde..32dd8c8ee 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh @@ -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 +template 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(); // 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(value_ptr)[tx] = cast(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 +template 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; } else { - return fused_norm_rope_flashmla; + return fused_norm_rope_flashmla; } } diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index 9fc1e29a5..ea3d8f729 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -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( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index f57486e68..bfbd82d66 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 0ad3e35e8..0c91ff41a 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index bc8b2111b..73a58e52f 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -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) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index fdbe68cc7..a499e429d 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -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( diff --git a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py index ca14cb262..34cbcde08 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py @@ -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( diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py new file mode 100644 index 000000000..1835798d2 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/env_gate.py @@ -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" diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py new file mode 100644 index 000000000..c6b9632c0 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py new file mode 100644 index 000000000..f46319e46 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode_indices.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py new file mode 100644 index 000000000..e384bc25e --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_prefill.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py new file mode 100644 index 000000000..437fc3e8b --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/runtime.py @@ -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, + ) diff --git a/python/sglang/srt/layers/fused_qk_norm_rope_store.py b/python/sglang/srt/layers/fused_qk_norm_rope_store.py index c6499902a..01a6ed8b7 100644 --- a/python/sglang/srt/layers/fused_qk_norm_rope_store.py +++ b/python/sglang/srt/layers/fused_qk_norm_rope_store.py @@ -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 diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 77a5a1387..4eddf503b 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -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, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index c32433edd..163acb873 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index deff2366b..d27ba11a9 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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,