@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b8cf275382

[AMD][Minimax-M3][perf] Enable packed LBHNC AITER QK-norm fusion for MiniMax-M3 on ROCm (#54535)

Signed-off-by: Wei-Ting Liao <wei-ting.liao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Shanshan Shen <467638484@qq.com>

weitliao committed Sep 21, 2026 at 01:38 UTC b8cf2753825d05d7f41eb4747480b78934953aac
3 files changed +636 -91
tests/kernels/attention/test_minimax_m3.py
+361 -10
@@ -19,7 +19,10 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
19 minimax_m3_sparse_attn,
20 minimax_m3_sparse_attn_decode,
21 )
22 -from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
22 +from vllm.models.minimax_m3.common.sparse_attention import (
23 + MiniMaxM3SparseTritonImpl,
24 + minimax_m3_rebase_slots_to_page16,
25 +)
26 from vllm.platforms import current_platform
27 from vllm.utils.torch_utils import set_random_seed
28 from vllm.v1.attention.backends.utils import record_kv_cache_layout
@@ -28,6 +31,7 @@ from vllm.v1.kv_cache_interface import (
31 KVCacheLayout,
32 MLAAttentionSpec,
33 compute_layer_kv_cache_shape_bytes,
34 + get_kv_quant_mode,
35 )
36
37 if not (current_platform.is_cuda() or current_platform.is_rocm()):
@@ -1672,8 +1676,9 @@ def test_main_cache_layout_contract():
1676 assert set(order) == set(range(len(order)))
1677
1678
1675 -def test_aiter_sparse_pa_layout_contract(monkeypatch):
1676 - """The shuffle-only AITER path retains separately contiguous K/V storage."""
1679 +@pytest.mark.parametrize("layout", [KVCacheLayout.LHBNC, KVCacheLayout.LBHNC])
1680 +def test_aiter_sparse_pa_layout_contract(monkeypatch, layout):
1681 + """Build the exact separate and packed page-16 views used by the model."""
1682 import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
1683
1684 monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
@@ -1685,21 +1690,43 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
1690
1691 backend = sparse_attn_mod.MiniMaxM3SparseBackend
1692 assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts()
1693 + assert KVCacheLayout.LBHNC in backend.supported_kv_cache_layouts()
1694
1689 - nb, h = 7, 1
1695 + nb, h = 3, 1
1696 + cache_dtype = current_platform.fp8_dtype()
1697 spec = FullAttentionSpec(
1698 block_size=BLOCK_SIZE,
1699 num_kv_heads=h,
1700 head_size=HEAD_DIM,
1694 - dtype=DTYPE,
1701 + dtype=cache_dtype,
1702 + kv_quant_mode=get_kv_quant_mode("fp8"),
1703 num_head_slots=2,
1696 - state_content_bytes=h * HEAD_DIM * DTYPE.itemsize,
1704 + state_content_bytes=h * HEAD_DIM * cache_dtype.itemsize,
1705 )
1706 raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
1699 - view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0]
1700 - key_cache, value_cache = view.unbind(1)
1701 - assert key_cache.is_contiguous()
1702 - assert value_cache.is_contiguous()
1707 + kv_cache = dense_kv_cache_views(raw, spec, nb, 1, layout)[0].view(cache_dtype)
1708 + key_cache, value_cache = kv_cache.unbind(1)
1709 + x = 16 // cache_dtype.itemsize
1710 + pages_per_side = BLOCK_SIZE // 16
1711 + if key_cache.is_contiguous() and value_cache.is_contiguous():
1712 + k_src, v_src = key_cache, value_cache
1713 + block_pages, v_page_offset = pages_per_side, 0
1714 + else:
1715 + assert kv_cache.is_contiguous()
1716 + k_src = v_src = kv_cache
1717 + block_pages, v_page_offset = 2 * pages_per_side, pages_per_side
1718 +
1719 + num_phys16 = nb * block_pages
1720 + k_view = k_src.view(num_phys16, h, HEAD_DIM // x, 16, x)
1721 + v_view = v_src.view(num_phys16, h, 16 // x, HEAD_DIM, x)[v_page_offset:]
1722 + assert k_view.is_contiguous()
1723 + assert v_view.is_contiguous()
1724 + if layout is KVCacheLayout.LHBNC:
1725 + assert k_view.shape == (24, 1, 8, 16, 16)
1726 + assert v_view.shape == (24, 1, 1, 128, 16)
1727 + else:
1728 + assert k_view.shape == (48, 1, 8, 16, 16)
1729 + assert v_view.shape == (40, 1, 1, 128, 16)
1730
1731
1732 def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
@@ -1769,6 +1796,330 @@ def test_aiter_indexer_requires_the_aiter_attend(monkeypatch):
1796 assert reason is None or "AITER sparse PA attend" not in reason
1797
1798
1799 +def test_aiter_consolidated_qknorm_enabled(monkeypatch):
1800 + from vllm._aiter_ops import rocm_aiter_ops
1801 +
1802 + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
1803 + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1804 +
1805 + scale = torch.tensor([0.25])
1806 + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1807 + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, scale)
1808 + assert rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e4m3", scale, scale)
1809 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8")
1810 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8", scale, None)
1811 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
1812 +
1813 + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", False)
1814 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1815 +
1816 + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1817 + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: False)
1818 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("auto")
1819 +
1820 +
1821 +def _fake_aiter_qknorm_args(kv_cache_dtype: str = "auto") -> dict:
1822 + tensor = torch.empty(1)
1823 + return {
1824 + "qkv": tensor,
1825 + "q_norm_weight": tensor,
1826 + "k_norm_weight": tensor,
1827 + "cos_sin_cache": tensor,
1828 + "positions": tensor,
1829 + "num_heads": 16,
1830 + "num_kv_heads": 1,
1831 + "rotary_dim": 64,
1832 + "eps": 1e-6,
1833 + "slot_mapping": tensor,
1834 + "kv_cache_k": tensor,
1835 + "kv_cache_v": tensor,
1836 + "q_out": tensor,
1837 + "kv_cache_dtype": kv_cache_dtype,
1838 + }
1839 +
1840 +
1841 +def _install_fake_aiter_qknorm(monkeypatch: pytest.MonkeyPatch) -> list[dict]:
1842 + """Stub ``aiter.fused_qknorm_idxrqknorm`` so mapping tests run without AITER."""
1843 + import sys
1844 + import types
1845 + from typing import Any
1846 +
1847 + calls: list[dict] = []
1848 +
1849 + def fake_op(*args, **kwargs):
1850 + calls.append(kwargs)
1851 +
1852 + existing = sys.modules.get("aiter")
1853 + if existing is not None:
1854 + monkeypatch.setattr(existing, "fused_qknorm_idxrqknorm", fake_op, raising=False)
1855 + else:
1856 + fake_mod: Any = types.ModuleType("aiter")
1857 + fake_mod.fused_qknorm_idxrqknorm = fake_op
1858 + monkeypatch.setitem(sys.modules, "aiter", fake_mod)
1859 + return calls
1860 +
1861 +
1862 +def test_aiter_consolidated_qknorm_packed_shuffle_unequal_spans(monkeypatch):
1863 + from vllm._aiter_ops import rocm_aiter_ops
1864 +
1865 + calls = _install_fake_aiter_qknorm(monkeypatch)
1866 + call_args = _fake_aiter_qknorm_args()
1867 + call_args["kv_cache_k"] = torch.empty(2, 1)
1868 + call_args["kv_cache_v"] = torch.empty(1, 1)
1869 +
1870 + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1871 + assert len(calls) == 1
1872 +
1873 +
1874 +@pytest.mark.parametrize(
1875 + ("kv_cache_dtype", "aiter_dtype", "passes_scales"),
1876 + [
1877 + ("auto", "auto", False),
1878 + ("fp8", "fp8_e4m3_static", True),
1879 + ("fp8_e4m3", "fp8_e4m3_static", True),
1880 + ],
1881 +)
1882 +def test_aiter_consolidated_qknorm_dtype_and_layout_mapping(
1883 + monkeypatch, kv_cache_dtype, aiter_dtype, passes_scales
1884 +):
1885 + from vllm._aiter_ops import rocm_aiter_ops
1886 +
1887 + calls = _install_fake_aiter_qknorm(monkeypatch)
1888 + call_args = _fake_aiter_qknorm_args(kv_cache_dtype)
1889 + k_scale = torch.tensor([0.25])
1890 + v_scale = torch.tensor([0.5])
1891 + call_args.update(k_scale=k_scale, v_scale=v_scale)
1892 +
1893 + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1894 + assert len(calls) == 1
1895 + kwargs = calls[0]
1896 + assert kwargs["kv_cache_dtype"] == aiter_dtype
1897 + assert kwargs["index_cache_dtype"] == "auto"
1898 + assert kwargs["block_size"] == 16
1899 + assert kwargs["asm_layout"] is True
1900 + if passes_scales:
1901 + assert kwargs["k_scale"] is k_scale
1902 + assert kwargs["v_scale"] is v_scale
1903 + else:
1904 + assert kwargs["k_scale"] is None
1905 + assert kwargs["v_scale"] is None
1906 +
1907 +
1908 +def test_aiter_consolidated_qknorm_fp8_index_dtype_mapping(monkeypatch):
1909 + from vllm._aiter_ops import rocm_aiter_ops
1910 +
1911 + calls = _install_fake_aiter_qknorm(monkeypatch)
1912 + fp8 = current_platform.fp8_dtype()
1913 + call_args = _fake_aiter_qknorm_args("fp8")
1914 + call_args.update(
1915 + k_scale=torch.tensor([0.25]),
1916 + v_scale=torch.tensor([0.5]),
1917 + index_cache=torch.empty(1, dtype=fp8),
1918 + index_q_out=torch.empty(1, dtype=fp8),
1919 + num_index_heads=1,
1920 + )
1921 +
1922 + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1923 + assert calls[0]["index_cache_dtype"] == "fp8"
1924 + assert calls[0]["kv_cache_dtype"] == "fp8_e4m3_static"
1925 +
1926 +
1927 +@pytest.mark.parametrize("skip_index_branch", [False, True])
1928 +def test_aiter_consolidated_qknorm_full_and_skip_index_args(
1929 + monkeypatch, skip_index_branch
1930 +):
1931 + from vllm._aiter_ops import rocm_aiter_ops
1932 +
1933 + calls = _install_fake_aiter_qknorm(monkeypatch)
1934 + call_args = _fake_aiter_qknorm_args()
1935 + if skip_index_branch:
1936 + call_args.update(skip_index_branch=True)
1937 + else:
1938 + tensor = torch.empty(1)
1939 + call_args.update(
1940 + index_q_norm_weight=tensor,
1941 + index_k_norm_weight=tensor,
1942 + num_index_heads=1,
1943 + index_cache=tensor,
1944 + index_q_out=tensor,
1945 + index_slot_mapping=tensor,
1946 + )
1947 +
1948 + rocm_aiter_ops.fused_qknorm_idxrqknorm(**call_args)
1949 + kwargs = calls[0]
1950 + assert kwargs["skip_index_branch"] is skip_index_branch
1951 + if skip_index_branch:
1952 + assert kwargs["index_q_norm_weight"] is None
1953 + assert kwargs["index_k_norm_weight"] is None
1954 + assert kwargs["index_cache"] is None
1955 + assert kwargs["index_q_out"] is None
1956 + assert kwargs["index_slot_mapping"] is None
1957 + else:
1958 + assert kwargs["index_q_norm_weight"] is not None
1959 + assert kwargs["index_k_norm_weight"] is not None
1960 + assert kwargs["index_cache"] is not None
1961 + assert kwargs["index_q_out"] is not None
1962 + assert kwargs["index_slot_mapping"] is not None
1963 +
1964 +
1965 +def test_aiter_consolidated_qknorm_rejects_fp8_e5m2(monkeypatch):
1966 + from vllm._aiter_ops import rocm_aiter_ops
1967 +
1968 + monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
1969 + monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
1970 + scale = torch.tensor([0.25])
1971 + assert not rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled("fp8_e5m2", scale, scale)
1972 +
1973 +
1974 +def test_aiter_sparse_pa_slot_rebase_preserves_page_offsets():
1975 + slots = torch.tensor([0, 15, 16, 127, 128, 129, -1])
1976 + actual = minimax_m3_rebase_slots_to_page16(slots, block_size=128)
1977 + expected = torch.tensor([0, 15, 16, 127, 256, 257, -1])
1978 + assert torch.equal(actual, expected)
1979 +
1980 +
1981 +@pytest.mark.skipif(
1982 + not current_platform.is_rocm(),
1983 + reason="The consolidated packed-SHUFFLE AITER op is ROCm-only",
1984 +)
1985 +@pytest.mark.parametrize("skip_index_branch", [False, True])
1986 +@pytest.mark.parametrize("fp8_index", [False, True])
1987 +def test_aiter_consolidated_qknorm_real_packed_shuffle(skip_index_branch, fp8_index):
1988 + from vllm._aiter_ops import rocm_aiter_ops
1989 +
1990 + torch.manual_seed(20260831)
1991 + num_tokens = 4
1992 + num_heads = 16
1993 + num_kv_heads = 1
1994 + num_index_heads = 1
1995 + rotary_dim = 64
1996 + head_dim = 128
1997 + logical_block_size = 128
1998 + num_blocks = 3
1999 + fp8_dtype = current_platform.fp8_dtype()
2000 +
2001 + qkv = torch.randn(
2002 + num_tokens,
2003 + (num_heads + 2 * num_kv_heads + num_index_heads + 1) * head_dim,
2004 + dtype=torch.bfloat16,
2005 + device="cuda",
2006 + )
2007 + weights = [
2008 + torch.randn(head_dim, dtype=torch.bfloat16, device="cuda") for _ in range(4)
2009 + ]
2010 + cos_sin_cache = torch.randn(16, rotary_dim, dtype=torch.bfloat16, device="cuda")
2011 + positions = torch.arange(num_tokens, dtype=torch.int64, device="cuda")
2012 + logical_slots = torch.tensor([0, 127, 128, -1], dtype=torch.int64, device="cuda")
2013 + slot_mapping = minimax_m3_rebase_slots_to_page16(logical_slots, logical_block_size)
2014 + index_slot_mapping = logical_slots.clone()
2015 +
2016 + packed = torch.zeros(
2017 + num_blocks,
2018 + 2,
2019 + logical_block_size,
2020 + head_dim,
2021 + dtype=fp8_dtype,
2022 + device="cuda",
2023 + )
2024 + x = 16 // packed.element_size()
2025 + pages_per_side = logical_block_size // 16
2026 + num_phys16 = num_blocks * 2 * pages_per_side
2027 + kv_cache_k = packed.view(num_phys16, num_kv_heads, head_dim // x, 16, x)
2028 + kv_cache_v = packed.view(num_phys16, num_kv_heads, 16 // x, head_dim, x)[
2029 + pages_per_side:
2030 + ]
2031 + assert kv_cache_k.shape == (48, 1, 8, 16, 16)
2032 + assert kv_cache_v.shape == (40, 1, 1, 128, 16)
2033 +
2034 + q_out = torch.zeros(
2035 + num_tokens, num_heads * head_dim, dtype=torch.bfloat16, device="cuda"
2036 + )
2037 + index_dtype = fp8_dtype if fp8_index else torch.bfloat16
2038 + if fp8_index:
2039 + index_q_out = torch.empty(
2040 + (num_tokens, num_index_heads * head_dim),
2041 + dtype=index_dtype,
2042 + device="cuda",
2043 + )
2044 + index_cache = torch.empty(
2045 + (num_blocks, logical_block_size, head_dim),
2046 + dtype=index_dtype,
2047 + device="cuda",
2048 + )
2049 + index_q_out.view(torch.uint8).fill_(7)
2050 + index_cache.view(torch.uint8).fill_(7)
2051 + else:
2052 + index_q_out = torch.full(
2053 + (num_tokens, num_index_heads * head_dim),
2054 + 7,
2055 + dtype=torch.bfloat16,
2056 + device="cuda",
2057 + )
2058 + index_cache = torch.full(
2059 + (num_blocks, logical_block_size, head_dim),
2060 + 7,
2061 + dtype=torch.bfloat16,
2062 + device="cuda",
2063 + )
2064 + scale = torch.tensor([0.02], dtype=torch.float32, device="cuda")
2065 +
2066 + kwargs = {}
2067 + if not skip_index_branch:
2068 + kwargs = {
2069 + "index_q_norm_weight": weights[2],
2070 + "index_k_norm_weight": weights[3],
2071 + "index_cache": index_cache,
2072 + "index_q_out": index_q_out,
2073 + "index_slot_mapping": index_slot_mapping,
2074 + }
2075 + rocm_aiter_ops.fused_qknorm_idxrqknorm(
2076 + qkv,
2077 + weights[0],
2078 + weights[1],
2079 + cos_sin_cache,
2080 + positions,
2081 + num_heads,
2082 + num_kv_heads,
2083 + rotary_dim,
2084 + 1e-6,
2085 + slot_mapping,
2086 + kv_cache_k,
2087 + kv_cache_v,
2088 + q_out,
2089 + "fp8",
2090 + scale,
2091 + scale,
2092 + num_index_heads=num_index_heads,
2093 + skip_index_branch=skip_index_branch,
2094 + **kwargs,
2095 + )
2096 + torch.accelerator.synchronize()
2097 +
2098 + assert torch.isfinite(q_out).all()
2099 + assert torch.count_nonzero(q_out) > 0
2100 + assert torch.count_nonzero(packed) > 0
2101 + if fp8_index:
2102 + index_q_untouched = torch.all(index_q_out.view(torch.uint8) == 7)
2103 + index_cache_untouched = torch.all(index_cache.view(torch.uint8) == 7)
2104 + index_cache_t0 = torch.all(index_cache[0, 0].view(torch.uint8) == 7)
2105 + index_cache_t127 = torch.all(index_cache[0, 127].view(torch.uint8) == 7)
2106 + index_cache_t128 = torch.all(index_cache[1, 0].view(torch.uint8) == 7)
2107 + else:
2108 + index_q_untouched = torch.all(index_q_out == 7)
2109 + index_cache_untouched = torch.all(index_cache == 7)
2110 + index_cache_t0 = torch.all(index_cache[0, 0] == 7)
2111 + index_cache_t127 = torch.all(index_cache[0, 127] == 7)
2112 + index_cache_t128 = torch.all(index_cache[1, 0] == 7)
2113 + if skip_index_branch:
2114 + assert index_q_untouched
2115 + assert index_cache_untouched
2116 + else:
2117 + assert not index_q_untouched
2118 + assert not index_cache_t0
2119 + assert not index_cache_t127
2120 + assert not index_cache_t128
2121 +
2122 +
2123 def test_indexer_cache_squeezes_to_contiguous_3d():
2124 """The indexer side cache is standardized 4D with H=1: under both layouts
2125 the allocator's logical view stays contiguous and squeezes (as
vllm/_aiter_ops.py
+112
@@ -35,10 +35,22 @@ except ImportError:
35 # which is a host op, so we cache it once here.
36 FP8_DTYPE = current_platform.fp8_dtype()
37 _HIPB_MM_INITIALIZED_DEVICES: set[int] = set()
38 +_FP8_E4M3_DTYPES = {
39 + dtype
40 + for dtype in (
41 + getattr(torch, "float8_e4m3fn", None),
42 + getattr(torch, "float8_e4m3fnuz", None),
43 + )
44 + if dtype is not None
45 +}
46 KB = 1024
47 MB = 1024 * KB
48
49
50 +def _is_fp8_e4m3_tensor(t: torch.Tensor | None) -> bool:
51 + return t is not None and t.dtype in _FP8_E4M3_DTYPES
52 +
53 +
54 def _get_or_create_aiter_qr_rmsnorm_comm(
55 device_comm: Any,
56 ) -> tuple[int, int, bool] | None:
@@ -2200,6 +2212,106 @@ class rocm_aiter_ops:
2212 def is_shuffle_kv_cache_enabled(cls) -> bool:
2213 return cls._SHUFFLE_KV_CACHE_ENABLED
2214
2215 + @classmethod
2216 + @if_aiter_supported
2217 + def fused_qknorm_idxrqknorm_enabled(
2218 + cls,
2219 + kv_cache_dtype: str,
2220 + k_scale: torch.Tensor | None = None,
2221 + v_scale: torch.Tensor | None = None,
2222 + ) -> bool:
2223 + """Whether MiniMax-M3 can use AITER's consolidated fused QK-norm.
2224 +
2225 + Requires AITER to be installed and enabled. bf16 (``auto``) always
2226 + qualifies; fp8 e4m3 also needs K/V scales. Other cache dtypes fall
2227 + back to vLLM's fused kernel.
2228 + """
2229 + if not cls._AITER_ENABLED:
2230 + return False
2231 + if kv_cache_dtype == "auto":
2232 + return True
2233 + if kv_cache_dtype in ("fp8", "fp8_e4m3"):
2234 + return k_scale is not None and v_scale is not None
2235 + return False
2236 +
2237 + @classmethod
2238 + def fused_qknorm_idxrqknorm(
2239 + cls,
2240 + qkv: torch.Tensor,
2241 + q_norm_weight: torch.Tensor,
2242 + k_norm_weight: torch.Tensor,
2243 + cos_sin_cache: torch.Tensor,
2244 + positions: torch.Tensor,
2245 + num_heads: int,
2246 + num_kv_heads: int,
2247 + rotary_dim: int,
2248 + eps: float,
2249 + slot_mapping: torch.Tensor,
2250 + kv_cache_k: torch.Tensor,
2251 + kv_cache_v: torch.Tensor,
2252 + q_out: torch.Tensor,
2253 + kv_cache_dtype: str,
2254 + k_scale: torch.Tensor | None = None,
2255 + v_scale: torch.Tensor | None = None,
2256 + index_q_norm_weight: torch.Tensor | None = None,
2257 + index_k_norm_weight: torch.Tensor | None = None,
2258 + num_index_heads: int = 0,
2259 + index_cache: torch.Tensor | None = None,
2260 + index_q_out: torch.Tensor | None = None,
2261 + index_slot_mapping: torch.Tensor | None = None,
2262 + skip_index_branch: bool = False,
2263 + ) -> None:
2264 + """Run consolidated MiniMax-M3 QK-norm fusion.
2265 +
2266 + Callers must check ``fused_qknorm_idxrqknorm_enabled`` first. Runtime
2267 + failures from the AITER op are deliberately propagated.
2268 + """
2269 + if kv_cache_dtype == "auto":
2270 + aiter_kv_cache_dtype = "auto"
2271 + aiter_k_scale = None
2272 + aiter_v_scale = None
2273 + elif kv_cache_dtype in ("fp8", "fp8_e4m3"):
2274 + aiter_kv_cache_dtype = "fp8_e4m3_static"
2275 + aiter_k_scale = k_scale
2276 + aiter_v_scale = v_scale
2277 + else:
2278 + raise ValueError(
2279 + "AITER fused QK-norm requires kv_cache_dtype 'auto', 'fp8', "
2280 + f"or 'fp8_e4m3', got {kv_cache_dtype!r}"
2281 + )
2282 + aiter_index_cache_dtype = "fp8" if _is_fp8_e4m3_tensor(index_cache) else "auto"
2283 +
2284 + from aiter import fused_qknorm_idxrqknorm
2285 +
2286 + fused_qknorm_idxrqknorm(
2287 + qkv,
2288 + q_norm_weight,
2289 + k_norm_weight,
2290 + cos_sin_cache,
2291 + positions,
2292 + num_heads,
2293 + num_kv_heads,
2294 + rotary_dim,
2295 + eps,
2296 + index_q_norm_weight=index_q_norm_weight,
2297 + index_k_norm_weight=index_k_norm_weight,
2298 + num_index_heads=num_index_heads,
2299 + slot_mapping=slot_mapping,
2300 + kv_cache_k=kv_cache_k,
2301 + kv_cache_v=kv_cache_v,
2302 + index_cache=index_cache,
2303 + block_size=16,
2304 + q_out=q_out,
2305 + index_q_out=index_q_out,
2306 + index_slot_mapping=index_slot_mapping,
2307 + kv_cache_dtype=aiter_kv_cache_dtype,
2308 + index_cache_dtype=aiter_index_cache_dtype,
2309 + k_scale=aiter_k_scale,
2310 + v_scale=aiter_v_scale,
2311 + asm_layout=True,
2312 + skip_index_branch=skip_index_branch,
2313 + )
2314 +
2315 @classmethod
2316 @if_aiter_supported
2317 def is_triton_unified_attn_enabled(cls) -> bool:
vllm/models/minimax_m3/amd/model.py
+163 -81
@@ -890,6 +890,30 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
890 self._ensure_aiter_sparse_pa_kv_cache()
891 return self.kv_cache_k, self.kv_cache_v
892
893 + def _get_aiter_sparse_pa_slot_mapping(
894 + self,
895 + slot_mapping: torch.Tensor,
896 + key_cache: torch.Tensor,
897 + value_cache: torch.Tensor,
898 + ) -> torch.Tensor:
899 + if value_cache.shape[0] == key_cache.shape[0]:
900 + return slot_mapping
901 +
902 + attn_metadata = get_forward_context().attn_metadata
903 + page16_slot_mapping = None
904 + if isinstance(attn_metadata, dict):
905 + main_md = attn_metadata[self.layer_name]
906 + assert isinstance(main_md, MiniMaxM3SparseMetadata)
907 + page16_slot_mapping = main_md.page16_slot_mapping
908 + if (
909 + page16_slot_mapping is None
910 + or page16_slot_mapping.shape != slot_mapping.shape
911 + ):
912 + page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
913 + slot_mapping, self.kv_cache.shape[2]
914 + )
915 + return page16_slot_mapping
916 +
917 def _insert_aiter_sparse_pa_kv(
918 self,
919 k: torch.Tensor,
@@ -897,6 +921,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
921 index_k: torch.Tensor | None,
922 slot_mapping: torch.Tensor,
923 index_slot_mapping: torch.Tensor | None,
924 + key_cache: torch.Tensor,
925 + value_cache: torch.Tensor,
926 ) -> None:
927 if self.kv_cache.numel() == 0:
928 return
@@ -906,22 +932,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
932 minimax_m3_insert_index_cache,
933 )
934
909 - key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
910 - if value_cache.shape[0] != key_cache.shape[0]:
911 - attn_metadata = get_forward_context().attn_metadata
912 - page16_slot_mapping = None
913 - if isinstance(attn_metadata, dict):
914 - main_md = attn_metadata[self.layer_name]
915 - assert isinstance(main_md, MiniMaxM3SparseMetadata)
916 - page16_slot_mapping = main_md.page16_slot_mapping
917 - if (
918 - page16_slot_mapping is None
919 - or page16_slot_mapping.shape != slot_mapping.shape
920 - ):
921 - page16_slot_mapping = minimax_m3_rebase_slots_to_page16(
922 - slot_mapping, self.kv_cache.shape[2]
923 - )
924 - slot_mapping = page16_slot_mapping
935 kv_cache_dtype = (
936 self.kv_cache_dtype
937 if is_quantized_kv_cache(self.kv_cache_dtype)
@@ -976,38 +986,78 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
986
987 main_slot_mapping = fwd_slot_mapping[self.layer_name]
988 q = qkv.new_empty((num_tokens, self.q_size))
989 + key_cache = value_cache = page16_slot_mapping = None
990 + k_scale = v_scale = None
991 + use_fused_qknorm = False
992 + if self.use_aiter_sparse_pa:
993 + key_cache, value_cache = self.get_aiter_sparse_pa_kv_cache()
994 + page16_slot_mapping = self._get_aiter_sparse_pa_slot_mapping(
995 + main_slot_mapping, key_cache, value_cache
996 + )
997 + k_scale = getattr(self, "_k_scale", None)
998 + v_scale = getattr(self, "_v_scale", None)
999 + use_fused_qknorm = rocm_aiter_ops.fused_qknorm_idxrqknorm_enabled(
1000 + self.kv_cache_dtype, k_scale, v_scale
1001 + )
1002 if self.skip_index_topk:
1003 index_q = None
1004 if self.use_aiter_sparse_pa:
982 - ops.fused_minimax_m3_qknorm_rope_kv_insert(
983 - qkv,
984 - self.q_norm.weight,
985 - self.k_norm.weight,
986 - cos_sin_cache,
987 - positions,
988 - self.num_heads,
989 - self.num_kv_heads,
990 - rotary_dim,
991 - eps,
992 - num_index_heads=self.num_idx_heads,
993 - q_out=q,
994 - skip_index_branch=True,
995 - )
996 - k_start = self.q_size
997 - v_start = k_start + self.kv_size
998 - k = qkv[:, k_start:v_start].view(
999 - num_tokens, self.num_kv_heads, self.head_dim
1000 - )
1001 - v = qkv[:, v_start : v_start + self.kv_size].view(
1002 - num_tokens, self.num_kv_heads, self.head_dim
1003 - )
1004 - self._insert_aiter_sparse_pa_kv(
1005 - k,
1006 - v,
1007 - None,
1008 - main_slot_mapping,
1009 - None,
1010 - )
1005 + assert key_cache is not None
1006 + assert value_cache is not None
1007 + assert page16_slot_mapping is not None
1008 + if use_fused_qknorm:
1009 + rocm_aiter_ops.fused_qknorm_idxrqknorm(
1010 + qkv,
1011 + self.q_norm.weight,
1012 + self.k_norm.weight,
1013 + cos_sin_cache,
1014 + positions,
1015 + self.num_heads,
1016 + self.num_kv_heads,
1017 + rotary_dim,
1018 + eps,
1019 + page16_slot_mapping,
1020 + key_cache,
1021 + value_cache,
1022 + q,
1023 + self.kv_cache_dtype,
1024 + k_scale,
1025 + v_scale,
1026 + num_index_heads=self.num_idx_heads,
1027 + skip_index_branch=True,
1028 + )
1029 + else:
1030 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
1031 + qkv,
1032 + self.q_norm.weight,
1033 + self.k_norm.weight,
1034 + cos_sin_cache,
1035 + positions,
1036 + self.num_heads,
1037 + self.num_kv_heads,
1038 + rotary_dim,
1039 + eps,
1040 + num_index_heads=self.num_idx_heads,
1041 + q_out=q,
1042 + skip_index_branch=True,
1043 + )
1044 + k_start = self.q_size
1045 + v_start = k_start + self.kv_size
1046 + k = qkv[:, k_start:v_start].view(
1047 + num_tokens, self.num_kv_heads, self.head_dim
1048 + )
1049 + v = qkv[:, v_start : v_start + self.kv_size].view(
1050 + num_tokens, self.num_kv_heads, self.head_dim
1051 + )
1052 + self._insert_aiter_sparse_pa_kv(
1053 + k,
1054 + v,
1055 + None,
1056 + page16_slot_mapping,
1057 + None,
1058 + key_cache,
1059 + value_cache,
1060 + )
1061 else:
1062 ops.fused_minimax_m3_qknorm_rope_kv_insert(
1063 qkv,
@@ -1037,42 +1087,74 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
1087 dtype=self.indexer.index_cache.dtype,
1088 )
1089 if self.use_aiter_sparse_pa:
1040 - ops.fused_minimax_m3_qknorm_rope_kv_insert(
1041 - qkv,
1042 - self.q_norm.weight,
1043 - self.k_norm.weight,
1044 - cos_sin_cache,
1045 - positions,
1046 - self.num_heads,
1047 - self.num_kv_heads,
1048 - rotary_dim,
1049 - eps,
1050 - self.index_q_norm.weight,
1051 - self.index_k_norm.weight,
1052 - self.num_idx_heads,
1053 - q_out=q,
1054 - index_q_out=index_q,
1055 - kv_cache_dtype=self.kv_cache_dtype,
1056 - )
1057 - k_start = self.q_size
1058 - v_start = k_start + self.kv_size
1059 - index_k_start = v_start + self.kv_size + self.index_q_size
1060 - k = qkv[:, k_start:v_start].view(
1061 - num_tokens, self.num_kv_heads, self.head_dim
1062 - )
1063 - v = qkv[:, v_start : v_start + self.kv_size].view(
1064 - num_tokens, self.num_kv_heads, self.head_dim
1065 - )
1066 - index_k = qkv[
1067 - :, index_k_start : index_k_start + self.idx_head_dim
1068 - ].view(num_tokens, self.idx_head_dim)
1069 - self._insert_aiter_sparse_pa_kv(
1070 - k,
1071 - v,
1072 - index_k,
1073 - main_slot_mapping,
1074 - index_slot_mapping,
1075 - )
1090 + assert key_cache is not None
1091 + assert value_cache is not None
1092 + assert page16_slot_mapping is not None
1093 + index_cache = self.indexer.index_cache.kv_cache
1094 + if use_fused_qknorm:
1095 + rocm_aiter_ops.fused_qknorm_idxrqknorm(
1096 + qkv,
1097 + self.q_norm.weight,
1098 + self.k_norm.weight,
1099 + cos_sin_cache,
1100 + positions,
1101 + self.num_heads,
1102 + self.num_kv_heads,
1103 + rotary_dim,
1104 + eps,
1105 + page16_slot_mapping,
1106 + key_cache,
1107 + value_cache,
1108 + q,
1109 + self.kv_cache_dtype,
1110 + k_scale,
1111 + v_scale,
1112 + self.index_q_norm.weight,
1113 + self.index_k_norm.weight,
1114 + self.num_idx_heads,
1115 + index_cache,
1116 + index_q,
1117 + index_slot_mapping,
1118 + )
1119 + else:
1120 + ops.fused_minimax_m3_qknorm_rope_kv_insert(
1121 + qkv,
1122 + self.q_norm.weight,
1123 + self.k_norm.weight,
1124 + cos_sin_cache,
1125 + positions,
1126 + self.num_heads,
1127 + self.num_kv_heads,
1128 + rotary_dim,
1129 + eps,
1130 + self.index_q_norm.weight,
1131 + self.index_k_norm.weight,
1132 + self.num_idx_heads,
1133 + q_out=q,
1134 + index_q_out=index_q,
1135 + kv_cache_dtype=self.kv_cache_dtype,
1136 + )
1137 + k_start = self.q_size
1138 + v_start = k_start + self.kv_size
1139 + index_k_start = v_start + self.kv_size + self.index_q_size
1140 + k = qkv[:, k_start:v_start].view(
1141 + num_tokens, self.num_kv_heads, self.head_dim
1142 + )
1143 + v = qkv[:, v_start : v_start + self.kv_size].view(
1144 + num_tokens, self.num_kv_heads, self.head_dim
1145 + )
1146 + index_k = qkv[
1147 + :, index_k_start : index_k_start + self.idx_head_dim
1148 + ].view(num_tokens, self.idx_head_dim)
1149 + self._insert_aiter_sparse_pa_kv(
1150 + k,
1151 + v,
1152 + index_k,
1153 + page16_slot_mapping,
1154 + index_slot_mapping,
1155 + key_cache,
1156 + value_cache,
1157 + )
1158 else:
1159 ops.fused_minimax_m3_qknorm_rope_kv_insert(
1160 qkv,