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/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,