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

[AMD] Auto-fallback NSA indexer to page_size=1 when aiter preshuffle gluon kernel is unavailable (Deepseek v3.2) (#25205)

Browse Source
This commit is contained in:
YC Yen-Ching Tseng
2026-05-14 00:01:56 -07:00
committed by GitHub
parent 142ed710b7
commit 22bfae0d1d
5 changed files with 137 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
+19 -5
View File
@@ -4,14 +4,19 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import get_bool_env_var, is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
# aiter cp_gather kernel with preshuffle=True is only valid when the indexer
# uses the page_size=64 preshuffle layout (i.e. when the matching MQA gluon path
# is also enabled).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter:
if _use_aiter_preshuffle:
from aiter.ops.cache import cp_gather_indexer_k_quant_cache
if TYPE_CHECKING:
@@ -167,7 +172,11 @@ class GetS:
class GetKAndS:
@classmethod
def execute(cls, *args, **kwargs):
if _use_aiter:
# The aiter path uses cp_gather_indexer_k_quant_cache(preshuffle=True),
# which only matches the layout produced when the rest of the indexer
# is on the page_size=64 preshuffle path. Otherwise fall back to the
# triton implementation (which works on the page_size=1 legacy layout).
if _use_aiter_preshuffle:
return cls.aiter(*args, **kwargs)
return cls.triton(*args, **kwargs)
@@ -419,9 +428,14 @@ def _set_k_and_s_triton(
assert index_head_dim == 128
assert scale_dim == 1
if _is_hip:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert (
page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {page_size}"
else:
assert page_size == 64
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+49 -16
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import contextlib
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
@@ -12,6 +13,11 @@ from sglang.jit_kernel.fused_store_index_cache import (
fused_store_index_k_cache,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
)
from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
@@ -29,6 +35,8 @@ from sglang.srt.utils import (
is_npu,
)
logger = logging.getLogger(__name__)
global _use_multi_stream
_is_cuda = is_cuda()
_is_hip = is_hip()
@@ -36,6 +44,16 @@ _is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_fp8_fnuz = is_fp8_fnuz()
_is_gfx95_supported = is_gfx95_supported()
# Whether the aiter preshuffle paged-MQA path (page_size=64 + Preshuffle=True +
# KVBlockSize=64) can be used. Falls back to the legacy page_size=1 / KVBlockSize=1
# path when the gluon kernel is unavailable (Triton<3.5 and no AOT bundle).
_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa()
if _use_aiter and not _use_aiter_preshuffle:
logger.warning(
"ROCm NSA indexer: aiter preshuffle paged-MQA path is unavailable "
"(needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1); "
"falling back to legacy page_size=1 / KVBlockSize=1 path."
)
if _is_cuda:
try:
import deep_gemm
@@ -55,10 +73,6 @@ from sglang.srt.distributed import (
)
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
is_nsa_prefill_cp_in_seq_split,
)
from sglang.srt.layers.communicator import ScatterMode
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -431,13 +445,21 @@ class Indexer(MultiPlatformOp):
page_size = forward_batch.token_to_kv_pool.page_size
# NOTE(dark): blocksize = 64 is hardcoded in deep_gemm
if _is_hip:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert (
page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {page_size}"
else:
assert page_size == 64, "only support page size 64"
# NOTE(dark): this support extend/decode/decode+graph
block_tables = metadata.get_page_table_64()
if _is_hip and not _use_aiter_preshuffle:
block_tables = metadata.get_page_table_1()
else:
block_tables = metadata.get_page_table_64()
max_seq_len = block_tables.shape[1] * page_size
kv_cache_fp8 = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
@@ -501,7 +523,7 @@ class Indexer(MultiPlatformOp):
seqlens_32,
block_tables,
max_seq_len,
Preshuffle=_use_aiter,
Preshuffle=_use_aiter_preshuffle,
KVBlockSize=block_kv,
)
else:
@@ -565,9 +587,14 @@ class Indexer(MultiPlatformOp):
page_size = forward_batch.token_to_kv_pool.page_size
if _is_hip:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
else:
assert (
page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {page_size}"
else:
assert page_size == 64, "only support page size 64"
@@ -578,7 +605,10 @@ class Indexer(MultiPlatformOp):
)
weights = weights.squeeze(-1)
block_tables = metadata.get_page_table_64()
if _is_hip and not _use_aiter_preshuffle:
block_tables = metadata.get_page_table_1()
else:
block_tables = metadata.get_page_table_64()
assert (
forward_batch.seq_lens_cpu is not None
@@ -1034,13 +1064,16 @@ class Indexer(MultiPlatformOp):
)
return
# Fast path: AITER fused quant + cache store (HIP, preshuffle)
# Fast path: AITER fused quant + cache store
# When _use_aiter_preshuffle is True we use the new MFMA 16x16 preshuffle
# layout (page_size>=16). Otherwise we fall back to the legacy row-major
# layout with page_size=1; the same kv_cache.view works for both cases
# because page_size is 1 there.
if _use_aiter:
page_size = forward_batch.token_to_kv_pool.page_size
buf = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer(
layer_id=layer_id
)
# Reshape from (num_pages, page_size*(128+4)) uint8 to (num_pages, page_size, 132) fp8
kv_cache = buf.view(-1, page_size, 132).view(fp8_dtype)
out_loc = forward_batch.out_cache_loc
if not out_loc.is_contiguous():
@@ -1051,7 +1084,7 @@ class Indexer(MultiPlatformOp):
out_loc,
self.block_size,
self.scale_fmt,
preshuffle=True,
preshuffle=_use_aiter_preshuffle,
)
return
python/sglang/srt/layers/attention/nsa/utils.py
+41
View File
@@ -1,3 +1,4 @@
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, Union
import torch
@@ -11,8 +12,48 @@ from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import get_bool_env_var, is_hip
from sglang.srt.utils.common import ceil_align, ceil_div
@lru_cache(maxsize=1)
def aiter_can_use_preshuffle_paged_mqa() -> bool:
"""Whether aiter's preshuffle paged MQA / cache kernels can be used on this runtime.
aiter's ``deepgemm_fp8_paged_mqa_logits`` only supports ``KVBlockSize > 1`` and
``Preshuffle=True`` on its gluon kernel path. The gluon path is enabled when
Triton >= 3.5.0, OR when ``AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1`` is set
(which additionally requires that the AOT gluon kernel artifacts ship inside
the aiter wheel/image). Otherwise aiter asserts ``KVBlockSize == 1`` and
refuses ``Preshuffle=True``.
sglang's NSA indexer uses this single decision to pick:
* ``page_size``: 64 (preshuffle) vs 1 (legacy) on ROCm
* ``Preshuffle`` / ``preshuffle`` flags on the aiter MQA + cache kernels
* ``get_page_table_64`` vs ``get_page_table_1`` on the metadata
* whether ``GetKAndS.execute`` uses the aiter or the triton implementation
The result is cached so the cost is paid once per process.
Set ``SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1`` to force the legacy path even when
the gluon kernel would otherwise be available (useful for CI bisection).
"""
if not is_hip():
return False
if not get_bool_env_var("SGLANG_USE_AITER"):
return False
if get_bool_env_var("SGLANG_NSA_HIP_DISABLE_PRESHUFFLE"):
return False
if get_bool_env_var("AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS"):
return True
try:
from packaging.version import Version
return Version(Version(triton.__version__).base_version) >= Version("3.5.0")
except Exception:
return False
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
python/sglang/srt/mem_cache/memory_pool.py
+9 -3
View File
@@ -45,6 +45,7 @@ from sglang.srt.layers.attention.nsa.quant_k_cache import (
quantize_k_cache,
quantize_k_cache_separate,
)
from sglang.srt.layers.attention.nsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.utils import (
@@ -2026,9 +2027,14 @@ class NSATokenToKVPool(MLATokenToKVPool):
assert index_head_dim == 128
if _is_hip:
assert (
self.page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {self.page_size}"
if aiter_can_use_preshuffle_paged_mqa():
assert (
self.page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {self.page_size}"
else:
assert (
self.page_size == 1
), f"HIP legacy NSA path requires page_size == 1, got {self.page_size}"
else:
assert self.page_size == 64
with (
python/sglang/srt/server_args.py
+19 -2
View File
@@ -1842,8 +1842,25 @@ class ServerArgs:
f"attn_tp_size={self.tp_size}, attention weights will be sharded across {self.tp_size} ranks."
)
self.page_size = 64
logger.warning("Setting page size to 64 for DeepSeek DSA.")
# Deferred import to avoid a circular import at module-load
# time (nsa.utils imports get_global_server_args).
from sglang.srt.layers.attention.nsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
)
if is_hip() and not aiter_can_use_preshuffle_paged_mqa():
# Legacy ROCm NSA path: aiter's gluon paged-MQA kernel is
# unavailable (Triton<3.5 and AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS
# not set, or SGLANG_NSA_HIP_DISABLE_PRESHUFFLE=1 / SGLANG_USE_AITER=0).
self.page_size = 1
logger.warning(
"Setting page size to 1 for DeepSeek DSA on ROCm "
"(aiter preshuffle paged-MQA path unavailable: "
"needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1)."
)
else:
self.page_size = 64
logger.warning("Setting page size to 64 for DeepSeek DSA.")
import torch