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

[AMD] Enable HiSparse on ROCm (#26639)

Browse Source 
Co-authored-by: clintg6 <7388379+clintg6@users.noreply.github.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Clint
2026-06-19 11:59:45 -07:00
committed by GitHub
co-authored by clintg6 HAI
parent ca88b7f1d2
commit c436a8161a
12 changed files with 971 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/hisparse.cuh
+104 -15
View File
@@ -8,14 +8,21 @@
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <cuda_runtime.h>
#include <stdexcept>
#include <stdint.h>
#include <string>
namespace {
#ifdef USE_ROCM
constexpr int WARP_SIZE = 64;
using BallotMask = uint64_t;
constexpr BallotMask FULL_WARP_MASK = 0xFFFFFFFFFFFFFFFFull;
#else
constexpr int WARP_SIZE = 32;
using BallotMask = unsigned int;
constexpr BallotMask FULL_WARP_MASK = 0xFFFFFFFFu;
#endif
constexpr int32_t TOKEN_HIT = 0xFFFFFFFF;
constexpr int32_t HASH_EMPTY = -1;
@@ -24,6 +31,25 @@ __device__ __forceinline__ int hash_slot(int32_t key, int hash_size) {
return ((uint32_t)key * 2654435761u) % (uint32_t)hash_size;
}
#ifdef USE_ROCM
__device__ __forceinline__ void transfer_item_warp(
int32_t lane_id, const void* __restrict__ src_addr, void* __restrict__ dst_addr, int64_t item_size_bytes) {
const auto src = static_cast<const char*>(src_addr);
auto dst = static_cast<char*>(dst_addr);
const int64_t word_count = item_size_bytes / static_cast<int64_t>(sizeof(uint64_t));
const auto src_words = reinterpret_cast<const uint64_t*>(src);
auto dst_words = reinterpret_cast<uint64_t*>(dst);
for (int64_t i = lane_id; i < word_count; i += WARP_SIZE) {
dst_words[i] = src_words[i];
}
const int64_t tail_start = word_count * static_cast<int64_t>(sizeof(uint64_t));
for (int64_t i = tail_start + lane_id; i < item_size_bytes; i += WARP_SIZE) {
dst[i] = src[i];
}
}
#else
__device__ __forceinline__ void
transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) {
// 128-bit bulk transfer via paired 64-bit loads (avoids alignment issues with uint4)
@@ -51,6 +77,15 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_
asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst8 + lane_id), "l"(tmp) : "memory");
}
}
#endif
__device__ __forceinline__ int popc_mask(BallotMask mask) {
#ifdef USE_ROCM
return __popcll(mask);
#else
return __popc(mask);
#endif
}
template <int BLOCK_SIZE>
__global__ __launch_bounds__(BLOCK_SIZE, 1) void transfer_cache_dsv4_mla_kernel(
@@ -113,15 +148,15 @@ __device__ __forceinline__ int warp_inclusive_scan(int* s_data, int lane_id, int
int val = (idx < count) ? s_data[idx] : 0;
#pragma unroll
for (int i = 1; i < 32; i *= 2) {
int n = __shfl_up_sync(0xffffffff, val, i);
for (int i = 1; i < WARP_SIZE; i *= 2) {
int n = __shfl_up_sync(FULL_WARP_MASK, val, i);
if (lane_id >= i) val += n;
}
val += accumulator;
if (idx < count) {
s_data[idx] = val;
}
accumulator = __shfl_sync(0xffffffff, val, 31);
accumulator = __shfl_sync(FULL_WARP_MASK, val, WARP_SIZE - 1);
return accumulator;
}
@@ -188,7 +223,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int tid = threadIdx.x;
const int warp_id = tid / WARP_SIZE;
const int lane_id = tid % WARP_SIZE;
const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1;
const BallotMask lanes_before = (BallotMask(1) << lane_id) - BallotMask(1);
const int64_t rid = req_pool_indices[bid];
const int64_t seq_len = seq_lens[bid];
@@ -316,22 +351,35 @@ __global__ void load_cache_to_device_buffer_kernel(
int local_hit_offset = 0;
int local_evict_offset = 0;
if (has_valid_chunk) {
const unsigned int hit_mask = __ballot_sync(0xFFFFFFFF, is_hit);
const unsigned int evict_mask = __ballot_sync(0xFFFFFFFF, is_evictable);
local_hit_offset = __popc(hit_mask & lanes_before);
local_evict_offset = __popc(evict_mask & lanes_before);
const BallotMask hit_mask = __ballot_sync(FULL_WARP_MASK, is_hit);
const BallotMask evict_mask = __ballot_sync(FULL_WARP_MASK, is_evictable);
local_hit_offset = popc_mask(hit_mask & lanes_before);
local_evict_offset = popc_mask(evict_mask & lanes_before);
if (lane_id == 0) {
s_chunk_offset[chunk_idx + 1] = __popc(hit_mask);
s_evict_chunk_offset[chunk_idx + 1] = __popc(evict_mask);
s_chunk_offset[chunk_idx + 1] = popc_mask(hit_mask);
s_evict_chunk_offset[chunk_idx + 1] = popc_mask(evict_mask);
}
}
__syncthreads();
if (warp_id == 0) {
#ifdef USE_ROCM
// ROCm wavefront64: WARP_SIZE (64) > NUM_WARPS (16 at block_size=1024),
// so the wide-count form below would let lanes beyond this iteration's
// NUM_WARPS-wide window write the accumulator into s_chunk_offset
// positions belonging to future iterations, corrupting their reads.
// Bound the scan window to NUM_WARPS lanes.
const int scan_offset = iter * NUM_WARPS + 1;
const int scan_count = min(scan_offset + NUM_WARPS, NUM_BUFFER_CHUNKS + 1);
total_hit_count = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_hit_count);
total_evict_count =
warp_inclusive_scan(s_evict_chunk_offset, lane_id, scan_offset, scan_count, total_evict_count);
#else
total_hit_count =
warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count);
total_evict_count =
warp_inclusive_scan(s_evict_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count);
#endif
if (tid == 0) {
s_total_hits = total_hit_count;
}
@@ -380,9 +428,9 @@ __global__ void load_cache_to_device_buffer_kernel(
}
if (has_valid_chunk) {
const unsigned int miss_mask = __ballot_sync(0xFFFFFFFF, is_miss);
local_miss_offset = __popc(miss_mask & lanes_before);
const int warp_miss_count = __popc(miss_mask);
const BallotMask miss_mask = __ballot_sync(FULL_WARP_MASK, is_miss);
local_miss_offset = popc_mask(miss_mask & lanes_before);
const int warp_miss_count = popc_mask(miss_mask);
if (lane_id == 0) {
s_chunk_offset[chunk_idx + 1] = warp_miss_count;
}
@@ -390,7 +438,13 @@ __global__ void load_cache_to_device_buffer_kernel(
__syncthreads();
if (warp_id == 0) {
#ifdef USE_ROCM
const int scan_offset = iter * NUM_WARPS + 1;
const int scan_count = min(scan_offset + NUM_WARPS, NUM_TOKEN_CHUNKS + 1);
total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_misses);
#else
total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses);
#endif
}
__syncthreads();
@@ -410,6 +464,24 @@ __global__ void load_cache_to_device_buffer_kernel(
// Write back LRU order: evictables at front (LRU), hits at back (MRU).
{
const int total_evictable = HOT_BUFFER_SIZE - s_total_hits;
#ifdef USE_ROCM
// ROCm: cap writeback threads at 512 for large kernels.
constexpr int LRU_WRITEBACK_THREADS = (BLOCK_SIZE > 512) ? 512 : BLOCK_SIZE;
if (tid < LRU_WRITEBACK_THREADS) {
for (int i = tid; i < HOT_BUFFER_SIZE; i += LRU_WRITEBACK_THREADS) {
if (i < total_misses) {
// Misses: just loaded from host, place right before hits
req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i];
} else if (i < total_evictable) {
// Remaining evictables: truly stale, dest at LRU front
req_lru_slots[i - total_misses] = s_lru_slots_out[HOT_BUFFER_SIZE - 1 - i];
} else {
// Hits: source at forward end, dest at MRU back
req_lru_slots[i] = s_lru_slots_out[i - total_evictable];
}
}
}
#else
for (int i = tid; i < HOT_BUFFER_SIZE; i += BLOCK_SIZE) {
if (i < total_misses) {
// Misses: just loaded from host, place right before hits
@@ -422,6 +494,7 @@ __global__ void load_cache_to_device_buffer_kernel(
req_lru_slots[i] = s_lru_slots_out[i - total_evictable];
}
}
#endif
}
// each warp copies one miss directly, can be separated into a new kernel if parallelism is a concern
@@ -433,7 +506,20 @@ __global__ void load_cache_to_device_buffer_kernel(
const int64_t dst_loc = static_cast<int64_t>(req_device_buffer_locs[evict_slot]);
if constexpr (IsDsv4Layout) {
// DSv4 path: page-padded device layout + page-padded host layout, K-only.
#ifdef USE_ROCM
// ROCm path: host cache and device buffer both use the page-padded C4
// layout (same as the write path and the CUDA branch). We can't reuse
// device::hisparse::transfer_item here because its warp logic is hardcoded
// to a 32-lane warp; on wavefront64 we use the warp-width-agnostic
// transfer_item_warp with paged source and destination addressing.
using namespace device::hisparse;
const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast<int32_t>(dst_loc));
const auto [src_value_ptr, src_scale_ptr] =
get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(src_loc));
transfer_item_warp(lane_id, src_value_ptr, dst_value_ptr, kValueBytes);
transfer_item_warp(lane_id, src_scale_ptr, dst_scale_ptr, kScaleBytes);
#else
// CUDA path: page-padded device layout + page-padded host layout, K-only.
// The host cache is pinned DRAM but uses the same row layout as the GPU C4
// cache, so use the page-padded address calculation for both ends.
device::hisparse::transfer_item(
@@ -441,6 +527,7 @@ __global__ void load_cache_to_device_buffer_kernel(
/*src_cache=*/const_cast<void*>(host_cache_k),
/*dst_index=*/static_cast<int32_t>(dst_loc),
/*src_index=*/static_cast<int32_t>(src_loc));
#endif
} else {
// Generic path: device + host both linear, stride = item_size_bytes.
const auto src_k = static_cast<const char*>(host_cache_k) + src_loc * item_size_bytes;
@@ -487,9 +574,11 @@ void load_cache_to_device_buffer(
// seq_lens and req_pool_indices; the correct combo is selected at runtime.
auto launch = [&](auto kernel_fn, const auto* seq_lens_ptr, const auto* req_pool_indices_ptr) {
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>::BYTES;
#ifndef USE_ROCM
if constexpr (smem_bytes > 48u * 1024u) {
cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
}
#endif
LaunchKernel(bs, BLOCK_SIZE, device, smem_bytes)(
kernel_fn,
static_cast<const int32_t*>(top_k_tokens.data_ptr()),
python/sglang/srt/arg_groups/hisparse_hook.py
+70 -23
View File
@@ -8,20 +8,34 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
# Backend/dtype pairing: flashmla_sparse only takes BF16 KV;
# flashmla_kv only supports FP8 (it always reads KV as FP8 via
# is_fp8_kvcache=True, inline-quantizing BF16 would defeat HiSparse).
_HISPARSE_ALLOWED_BACKENDS_BY_DTYPE = {
HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = {
"bfloat16": {"flashmla_sparse"},
"fp8_e4m3": {"flashmla_kv"},
}
HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "aiter"}
HISPARSE_KV_CACHE_DTYPES = ("bfloat16", "fp8_e4m3")
def _is_hip() -> bool:
from sglang.srt.server_args import is_hip
return is_hip()
def _hisparse_default_backend(kv_cache_dtype: str) -> str:
if _is_hip():
return "tilelang"
return "flashmla_kv" if kv_cache_dtype == "fp8_e4m3" else "flashmla_sparse"
def _hisparse_allowed_backends(kv_cache_dtype: str) -> set[str]:
if _is_hip():
return HISPARSE_ROCM_DSA_BACKENDS
return HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE.get(
kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"}
)
def apply_hisparse_dsa_backend_defaults(
server_args: ServerArgs,
user_set_prefill: bool,
@@ -30,8 +44,8 @@ def apply_hisparse_dsa_backend_defaults(
) -> bool:
"""Pick DSA backends for --enable-hisparse based on KV dtype.
BF16 KV -> flashmla_sparse, FP8 KV -> flashmla_kv. Returns True if hisparse
handled backend selection (caller should skip its own default logic).
CUDA uses dtype-specific FlashMLA backends; ROCm uses TileLang. Returns
True if hisparse handled backend selection.
"""
if not server_args.enable_hisparse:
return False
@@ -48,6 +62,35 @@ def apply_hisparse_dsa_backend_defaults(
return True
def validate_hisparse_dsa_backend(
server_args: ServerArgs, attr: str, label: str
) -> None:
backend = getattr(server_args, attr)
allowed_backends = _hisparse_allowed_backends(server_args.kv_cache_dtype)
if backend is not None and backend not in allowed_backends:
raise ValueError(
f"HiSparse supports DSA {label} backend(s) {sorted(allowed_backends)} "
f"on this platform with --kv-cache-dtype={server_args.kv_cache_dtype}, "
f"but got --dsa-{label}-backend={backend}. "
f"Please use --dsa-{label}-backend="
f"{_hisparse_default_backend(server_args.kv_cache_dtype)} "
"or omit it."
)
def validate_hisparse_kv_cache_dtype(server_args: ServerArgs) -> None:
if server_args.kv_cache_dtype in HISPARSE_KV_CACHE_DTYPES:
return
choices = " or ".join(
f"--kv-cache-dtype={dtype}" for dtype in HISPARSE_KV_CACHE_DTYPES
)
raise ValueError(
f"HiSparse requires one of {HISPARSE_KV_CACHE_DTYPES} KV cache dtypes, "
f"but got --kv-cache-dtype={server_args.kv_cache_dtype}. Please use {choices}."
)
def validate_hisparse(server_args: ServerArgs) -> None:
"""Validate --enable-hisparse constraints (model class, radix cache, DSA backend)."""
if not server_args.enable_hisparse:
@@ -60,6 +103,7 @@ def validate_hisparse(server_args: ServerArgs) -> None:
hf_config = server_args.get_model_config().hf_config
is_v4_hisparse = is_deepseek_v4(hf_config)
is_hip = _is_hip()
assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
@@ -71,27 +115,30 @@ def validate_hisparse(server_args: ServerArgs) -> None:
# DSv4 hisparse handles its own dtype/backend pairing elsewhere; the dtype-
# aware checks below only apply to the DSA hisparse path.
if is_v4_hisparse:
if is_hip and is_v4_hisparse:
# TEMPORARY GUARD: DSv4 HiSparse is not supported on the unified-KV path.
# In unified-KV mode c4_kv_pool is None, so DeepSeekV4HiSparseTokenToKVPoolAllocator
# cannot attach and pool init dies with a cryptic AssertionError. Fail fast
# at startup with a clear message instead. Remove once unified-KV HiSparse lands.
from sglang.srt.layers.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_triton,
)
if is_unified_kv_triton():
raise ValueError(
"--enable-hisparse is not supported with the unified-KV path on ROCm"
"(SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton) for DeepSeek-V4: "
"HiSparse currently requires the separate packed KV layout. "
"Either set SGLANG_HACK_FLASHMLA_BACKEND=triton, or run without "
"--enable-hisparse."
)
return
if server_args.kv_cache_dtype not in ("bfloat16", "auto", "fp8_e4m3"):
raise ValueError(
f"HiSparse requires bfloat16 or fp8_e4m3 KV cache, "
f"but got --kv-cache-dtype={server_args.kv_cache_dtype}. "
f"Please use --kv-cache-dtype=bfloat16 or fp8_e4m3."
)
validate_hisparse_kv_cache_dtype(server_args)
allowed_backends = _HISPARSE_ALLOWED_BACKENDS_BY_DTYPE.get(
server_args.kv_cache_dtype, {"flashmla_sparse", "flashmla_kv"}
)
for attr, label in [
("dsa_prefill_backend", "prefill"),
("dsa_decode_backend", "decode"),
]:
backend = getattr(server_args, attr)
if backend is not None and backend not in allowed_backends:
raise ValueError(
f"HiSparse with --kv-cache-dtype={server_args.kv_cache_dtype} requires "
f"--dsa-{label}-backend in {sorted(allowed_backends)}, "
f"but got {backend}."
)
validate_hisparse_dsa_backend(server_args, attr, label)
python/sglang/srt/layers/attention/dsa_backend.py
+208 -6
View File
@@ -65,10 +65,13 @@ _is_hip = is_hip()
if _is_hip:
from sglang.srt.layers.attention.dsa.triton_kernel import get_valid_kv_indices
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
try:
from aiter import ( # noqa: F401
flash_attn_varlen_func,
get_mla_metadata_info_v1,
get_mla_metadata_v1,
mha_batch_prefill_func,
paged_attention_ragged,
)
@@ -361,6 +364,24 @@ class DeepseekSparseAttnBackend(
self.head_repeat_factor = (
16 // self.num_q_heads if self.num_q_heads < 16 else 1
)
self.num_head_padded = self.num_q_heads * self.head_repeat_factor
self.aiter_dsa_max_split_per_batch = 64
self.aiter_dsa_metadata_capacity = 0
self.aiter_dsa_metadata_max_seqlen_q = 0
self.aiter_dsa_metadata_q_dtype = None
self.aiter_dsa_metadata_kv_dtype = None
self.aiter_dsa_kv_last_page_lens = None
self.aiter_dsa_work_metadata = None
if (
self.dsa_prefill_impl == "aiter" or self.dsa_decode_impl == "aiter"
) and model_runner.kv_cache_dtype == fp8_dtype:
self._ensure_aiter_dsa_decode_metadata_buffer(
max_seqlen_q=1,
batch_size=max_bs,
q_dtype=torch.bfloat16,
kv_dtype=fp8_dtype,
)
# Speculative decoding
self.topk = model_runner.server_args.speculative_eagle_topk or 0
@@ -387,6 +408,145 @@ class DeepseekSparseAttnBackend(
else:
self.workspace_buffer = None
def _make_aiter_dsa_decode_metadata_buffer(
self,
max_seqlen_q: int,
batch_size: int,
q_dtype: torch.dtype,
kv_dtype: torch.dtype,
):
(
(work_metadata_size, work_metadata_type),
(work_indptr_size, work_indptr_type),
(work_info_set_size, work_info_set_type),
(reduce_indptr_size, reduce_indptr_type),
(reduce_final_map_size, reduce_final_map_type),
(reduce_partial_map_size, reduce_partial_map_type),
) = get_mla_metadata_info_v1(
batch_size,
max_seqlen_q,
self.num_head_padded,
q_dtype,
kv_dtype,
is_sparse=True,
fast_mode=False,
num_kv_splits=self.aiter_dsa_max_split_per_batch,
intra_batch_mode=True,
)
return (
torch.empty(
work_metadata_size, dtype=work_metadata_type, device=self.device
),
torch.empty(work_indptr_size, dtype=work_indptr_type, device=self.device),
torch.empty(
work_info_set_size, dtype=work_info_set_type, device=self.device
),
torch.empty(
reduce_indptr_size, dtype=reduce_indptr_type, device=self.device
),
torch.empty(
reduce_final_map_size, dtype=reduce_final_map_type, device=self.device
),
torch.empty(
reduce_partial_map_size,
dtype=reduce_partial_map_type,
device=self.device,
),
)
def _ensure_aiter_dsa_decode_metadata_buffer(
self,
max_seqlen_q: int,
batch_size: int,
q_dtype: torch.dtype,
kv_dtype: torch.dtype,
) -> None:
if (
self.aiter_dsa_work_metadata is not None
and self.aiter_dsa_metadata_capacity >= batch_size
and self.aiter_dsa_metadata_max_seqlen_q == max_seqlen_q
and self.aiter_dsa_metadata_q_dtype == q_dtype
and self.aiter_dsa_metadata_kv_dtype == kv_dtype
):
return
(
self.aiter_dsa_work_metadata,
self.aiter_dsa_work_indptr,
self.aiter_dsa_work_info_set,
self.aiter_dsa_reduce_indptr,
self.aiter_dsa_reduce_final_map,
self.aiter_dsa_reduce_partial_map,
) = self._make_aiter_dsa_decode_metadata_buffer(
max_seqlen_q=max_seqlen_q,
batch_size=batch_size,
q_dtype=q_dtype,
kv_dtype=kv_dtype,
)
self.aiter_dsa_kv_last_page_lens = torch.ones(
(batch_size,), dtype=torch.int32, device=self.device
)
self.aiter_dsa_metadata_capacity = batch_size
self.aiter_dsa_metadata_max_seqlen_q = max_seqlen_q
self.aiter_dsa_metadata_q_dtype = q_dtype
self.aiter_dsa_metadata_kv_dtype = kv_dtype
def _prepare_aiter_dsa_decode_metadata(
self,
qo_indptr: torch.Tensor,
kv_indptr: torch.Tensor,
bs: int,
max_seqlen_q: int,
q_dtype: torch.dtype,
kv_dtype: torch.dtype,
) -> dict:
self._ensure_aiter_dsa_decode_metadata_buffer(
max_seqlen_q=max_seqlen_q,
batch_size=bs,
q_dtype=q_dtype,
kv_dtype=kv_dtype,
)
self.aiter_dsa_kv_last_page_lens[:bs].fill_(1)
kv_last_page_lens = self.aiter_dsa_kv_last_page_lens[:bs]
get_mla_metadata_v1(
qo_indptr,
kv_indptr,
kv_last_page_lens,
self.num_head_padded,
1,
False,
self.aiter_dsa_work_metadata,
self.aiter_dsa_work_info_set,
self.aiter_dsa_work_indptr,
self.aiter_dsa_reduce_indptr,
self.aiter_dsa_reduce_final_map,
self.aiter_dsa_reduce_partial_map,
page_size=1,
kv_granularity=16,
max_seqlen_qo=max_seqlen_q,
uni_seqlen_qo=max_seqlen_q,
fast_mode=False,
topk=self.dsa_index_topk,
max_split_per_batch=self.aiter_dsa_max_split_per_batch,
intra_batch_mode=True,
dtype_q=q_dtype,
dtype_kv=kv_dtype,
)
return {
"kv_last_page_lens": kv_last_page_lens,
"work_meta_data": self.aiter_dsa_work_metadata,
"work_indptr": self.aiter_dsa_work_indptr,
"work_info_set": self.aiter_dsa_work_info_set,
"reduce_indptr": self.aiter_dsa_reduce_indptr,
"reduce_final_map": self.aiter_dsa_reduce_final_map,
"reduce_partial_map": self.aiter_dsa_reduce_partial_map,
"intra_batch_mode": True,
"num_kv_splits": self.aiter_dsa_max_split_per_batch,
}
def _build_paged_mqa_schedule_2d_ctx_lens(
self,
forward_mode: ForwardMode,
@@ -420,13 +580,13 @@ class DeepseekSparseAttnBackend(
f"Unsupported {self.dsa_topk_backend = } for SGLANG_DSA_FUSE_TOPK."
)
def get_device_int32_arange(self, l: int) -> torch.Tensor:
if l > len(self._arange_buf):
next_pow_of_2 = 1 << (l - 1).bit_length()
def get_device_int32_arange(self, length: int) -> torch.Tensor:
if length > len(self._arange_buf):
next_pow_of_2 = 1 << (length - 1).bit_length()
self._arange_buf = torch.arange(
next_pow_of_2, device=self.device, dtype=torch.int32
)
return self._arange_buf[:l]
return self._arange_buf[:length]
def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor:
page_size = self.real_page_size
@@ -1982,6 +2142,12 @@ class DeepseekSparseAttnBackend(
q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim)
o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim)
q_scale = None
kv_scale = None
aiter_persistent_kwargs = {}
if kv_cache.dtype == fp8_dtype:
kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device)
kv_indptr = self.kv_indptr
non_minus1_mask = page_table_1 != -1
@@ -1991,6 +2157,18 @@ class DeepseekSparseAttnBackend(
kv_indices = self.kv_indices
get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs)
kv_last_page_lens = metadata.cu_seqlens_q
if kv_cache.dtype == fp8_dtype:
aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata(
metadata.cu_seqlens_q,
kv_indptr,
bs,
metadata.max_seq_len_q,
q_kernel.dtype,
kv_cache.dtype,
)
kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens")
mla_decode_fwd(
q_kernel,
kv_cache.view(-1, 1, 1, layer.head_dim),
@@ -1998,10 +2176,13 @@ class DeepseekSparseAttnBackend(
metadata.cu_seqlens_q,
kv_indptr,
kv_indices,
metadata.cu_seqlens_q,
kv_last_page_lens,
metadata.max_seq_len_q,
sm_scale=layer.scaling,
logit_cap=layer.logit_cap,
q_scale=q_scale,
kv_scale=kv_scale,
**aiter_persistent_kwargs,
)
if self.need_pad_heads:
@@ -2039,6 +2220,12 @@ class DeepseekSparseAttnBackend(
q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim)
o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim)
q_scale = None
kv_scale = None
aiter_persistent_kwargs = {}
if kv_cache.dtype == fp8_dtype:
kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device)
non_minus1_mask = page_table_1 != -1
non_minus1_counts = non_minus1_mask.sum(dim=1)
@@ -2058,6 +2245,18 @@ class DeepseekSparseAttnBackend(
cu_seqlens_q = torch.arange(
0, num_tokens + 1, dtype=torch.int32, device=self.device
)
kv_last_page_lens = cu_seqlens_q
if kv_cache.dtype == fp8_dtype:
aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata(
cu_seqlens_q,
kv_indptr,
num_tokens,
1,
q_kernel.dtype,
kv_cache.dtype,
)
kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens")
# TODO support more forward_mode
mla_decode_fwd(
q_kernel,
@@ -2066,10 +2265,13 @@ class DeepseekSparseAttnBackend(
cu_seqlens_q,
kv_indptr,
kv_indices,
cu_seqlens_q,
kv_last_page_lens,
1, # max_seq_len_q = 1 for per-token attention
sm_scale=layer.scaling,
logit_cap=layer.logit_cap,
q_scale=q_scale,
kv_scale=kv_scale,
**aiter_persistent_kwargs,
)
if self.need_pad_heads:
python/sglang/srt/managers/hisparse_coordinator.py
+23 -10
View File
@@ -5,6 +5,10 @@ from typing import List, NamedTuple, Union
import torch
from sglang.jit_kernel.hisparse import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
)
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.allocator.hisparse import (
DeepSeekV4HiSparseTokenToKVPoolAllocator,
@@ -13,19 +17,16 @@ from sglang.srt.mem_cache.allocator.hisparse import (
from sglang.srt.mem_cache.hisparse_memory_pool import (
HiSparseDSATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.memory_pool_host import (
DeepSeekV4PagedHostPool,
MLATokenToKVPoolHost,
)
from sglang.srt.utils import get_device_module
from sglang.srt.utils import get_device_module, is_hip
device_module = get_device_module()
from sglang.jit_kernel.hisparse import (
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
_is_hip = is_hip()
logger = logging.getLogger(__name__)
@@ -471,13 +472,25 @@ class HiSparseCoordinator:
:, req_pool_indices, self.device_buffer_size
] = reserved_buffer_loc.to(torch.int32)
# No need to clear prior mappings: the only consumer of the mapping
# for past tokens is the swap-in kernel, and it goes through
# top_k_device_locs returned by swap_in_selected_pages -- not via
# mapping[old_out_cache_loc] -- so stale entries are harmless.
compressed_locs = self.token_to_kv_pool_allocator.get_last_loc_compressed(
out_cache_loc
)
# ROCm: the decode remap creates a temporary hisparse device slot per
# new token (via the page_size==1 allocator path). Free the stale
# slot before pointing the mapping at the reserved device-buffer slot,
# otherwise the temporary slots leak and corrupt later swap-in lookups.
# CUDA keeps the original behavior: the swap-in kernel consumes only
# top_k_device_locs, so stale mapping entries are harmless there.
if _is_hip:
previous_locs = self.mem_pool_device._translate_loc_to_hisparse_device(
compressed_locs
)
stale_locs = previous_locs[
(previous_locs > 0) & (previous_locs != reserved_buffer_loc)
]
if stale_locs.numel() > 0:
self.token_to_kv_pool_allocator.free_hisparse_indices(stale_locs)
self.mem_pool_device.full_to_hisparse_device_index_mapping[
compressed_locs
] = reserved_buffer_loc
python/sglang/srt/mem_cache/allocator/hisparse.py
+17 -14
View File
@@ -86,10 +86,23 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return self._kvcache
def alloc(self, need_size: int):
raise NotImplementedError(
"HiSparse allocator does not support direct token allocation; "
"use alloc_extend or alloc_decode instead."
)
if self.page_size != 1:
raise NotImplementedError(
"HiSparse generic allocation is only supported for page_size=1. "
"Use alloc_extend for paged allocation."
)
logical_indices = self.logical_attn_allocator.alloc(need_size)
if logical_indices is None:
return None
hisparse_indices = self.hisparse_attn_allocator.alloc(need_size)
if hisparse_indices is None:
self.logical_attn_allocator.free(logical_indices)
return None
self.full_to_hisparse_device_index_mapping[logical_indices] = hisparse_indices
return logical_indices
def alloc_logical_only(
self,
@@ -172,8 +185,6 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
):
assert self.page_size > 1
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu, page_size=self.page_size, prefix_lens=prefix_lens_cpu
)
@@ -555,11 +566,3 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.logical_attn_allocator.free(free_index)
else:
self.free_group.append(free_index)
assert (
self.logical_attn_allocator.available_size()
<= self.logical_attn_allocator.size
)
assert (
self.hisparse_attn_allocator.available_size()
<= self.hisparse_attn_allocator.size
)
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
+6 -5
View File
@@ -227,17 +227,18 @@ class ModelRunnerKVCacheMixin:
):
return kv_cache_dim
# On HIP with TileLang backend, keep the default MLA KV cache dimension.
# FP8 attention uses the nope(512 fp8) + rope(64 fp8) layout, without extra per-block scales.
# On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout:
# nope(512 fp8) + rope(64 fp8), without extra per-block scales.
if _is_hip and (
self.server_args.dsa_prefill_backend == "tilelang"
or self.server_args.dsa_decode_backend == "tilelang"
self.server_args.dsa_prefill_backend in ("tilelang", "aiter")
or self.server_args.dsa_decode_backend in ("tilelang", "aiter")
):
return kv_cache_dim
quant_block_size = DSATokenToKVPool.quant_block_size
rope_storage_dtype = DSATokenToKVPool.rope_storage_dtype
# Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout (excluding TRTLLM and HIP+TileLang).
# Calculate override_kv_cache_dim for FP8 storage in backends that use scaled KV layout
# (excluding TRTLLM and HIP raw-layout kernels).
# kv_lora_rank + scale storage (kv_lora_rank // quant_block_size * 4 bytes) + rope dimension storage
# Note: rope dimension is stored in original dtype (bf16), not quantized to fp8
if kv_cache_dtype == torch.float8_e4m3fn:
python/sglang/srt/server_args.py
+10
View File
@@ -1946,6 +1946,16 @@ class ServerArgs:
f"Set DSA backends for {self.kv_cache_dtype} KV Cache: prefill={self.dsa_prefill_backend}, decode={self.dsa_decode_backend}."
)
def _validate_hisparse_dsa_backend(self, attr: str, label: str):
from sglang.srt.arg_groups.hisparse_hook import validate_hisparse_dsa_backend
validate_hisparse_dsa_backend(self, attr, label)
def _validate_hisparse_kv_cache_dtype(self):
from sglang.srt.arg_groups.hisparse_hook import validate_hisparse_kv_cache_dtype
validate_hisparse_kv_cache_dtype(self)
def _handle_model_specific_adjustments(self):
from sglang.srt.configs.model_config import (
get_mimo_v2_fused_qkv_expected_tp_size,