[ROCm][MiniMax-M3] MXFP4 mode for the fused norm / RoPE / KV insert

tests/config/test_kv_cache_dtype.py

@@ -119,6 +119,20 @@ def test_backend_sentinel_properties():
     assert mode.is_nvfp4 is False
 
 
+def test_mxfp4_kv_cache_dtype():
+    """The per-head MXFP4 KV cache: its own quant mode, uint8 pages, 68 B rows."""
+    from vllm.utils.torch_utils import mxfp4_kv_cache_full_dim
+
+    assert is_known_kv_cache_dtype("mxfp4")
+    assert CacheConfig(cache_dtype="mxfp4").cache_dtype == "mxfp4"
+    assert STR_DTYPE_TO_TORCH_DTYPE["mxfp4"] == torch.uint8
+    mode = get_kv_quant_mode("mxfp4")
+    assert mode == KVQuantMode.MXFP4 and mode.is_mxfp4
+    assert not mode.is_nvfp4
+    assert is_quantized_kv_cache("mxfp4")
+    assert mxfp4_kv_cache_full_dim(128) == 64 + 4
+
+
 def test_no_snapshot_of_shared_dtype_registry():
     """The registration-time-mutable registries must never be snapshotted.
 

tests/kernels/attention/minimax_m3_nvfp4_ref.py

@@ -8,6 +8,9 @@
 order, V data ``[token octet][dim octet][8 dims][4 B]`` (a byte holds two
 consecutive tokens of one dim), V scales ``[group][token]``. An index page is
 ``[128 x 64 B data][128 x 8 B scales]``.
+
+MXFP4 pages (``mxfp4=True``) keep the data layout with E8M0 scales per 32 dims
+(512 B per head): K and index-K scales ``[token][4]``, V scales ``[group][token]``.
 """
 
 import torch
@@ -31,16 +34,32 @@ def quantize_nvfp4(x: torch.Tensor, global_scale: float):
     scale_e4m3 = scale.to(torch.float8_e4m3fn)
     scale_value = scale_e4m3.float()
     quant = torch.where(scale_value != 0, 1.0 / (scale_value * gs), 0.0)
-    y = groups * quant[..., None]
-    grid = torch.tensor(E2M1_GRID, device=x.device)
+    codes = _e2m1_codes(groups * quant[..., None])
+    return codes.view(-1, D), scale_e4m3.view(torch.uint8).view(-1, 8)
+
+
+def quantize_mxfp4(x: torch.Tensor):
+    """E2M1 codes ``[T, 128]`` and E8M0 scale bytes ``[T, 4]`` of fp32 rows: the
+    scale is 2^ceil(log2(amax / 6)), clamped to 2^-126, and 2^0 for zeros."""
+    groups = x.view(-1, 4, 32)
+    s = groups.abs().amax(-1) / 6.0
+    mant, exp = torch.frexp(s)
+    exp = torch.where(mant == 0.5, exp - 1, exp).clamp(min=-126)
+    exp = torch.where(s == 0, 0, exp)
+    codes = _e2m1_codes(groups / torch.exp2(exp.float())[..., None])
+    return codes.view(-1, D), (exp + 127).to(torch.uint8)
+
+
+def _e2m1_codes(y: torch.Tensor) -> torch.Tensor:
+    """Nearest E2M1 code of each value, ties to the even code."""
+    grid = torch.tensor(E2M1_GRID, device=y.device)
     dist = (y.abs()[..., None] - grid).abs()
     is_min = dist == dist.amin(-1, keepdim=True)
-    ids = torch.arange(8, device=x.device)
+    ids = torch.arange(8, device=y.device)
     # Neighbouring codes alternate parity, so every tie has an even candidate.
     even_min = torch.where(is_min & (ids % 2 == 0), ids, 8).amin(-1)
     mag = torch.where(even_min < 8, even_min, dist.argmin(-1))
-    codes = mag | (torch.signbit(y).long() << 3)
-    return codes.view(-1, D), scale_e4m3.view(torch.uint8).view(-1, 8)
+    return mag | (torch.signbit(y).long() << 3)
 
 
 def dequant(data: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
@@ -52,33 +71,37 @@ def dequant(data: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
     return (mag.unflatten(-1, (8, 16)) * s[..., None]).flatten(-2)
 
 
-def head_offsets(side: int, device):
+def head_offsets(side: int, device, mxfp4: bool = False):
     """Byte offsets inside one head's regions of side 0 (K) or 1 (V).
 
     Returns the data byte and nibble shift of each (token, dim) ``[128, 128]``
-    and the scale byte of each (token, group) ``[128, 8]``.
+    and the scale byte of each (token, group) ``[128, 8 (MXFP4: 4)]``.
     """
+    ng = 4 if mxfp4 else 8
     t = torch.arange(BLOCK, device=device)[:, None]
     d = torch.arange(D, device=device)[None, :]
-    g = torch.arange(8, device=device)[None, :]
+    g = torch.arange(ng, device=device)[None, :]
     if side == 0:
         r = t % 32
         n = 4 * (r // 8) + r % 4
         data = t * 64 + d // 2
         shift = (d % 2) * 4 + 0 * t
         scale = ((((t // 32) * 4 + g // 2) * 16 + n) * 2 + (r % 8) // 4) * 2 + g % 2
+        if mxfp4:
+            scale = t * ng + g
     else:
         data = ((t // 8 * 16 + d // 8) * 8 + d % 8) * 4 + (t % 8) // 2
         shift = (t % 2) * 4 + 0 * d
         scale = g * BLOCK + t
     return data, shift, scale
 
 
-def _region_bases(side: int, heads: int, device):
+def _region_bases(side: int, heads: int, device, mxfp4: bool = False):
     """``[H, 1, 1]`` page offsets of each head's data and scale regions."""
+    ng = 4 if mxfp4 else 8
     h = torch.arange(heads, device=device)[:, None, None]
-    side0 = side * heads * BLOCK * 72
-    return side0 + h * BLOCK * 64, side0 + heads * BLOCK * 64 + h * BLOCK * 8
+    side0 = side * heads * BLOCK * (64 + ng)
+    return side0 + h * BLOCK * 64, side0 + heads * BLOCK * 64 + h * BLOCK * ng
 
 
 def _flat(cache: torch.Tensor) -> torch.Tensor:
@@ -87,16 +110,20 @@ def _flat(cache: torch.Tensor) -> torch.Tensor:
     return cache.as_strided((n,), (1,))
 
 
-def _insert_side(cache, rows, slots, side, global_scale):
+def _quantize(x, global_scale, mxfp4):
+    return quantize_mxfp4(x) if mxfp4 else quantize_nvfp4(x, global_scale)
+
+
+def _insert_side(cache, rows, slots, side, global_scale, mxfp4=False):
     tokens, heads, _ = rows.shape
-    codes, scales = quantize_nvfp4(rows.reshape(-1, D).float(), global_scale)
+    codes, scales = _quantize(rows.reshape(-1, D).float(), global_scale, mxfp4)
     slots = slots.long()
     keep = slots >= 0
     codes = codes.view(tokens, heads, D)[keep].to(torch.uint8)
-    scales = scales.view(tokens, heads, 8)[keep]
+    scales = scales.view(tokens, heads, -1)[keep]
     blk, pos = slots[keep] // BLOCK, slots[keep] % BLOCK
-    data_off, shift, scale_off = head_offsets(side, rows.device)
-    data_base, scale_base = _region_bases(side, heads, rows.device)
+    data_off, shift, scale_off = head_offsets(side, rows.device, mxfp4)
+    data_base, scale_base = _region_bases(side, heads, rows.device, mxfp4)
     page = (blk * cache.stride(0))[:, None, None]
     byte = page + data_base[:, 0] + data_off[pos][:, None]  # [n, H, 128]
     sh = shift[pos][:, None].expand_as(byte)
@@ -107,15 +134,16 @@ def _insert_side(cache, rows, slots, side, global_scale):
     flat[page + scale_base[:, 0] + scale_off[pos][:, None]] = scales
 
 
-def insert_nvfp4_kv(k, v, kv_cache, slots, k_scale, v_scale):
+def insert_nvfp4_kv(k, v, kv_cache, slots, k_scale, v_scale, mxfp4=False):
     """Quantize K / V ``[T, Hkv, 128]`` into their pages; slots < 0 skip."""
-    _insert_side(kv_cache, k, slots, 0, float(k_scale))
-    _insert_side(kv_cache, v, slots, 1, float(v_scale))
+    _insert_side(kv_cache, k, slots, 0, float(k_scale), mxfp4)
+    _insert_side(kv_cache, v, slots, 1, float(v_scale), mxfp4)
 
 
-def insert_nvfp4_index(index_cache, keys, slots):
+def insert_nvfp4_index(index_cache, keys, slots, mxfp4=False):
     """Quantize index keys ``[T, 128]`` into their pages (no global scale)."""
-    codes, scales = quantize_nvfp4(keys.float(), 1.0)
+    codes, scales = _quantize(keys.float(), 1.0, mxfp4)
+    ng = scales.shape[1]
     slots = slots.long()
     keep = slots >= 0
     codes, scales = codes[keep].to(torch.uint8), scales[keep]
@@ -125,9 +153,9 @@ def insert_nvfp4_index(index_cache, keys, slots):
     flat[base + pos[:, None] * 64 + torch.arange(64, device=keys.device)] = codes[
         :, 0::2
     ] | (codes[:, 1::2] << 4)
-    flat[base + BLOCK * 64 + pos[:, None] * 8 + torch.arange(8, device=keys.device)] = (
-        scales
-    )
+    flat[
+        base + BLOCK * 64 + pos[:, None] * ng + torch.arange(ng, device=keys.device)
+    ] = scales
 
 
 def nvfp4_kv_cache_token_major(kv_cache: torch.Tensor):

tests/kernels/attention/test_minimax_m3_norm_rope_nvfp4_insert.py

@@ -46,9 +46,11 @@ def _request_slots(starts, lens, num_pages, gen):
 
 
 class _Batch:
-    def __init__(self, starts, lens, pad, heads, scales, seed=0):
+    def __init__(self, starts, lens, pad, heads, scales, seed=0, mxfp4=False):
         g = torch.Generator().manual_seed(seed)
         self.nq, self.nkv, self.niq = heads
+        self.mxfp4 = mxfp4
+        W = self.W = 68 if mxfp4 else 72
         n = self.n = sum(lens) + pad
         row = (self.nq + 2 * self.nkv + self.niq + 1) * D
         # magnitudes 1e-3 .. 3e3 per 16-value group, some all-zero groups
@@ -72,19 +74,20 @@ def __init__(self, starts, lens, pad, heads, scales, seed=0):
             [_request_slots(starts, lens, pages, g), pad_slots]
         ).cuda()
         # stale bytes; the layer views have a block stride past one page
-        main = 2 * self.nkv * BLK * 72
+        main = 2 * self.nkv * BLK * W
         self.kv_full = torch.randint(0, 256, (pages, main + 256), generator=g)
         self.kv_full = self.kv_full.to(torch.uint8).cuda()
-        self.idx_full = torch.randint(0, 256, (pages, BLK * 72 + 64), generator=g)
+        self.idx_full = torch.randint(0, 256, (pages, BLK * W + 64), generator=g)
         self.idx_full = self.idx_full.to(torch.uint8).cuda()
         self.ks = torch.tensor(scales[0], dtype=torch.float32, device="cuda")
         self.vs = torch.tensor(scales[1], dtype=torch.float32, device="cuda")
 
     def kv(self, full):
-        return full[:, : 2 * self.nkv * BLK * 72].view(-1, 2 * self.nkv, BLK, 72)
+        W = self.W
+        return full[:, : 2 * self.nkv * BLK * W].view(-1, 2 * self.nkv, BLK, W)
 
     def idx(self, full):
-        return full[:, : BLK * 72].view(-1, BLK, 72)
+        return full[:, : BLK * self.W].view(-1, BLK, self.W)
 
     def reference(self):
         qkv = self.qkv.clone()
@@ -104,16 +107,19 @@ def reference(self):
             self.slots,
             self.ks,
             self.vs,
+            self.mxfp4,
+        )
+        insert_nvfp4_index(
+            self.idx(idx_full), qkv[:, ik0 : ik0 + D], self.islots, self.mxfp4
         )
-        insert_nvfp4_index(self.idx(idx_full), qkv[:, ik0 : ik0 + D], self.islots)
         return q, iq, kv_full, idx_full
 
     def run(self, qkv, pos, slots, islots, kv_full, idx_full, q, iq):
         minimax_m3_norm_rope_nvfp4_insert(
             qkv, self.w[0], self.w[1], self.w[2], self.w[3], self.cs, pos,
             slots, islots, self.kv(kv_full), self.idx(idx_full), self.ks,
             self.vs, q, iq, num_heads=self.nq, num_kv_heads=self.nkv,
-            num_index_heads=self.niq, eps=EPS,
+            num_index_heads=self.niq, eps=EPS, mxfp4=self.mxfp4,
         )  # fmt: skip
 
     def fused(self, iq):
@@ -161,9 +167,12 @@ def _check(batch, iq_mode="contiguous"):
         ([0, 100], [4096, 4000], 3),  # the prefill tier
     ],
 )
-@pytest.mark.parametrize("scales", [(1.0, 1.0), (0.002, 3.0)])
-def test_matches_two_kernel_path(heads, starts, lens, pad, scales):
-    _check(_Batch(starts, lens, pad, heads, scales))
+@pytest.mark.parametrize(
+    ("scales", "mxfp4"),
+    [((1.0, 1.0), False), ((0.002, 3.0), False), ((1.0, 1.0), True)],
+)
+def test_matches_two_kernel_path(heads, starts, lens, pad, scales, mxfp4):
+    _check(_Batch(starts, lens, pad, heads, scales, mxfp4=mxfp4))
 
 
 @pytest.mark.parametrize("iq_mode", ["rank", "cols"])
@@ -176,19 +185,21 @@ def test_64k_tokens():
     _check(_Batch([0, 777], [32768, 32768], 0, TP4, (1.0, 1.0)))
 
 
-def test_all_padding():
+@pytest.mark.parametrize("mxfp4", [False, True])
+def test_all_padding(mxfp4):
     """A graph-capture dummy batch: every slot -1, the pages stay as they were."""
-    batch = _Batch([], [], 64, TP4, (1.0, 1.0))
+    batch = _Batch([], [], 64, TP4, (1.0, 1.0), mxfp4=mxfp4)
     _check(batch)
     iq = torch.empty(batch.n, D, dtype=torch.float8_e4m3fn, device="cuda")
     _, kv, idx = batch.fused(iq)
     assert torch.equal(kv, batch.kv_full) and torch.equal(idx, batch.idx_full)
 
 
-def test_cuda_graph_replay():
+@pytest.mark.parametrize("mxfp4", [False, True])
+def test_cuda_graph_replay(mxfp4):
     """Captured once, replayed on new qkv / positions / slots in place."""
-    a = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=1)
-    b = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=2)
+    a = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=1, mxfp4=mxfp4)
+    b = _Batch(*DSPARK, 4, TP4, (1.0, 1.0), seed=2, mxfp4=mxfp4)
     b.w, b.cs = a.w, a.cs  # captured with the layer's weights and RoPE cache
     qkv, pos = a.qkv.clone(), a.pos.clone()
     slots, islots = a.slots.clone(), a.islots.clone()

vllm/config/cache.py

@@ -64,6 +64,7 @@ def _layout_from_name(layout_name: str) -> KVCacheLayout:
     "fp8_per_token_head",
     "nvfp4",
     "nvfp4_4over6",
+    "mxfp4",
 ]
 
 # Mutable mirror of the CacheDType Literal. Custom dtypes registered by

vllm/models/minimax_m3/amd/ops/norm_rope_nvfp4_insert/host.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Launch wrapper of the fused norm / RoPE / NVFP4 insert."""
+"""Launch wrapper of the fused norm / RoPE / NVFP4 or MXFP4 insert."""
 
 import functools
 import struct
@@ -10,13 +10,14 @@
 from vllm.models.minimax_m3.amd.ops.flydsl_utils import (
     run_compiled,
 )
+from vllm.utils.torch_utils import mxfp4_kv_cache_full_dim, nvfp4_kv_cache_full_dim
 
 from . import kernel as _kernel
 
 
 @functools.cache
-def _get_kernel(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int):
-    return _kernel.compile_norm_rope_insert(nq, nkv, niq, rotary_dim, tp)
+def _get_kernel(nq: int, nkv: int, niq: int, rotary_dim: int, tp: int, mx: bool):
+    return _kernel.compile_norm_rope_insert(nq, nkv, niq, rotary_dim, tp, mx)
 
 
 def _token_pairs_per_wave(n: int) -> int:
@@ -33,23 +34,26 @@ def minimax_m3_norm_rope_nvfp4_insert(
     positions: torch.Tensor,  # [n] int64
     slot_mapping: torch.Tensor,  # [n] int64, -1 skips the token
     index_slot_mapping: torch.Tensor,  # [n] int64, -1 skips the token
-    kv_cache: torch.Tensor,  # [blocks, 2 nkv, 128, 72] uint8
-    index_cache: torch.Tensor,  # [blocks, 128, 72] uint8
-    k_scale: torch.Tensor,  # fp32 global scales
+    kv_cache: torch.Tensor,  # [blocks, 2 nkv, 128, 72 (MXFP4: 68)] uint8
+    index_cache: torch.Tensor,  # [blocks, 128, 72 (MXFP4: 68)] uint8
+    k_scale: torch.Tensor,  # fp32 global scales, unused for MXFP4
     v_scale: torch.Tensor,
     q_out: torch.Tensor,  # [n, nq * 128] bf16
     index_q_out: torch.Tensor,  # [n, niq * 128] e4m3, row stride free
     num_heads: int,
     num_kv_heads: int,
     num_index_heads: int,
     eps: float,
+    mxfp4: bool = False,
 ) -> None:
-    """Norm + RoPE of q / k / index-q / index-k and the NVFP4 K/V and index-K
-    inserts in one launch; k and index-k are not written back into ``qkv``."""
+    """Norm + RoPE of q / k / index-q / index-k and the NVFP4 (``mxfp4``: MXFP4)
+    K/V and index-K inserts in one launch; k and index-k are not written back
+    into ``qkv``."""
     n = qkv.shape[0]
     if n == 0:
         return
     nq, nkv, niq = num_heads, num_kv_heads, num_index_heads
+    slot_w = (mxfp4_kv_cache_full_dim if mxfp4 else nvfp4_kv_cache_full_dim)(_kernel.D)
     rotary_dim = cos_sin_cache.shape[1]
     assert eps >= 2.0**-126  # the kernel's rsqrt skips ocml's denormal scaling
     assert qkv.dtype == torch.bfloat16 and qkv.is_contiguous()
@@ -60,12 +64,12 @@ def minimax_m3_norm_rope_nvfp4_insert(
     assert q_out.dtype == torch.bfloat16 and q_out.is_contiguous()
     assert index_q_out.dtype == torch.float8_e4m3fn and index_q_out.stride(1) == 1
     assert index_q_out.stride(0) % 4 == 0
-    assert kv_cache.shape[1:] == (2 * nkv, _kernel.BLK, 72)
-    assert kv_cache.stride(1) == _kernel.BLK * 72 and kv_cache.stride(3) == 1
-    assert index_cache.shape[1:] == (_kernel.BLK, 72) and index_cache.stride(2) == 1
-    assert index_cache.stride(1) == 72
+    assert kv_cache.shape[1:] == (2 * nkv, _kernel.BLK, slot_w)
+    assert kv_cache.stride(1) == _kernel.BLK * slot_w and kv_cache.stride(3) == 1
+    assert index_cache.shape[1:] == (_kernel.BLK, slot_w) and index_cache.stride(2) == 1
+    assert index_cache.stride(1) == slot_w
     run_compiled(
-        _get_kernel(nq, nkv, niq, rotary_dim, _token_pairs_per_wave(n)),
+        _get_kernel(nq, nkv, niq, rotary_dim, _token_pairs_per_wave(n), mxfp4),
         qkv.data_ptr(),
         q_out.data_ptr(),
         index_q_out.data_ptr(),