From 1c73ff8ad3fd365420de8088594be15ac61ffb53 Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Mon, 8 Jun 2026 11:06:30 +0800 Subject: [PATCH] [AMD] Optimize gpt-oss-120B performance (#27063) Co-authored-by: wunhuang --- python/sglang/srt/environ.py | 20 + .../srt/layers/attention/aiter_backend.py | 84 +- .../srt/layers/attention/aiter_utils.py | 309 +++++ python/sglang/srt/layers/attention/utils.py | 1185 +++++++++++++++++ python/sglang/srt/layers/layernorm.py | 35 + .../sglang/srt/layers/quantization/mxfp4.py | 13 +- .../srt/layers/rotary_embedding/base.py | 19 +- python/sglang/srt/mem_cache/memory_pool.py | 136 +- python/sglang/srt/models/gpt_oss.py | 79 +- python/sglang/srt/models/utils.py | 32 +- python/sglang/srt/server_args.py | 18 +- 11 files changed, 1875 insertions(+), 55 deletions(-) create mode 100644 python/sglang/srt/layers/attention/aiter_utils.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7239ccf23..4c6684c74 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index ab92eb224..3f7d40f9b 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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) - page_table = self.forward_metadata.kv_indices + k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) - if layer.sliding_window_size is not None and layer.sliding_window_size > -1: - window_size = (layer.sliding_window_size, -1) - if self.forward_metadata.swa_page_table is not None: - page_table = self.forward_metadata.swa_page_table + page_table = self.forward_metadata.kv_indices + if ( + layer.sliding_window_size is not None + 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 diff --git a/python/sglang/srt/layers/attention/aiter_utils.py b/python/sglang/srt/layers/attention/aiter_utils.py new file mode 100644 index 000000000..844ccff14 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index 489f5c325..d31004baa 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -1,4 +1,6 @@ import torch +import triton +import triton.language as tl from sglang.jit_kernel.utils import is_arch_support_pdl from sglang.srt.layers.attention.triton_ops.cache_ops import ( @@ -176,6 +178,1189 @@ def concat_mla_absorb_q_general(q_nope, q_rope): return torch.cat([q_nope, q_rope], dim=-1) +@triton.jit +def reshape_and_cache_flash( + key_ptr, + value_ptr, + key_cache_ptr, + value_cache_ptr, + slot_mapping_ptr, + swa_slot_mapping_ptr, + k_scale_ptr, + v_scale_ptr, + block_stride, + key_stride, + value_stride, + num_heads, + head_size, + block_size, + HEAD_BLOCK: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_SWA: tl.constexpr, + USE_SCALE: tl.constexpr, +): + """ + Triton kernel for reshaping per-token K/V tensors into paged KV cache layout. + + Source layout: + key/value: [num_tokens, num_heads, head_size] + + Target cache layout: + cache: [num_blocks, block_size, num_heads, head_size] + + Each Triton program instance handles: + - one token (program_id(0)) + - one block of heads (program_id(1)) + + Features: + - optional SWA slot remapping + - optional FP8 scale dequantization before cache write + + Args: + key_ptr: Pointer to source key tensor. + value_ptr: Pointer to source value tensor. + key_cache_ptr: Pointer to destination key cache tensor. + value_cache_ptr: Pointer to destination value cache tensor. + slot_mapping_ptr: Maps token -> cache slot. + swa_slot_mapping_ptr: Optional second-stage slot remap for SWA mode. + k_scale_ptr: Optional key scaling factor pointer. + v_scale_ptr: Optional value scaling factor pointer. + block_stride: Stride between cache blocks. + key_stride: Stride between source key tokens. + value_stride: Stride between source value tokens. + num_heads: Number of attention heads. + head_size: Hidden dimension per head. + block_size: Number of slots per cache block. + HEAD_BLOCK: Number of heads processed per program. + BLOCK_D: Vectorized dimension size (power-of-2 padded). + HAS_SWA: Enable SWA remapping. + USE_SCALE: Enable scale division before storing. + """ + + # ---------------------------------- + # program ids + # pid0 = token + # pid1 = head block + # ---------------------------------- + token_idx = tl.program_id(0) + head_block_idx = tl.program_id(1) + + # ---------------------------------- + # slot mapping + # ---------------------------------- + slot_idx = tl.load(slot_mapping_ptr + token_idx) + + if HAS_SWA: + slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx) + + if slot_idx < 0: + return + + block_idx = slot_idx // block_size + block_offset = slot_idx % block_size + + # ---------------------------------- + # head range + # ---------------------------------- + head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK) + + head_mask = head_idx < num_heads + + dim_idx = tl.arange(0, BLOCK_D) + + # shape = [HEAD_BLOCK, BLOCK_D] + offs = head_idx[:, None] * head_size + dim_idx[None, :] + + mask = head_mask[:, None] & (dim_idx[None, :] < head_size) + + # ---------------------------------- + # source load + # ---------------------------------- + src_key = token_idx * key_stride + offs + src_value = token_idx * value_stride + offs + + k = tl.load(key_ptr + src_key, mask=mask) + v = tl.load(value_ptr + src_value, mask=mask) + + # ---------------------------------- + # optional scale + # ---------------------------------- + if USE_SCALE: + k_scale = tl.load(k_scale_ptr) + v_scale = tl.load(v_scale_ptr) + + k = k / k_scale + v = v / v_scale + + # ---------------------------------- + # target layout + # [block_idx, block_offset, head, dim] + # ---------------------------------- + tgt = block_idx * block_stride + block_offset * num_heads * head_size + offs + + tl.store(key_cache_ptr + tgt, k, mask=mask) + tl.store(value_cache_ptr + tgt, v, mask=mask) + + +def launch_reshape_and_cache_flash( + key, + value, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping=None, + k_scale=None, + v_scale=None, +): + """ + Launch wrapper for reshape_and_cache_flash Triton kernel. + + This wrapper prepares launch configuration and dispatches the Triton kernel + that writes token-major K/V tensors into paged KV cache layout. + + Args: + key: Source key tensor [num_tokens, num_heads, head_size] + value: Source value tensor [num_tokens, num_heads, head_size] + key_cache: Destination key cache [num_blocks, block_size, num_heads, head_size] + value_cache: Destination value cache [num_blocks, block_size, num_heads, head_size] + slot_mapping: Token-to-cache slot mapping + swa_slot_mapping: Optional SWA remapping table + k_scale: Optional key scaling factor + v_scale: Optional value scaling factor + """ + + num_tokens = key.shape[0] + num_heads = key.shape[1] + head_size = key.shape[2] + + HEAD_BLOCK = 4 + + BLOCK_D = triton.next_power_of_2(head_size) + + grid = ( + num_tokens, + triton.cdiv(num_heads, HEAD_BLOCK), + ) + + reshape_and_cache_flash[grid]( + key, + value, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping, + k_scale if k_scale is not None else key, + v_scale if v_scale is not None else key, + key_cache.stride(0), + key.stride(0), + value.stride(0), + num_heads, + head_size, + key_cache.shape[1], + HEAD_BLOCK=HEAD_BLOCK, + BLOCK_D=BLOCK_D, + HAS_SWA=(swa_slot_mapping is not None), + USE_SCALE=(k_scale is not None), + ) + + +@triton.jit +def reshape_and_cache_shuffle_5d( + key_ptr, + value_ptr, + key_cache_ptr, + value_cache_ptr, + slot_mapping_ptr, + swa_slot_mapping_ptr, + key_stride_token, + value_stride_token, + num_heads, + head_size, + block_size, + X: tl.constexpr, + HEAD_BLOCK: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_SWA: tl.constexpr, +): + """Scatter per-token (num_tokens, num_heads, head_size) K/V into the + SHUFFLE 5D "vectorized" KV cache layout used by aiter CK + `mha_batch_prefill_func` and aiter `pa_decode_gluon`. + + K cache shape: (num_blocks, num_heads, head_size // X, block_size, X) + V cache shape: (num_blocks, num_heads, block_size // X, head_size, X) + where X = 16 // element_size (=8 for bf16/fp16, =16 for fp8). + block_size must be divisible by X, and head_size must be divisible by X. + + Each program handles one token and a HEAD_BLOCK-wide slice of heads. + """ + token_idx = tl.program_id(0) + head_block_idx = tl.program_id(1) + + slot_idx = tl.load(slot_mapping_ptr + token_idx) + if HAS_SWA: + slot_idx = tl.load(swa_slot_mapping_ptr + slot_idx) + if slot_idx < 0: + return + + block_idx = slot_idx // block_size + slot_in_page = slot_idx % block_size + page_outer = slot_in_page // X + page_inner = slot_in_page % X + + head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK) + head_mask = head_idx < num_heads + d = tl.arange(0, BLOCK_D) + d_mask = d < head_size + d_outer = d // X + d_inner = d % X + + src_off = token_idx * key_stride_token + head_idx[:, None] * head_size + d[None, :] + src_mask = head_mask[:, None] & d_mask[None, :] + k = tl.load(key_ptr + src_off, mask=src_mask) + src_off_v = ( + token_idx * value_stride_token + head_idx[:, None] * head_size + d[None, :] + ) + v = tl.load(value_ptr + src_off_v, mask=src_mask) + + layer_stride = num_heads * head_size * block_size + head_stride = head_size * block_size + + k_tgt = ( + block_idx * layer_stride + + head_idx[:, None] * head_stride + + d_outer[None, :] * block_size * X + + slot_in_page * X + + d_inner[None, :] + ) + tl.store(key_cache_ptr + k_tgt, k, mask=src_mask) + + v_tgt = ( + block_idx * layer_stride + + head_idx[:, None] * head_stride + + page_outer * head_size * X + + d[None, :] * X + + page_inner + ) + tl.store(value_cache_ptr + v_tgt, v, mask=src_mask) + + +def launch_reshape_and_cache_shuffle_5d( + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + swa_slot_mapping=None, +): + """Launcher for reshape_and_cache_shuffle_5d. + + Args: + key/value: (num_tokens, num_heads, head_size) source tensors + key_cache: (num_blocks, num_heads, head_size//X, block_size, X) + value_cache: (num_blocks, num_heads, block_size//X, head_size, X) + slot_mapping: per-token destination slot in [0, num_blocks*block_size) + """ + num_tokens, num_heads, head_size = key.shape + assert value.shape == key.shape, "K/V must share token-major shape" + assert key_cache.dim() == 5 and value_cache.dim() == 5 + num_blocks, kc_H, kc_D_over_X, block_size, X = key_cache.shape + assert kc_H == num_heads and kc_D_over_X * X == head_size + vb_blocks, vc_H, vc_page_over_X, vc_D, vc_X = value_cache.shape + assert ( + vc_H == num_heads + and vc_page_over_X * X == block_size + and vc_D == head_size + and vc_X == X + ) + assert block_size % X == 0 and head_size % X == 0 + + HEAD_BLOCK = min(4, triton.next_power_of_2(num_heads)) + BLOCK_D = triton.next_power_of_2(head_size) + grid = (num_tokens, triton.cdiv(num_heads, HEAD_BLOCK)) + + reshape_and_cache_shuffle_5d[grid]( + key, + value, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping if swa_slot_mapping is not None else slot_mapping, + key.stride(0), + value.stride(0), + num_heads, + head_size, + block_size, + X=X, + HEAD_BLOCK=HEAD_BLOCK, + BLOCK_D=BLOCK_D, + HAS_SWA=(swa_slot_mapping is not None), + ) + + +@triton.jit +def gather_shuffle_5d_to_linear( + key_cache_ptr, + value_cache_ptr, + key_out_ptr, # (T, num_heads, head_size), store dtype + value_out_ptr, # (T, num_heads, head_size), store dtype + slot_mapping_ptr, # (T,) absolute pool slot id per token + key_out_stride_token, + value_out_stride_token, + num_heads, + head_size, + block_size, + X: tl.constexpr, + HEAD_BLOCK: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Inverse of :func:`reshape_and_cache_shuffle_5d`. + + Gather one token's K/V from the SHUFFLE 5D paged cache into the + canonical (T, H, D) layout that aiter's ``mha_batch_prefill_func`` + expects in LINEAR mode. Source addressing is identical to the + writer kernel so any bit-exact round-trip is guaranteed. + """ + token_idx = tl.program_id(0) + head_block_idx = tl.program_id(1) + + slot_idx = tl.load(slot_mapping_ptr + token_idx) + + block_idx = slot_idx // block_size + slot_in_page = slot_idx % block_size + page_outer = slot_in_page // X + page_inner = slot_in_page % X + + head_idx = head_block_idx * HEAD_BLOCK + tl.arange(0, HEAD_BLOCK) + head_mask = head_idx < num_heads + d = tl.arange(0, BLOCK_D) + d_mask = d < head_size + d_outer = d // X + d_inner = d % X + + layer_stride = num_heads * head_size * block_size + head_stride = head_size * block_size + + src_mask = head_mask[:, None] & d_mask[None, :] + k_src = ( + block_idx * layer_stride + + head_idx[:, None] * head_stride + + d_outer[None, :] * block_size * X + + slot_in_page * X + + d_inner[None, :] + ) + k = tl.load(key_cache_ptr + k_src, mask=src_mask) + v_src = ( + block_idx * layer_stride + + head_idx[:, None] * head_stride + + page_outer * head_size * X + + d[None, :] * X + + page_inner + ) + v = tl.load(value_cache_ptr + v_src, mask=src_mask) + + dst_k = ( + token_idx * key_out_stride_token + head_idx[:, None] * head_size + d[None, :] + ) + tl.store(key_out_ptr + dst_k, k, mask=src_mask) + dst_v = ( + token_idx * value_out_stride_token + head_idx[:, None] * head_size + d[None, :] + ) + tl.store(value_out_ptr + dst_v, v, mask=src_mask) + + +def launch_gather_shuffle_5d_to_linear( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, +): + """Inverse of :func:`launch_reshape_and_cache_shuffle_5d`. + + Returns ``(key_out, value_out)`` each shaped + ``(T, num_heads, head_size)`` in ``key_cache.dtype`` / + ``value_cache.dtype``. The caller is responsible for passing the + right per-tensor descales downstream when ``store_dtype`` is fp8. + + Args: + key_cache: (num_blocks, num_heads, head_size // X, block_size, X) + value_cache: (num_blocks, num_heads, block_size // X, head_size, X) + slot_mapping: (T,) per-token absolute slot id in + ``[0, num_blocks * block_size)`` + """ + assert key_cache.dim() == 5 and value_cache.dim() == 5 + num_blocks, num_heads, kc_D_over_X, block_size, X = key_cache.shape + vc_blocks, vc_H, vc_page_over_X, vc_D, vc_X = value_cache.shape + assert vc_blocks == num_blocks and vc_H == num_heads + assert vc_page_over_X * X == block_size and vc_X == X + head_size = kc_D_over_X * X + assert vc_D == head_size + + num_tokens = slot_mapping.numel() + key_out = torch.empty( + (num_tokens, num_heads, head_size), + dtype=key_cache.dtype, + device=key_cache.device, + ) + value_out = torch.empty( + (num_tokens, num_heads, head_size), + dtype=value_cache.dtype, + device=value_cache.device, + ) + + HEAD_BLOCK = min(4, triton.next_power_of_2(num_heads)) + BLOCK_D = triton.next_power_of_2(head_size) + grid = (num_tokens, triton.cdiv(num_heads, HEAD_BLOCK)) + + gather_shuffle_5d_to_linear[grid]( + key_cache, + value_cache, + key_out, + value_out, + slot_mapping, + key_out.stride(0), + value_out.stride(0), + num_heads, + head_size, + block_size, + X=X, + HEAD_BLOCK=HEAD_BLOCK, + BLOCK_D=BLOCK_D, + ) + return key_out, value_out + + +@triton.jit +def _get_gptj_rotated_x( + x, + x_rotated_mask, + BLOCK_D: tl.constexpr, + BLOCK_D_HALF: tl.constexpr, +): + # GPT-J rotary layout: + # Pair adjacent dimensions and apply: + # [x0, x1, x2, x3] -> [-x1, x0, -x3, x2] + + # Apply sign inversion on odd positions. + x_rotated = tl.where(x_rotated_mask, x, -x) + # Reshape into (D/2, 2) pairs. + x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2)) + # Swap each pair. + x_rotated = tl.flip(x_rotated, 1) + # Flatten back to original shape. + x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) + return x_rotated + + +@triton.jit +def _get_neox_rotated_x( + x, + x_rotated_mask, + BLOCK_D: tl.constexpr, + BLOCK_D_HALF: tl.constexpr, +): + # GPT-NeoX rotary layout: + # Split head dimension into two halves: + # [x0, x1, x2, x3] -> [-x2, -x3, x0, x1] + + # Keep first half positive, second half negative. + x_rotated = tl.where(x_rotated_mask, x, -x) + # Reshape into (2, D/2). + x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF)) + # Reverse each half. + x_rotated = tl.flip(x_rotated, 1) + # Flatten and reverse full vector. + x_rotated = tl.reshape(x_rotated, (BLOCK_D,)) + x_rotated = tl.flip(x_rotated, 0) + return x_rotated + + +@triton.jit +def _unit_rope( + x_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX: tl.constexpr, + BLOCK_D_pe: tl.constexpr, + BLOCK_D_HALF_pe: tl.constexpr, +): + # Load one full attention head vector. + x_pe = tl.load(x_ptrs) + + # Stage 1: Build rotated vector according to rotary layout. + if IS_NEOX: + x_rotated_mask = d_pe_offs < BLOCK_D_HALF_pe + x_pe_rotated = _get_neox_rotated_x( + x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe + ) + else: + x_rotated_mask = d_pe_offs % 2 == 0 + x_pe_rotated = _get_gptj_rotated_x( + x_pe, x_rotated_mask, BLOCK_D_pe, BLOCK_D_HALF_pe + ) + + # Stage 2: Apply RoPE transform: + # x' = x*cos + rotate(x)*sin + x_pe = x_pe * cos + x_pe_rotated * sin + + return x_pe + + +@triton.jit +def _load_cos_sin( + cos_sin_ptr, + pos, + d_cos_offs, + stride_t, + stride_d, + freq_dim, +): + base = pos * stride_t + cos = tl.load(cos_sin_ptr + base + d_cos_offs * stride_d) + sin = tl.load(cos_sin_ptr + base + (d_cos_offs + freq_dim) * stride_d) + return cos, sin + + +@triton.jit +def _fused_qk_rope_reshape_and_cache_kernel( + q_ptr, + k_ptr, + v_ptr, + pos_ptr, + cos_sin_ptr, + offs_ptr, + key_cache_ptr, + value_cache_ptr, + slot_mapping_ptr, + swa_slot_mapping_ptr, + q_out_ptr, + k_out_ptr, + zeros_out_ptr, + T, + T_slot, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + v_stride_t, + v_stride_h, + v_stride_d, + cos_sin_stride_t, + cos_sin_stride_d, + q_out_stride_t, + q_out_stride_h, + q_out_stride_d, + k_out_stride_t, + k_out_stride_h, + k_out_stride_d, + key_cache_stride_t, + key_cache_stride_h, + key_cache_stride_d, + key_cache_stride_b, + key_cache_stride_x, + value_cache_stride_t, + value_cache_stride_h, + value_cache_stride_d, + value_cache_stride_b, + value_cache_stride_slot_chunk, + value_cache_stride_x, + zeros_out_stride_t, + zeros_out_stride_h, + zeros_out_stride_d, + k_scale_ptr, + v_scale_ptr, + QH_PER_KH: tl.constexpr, + QH: tl.constexpr, + KH: tl.constexpr, + REUSE_FREQS_FRONT_PART: tl.constexpr, + IS_NEOX: tl.constexpr, + BLOCK_D_pe: tl.constexpr, + BLOCK_D_HALF_pe: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + X_SIZE: tl.constexpr, + FLASH_LAYOUT: tl.constexpr, + VALUE_SHUFFLE_LAYOUT: tl.constexpr = False, + HAVE_POS: tl.constexpr = False, + HAVE_K_SCALE: tl.constexpr = False, + HAVE_V_SCALE: tl.constexpr = False, + HAVE_ZEROS: tl.constexpr = False, + HAS_SWA: tl.constexpr = False, +): + # ============================================================ + # Stage 0: Static stride assumptions for Triton compiler + # + # These assumptions help Triton optimize pointer arithmetic and + # simplify generated address calculations. + # ============================================================ + + tl.assume(q_stride_t >= 0) + tl.assume(q_stride_h >= 0) + tl.assume(q_stride_d >= 0) + tl.assume(k_stride_t >= 0) + tl.assume(k_stride_h >= 0) + tl.assume(k_stride_d >= 0) + tl.assume(v_stride_t >= 0) + tl.assume(v_stride_h >= 0) + tl.assume(v_stride_d >= 0) + tl.assume(cos_sin_stride_t >= 0) + tl.assume(cos_sin_stride_d >= 0) + tl.assume(q_out_stride_t >= 0) + tl.assume(q_out_stride_h >= 0) + tl.assume(q_out_stride_d >= 0) + tl.assume(k_out_stride_t >= 0) + tl.assume(k_out_stride_h >= 0) + tl.assume(k_out_stride_d >= 0) + tl.assume(key_cache_stride_t >= 0) + tl.assume(key_cache_stride_h >= 0) + tl.assume(key_cache_stride_d >= 0) + tl.assume(key_cache_stride_b >= 0) + tl.assume(key_cache_stride_x >= 0) + tl.assume(value_cache_stride_t >= 0) + tl.assume(value_cache_stride_h >= 0) + tl.assume(value_cache_stride_d >= 0) + tl.assume(value_cache_stride_b >= 0) + tl.assume(value_cache_stride_slot_chunk >= 0) + tl.assume(value_cache_stride_x >= 0) + tl.assume(zeros_out_stride_t >= 0) + tl.assume(zeros_out_stride_h >= 0) + tl.assume(zeros_out_stride_d >= 0) + + # ============================================================ + # Stage 1: Program instance mapping + # + # Each program handles: + # - one (token, q_head) for Q path + # - selected KV ownership for cache write path + # + # pid layout: + # [0, T*QH) -> decode Q path + # [T*QH, extra KV) -> KV-only path + # ============================================================ + + pid = tl.program_id(0) + tl.assume(pid >= 0) + + d_pe_offs = tl.arange(0, BLOCK_D_pe).to(tl.int64) + + # ============================================================ + # Stage 2: Main decode path (Q always active) + # ============================================================ + + if pid < T * QH: + pid_t = pid // QH + pid_hq = pid % QH + + # -------------------------------------------------------- + # Stage 2.1: Compute rotary frequency offsets + # + # RoPE frequencies may be stored as: + # D/2 frequencies (shared front-half) + # D frequencies (full explicit) + # -------------------------------------------------------- + + if REUSE_FREQS_FRONT_PART: + if IS_NEOX: + d_cos_offs = d_pe_offs + d_cos_offs = tl.where( + (d_cos_offs >= BLOCK_D_HALF_pe) & (d_cos_offs < BLOCK_D_pe), + d_cos_offs - BLOCK_D_HALF_pe, + d_cos_offs, + ).to(d_cos_offs.dtype) + # d_cos_mask = d_cos_offs < BLOCK_D_pe + else: + d_cos_offs = d_pe_offs // 2 + # d_cos_mask = d_cos_offs < BLOCK_D_HALF_pe + else: + d_cos_offs = d_pe_offs + # d_cos_mask = d_cos_offs < BLOCK_D_pe + + # -------------------------------------------------------- + # Stage 2.2: Load token position and optional offset + # + # offs_ptr is used by chunked prefill / sliding-window decode. + # -------------------------------------------------------- + pos = tl.load(pos_ptr + pid_t) + if HAVE_POS: + offset = tl.load(offs_ptr + pid_t) + pos = pos + offset + + # -------------------------------------------------------- + # Stage 2.3: Load cosine / sine table + # -------------------------------------------------------- + # cos_offs = pos * cos_stride_t + d_cos_offs * cos_stride_d + # cos = tl.load(cos_ptr + cos_offs) + # sin = tl.load(sin_ptr + cos_offs) + + freq_dim = BLOCK_D_HALF_pe if REUSE_FREQS_FRONT_PART else BLOCK_D_pe + + cos, sin = _load_cos_sin( + cos_sin_ptr, + pos, + d_cos_offs, + cos_sin_stride_t, + cos_sin_stride_d, + freq_dim, + ) + + # -------------------------------------------------------- + # Stage 2.4: Apply RoPE to Q + # -------------------------------------------------------- + q_ptrs = ( + q_ptr + pid_t * q_stride_t + pid_hq * q_stride_h + d_pe_offs * q_stride_d + ) + q_pe = _unit_rope( + q_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX, + BLOCK_D_pe, + BLOCK_D_HALF_pe, + ) + + # Store rotated Q output. + q_out_ptrs = ( + q_out_ptr + + pid_t * q_out_stride_t + + pid_hq * q_out_stride_h + + d_pe_offs * q_out_stride_d + ) + tl.store(q_out_ptrs, q_pe.to(q_out_ptr.dtype.element_ty)) + + if HAVE_ZEROS: + z = tl.zeros((BLOCK_D_pe,), dtype=zeros_out_ptr.dtype.element_ty) + zeros_out_ptrs = ( + zeros_out_ptr + + pid_t * zeros_out_stride_t + + pid_hq * zeros_out_stride_h + + d_pe_offs * zeros_out_stride_d + ) + tl.store(zeros_out_ptrs, z) + + # ======================================================== + # Stage 3: KV ownership path + # + # Only one Q group leader writes KV: + # pid_hq % QH_PER_KH == 0 + # + # This prevents duplicated KV cache writes. + # ======================================================== + + if pid_hq % QH_PER_KH == 0: + # ---------------------------------------------------- + # Stage 3.1: Resolve cache slot + # ---------------------------------------------------- + pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) + if HAS_SWA: + pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) + + # ------------------------------------------------ + # Stage 3.2: Apply RoPE to K + # ------------------------------------------------ + if pid_slot >= 0: + pid_t_slot = pid_slot // BLOCK_SIZE + pid_b = pid_slot % BLOCK_SIZE + pid_hk = pid_hq // QH_PER_KH + if HAVE_K_SCALE: + k_scale = tl.load(k_scale_ptr) + else: + k_scale = 1 + k_ptrs = ( + k_ptr + + pid_t * k_stride_t + + pid_hk * k_stride_h + + d_pe_offs * k_stride_d + ) + k_pe = _unit_rope( + k_ptrs, + cos, + sin, + d_pe_offs, + IS_NEOX, + BLOCK_D_pe, + BLOCK_D_HALF_pe, + ) + + k_out_ptrs = ( + k_out_ptr + + pid_t * k_out_stride_t + + pid_hk * k_out_stride_h + + d_pe_offs * k_out_stride_d + ) + tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) + + # ------------------------------------------------ + # Stage 3.3: Optional fp8 scaling before cache + # ------------------------------------------------ + + k_scale_rcprl = 1 / k_scale + k_pe = k_pe * k_scale_rcprl + + # ------------------------------------------------ + # Stage 3.4: Write K cache + # + # Two layouts supported: + # FLASH_LAYOUT + # paged KV layout + # ------------------------------------------------ + + if FLASH_LAYOUT: + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_b * key_cache_stride_b + + pid_hk * key_cache_stride_h + + d_pe_offs * key_cache_stride_d + ) + else: + k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) + dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) + x_offs = tl.arange(0, X_SIZE).to(tl.int64) + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_hk * key_cache_stride_h + + dx_offs[:, None] * key_cache_stride_d + + pid_b * key_cache_stride_b + + x_offs[None, :] * key_cache_stride_x + ) + + tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) + + # ------------------------------------------------ + # Stage 3.5: Write V cache + # + # Supports: + # normal layout + # shuffle layout + # ------------------------------------------------ + + v_ptrs = ( + v_ptr + + pid_t * v_stride_t + + pid_hk * v_stride_h + + d_pe_offs * v_stride_d + ) + if HAVE_V_SCALE: + v_scale = tl.load(v_scale_ptr) + else: + v_scale = 1 + v_scale_rcprl = 1 / v_scale + v = tl.load(v_ptrs) * v_scale_rcprl + if VALUE_SHUFFLE_LAYOUT: + slot_chunk = pid_b // X_SIZE + x_off = pid_b % X_SIZE + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + slot_chunk * value_cache_stride_slot_chunk + + d_pe_offs.to(tl.int64) * value_cache_stride_d + + x_off * value_cache_stride_x + ) + else: + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + d_pe_offs.to(tl.int64) * value_cache_stride_d + + pid_b * value_cache_stride_b + ) + tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) + # ============================================================ + # Stage 4: Extra KV-only path + # + # Handles tokens that only require cache update: + # T_slot > T + # + # No Q / no RoPE on Q branch. + # ============================================================ + else: + pid = pid - T * QH + T * KH + if pid < T_slot * KH: + pid_t = pid // KH + pid_hk = pid % KH + pid_slot = tl.load(slot_mapping_ptr + pid_t).to(tl.int64) + if HAS_SWA: + pid_slot = tl.load(swa_slot_mapping_ptr + pid_slot) + + if pid_slot >= 0: + pid_t_slot = pid_slot // BLOCK_SIZE + pid_b = pid_slot % BLOCK_SIZE + if HAVE_K_SCALE: + k_scale = tl.load(k_scale_ptr) + else: + k_scale = 1 + k_ptrs = ( + k_ptr + + pid_t * k_stride_t + + pid_hk * k_stride_h + + d_pe_offs * k_stride_d + ) + + k_pe = tl.load(k_ptrs) + + k_out_ptrs = ( + k_out_ptr + + pid_t * k_out_stride_t + + pid_hk * k_out_stride_h + + d_pe_offs * k_out_stride_d + ) + tl.store(k_out_ptrs, k_pe.to(k_out_ptr.dtype.element_ty)) + + k_scale_rcprl = 1 / k_scale + k_pe = k_pe * k_scale_rcprl + + if FLASH_LAYOUT: + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + d_pe_offs * key_cache_stride_d + + pid_b * key_cache_stride_b + + pid_hk * key_cache_stride_h + ) + else: + k_pe = tl.reshape(k_pe, (BLOCK_D_pe // X_SIZE, X_SIZE)) + dx_offs = tl.arange(0, BLOCK_D_pe // X_SIZE).to(tl.int64) + x_offs = tl.arange(0, X_SIZE).to(tl.int64) + k_out_ptrs = ( + key_cache_ptr + + pid_t_slot * key_cache_stride_t + + pid_hk * key_cache_stride_h + + dx_offs[:, None] * key_cache_stride_d + + pid_b * key_cache_stride_b + + x_offs[None, :] * key_cache_stride_x + ) + tl.store(k_out_ptrs, k_pe.to(key_cache_ptr.dtype.element_ty)) + + v_ptrs = ( + v_ptr + + pid_t * v_stride_t + + pid_hk * v_stride_h + + d_pe_offs * v_stride_d + ) + if HAVE_V_SCALE: + v_scale = tl.load(v_scale_ptr) + else: + v_scale = 1 + v_scale_rcprl = 1 / v_scale + v = tl.load(v_ptrs) * v_scale_rcprl + if VALUE_SHUFFLE_LAYOUT: + slot_chunk = pid_b // X_SIZE + x_off = pid_b % X_SIZE + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + slot_chunk * value_cache_stride_slot_chunk + + d_pe_offs * value_cache_stride_d + + x_off * value_cache_stride_x + ) + else: + v_out_ptrs = ( + value_cache_ptr + + pid_t_slot * value_cache_stride_t + + pid_hk * value_cache_stride_h + + d_pe_offs * value_cache_stride_d + + pid_b * value_cache_stride_b + ) + tl.store(v_out_ptrs, v.to(value_cache_ptr.dtype.element_ty)) + + +def fused_qk_rope_reshape_and_cache( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + pos: torch.Tensor, + cos_sin: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + is_neox: bool, + flash_layout: bool, + apply_scale: bool = True, + offs: torch.Tensor = None, + q_out: torch.Tensor = None, + k_out: torch.Tensor = None, + output_zeros: bool = True, + zeros_out: torch.Tensor = None, + swa_slot_mapping=None, +): + """ + Perform RoPE on q and k and along the last dimension and copy k and v in to key_cache and value_cache inplace + + Key parameters: + - q: shape (T, QH, D). + - k: shape (T_slot, KH, D). + - v: shape (T_slot, KH, D). + - if flash_layout: + - key_cache: shape (T_cache, block_size, KH, D). + - value_cache: shape (T_cache, block_size, KH, D). + - else: + - key_cache: shape (T_cache, KH, D // x, block_size, x). + - value_cache: shape (T_cache, KH, D, block_size). + - slot_mapping: shape (T_slot, ). + + T is the number of decode tokens, T_cahce * block_size is the max number of tokens of kv_cache + QH must be multiple of KH + + Returns: + - q_out: same shape as input q. + - k_out: same shape as input k. + - key_cache: same shape as input key_cache (inplace). + - value_cache: same shape as input value_cache (inplace). + - zeros_out: same shape as input q. + """ + + t, qh, d = q.shape + tk, kh, dk = k.shape + tv, vh, dv = v.shape + if flash_layout: + t_cache, block_size, kh_cache, dk_cache = key_cache.shape + t_cache_v, block_size_v, vh_cache, dv_cache = value_cache.shape + value_shuffle_layout = False + else: + t_cache, kh_cache, dkx_cache, block_size, x_cache = key_cache.shape + if value_cache.ndim == 5: + # value_cache shuffle: (num_blocks, num_kv_heads, block_size // x, head_size, x) + t_cache_v, vh_cache, slot_chunk_v, dv_cache, x_v = value_cache.shape + value_shuffle_layout = True + block_size_v = slot_chunk_v * x_v + assert block_size_v == block_size and x_v == x_cache, ( + f"value_cache shuffle (T,KH,block_size//x,D,x) must match key: " + f"{block_size_v=} {block_size=} {x_v=} {x_cache=}" + ) + else: + t_cache_v, vh_cache, dv_cache, block_size_v = value_cache.shape + value_shuffle_layout = False + (t_slot,) = slot_mapping.shape + + assert ( + t == tk == tv and t_slot <= tk + ), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}" + assert ( + block_size == block_size_v + ), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}" + assert ( + kh == vh == kh_cache == vh_cache + ), "KV head should be identical for k, v, key_cache, and value_cache" + assert ( + t_cache == t_cache_v + ), "Number of tokens should be identical for key_cache, and value_cache" + if flash_layout: + assert ( + d == dk == dv == dk_cache == dv_cache + ), "D dimension should be identical for q, k, and v" + else: + assert ( + d == dk == dv == dkx_cache * x_cache == dv_cache + ), "D dimension should be identical for q, k, and v" + assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2" + + assert d == triton.next_power_of_2(d), "D dimension should be power of 2" + assert block_size == triton.next_power_of_2( + block_size + ), "block_size should be power of 2" + assert qh % kh == 0, "Q heads must be multiple of H heads" + d_freq = cos_sin.shape[-1] // 2 + assert (d_freq == d // 2) or ( + d_freq == d + ), "cos/sin last dim should be the same or half of the qk last dim" + reuse_freqs_front_part = d_freq == d // 2 + + if q_out is None: + q_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) + + if k_out is None: + k_out = torch.empty((tk, kh, dk), dtype=k.dtype, device=q.device) + + if zeros_out is not None: + tz, qhz, dz = zeros_out.shape + assert ( + t == tz and qh == qhz and d == dz + ), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}" + output_zeros = True + elif output_zeros: + zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device) + else: + zeros_out = None + + n_pid = t * qh + (t_slot - t) * kh if t_slot >= t else t * qh + grid = (n_pid, 1, 1) + _fused_qk_rope_reshape_and_cache_kernel[grid]( + q, + k, + v, + pos, + cos_sin, + offs, + key_cache, + value_cache, + slot_mapping, + swa_slot_mapping, + q_out, + k_out, + zeros_out, + t, + t_slot, + *q.stride(), + *k.stride(), + *v.stride(), + cos_sin.stride(0), + cos_sin.stride(-1), + *q_out.stride(), + *k_out.stride(), + key_cache.stride(0) if not flash_layout else key_cache.stride(0), + key_cache.stride(1) if not flash_layout else key_cache.stride(2), + key_cache.stride(2) if not flash_layout else key_cache.stride(3), + key_cache.stride(3) if not flash_layout else key_cache.stride(1), + key_cache.stride(4) if not flash_layout else 0, + value_cache.stride(0) if not flash_layout else value_cache.stride(0), + value_cache.stride(1) if not flash_layout else value_cache.stride(2), + ( + value_cache.stride(3) + if (not flash_layout and value_shuffle_layout) + else (value_cache.stride(2) if not flash_layout else value_cache.stride(3)) + ), + ( + 0 + if (not flash_layout and value_shuffle_layout) + else (value_cache.stride(3) if not flash_layout else value_cache.stride(1)) + ), + value_cache.stride(2) if (not flash_layout and value_shuffle_layout) else 0, + value_cache.stride(4) if (not flash_layout and value_shuffle_layout) else 0, + zeros_out.stride(0) if zeros_out is not None else 0, + zeros_out.stride(1) if zeros_out is not None else 0, + zeros_out.stride(2) if zeros_out is not None else 0, + k_scale_ptr=k_scale, + v_scale_ptr=v_scale, + QH_PER_KH=qh // kh, + QH=qh, + KH=kh, + REUSE_FREQS_FRONT_PART=reuse_freqs_front_part, + IS_NEOX=is_neox, + BLOCK_D_pe=d, + BLOCK_D_HALF_pe=d // 2, + BLOCK_SIZE=block_size, + X_SIZE=x_cache if not flash_layout else 0, + FLASH_LAYOUT=flash_layout, + VALUE_SHUFFLE_LAYOUT=value_shuffle_layout, + HAVE_POS=(offs is not None), + HAVE_K_SCALE=(k_scale is not None and apply_scale), + HAVE_V_SCALE=(v_scale is not None and apply_scale), + HAVE_ZEROS=output_zeros, + HAS_SWA=(swa_slot_mapping is not None), + num_warps=1, + ) + + if zeros_out is not None: + return q_out.view(-1, qh * d), k_out, key_cache, value_cache, zeros_out + return q_out.view(-1, qh * d), k_out, key_cache, value_cache + + def assert_buffer_fits(used: int, capacity: int, what: str, **context) -> None: """Safety guard: a preallocated cuda-graph buffer must hold the runtime write. diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 6cf269b8d..89be4a16e 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index a29f3faf9..1a2577b0d 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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( - x, (0, self.hidden_pad), mode="constant", value=0.0 - ) + # Skip the explicit pad if x already arrives at the padded + # 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, diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index 2928812b2..0770a5786 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/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_shape = fused_set_kv_buffer_arg["key_cache"].shape - qk_head_dim = k_cache_shape[-1] - tp_k_head_num = k_cache_shape[-2] + k_cache = fused_set_kv_buffer_arg["key_cache"] + # 5D SHUFFLE pool feeds raw (N, H, D/x, page, x) K cache; + # NHD 3D pool feeds the legacy 4D paged view. Auto-detect. + 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, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 4911a1492..24abecec0 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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 - # 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) - ] + if self.kv_cache_layout == "vectorized_5d": + total_slots = self.size + self.page_size + num_blocks = total_slots // self.page_size + x = self._kv_vector_x + # K: (num_blocks, H, D_k // X, page, X) + self.k_buffer = [ + torch.zeros( + ( + num_blocks, + self.head_num, + self.head_dim // x, + self.page_size, + x, + ), + dtype=self.store_dtype, + device=self.device, + ) + 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, diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index b979a9264..c2a2419d3 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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) - topk_output = self.topk(hidden_states, router_logits) + router_logits, _ = self.router(router_input) + 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( diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 341f7b458..fa77d4afa 100644 --- a/python/sglang/srt/models/utils.py +++ b/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( - -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 - ), + "key_cache": key_cache, + "value_cache": value_cache, "slot_mapping": forward_batch.out_cache_loc, "swa_slot_mapping": slot_mapping_swa, } diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 690f9d518..e8b2a360e 100644 --- a/python/sglang/srt/server_args.py +++ b/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