[AMD] Optimize gpt-oss-120B performance (#27063)

Co-authored-by: wunhuang <wunhuang@amd.com>

python/sglang/srt/environ.py

@@ -407,6 +407,26 @@ class Envs:
     # (matches `gate_mode="separated"`, the layout used by gptoss_fp4 tuned
     # configs and by Mxfp4MoEMethod's post-fix weight shuffle).
     SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)
+    # Fuse the `residual_add + RMSNorm + zero-pad` triplet that appears
+    # before the MoE block for models whose MoE input hidden_size must be
+    # padded up to a stride (e.g. GPT-OSS MXFP4 needs pad to multiple of
+    # 256). When False (default) the pad runs as a separate
+    # torch.nn.functional.pad call inside the MoE method. When True, the
+    # aiter Triton kernel `fused_add_rmsnorm_pad` produces a padded
+    # post-attention layernorm output in one launch and the MoE method
+    # skips the explicit pad. Currently only takes effect on the
+    # post_attention_layernorm path with aiter backend and TP=1.
+    SGLANG_AITER_FUSE_RMSNORM_PAD = EnvBool(False)
+    # Physical layout for MHA KV cache. "nhd" (default) keeps the existing
+    # (size, head_num, head_dim) per-token storage that
+    # `aiter.mha.mha_batch_prefill_func`/`unified_attention` consume directly.
+    # "vectorized_5d" allocates K as (num_blocks, H_kv, head_dim/x, page_size, x)
+    # and V as (num_blocks, H_kv, page_size/x, head_dim, x) (x = 16 / dtype_size),
+    # matching the SHUFFLE layout that aiter's CK FmhaBatchPrefill kernel and
+    # `aiter.ops.triton.gluon.pa_decode_gluon` both consume natively. This is
+    # the SHUFFLE KV layout that enables pa_decode_gluon for full-attn
+    # decode without runtime permutes.
+    SGLANG_AITER_KV_CACHE_LAYOUT = EnvStr("nhd")
     SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)
     SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False)
     SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(4096)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -60,6 +60,10 @@ except ImportError:
     )
 
 from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.layers.attention.aiter_utils import (
+    forward_decode_vectorized_5d,
+    forward_extend_vectorized_5d,
+)
 from sglang.srt.layers.attention.utils import (
     launch_reshape_and_cache_flash,
     pad_sequence_with_mask,
@@ -222,6 +226,21 @@ class AiterAttnBackend(AttentionBackend):
             and model_runner.token_to_kv_pool.swa_layer_nums > 0
         )
 
+        # Detect SHUFFLE 5D ("vectorized") KV cache layout. When active
+        # we (a) skip the launch_reshape_and_cache_flash shortcut and always go
+        # through `set_kv_buffer` (which dispatches to the 5D Triton writer),
+        # and (b) route the decode attention through pa_decode_gluon (see the
+        # corresponding branch in forward_decode), since unified_attention's
+        # 4D `.view(-1, page, H, D)` cannot be applied to a 5D pool.
+        def _pool_is_vec5d(pool):
+            if isinstance(pool, SWAKVPool):
+                return getattr(pool.full_kv_pool, "kv_cache_layout", "nhd") == (
+                    "vectorized_5d"
+                )
+            return getattr(pool, "kv_cache_layout", "nhd") == "vectorized_5d"
+
+        self.kv_cache_is_vectorized_5d = _pool_is_vec5d(model_runner.token_to_kv_pool)
+
         if self.use_sliding_window_kv_pool:
             self.use_triton_unified_attention = True
         else:
@@ -1997,12 +2016,19 @@ class AiterAttnBackend(AttentionBackend):
         if k is not None:
             assert v is not None
             if save_kv_cache:
+                # 5D pool cannot be reshaped to the 4D paged view used by
+                # launch_reshape_and_cache_flash; always route through
+                # set_kv_buffer which dispatches to the SHUFFLE 5D writer.
+                if self.kv_cache_is_vectorized_5d:
+                    self.token_to_kv_pool.set_kv_buffer(
+                        layer, cache_loc, k, v, k_descale, v_descale
+                    )
                 # Only use SWA-specific kv cache write (reshape_and_cache_flash) when
                 # both unified attention and sliding window kv pool are active.
                 # Non-SWA models (e.g. Qwen3-VL) enabled via SGLANG_USE_AITER_UNIFIED_ATTN
                 # use standard set_kv_buffer, as they lack SWA-specific attributes
                 # like full_to_swa_index_mapping.
-                if (
+                elif (
                     self.use_triton_unified_attention
                     and self.use_sliding_window_kv_pool
                 ):
@@ -2372,26 +2398,43 @@ class AiterAttnBackend(AttentionBackend):
                 )
                 return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
-            k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-
             bs0 = forward_batch.batch_size + 1
-
-            # To keep the mha_batch_prefill_func function parameters
-            # declare the necessary parameter and assign None as default value
             q_descale = None
 
+            window_size = (-1, -1)
+            if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
+                window_size = (layer.sliding_window_size, -1)
+
+            if self.kv_cache_is_vectorized_5d:
+                return forward_extend_vectorized_5d(
+                    self,
+                    q,
+                    k,
+                    v,
+                    layer,
+                    forward_batch,
+                    bs0,
+                    window_size,
+                    sinks,
+                )
+
+            # NHD path — original aiter paged batch_prefill.
             # TODO kkhuang-amd need to remove it when mha_batch_prefill_func support fp8-kv
             if self.kv_cache_dtype == fp8_dtype:
                 q = q.to(fp8_dtype)
                 q_descale = layer.k_scale if layer.k_scale is not None else self.k_scale
 
-            window_size = (-1, -1)
+            k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-            page_table = self.forward_metadata.kv_indices
 
-            if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
+            page_table = self.forward_metadata.kv_indices
-                window_size = (layer.sliding_window_size, -1)
+            if (
-                if self.forward_metadata.swa_page_table is not None:
+                layer.sliding_window_size is not None
-                    page_table = self.forward_metadata.swa_page_table
+                and layer.sliding_window_size > -1
+                and self.forward_metadata.swa_page_table is not None
+            ):
+                page_table = self.forward_metadata.swa_page_table
+
+            extra_kwargs = {}
 
             o = mha_batch_prefill_func(
                 q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
@@ -2412,6 +2455,7 @@ class AiterAttnBackend(AttentionBackend):
                 q_descale=q_descale,
                 k_descale=k_descale,
                 v_descale=v_descale,
+                **extra_kwargs,
             )
 
             # The fp8bf16 aiter prefill kernel returns bf16 even when the
@@ -2441,12 +2485,17 @@ class AiterAttnBackend(AttentionBackend):
             v_descale = layer.v_scale if layer.v_scale is not None else self.k_scale
 
         if save_kv_cache:
+            # SHUFFLE 5D pool path — see forward_extend for rationale.
+            if self.kv_cache_is_vectorized_5d:
+                self.token_to_kv_pool.set_kv_buffer(
+                    layer, forward_batch.out_cache_loc, k, v, k_descale, v_descale
+                )
             # Only use SWA-specific kv cache write (reshape_and_cache_flash) when
             # both unified attention and sliding window kv pool are active.
             # Non-SWA models (e.g. Qwen3-VL) enabled via SGLANG_USE_AITER_UNIFIED_ATTN
             # use standard set_kv_buffer, as they lack SWA-specific attributes
             # like full_to_swa_index_mapping.
-            if self.use_triton_unified_attention and self.use_sliding_window_kv_pool:
+            elif self.use_triton_unified_attention and self.use_sliding_window_kv_pool:
                 token_to_kv_pool = self.token_to_kv_pool
                 k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
                 slot_mapping_swa = token_to_kv_pool.full_to_swa_index_mapping
@@ -2535,7 +2584,14 @@ class AiterAttnBackend(AttentionBackend):
             else:
                 o = torch.empty_like(q, dtype=self.input_dtype)
 
-            if self.use_triton_unified_attention:
+            if self.kv_cache_is_vectorized_5d:
+                # SHUFFLE 5D pool: pa_decode_gluon for full + SWA layers
+                # (see :func:`aiter_utils.forward_decode_vectorized_5d`
+                # for the dispatch rationale).
+                forward_decode_vectorized_5d(
+                    self, q, layer, forward_batch, k_cache, v_cache, o, sinks
+                )
+            elif self.use_triton_unified_attention:
                 bs = forward_batch.batch_size
                 window_size = (-1, -1)
                 page_table = self.forward_metadata.kv_indices

python/sglang/srt/layers/attention/aiter_utils.py

@@ -0,0 +1,309 @@
+"""SHUFFLE 5D KV pool helpers for the AITER attention backend.
+
+This module hosts the attention pathways that are specific to the
+``SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d`` (SHUFFLE 5D) physical layout.
+They live here rather than inline in
+:mod:`sglang.srt.layers.attention.aiter_backend` so the main backend
+file keeps focused on the legacy NHD path and on dispatch wiring.
+
+Each entry point takes the :class:`AiterAttnBackend` instance as its
+first argument so it can reach the shared per-step metadata
+(``forward_metadata``, ``qo_indptr``, ``input_dtype``, …) without
+needing to be a method on the class.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+try:
+    # `mha_batch_prefill_func` is re-exported at the aiter top level via
+    # `aiter/__init__.py` (`from .ops.mha import *`). Note: a bare
+    # `from aiter.mha import ...` does NOT work — that module path only
+    # exists as `aiter.ops.mha`.
+    from aiter import mha_batch_prefill_func
+    from aiter.ops.triton.gluon.pa_decode_gluon import (
+        get_recommended_splits,
+        pa_decode_gluon,
+    )
+except ImportError:  # pragma: no cover - import-time guard mirrors aiter_backend
+    mha_batch_prefill_func = None
+    pa_decode_gluon = None
+    get_recommended_splits = None
+
+from sglang.srt.layers.attention.utils import launch_gather_shuffle_5d_to_linear
+from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def forward_extend_vectorized_5d(
+    backend: "AiterAttnBackend",
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    layer: "RadixAttention",
+    forward_batch: "ForwardBatch",
+    bs0: int,
+    window_size,
+    sinks,
+) -> torch.Tensor:
+    """``forward_extend`` specialization for the SHUFFLE 5D KV pool.
+
+    Two sub-paths, both routing through aiter's 3D LINEAR-mode
+    ``mha_batch_prefill_func`` (page_size=1):
+
+    1. Fresh-prompt shortcut: when every request in the batch has zero
+       ``extend_prefix_lens`` (first chunk of a fresh prompt, or any
+       path bypassing prefix reuse) the fresh ``(k, v)`` inputs ARE the
+       full KV stream — skip pool reads entirely and run on bf16
+       ``(k, v)`` directly. No descales needed since no data is read
+       from the (possibly fp8) cache.
+
+    2. Gather-and-linearize: otherwise gather the per-token K/V from the
+       SHUFFLE 5D pool via ``launch_gather_shuffle_5d_to_linear``
+       (triton inverse of the SHUFFLE writer) into a contiguous
+       ``(T, H, D)`` buffer in the cache's ``store_dtype``, then run the
+       same LINEAR prefill. fp8-store layers are forwarded to aiter as
+       raw fp8 with the per-tensor descales — aiter's LINEAR-mode kernel
+       supports fp8 K/V/Q natively, so no host-side dequant is needed.
+
+    The fallback exists because aiter's paged ``mha_batch_prefill_func``
+    lacks a compiled kernel for our
+    ``(page_size=64, bf16/fp8, SHUFFLE 5D)`` configuration; calling it
+    from the 5D pool aborts with ``"no matching kernel found"``.
+
+    Returns the ``(T, H_q * D_v)`` attention output, ready to be
+    returned from ``AiterAttnBackend.forward_extend``.
+    """
+    # Path 1: fresh-prompt shortcut.
+    extend_no_prefix = forward_batch.extend_prefix_lens_cpu is not None and not any(
+        forward_batch.extend_prefix_lens_cpu
+    )
+    if extend_no_prefix:
+        k_lin = k.contiguous().view(-1, layer.tp_k_head_num, layer.qk_head_dim)
+        v_lin = v.contiguous().view(-1, layer.tp_v_head_num, layer.v_head_dim)
+        total_tokens = k_lin.shape[0]
+        kv_indices_lin = torch.arange(
+            total_tokens, dtype=torch.int32, device=k_lin.device
+        )
+        kv_indptr_lin = backend.qo_indptr[:bs0]
+        max_q = int(backend.forward_metadata.max_q_len)
+        o = mha_batch_prefill_func(
+            q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+            k_lin,
+            v_lin,
+            backend.qo_indptr[:bs0],
+            kv_indptr_lin,
+            kv_indices_lin,
+            max_q,
+            max_q,
+            causal=True,
+            logits_soft_cap=backend.logits_soft_cap,
+            alibi_slopes=None,
+            return_lse=False,
+            return_attn_probs=False,
+            window_size=window_size,
+            sink_ptr=sinks,
+        )
+        if o.dtype != backend.input_dtype:
+            o = o.to(backend.input_dtype)
+        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+    # Path 2: gather-and-linearize.
+    # SWA layers gather from the SWA sub-pool via swa_page_table;
+    # full-attn layers gather from the full sub-pool via kv_indices.
+    # Both are per-TOKEN slot id lists populated by
+    # ``create_flashinfer_kv_indices_triton`` from ``req_to_token`` (one
+    # slot id per logical token), so the first ``seq_lens_sum`` entries
+    # of either tensor are exactly the per-token absolute pool slot ids
+    # in request-major order — no per-token gather metadata to build on
+    # host.
+    is_swa_layer = (
+        layer.sliding_window_size is not None
+        and layer.sliding_window_size > -1
+        and backend.forward_metadata.swa_page_table is not None
+    )
+    total_kv = int(forward_batch.seq_lens_sum)
+    if is_swa_layer:
+        slot_ids = backend.forward_metadata.swa_page_table[:total_kv]
+    else:
+        slot_ids = backend.forward_metadata.kv_indices[:total_kv]
+
+    # Resolve the raw 5D K/V buffer for this layer (going through the
+    # SWA→sub-pool mapping when applicable).
+    pool = backend.token_to_kv_pool
+    if hasattr(pool, "layers_mapping"):
+        sub_layer_id, sub_is_swa = pool.layers_mapping[layer.layer_id]
+        sub_pool = pool.swa_kv_pool if sub_is_swa else pool.full_kv_pool
+    else:
+        sub_pool = pool
+        sub_layer_id = layer.layer_id
+    k_buf = sub_pool.k_buffer[sub_layer_id - sub_pool.start_layer]
+    v_buf = sub_pool.v_buffer[sub_layer_id - sub_pool.start_layer]
+
+    k_lin, v_lin = launch_gather_shuffle_5d_to_linear(k_buf, v_buf, slot_ids)
+    # k_lin / v_lin come out in ``store_dtype`` (uint8 for fp8 pools
+    # because ``Tensor.index_put`` isn't implemented for fp8 — see
+    # ``MHATokenToKVPool`` ctor). Reinterpret them back to the compute
+    # dtype so aiter sees matching q/k/v dtypes. The bytes are
+    # identical; this is a zero-copy view.
+    if sub_pool.store_dtype != sub_pool.dtype:
+        k_lin = k_lin.view(sub_pool.dtype)
+        v_lin = v_lin.view(sub_pool.dtype)
+
+    # For fp8 K/V we hand the raw fp8 tensors and the layer's per-tensor
+    # descales straight to aiter.
+    if sub_pool.dtype == fp8_dtype:
+        q_local = q.to(fp8_dtype)
+        q_descale_local = (
+            layer.k_scale if layer.k_scale is not None else backend.k_scale
+        )
+        k_descale_local = (
+            layer.k_scale if layer.k_scale is not None else backend.k_scale
+        )
+        v_descale_local = (
+            layer.v_scale if layer.v_scale is not None else backend.v_scale
+        )
+    else:
+        q_local = q
+        q_descale_local = None
+        k_descale_local = None
+        v_descale_local = None
+
+    kv_indptr_lin = backend.forward_metadata.kv_indptr[:bs0]
+    kv_indices_lin = torch.arange(total_kv, dtype=torch.int32, device=k_lin.device)
+    max_kv = int(backend.forward_metadata.max_kv_len)
+    max_q = int(backend.forward_metadata.max_q_len)
+
+    o = mha_batch_prefill_func(
+        q_local.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+        k_lin,
+        v_lin,
+        backend.qo_indptr[:bs0],
+        kv_indptr_lin,
+        kv_indices_lin,
+        max_q,
+        max_kv,
+        causal=True,
+        logits_soft_cap=backend.logits_soft_cap,
+        alibi_slopes=None,
+        return_lse=False,
+        return_attn_probs=False,
+        window_size=window_size,
+        sink_ptr=sinks,
+        q_descale=q_descale_local,
+        k_descale=k_descale_local,
+        v_descale=v_descale_local,
+    )
+    if o.dtype != backend.input_dtype:
+        o = o.to(backend.input_dtype)
+    return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+
+
+def forward_decode_vectorized_5d(
+    backend: "AiterAttnBackend",
+    q: torch.Tensor,
+    layer: "RadixAttention",
+    forward_batch: "ForwardBatch",
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    o: torch.Tensor,
+    sinks,
+) -> None:
+    """``forward_decode`` specialization for the SHUFFLE 5D KV pool.
+
+    Runs ``pa_decode_gluon`` for both full-attention and sliding-window
+    layers — when SHUFFLE 5D is active the SWA sub-pool is also
+    allocated 5D (see ``SWAKVPool`` ctor), so we keep one decode kernel
+    instead of falling back to ``unified_attention`` for SWA layers.
+
+    The choice between the two layer kinds is purely metadata:
+
+    * Full-attn  → ``kv_indices`` page table + ``sliding_window=0`` +
+      ``max_part_num`` recommended by aiter heuristics.
+    * SWA layer  → ``swa_page_table`` + ``sliding_window=layer.sliding_window_size``
+      + ``max_part_num=1`` (SWA windows are small enough that
+      splitting does not help).
+
+    fp8 KV requires per-tensor ``key_scale`` / ``value_scale`` to be
+    forwarded; without them the kernel reads the fp8 bytes as fp8
+    values without any dequant and produces garbage logits.
+
+    Writes the attention output into ``o`` in place (via a stride-0
+    safe ``o.view``).
+    """
+    bs = forward_batch.batch_size
+    num_kv_heads = layer.tp_k_head_num
+    num_q_heads = layer.tp_q_head_num
+    q_group = num_q_heads // num_kv_heads
+    is_swa_layer = (
+        layer.sliding_window_size is not None and layer.sliding_window_size > -1
+    )
+
+    if is_swa_layer:
+        block_tables_pa = (
+            backend.forward_metadata.swa_page_table
+            if backend.forward_metadata.swa_page_table is not None
+            else backend.forward_metadata.kv_indices
+        )
+        ctx_part = 256
+        max_part_num = 1
+        sliding_window_arg = int(layer.sliding_window_size)
+    else:
+        block_tables_pa = backend.forward_metadata.kv_indices
+        ctx_part = 256
+        max_part_num = get_recommended_splits(bs, num_kv_heads)
+        sliding_window_arg = 0
+
+    q_in = q.view(-1, num_q_heads, layer.qk_head_dim)
+    # Direct view of o as kernel output — saves a per-layer o.copy_ of
+    # bs * H_q * D bf16 elementwise.
+    o_view = o.view(-1, num_q_heads, layer.v_head_dim)
+    exp_sums = torch.empty(
+        (bs, num_kv_heads, max_part_num, q_group),
+        dtype=torch.float32,
+        device=q_in.device,
+    )
+    max_logits = torch.empty_like(exp_sums)
+    temporary_output = torch.empty(
+        (bs, num_kv_heads, max_part_num, q_group, layer.qk_head_dim),
+        dtype=q_in.dtype,
+        device=q_in.device,
+    )
+
+    # For fp8 KV cache the kernel needs per-tensor dequant scales
+    # (key_scale / value_scale). Without them the fp8 bytes are
+    # interpreted as fp8 values with no dequant.
+    key_scale = None
+    value_scale = None
+    if backend.kv_cache_dtype == fp8_dtype:
+        key_scale = layer.k_scale if layer.k_scale is not None else backend.k_scale
+        value_scale = layer.v_scale if layer.v_scale is not None else backend.v_scale
+
+    pa_decode_gluon(
+        output=o_view,
+        query=q_in,
+        key_cache=k_cache,
+        value_cache=v_cache,
+        context_lengths=forward_batch.seq_lens,
+        block_tables=block_tables_pa,
+        softmax_scale=layer.scaling,
+        query_length=1,
+        max_context_partition_num=max_part_num,
+        context_partition_size=ctx_part,
+        compute_type=backend.input_dtype,
+        key_scale=key_scale,
+        value_scale=value_scale,
+        exp_sums=exp_sums,
+        max_logits=max_logits,
+        temporary_output=temporary_output,
+        sinks=sinks,
+        sliding_window=sliding_window_arg,
+        ps=True,
+    )

python/sglang/srt/layers/layernorm.py

@@ -208,6 +208,7 @@ class RMSNorm(MultiPlatformOp):
         has_weight: bool = True,
         weight_dtype: Optional = None,
         override_orig_dtype: Optional = None,
+        x_pad_to_multiple: int = 0,
     ) -> None:
         super().__init__()
         self.has_weight = has_weight
@@ -223,7 +224,25 @@ class RMSNorm(MultiPlatformOp):
         self.variance_size_override = (
             None if var_hidden_size == hidden_size else var_hidden_size
         )
+        # When > 0, fuse a zero-pad of the last dim out to a multiple of
+        # this value into the rmsnorm kernel via aiter's
+        # `fused_add_rmsnorm_pad` Triton kernel. The padded output has
+        # shape (M, ceil(N/x_pad_to_multiple)*x_pad_to_multiple); the
+        # residual_out stays at the original (M, N) shape.
+
         if _use_aiter:
+            self.x_pad_to_multiple = x_pad_to_multiple
+            self._fused_pad_kernel = None
+
+            if x_pad_to_multiple > 0:
+                try:
+                    from aiter.ops.triton.fused_add_rmsnorm_pad import (
+                        fused_add_rmsnorm_pad as _fused_add_rmsnorm_pad,
+                    )
+
+                    self._fused_pad_kernel = _fused_add_rmsnorm_pad
+                except ImportError:
+                    self._fused_pad_kernel = None
             self._forward_method = self.forward_aiter
 
     def forward_cuda(
@@ -345,6 +364,22 @@ class RMSNorm(MultiPlatformOp):
             x = x.contiguous().reshape(-1, original_shape[-1])
         elif not x.is_contiguous():
             x = x.contiguous()
+        # Fused (add +) rmsnorm + zero-pad path. Triggered when caller
+        # constructed RMSNorm with x_pad_to_multiple > 0. Output last
+        # dim is padded up; residual_out stays at original width. Used
+        # by callers (e.g. GPT-OSS MXFP4 MoE) whose immediate consumer
+        # needs a padded hidden_size — folding the pad in here removes a
+        # separate launch.
+        if self._fused_pad_kernel is not None and self.x_pad_to_multiple > 0:
+            if post_residual_addition is not None and residual is not None:
+                residual = residual + post_residual_addition
+            return self._fused_pad_kernel(
+                x,
+                self.weight.data,
+                self.variance_epsilon,
+                residual,
+                self.x_pad_to_multiple,
+            )
         if residual is not None:
             residual_out = torch.empty_like(x)
             output = torch.empty_like(x)

python/sglang/srt/layers/quantization/mxfp4.py

@@ -1267,9 +1267,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 w13_weight.is_shuffled = True
                 w2_weight.is_shuffled = True
 
-            x_padded = torch.nn.functional.pad(
+            # Skip the explicit pad if x already arrives at the padded
-                x, (0, self.hidden_pad), mode="constant", value=0.0
+            # hidden_size (the upstream RMSNorm fused the pad into its
-            )
+            # output — see RMSNorm.x_pad_to_multiple). Saves a separate
+            # zero-pad kernel launch per layer.
+            if x.shape[-1] == self.hidden_size:
+                x_padded = x
+            else:
+                x_padded = torch.nn.functional.pad(
+                    x, (0, self.hidden_pad), mode="constant", value=0.0
+                )
             quant_info = AiterMoeQuantInfo(
                 w13_weight=w13_weight,
                 w2_weight=w2_weight,

python/sglang/srt/layers/rotary_embedding/base.py

@@ -364,15 +364,20 @@ class RotaryEmbedding(MultiPlatformOp):
 
             if fused_set_kv_buffer_arg is not None and _is_hip:
                 extra_args = fused_set_kv_buffer_arg
-
+                k_cache = fused_set_kv_buffer_arg["key_cache"]
-                k_cache_shape = fused_set_kv_buffer_arg["key_cache"].shape
+                # 5D SHUFFLE pool feeds raw (N, H, D/x, page, x) K cache;
-                qk_head_dim = k_cache_shape[-1]
+                # NHD 3D pool feeds the legacy 4D paged view. Auto-detect.
-                tp_k_head_num = k_cache_shape[-2]
+                is_shuffle_5d = k_cache.ndim == 5
+                if is_shuffle_5d:
+                    # K shape (num_blocks, H_kv, D//x, page, x): D = D//x * x
+                    qk_head_dim = k_cache.shape[2] * k_cache.shape[4]
+                    tp_k_head_num = k_cache.shape[1]
+                else:
+                    qk_head_dim = k_cache.shape[-1]
+                    tp_k_head_num = k_cache.shape[-2]
 
                 key = key.view(-1, tp_k_head_num, qk_head_dim)
-
                 tokens = key.shape[0]
-
                 query = query.view(tokens, -1, qk_head_dim)
 
                 query, key, k_cache, v_cache = fused_qk_rope_reshape_and_cache(
@@ -381,7 +386,7 @@ class RotaryEmbedding(MultiPlatformOp):
                     pos=positions,
                     cos_sin=self.cos_sin_cache,
                     is_neox=self.is_neox_style,
-                    flash_layout=True,
+                    flash_layout=not is_shuffle_5d,
                     offs=None,
                     q_out=query,
                     k_out=key,

python/sglang/srt/mem_cache/memory_pool.py

@@ -81,6 +81,11 @@ _is_cpu = is_cpu()
 _cpu_has_amx_support = cpu_has_amx_support()
 _is_hip = is_hip()
 _is_fp8_fnuz = is_fp8_fnuz()
+# `SGLANG_AITER_KV_CACHE_LAYOUT` is only meaningful on the ROCm AITER backend
+# (HIP + --enable-aiter / SGLANG_USE_AITER=1). On any other platform / backend
+# the SHUFFLE 5D pool layout has no consumer kernels, so the env var is
+# silently ignored and the legacy NHD layout is used.
+_use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip
 
 
 def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]):
@@ -892,6 +897,40 @@ class MHATokenToKVPool(KVCache):
             else v_head_dim if v_head_dim is not None else head_dim
         )
 
+        # Optional SHUFFLE 5D ("vectorized") physical layout for K/V.
+        # Selected by `SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d` on the ROCm
+        # AITER backend (HIP + SGLANG_USE_AITER=1). When active:
+        #   K shape: (num_blocks, H, D_k // X, page, X)
+        #   V shape: (num_blocks, H, page // X, D_v, X)   where X = 16 / dtype_bytes
+        # aiter `mha_batch_prefill_func` consumes these 5D shapes natively and
+        # aiter `pa_decode_gluon` reads SHUFFLE blocks directly during decode.
+        # An explicit `kv_cache_layout=` argument always wins (e.g. SWAKVPool
+        # passes "nhd" to keep its SWA sub-pool on the legacy layout); on
+        # non-AITER platforms the env var is ignored and NHD is forced since
+        # no consumer kernel exists for SHUFFLE 5D outside the AITER backend.
+        self.kv_cache_layout = "nhd"
+        if _use_aiter:
+            layout = envs.SGLANG_AITER_KV_CACHE_LAYOUT.get().lower()
+            if layout not in ("nhd", "vectorized_5d"):
+                raise ValueError(
+                    f"Unsupported SGLANG_AITER_KV_CACHE_LAYOUT={layout!r}; "
+                    "expected 'nhd' or 'vectorized_5d'."
+                )
+            self.kv_cache_layout = layout
+            if layout == "vectorized_5d":
+                # X is the inner vectorization width in the SHUFFLE layout,
+                # determined by the STORAGE dtype (not the compute dtype) since
+                # it controls how many elements fit in 16 bytes of the on-pool
+                # tensor. For fp8 storage X=16, for bf16/fp16 X=8.
+                self._kv_vector_x = 16 // self.store_dtype.itemsize
+                assert (self.size + self.page_size) % self.page_size == 0
+                assert self.page_size % self._kv_vector_x == 0, (
+                    f"page_size={self.page_size} must be divisible by "
+                    f"X={self._kv_vector_x} for vectorized_5d layout"
+                )
+                assert self.head_dim % self._kv_vector_x == 0
+                assert self.v_head_dim % self._kv_vector_x == 0
+
         self._create_buffers()
 
         self.device_module = torch.get_device_module(self.device)
@@ -973,24 +1012,63 @@ class MHATokenToKVPool(KVCache):
                 if self.enable_custom_mem_pool
                 else nullcontext()
             ):
-                # [size, head_num, head_dim] for each layer
+                if self.kv_cache_layout == "vectorized_5d":
-                # The padded slot 0 is used for writing dummy outputs from padded tokens.
+                    total_slots = self.size + self.page_size
-                self.k_buffer = [
+                    num_blocks = total_slots // self.page_size
-                    torch.zeros(
+                    x = self._kv_vector_x
-                        (self.size + self.page_size, self.head_num, self.head_dim),
+                    # K: (num_blocks, H, D_k // X, page, X)
-                        dtype=self.store_dtype,
+                    self.k_buffer = [
-                        device=self.device,
+                        torch.zeros(
-                    )
+                            (
-                    for _ in range(self.layer_num)
+                                num_blocks,
-                ]
+                                self.head_num,
-                self.v_buffer = [
+                                self.head_dim // x,
-                    torch.zeros(
+                                self.page_size,
-                        (self.size + self.page_size, self.head_num, self.v_head_dim),
+                                x,
-                        dtype=self.store_dtype,
+                            ),
-                        device=self.device,
+                            dtype=self.store_dtype,
-                    )
+                            device=self.device,
-                    for _ in range(self.layer_num)
+                        )
-                ]
+                        for _ in range(self.layer_num)
+                    ]
+                    # V: (num_blocks, H, page // X, D_v, X)
+                    self.v_buffer = [
+                        torch.zeros(
+                            (
+                                num_blocks,
+                                self.head_num,
+                                self.page_size // x,
+                                self.v_head_dim,
+                                x,
+                            ),
+                            dtype=self.store_dtype,
+                            device=self.device,
+                        )
+                        for _ in range(self.layer_num)
+                    ]
+                else:
+                    # [size, head_num, head_dim] for each layer
+                    # The padded slot 0 is used for writing dummy outputs from padded tokens.
+                    self.k_buffer = [
+                        torch.zeros(
+                            (self.size + self.page_size, self.head_num, self.head_dim),
+                            dtype=self.store_dtype,
+                            device=self.device,
+                        )
+                        for _ in range(self.layer_num)
+                    ]
+                    self.v_buffer = [
+                        torch.zeros(
+                            (
+                                self.size + self.page_size,
+                                self.head_num,
+                                self.v_head_dim,
+                            ),
+                            dtype=self.store_dtype,
+                            device=self.device,
+                        )
+                        for _ in range(self.layer_num)
+                    ]
 
         self.k_data_ptrs = torch.tensor(
             [x.data_ptr() for x in self.k_buffer],
@@ -1145,6 +1223,28 @@ class MHATokenToKVPool(KVCache):
             cache_k = cache_k.view(self.store_dtype)
             cache_v = cache_v.view(self.store_dtype)
 
+        if self.kv_cache_layout == "vectorized_5d":
+            # Late-import to keep the NHD path import-clean.
+            from sglang.srt.layers.attention.utils import (
+                launch_reshape_and_cache_shuffle_5d,
+            )
+
+            # The writer kernel uses key.stride(0) directly as the source
+            # token stride; head/dim are assumed contiguous within each
+            # token (stride(1)=head_size, stride(2)=1). Both hold for K/V
+            # produced by QKV split + RoPE in upstream attention even when
+            # the outer per-token stride is non-canonical, so we skip the
+            # protective .contiguous() copies that would otherwise fire
+            # large per-layer elementwise kernels.
+            launch_reshape_and_cache_shuffle_5d(
+                cache_k,
+                cache_v,
+                self.k_buffer[layer_id - self.start_layer],
+                self.v_buffer[layer_id - self.start_layer],
+                loc,
+            )
+            return
+
         _set_kv_buffer_impl(
             cache_k,
             cache_v,

python/sglang/srt/models/gpt_oss.py

@@ -84,6 +84,7 @@ from sglang.srt.utils import (
     is_cpu,
     is_cuda,
     is_flashinfer_available,
+    is_hip,
     is_npu,
     is_sm90_supported,
     make_layers,
@@ -92,6 +93,7 @@ from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cpu = is_cpu()
 _is_npu = is_npu()
+_is_hip = is_hip()
 _is_cuda = is_cuda()
 _is_tinygemm_supported = (
     _is_cuda
@@ -165,6 +167,36 @@ class TinyGemmLinear(ReplicatedLinear):
         return super().forward(x)
 
 
+def _resolve_moe_input_pad_multiple(
+    quant_config: Optional[QuantizationConfig],
+) -> int:
+    """Return the alignment the MoE backend requires on its input
+    hidden_size, or 0 when no fused pad should be inserted into the
+    preceding layernorm. See post_attention_layernorm construction in
+    GptOssDecoderLayer for the safety preconditions."""
+    if quant_config is None:
+        return 0
+    from sglang.srt.environ import envs
+
+    if not envs.SGLANG_AITER_FUSE_RMSNORM_PAD.get():
+        return 0
+    if not (_is_hip and envs.SGLANG_USE_AITER.get()):
+        return 0
+    # Only the MXFP4 path needs the 256-multiple pad on hidden_size; other
+    # quant methods (or unquantized bf16) consume the unpadded layernorm
+    # output directly.
+    if quant_config.get_name() != "mxfp4":
+        return 0
+    if get_tensor_model_parallel_world_size() != 1:
+        # Mid-layer hidden_states still flow through CommunicateWith...
+        # AllReduceAndLayerNormFn helpers other than `_simple` when
+        # attn_tp_size > 1; those helpers haven't been updated to handle
+        # a padded layernorm output. Keep the optimisation off to stay
+        # correct.
+        return 0
+    return 256
+
+
 class GptOssSparseMoeBlock(nn.Module):
     def __init__(
         self,
@@ -250,18 +282,44 @@ class GptOssSparseMoeBlock(nn.Module):
         hidden_states: torch.Tensor,
         should_allreduce_fusion: bool = False,
     ) -> torch.Tensor:
-        num_tokens, hidden_dim = hidden_states.shape
+        # `hidden_states` may arrive pre-padded along the last dim when the
+        # preceding RMSNorm fused the MoE input pad (gated by
+        # SGLANG_AITER_FUSE_RMSNORM_PAD). Router/topk are computed on the
+        # unpadded slice so the small bf16 router GEMM dimensions stay
+        # untouched, while the experts call gets to keep the padded view
+        # and skip the duplicate pad inside the MXFP4 method. The output
+        # is then trimmed back to the unpadded width so postprocess_layer
+        # can pair it with the (M, hidden_dim_unpadded) residual.
+        num_tokens = hidden_states.shape[0]
+        hidden_dim_unpadded = self.experts.hidden_size
+        is_prepadded = hidden_states.shape[-1] != hidden_dim_unpadded
+        if is_prepadded:
+            router_input = hidden_states[..., :hidden_dim_unpadded]
+        else:
+            router_input = hidden_states
+
         if is_in_piecewise_cuda_graph():
             final_hidden_states = moe_impl(self.layer_id, hidden_states)
         else:
-            router_logits, _ = self.router(hidden_states)
+            router_logits, _ = self.router(router_input)
-            topk_output = self.topk(hidden_states, router_logits)
+            topk_output = self.topk(router_input, router_logits)
             final_hidden_states = self.experts(hidden_states, topk_output)
 
         if self.tp_size > 1 and not should_allreduce_fusion:
             final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
 
-        ans = final_hidden_states.view(num_tokens, hidden_dim)
+        # When input was pre-padded, FusedMoE.forward_impl captured the
+        # padded width as `origin_hidden_states_dim` and skipped its own
+        # output-trim contiguous() — so the experts output is still
+        # (M, hidden_dim_padded). Drop the pad columns here. When input
+        # was unpadded (default code path), FusedMoE.forward_impl already
+        # produced a contiguous (M, hidden_dim_unpadded) tensor, so the
+        # view is a no-op and matches the pre-fusion behavior bit-for-bit.
+        if is_prepadded:
+            ans = final_hidden_states[..., :hidden_dim_unpadded].contiguous()
+            ans = ans.view(num_tokens, hidden_dim_unpadded)
+        else:
+            ans = final_hidden_states.view(num_tokens, hidden_dim_unpadded)
         return ans
 
 
@@ -505,8 +563,19 @@ class GptOssDecoderLayer(nn.Module):
                 "Please use GptOssSparseMoeBlock instead."
             )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Optionally fuse the MoE-input zero-pad into post_attention_layernorm
+        # via aiter's `fused_add_rmsnorm_pad`. Only enabled when:
+        #   * SGLANG_AITER_FUSE_RMSNORM_PAD=1
+        #   * Quant method is MXFP4 (the only path that demands a 256-pad)
+        #   * Communication path between layernorm and MoE is the no-op
+        #     `_simple` route (attn_tp_size == 1) — otherwise the padded
+        #     hidden_states would have to survive an AllReduce/scatter that
+        #     hasn't been taught about the extra columns yet.
+        post_attn_pad_multiple = _resolve_moe_input_pad_multiple(quant_config)
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            x_pad_to_multiple=post_attn_pad_multiple,
         )
 
         self.layer_communicator = LayerCommunicator(

python/sglang/srt/models/utils.py

@@ -275,7 +275,15 @@ class AutoWeightsLoader:
 
 
 def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
-    """Enable fused set_kv_buffer only on CUDA with bfloat16 KV cache."""
+    """Enable fused set_kv_buffer on CUDA with bfloat16 KV cache and HIP with bf16/fp16/fp8 KV cache.
+
+    SHUFFLE 5D pools on HIP also work — the underlying triton kernel
+    (`fused_qk_rope_reshape_and_cache`) natively supports the 5D
+    SHUFFLE layout (key_cache.ndim==5, value_cache.ndim==5). We just need
+    the per-layer arg builder to pass the raw 5D buffers without the
+    `.view(-> 4D NHD)` reshape, and let the rotary forward pass
+    `flash_layout=False`. See `create_fused_set_kv_buffer_arg` below.
+    """
     pool = get_token_to_kv_pool()
     return (
         _is_cuda
@@ -313,16 +321,26 @@ def create_fused_set_kv_buffer_arg(
             if layer.sliding_window_size > 0
             else None
         )
+        # SHUFFLE 5D pools (k_buffer.ndim == 5) consumed natively by
+        # fused_qk_rope_reshape_and_cache via flash_layout=False. For the
+        # legacy NHD 3D pool we reshape to the (num_blocks, page_size, H, D)
+        # paged view the kernel expects under flash_layout=True.
+        if k_buffer.ndim == 5:
+            key_cache = k_buffer
+            value_cache = v_buffer
+        else:
+            key_cache = k_buffer.view(
+                -1, page_size, layer.tp_k_head_num, layer.qk_head_dim
+            )
+            value_cache = v_buffer.view(
+                -1, page_size, layer.tp_v_head_num, layer.v_head_dim
+            )
         return {
             "v": value.view(-1, layer.tp_v_head_num, layer.v_head_dim),
             "k_scale": layer.k_scale,
             "v_scale": layer.v_scale,
-            "key_cache": k_buffer.view(
+            "key_cache": key_cache,
-                -1, page_size, layer.tp_k_head_num, layer.qk_head_dim
+            "value_cache": value_cache,
-            ),
-            "value_cache": v_buffer.view(
-                -1, page_size, layer.tp_v_head_num, layer.v_head_dim
-            ),
             "slot_mapping": forward_batch.out_cache_loc,
             "swa_slot_mapping": slot_mapping_swa,
         }

python/sglang/srt/server_args.py

@@ -3139,7 +3139,23 @@ class ServerArgs:
 
     def _handle_page_size(self):
         if self.page_size is None:
-            if not is_musa():
+            # SHUFFLE 5D vectorized KV layout (aiter backend + pa_decode_gluon)
+            # is tuned for and prefers page_size=64 — making it the default
+            # when the layout flag is set avoids users having to pass
+            # --page-size 64 explicitly. The env var is only consumed by the
+            # ROCm AITER backend, so the auto-bump is gated on HIP; on other
+            # platforms the SHUFFLE 5D pool has no consumer kernels and the
+            # env var is silently ignored (see MHATokenToKVPool).
+            if (
+                is_hip()
+                and envs.SGLANG_AITER_KV_CACHE_LAYOUT.get().lower() == "vectorized_5d"
+            ):
+                self.page_size = 64
+                logger.info(
+                    "Setting page_size=64 as default for "
+                    "SGLANG_AITER_KV_CACHE_LAYOUT=vectorized_5d."
+                )
+            elif not is_musa():
                 self.page_size = 1
             else:
                 self.page_size = 64