Skip to content

[ROCm][Perf][DSv4.1] TP1 MXFP8 GEMM tiles and a pre-quantized wo_b input for DEP8 on MI355X #63821

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maslan-coop26 wants to merge 10 commits into
base: main
Choose a base branch
from
+560 −43
Open

[ROCm][Perf][DSv4.1] TP1 MXFP8 GEMM tiles and a pre-quantized wo_b input for DEP8 on MI355X #63821

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
1ae3465
[ROCm][Perf][DSv4.1] TP1 MXFP8 GEMM tiles and a pre-quantized wo_b in…
maslan-coop26 Oct 26, 2026
621c4e5
[ROCm][DSv4.1] Tie the wo_b hand-off to its kernel and test every fus…
maslan-coop26 Oct 29, 2026
825d04d
Merge remote-tracking branch 'origin/main' into rocm-dsv41-tp1-oproj-…
maslan-coop26 Oct 29, 2026
6392134
Merge branch 'main' into rocm-dsv41-tp1-oproj-gemms
Fangzhou-Ai Oct 30, 2026
98c7dcb
[ROCm][DSv4.1] Keep wo_b's fused quant where the wo_a tile cannot qua…
maslan-coop26 Oct 30, 2026
8902eb3
Merge branch 'main' into rocm-dsv41-tp1-oproj-gemms
mergify[bot] Oct 31, 2026
3d2b5cd
Merge branch 'main' into rocm-dsv41-tp1-oproj-gemms
vllm-agent Oct 31, 2026
fb01225
Merge branch 'main' into rocm-dsv41-tp1-oproj-gemms
mergify[bot] Nov 2, 2026
82daaac
[ROCm][Perf][DSv4.1] block32 tiers for the DSpark drafter's main_proj
maslan-coop26 Nov 6, 2026
1d99585
Merge branch 'main' into rocm-dsv41-tp1-oproj-gemms
mergify[bot] Nov 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .buildkite/test_areas/kernels.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,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
Expand Down
145 changes: 133 additions & 12 deletions tests/kernels/attention/test_rocm_triton_attn_dsv4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
)
Expand All @@ -1541,6 +1545,123 @@ 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(

@okorzh-amd okorzh-amd Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The description says these tests run in ":amd: (MI355 DPX) MLA Kernels", but that step's AMD mirror (.buildkite/test_areas/kernels.yaml:437) runs only test_dsv41_adaptive_mla_swa_metadata_graph_replay from this file, and the MI300 Attention shards skip them (requires_gfx950). The only gfx950 step that collects test_rocm_mxfp8_wo_a_bmm and this new test is the optional ":amd: (MI355 DPX) Attention Kernels Shard" (test-amd.yaml:4469), so PR CI never runs the epilogue. Could you add pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py -k wo_a_bmm to that mirror's commands (its source deps already list both changed files), or correct the description? The existing wo_a test has the same gap on main

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 621c4e5. The MI355 MLA Kernels mirror now runs the wo_a tests and the new gate test from this file.

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)


@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()
Expand Down
8 changes: 6 additions & 2 deletions tests/kernels/quantization/test_rocm_mxfp8_linear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,11 +180,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),
Expand All @@ -211,7 +215,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.
Expand Down
61 changes: 60 additions & 1 deletion vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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): [

@okorzh-amd okorzh-amd Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#62533 adds the same (4608, 5120), (5120, 8192) and (32768, 1280) keys at other positions in _TUNED. The two PRs merge with no git conflict so the second to land leaves repeated keys: ruff F601 then fails pre-commit on main, and at runtime the later literal silently wins. One of the two needs to drop its copies, or fold the tables into #62533 as you offered. cc karverma-amd

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#62533 was closed, so that clash is gone. But #64459 (opened today) adds the same four shapes. I'm benchmarking both tables and will propose one merged table, so only one PR carries these keys.

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Compared both tables (numbers in the description and on #64459): #64459 is faster at small batches, this PR past 512 tokens. This PR keeps its table; both add the same four keys, so whichever lands second drops its copies.

@okorzh-amd okorzh-amd Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Per the table in the description the two win in different ranges: #64459 up to 2.01x at 1-32 tokens, this PR 1.08-1.19x past 640. Dropping either copy loses one side; could the second to land merge the tiers by M range instead? (Your 1-32 tiers equal main's, so the 1.08x there is noise.)

@maslan-coop26 maslan-coop26 Oct 30, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed: the second to land merges the tiers by M range instead of dropping either copy. We have the per-M merged table from the side-by-side run. And yes, our 1–32 tiers are main's defaults, so the 1.08x there is noise.

(8, _packed(8, 16, 256, 2, 2, 3, 16, 0)),

@okorzh-amd okorzh-amd Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

description needs perf and sccuracy check for each kernel you are adding on a microbench level

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a per-kernel table to the description (one MI355X, cold weights). At our decode sizes (224–512 tokens per rank), every new tile is faster than main, up to 2.2x. The only slower point is wo_b at 1024 tokens (1% slower). Accuracy is the same as main: max relative error 1.7e-3 everywhere, and the fused wo_b input is bit-identical.

(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)),
Expand Down Expand Up @@ -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)),
Expand All @@ -629,6 +657,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)),
Expand Down Expand Up @@ -663,6 +708,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)),
],
}


Expand Down
51 changes: 51 additions & 0 deletions vllm/models/deepseek_v41/amd/rocm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (

@okorzh-amd okorzh-amd Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The output tok/s (17,557) and tok/s per user (91.12) gains each come from one run. The excluded run measured 16,711, below all three baselines, and the "wo_b tiles, wo_a tiles and hand-off only" run (17,236) matches the best baseline (17,238). Median TPOT (7.14 and 7.11 vs 7.34-7.39 ms) is the only end-to-end metric that separates on both runs. Could you add two more runs of this PR, or lead with TPOT?

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. We don't have more runs of this exact build, so the description now leads with TPOT: 7.14 and 7.11 ms on both runs of this PR, against 7.34–7.39 ms on all three baselines. The tok/s numbers are now marked as one run.

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 (
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -931,6 +935,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():

@okorzh-amd okorzh-amd Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now that #63179 is in the base, wo_b's own linear fuses the activation quant into the block32 GEMM at M <= 8 (rocm_native.py:457, eligibility at rocm_block32_gemm.py:1032). When this gate holds, wo_b's linear is never called. At 1-8 tokens the wo_a tile is narrower than 32 columns, so rocm_dsv41_mxfp8_wo_a_bmm_mxfp8 runs the bf16 BMM plus a separate mxfp8_e4m3_quantize (rocm_aiter_mla_sparse.py:2179), and this branch then calls the unfused rocm_mxfp8_block32_gemm: one launch per layer more than main. That covers every TP1 layer at 8 or fewer tokens per rank (e.g. DP8 with one request per rank and DSpark K=6), and the TP2/TP4/TP8 monokernel layers that run only the FFN launch (model.py:288), where TP4 wo_b has a tuned fused tile (rocm_block32_gemm.py:766). The description says the opposite ("Below that, wo_b's linear quantizes the activation as it does on main"). Suggest moving the wo_b GEMM into the opaque op so its narrow-tile branch calls rocm_mxfp8_block32_gemm_quant on the bf16 wo_a output. Otherwise, could you show an A/B against main at 1-8 tokens?

@maslan-coop26 maslan-coop26 Oct 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 98c7dcb. o_proj now runs one opaque op that takes the hand-off only where the wo_a tile can quantize (T·G > 128). Below that, the bf16 wo_a output goes to wo_b's own rocm_mxfp8_block32_gemm_quant, so #63179's fused quant still applies at tiny batches. The new test checks the path on both sides of T·G = 128 and that the output is bit-identical to wo_b's linear. I also corrected the description.

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,
Expand All @@ -953,6 +972,38 @@ def _o_proj(
)
return self._wo_b_after_wo_a(z.flatten(1))

def _wo_b_takes_wo_a_mxfp8(self) -> bool:

@okorzh-amd okorzh-amd Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This gate re-derives wo_b's kernel dispatch (the block32 branch at rocm_native.py:277) from dtypes and shapes, and nothing ties the two: no test reaches this branch or the gate. It gives the right answer for every ROCm MXFP8 kernel today, but a kernel-side relayout that keeps an e4m3 [N, K] weight with uint8 [N/32, K/32] scales still passes it. For example, draft #60169's in-place (16, 16) shuffle in RocmDotScaledMxfp8LinearKernel.process_weights_after_loading would hand rocm_mxfp8_block32_gemm a permuted wo_b on every layer, with no error. Check the consumer the way can_fuse_query_quant does (common/ops/query_quant.py:146-150): and type(getattr(getattr(self.wo_b, "quant_method", None), "kernel", None)) is RocmDotScaledMxfp8LinearKernel

@maslan-coop26 maslan-coop26 Oct 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in 621c4e5. The gate now checks the kernel type, as you suggested. A new test covers the gate, including a different kernel class with the same weight layout: it does not get the fused input.

"""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(
Expand Down
Loading
Loading