From 1ae3465495a96facab650b931d35edc35b7043fa Mon Sep 17 00:00:00 2001 From: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:34:10 -0500 Subject: [PATCH 1/4] [ROCm][Perf][DSv4.1] TP1 MXFP8 GEMM tiles and a pre-quantized wo_b input for DEP8 on MI355X At TP1 (DP attention, e.g. DEP8) each rank runs the full DeepSeek-V4.1-Flash projections, whose shapes rocm_block32_gemm tunes only at TP2/TP4, so wo_b, wq_b and the shared expert's gate_up/down ran _default_config, which drops to a 128x128x128 tile past M = 256. Add TP1 tiers for those four shapes, swept at every CUDA graph size 48-4096, and 8-group tiers for the wo_a BMM past 256 tokens. When wo_b would run exactly the block32 MXFP8 GEMM, the wo_a BMM now quantizes its bf16 tile to MXFP8 in the epilogue and hands wo_b the e4m3 values and E8M0 scales, which is bit-identical to quantizing the bf16 output in a separate launch. Signed-off-by: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 77 ++++++-- .../quantization/test_rocm_mxfp8_linear.py | 6 +- .../kernels/linear/mxfp8/rocm_block32_gemm.py | 61 +++++- vllm/models/deepseek_v41/amd/rocm.py | 42 +++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 176 +++++++++++++++--- 5 files changed, 320 insertions(+), 42 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 66414f723ee2..8bd40d9cb949 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1486,24 +1486,28 @@ def test_sparse_attn_decode_mxfp8_output() -> None: @requires_gfx950 @pytest.mark.parametrize( - "num_tokens, n_groups", + "num_tokens, n_groups, o_lora_rank", # One case per tile tier of _mxfp8_wo_a_bmm_config, with partial M tiles. [ - (1, 4), - (20, 4), - (48, 4), - (77, 4), - (130, 4), - (300, 4), - (700, 4), - (1000, 4), - (1100, 8), + (1, 4, 128), + (20, 4, 128), + (48, 4, 128), + (77, 4, 128), + (130, 4, 128), + (300, 4, 128), + (700, 4, 128), + (1000, 4, 128), + (300, 8, 128), + (700, 8, 128), + (1100, 8, 128), + (1600, 8, 128), + (1600, 8, 256), ], ) @pytest.mark.parametrize("block_scales", [False, True]) @torch.inference_mode() def test_rocm_mxfp8_wo_a_bmm( - num_tokens: int, n_groups: int, block_scales: bool + num_tokens: int, n_groups: int, o_lora_rank: int, block_scales: bool ) -> None: from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( _as_block32_scale, @@ -1515,7 +1519,7 @@ def test_rocm_mxfp8_wo_a_bmm( device = torch.device("cuda") torch.manual_seed(3) - o_lora_rank, group_dim = 128, 2048 + group_dim = 2048 a, a_scale = _mxfp8_e4m3_quantize_torch( torch.randn(num_tokens, n_groups * group_dim, device=device) ) @@ -1541,6 +1545,55 @@ def test_rocm_mxfp8_wo_a_bmm( torch.testing.assert_close(out.float(), expected.flatten(1), atol=5e-2, rtol=1e-2) +@requires_gfx950 +@pytest.mark.parametrize( + "num_tokens, n_groups, o_lora_rank", + # A tile narrower than an MX block (separate quant), then fused epilogues + # on the G = 4 and G = 8 tiers, including the 256-wide G = 8 tile. + [(1, 4, 128), (77, 4, 128), (300, 8, 128), (1100, 8, 128), (1600, 8, 256)], +) +@pytest.mark.parametrize("block_scales", [False, True]) +@torch.inference_mode() +def test_rocm_mxfp8_wo_a_bmm_mxfp8_matches_quantized_bmm( + num_tokens: int, n_groups: int, o_lora_rank: int, block_scales: bool +) -> None: + """wo_b's MXFP8 input from the wo_a epilogue is bit-identical to quantizing + the bf16 wo_a output, so the fused o_proj path does not change results.""" + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _as_block32_scale, + ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + _mxfp8_e4m3_quantize_torch, + mxfp8_e4m3_quantize, + ) + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + rocm_mxfp8_wo_a_bmm, + rocm_mxfp8_wo_a_bmm_mxfp8, + ) + + device = torch.device("cuda") + torch.manual_seed(5) + group_dim = 2048 + a, a_scale = _mxfp8_e4m3_quantize_torch( + torch.randn(num_tokens, n_groups * group_dim, device=device) + ) + w, w_scale = _mxfp8_e4m3_quantize_torch( + torch.randn(n_groups * o_lora_rank, group_dim, device=device) + ) + if block_scales: + w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0)) + assert w_scale is not None + wo_a = SimpleNamespace(weight=w, weight_scale=w_scale) + + q, q_scale = rocm_mxfp8_wo_a_bmm_mxfp8(a, a_scale, wo_a, n_groups, o_lora_rank) + + ref_q, ref_scale = mxfp8_e4m3_quantize( + rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank) + ) + assert torch.equal(q.view(torch.uint8), ref_q.view(torch.uint8)) + assert torch.equal(q_scale, ref_scale) + + @requires_gfx950 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) @torch.inference_mode() diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index d5e48f2aafa2..7bde5576a5cf 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -117,11 +117,15 @@ def test_rocm_mxfp8_linear_handles_extreme_magnitudes(): assert (out[0] == 0).all(), "an all-zero activation row must give a zero row" -# (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=2 and TP=4. The +# (name, N, K) of the DeepSeek-V4.1-Flash MXFP8 linears at TP=1, 2 and 4. The # checkpoint stores their scales in 32x32 blocks; shared_down at TP=4 has # K % 128 != 0, which only the block-scaled GEMM runs natively. V41_BLOCK32_SHAPES = [ ("fused_wqa_wkv", 1280 + 512, 5120), + ("wq_b_tp1", 64 * 512, 1280), + ("wo_b_tp1", 5120, 8192), + ("shared_gate_up_tp1", 2 * 2304, 5120), + ("shared_down_tp1", 5120, 2304), ("wq_b_tp2", 64 * 512 // 2, 1280), ("indexer_wq_b", 32 * 128, 1280), ("wo_b_tp2", 5120, 8192 // 2), diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py index 3d5e9c8357ca..54c010e6e559 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py @@ -332,7 +332,7 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf # (N, K) -> [(max M, config), ...] for the DeepSeek-V4.1-Flash MXFP8 linears at -# TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A +# TP1/TP2/TP4, tuned on MI355X under HIP graphs with weights cold in the LLC. A # tier serves every M up to its bound; larger M reuse the last tier. _TUNED: dict[tuple[int, int], list[tuple[int, _Config]]] = { # shared expert gate_up, TP4 @@ -426,6 +426,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf (4096, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), (8192, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)), ], + # shared expert gate_up, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096 + (4608, 5120): [ + (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)), + (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)), + (64, _tiled(64, 64, 512, 1, False, 4, 2, 32, 0, 2)), + (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 2)), + (128, _tiled(64, 64, 256, 1, False, 4, 2, 32, 2, 2)), + (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2)), + (448, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)), + (896, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0)), + (1792, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)), + (2560, _tiled(128, 128, 256, 1, True, 8, 2, 32, 0)), + (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)), + ], # shared expert down, TP4 (5120, 576): [ (8, _packed(8, 32, 256, 2, 2, 2, 16, 1)), @@ -483,6 +497,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf (4096, _tiled(128, 128, 128, 1, True, 8, 2, 32, 1)), (8192, _tiled(128, 128, 128, 1, True, 2, 2, 32, 1)), ], + # shared expert down, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096 + (5120, 2304): [ + (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)), + (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)), + (64, _tiled(64, 64, 256, 1, True, 8, 3, 16, 0, 4)), + (96, _tiled(32, 64, 512, 1, True, 8, 3, 16, 2)), + (192, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2)), + (288, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)), + (384, _tiled(128, 64, 256, 1, True, 8, 2, 32, 0)), + (768, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)), + (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)), + (3072, _tiled(128, 256, 128, 1, True, 8, 3, 32, 2)), + (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), + ], # wo_b, TP2 (5120, 4096): [ (4, _packed(4, 32, 512, 4, 2, 2, 16, 0)), @@ -505,6 +533,23 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf (4096, _tiled(128, 128, 128, 1, True, 8, 2, 16, 1)), (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), ], + # wo_b, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096 + (5120, 8192): [ + (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)), + (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)), + (48, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2, 4)), + (64, _tiled(64, 64, 512, 1, True, 4, 2, 32, 2, 2)), + (96, _tiled(32, 64, 512, 1, True, 8, 2, 16, 0)), + (128, _tiled(64, 64, 256, 1, True, 4, 2, 32, 0, 2)), + (192, _tiled(64, 64, 512, 1, True, 4, 2, 16, 0)), + (384, _tiled(128, 128, 256, 2, True, 4, 2, 16, 0)), + (768, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0)), + (1024, _tiled(128, 128, 256, 2, True, 4, 2, 16, 0)), + (1536, _tiled(128, 128, 256, 1, True, 8, 2, 32, 0)), + (2048, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)), + (3072, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)), + (8192, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)), + ], # wq_b, TP4 (8192, 1280): [ (1, _packed(8, 32, 256, 2, 4, 3, 16, 1)), @@ -539,6 +584,20 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf (2048, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), (8192, _tiled(128, 256, 128, 1, True, 2, 2, 32, 0)), ], + # wq_b, TP1 (DP attention), tuned at the CUDA graph sizes 48-4096 + (32768, 1280): [ + (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)), + (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)), + (56, _tiled(64, 64, 256, 1, True, 4, 2, 32, 2)), + (64, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2)), + (128, _tiled(128, 128, 256, 1, False, 8, 2, 32, 0)), + (160, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)), + (224, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)), + (288, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), + (448, _tiled(128, 128, 128, 1, True, 8, 2, 32, 0)), + (1280, _tiled(128, 128, 128, 1, True, 4, 2, 32, 0)), + (8192, _tiled(128, 128, 128, 1, True, 4, 2, 32, 2)), + ], } diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 08ce0b8c2ce7..e1592812b7cc 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -14,6 +14,9 @@ ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( + rocm_mxfp8_block32_gemm, +) from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import ( @@ -57,6 +60,7 @@ rocm_inverse_rope_mxfp8_rows, rocm_inverse_rope_rows_, rocm_mxfp8_wo_a_bmm, + rocm_mxfp8_wo_a_bmm_mxfp8, rocm_sparse_attn_decode, rocm_sparse_attn_prefill, ) @@ -932,6 +936,21 @@ def _o_proj( self, attn_out: torch.Tensor | QuantizedActivation, positions: torch.Tensor ) -> torch.Tensor: if isinstance(attn_out, QuantizedActivation): + if self._wo_b_takes_wo_a_mxfp8(): + q, q_scale = rocm_mxfp8_wo_a_bmm_mxfp8( + attn_out.data, + attn_out.scale, + self.wo_a, + self.n_local_groups, + self.o_lora_rank, + ) + return rocm_mxfp8_block32_gemm( + q, + q_scale, + self.wo_b.weight, + self.wo_b.weight_scale, + attn_out.orig_dtype, + ) z = rocm_mxfp8_wo_a_bmm( attn_out.data, attn_out.scale, @@ -954,6 +973,29 @@ def _o_proj( ) return self._wo_b_after_wo_a(z.flatten(1)) + def _wo_b_takes_wo_a_mxfp8(self) -> bool: + """Whether wo_b's linear runs exactly a block32 MXFP8 GEMM on the + MXFP8-quantized wo_a output, so wo_a can hand that over directly: + 32x32 block-scaled e4m3 weight, no bias, no tensor-parallel reduce.""" + if self._wo_b_scale is not None: + return False + w = getattr(self.wo_b, "weight", None) + ws = getattr(self.wo_b, "weight_scale", None) + return ( + w is not None + and ws is not None + and w.dtype == torch.float8_e4m3fn + and ws.dtype == torch.uint8 + and w.dim() == 2 + and ws.dim() == 2 + and ws.shape[0] != w.shape[0] + and getattr(self.wo_b, "bias", None) is None + and ( + not getattr(self.wo_b, "reduce_results", False) + or getattr(self.wo_b, "tp_size", 1) == 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: return self._bpre_attn_gemm( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 2b882dc298c7..972b9fe1bef8 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1795,19 +1795,22 @@ def _mxfp8_wo_a_bmm_kernel( as_ptr, # [T, G * K // 32] E8M0 w_ptr, # [G * R, K] e4m3 ws_ptr, # [G * R // SCALE_ROWS, K // 32] E8M0 - out_ptr, # [T, G * R] + out_ptr, # [T, G * R] bf16, or e4m3 with QUANT_OUT + out_scale_ptr, # [T, G * R // 32] E8M0 with QUANT_OUT, else unused num_tokens, stride_at, stride_ast, stride_wn, stride_wsn, stride_out, + stride_out_scale, R: tl.constexpr, K: tl.constexpr, SCALE_ROWS: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + QUANT_OUT: 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. @@ -1839,6 +1842,21 @@ def _mxfp8_wo_a_bmm_kernel( w_ptrs += BLOCK_K ws_ptrs += BLOCK_K // 32 + if QUANT_OUT: + # mxfp8_e4m3_quantize of the bf16-rounded tile: the same values and + # scales as quantizing the bf16 output in a separate launch. + x = tl.reshape(acc.to(tl.bfloat16).to(tl.float32), (BLOCK_M, BLOCK_N // 32, 32)) + amax = tl.maximum(tl.max(tl.abs(x), axis=2), 1.1754943508222875e-38) + bits = tl.minimum( + tl.maximum(tl.ceil(tl.log2(amax / 448.0)) + 127.0, 0.0), 254.0 + ) + acc = tl.reshape(x * tl.exp2(127.0 - bits)[:, :, None], (BLOCK_M, BLOCK_N)) + offs_sn = pid_n * (BLOCK_N // 32) + tl.arange(0, BLOCK_N // 32) + tl.store( + out_scale_ptr + offs_m[:, None] * stride_out_scale + offs_sn[None, :], + bits.to(tl.uint8), + mask=m_mask[:, None], + ) tl.store( out_ptr + offs_m[:, None] * stride_out + offs_n[None, :], acc.to(out_ptr.dtype.element_ty), @@ -1846,76 +1864,110 @@ def _mxfp8_wo_a_bmm_kernel( ) -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. +def _mxfp8_wo_a_bmm_config( + num_tokens: int, n_groups: int, o_lora_rank: int +) -> tuple[tuple[int, int, int, int, int], dict[str, int]]: + """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for gfx950, and the + compiler hints that tile was tuned with. 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. + keyed on it. G = 8 (TP1, DP attention) was tuned the same way at the CUDA + graph sizes 48-4096; past 256 tokens it has its own tiers. This will be replaced after new GEMM kernel from AITER with proper 32x32 scale shape GEMM fp8 enabled. """ + if n_groups == 8 and num_tokens > 256: + if num_tokens <= 512: + return (128, 128, 256, 4, 2), {} + if num_tokens <= 1024: + return (128, 128, 128, 4, 2), {} + if num_tokens > 1536 and o_lora_rank % 256 == 0: + return (128, 256, 128, 8, 3), { + "waves_per_eu": 2, + "matrix_instr_nonkdim": 16, + } + return (128, 128, 256, 4, 2), {} work = num_tokens * n_groups if work <= 64: - return 16, 16, 1024, 2, 3 + return (16, 16, 1024, 2, 3), {} if work <= 128: - return 32, 16, 1024, 2, 3 + return (32, 16, 1024, 2, 3), {} if work <= 256: - return 32, 32, 512, 2, 3 + return (32, 32, 512, 2, 3), {} if work <= 512: - return 64, 32, 512, 2, 3 + return (64, 32, 512, 2, 3), {} if work <= 1024: - return 64, 64, 512, 4, 2 + return (64, 64, 512, 4, 2), {} if work <= 2048: - return 64, 64, 256, 4, 2 + return (64, 64, 256, 4, 2), {} if work <= 3072: - return 64, 64, 256, 2, 1 + return (64, 64, 256, 2, 1), {} if work <= 4096: - return 128, 128, 256, 8, 2 - return 128, 128, 128, 4, 2 + return (128, 128, 256, 8, 2), {} + return (128, 128, 128, 4, 2), {} -def _rocm_mxfp8_wo_a_bmm_impl( +def _launch_mxfp8_wo_a_bmm( 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)) + tile: tuple[int, int, int, int, int], + hints: dict[str, int], + out: torch.Tensor, + out_scale: torch.Tensor | None = None, +) -> None: + block_m, block_n, block_k, num_warps, num_stages = tile + grid = (n_groups * o_lora_rank // block_n, triton.cdiv(a.shape[0], block_m)) _mxfp8_wo_a_bmm_kernel[grid]( a, a_scale, weight, weight_scale, out, - num_tokens, + out if out_scale is None else out_scale, + a.shape[0], a.stride(0), a_scale.stride(0), weight.stride(0), weight_scale.stride(0), out.stride(0), + 0 if out_scale is None else out_scale.stride(0), R=o_lora_rank, - K=group_dim, + K=a.shape[1] // n_groups, SCALE_ROWS=weight.shape[0] // weight_scale.shape[0], BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, + QUANT_OUT=out_scale is not None, num_warps=num_warps, num_stages=num_stages, + **hints, + ) + + +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] + out = torch.empty( + (num_tokens, n_groups * o_lora_rank), dtype=torch.bfloat16, device=a.device + ) + if num_tokens == 0: + return out + tile, hints = _mxfp8_wo_a_bmm_config(num_tokens, n_groups, o_lora_rank) + _launch_mxfp8_wo_a_bmm( + a, a_scale, weight, weight_scale, n_groups, o_lora_rank, tile, hints, out ) return out @@ -1961,6 +2013,74 @@ def rocm_mxfp8_wo_a_bmm( ) +def _rocm_mxfp8_wo_a_bmm_mxfp8_impl( + a: torch.Tensor, + a_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + n_groups: int, + o_lora_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + num_tokens = a.shape[0] + width = n_groups * o_lora_rank + q = torch.empty((num_tokens, width), dtype=torch.float8_e4m3fn, device=a.device) + q_scale = torch.empty((num_tokens, width // 32), dtype=torch.uint8, device=a.device) + if num_tokens == 0: + return q, q_scale + tile, hints = _mxfp8_wo_a_bmm_config(num_tokens, n_groups, o_lora_rank) + if tile[1] % 32: + # A tile narrower than one MX block cannot quantize in its epilogue. + out = _rocm_mxfp8_wo_a_bmm_impl( + a, a_scale, weight, weight_scale, n_groups, o_lora_rank + ) + return mxfp8_e4m3_quantize(out) + _launch_mxfp8_wo_a_bmm( + a, a_scale, weight, weight_scale, n_groups, o_lora_rank, tile, hints, q, q_scale + ) + return q, q_scale + + +def _rocm_mxfp8_wo_a_bmm_mxfp8_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + n_groups: int, + o_lora_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + width = n_groups * o_lora_rank + return ( + a.new_empty((a.shape[0], width), dtype=torch.float8_e4m3fn), + a.new_empty((a.shape[0], width // 32), dtype=torch.uint8), + ) + + +direct_register_custom_op( + op_name="rocm_dsv41_mxfp8_wo_a_bmm_mxfp8", + op_func=_rocm_mxfp8_wo_a_bmm_mxfp8_impl, + fake_impl=_rocm_mxfp8_wo_a_bmm_mxfp8_fake, +) + + +def rocm_mxfp8_wo_a_bmm_mxfp8( + a: torch.Tensor, + a_scale: torch.Tensor, + wo_a: torch.nn.Module, + n_groups: int, + o_lora_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """rocm_mxfp8_wo_a_bmm with wo_b's MXFP8 input as output: the [T, G * R] + e4m3 values and their [T, G * R // 32] E8M0 scales, bit-identical to + mxfp8_e4m3_quantize of the bf16 result.""" + return torch.ops.vllm.rocm_dsv41_mxfp8_wo_a_bmm_mxfp8( + 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 From 621c4e5c3f8f0b36fc94640f9dfe0b36722eb913 Mon Sep 17 00:00:00 2001 From: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:32:42 -0500 Subject: [PATCH 2/4] [ROCm][DSv4.1] Tie the wo_b hand-off to its kernel and test every fused tile - The wo_a -> wo_b MXFP8 hand-off now also requires wo_b's linear kernel to be RocmDotScaledMxfp8LinearKernel, so a kernel that relays out the weight falls back to the separate quant. test_wo_b_takes_wo_a_mxfp8_gate covers the gate. - The bit-exact wo_a test covers all nine fused-epilogue tiles, and the block32 sweep adds 96, 300 and 2100 tokens. - The MI355 MLA Kernels step runs the wo_a and gate tests. Signed-off-by: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 1 + .../attention/test_rocm_triton_attn_dsv4.py | 74 ++++++++++++++++++- .../quantization/test_rocm_mxfp8_linear.py | 2 +- vllm/models/deepseek_v41/amd/rocm.py | 11 ++- 4 files changed, 83 insertions(+), 5 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 9af2dc4ee2aa..ea0e54d88118 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -435,6 +435,7 @@ steps: - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py - pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py::test_dsv41_adaptive_mla_swa_metadata_graph_replay + - pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py -k "wo_a_bmm or wo_b_takes_wo_a_mxfp8_gate" - label: ":nvidia: (L4) Quantization Kernels Shard %N" key: kernels-quantization-test diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 8bd40d9cb949..e8203f01b059 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1548,9 +1548,21 @@ def test_rocm_mxfp8_wo_a_bmm( @requires_gfx950 @pytest.mark.parametrize( "num_tokens, n_groups, o_lora_rank", - # A tile narrower than an MX block (separate quant), then fused epilogues - # on the G = 4 and G = 8 tiers, including the 256-wide G = 8 tile. - [(1, 4, 128), (77, 4, 128), (300, 8, 128), (1100, 8, 128), (1600, 8, 256)], + # A tile narrower than an MX block (separate quant), then each of the nine + # tiles with a fused epilogue: the work tiers at G = 4 and 8, the G = 8 + # tiers past 256 tokens and the 256-wide G = 8 tile. + [ + (1, 4, 128), + (77, 4, 128), + (600, 4, 128), + (900, 4, 128), + (24, 8, 128), + (100, 8, 128), + (224, 8, 128), + (300, 8, 128), + (800, 8, 128), + (1600, 8, 256), + ], ) @pytest.mark.parametrize("block_scales", [False, True]) @torch.inference_mode() @@ -1594,6 +1606,62 @@ def test_rocm_mxfp8_wo_a_bmm_mxfp8_matches_quantized_bmm( assert torch.equal(q_scale, ref_scale) +@pytest.mark.parametrize( + "case, expected", + [ + ("block32_gemm", True), + ("reduce_at_tp1", True), + ("other_mxfp8_kernel", False), + ("per_row_scale", False), + ("bf16_weight", False), + ("bias", False), + ("tp_reduce", False), + ("bpreshuffled_scale", False), + ], +) +def test_wo_b_takes_wo_a_mxfp8_gate(case: str, expected: bool) -> None: + """wo_a hands wo_b its MXFP8 output only when wo_b's linear is the block32 + MXFP8 GEMM on a 32x32 block-scaled e4m3 weight, without a bias or a + tensor-parallel reduce. Another kernel that keeps the same weight dtypes + and shapes does not get the fused input.""" + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention + + class OtherMxfp8Kernel(RocmDotScaledMxfp8LinearKernel): + pass + + n, k = 64, 128 + wo_b = SimpleNamespace( + quant_method=SimpleNamespace( + kernel=object.__new__(RocmDotScaledMxfp8LinearKernel) + ), + weight=torch.zeros(n, k, dtype=torch.float8_e4m3fn), + weight_scale=torch.zeros(n // 32, k // 32, dtype=torch.uint8), + bias=None, + reduce_results=False, + tp_size=1, + ) + attn = SimpleNamespace(_wo_b_scale=None, wo_b=wo_b) + if case == "reduce_at_tp1": + wo_b.reduce_results = True + elif case == "other_mxfp8_kernel": + wo_b.quant_method.kernel = object.__new__(OtherMxfp8Kernel) + elif case == "per_row_scale": + wo_b.weight_scale = torch.zeros(n, k // 32, dtype=torch.uint8) + elif case == "bf16_weight": + wo_b.weight = torch.zeros(n, k, dtype=torch.bfloat16) + elif case == "bias": + wo_b.bias = torch.zeros(n, dtype=torch.bfloat16) + elif case == "tp_reduce": + wo_b.reduce_results, wo_b.tp_size = True, 2 + elif case == "bpreshuffled_scale": + attn._wo_b_scale = torch.ones(1) + + assert DeepseekV41ROCMAiterMLAAttention._wo_b_takes_wo_a_mxfp8(attn) == expected + + @requires_gfx950 @pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) @torch.inference_mode() diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index 7bde5576a5cf..77817a69da6c 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -152,7 +152,7 @@ def _make_block32_layer(n: int, k: int, device: str) -> torch.nn.Module: @pytest.mark.parametrize("name,n,k", V41_BLOCK32_SHAPES, ids=lambda v: str(v)) -@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 192, 1024]) +@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 96, 192, 300, 1024, 2100]) @torch.inference_mode() def test_rocm_mxfp8_linear_block32_matches_dequant_reference(name, n, k, num_tokens): """32x32 block scales stay compact and the block-scaled GEMM is exact. diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index e1592812b7cc..e1e0643fd8b4 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -979,10 +979,19 @@ def _wo_b_takes_wo_a_mxfp8(self) -> bool: 32x32 block-scaled e4m3 weight, no bias, no tensor-parallel reduce.""" if self._wo_b_scale is not None: return False + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + + # The kernel owns wo_b's weight layout: a kernel that relays out an + # e4m3 [N, K] weight with uint8 [N/32, K/32] scales passes every + # dtype and shape check below. + kernel = getattr(getattr(self.wo_b, "quant_method", None), "kernel", None) w = getattr(self.wo_b, "weight", None) ws = getattr(self.wo_b, "weight_scale", None) return ( - w is not None + type(kernel) is RocmDotScaledMxfp8LinearKernel + and w is not None and ws is not None and w.dtype == torch.float8_e4m3fn and ws.dtype == torch.uint8 From 98c7dcbb34ad0fc4a9e0e8fab3bb608f0f470c43 Mon Sep 17 00:00:00 2001 From: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:56:34 -0500 Subject: [PATCH 3/4] [ROCm][DSv4.1] Keep wo_b's fused quant where the wo_a tile cannot quantize Since #63179, wo_b's linear quantizes its bf16 input inside the block32 GEMM at M <= 8. The wo_a -> wo_b hand-off replaced wo_b's linear at every batch size, and where the wo_a tile is narrower than one MX block (T * G <= 128) it quantized separately and ran the unfused GEMM: one launch more per layer than main at tiny batches. o_proj now runs one opaque op, rocm_dsv41_mxfp8_o_proj, that picks per call: wo_a's MXFP8 epilogue and the block32 GEMM where the tile can quantize, and wo_a's bf16 output into rocm_mxfp8_block32_gemm_quant below that, as wo_b's linear does. The branch is inside the op so the compiled o_proj does not freeze the token count. test_rocm_mxfp8_wo_a_bmm_o_proj_matches_wo_b_linear checks the path on both sides of T * G = 128 and that the output is bit-identical to wo_b's linear. Signed-off-by: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> --- .../attention/test_rocm_triton_attn_dsv4.py | 73 ++++++++++++++++ vllm/models/deepseek_v41/amd/rocm.py | 14 +--- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 84 +++++++++++++++++++ 3 files changed, 160 insertions(+), 11 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index e8203f01b059..f802aea36d29 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1606,6 +1606,79 @@ def test_rocm_mxfp8_wo_a_bmm_mxfp8_matches_quantized_bmm( assert torch.equal(q_scale, ref_scale) +@requires_gfx950 +@pytest.mark.parametrize( + "num_tokens, n_groups", + # On both sides of the narrowest wo_a tile that quantizes in its epilogue + # (T * G > 128), and at the per-rank decode sizes of DP8 at conc 256. + [(1, 8), (8, 8), (16, 8), (17, 8), (224, 8), (300, 8), (8, 4), (32, 4), (33, 4)], +) +@torch.inference_mode() +def test_rocm_mxfp8_wo_a_bmm_o_proj_matches_wo_b_linear( + num_tokens: int, n_groups: int, monkeypatch: pytest.MonkeyPatch +) -> None: + """The o_proj op hands wo_b its MXFP8 input from wo_a's epilogue only where + the wo_a tile can quantize; below that wo_b's own quantizing GEMM runs, so a + tiny batch keeps the block32 GEMM's fused quant. Either way the output is + bit-identical to wo_b's linear on the bf16 wo_a output.""" + import vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm as block32 + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _as_block32_scale, + ) + 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_o_proj, + rocm_mxfp8_wo_a_bmm, + ) + + device = torch.device("cuda") + torch.manual_seed(7) + group_dim, o_lora_rank, n = 2048, 128, 512 + + def block_scaled(rows: int, cols: int) -> SimpleNamespace: + w, w_scale = _mxfp8_e4m3_quantize_torch(torch.randn(rows, cols, device=device)) + w_scale = _as_block32_scale(w_scale[::32].repeat_interleave(32, dim=0)) + assert w_scale is not None + return SimpleNamespace(weight=w, weight_scale=w_scale) + + a, a_scale = _mxfp8_e4m3_quantize_torch( + torch.randn(num_tokens, n_groups * group_dim, device=device) + ) + wo_a = block_scaled(n_groups * o_lora_rank, group_dim) + wo_b = block_scaled(n, n_groups * o_lora_rank) + expected = block32.rocm_mxfp8_block32_gemm_quant( + rocm_mxfp8_wo_a_bmm(a, a_scale, wo_a, n_groups, o_lora_rank), + wo_b.weight, + wo_b.weight_scale, + torch.bfloat16, + ) + + calls: list[str] = [] + + def counted(name: str): + fn = getattr(block32, name) + + def wrapper(*args): + calls.append(name) + return fn(*args) + + return wrapper + + for name in ("rocm_mxfp8_block32_gemm", "rocm_mxfp8_block32_gemm_quant"): + monkeypatch.setattr(block32, name, counted(name)) + out = rocm_mxfp8_o_proj( + a, a_scale, wo_a, n_groups, o_lora_rank, wo_b, torch.bfloat16 + ) + + fused = num_tokens * n_groups > 128 + assert calls == [ + "rocm_mxfp8_block32_gemm" if fused else "rocm_mxfp8_block32_gemm_quant" + ] + assert torch.equal(out, expected) + + @pytest.mark.parametrize( "case, expected", [ diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 0bb607f894ca..5ab439f78d87 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -14,9 +14,6 @@ ) from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( - rocm_mxfp8_block32_gemm, -) from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import ( @@ -59,8 +56,8 @@ rocm_inv_rope_einsum, rocm_inverse_rope_mxfp8_rows, rocm_inverse_rope_rows_, + rocm_mxfp8_o_proj, rocm_mxfp8_wo_a_bmm, - rocm_mxfp8_wo_a_bmm_mxfp8, rocm_sparse_attn_decode, rocm_sparse_attn_prefill, ) @@ -936,18 +933,13 @@ def _o_proj( ) -> torch.Tensor: if isinstance(attn_out, QuantizedActivation): if self._wo_b_takes_wo_a_mxfp8(): - q, q_scale = rocm_mxfp8_wo_a_bmm_mxfp8( + return rocm_mxfp8_o_proj( attn_out.data, attn_out.scale, self.wo_a, self.n_local_groups, self.o_lora_rank, - ) - return rocm_mxfp8_block32_gemm( - q, - q_scale, - self.wo_b.weight, - self.wo_b.weight_scale, + self.wo_b, attn_out.orig_dtype, ) z = rocm_mxfp8_wo_a_bmm( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 99d57db5a644..02020ae38f13 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -2225,6 +2225,90 @@ def rocm_mxfp8_wo_a_bmm_mxfp8( ) +def _rocm_mxfp8_o_proj_impl( + a: torch.Tensor, + a_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + n_groups: int, + o_lora_rank: int, + wo_b_weight: torch.Tensor, + wo_b_weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + from vllm.model_executor.kernels.linear.mxfp8.rocm_block32_gemm import ( + rocm_mxfp8_block32_gemm, + rocm_mxfp8_block32_gemm_quant, + ) + + tile, _ = _mxfp8_wo_a_bmm_config(a.shape[0], n_groups, o_lora_rank) + if tile[1] % 32: + # A tile narrower than one MX block cannot quantize in its epilogue, so + # wo_b quantizes the bf16 wo_a output as its own linear does, inside + # the GEMM at tiny M. + out = _rocm_mxfp8_wo_a_bmm_impl( + a, a_scale, weight, weight_scale, n_groups, o_lora_rank + ) + return rocm_mxfp8_block32_gemm_quant( + out, wo_b_weight, wo_b_weight_scale, out_dtype + ) + q, q_scale = _rocm_mxfp8_wo_a_bmm_mxfp8_impl( + a, a_scale, weight, weight_scale, n_groups, o_lora_rank + ) + return rocm_mxfp8_block32_gemm( + q, q_scale, wo_b_weight, wo_b_weight_scale, out_dtype + ) + + +def _rocm_mxfp8_o_proj_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + n_groups: int, + o_lora_rank: int, + wo_b_weight: torch.Tensor, + wo_b_weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + return a.new_empty((a.shape[0], wo_b_weight.shape[0]), dtype=out_dtype) + + +# An opaque op: whether the wo_a tile can hand wo_b an MXFP8 input depends on +# the token count, which the compiled o_proj must not freeze. +direct_register_custom_op( + op_name="rocm_dsv41_mxfp8_o_proj", + op_func=_rocm_mxfp8_o_proj_impl, + fake_impl=_rocm_mxfp8_o_proj_fake, +) + + +def rocm_mxfp8_o_proj( + a: torch.Tensor, + a_scale: torch.Tensor, + wo_a: torch.nn.Module, + n_groups: int, + o_lora_rank: int, + wo_b: torch.nn.Module, + out_dtype: torch.dtype, +) -> torch.Tensor: + """wo_b(wo_a(a)) when wo_b's linear is the block32 MXFP8 GEMM. Where the + wo_a tile is at least one MX block wide, its epilogue hands wo_b the MXFP8 + input (rocm_mxfp8_wo_a_bmm_mxfp8); below that, wo_b quantizes the bf16 + wo_a output itself, as its linear would.""" + return torch.ops.vllm.rocm_dsv41_mxfp8_o_proj( + a, + a_scale, + wo_a.weight, + wo_a.weight_scale, + n_groups, + o_lora_rank, + wo_b.weight, + wo_b.weight_scale, + out_dtype, + ) + + _DSV4_SPARSE_NOPE_DIM = 448 _DSV4_SPARSE_ROPE_DIM = 64 From 82daaacb598bb24a562d032c797b39180bc4c931 Mon Sep 17 00:00:00 2001 From: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:53:48 -0600 Subject: [PATCH 4/4] [ROCm][Perf][DSv4.1] block32 tiers for the DSpark drafter's main_proj main_proj (ReplicatedLinear, N=5120, K=15360) projects the target's aux hidden states once per draft step, 7 rows a request, so a DEP8 rank runs it at M = 56-336 (a DEP4 rank at about twice that). (5120, 15360) has no _TUNED entry, so _default_config served it. This adds a tier list, the best configs of a block32 sweep at M = 56-336, keeping the default packed configs at M <= 32, and the shape to the block32 accuracy test. MI355X, weights cold over 4 rotating copies, HIP graphs, us per call: M 14 42 56 84 112 126 133 168 224 280 336 before 57.6 61.4 62.9 60.3 61.6 62.8 60.2 61.0 62.8 137.5 143.3 after 58.1 28.2 27.8 30.3 31.9 32.6 46.1 42.3 51.9 55.7 55.9 Against an fp32 matmul of the dequantized inputs the new tiers reach a max relative error of 2.95e-3 over M = 7-448, as the other tuned shapes do. Signed-off-by: Mehmet Aslan <286132399+maslan-coop26@users.noreply.github.com> --- tests/kernels/quantization/test_rocm_mxfp8_linear.py | 1 + .../kernels/linear/mxfp8/rocm_block32_gemm.py | 11 +++++++++++ 2 files changed, 12 insertions(+) diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index 997966510094..67efd896ffd7 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -195,6 +195,7 @@ def test_rocm_mxfp8_linear_handles_extreme_magnitudes(): ("shared_gate_up_tp2", 2 * 2304 // 2, 5120), ("shared_down_tp2", 5120, 2304 // 2), ("shared_down_tp4", 5120, 2304 // 4), + ("dspark_main_proj", 5120, 3 * 5120), ] diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py index 0983182ed74a..9617926b575c 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py @@ -674,6 +674,17 @@ def _packed(bm, bn, bk, k_pack, warps, stages, nonkdim, waves, fused=1) -> _Conf (3072, _tiled(128, 128, 128, 1, True, 4, 2, 16, 0)), (8192, _tiled(128, 128, 256, 1, True, 8, 2, 32, 2)), ], + # DSpark main_proj (15360 -> 5120, replicated): 7 rows a request per step + (5120, 15360): [ + (8, _packed(8, 16, 256, 2, 2, 3, 16, 0)), + (32, _packed(32, 32, 256, 1, 2, 3, 16, 0)), + (56, _tiled(64, 64, 512, 1, True, 4, 2, 16, 2, 2)), + (126, _tiled(128, 128, 256, 1, True, 4, 2, 16, 2, 4)), + (133, _tiled(64, 64, 256, 1, True, 4, 2, 16, 2, 2)), + (168, _tiled(64, 64, 512, 1, True, 4, 2, 16, 2)), + (224, _tiled(128, 128, 256, 1, True, 4, 2, 16, 0, 2)), + (8192, _tiled(128, 128, 256, 1, True, 4, 2, 16, 2, 2)), + ], # wq_b, TP4 (8192, 1280): [ (1, _packed(8, 32, 256, 2, 4, 3, 16, 1)),