From 22bfae0d1de4a6301e23b93f690c3a881b95a3a3 Mon Sep 17 00:00:00 2001 From: YC Yen-Ching Tseng Date: Thu, 14 May 2026 15:01:56 +0800 Subject: [PATCH] [AMD] Auto-fallback NSA indexer to page_size=1 when aiter preshuffle gluon kernel is unavailable (Deepseek v3.2) (#25205) --- .../attention/nsa/index_buf_accessor.py | 24 +++++-- .../srt/layers/attention/nsa/nsa_indexer.py | 65 ++++++++++++++----- .../sglang/srt/layers/attention/nsa/utils.py | 41 ++++++++++++ python/sglang/srt/mem_cache/memory_pool.py | 12 +++- python/sglang/srt/server_args.py | 21 +++++- 5 files changed, 137 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index 3bfa1bfad..b5b97147e 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py +++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index d232a8681..f3c2c295d 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 0d2c7ccdb..b49a3e261 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 23d15f71a..9bf36c317 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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 ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 35a5e4618..d0e1b1b91 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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