[AMD] Optimize gpt-oss-120B performance (#27063)
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user