Skip to content

[ROCm][PERF] Enable AITER PA gluon decode for MiniMax-M3 MTP and dense layers #52849

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ukannika wants to merge 5 commits into
base: main
Choose a base branch
from
+260 −58
Open

[ROCm][PERF] Enable AITER PA gluon decode for MiniMax-M3 MTP and dense layers #52849

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4360a28
enable pa_gluon kernel in rocm_aiter_fa
ukannika Aug 19, 2026
f71d75f
Fix the uniform-decode assumption from PA decode gluon
ykamiset Aug 19, 2026
5cfdff8
fix drafting function for missing variables
ykamiset Aug 19, 2026
9842652
remove blank line
Aug 19, 2026
c2b78e4
[Bugfix][ROCm] Fix KV block zeroer dim classification for permuted la…
Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
278 changes: 219 additions & 59 deletions vllm/v1/attention/backends/rocm_aiter_fa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0

​# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Attention layer with AiterFlashAttention."""

Expand Down Expand Up @@ -33,9 +34,29 @@
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
from vllm.v1.kv_cache_interface import AttentionSpec

_PA_GLUON_MAX_QUERY_LEN = 4
_PA_GLUON_MAX_QUERY_GROUP_SIZE = 64
# Query group sizes the gluon paged-attention decode kernel is validated for: 8, 16
_PA_GLUON_QUERY_GROUP_SIZES = (8, 16)

# Enable PA decode gluon when the shuffle KV cache layout is on (the kernel
# reads K/V in that layout) and the head config is one of the above.
ENABLE_PA_GLUON = lambda num_heads_q, num_heads_kv: (
rocm_aiter_ops.is_shuffle_kv_cache_enabled()
and num_heads_q % num_heads_kv == 0
and num_heads_q // num_heads_kv in _PA_GLUON_QUERY_GROUP_SIZES
)

_PARTITION_SIZE_ROCM = 256
_CP_TOKENS_PER_ITER_ROCM = 32 * 1024
if current_platform.is_rocm():
from aiter.ops.triton.gluon.pa_decode_gluon import (
get_recommended_splits,
)
from aiter.ops.triton.gluon.pa_decode_gluon import (
pa_decode_gluon as _pa_decode_gluon,
)

from vllm.triton_utils import tl, triton

def block_size(x, head_dim):
Expand Down Expand Up @@ -318,6 +339,7 @@ def reshape_and_cache_shuffle_triton(
@dataclass
class AiterFlashAttentionDecodeMetadata:
max_query_len: int
uniform_query_len: int | None


@dataclass
Expand Down Expand Up @@ -474,10 +496,11 @@ def build(
and self.scale.numel() == 1
and is_quantized_kv_cache(self.vllm_config.cache_config.cache_dtype)
):
layers = get_layers_from_vllm_config(self.vllm_config, Attention)
first_layer_name = [k for k in layers][0]
# Size the scales from a layer this builder owns. The draft model
# runs its own builder over its own KV cache, so the first layer of
# the whole config can carry an unrelated block count.
kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[
first_layer_name
self.layer_names[0]
].kv_cache.shape
num_blocks = kv_cache_shape[0]
self.scale = torch.ones(
Expand Down Expand Up @@ -508,8 +531,15 @@ def build(

decode_metadata = None
if num_decodes > 0:
decode_max_query_len = query_lens_cpu[:num_decodes].max().item()
uniform_query_len = (
decode_max_query_len
if num_decode_tokens == num_decodes * decode_max_query_len
else None
)
decode_metadata = AiterFlashAttentionDecodeMetadata(
max_query_len=query_lens_cpu[:num_decodes].max().item(),
max_query_len=decode_max_query_len,
uniform_query_len=uniform_query_len,
)

prefill_metadata = None
Expand Down Expand Up @@ -678,11 +708,23 @@ def build_for_drafting(
skip split_decodes_prefills_and_extends() and avoid all .cpu() /
.item() calls that would otherwise break CUDA graph capture.
"""
num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens
# Uniform-decode assumption does not hold for the
# drafter's first forward after a target step: it inherits the target's
# per-request query lengths, so rows can be longer than gluon's limit or
# ragged. Those batches need the real split, which costs a sync.
if rocm_aiter_ops.is_shuffle_kv_cache_enabled() and (
max_query_len > _PA_GLUON_MAX_QUERY_LEN
or num_tokens != num_reqs * max_query_len
):
return self.build(
common_prefix_len=0, common_attn_metadata=common_attn_metadata
)

decode_metadata = AiterFlashAttentionDecodeMetadata(
max_query_len=common_attn_metadata.max_query_len,
max_query_len=max_query_len,
uniform_query_len=(
max_query_len if num_tokens == num_reqs * max_query_len else None
),
)

return AiterFlashAttentionMetadata(
Expand Down Expand Up @@ -737,6 +779,8 @@ def supports_attn_type(cls, attn_type: str) -> bool:

@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
return [MultipleOf(16)]
return [16, 32]

@classmethod
Expand Down Expand Up @@ -771,9 +815,25 @@ def get_kv_cache_shape(
) -> tuple[int, ...]:
if block_size % 16 != 0:
raise ValueError("Block size must be a multiple of 16.")

if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
return (num_blocks, 2, block_size, num_kv_heads, head_size)

@tjtanaa tjtanaa Aug 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we update this?

get_kv_cache_shape is used by gpu runner to determine how the kvcache is packed in the kvcache buffers/pool/management system. It is not referring to the kvcache shape expect by the kernels.

The KVCACHE is reshaped on the fly in the attention forward pass from (num_blocks, num_kv_heads, block_size, 2 * head_size) to (num_blocks, 2, block_size, num_kv_heads, head_size) every time.

                    num_blocks, block_size, num_kv_heads, _ = key_cache.shape
                    x = 16 // key_cache.element_size()
                    new_key_cache = key_cache.reshape(
                        num_blocks, num_kv_heads, head_size // x, block_size, x
                    )
                    new_value_cache = value_cache.reshape(
                        num_blocks, num_kv_heads, block_size // x, head_size, x
                    )

However, the kvcache stored in the kvcache management system is still (num_blocks, num_kv_heads, block_size, 2 * head_size)

@ukannika ukannika Aug 21, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tjtanaa we see an issue if we don't define the shape and stride order when the shuffle layout is enabled: the pool is then labeled (num_blocks, num_kv_heads, block_size, 2*head_size), where each tokens K is immediately followed by its own V, head_size elements apart. The cache write kernel still writes each blocks K in the shuffle layout because of shuffle layout flag enabled, which assumes K and V are kept apart. Looking into how is it working with shuffle layout enabled for asm pa kernel before this change.

@ukannika ukannika Aug 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tjtanaa Enabled the shuffle layout flag and tested the Llama2 70B model using the asm_pa code path. I am seeing an accuracy issue, and the current state of this file is broken due to the layout changes introduced in this PR #44455. AITER assembly paged-attention kernels require independently contiguous K and V storage. This PR addresses the accuracy issue as well.
Here's the command to reproduce accuracy issue

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1

vllm serve /model/llama2-70b-chat-hf/fp4_quantized_gptq \
  --dtype bfloat16 \
  --quantization quark \
  --tensor-parallel-size 1 \
  --max-model-len 2048 \
  --gpu-memory-utilization 0.94 \
  --attention-backend ROCM_AITER_FA \
  --kv-cache-dtype fp8 \
  --block-size 16 \
  --max-num-batched-tokens 32768 \
  --max-num-seqs 6400 \
  --enable-chunked-prefill \
  --async-scheduling \
  --host 0.0.0.0 \
  --port 8000

curl -s http://127.0.0.1:8000/v1/completions   -H "Content-Type: application/json"   -d '{
    "model": "/model/llama2-70b-chat-hf/fp4_quantized_gptq",
    "prompt": "The capital of France is",
    "max_tokens": 32,
    "temperature": 0
  }'```

# K and V are packed into the content dim: logical (B, H, N, 2*hs).
return (num_blocks, num_kv_heads, block_size, 2 * head_size)

@staticmethod
def get_kv_cache_stride_order(
include_num_layers_dimension: bool = False,
) -> tuple[int, ...]:
if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
# Physical layout matches the logical packed shape.
raise NotImplementedError
if include_num_layers_dimension:
raise NotImplementedError
# Hoist the K/V dim out so kv_cache[:, 0] and kv_cache[:, 1] are each a
# contiguous (num_blocks, block_size, num_kv_heads, head_size) range.
return (1, 0, 2, 3, 4)

@classmethod
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
from vllm.platforms.rocm import get_cdna_version
Expand Down Expand Up @@ -1154,11 +1214,22 @@ def forward(
if num_decodes > 0:
assert attn_metadata.decode_metadata is not None
decode_max_query_len = attn_metadata.decode_metadata.max_query_len

# Use unified_attention for speculative decoding (multi-token),
# sliding window, or sinks
# (pa_fwd_asm and paged_attention_v1 don't support sinks)
if (
decode_query_len = attn_metadata.decode_metadata.uniform_query_len

# check if we can use the gluon paged-attention decode kernel
use_gluon = (
ENABLE_PA_GLUON(self.num_heads, self.num_kv_heads)
and decode_query_len is not None
and decode_query_len <= _PA_GLUON_MAX_QUERY_LEN
and decode_query_len * (self.num_heads // self.num_kv_heads)
<= _PA_GLUON_MAX_QUERY_GROUP_SIZE
and (decode_query_len == 1 or attn_metadata.causal)
)
# Use unified_attention for the decodes the paged kernels can't
# take: sliding window, sinks, or a multi-token batch that gluon
# declined (pa_fwd_asm and paged_attention_v1 don't support
# sinks).
if not use_gluon and (
self.sliding_window[0] != -1
or decode_max_query_len > 1
or self.sinks is not None
Expand Down Expand Up @@ -1281,21 +1352,6 @@ def forward(
)
elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
_, num_heads, head_size = query.shape
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
tmp_out = torch.empty(
(num_seqs, num_heads, max_num_partitions, head_size),
dtype=query.dtype,
device=query.device,
)
exp_sums = torch.empty(
(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32,
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
num_blocks, block_size, num_kv_heads, _ = key_cache.shape
x = 16 // key_cache.element_size()
new_key_cache = key_cache.reshape(
Expand All @@ -1304,37 +1360,136 @@ def forward(
new_value_cache = value_cache.reshape(
num_blocks, num_kv_heads, block_size // x, head_size, x
)
k_qscale = (
layer._k_scale
if attn_metadata.k_scale is None
else attn_metadata.k_scale
)
v_qscale = (
layer._v_scale
if attn_metadata.v_scale is None
else attn_metadata.v_scale
)
rocm_aiter_ops.paged_attention_common(
Q=query[:num_decode_tokens],
K=new_key_cache,
V=new_value_cache,
tmp_out=tmp_out,
max_logits=max_logits,
exp_sums=exp_sums,
max_seq_len=attn_metadata.max_seq_len,
block_tables=attn_metadata.block_table[:num_decodes],
context_lens=attn_metadata.seq_lens[:num_decodes],
block_tables_stride0=attn_metadata.block_table[
:num_decodes
].stride(0),
scale=self.scale,
K_QScale_hip=k_qscale,
V_QScale_hip=v_qscale,
K_QScale_asm=k_qscale,
V_QScale_asm=v_qscale,
out_=output[:num_decode_tokens],
kv_cache_dtype=self.kv_cache_dtype,
)

if use_gluon:
is_fp8_kv = is_quantized_kv_cache(self.kv_cache_dtype)
# Per-tensor descale, as a float32 [1] tensor.
k_scale_gluon = (
layer._k_scale.reshape(1).to(torch.float32)
if is_fp8_kv
else None
)
v_scale_gluon = (
layer._v_scale.reshape(1).to(torch.float32)
if is_fp8_kv
else None
)
compute_type = (
current_platform.fp8_dtype() if is_fp8_kv else query.dtype
)
# The kernel folds the query positions into the group
# dim, so the intermediate buffers are sized by the
# combined extent.
query_group_size = decode_query_len * (
num_heads // num_kv_heads
)

sliding_window_int = (
self.sliding_window[0] + 1
if self.sliding_window[0] > 0
else 0
)
if sliding_window_int > 0:
max_context_partition_num = 1
context_partition_size = 128
else:
max_context_partition_num = get_recommended_splits(
num_decodes, num_kv_heads
)
context_partition_size = _PARTITION_SIZE_ROCM

intermediate_shape = (
num_decodes,
num_kv_heads,
max_context_partition_num,
query_group_size,
)
exp_sums = torch.empty(
intermediate_shape,
dtype=torch.float32,
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
temporary_output = torch.empty(
(*intermediate_shape, head_size),
dtype=output.dtype,
device=query.device,
)

_pa_decode_gluon(
output=output[:num_decode_tokens],
query=query[:num_decode_tokens],
key_cache=new_key_cache,
value_cache=new_value_cache,
context_lengths=attn_metadata.seq_lens[:num_decodes].to(
torch.int32
),
block_tables=attn_metadata.block_table[:num_decodes].to(
torch.int32
),
softmax_scale=self.scale,
query_length=decode_query_len,
max_context_partition_num=max_context_partition_num,
context_partition_size=context_partition_size,
compute_type=compute_type,
query_scale=None,
key_scale=k_scale_gluon,
value_scale=v_scale_gluon,
exp_sums=exp_sums,
max_logits=max_logits,
temporary_output=temporary_output,
alibi_slopes=self.alibi_slopes,
sinks=self.sinks,
sliding_window=sliding_window_int,
ps=True,
)
else:
num_seqs = attn_metadata.seq_lens.shape[0]
max_num_partitions = (
attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
) // _PARTITION_SIZE_ROCM
tmp_out = torch.empty(
(num_seqs, num_heads, max_num_partitions, head_size),
dtype=query.dtype,
device=query.device,
)
exp_sums = torch.empty(
(num_seqs, num_heads, max_num_partitions),
dtype=torch.float32,
device=query.device,
)
max_logits = torch.empty_like(exp_sums)
k_qscale = (
layer._k_scale
if attn_metadata.k_scale is None
else attn_metadata.k_scale
)
v_qscale = (
layer._v_scale
if attn_metadata.v_scale is None
else attn_metadata.v_scale
)
rocm_aiter_ops.paged_attention_common(
Q=query[:num_decode_tokens],
K=new_key_cache,
V=new_value_cache,
tmp_out=tmp_out,
max_logits=max_logits,
exp_sums=exp_sums,
max_seq_len=attn_metadata.max_seq_len,
block_tables=attn_metadata.block_table[:num_decodes],
context_lens=attn_metadata.seq_lens[:num_decodes],
block_tables_stride0=attn_metadata.block_table[
:num_decodes
].stride(0),
scale=self.scale,
K_QScale_hip=k_qscale,
V_QScale_hip=v_qscale,
K_QScale_asm=k_qscale,
V_QScale_asm=v_qscale,
out_=output[:num_decode_tokens],
kv_cache_dtype=self.kv_cache_dtype,
)
else:
_, num_heads, head_size = query.shape
nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8
Expand Down Expand Up @@ -1387,6 +1542,11 @@ def forward(
def _split_kv_cache(
self, kv_cache: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
# (B, 2, N, H, hs) -> two contiguous (B, N, H, hs), which is what
# the shuffle read/write kernels reinterpret in place.
key_cache, value_cache = kv_cache.unbind(1)
return key_cache, value_cache
# (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

Expand Down
Loading