iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][DSv4.1][Perf] Emit MXFP8 from the sparse decode reduce and run wo_a as a grouped FP8 GEMM (#58456)

Browse Source 
Signed-off-by: fai <fangzhouai@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Fangzhou Ai
2026-09-24 11:58:14 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent dcfc17e0b1
commit 721d0e5c11
3 changed files with 664 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+177 -2
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools
from types import SimpleNamespace
import pytest
@@ -261,14 +262,19 @@ def _launch_sparse_decode_reduce(
adaptive_splits: bool,
positions: torch.Tensor | None = None,
cos_sin_cache: torch.Tensor | None = None,
out_dtype: torch.dtype = torch.bfloat16,
out_scale: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run the reduce kernel, with the inverse-RoPE epilogue if given positions."""
"""Run the reduce kernel, with the inverse-RoPE epilogue if given positions.
Passing ``out_scale`` selects the MXFP8 epilogue (``out_dtype`` e4m3).
"""
from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
num_queries, num_splits, num_heads = part_m.shape
out = torch.empty(
(num_queries, num_heads, HEAD_DIM),
dtype=torch.bfloat16,
dtype=out_dtype,
device=part_m.device,
)
attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device)
@@ -280,8 +286,11 @@ def _launch_sparse_decode_reduce(
out,
positions,
cos_sin_cache,
out_scale,
out.stride(0),
out.stride(1),
out_scale.stride(0) if out_scale is not None else 0,
HEAD_DIM // 32,
part_m.stride(0),
part_m.stride(1),
part_acc.stride(0),
@@ -298,6 +307,7 @@ def _launch_sparse_decode_reduce(
FUSE_INV_ROPE=positions is not None,
NOPE=NOPE_HEAD_DIM,
HALF=ROPE_HEAD_DIM // 2,
QUANT_OUT=out_scale is not None,
num_warps=4,
)
return out
@@ -1057,6 +1067,171 @@ def test_sparse_attn_decode_reduce_inverse_rope_epilogue(adaptive_splits: bool)
assert not torch.equal(actual[..., NOPE_HEAD_DIM:], unfused[..., NOPE_HEAD_DIM:])
def _mxfp8_dequant(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
blocks = data.float().view(*data.shape[:-1], -1, 32)
return (blocks * torch.exp2(scale.float() - 127.0)[..., None]).view(data.shape)
def _random_cos_sin_cache(max_pos: int, device: torch.device) -> torch.Tensor:
angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
return torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
@requires_gfx950
@torch.inference_mode()
def test_sparse_attn_decode_reduce_mxfp8_epilogue() -> None:
"""QUANT_OUT must equal MXFP8-quantizing the reduce's fp32 output."""
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
_mxfp8_e4m3_quantize_torch,
)
device = torch.device("cuda")
torch.manual_seed(7)
num_queries, num_splits, num_heads = 5, 4, 8
shape = (num_queries, num_splits, num_heads)
part_m = torch.randn(shape, device=device)
part_l = torch.rand(shape, device=device) + 0.5
part_acc = torch.randn((*shape, HEAD_DIM), device=device)
positions = torch.randint(0, 64, (num_queries,), device=device)
cos_sin_cache = _random_cos_sin_cache(64, device)
args = (part_m, part_l, part_acc, False, positions, cos_sin_cache)
rotated = _launch_sparse_decode_reduce(*args, out_dtype=torch.float32)
expected_data, expected_scale = _mxfp8_e4m3_quantize_torch(
rotated.view(num_queries, -1)
)
scale = torch.empty_like(expected_scale)
data = _launch_sparse_decode_reduce(
*args, out_dtype=torch.float8_e4m3fn, out_scale=scale
)
assert torch.equal(scale, expected_scale)
assert torch.equal(
data.view(num_queries, -1).view(torch.uint8), expected_data.view(torch.uint8)
)
@requires_gfx950
@torch.inference_mode()
def test_inverse_rope_mxfp8_rows() -> None:
"""Prefill rows get the same inverse RoPE + MXFP8 as the decode epilogue."""
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
_mxfp8_e4m3_quantize_torch,
)
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
rocm_inverse_rope_mxfp8_rows,
)
device = torch.device("cuda")
torch.manual_seed(5)
num_tokens, num_heads = 7, 4
o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).bfloat16()
positions = torch.randint(0, 64, (num_tokens,), device=device)
cos_sin_cache = _random_cos_sin_cache(64, device)
ref = o.float()
cos, sin = cos_sin_cache[positions, None].chunk(2, dim=-1)
even, odd = ref[..., NOPE_HEAD_DIM::2].clone(), ref[..., NOPE_HEAD_DIM + 1 :: 2]
ref[..., NOPE_HEAD_DIM::2] = even * cos + odd * sin
ref[..., NOPE_HEAD_DIM + 1 :: 2] = odd * cos - even * sin
ref = ref.view(num_tokens, -1)
_, expected_scale = _mxfp8_e4m3_quantize_torch(ref)
data = torch.empty_like(ref, dtype=torch.float8_e4m3fn)
scale = torch.empty_like(expected_scale)
rocm_inverse_rope_mxfp8_rows(
o, positions, cos_sin_cache, ROPE_HEAD_DIM, data, scale
)
assert torch.equal(scale, expected_scale)
# Within half an e4m3 step at the top of each block's range.
half_step = 16.0 * torch.exp2(scale.float() - 127.0).repeat_interleave(32, -1)
assert bool(((_mxfp8_dequant(data, scale) - ref).abs() <= half_step).all())
@requires_gfx950
@torch.inference_mode()
def test_sparse_attn_decode_mxfp8_output() -> None:
"""``out_mxfp8`` routes the real decode through the MXFP8 epilogue."""
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
_rocm_sparse_attn_decode_ragged_triton,
)
device = torch.device("cuda")
torch.manual_seed(1)
num_heads = 16
q = torch.randn(2, num_heads, HEAD_DIM, device=device).bfloat16() * 0.125
kv = torch.randn(6, HEAD_DIM, device=device).bfloat16() * 0.125
decode = functools.partial(
_rocm_sparse_attn_decode_ragged_triton,
q=q,
main_cache=_pack_fp8_ds_mla_cache(kv, 4, current_platform.is_fp8_fnuz()),
main_indices=torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device),
main_indptr=torch.tensor([0, 2, 4], dtype=torch.int32, device=device),
scale=HEAD_DIM**-0.5,
attn_sink=None,
nope_head_dim=NOPE_HEAD_DIM,
rope_head_dim=ROPE_HEAD_DIM,
inv_rope_positions=torch.tensor([3, 9], device=device),
inv_rope_cos_sin_cache=_random_cos_sin_cache(16, device),
)
expected = decode().float().view(2, -1)
data = torch.empty(
2, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
)
scale = torch.empty(2, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device)
decode(out_mxfp8=(data, scale))
torch.testing.assert_close(
_mxfp8_dequant(data, scale), expected, atol=1e-3, rtol=2**-4
)
@requires_gfx950
@pytest.mark.parametrize(
"num_tokens, n_groups",
# One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles.
[
(1, 4),
(20, 4),
(48, 4),
(77, 4),
(130, 4),
(300, 4),
(700, 4),
(1000, 4),
(1100, 8),
],
)
@torch.inference_mode()
def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
_mxfp8_e4m3_quantize_torch,
)
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_mxfp8_wo_a_bmm
device = torch.device("cuda")
torch.manual_seed(3)
o_lora_rank, group_dim = 128, 2048
a, a_scale = _mxfp8_e4m3_quantize_torch(
torch.randn(num_tokens, n_groups * group_dim, device=device)
)
w, w_scale = _mxfp8_e4m3_quantize_torch(
torch.randn(n_groups * o_lora_rank, group_dim, device=device)
)
wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
out = rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
expected = torch.einsum(
"tgd,grd->tgr",
_mxfp8_dequant(a, a_scale).view(num_tokens, n_groups, group_dim),
_mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim),
)
torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
@requires_gfx950
@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
@torch.inference_mode()
vllm/models/deepseek_v41/amd/rocm.py
+136 -22
View File
@@ -13,6 +13,8 @@ from vllm.distributed import (
)
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
from vllm.models.deepseek_v41.attention import (
DeepseekV4Attention,
_replace_layer_index,
@@ -37,7 +39,9 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
build_ragged_indices_from_dense,
rocm_inv_rope_einsum,
rocm_inverse_rope_mxfp8_rows,
rocm_inverse_rope_rows_,
rocm_mxfp8_wo_a_bmm,
rocm_sparse_attn_decode,
rocm_sparse_attn_prefill,
)
@@ -663,7 +667,40 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
transpose_scale=False,
)
def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
def _alloc_attn_out(
self, num_tokens: int, hidden_states: torch.Tensor
) -> torch.Tensor | QuantizedActivation:
if not _ON_GFX950:
return super()._alloc_attn_out(num_tokens, hidden_states)
# wo_a's MXFP8 input: the decode reduce writes it directly, prefill
# rows are rotated and quantized after their bf16 attention.
width = self.n_local_heads * self.head_dim
data = torch.empty(
(num_tokens, width), dtype=torch.float8_e4m3fn, device=hidden_states.device
)
scale = torch.empty(
(num_tokens, width // 32), dtype=torch.uint8, device=hidden_states.device
)
return QuantizedActivation(
data=data,
scale=scale,
orig_dtype=hidden_states.dtype,
orig_shape=data.shape,
quant_key=kMxfp8Dynamic,
)
def _o_proj(
self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor
) -> torch.Tensor:
if isinstance(attn_out, QuantizedActivation):
z = rocm_mxfp8_wo_a_bmm(
attn_out.data,
attn_out.scale,
self.wo_a,
self.n_local_groups,
self.o_lora_rank,
)
return self._wo_b_after_wo_a(z)
o = attn_out[:, : self.n_local_heads, :]
# ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
z = rocm_inv_rope_einsum(
@@ -676,26 +713,29 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
self.wo_a,
inverse_rope=False,
)
zf = z.flatten(1)
return self._wo_b_after_wo_a(z.flatten(1))
def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor:
if self._wo_b_scale is not None and zf.dim() == 2:
result = self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
else:
result = self.wo_b(zf)
return result
return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
return self.wo_b(zf)
def forward_mqa(
self,
q: torch.Tensor,
kv: torch.Tensor,
positions: torch.Tensor,
output: torch.Tensor,
output: torch.Tensor | QuantizedActivation,
) -> None:
assert output.shape == q.shape, (
f"output buffer shape {output.shape} must match q shape {q.shape}"
)
assert output.dtype == q.dtype, (
f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
)
mxfp8_out = output if isinstance(output, QuantizedActivation) else None
if mxfp8_out is None:
assert isinstance(output, torch.Tensor)
assert output.shape == q.shape, (
f"output buffer shape {output.shape} must match q shape {q.shape}"
)
assert output.dtype == q.dtype, (
f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
)
forward_context = get_forward_context()
attn_metadata = forward_context.attn_metadata
@@ -713,9 +753,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
)
M = N + self.window_size + self.max_num_batched_tokens
current_workspace_manager().get_simultaneous(
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
*self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q)
)
output.zero_()
if mxfp8_out is None:
output.zero_()
else:
mxfp8_out.data.zero_()
mxfp8_out.scale.zero_()
return
assert isinstance(attn_metadata, dict)
@@ -745,10 +789,41 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
positions=positions[num_decode_tokens:],
compressed_k_cache=self_kv_cache,
swa_k_cache=swa_kv_cache,
output=output[num_decode_tokens:],
output=(
output[num_decode_tokens:]
if isinstance(output, torch.Tensor)
else None
),
attn_metadata=rocm_metadata,
swa_metadata=swa_metadata,
mxfp8_out=(
None
if mxfp8_out is None
else (
mxfp8_out.data[num_decode_tokens:],
mxfp8_out.scale[num_decode_tokens:],
)
),
)
if mxfp8_out is not None:
# The reduce epilogue rotates and quantizes every decode row and
# the prefill path already did its own, so nothing is owed here.
if num_decodes > 0:
self._forward_decode(
q=q[:num_decode_tokens],
positions=positions[:num_decode_tokens],
kv_cache=self_kv_cache,
swa_metadata=swa_metadata,
attn_metadata=rocm_metadata,
swa_only=swa_only,
output=None,
output_mxfp8=(
mxfp8_out.data[:num_decode_tokens],
mxfp8_out.scale[:num_decode_tokens],
),
)
return
assert isinstance(output, torch.Tensor)
rotated = 0
if num_decodes > 0:
rotated = self._forward_decode(
@@ -816,7 +891,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_only: bool,
output: torch.Tensor,
output: torch.Tensor | None,
output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> int:
"""Returns how many leading rows the decode epilogue inverse-RoPE'd."""
num_decodes = swa_metadata.num_decodes
@@ -856,6 +932,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
nope_head_dim=self.nope_head_dim,
rope_head_dim=self.rope_head_dim,
output=output,
output_mxfp8=output_mxfp8,
inv_rope_positions=positions,
inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
@@ -871,10 +948,16 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
positions: torch.Tensor,
compressed_k_cache: torch.Tensor | None,
swa_k_cache: torch.Tensor,
output: torch.Tensor,
output: torch.Tensor | None,
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> None:
"""Sparse prefill into bf16 ``output``, or into ``mxfp8_out``.
With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
inverse-RoPE'd and quantized into it in one pass.
"""
swa_only = attn_metadata is None
num_prefills = swa_metadata.num_prefills
@@ -910,10 +993,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
self.PREFILL_CHUNK_SIZE
)
workspace_manager = current_workspace_manager()
kv = workspace_manager.get_simultaneous(
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
)[0]
workspace = current_workspace_manager().get_simultaneous(
*self._prefill_workspace_shapes(M, num_prefill_tokens, q)
)
kv = workspace[0]
if mxfp8_out is not None:
output = workspace[1]
assert output is not None
for chunk_idx in range(num_chunks):
chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
@@ -978,3 +1064,31 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
attn_sink=self.attn_sink,
output=output[query_start:query_end],
)
if mxfp8_out is not None:
rocm_inverse_rope_mxfp8_rows(
output,
positions[:num_prefill_tokens],
self.rotary_emb.cos_sin_cache,
self.rope_head_dim,
mxfp8_out[0][:num_prefill_tokens],
mxfp8_out[1][:num_prefill_tokens],
)
def _prefill_workspace_shapes(
self, M: int, num_prefill_tokens: int, q: torch.Tensor
) -> list[tuple[tuple[int, ...], torch.dtype]]:
"""The prefill gather buffer, plus bf16 rows for the MXFP8 output.
One ``get_simultaneous`` call: separate calls alias the same memory.
"""
shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
]
if _ON_GFX950:
shapes.append(
(
(num_prefill_tokens, self.n_local_heads, self.head_dim),
torch.bfloat16,
)
)
return shapes
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+351 -5
View File
@@ -16,7 +16,7 @@ from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.torch_utils import LayerNameType
from vllm.utils.torch_utils import LayerNameType, direct_register_custom_op
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
from vllm.v1.worker.workspace import current_workspace_manager
@@ -1637,6 +1637,90 @@ def rocm_inverse_rope_rows_(
_fused_inverse_rope_gptj(o, positions, cos_sin_cache, rope_head_dim, out=o)
@triton.jit
def _inverse_rope_mxfp8_quant_kernel(
o_ptr, # [T, H, D] bf16
q_ptr, # [T, H * D] e4m3
s_ptr, # [T, H * D // 32] uint8 (E8M0)
pos_ptr,
cos_sin_ptr,
s_t,
s_h,
qs_t,
ss_t,
cs_stride,
HEAD_DIM: tl.constexpr,
NOPE: tl.constexpr,
HALF: tl.constexpr,
):
t = tl.program_id(0)
h = tl.program_id(1)
lanes = tl.arange(0, HEAD_DIM)
x = tl.load(o_ptr + t * s_t + h * s_h + lanes).to(tl.float32)
# Same lane-pair formulation as the reduce epilogue: NoPE pairs take
# cos=1 / sin=0, so one expression rotates the whole row.
pos = tl.load(pos_ptr + t).to(tl.int64)
pair_idx = tl.arange(0, HEAD_DIM // 2) - (NOPE // 2)
is_rope = pair_idx >= 0
k = tl.where(is_rope, pair_idx, 0)
cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
even, odd = tl.split(tl.reshape(x, (HEAD_DIM // 2, 2)))
x = tl.reshape(
tl.join(even * cos + odd * sin, odd * cos - even * sin), (1, HEAD_DIM)
)
xq, bits = _mxfp8_quantize_rows(x, 1, HEAD_DIM)
tl.store(
q_ptr + t * qs_t + h * HEAD_DIM + lanes,
tl.reshape(xq, (HEAD_DIM,)).to(q_ptr.dtype.element_ty),
)
scale_offsets = tl.arange(0, HEAD_DIM // 32)
tl.store(
s_ptr + t * ss_t + h * (HEAD_DIM // 32) + scale_offsets,
tl.reshape(bits, (HEAD_DIM // 32,)).to(tl.uint8),
)
def rocm_inverse_rope_mxfp8_rows(
o: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
rope_head_dim: int,
out_data: torch.Tensor,
out_scale: torch.Tensor,
) -> None:
"""Inverse-RoPE bf16 attention rows and MXFP8-quantize them for wo_a.
The counterpart of ``rocm_inverse_rope_rows_`` for layers whose attention
output is MXFP8: rows the decode reduce did not emit (prefill) go through
here. ``o`` is [T, H, D]; ``out_data`` [T, H * D] e4m3 and ``out_scale``
[T, H * D // 32] E8M0, the layout the reduce epilogue writes.
"""
num_tokens, num_heads, head_dim = o.shape
if num_tokens == 0:
return
assert o.stride(-1) == 1 and out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
assert out_data.shape == (num_tokens, num_heads * head_dim)
assert out_scale.shape == (num_tokens, num_heads * head_dim // 32)
assert cos_sin_cache.shape[-1] == rope_head_dim
_inverse_rope_mxfp8_quant_kernel[(num_tokens, num_heads)](
o,
out_data,
out_scale,
positions,
cos_sin_cache,
o.stride(0),
o.stride(1),
out_data.stride(0),
out_scale.stride(0),
cos_sin_cache.stride(0),
HEAD_DIM=head_dim,
NOPE=head_dim - rope_head_dim,
HALF=rope_head_dim // 2,
num_warps=4,
)
def _get_cached_wo_a_bf16(
wo_a: torch.nn.Module,
n_local_groups: int,
@@ -1722,6 +1806,175 @@ def rocm_inv_rope_einsum(
return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
@triton.jit
def _mxfp8_wo_a_bmm_kernel(
a_ptr, # [T, G * K] e4m3
as_ptr, # [T, G * K // 32] E8M0
w_ptr, # [G * R, K] e4m3
ws_ptr, # [G * R, K // 32] E8M0
out_ptr, # [T, G * R]
num_tokens,
stride_at,
stride_ast,
stride_wn,
stride_wsn,
stride_out,
R: tl.constexpr,
K: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
BLOCK_K: tl.constexpr,
):
# Output columns are group-major, so a BLOCK_N tile never straddles two
# groups (R % BLOCK_N == 0) and picks its group's K slice of A.
pid_n = tl.program_id(0)
pid_m = tl.program_id(1)
group = (pid_n * BLOCK_N) // R
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
m_mask = offs_m < num_tokens
offs_k = tl.arange(0, BLOCK_K)
offs_sk = tl.arange(0, BLOCK_K // 32)
a_ptrs = a_ptr + offs_m[:, None] * stride_at + (group * K + offs_k)[None, :]
as_ptrs = (
as_ptr + offs_m[:, None] * stride_ast + (group * (K // 32) + offs_sk)[None, :]
)
w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :]
ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :]
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for _ in range(K // BLOCK_K):
a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
a_s = tl.load(as_ptrs, mask=m_mask[:, None], other=127)
w = tl.load(w_ptrs)
w_s = tl.load(ws_ptrs)
acc += tl.dot_scaled(a, a_s, "e4m3", w.T, w_s, "e4m3")
a_ptrs += BLOCK_K
as_ptrs += BLOCK_K // 32
w_ptrs += BLOCK_K
ws_ptrs += BLOCK_K // 32
tl.store(
out_ptr + offs_m[:, None] * stride_out + offs_n[None, :],
acc.to(out_ptr.dtype.element_ty),
mask=m_mask[:, None],
)
def _mxfp8_wo_a_bmm_config(num_tokens: int, n_groups: int) -> tuple[int, ...]:
"""(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx950.
Tuned under HIP graphs with a cold weight at G = 4 and 2, over every
decode shape of conc 1-128 x 0-5 spec tokens plus prefill chunks up to
8K tokens. The best tile tracks the total work T * G, so the tiers are
keyed on it.
This will be replaced after new GEMM kernel from AITER with proper 32x32 scale
shape GEMM fp8 enabled.
"""
work = num_tokens * n_groups
if work <= 64:
return 16, 16, 1024, 2, 3
if work <= 128:
return 32, 16, 1024, 2, 3
if work <= 256:
return 32, 32, 512, 2, 3
if work <= 512:
return 64, 32, 512, 2, 3
if work <= 1024:
return 64, 64, 512, 4, 2
if work <= 2048:
return 64, 64, 256, 4, 2
if work <= 3072:
return 64, 64, 256, 2, 1
if work <= 4096:
return 128, 128, 256, 8, 2
return 128, 128, 128, 4, 2
def _rocm_mxfp8_wo_a_bmm_impl(
a: torch.Tensor,
a_scale: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
n_groups: int,
o_lora_rank: int,
) -> torch.Tensor:
num_tokens = a.shape[0]
group_dim = a.shape[1] // n_groups
out = torch.empty(
(num_tokens, n_groups * o_lora_rank), dtype=torch.bfloat16, device=a.device
)
if num_tokens == 0:
return out
block_m, block_n, block_k, num_warps, num_stages = _mxfp8_wo_a_bmm_config(
num_tokens, n_groups
)
grid = (n_groups * o_lora_rank // block_n, triton.cdiv(num_tokens, block_m))
_mxfp8_wo_a_bmm_kernel[grid](
a,
a_scale,
weight,
weight_scale,
out,
num_tokens,
a.stride(0),
a_scale.stride(0),
weight.stride(0),
weight_scale.stride(0),
out.stride(0),
R=o_lora_rank,
K=group_dim,
BLOCK_M=block_m,
BLOCK_N=block_n,
BLOCK_K=block_k,
num_warps=num_warps,
num_stages=num_stages,
)
return out
def _rocm_mxfp8_wo_a_bmm_fake(
a: torch.Tensor,
a_scale: torch.Tensor,
weight: torch.Tensor,
weight_scale: torch.Tensor,
n_groups: int,
o_lora_rank: int,
) -> torch.Tensor:
return a.new_empty((a.shape[0], n_groups * o_lora_rank), dtype=torch.bfloat16)
# An opaque op: the tile choice branches on the token count, which the
# compiled o_proj must not freeze at its trace-time value.
direct_register_custom_op(
op_name="rocm_dsv41_mxfp8_wo_a_bmm",
op_func=_rocm_mxfp8_wo_a_bmm_impl,
fake_impl=_rocm_mxfp8_wo_a_bmm_fake,
)
def rocm_mxfp8_wo_a_bmm(
a: torch.Tensor,
a_scale: torch.Tensor,
wo_a: torch.nn.Module,
n_groups: int,
o_lora_rank: int,
) -> torch.Tensor:
"""Grouped MXFP8 wo_a: ``out[t, g, :] = a[t, g, :] @ W[g].T``, bf16 out.
``a`` is the [T, G * K] e4m3 attention output and ``a_scale`` its
[T, G * K // 32] E8M0 scales, as the sparse decode reduce writes them.
The weight is the checkpoint's MXFP8 ``wo_a`` as loaded, [G * R, K] with
[G * R, K // 32] scales, so there is no dequantized copy to keep.
Returns [T, G * R].
"""
return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(
a, a_scale, wo_a.weight, wo_a.weight_scale, n_groups, o_lora_rank
)
_DSV4_SPARSE_NOPE_DIM = 448
_DSV4_SPARSE_ROPE_DIM = 64
@@ -2922,6 +3175,33 @@ def _sparse_attn_decode_gfx950_partial_kernel(
)
@triton.jit
def _mxfp8_scale_bits(amax):
"""Biased E8M0 exponent that puts ``amax`` at the top of the e4m3 range.
Same rounding as ``mxfp8_e4m3_quantize``, so the output is bit-identical to
quantizing the tensor there.
"""
amax = tl.maximum(amax, 1.1754943508222875e-38)
bits = tl.ceil(tl.log2(amax / 448.0)) + 127.0
return tl.minimum(tl.maximum(bits, 0.0), 254.0)
@triton.jit
def _mxfp8_quantize_rows(x, ROWS: tl.constexpr, COLS: tl.constexpr):
"""MXFP8-quantize ``x`` [ROWS, COLS] in registers, one scale per 32 lanes.
Returns the rescaled fp32 values (to be cast to e4m3 on store) and the
[ROWS, COLS // 32] biased E8M0 exponents.
"""
blocks = tl.reshape(x, (ROWS, COLS // 32, 32))
bits = _mxfp8_scale_bits(tl.max(tl.abs(blocks), axis=2))
# Multiply by the reciprocal: a divisor of 2**-127 would be subnormal and
# flush to zero, turning an all-zero block into NaN.
q = blocks * tl.exp2(127.0 - bits)[:, :, None]
return tl.reshape(q, (ROWS, COLS)), bits
@triton.jit
def _sparse_attn_decode_reduce_kernel(
part_m_ptr,
@@ -2931,8 +3211,11 @@ def _sparse_attn_decode_reduce_kernel(
out_ptr,
pos_ptr,
cos_sin_ptr,
out_scale_ptr,
out_stride0,
out_stride1,
os_stride0,
os_stride1,
pm_stride0,
pm_stride_s,
pa_stride0,
@@ -2949,6 +3232,7 @@ def _sparse_attn_decode_reduce_kernel(
FUSE_INV_ROPE: tl.constexpr,
NOPE: tl.constexpr,
HALF: tl.constexpr,
QUANT_OUT: tl.constexpr,
):
query_idx = tl.program_id(0)
pid_h = tl.program_id(1)
@@ -3061,9 +3345,22 @@ def _sparse_attn_decode_reduce_kernel(
out_row_ptr = (
out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
)
if QUANT_OUT:
# Emit wo_a's MXFP8 input directly instead of a bf16 row that a
# separate pass would read back to quantize.
out, scale_bits = _mxfp8_quantize_rows(out, BLOCK_H, COMB_DIM)
scale_offsets = tl.arange(0, COMB_DIM // 32)
tl.store(
out_scale_ptr
+ query_idx * os_stride0
+ head_offsets[:, None] * os_stride1
+ scale_offsets[None, :],
scale_bits.to(tl.uint8),
mask=head_mask[:, None],
)
tl.store(
out_row_ptr + comb_offsets[None, :],
out,
out.to(out_ptr.dtype.element_ty),
mask=head_mask[:, None],
)
@@ -3380,7 +3677,14 @@ def _rocm_sparse_attn_decode_ragged_triton(
adaptive_splits: bool = False,
inv_rope_positions: torch.Tensor | None = None,
inv_rope_cos_sin_cache: torch.Tensor | None = None,
out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
"""Split-K sparse decode; returns the attention output.
With ``out_mxfp8 = (data, scale)`` the reduce writes MXFP8 instead of
bf16: ``data`` is [b, h * d] e4m3 and ``scale`` [b, h * d // 32] E8M0, and
``data`` viewed as [b, h, d] is returned.
"""
assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
assert main_cache.ndim == 3, (
f"expected main_cache=[blocks,block,bytes], got {main_cache.shape}"
@@ -3445,7 +3749,28 @@ def _rocm_sparse_attn_decode_ragged_triton(
extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
block_h = 16
if out is None:
out_scale = None
if out_mxfp8 is not None:
assert out is None, "out and out_mxfp8 are mutually exclusive"
assert _ON_GFX950, "the MXFP8 reduce epilogue is gfx950-only"
assert inv_rope_positions is not None, (
"the MXFP8 output feeds wo_a, so it must be inverse-RoPE'd first"
)
out_data, out_scale = out_mxfp8
assert out_data.dtype == torch.float8_e4m3fn and out_scale.dtype == (
torch.uint8
), f"expected e4m3/uint8 MXFP8 buffers, got {out_data.dtype}/{out_scale.dtype}"
assert out_data.shape == (num_queries, num_heads * head_dim), (
f"expected MXFP8 data [{num_queries}, {num_heads * head_dim}], "
f"got {tuple(out_data.shape)}"
)
assert out_scale.shape == (num_queries, num_heads * head_dim // 32), (
f"expected MXFP8 scale [{num_queries}, {num_heads * head_dim // 32}], "
f"got {tuple(out_scale.shape)}"
)
assert out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
out = out_data.view(num_queries, num_heads, head_dim)
elif out is None:
out = torch.empty_like(q, dtype=torch.bfloat16)
else:
assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}"
@@ -3643,8 +3968,11 @@ def _rocm_sparse_attn_decode_ragged_triton(
out,
inv_rope_positions,
inv_rope_cos_sin_cache,
out_scale,
out.stride(0),
out.stride(1),
out_scale.stride(0) if out_scale is not None else 0,
head_dim // 32,
part_m.stride(0),
part_m.stride(1),
part_acc.stride(0),
@@ -3661,6 +3989,7 @@ def _rocm_sparse_attn_decode_ragged_triton(
FUSE_INV_ROPE=fuse_inv_rope,
NOPE=nope_head_dim,
HALF=rope_head_dim // 2,
QUANT_OUT=out_scale is not None,
num_warps=4,
)
return out
@@ -3687,6 +4016,7 @@ def _rocm_sparse_attn_decode_triton(
adaptive_splits: bool = False,
inv_rope_positions: torch.Tensor | None = None,
inv_rope_cos_sin_cache: torch.Tensor | None = None,
out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
if main_ragged_indices is None or main_ragged_indptr is None:
main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
@@ -3727,6 +4057,7 @@ def _rocm_sparse_attn_decode_triton(
adaptive_splits=adaptive_splits,
inv_rope_positions=inv_rope_positions,
inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
out_mxfp8=out_mxfp8,
)
@@ -3825,11 +4156,12 @@ def rocm_sparse_attn_decode(
head_dim: int,
nope_head_dim: int,
rope_head_dim: int,
output: torch.Tensor,
output: torch.Tensor | None,
extra_cache_nan_free: bool = False,
adaptive_splits: bool = False,
inv_rope_positions: torch.Tensor | None = None,
inv_rope_cos_sin_cache: torch.Tensor | None = None,
output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> int:
"""Run sparse MLA decode into ``output``.
@@ -3837,6 +4169,11 @@ def rocm_sparse_attn_decode(
epilogue. Returns how many leading rows of ``output`` came back rotated,
so a caller mixing in a decode path that does not fuse still knows what it
owes the standalone pass. Read it from the eager attention segment only.
``output_mxfp8 = (data, scale)`` replaces ``output``: the reduce also
MXFP8-quantizes the rotated rows for the FP8 wo_a (see
``_rocm_sparse_attn_decode_ragged_triton``). It needs gfx950 and the fused
inverse RoPE, and always covers every row.
"""
assert swa_k_cache.dtype == torch.uint8, (
"ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
@@ -3866,7 +4203,12 @@ def rocm_sparse_attn_decode(
if topk_indices is not None:
extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
if output_mxfp8 is not None:
assert output is None, "output and output_mxfp8 are mutually exclusive"
direct_out = None
else:
assert output is not None
direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
attn_out = _rocm_sparse_attn_decode_triton(
q=q,
main_cache=swa_k_cache,
@@ -3888,7 +4230,11 @@ def rocm_sparse_attn_decode(
adaptive_splits=adaptive_splits,
inv_rope_positions=inv_rope_positions,
inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
out_mxfp8=output_mxfp8,
)
if output_mxfp8 is not None:
return q.shape[0]
assert output is not None
if direct_out is None:
output.copy_(attn_out.to(output.dtype))
return output.shape[0] if inv_rope_positions is not None else 0