@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 721d0e5c11

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

Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Fangzhou Ai committed Sep 24, 2026 at 07:58 UTC 721d0e5c112a2e4e0bed9078fc8d953eaaf5abc2
3 files changed +664 -29
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+177 -2
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import functools
5 from types import SimpleNamespace
6
7 import pytest
@@ -261,14 +262,19 @@ def _launch_sparse_decode_reduce(
262 adaptive_splits: bool,
263 positions: torch.Tensor | None = None,
264 cos_sin_cache: torch.Tensor | None = None,
265 + out_dtype: torch.dtype = torch.bfloat16,
266 + out_scale: torch.Tensor | None = None,
267 ) -> torch.Tensor:
265 - """Run the reduce kernel, with the inverse-RoPE epilogue if given positions."""
268 + """Run the reduce kernel, with the inverse-RoPE epilogue if given positions.
269 +
270 + Passing ``out_scale`` selects the MXFP8 epilogue (``out_dtype`` e4m3).
271 + """
272 from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
273
274 num_queries, num_splits, num_heads = part_m.shape
275 out = torch.empty(
276 (num_queries, num_heads, HEAD_DIM),
271 - dtype=torch.bfloat16,
277 + dtype=out_dtype,
278 device=part_m.device,
279 )
280 attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device)
@@ -280,8 +286,11 @@ def _launch_sparse_decode_reduce(
286 out,
287 positions,
288 cos_sin_cache,
289 + out_scale,
290 out.stride(0),
291 out.stride(1),
292 + out_scale.stride(0) if out_scale is not None else 0,
293 + HEAD_DIM // 32,
294 part_m.stride(0),
295 part_m.stride(1),
296 part_acc.stride(0),
@@ -298,6 +307,7 @@ def _launch_sparse_decode_reduce(
307 FUSE_INV_ROPE=positions is not None,
308 NOPE=NOPE_HEAD_DIM,
309 HALF=ROPE_HEAD_DIM // 2,
310 + QUANT_OUT=out_scale is not None,
311 num_warps=4,
312 )
313 return out
@@ -1057,6 +1067,171 @@ def test_sparse_attn_decode_reduce_inverse_rope_epilogue(adaptive_splits: bool)
1067 assert not torch.equal(actual[..., NOPE_HEAD_DIM:], unfused[..., NOPE_HEAD_DIM:])
1068
1069
1070 +def _mxfp8_dequant(data: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
1071 + blocks = data.float().view(*data.shape[:-1], -1, 32)
1072 + return (blocks * torch.exp2(scale.float() - 127.0)[..., None]).view(data.shape)
1073 +
1074 +
1075 +def _random_cos_sin_cache(max_pos: int, device: torch.device) -> torch.Tensor:
1076 + angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device)
1077 + return torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous()
1078 +
1079 +
1080 +@requires_gfx950
1081 +@torch.inference_mode()
1082 +def test_sparse_attn_decode_reduce_mxfp8_epilogue() -> None:
1083 + """QUANT_OUT must equal MXFP8-quantizing the reduce's fp32 output."""
1084 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1085 + _mxfp8_e4m3_quantize_torch,
1086 + )
1087 +
1088 + device = torch.device("cuda")
1089 + torch.manual_seed(7)
1090 + num_queries, num_splits, num_heads = 5, 4, 8
1091 + shape = (num_queries, num_splits, num_heads)
1092 + part_m = torch.randn(shape, device=device)
1093 + part_l = torch.rand(shape, device=device) + 0.5
1094 + part_acc = torch.randn((*shape, HEAD_DIM), device=device)
1095 + positions = torch.randint(0, 64, (num_queries,), device=device)
1096 + cos_sin_cache = _random_cos_sin_cache(64, device)
1097 + args = (part_m, part_l, part_acc, False, positions, cos_sin_cache)
1098 +
1099 + rotated = _launch_sparse_decode_reduce(*args, out_dtype=torch.float32)
1100 + expected_data, expected_scale = _mxfp8_e4m3_quantize_torch(
1101 + rotated.view(num_queries, -1)
1102 + )
1103 + scale = torch.empty_like(expected_scale)
1104 + data = _launch_sparse_decode_reduce(
1105 + *args, out_dtype=torch.float8_e4m3fn, out_scale=scale
1106 + )
1107 +
1108 + assert torch.equal(scale, expected_scale)
1109 + assert torch.equal(
1110 + data.view(num_queries, -1).view(torch.uint8), expected_data.view(torch.uint8)
1111 + )
1112 +
1113 +
1114 +@requires_gfx950
1115 +@torch.inference_mode()
1116 +def test_inverse_rope_mxfp8_rows() -> None:
1117 + """Prefill rows get the same inverse RoPE + MXFP8 as the decode epilogue."""
1118 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1119 + _mxfp8_e4m3_quantize_torch,
1120 + )
1121 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
1122 + rocm_inverse_rope_mxfp8_rows,
1123 + )
1124 +
1125 + device = torch.device("cuda")
1126 + torch.manual_seed(5)
1127 + num_tokens, num_heads = 7, 4
1128 + o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).bfloat16()
1129 + positions = torch.randint(0, 64, (num_tokens,), device=device)
1130 + cos_sin_cache = _random_cos_sin_cache(64, device)
1131 +
1132 + ref = o.float()
1133 + cos, sin = cos_sin_cache[positions, None].chunk(2, dim=-1)
1134 + even, odd = ref[..., NOPE_HEAD_DIM::2].clone(), ref[..., NOPE_HEAD_DIM + 1 :: 2]
1135 + ref[..., NOPE_HEAD_DIM::2] = even * cos + odd * sin
1136 + ref[..., NOPE_HEAD_DIM + 1 :: 2] = odd * cos - even * sin
1137 + ref = ref.view(num_tokens, -1)
1138 + _, expected_scale = _mxfp8_e4m3_quantize_torch(ref)
1139 +
1140 + data = torch.empty_like(ref, dtype=torch.float8_e4m3fn)
1141 + scale = torch.empty_like(expected_scale)
1142 + rocm_inverse_rope_mxfp8_rows(
1143 + o, positions, cos_sin_cache, ROPE_HEAD_DIM, data, scale
1144 + )
1145 +
1146 + assert torch.equal(scale, expected_scale)
1147 + # Within half an e4m3 step at the top of each block's range.
1148 + half_step = 16.0 * torch.exp2(scale.float() - 127.0).repeat_interleave(32, -1)
1149 + assert bool(((_mxfp8_dequant(data, scale) - ref).abs() <= half_step).all())
1150 +
1151 +
1152 +@requires_gfx950
1153 +@torch.inference_mode()
1154 +def test_sparse_attn_decode_mxfp8_output() -> None:
1155 + """``out_mxfp8`` routes the real decode through the MXFP8 epilogue."""
1156 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
1157 + _rocm_sparse_attn_decode_ragged_triton,
1158 + )
1159 +
1160 + device = torch.device("cuda")
1161 + torch.manual_seed(1)
1162 + num_heads = 16
1163 + q = torch.randn(2, num_heads, HEAD_DIM, device=device).bfloat16() * 0.125
1164 + kv = torch.randn(6, HEAD_DIM, device=device).bfloat16() * 0.125
1165 + decode = functools.partial(
1166 + _rocm_sparse_attn_decode_ragged_triton,
1167 + q=q,
1168 + main_cache=_pack_fp8_ds_mla_cache(kv, 4, current_platform.is_fp8_fnuz()),
1169 + main_indices=torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device),
1170 + main_indptr=torch.tensor([0, 2, 4], dtype=torch.int32, device=device),
1171 + scale=HEAD_DIM**-0.5,
1172 + attn_sink=None,
1173 + nope_head_dim=NOPE_HEAD_DIM,
1174 + rope_head_dim=ROPE_HEAD_DIM,
1175 + inv_rope_positions=torch.tensor([3, 9], device=device),
1176 + inv_rope_cos_sin_cache=_random_cos_sin_cache(16, device),
1177 + )
1178 +
1179 + expected = decode().float().view(2, -1)
1180 + data = torch.empty(
1181 + 2, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device
1182 + )
1183 + scale = torch.empty(2, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device)
1184 + decode(out_mxfp8=(data, scale))
1185 +
1186 + torch.testing.assert_close(
1187 + _mxfp8_dequant(data, scale), expected, atol=1e-3, rtol=2**-4
1188 + )
1189 +
1190 +
1191 +@requires_gfx950
1192 +@pytest.mark.parametrize(
1193 + "num_tokens, n_groups",
1194 + # One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles.
1195 + [
1196 + (1, 4),
1197 + (20, 4),
1198 + (48, 4),
1199 + (77, 4),
1200 + (130, 4),
1201 + (300, 4),
1202 + (700, 4),
1203 + (1000, 4),
1204 + (1100, 8),
1205 + ],
1206 +)
1207 +@torch.inference_mode()
1208 +def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None:
1209 + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1210 + _mxfp8_e4m3_quantize_torch,
1211 + )
1212 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import rocm_mxfp8_wo_a_bmm
1213 +
1214 + device = torch.device("cuda")
1215 + torch.manual_seed(3)
1216 + o_lora_rank, group_dim = 128, 2048
1217 + a, a_scale = _mxfp8_e4m3_quantize_torch(
1218 + torch.randn(num_tokens, n_groups * group_dim, device=device)
1219 + )
1220 + w, w_scale = _mxfp8_e4m3_quantize_torch(
1221 + torch.randn(n_groups * o_lora_rank, group_dim, device=device)
1222 + )
1223 + wo_a = SimpleNamespace(weight=w, weight_scale=w_scale)
1224 +
1225 + out = rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank)
1226 +
1227 + expected = torch.einsum(
1228 + "tgd,grd->tgr",
1229 + _mxfp8_dequant(a, a_scale).view(num_tokens, n_groups, group_dim),
1230 + _mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim),
1231 + )
1232 + torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2)
1233 +
1234 +
1235 @requires_gfx950
1236 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65])
1237 @torch.inference_mode()
vllm/models/deepseek_v41/amd/rocm.py
+136 -22
@@ -13,6 +13,8 @@ from vllm.distributed import (
13 )
14 from vllm.forward_context import get_forward_context
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
17 +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
18 from vllm.models.deepseek_v41.attention import (
19 DeepseekV4Attention,
20 _replace_layer_index,
@@ -37,7 +39,9 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
39 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
40 build_ragged_indices_from_dense,
41 rocm_inv_rope_einsum,
42 + rocm_inverse_rope_mxfp8_rows,
43 rocm_inverse_rope_rows_,
44 + rocm_mxfp8_wo_a_bmm,
45 rocm_sparse_attn_decode,
46 rocm_sparse_attn_prefill,
47 )
@@ -663,7 +667,40 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
667 transpose_scale=False,
668 )
669
666 - def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
670 + def _alloc_attn_out(
671 + self, num_tokens: int, hidden_states: torch.Tensor
672 + ) -> torch.Tensor | QuantizedActivation:
673 + if not _ON_GFX950:
674 + return super()._alloc_attn_out(num_tokens, hidden_states)
675 + # wo_a's MXFP8 input: the decode reduce writes it directly, prefill
676 + # rows are rotated and quantized after their bf16 attention.
677 + width = self.n_local_heads * self.head_dim
678 + data = torch.empty(
679 + (num_tokens, width), dtype=torch.float8_e4m3fn, device=hidden_states.device
680 + )
681 + scale = torch.empty(
682 + (num_tokens, width // 32), dtype=torch.uint8, device=hidden_states.device
683 + )
684 + return QuantizedActivation(
685 + data=data,
686 + scale=scale,
687 + orig_dtype=hidden_states.dtype,
688 + orig_shape=data.shape,
689 + quant_key=kMxfp8Dynamic,
690 + )
691 +
692 + def _o_proj(
693 + self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor
694 + ) -> torch.Tensor:
695 + if isinstance(attn_out, QuantizedActivation):
696 + z = rocm_mxfp8_wo_a_bmm(
697 + attn_out.data,
698 + attn_out.scale,
699 + self.wo_a,
700 + self.n_local_groups,
701 + self.o_lora_rank,
702 + )
703 + return self._wo_b_after_wo_a(z)
704 o = attn_out[:, : self.n_local_heads, :]
705 # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
706 z = rocm_inv_rope_einsum(
@@ -676,26 +713,29 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
713 self.wo_a,
714 inverse_rope=False,
715 )
679 - zf = z.flatten(1)
716 + return self._wo_b_after_wo_a(z.flatten(1))
717 +
718 + def _wo_b_after_wo_a(self, zf: torch.Tensor) -> torch.Tensor:
719 if self._wo_b_scale is not None and zf.dim() == 2:
681 - result = self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
682 - else:
683 - result = self.wo_b(zf)
684 - return result
720 + return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
721 + return self.wo_b(zf)
722
723 def forward_mqa(
724 self,
725 q: torch.Tensor,
726 kv: torch.Tensor,
727 positions: torch.Tensor,
691 - output: torch.Tensor,
728 + output: torch.Tensor | QuantizedActivation,
729 ) -> None:
693 - assert output.shape == q.shape, (
694 - f"output buffer shape {output.shape} must match q shape {q.shape}"
695 - )
696 - assert output.dtype == q.dtype, (
697 - f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
698 - )
730 + mxfp8_out = output if isinstance(output, QuantizedActivation) else None
731 + if mxfp8_out is None:
732 + assert isinstance(output, torch.Tensor)
733 + assert output.shape == q.shape, (
734 + f"output buffer shape {output.shape} must match q shape {q.shape}"
735 + )
736 + assert output.dtype == q.dtype, (
737 + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}"
738 + )
739
740 forward_context = get_forward_context()
741 attn_metadata = forward_context.attn_metadata
@@ -713,9 +753,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
753 )
754 M = N + self.window_size + self.max_num_batched_tokens
755 current_workspace_manager().get_simultaneous(
716 - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
756 + *self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q)
757 )
718 - output.zero_()
758 + if mxfp8_out is None:
759 + output.zero_()
760 + else:
761 + mxfp8_out.data.zero_()
762 + mxfp8_out.scale.zero_()
763 return
764
765 assert isinstance(attn_metadata, dict)
@@ -745,10 +789,41 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
789 positions=positions[num_decode_tokens:],
790 compressed_k_cache=self_kv_cache,
791 swa_k_cache=swa_kv_cache,
748 - output=output[num_decode_tokens:],
792 + output=(
793 + output[num_decode_tokens:]
794 + if isinstance(output, torch.Tensor)
795 + else None
796 + ),
797 attn_metadata=rocm_metadata,
798 swa_metadata=swa_metadata,
799 + mxfp8_out=(
800 + None
801 + if mxfp8_out is None
802 + else (
803 + mxfp8_out.data[num_decode_tokens:],
804 + mxfp8_out.scale[num_decode_tokens:],
805 + )
806 + ),
807 )
808 + if mxfp8_out is not None:
809 + # The reduce epilogue rotates and quantizes every decode row and
810 + # the prefill path already did its own, so nothing is owed here.
811 + if num_decodes > 0:
812 + self._forward_decode(
813 + q=q[:num_decode_tokens],
814 + positions=positions[:num_decode_tokens],
815 + kv_cache=self_kv_cache,
816 + swa_metadata=swa_metadata,
817 + attn_metadata=rocm_metadata,
818 + swa_only=swa_only,
819 + output=None,
820 + output_mxfp8=(
821 + mxfp8_out.data[:num_decode_tokens],
822 + mxfp8_out.scale[:num_decode_tokens],
823 + ),
824 + )
825 + return
826 + assert isinstance(output, torch.Tensor)
827 rotated = 0
828 if num_decodes > 0:
829 rotated = self._forward_decode(
@@ -816,7 +891,8 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
891 swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
892 attn_metadata: DeepseekV4FlashMLAMetadata | None,
893 swa_only: bool,
819 - output: torch.Tensor,
894 + output: torch.Tensor | None,
895 + output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
896 ) -> int:
897 """Returns how many leading rows the decode epilogue inverse-RoPE'd."""
898 num_decodes = swa_metadata.num_decodes
@@ -856,6 +932,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
932 nope_head_dim=self.nope_head_dim,
933 rope_head_dim=self.rope_head_dim,
934 output=output,
935 + output_mxfp8=output_mxfp8,
936 inv_rope_positions=positions,
937 inv_rope_cos_sin_cache=self.rotary_emb.cos_sin_cache,
938 extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free(
@@ -871,10 +948,16 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
948 positions: torch.Tensor,
949 compressed_k_cache: torch.Tensor | None,
950 swa_k_cache: torch.Tensor,
874 - output: torch.Tensor,
951 + output: torch.Tensor | None,
952 attn_metadata: DeepseekV4FlashMLAMetadata | None,
953 swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
954 + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
955 ) -> None:
956 + """Sparse prefill into bf16 ``output``, or into ``mxfp8_out``.
957 +
958 + With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
959 + inverse-RoPE'd and quantized into it in one pass.
960 + """
961 swa_only = attn_metadata is None
962
963 num_prefills = swa_metadata.num_prefills
@@ -910,10 +993,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
993 self.PREFILL_CHUNK_SIZE
994 )
995
913 - workspace_manager = current_workspace_manager()
914 - kv = workspace_manager.get_simultaneous(
915 - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
916 - )[0]
996 + workspace = current_workspace_manager().get_simultaneous(
997 + *self._prefill_workspace_shapes(M, num_prefill_tokens, q)
998 + )
999 + kv = workspace[0]
1000 + if mxfp8_out is not None:
1001 + output = workspace[1]
1002 + assert output is not None
1003 for chunk_idx in range(num_chunks):
1004 chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
1005 chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
@@ -978,3 +1064,31 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
1064 attn_sink=self.attn_sink,
1065 output=output[query_start:query_end],
1066 )
1067 + if mxfp8_out is not None:
1068 + rocm_inverse_rope_mxfp8_rows(
1069 + output,
1070 + positions[:num_prefill_tokens],
1071 + self.rotary_emb.cos_sin_cache,
1072 + self.rope_head_dim,
1073 + mxfp8_out[0][:num_prefill_tokens],
1074 + mxfp8_out[1][:num_prefill_tokens],
1075 + )
1076 +
1077 + def _prefill_workspace_shapes(
1078 + self, M: int, num_prefill_tokens: int, q: torch.Tensor
1079 + ) -> list[tuple[tuple[int, ...], torch.dtype]]:
1080 + """The prefill gather buffer, plus bf16 rows for the MXFP8 output.
1081 +
1082 + One ``get_simultaneous`` call: separate calls alias the same memory.
1083 + """
1084 + shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
1085 + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
1086 + ]
1087 + if _ON_GFX950:
1088 + shapes.append(
1089 + (
1090 + (num_prefill_tokens, self.n_local_heads, self.head_dim),
1091 + torch.bfloat16,
1092 + )
1093 + )
1094 + return shapes
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+351 -5
@@ -16,7 +16,7 @@ from vllm.forward_context import get_forward_context
16 from vllm.logger import init_logger
17 from vllm.platforms import current_platform
18 from vllm.triton_utils import tl, triton
19 -from vllm.utils.torch_utils import LayerNameType
19 +from vllm.utils.torch_utils import LayerNameType, direct_register_custom_op
20 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
21 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
22 from vllm.v1.worker.workspace import current_workspace_manager
@@ -1637,6 +1637,90 @@ def rocm_inverse_rope_rows_(
1637 _fused_inverse_rope_gptj(o, positions, cos_sin_cache, rope_head_dim, out=o)
1638
1639
1640 +@triton.jit
1641 +def _inverse_rope_mxfp8_quant_kernel(
1642 + o_ptr, # [T, H, D] bf16
1643 + q_ptr, # [T, H * D] e4m3
1644 + s_ptr, # [T, H * D // 32] uint8 (E8M0)
1645 + pos_ptr,
1646 + cos_sin_ptr,
1647 + s_t,
1648 + s_h,
1649 + qs_t,
1650 + ss_t,
1651 + cs_stride,
1652 + HEAD_DIM: tl.constexpr,
1653 + NOPE: tl.constexpr,
1654 + HALF: tl.constexpr,
1655 +):
1656 + t = tl.program_id(0)
1657 + h = tl.program_id(1)
1658 + lanes = tl.arange(0, HEAD_DIM)
1659 + x = tl.load(o_ptr + t * s_t + h * s_h + lanes).to(tl.float32)
1660 + # Same lane-pair formulation as the reduce epilogue: NoPE pairs take
1661 + # cos=1 / sin=0, so one expression rotates the whole row.
1662 + pos = tl.load(pos_ptr + t).to(tl.int64)
1663 + pair_idx = tl.arange(0, HEAD_DIM // 2) - (NOPE // 2)
1664 + is_rope = pair_idx >= 0
1665 + k = tl.where(is_rope, pair_idx, 0)
1666 + cos = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + k), 1.0)
1667 + sin = tl.where(is_rope, tl.load(cos_sin_ptr + pos * cs_stride + HALF + k), 0.0)
1668 + even, odd = tl.split(tl.reshape(x, (HEAD_DIM // 2, 2)))
1669 + x = tl.reshape(
1670 + tl.join(even * cos + odd * sin, odd * cos - even * sin), (1, HEAD_DIM)
1671 + )
1672 + xq, bits = _mxfp8_quantize_rows(x, 1, HEAD_DIM)
1673 + tl.store(
1674 + q_ptr + t * qs_t + h * HEAD_DIM + lanes,
1675 + tl.reshape(xq, (HEAD_DIM,)).to(q_ptr.dtype.element_ty),
1676 + )
1677 + scale_offsets = tl.arange(0, HEAD_DIM // 32)
1678 + tl.store(
1679 + s_ptr + t * ss_t + h * (HEAD_DIM // 32) + scale_offsets,
1680 + tl.reshape(bits, (HEAD_DIM // 32,)).to(tl.uint8),
1681 + )
1682 +
1683 +
1684 +def rocm_inverse_rope_mxfp8_rows(
1685 + o: torch.Tensor,
1686 + positions: torch.Tensor,
1687 + cos_sin_cache: torch.Tensor,
1688 + rope_head_dim: int,
1689 + out_data: torch.Tensor,
1690 + out_scale: torch.Tensor,
1691 +) -> None:
1692 + """Inverse-RoPE bf16 attention rows and MXFP8-quantize them for wo_a.
1693 +
1694 + The counterpart of ``rocm_inverse_rope_rows_`` for layers whose attention
1695 + output is MXFP8: rows the decode reduce did not emit (prefill) go through
1696 + here. ``o`` is [T, H, D]; ``out_data`` [T, H * D] e4m3 and ``out_scale``
1697 + [T, H * D // 32] E8M0, the layout the reduce epilogue writes.
1698 + """
1699 + num_tokens, num_heads, head_dim = o.shape
1700 + if num_tokens == 0:
1701 + return
1702 + assert o.stride(-1) == 1 and out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
1703 + assert out_data.shape == (num_tokens, num_heads * head_dim)
1704 + assert out_scale.shape == (num_tokens, num_heads * head_dim // 32)
1705 + assert cos_sin_cache.shape[-1] == rope_head_dim
1706 + _inverse_rope_mxfp8_quant_kernel[(num_tokens, num_heads)](
1707 + o,
1708 + out_data,
1709 + out_scale,
1710 + positions,
1711 + cos_sin_cache,
1712 + o.stride(0),
1713 + o.stride(1),
1714 + out_data.stride(0),
1715 + out_scale.stride(0),
1716 + cos_sin_cache.stride(0),
1717 + HEAD_DIM=head_dim,
1718 + NOPE=head_dim - rope_head_dim,
1719 + HALF=rope_head_dim // 2,
1720 + num_warps=4,
1721 + )
1722 +
1723 +
1724 def _get_cached_wo_a_bf16(
1725 wo_a: torch.nn.Module,
1726 n_local_groups: int,
@@ -1722,6 +1806,175 @@ def rocm_inv_rope_einsum(
1806 return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
1807
1808
1809 +@triton.jit
1810 +def _mxfp8_wo_a_bmm_kernel(
1811 + a_ptr, # [T, G * K] e4m3
1812 + as_ptr, # [T, G * K // 32] E8M0
1813 + w_ptr, # [G * R, K] e4m3
1814 + ws_ptr, # [G * R, K // 32] E8M0
1815 + out_ptr, # [T, G * R]
1816 + num_tokens,
1817 + stride_at,
1818 + stride_ast,
1819 + stride_wn,
1820 + stride_wsn,
1821 + stride_out,
1822 + R: tl.constexpr,
1823 + K: tl.constexpr,
1824 + BLOCK_M: tl.constexpr,
1825 + BLOCK_N: tl.constexpr,
1826 + BLOCK_K: tl.constexpr,
1827 +):
1828 + # Output columns are group-major, so a BLOCK_N tile never straddles two
1829 + # groups (R % BLOCK_N == 0) and picks its group's K slice of A.
1830 + pid_n = tl.program_id(0)
1831 + pid_m = tl.program_id(1)
1832 + group = (pid_n * BLOCK_N) // R
1833 + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
1834 + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
1835 + m_mask = offs_m < num_tokens
1836 + offs_k = tl.arange(0, BLOCK_K)
1837 + offs_sk = tl.arange(0, BLOCK_K // 32)
1838 +
1839 + a_ptrs = a_ptr + offs_m[:, None] * stride_at + (group * K + offs_k)[None, :]
1840 + as_ptrs = (
1841 + as_ptr + offs_m[:, None] * stride_ast + (group * (K // 32) + offs_sk)[None, :]
1842 + )
1843 + w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :]
1844 + ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :]
1845 +
1846 + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
1847 + for _ in range(K // BLOCK_K):
1848 + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
1849 + a_s = tl.load(as_ptrs, mask=m_mask[:, None], other=127)
1850 + w = tl.load(w_ptrs)
1851 + w_s = tl.load(ws_ptrs)
1852 + acc += tl.dot_scaled(a, a_s, "e4m3", w.T, w_s, "e4m3")
1853 + a_ptrs += BLOCK_K
1854 + as_ptrs += BLOCK_K // 32
1855 + w_ptrs += BLOCK_K
1856 + ws_ptrs += BLOCK_K // 32
1857 +
1858 + tl.store(
1859 + out_ptr + offs_m[:, None] * stride_out + offs_n[None, :],
1860 + acc.to(out_ptr.dtype.element_ty),
1861 + mask=m_mask[:, None],
1862 + )
1863 +
1864 +
1865 +def _mxfp8_wo_a_bmm_config(num_tokens: int, n_groups: int) -> tuple[int, ...]:
1866 + """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx950.
1867 +
1868 + Tuned under HIP graphs with a cold weight at G = 4 and 2, over every
1869 + decode shape of conc 1-128 x 0-5 spec tokens plus prefill chunks up to
1870 + 8K tokens. The best tile tracks the total work T * G, so the tiers are
1871 + keyed on it.
1872 +
1873 + This will be replaced after new GEMM kernel from AITER with proper 32x32 scale
1874 + shape GEMM fp8 enabled.
1875 + """
1876 + work = num_tokens * n_groups
1877 + if work <= 64:
1878 + return 16, 16, 1024, 2, 3
1879 + if work <= 128:
1880 + return 32, 16, 1024, 2, 3
1881 + if work <= 256:
1882 + return 32, 32, 512, 2, 3
1883 + if work <= 512:
1884 + return 64, 32, 512, 2, 3
1885 + if work <= 1024:
1886 + return 64, 64, 512, 4, 2
1887 + if work <= 2048:
1888 + return 64, 64, 256, 4, 2
1889 + if work <= 3072:
1890 + return 64, 64, 256, 2, 1
1891 + if work <= 4096:
1892 + return 128, 128, 256, 8, 2
1893 + return 128, 128, 128, 4, 2
1894 +
1895 +
1896 +def _rocm_mxfp8_wo_a_bmm_impl(
1897 + a: torch.Tensor,
1898 + a_scale: torch.Tensor,
1899 + weight: torch.Tensor,
1900 + weight_scale: torch.Tensor,
1901 + n_groups: int,
1902 + o_lora_rank: int,
1903 +) -> torch.Tensor:
1904 + num_tokens = a.shape[0]
1905 + group_dim = a.shape[1] // n_groups
1906 + out = torch.empty(
1907 + (num_tokens, n_groups * o_lora_rank), dtype=torch.bfloat16, device=a.device
1908 + )
1909 + if num_tokens == 0:
1910 + return out
1911 + block_m, block_n, block_k, num_warps, num_stages = _mxfp8_wo_a_bmm_config(
1912 + num_tokens, n_groups
1913 + )
1914 + grid = (n_groups * o_lora_rank // block_n, triton.cdiv(num_tokens, block_m))
1915 + _mxfp8_wo_a_bmm_kernel[grid](
1916 + a,
1917 + a_scale,
1918 + weight,
1919 + weight_scale,
1920 + out,
1921 + num_tokens,
1922 + a.stride(0),
1923 + a_scale.stride(0),
1924 + weight.stride(0),
1925 + weight_scale.stride(0),
1926 + out.stride(0),
1927 + R=o_lora_rank,
1928 + K=group_dim,
1929 + BLOCK_M=block_m,
1930 + BLOCK_N=block_n,
1931 + BLOCK_K=block_k,
1932 + num_warps=num_warps,
1933 + num_stages=num_stages,
1934 + )
1935 + return out
1936 +
1937 +
1938 +def _rocm_mxfp8_wo_a_bmm_fake(
1939 + a: torch.Tensor,
1940 + a_scale: torch.Tensor,
1941 + weight: torch.Tensor,
1942 + weight_scale: torch.Tensor,
1943 + n_groups: int,
1944 + o_lora_rank: int,
1945 +) -> torch.Tensor:
1946 + return a.new_empty((a.shape[0], n_groups * o_lora_rank), dtype=torch.bfloat16)
1947 +
1948 +
1949 +# An opaque op: the tile choice branches on the token count, which the
1950 +# compiled o_proj must not freeze at its trace-time value.
1951 +direct_register_custom_op(
1952 + op_name="rocm_dsv41_mxfp8_wo_a_bmm",
1953 + op_func=_rocm_mxfp8_wo_a_bmm_impl,
1954 + fake_impl=_rocm_mxfp8_wo_a_bmm_fake,
1955 +)
1956 +
1957 +
1958 +def rocm_mxfp8_wo_a_bmm(
1959 + a: torch.Tensor,
1960 + a_scale: torch.Tensor,
1961 + wo_a: torch.nn.Module,
1962 + n_groups: int,
1963 + o_lora_rank: int,
1964 +) -> torch.Tensor:
1965 + """Grouped MXFP8 wo_a: ``out[t, g, :] = a[t, g, :] @ W[g].T``, bf16 out.
1966 +
1967 + ``a`` is the [T, G * K] e4m3 attention output and ``a_scale`` its
1968 + [T, G * K // 32] E8M0 scales, as the sparse decode reduce writes them.
1969 + The weight is the checkpoint's MXFP8 ``wo_a`` as loaded, [G * R, K] with
1970 + [G * R, K // 32] scales, so there is no dequantized copy to keep.
1971 + Returns [T, G * R].
1972 + """
1973 + return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm(
1974 + a, a_scale, wo_a.weight, wo_a.weight_scale, n_groups, o_lora_rank
1975 + )
1976 +
1977 +
1978 _DSV4_SPARSE_NOPE_DIM = 448
1979 _DSV4_SPARSE_ROPE_DIM = 64
1980
@@ -2922,6 +3175,33 @@ def _sparse_attn_decode_gfx950_partial_kernel(
3175 )
3176
3177
3178 +@triton.jit
3179 +def _mxfp8_scale_bits(amax):
3180 + """Biased E8M0 exponent that puts ``amax`` at the top of the e4m3 range.
3181 +
3182 + Same rounding as ``mxfp8_e4m3_quantize``, so the output is bit-identical to
3183 + quantizing the tensor there.
3184 + """
3185 + amax = tl.maximum(amax, 1.1754943508222875e-38)
3186 + bits = tl.ceil(tl.log2(amax / 448.0)) + 127.0
3187 + return tl.minimum(tl.maximum(bits, 0.0), 254.0)
3188 +
3189 +
3190 +@triton.jit
3191 +def _mxfp8_quantize_rows(x, ROWS: tl.constexpr, COLS: tl.constexpr):
3192 + """MXFP8-quantize ``x`` [ROWS, COLS] in registers, one scale per 32 lanes.
3193 +
3194 + Returns the rescaled fp32 values (to be cast to e4m3 on store) and the
3195 + [ROWS, COLS // 32] biased E8M0 exponents.
3196 + """
3197 + blocks = tl.reshape(x, (ROWS, COLS // 32, 32))
3198 + bits = _mxfp8_scale_bits(tl.max(tl.abs(blocks), axis=2))
3199 + # Multiply by the reciprocal: a divisor of 2**-127 would be subnormal and
3200 + # flush to zero, turning an all-zero block into NaN.
3201 + q = blocks * tl.exp2(127.0 - bits)[:, :, None]
3202 + return tl.reshape(q, (ROWS, COLS)), bits
3203 +
3204 +
3205 @triton.jit
3206 def _sparse_attn_decode_reduce_kernel(
3207 part_m_ptr,
@@ -2931,8 +3211,11 @@ def _sparse_attn_decode_reduce_kernel(
3211 out_ptr,
3212 pos_ptr,
3213 cos_sin_ptr,
3214 + out_scale_ptr,
3215 out_stride0,
3216 out_stride1,
3217 + os_stride0,
3218 + os_stride1,
3219 pm_stride0,
3220 pm_stride_s,
3221 pa_stride0,
@@ -2949,6 +3232,7 @@ def _sparse_attn_decode_reduce_kernel(
3232 FUSE_INV_ROPE: tl.constexpr,
3233 NOPE: tl.constexpr,
3234 HALF: tl.constexpr,
3235 + QUANT_OUT: tl.constexpr,
3236 ):
3237 query_idx = tl.program_id(0)
3238 pid_h = tl.program_id(1)
@@ -3061,9 +3345,22 @@ def _sparse_attn_decode_reduce_kernel(
3345 out_row_ptr = (
3346 out_ptr + query_idx * out_stride0 + head_offsets[:, None] * out_stride1
3347 )
3348 + if QUANT_OUT:
3349 + # Emit wo_a's MXFP8 input directly instead of a bf16 row that a
3350 + # separate pass would read back to quantize.
3351 + out, scale_bits = _mxfp8_quantize_rows(out, BLOCK_H, COMB_DIM)
3352 + scale_offsets = tl.arange(0, COMB_DIM // 32)
3353 + tl.store(
3354 + out_scale_ptr
3355 + + query_idx * os_stride0
3356 + + head_offsets[:, None] * os_stride1
3357 + + scale_offsets[None, :],
3358 + scale_bits.to(tl.uint8),
3359 + mask=head_mask[:, None],
3360 + )
3361 tl.store(
3362 out_row_ptr + comb_offsets[None, :],
3066 - out,
3363 + out.to(out_ptr.dtype.element_ty),
3364 mask=head_mask[:, None],
3365 )
3366
@@ -3380,7 +3677,14 @@ def _rocm_sparse_attn_decode_ragged_triton(
3677 adaptive_splits: bool = False,
3678 inv_rope_positions: torch.Tensor | None = None,
3679 inv_rope_cos_sin_cache: torch.Tensor | None = None,
3680 + out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
3681 ) -> torch.Tensor:
3682 + """Split-K sparse decode; returns the attention output.
3683 +
3684 + With ``out_mxfp8 = (data, scale)`` the reduce writes MXFP8 instead of
3685 + bf16: ``data`` is [b, h * d] e4m3 and ``scale`` [b, h * d // 32] E8M0, and
3686 + ``data`` viewed as [b, h, d] is returned.
3687 + """
3688 assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}"
3689 assert main_cache.ndim == 3, (
3690 f"expected main_cache=[blocks,block,bytes], got {main_cache.shape}"
@@ -3445,7 +3749,28 @@ def _rocm_sparse_attn_decode_ragged_triton(
3749 extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32)
3750
3751 block_h = 16
3448 - if out is None:
3752 + out_scale = None
3753 + if out_mxfp8 is not None:
3754 + assert out is None, "out and out_mxfp8 are mutually exclusive"
3755 + assert _ON_GFX950, "the MXFP8 reduce epilogue is gfx950-only"
3756 + assert inv_rope_positions is not None, (
3757 + "the MXFP8 output feeds wo_a, so it must be inverse-RoPE'd first"
3758 + )
3759 + out_data, out_scale = out_mxfp8
3760 + assert out_data.dtype == torch.float8_e4m3fn and out_scale.dtype == (
3761 + torch.uint8
3762 + ), f"expected e4m3/uint8 MXFP8 buffers, got {out_data.dtype}/{out_scale.dtype}"
3763 + assert out_data.shape == (num_queries, num_heads * head_dim), (
3764 + f"expected MXFP8 data [{num_queries}, {num_heads * head_dim}], "
3765 + f"got {tuple(out_data.shape)}"
3766 + )
3767 + assert out_scale.shape == (num_queries, num_heads * head_dim // 32), (
3768 + f"expected MXFP8 scale [{num_queries}, {num_heads * head_dim // 32}], "
3769 + f"got {tuple(out_scale.shape)}"
3770 + )
3771 + assert out_data.stride(-1) == 1 and out_scale.stride(-1) == 1
3772 + out = out_data.view(num_queries, num_heads, head_dim)
3773 + elif out is None:
3774 out = torch.empty_like(q, dtype=torch.bfloat16)
3775 else:
3776 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(
3968 out,
3969 inv_rope_positions,
3970 inv_rope_cos_sin_cache,
3971 + out_scale,
3972 out.stride(0),
3973 out.stride(1),
3974 + out_scale.stride(0) if out_scale is not None else 0,
3975 + head_dim // 32,
3976 part_m.stride(0),
3977 part_m.stride(1),
3978 part_acc.stride(0),
@@ -3661,6 +3989,7 @@ def _rocm_sparse_attn_decode_ragged_triton(
3989 FUSE_INV_ROPE=fuse_inv_rope,
3990 NOPE=nope_head_dim,
3991 HALF=rope_head_dim // 2,
3992 + QUANT_OUT=out_scale is not None,
3993 num_warps=4,
3994 )
3995 return out
@@ -3687,6 +4016,7 @@ def _rocm_sparse_attn_decode_triton(
4016 adaptive_splits: bool = False,
4017 inv_rope_positions: torch.Tensor | None = None,
4018 inv_rope_cos_sin_cache: torch.Tensor | None = None,
4019 + out_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
4020 ) -> torch.Tensor:
4021 if main_ragged_indices is None or main_ragged_indptr is None:
4022 main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense(
@@ -3727,6 +4057,7 @@ def _rocm_sparse_attn_decode_triton(
4057 adaptive_splits=adaptive_splits,
4058 inv_rope_positions=inv_rope_positions,
4059 inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
4060 + out_mxfp8=out_mxfp8,
4061 )
4062
4063
@@ -3825,11 +4156,12 @@ def rocm_sparse_attn_decode(
4156 head_dim: int,
4157 nope_head_dim: int,
4158 rope_head_dim: int,
3828 - output: torch.Tensor,
4159 + output: torch.Tensor | None,
4160 extra_cache_nan_free: bool = False,
4161 adaptive_splits: bool = False,
4162 inv_rope_positions: torch.Tensor | None = None,
4163 inv_rope_cos_sin_cache: torch.Tensor | None = None,
4164 + output_mxfp8: tuple[torch.Tensor, torch.Tensor] | None = None,
4165 ) -> int:
4166 """Run sparse MLA decode into ``output``.
4167
@@ -3837,6 +4169,11 @@ def rocm_sparse_attn_decode(
4169 epilogue. Returns how many leading rows of ``output`` came back rotated,
4170 so a caller mixing in a decode path that does not fuse still knows what it
4171 owes the standalone pass. Read it from the eager attention segment only.
4172 +
4173 + ``output_mxfp8 = (data, scale)`` replaces ``output``: the reduce also
4174 + MXFP8-quantizes the rotated rows for the FP8 wo_a (see
4175 + ``_rocm_sparse_attn_decode_ragged_triton``). It needs gfx950 and the fused
4176 + inverse RoPE, and always covers every row.
4177 """
4178 assert swa_k_cache.dtype == torch.uint8, (
4179 "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, "
@@ -3866,7 +4203,12 @@ def rocm_sparse_attn_decode(
4203 if topk_indices is not None:
4204 extra_indices = topk_indices.reshape(topk_indices.shape[0], -1)
4205
3869 - direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
4206 + if output_mxfp8 is not None:
4207 + assert output is None, "output and output_mxfp8 are mutually exclusive"
4208 + direct_out = None
4209 + else:
4210 + assert output is not None
4211 + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None
4212 attn_out = _rocm_sparse_attn_decode_triton(
4213 q=q,
4214 main_cache=swa_k_cache,
@@ -3888,7 +4230,11 @@ def rocm_sparse_attn_decode(
4230 adaptive_splits=adaptive_splits,
4231 inv_rope_positions=inv_rope_positions,
4232 inv_rope_cos_sin_cache=inv_rope_cos_sin_cache,
4233 + out_mxfp8=output_mxfp8,
4234 )
4235 + if output_mxfp8 is not None:
4236 + return q.shape[0]
4237 + assert output is not None
4238 if direct_out is None:
4239 output.copy_(attn_out.to(output.dtype))
4240 return output.shape[0] if inv_rope_positions is not None else 0