diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 2345d78e82be..a12f0cf5da49 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -523,6 +523,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" 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 0c37a127c5a1..4d2457ab4a3d 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1593,24 +1593,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, @@ -1622,7 +1626,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) ) @@ -1648,6 +1652,196 @@ 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 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() +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( + "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", + [ + ("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 b1f1302957fc..9c5876ed1079 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -280,17 +280,22 @@ 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), ("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), ] @@ -311,7 +316,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/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py index 9effa2b50bd2..9617926b575c 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py @@ -456,7 +456,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. A tier serves +# linears at TP1/TP2/TP4, tuned on MI355X under HIP graphs. 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 @@ -550,6 +550,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)), @@ -607,6 +621,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)), @@ -629,6 +657,34 @@ 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)), + ], + # 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)), @@ -663,6 +719,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 3d694b2b9106..416acf1ae5d1 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -56,6 +56,7 @@ rocm_inv_rope_einsum, rocm_inverse_rope_mxfp8_rows, rocm_inverse_rope_rows_, + rocm_mxfp8_o_proj, rocm_mxfp8_wo_a_bmm, rocm_sparse_attn_decode, rocm_sparse_attn_prefill, @@ -931,6 +932,16 @@ 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(): + return rocm_mxfp8_o_proj( + attn_out.data, + attn_out.scale, + self.wo_a, + self.n_local_groups, + self.o_lora_rank, + self.wo_b, + attn_out.orig_dtype, + ) z = rocm_mxfp8_wo_a_bmm( attn_out.data, attn_out.scale, @@ -953,6 +964,38 @@ 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 + 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 ( + 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 + 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 93f509240849..6be06dc153a7 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1967,19 +1967,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. @@ -2011,6 +2014,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), @@ -2018,76 +2036,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 @@ -2133,6 +2185,158 @@ 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 + ) + + +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