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

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:
+        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

@@ -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,12 +445,20 @@ 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:
+            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
+        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
@@ -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:
+            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,6 +605,9 @@ class Indexer(MultiPlatformOp):
         )
         weights = weights.squeeze(-1)
 
+        if _is_hip and not _use_aiter_preshuffle:
+            block_tables = metadata.get_page_table_1()
+        else:
             block_tables = metadata.get_page_table_64()
 
         assert (
@@ -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

@@ -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

@@ -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:
+            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

@@ -1842,6 +1842,23 @@ class ServerArgs:
                                 f"attn_tp_size={self.tp_size}, attention weights will be sharded across {self.tp_size} ranks."
                             )
 
+                    # 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.")