From e336969fe34c1c8c9ffc51aba04d2243b4c863fa Mon Sep 17 00:00:00 2001 From: fai Date: Wed, 23 Sep 2026 21:21:07 +0000 Subject: [PATCH 1/4] [ROCm][DSv4.1][Perf] Emit MXFP8 from the sparse decode reduce and run wo_a as a grouped FP8 GEMM On gfx950 the sparse decode reduce now quantizes its inverse-RoPE'd output to MXFP8 (e4m3 + per-32 E8M0 scales) in the epilogue, and wo_a consumes it through a grouped tl.dot_scaled GEMM that reads the checkpoint's MXFP8 weight as loaded. This replaces the bf16 attention output + bf16 hipBLASLt einsum on a dequantized weight copy. Prefill rows are rotated and quantized in one pass from a bf16 workspace. The bf16 path remains the fallback when the loaded wo_a layout or the platform does not support the MXFP8 GEMM. Co-Authored-By: Claude Opus 5.5 Signed-off-by: fai --- .../attention/test_rocm_triton_attn_dsv4.py | 339 +++++++++++++++- vllm/models/deepseek_v41/amd/rocm.py | 177 +++++++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 383 +++++++++++++++++- 3 files changed, 870 insertions(+), 29 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index a966fb60a49a..8ab6403238d6 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -261,16 +261,72 @@ 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, ) -> torch.Tensor: """Run the reduce kernel, with the inverse-RoPE epilogue if given positions.""" + out, _ = _run_sparse_decode_reduce( + part_m, + part_l, + part_acc, + adaptive_splits, + positions, + cos_sin_cache, + out_dtype, + quant=False, + ) + return out + + +def _launch_sparse_decode_reduce_mxfp8( + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + adaptive_splits: bool, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 epilogue: ([Q, H * D] e4m3, [Q, H * D // 32] E8M0).""" + out, out_scale = _run_sparse_decode_reduce( + part_m, + part_l, + part_acc, + adaptive_splits, + positions, + cos_sin_cache, + torch.float8_e4m3fn, + quant=True, + ) + assert out_scale is not None + return out.view(out.shape[0], -1), out_scale + + +def _run_sparse_decode_reduce( + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + adaptive_splits: bool, + positions: torch.Tensor | None, + cos_sin_cache: torch.Tensor | None, + out_dtype: torch.dtype, + quant: bool, +) -> tuple[torch.Tensor, torch.Tensor | None]: 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, ) + out_scale = ( + torch.empty( + (num_queries, num_heads * HEAD_DIM // 32), + dtype=torch.uint8, + device=part_m.device, + ) + if quant + else None + ) attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device) mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -280,8 +336,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,9 +357,10 @@ def _launch_sparse_decode_reduce( FUSE_INV_ROPE=positions is not None, NOPE=NOPE_HEAD_DIM, HALF=ROPE_HEAD_DIM // 2, + QUANT_OUT=quant, num_warps=4, ) - return out + return out, out_scale @torch.inference_mode() @@ -1057,6 +1117,281 @@ 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 _random_reduce_inputs( + num_queries: int, num_splits: int, num_heads: int, max_pos: int, seed: int +) -> tuple[torch.Tensor, ...]: + device = torch.device("cuda") + torch.manual_seed(seed) + shape = (num_queries, num_splits, num_heads) + part_m = torch.randn(shape, dtype=torch.float32, device=device) + part_l = torch.rand(shape, dtype=torch.float32, device=device) + 0.5 + part_acc = torch.randn((*shape, HEAD_DIM), dtype=torch.float32, device=device) + # Spread block amaxes over many binades so the E8M0 scales vary. + block_exp = torch.randint(-12, 6, (*shape, HEAD_DIM // 32), device=device) + part_acc *= torch.exp2(block_exp.float()).repeat_interleave(32, dim=-1) + positions = torch.randint( + 0, max_pos, (num_queries,), dtype=torch.int64, device=device + ) + angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device) + cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() + return part_m, part_l, part_acc, positions, cos_sin_cache + + +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) + + +@requires_gfx950 +@pytest.mark.parametrize("adaptive_splits", [False, True]) +@torch.inference_mode() +def test_sparse_attn_decode_reduce_mxfp8_epilogue(adaptive_splits: bool) -> None: + """QUANT_OUT must equal MXFP8-quantizing the rotated fp32 reduce output. + + Checked bit for bit against ``_mxfp8_e4m3_quantize_torch``, the reference + the ROCm MXFP8 linear kernels are held to, so wo_a sees exactly what a + standalone quant of the reduce output would have given it. + """ + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + ) + + part_m, part_l, part_acc, positions, cos_sin_cache = _random_reduce_inputs( + 6, 4, 8, 64, seed=11 + ) + rotated_fp32 = _launch_sparse_decode_reduce( + part_m, + part_l, + part_acc, + adaptive_splits, + positions, + cos_sin_cache, + out_dtype=torch.float32, + ) + expected_data, expected_scale = _mxfp8_e4m3_quantize_torch( + rotated_fp32.view(rotated_fp32.shape[0], -1) + ) + data, scale = _launch_sparse_decode_reduce_mxfp8( + part_m, part_l, part_acc, adaptive_splits, positions, cos_sin_cache + ) + + assert data.dtype == torch.float8_e4m3fn and scale.dtype == torch.uint8 + assert torch.equal(scale, expected_scale) + assert torch.equal(data.view(torch.uint8), expected_data.view(torch.uint8)) + assert scale.unique().numel() > 4 + + +@requires_gfx950 +@pytest.mark.parametrize("num_tokens", [1, 7, 33]) +@torch.inference_mode() +def test_inverse_rope_mxfp8_rows_matches_reference(num_tokens: int) -> None: + """The prefill rotate+quant pass must write what the reduce epilogue writes. + + Both have to agree with inverse-RoPE-then-MXFP8 in fp32, since one wo_a + GEMM consumes decode and prefill rows side by side. + """ + 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_heads, max_pos = 4, 128 + o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).to(torch.bfloat16) + positions = torch.randint( + 0, max_pos, (num_tokens,), dtype=torch.int64, device=device + ) + angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device) + cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() + + ref = o.float().clone() + cos = cos_sin_cache[positions, : ROPE_HEAD_DIM // 2][:, None, :] + sin = cos_sin_cache[positions, ROPE_HEAD_DIM // 2 :][:, None, :] + even = o[..., NOPE_HEAD_DIM::2].float() + odd = o[..., NOPE_HEAD_DIM + 1 :: 2].float() + ref[..., NOPE_HEAD_DIM::2] = even * cos + odd * sin + ref[..., NOPE_HEAD_DIM + 1 :: 2] = odd * cos - even * sin + _, expected_scale = _mxfp8_e4m3_quantize_torch(ref.view(num_tokens, -1)) + + data = torch.empty( + num_tokens, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device + ) + scale = torch.empty( + num_tokens, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device + ) + rocm_inverse_rope_mxfp8_rows( + o, positions, cos_sin_cache, ROPE_HEAD_DIM, data, scale + ) + + assert torch.equal(scale, expected_scale) + # Half an e4m3 step at the top of each block's range. + bound = 16.0 * torch.exp2(expected_scale.float() - 127.0).repeat_interleave(32, -1) + err = (_mxfp8_dequant(data, scale) - ref.view(num_tokens, -1)).abs() + assert bool((err <= bound).all()) + + +def _random_mxfp8( + rows: int, cols: int, device: torch.device +) -> tuple[torch.Tensor, ...]: + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + ) + + return _mxfp8_e4m3_quantize_torch(torch.randn(rows, cols, device=device)) + + +class _FakeMxfp8WoA(torch.nn.Module): + """wo_a as the ROCm MXFP8 linear leaves it: e4m3 weight, per-row E8M0.""" + + def __init__(self, weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + super().__init__() + self.weight = weight + self.weight_scale = weight_scale + + +@requires_gfx950 +@pytest.mark.parametrize( + "num_tokens, n_groups", + # Covers every tile tier of _mxfp8_wo_a_bmm_config, plus partial M tiles. + [ + (1, 4), + (5, 2), + (20, 4), + (48, 4), + (77, 4), + (130, 4), + (300, 2), + (700, 4), + (1025, 4), + (1100, 8), + ], +) +@torch.inference_mode() +def test_rocm_mxfp8_wo_a_bmm(num_tokens: int, n_groups: int) -> None: + 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 = _random_mxfp8(num_tokens, n_groups * group_dim, device) + w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) + + out = rocm_mxfp8_wo_a_bmm( + a, a_scale, _FakeMxfp8WoA(w, w_scale), 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), + ).reshape(num_tokens, -1) + assert out.shape == (num_tokens, n_groups * o_lora_rank) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), expected, atol=5e-2, rtol=1e-2) + + +@requires_gfx950 +def test_mxfp8_wo_a_bmm_supported_checks_loaded_layout() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import mxfp8_wo_a_bmm_supported + + device = torch.device("cuda") + n_groups, o_lora_rank, group_dim = 2, 128, 4096 + w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) + assert mxfp8_wo_a_bmm_supported( + _FakeMxfp8WoA(w, w_scale), n_groups, o_lora_rank, group_dim + ) + # A weight an emulated kernel already dequantized. + assert not mxfp8_wo_a_bmm_supported( + _FakeMxfp8WoA(w.to(torch.bfloat16), w_scale), n_groups, o_lora_rank, group_dim + ) + # Checkpoint-style 32x32 block scales rather than the per-row expansion. + assert not mxfp8_wo_a_bmm_supported( + _FakeMxfp8WoA(w, w_scale[::32].contiguous()), n_groups, o_lora_rank, group_dim + ) + assert not mxfp8_wo_a_bmm_supported(_FakeWoA(w), n_groups, o_lora_rank, group_dim) + + +@requires_gfx950 +@torch.inference_mode() +def test_sparse_decode_mxfp8_wo_a_matches_bf16_path() -> None: + """Decode -> wo_a through the MXFP8 path against today's bf16 path. + + Runs the real split-K decode kernels both ways on identical inputs: bf16 + output + inverse RoPE + bf16 einsum, and the MXFP8 epilogue + grouped FP8 + GEMM. The FP8 result must match the einsum of its own dequantized + activations tightly (layout / group mapping) and the bf16 path to within + the MXFP8 activation error. + """ + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_triton, + rocm_mxfp8_wo_a_bmm, + ) + + device = torch.device("cuda") + torch.manual_seed(9) + block_size, num_queries, num_heads, n_groups, o_lora_rank = 4, 5, 16, 2, 64 + group_dim = num_heads * HEAD_DIM // n_groups + q = torch.randn(num_queries, num_heads, HEAD_DIM, device=device).to(torch.bfloat16) + kv = torch.randn(40, HEAD_DIM, device=device).to(torch.bfloat16) + cache = _pack_fp8_ds_mla_cache( + kv, block_size, use_fnuz=current_platform.is_fp8_fnuz() + ) + lens = torch.tensor([3, 17, 40, 9, 26], dtype=torch.int32, device=device) + indptr = torch.zeros(num_queries + 1, dtype=torch.int32, device=device) + torch.cumsum(lens, 0, out=indptr[1:]) + indices = torch.cat( + [torch.randperm(40, device=device)[:n] for n in lens.tolist()] + ).to(torch.int32) + positions = torch.randint(0, 64, (num_queries,), dtype=torch.int64, device=device) + angle = torch.randn(64, ROPE_HEAD_DIM // 2, device=device) + cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() + attn_sink = torch.randn(num_heads, device=device) + w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) + w_ref = _mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim) + common = dict( + q=q, + main_cache=cache, + main_indices=indices, + main_indptr=indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + inv_rope_positions=positions, + inv_rope_cos_sin_cache=cos_sin_cache, + ) + + bf16_out = _rocm_sparse_attn_decode_ragged_triton(**common) + z_bf16 = torch.einsum( + "tgd,grd->tgr", + bf16_out.view(num_queries, n_groups, group_dim), + w_ref.to(torch.bfloat16), + ).reshape(num_queries, -1) + + data = torch.empty( + num_queries, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device + ) + scale = torch.empty( + num_queries, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device + ) + _rocm_sparse_attn_decode_ragged_triton(**common, out_mxfp8=(data, scale)) + z_fp8 = rocm_mxfp8_wo_a_bmm( + data, scale, _FakeMxfp8WoA(w, w_scale), n_groups, o_lora_rank + ) + + z_own = torch.einsum( + "tgd,grd->tgr", + _mxfp8_dequant(data, scale).view(num_queries, n_groups, group_dim), + w_ref, + ).reshape(num_queries, -1) + torch.testing.assert_close(z_fp8.float(), z_own, atol=2e-2, rtol=1e-2) + rel = (z_fp8.float() - z_bf16.float()).norm() / z_bf16.float().norm() + assert rel < 4e-2, f"MXFP8 wo_a drifted {rel:.3e} from the bf16 path" + + @requires_gfx950 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) @torch.inference_mode() diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index dee7237ff46f..261d88567780 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -13,6 +13,8 @@ ) 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, @@ -36,8 +38,11 @@ ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( build_ragged_indices_from_dense, + mxfp8_wo_a_bmm_supported, 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 +668,58 @@ def _split_qkv_and_norm( transpose_scale=False, ) - def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + @functools.cached_property + def _use_mxfp8_wo_a(self) -> bool: + """Whether attention emits MXFP8 and wo_a runs as a grouped FP8 GEMM. + + Resolved at the first forward, after the MXFP8 linear has finished + post-processing the wo_a weight it checks. + """ + return ( + self.padded_heads == self.n_local_heads + and self.head_dim % 32 == 0 + and mxfp8_wo_a_bmm_supported( + self.wo_a, + self.n_local_groups, + self.o_lora_rank, + self.n_local_heads * self.head_dim // self.n_local_groups, + ) + ) + + def _alloc_attn_out( + self, num_tokens: int, hidden_states: torch.Tensor + ) -> torch.Tensor | QuantizedActivation: + if not self._use_mxfp8_wo_a: + 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 +732,29 @@ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tens 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 +772,13 @@ def forward_mqa( ) 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 +808,41 @@ def forward_mqa( 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 +910,8 @@ def _forward_decode( 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 +951,7 @@ def _forward_decode( 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 +967,16 @@ def _forward_prefill( 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 +1012,13 @@ def _forward_prefill( 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 +1083,31 @@ def _forward_prefill( 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 self._use_mxfp8_wo_a: + shapes.append( + ( + (num_prefill_tokens, self.n_local_heads, self.head_dim), + torch.bfloat16, + ) + ) + return shapes diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index d14eedb1e1d5..2af079b6d55f 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -16,7 +16,7 @@ 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,202 @@ 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 mxfp8_wo_a_bmm_supported( + wo_a: torch.nn.Module, n_groups: int, o_lora_rank: int, group_dim: int +) -> bool: + """Whether ``rocm_mxfp8_wo_a_bmm`` can run this ``wo_a`` as loaded. + + Needs native MX matrix cores (gfx950) and the one-byte MXFP8 weight with + per-row E8M0 scales that the ROCm MXFP8 linear keeps after loading. + """ + if not (_ON_GFX950 and current_platform.supports_mx()): + return False + weight = getattr(wo_a, "weight", None) + scale = getattr(wo_a, "weight_scale", None) + if weight is None or scale is None: + return False + n = n_groups * o_lora_rank + return ( + weight.dtype == torch.float8_e4m3fn + and tuple(weight.shape) == (n, group_dim) + and weight.is_contiguous() + and scale.dtype == torch.uint8 + and tuple(scale.shape) == (n, group_dim // 32) + and scale.is_contiguous() + and o_lora_rank % 64 == 0 + and group_dim % 1024 == 0 + ) + + +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 +3202,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 +3238,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 +3259,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 +3372,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 +3704,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 +3776,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 +3995,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 +4016,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 +4043,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 +4084,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 +4183,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 +4196,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 +4230,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 +4257,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 From fb009a292831b0b5b2550b13594259086c2f8d28 Mon Sep 17 00:00:00 2001 From: fai Date: Wed, 23 Sep 2026 21:30:04 +0000 Subject: [PATCH 2/4] Simplify the MXFP8 wo_a tests Co-Authored-By: Claude Opus 5.5 Signed-off-by: fai --- .../attention/test_rocm_triton_attn_dsv4.py | 356 +++++------------- 1 file changed, 98 insertions(+), 258 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 8ab6403238d6..6c10f169bee4 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -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 @@ -262,54 +263,12 @@ def _launch_sparse_decode_reduce( 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.""" - out, _ = _run_sparse_decode_reduce( - part_m, - part_l, - part_acc, - adaptive_splits, - positions, - cos_sin_cache, - out_dtype, - quant=False, - ) - return out - - -def _launch_sparse_decode_reduce_mxfp8( - part_m: torch.Tensor, - part_l: torch.Tensor, - part_acc: torch.Tensor, - adaptive_splits: bool, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """The MXFP8 epilogue: ([Q, H * D] e4m3, [Q, H * D // 32] E8M0).""" - out, out_scale = _run_sparse_decode_reduce( - part_m, - part_l, - part_acc, - adaptive_splits, - positions, - cos_sin_cache, - torch.float8_e4m3fn, - quant=True, - ) - assert out_scale is not None - return out.view(out.shape[0], -1), out_scale - + """Run the reduce kernel, with the inverse-RoPE epilogue if given positions. -def _run_sparse_decode_reduce( - part_m: torch.Tensor, - part_l: torch.Tensor, - part_acc: torch.Tensor, - adaptive_splits: bool, - positions: torch.Tensor | None, - cos_sin_cache: torch.Tensor | None, - out_dtype: torch.dtype, - quant: bool, -) -> tuple[torch.Tensor, torch.Tensor | None]: + 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 @@ -318,15 +277,6 @@ def _run_sparse_decode_reduce( dtype=out_dtype, device=part_m.device, ) - out_scale = ( - torch.empty( - (num_queries, num_heads * HEAD_DIM // 32), - dtype=torch.uint8, - device=part_m.device, - ) - if quant - else None - ) attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device) mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -357,10 +307,10 @@ def _run_sparse_decode_reduce( FUSE_INV_ROPE=positions is not None, NOPE=NOPE_HEAD_DIM, HALF=ROPE_HEAD_DIM // 2, - QUANT_OUT=quant, + QUANT_OUT=out_scale is not None, num_warps=4, ) - return out, out_scale + return out @torch.inference_mode() @@ -1117,79 +1067,54 @@ 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 _random_reduce_inputs( - num_queries: int, num_splits: int, num_heads: int, max_pos: int, seed: int -) -> tuple[torch.Tensor, ...]: - device = torch.device("cuda") - torch.manual_seed(seed) - shape = (num_queries, num_splits, num_heads) - part_m = torch.randn(shape, dtype=torch.float32, device=device) - part_l = torch.rand(shape, dtype=torch.float32, device=device) + 0.5 - part_acc = torch.randn((*shape, HEAD_DIM), dtype=torch.float32, device=device) - # Spread block amaxes over many binades so the E8M0 scales vary. - block_exp = torch.randint(-12, 6, (*shape, HEAD_DIM // 32), device=device) - part_acc *= torch.exp2(block_exp.float()).repeat_interleave(32, dim=-1) - positions = torch.randint( - 0, max_pos, (num_queries,), dtype=torch.int64, device=device - ) - angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device) - cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() - return part_m, part_l, part_acc, positions, cos_sin_cache - - 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 -@pytest.mark.parametrize("adaptive_splits", [False, True]) @torch.inference_mode() -def test_sparse_attn_decode_reduce_mxfp8_epilogue(adaptive_splits: bool) -> None: - """QUANT_OUT must equal MXFP8-quantizing the rotated fp32 reduce output. - - Checked bit for bit against ``_mxfp8_e4m3_quantize_torch``, the reference - the ROCm MXFP8 linear kernels are held to, so wo_a sees exactly what a - standalone quant of the reduce output would have given it. - """ +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, ) - part_m, part_l, part_acc, positions, cos_sin_cache = _random_reduce_inputs( - 6, 4, 8, 64, seed=11 - ) - rotated_fp32 = _launch_sparse_decode_reduce( - part_m, - part_l, - part_acc, - adaptive_splits, - positions, - cos_sin_cache, - out_dtype=torch.float32, - ) + 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_fp32.view(rotated_fp32.shape[0], -1) + rotated.view(num_queries, -1) ) - data, scale = _launch_sparse_decode_reduce_mxfp8( - part_m, part_l, part_acc, adaptive_splits, positions, cos_sin_cache + scale = torch.empty_like(expected_scale) + data = _launch_sparse_decode_reduce( + *args, out_dtype=torch.float8_e4m3fn, out_scale=scale ) - assert data.dtype == torch.float8_e4m3fn and scale.dtype == torch.uint8 assert torch.equal(scale, expected_scale) - assert torch.equal(data.view(torch.uint8), expected_data.view(torch.uint8)) - assert scale.unique().numel() > 4 + assert torch.equal( + data.view(num_queries, -1).view(torch.uint8), expected_data.view(torch.uint8) + ) @requires_gfx950 -@pytest.mark.parametrize("num_tokens", [1, 7, 33]) @torch.inference_mode() -def test_inverse_rope_mxfp8_rows_matches_reference(num_tokens: int) -> None: - """The prefill rotate+quant pass must write what the reduce epilogue writes. - - Both have to agree with inverse-RoPE-then-MXFP8 in fp32, since one wo_a - GEMM consumes decode and prefill rows side by side. - """ +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, ) @@ -1199,197 +1124,112 @@ def test_inverse_rope_mxfp8_rows_matches_reference(num_tokens: int) -> None: device = torch.device("cuda") torch.manual_seed(5) - num_heads, max_pos = 4, 128 - o = torch.randn(num_tokens, num_heads, HEAD_DIM, device=device).to(torch.bfloat16) - positions = torch.randint( - 0, max_pos, (num_tokens,), dtype=torch.int64, device=device - ) - angle = torch.randn(max_pos, ROPE_HEAD_DIM // 2, device=device) - cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() - - ref = o.float().clone() - cos = cos_sin_cache[positions, : ROPE_HEAD_DIM // 2][:, None, :] - sin = cos_sin_cache[positions, ROPE_HEAD_DIM // 2 :][:, None, :] - even = o[..., NOPE_HEAD_DIM::2].float() - odd = o[..., NOPE_HEAD_DIM + 1 :: 2].float() + 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 - _, expected_scale = _mxfp8_e4m3_quantize_torch(ref.view(num_tokens, -1)) + ref = ref.view(num_tokens, -1) + _, expected_scale = _mxfp8_e4m3_quantize_torch(ref) - data = torch.empty( - num_tokens, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device - ) - scale = torch.empty( - num_tokens, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device - ) + 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) - # Half an e4m3 step at the top of each block's range. - bound = 16.0 * torch.exp2(expected_scale.float() - 127.0).repeat_interleave(32, -1) - err = (_mxfp8_dequant(data, scale) - ref.view(num_tokens, -1)).abs() - assert bool((err <= bound).all()) + # 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()) -def _random_mxfp8( - rows: int, cols: int, device: torch.device -) -> tuple[torch.Tensor, ...]: - from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( - _mxfp8_e4m3_quantize_torch, +@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, ) - return _mxfp8_e4m3_quantize_torch(torch.randn(rows, cols, device=device)) - + 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), + ) -class _FakeMxfp8WoA(torch.nn.Module): - """wo_a as the ROCm MXFP8 linear leaves it: e4m3 weight, per-row E8M0.""" + 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)) - def __init__(self, weight: torch.Tensor, weight_scale: torch.Tensor) -> None: - super().__init__() - self.weight = weight - self.weight_scale = weight_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", - # Covers every tile tier of _mxfp8_wo_a_bmm_config, plus partial M tiles. + # One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles. [ (1, 4), - (5, 2), (20, 4), (48, 4), (77, 4), (130, 4), - (300, 2), + (300, 4), (700, 4), - (1025, 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 = _random_mxfp8(num_tokens, n_groups * group_dim, device) - w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) - - out = rocm_mxfp8_wo_a_bmm( - a, a_scale, _FakeMxfp8WoA(w, w_scale), n_groups, o_lora_rank + 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), - ).reshape(num_tokens, -1) - assert out.shape == (num_tokens, n_groups * o_lora_rank) - assert out.dtype == torch.bfloat16 - torch.testing.assert_close(out.float(), expected, atol=5e-2, rtol=1e-2) - - -@requires_gfx950 -def test_mxfp8_wo_a_bmm_supported_checks_loaded_layout() -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import mxfp8_wo_a_bmm_supported - - device = torch.device("cuda") - n_groups, o_lora_rank, group_dim = 2, 128, 4096 - w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) - assert mxfp8_wo_a_bmm_supported( - _FakeMxfp8WoA(w, w_scale), n_groups, o_lora_rank, group_dim - ) - # A weight an emulated kernel already dequantized. - assert not mxfp8_wo_a_bmm_supported( - _FakeMxfp8WoA(w.to(torch.bfloat16), w_scale), n_groups, o_lora_rank, group_dim - ) - # Checkpoint-style 32x32 block scales rather than the per-row expansion. - assert not mxfp8_wo_a_bmm_supported( - _FakeMxfp8WoA(w, w_scale[::32].contiguous()), n_groups, o_lora_rank, group_dim - ) - assert not mxfp8_wo_a_bmm_supported(_FakeWoA(w), n_groups, o_lora_rank, group_dim) - - -@requires_gfx950 -@torch.inference_mode() -def test_sparse_decode_mxfp8_wo_a_matches_bf16_path() -> None: - """Decode -> wo_a through the MXFP8 path against today's bf16 path. - - Runs the real split-K decode kernels both ways on identical inputs: bf16 - output + inverse RoPE + bf16 einsum, and the MXFP8 epilogue + grouped FP8 - GEMM. The FP8 result must match the einsum of its own dequantized - activations tightly (layout / group mapping) and the bf16 path to within - the MXFP8 activation error. - """ - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _rocm_sparse_attn_decode_ragged_triton, - rocm_mxfp8_wo_a_bmm, - ) - - device = torch.device("cuda") - torch.manual_seed(9) - block_size, num_queries, num_heads, n_groups, o_lora_rank = 4, 5, 16, 2, 64 - group_dim = num_heads * HEAD_DIM // n_groups - q = torch.randn(num_queries, num_heads, HEAD_DIM, device=device).to(torch.bfloat16) - kv = torch.randn(40, HEAD_DIM, device=device).to(torch.bfloat16) - cache = _pack_fp8_ds_mla_cache( - kv, block_size, use_fnuz=current_platform.is_fp8_fnuz() - ) - lens = torch.tensor([3, 17, 40, 9, 26], dtype=torch.int32, device=device) - indptr = torch.zeros(num_queries + 1, dtype=torch.int32, device=device) - torch.cumsum(lens, 0, out=indptr[1:]) - indices = torch.cat( - [torch.randperm(40, device=device)[:n] for n in lens.tolist()] - ).to(torch.int32) - positions = torch.randint(0, 64, (num_queries,), dtype=torch.int64, device=device) - angle = torch.randn(64, ROPE_HEAD_DIM // 2, device=device) - cos_sin_cache = torch.cat((angle.cos(), angle.sin()), dim=-1).contiguous() - attn_sink = torch.randn(num_heads, device=device) - w, w_scale = _random_mxfp8(n_groups * o_lora_rank, group_dim, device) - w_ref = _mxfp8_dequant(w, w_scale).view(n_groups, o_lora_rank, group_dim) - common = dict( - q=q, - main_cache=cache, - main_indices=indices, - main_indptr=indptr, - scale=HEAD_DIM**-0.5, - attn_sink=attn_sink, - nope_head_dim=NOPE_HEAD_DIM, - rope_head_dim=ROPE_HEAD_DIM, - inv_rope_positions=positions, - inv_rope_cos_sin_cache=cos_sin_cache, ) - - bf16_out = _rocm_sparse_attn_decode_ragged_triton(**common) - z_bf16 = torch.einsum( - "tgd,grd->tgr", - bf16_out.view(num_queries, n_groups, group_dim), - w_ref.to(torch.bfloat16), - ).reshape(num_queries, -1) - - data = torch.empty( - num_queries, num_heads * HEAD_DIM, dtype=torch.float8_e4m3fn, device=device - ) - scale = torch.empty( - num_queries, num_heads * HEAD_DIM // 32, dtype=torch.uint8, device=device - ) - _rocm_sparse_attn_decode_ragged_triton(**common, out_mxfp8=(data, scale)) - z_fp8 = rocm_mxfp8_wo_a_bmm( - data, scale, _FakeMxfp8WoA(w, w_scale), n_groups, o_lora_rank - ) - - z_own = torch.einsum( - "tgd,grd->tgr", - _mxfp8_dequant(data, scale).view(num_queries, n_groups, group_dim), - w_ref, - ).reshape(num_queries, -1) - torch.testing.assert_close(z_fp8.float(), z_own, atol=2e-2, rtol=1e-2) - rel = (z_fp8.float() - z_bf16.float()).norm() / z_bf16.float().norm() - assert rel < 4e-2, f"MXFP8 wo_a drifted {rel:.3e} from the bf16 path" + torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2) @requires_gfx950 From 3857a7d7150abbbe0f85f5af5e6cd1b2663c203c Mon Sep 17 00:00:00 2001 From: fai Date: Wed, 23 Sep 2026 21:46:28 +0000 Subject: [PATCH 3/4] Gate the MXFP8 wo_a path on gfx950 only Co-Authored-By: Claude Opus 5.5 Signed-off-by: fai --- vllm/models/deepseek_v41/amd/rocm.py | 25 +++-------------- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 27 ------------------- 2 files changed, 3 insertions(+), 49 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 261d88567780..e54294df0ce6 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -27,7 +27,7 @@ DeepseekV41SparseSWAMetadataBuilder, ) from vllm.platforms import current_platform -from vllm.platforms.rocm import _ON_GFX950 +from vllm.platforms.rocm import _ON_GFX950, on_gfx950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -38,7 +38,6 @@ ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( build_ragged_indices_from_dense, - mxfp8_wo_a_bmm_supported, rocm_inv_rope_einsum, rocm_inverse_rope_mxfp8_rows, rocm_inverse_rope_rows_, @@ -668,28 +667,10 @@ def _split_qkv_and_norm( transpose_scale=False, ) - @functools.cached_property - def _use_mxfp8_wo_a(self) -> bool: - """Whether attention emits MXFP8 and wo_a runs as a grouped FP8 GEMM. - - Resolved at the first forward, after the MXFP8 linear has finished - post-processing the wo_a weight it checks. - """ - return ( - self.padded_heads == self.n_local_heads - and self.head_dim % 32 == 0 - and mxfp8_wo_a_bmm_supported( - self.wo_a, - self.n_local_groups, - self.o_lora_rank, - self.n_local_heads * self.head_dim // self.n_local_groups, - ) - ) - def _alloc_attn_out( self, num_tokens: int, hidden_states: torch.Tensor ) -> torch.Tensor | QuantizedActivation: - if not self._use_mxfp8_wo_a: + 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. @@ -1103,7 +1084,7 @@ def _prefill_workspace_shapes( shapes: list[tuple[tuple[int, ...], torch.dtype]] = [ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16) ] - if self._use_mxfp8_wo_a: + if on_gfx950(): shapes.append( ( (num_prefill_tokens, self.n_local_heads, self.head_dim), diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 2af079b6d55f..c7cef64ddfa2 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1893,33 +1893,6 @@ def _mxfp8_wo_a_bmm_config(num_tokens: int, n_groups: int) -> tuple[int, ...]: return 128, 128, 128, 4, 2 -def mxfp8_wo_a_bmm_supported( - wo_a: torch.nn.Module, n_groups: int, o_lora_rank: int, group_dim: int -) -> bool: - """Whether ``rocm_mxfp8_wo_a_bmm`` can run this ``wo_a`` as loaded. - - Needs native MX matrix cores (gfx950) and the one-byte MXFP8 weight with - per-row E8M0 scales that the ROCm MXFP8 linear keeps after loading. - """ - if not (_ON_GFX950 and current_platform.supports_mx()): - return False - weight = getattr(wo_a, "weight", None) - scale = getattr(wo_a, "weight_scale", None) - if weight is None or scale is None: - return False - n = n_groups * o_lora_rank - return ( - weight.dtype == torch.float8_e4m3fn - and tuple(weight.shape) == (n, group_dim) - and weight.is_contiguous() - and scale.dtype == torch.uint8 - and tuple(scale.shape) == (n, group_dim // 32) - and scale.is_contiguous() - and o_lora_rank % 64 == 0 - and group_dim % 1024 == 0 - ) - - def _rocm_mxfp8_wo_a_bmm_impl( a: torch.Tensor, a_scale: torch.Tensor, From 44954aa7db526227ce9d5263763feb65e72b7928 Mon Sep 17 00:00:00 2001 From: fai Date: Thu, 24 Sep 2026 05:53:52 +0000 Subject: [PATCH 4/4] Use _ON_GFX950 uniformly in the V4.1 ROCm attention Co-Authored-By: Claude Opus 5.5 Signed-off-by: fai --- vllm/models/deepseek_v41/amd/rocm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index e54294df0ce6..e1588b55e716 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -27,7 +27,7 @@ DeepseekV41SparseSWAMetadataBuilder, ) from vllm.platforms import current_platform -from vllm.platforms.rocm import _ON_GFX950, on_gfx950 +from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -670,7 +670,7 @@ def _split_qkv_and_norm( def _alloc_attn_out( self, num_tokens: int, hidden_states: torch.Tensor ) -> torch.Tensor | QuantizedActivation: - if not on_gfx950(): + 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. @@ -1084,7 +1084,7 @@ def _prefill_workspace_shapes( shapes: list[tuple[tuple[int, ...], torch.dtype]] = [ ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16) ] - if on_gfx950(): + if _ON_GFX950: shapes.append( ( (num_prefill_tokens, self.n_local_heads, self.head_dim),