diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 6849c4e3ebde..565cacdecfb4 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -49,18 +49,98 @@ def enabled(aiter: bool, flag: bool) -> bool: rocm_aiter_ops.refresh_env_variables() +def test_triton_sparse_mla_gate_gfx942(monkeypatch) -> None: + """gfx942 takes the flag only for callers that opt in, and only with an aiter + whose kernel lists gfx942.""" + import vllm._aiter_ops as aiter_ops + import vllm.platforms.rocm as rocm + from vllm._aiter_ops import rocm_aiter_ops + + warnings: list[str] = [] + monkeypatch.setattr( + aiter_ops.logger, + "warning_once", + lambda msg, *args, **kwargs: warnings.append(msg % args), + ) + monkeypatch.setattr(rocm, "on_gfx950", lambda: False) + monkeypatch.setattr(rocm, "on_gfx942", lambda: True) + archs = {"SUPPORTED_ARCHS": ("gfx942", "gfx950")} + monkeypatch.setattr( + aiter_ops, "_aiter_sparse_mla_archs", lambda name: archs.get(name, ()) + ) + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "1") + rocm_aiter_ops.refresh_env_variables() + try: + assert rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True) + assert not warnings + # DeepSeek V4 does not opt in: its packed caches are gfx950-only. + assert not rocm_aiter_ops.is_triton_sparse_mla_enabled() + assert "gfx950-only" in warnings.pop() + # An aiter from before ROCm/aiter#5721 has no gfx942 kernel. + archs["SUPPORTED_ARCHS"] = ("gfx950",) + assert not rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True) + assert "ROCm/aiter#5721" in warnings.pop() + finally: + monkeypatch.undo() + rocm_aiter_ops.refresh_env_variables() + + +@pytest.mark.parametrize( + ("rope_dim", "kv_cache_dtype", "reads_fp8", "reason"), + [ + (0, "auto", False, None), + (0, "fp8", True, None), + (0, "fp8", False, "ROCm/aiter#6199"), + (64, "auto", True, "rope-free"), + ], + ids=["bf16", "fp8", "fp8_old_aiter", "rope"], +) +def test_aiter_sparse_mla_gfx942_scope( + monkeypatch, rope_dim: int, kv_cache_dtype: str, reads_fp8: bool, reason +) -> None: + """On gfx942 the backend takes aiter's kernel for rope-free sparse MLA only, + and with an fp8 cache only if the installed aiter reads it there.""" + import vllm.platforms.rocm as rocm + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseImpl, + ) + + monkeypatch.setattr(rocm, "on_gfx942", lambda: True) + monkeypatch.setattr( + rocm_aiter_ops, "triton_sparse_mla_reads_fp8_on_gfx942", lambda: reads_fp8 + ) + impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl) + impl.kv_cache_dtype = kv_cache_dtype + impl.qk_rope_head_dim = rope_dim + impl.dcp_world_size = impl.pcp_world_size = 1 + vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.bfloat16)) + + got = impl._aiter_sparse_mla_unsupported_reason(vllm_config) + if reason is None: + assert got is None + else: + assert got is not None and reason in got + + @pytest.mark.parametrize( "rope_dim", [0, 64], ids=["rope_free", "appended_rope"], # GLM-5.3-Flash; GLM-5.1/5.2, V3.2 ) -@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize( + ("kv_cache_dtype", "quantizes_q"), + [("auto", False), ("fp8", False), ("fp8", True)], + ids=["bf16", "fp8", "fp8_bf16_q"], # fp8_bf16_q: gfx942's kernel +) @torch.inference_mode() def test_forward_mqa_prepares_triton_sparse_mla_inputs( - monkeypatch, rope_dim: int, kv_cache_dtype: str + monkeypatch, rope_dim: int, kv_cache_dtype: str, quantizes_q: bool ) -> None: """What forward_mqa hands aiter: global slots built from the request-local - top-k rows, q quantized with the layer's scale, and the cache as stored.""" + top-k rows, q quantized with the layer's scale (or, where the kernel + quantizes q itself, q as is with fp8 dots), and the cache as stored.""" from vllm._aiter_ops import rocm_aiter_ops from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( ROCMAiterMLASparseImpl, @@ -99,6 +179,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( impl.owns_indexer = True impl.index_epoch = 0 impl.use_aiter_sparse_mla = True + impl.aiter_sparse_mla_quantizes_q = quantizes_q layer = SimpleNamespace( _q_scale=torch.tensor([0.02], device=device), _k_scale=torch.tensor([0.01], device=device), @@ -130,7 +211,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( (q_in, kv_in, o, sm_scale, kv_indptr, kv_indices), kwargs = calls[0] assert _rows_from_ragged(kv_indices, kv_indptr) == rows assert q_in.shape[0] == num_tokens - if fp8: + if fp8 and not quantizes_q: assert q_in.dtype == cache_dtype torch.testing.assert_close( q_in.float() * layer._q_scale, q[:num_tokens].float(), atol=1e-3, rtol=0.07 @@ -146,6 +227,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( assert kwargs["q_scale"] is layer._q_scale assert kwargs["kv_scale"] is layer._k_scale assert kwargs["attn_sink"] is impl.sinks + assert kwargs["dot_precision"] == ("fp8" if quantizes_q else None) assert lse is None diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0817c651a0f4..2db392dfff05 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -673,6 +673,20 @@ def _rocm_aiter_gather_kv_b_proj_fake( pass +@functools.cache +def _aiter_sparse_mla_archs(name: str) -> tuple[str, ...]: + """The archs the installed aiter's sparse_mla lists under name, or (). + + gfx942 joins SUPPORTED_ARCHS in ROCm/aiter#5721 and FP8_SCALAR_ARCHS, the + archs that read a per-tensor fp8 cache, in ROCm/aiter#6199. + """ + try: + from aiter.ops.triton.attention import sparse_mla + except ImportError: + return () + return tuple(getattr(sparse_mla, name, ())) + + def _rocm_aiter_mla_decode_fwd_impl( q: torch.Tensor, kv_buffer: torch.Tensor, @@ -1984,7 +1998,8 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. - VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950). + VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950, + and rope-free sparse MLA on gfx942). VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation dtype (a16w4 | a8w4 | a4w4). @@ -2455,15 +2470,28 @@ def is_triton_rotary_embed_enabled(cls) -> bool: @classmethod @if_aiter_supported - def is_triton_sparse_mla_enabled(cls) -> bool: + def is_triton_sparse_mla_enabled(cls, gfx942_ok: bool = False) -> bool: + """gfx942_ok: the caller passes only flat bf16 or per-tensor fp8 caches, + which aiter's kernel also runs on gfx942. DeepSeek V4's packed caches + are gfx950-only.""" if not cls._TRITON_SPARSE_MLA: return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx942, on_gfx950 if not cls._AITER_ENABLED: reason = "VLLM_ROCM_USE_AITER is off" - elif not on_gfx950(): - reason = "the kernel is gfx950-only" + elif on_gfx950(): + return True + elif not (gfx942_ok and on_gfx942()): + reason = ( + "the kernel is gfx942/gfx950-only" + if gfx942_ok + else "the kernel is gfx950-only" + ) + elif "gfx942" not in _aiter_sparse_mla_archs("SUPPORTED_ARCHS"): + reason = ( + "the installed aiter has no gfx942 sparse MLA kernel (ROCm/aiter#5721)" + ) else: return True logger.warning_once( @@ -3451,6 +3479,11 @@ def fused_topk( ) -> tuple[torch.Tensor, torch.Tensor]: return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up) + @staticmethod + def triton_sparse_mla_reads_fp8_on_gfx942() -> bool: + """Whether aiter's sparse MLA kernel reads an fp8 cache on gfx942.""" + return "gfx942" in _aiter_sparse_mla_archs("FP8_SCALAR_ARCHS") + @staticmethod def triton_sparse_mla_fwd( q: torch.Tensor, @@ -3468,6 +3501,7 @@ def triton_sparse_mla_fwd( extra_kv_indptr: torch.Tensor | None = None, extra_kv_indices: torch.Tensor | None = None, has_invalid: bool = True, + dot_precision: str | None = None, ) -> None: """Sparse MLA read straight from the KV cache, for prefill and decode. @@ -3476,13 +3510,17 @@ def triton_sparse_mla_fwd( negative slot passes False, which lets the fp8 kernel stage K straight into LDS. The cache format (bf16, per-tensor fp8, or DeepSeek V4's paged fp8_ds_mla) is inferred from kv_buffer and kv_scale. An fp8 q - must come with its q_scale and runs both dots in fp8. The extra segment - is DeepSeek V4's second cache: SWA window in kv_buffer, top-k + must come with its q_scale and runs both dots in fp8. dot_precision + overrides that choice: gfx942's kernel takes no fp8 q, so with an fp8 + cache it gets a bf16 q and "fp8", and quantizes q itself. The extra + segment is DeepSeek V4's second cache: SWA window in kv_buffer, top-k compressed tokens here. """ from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd fp8_q = q.dtype == FP8_DTYPE + if dot_precision is None: + dot_precision = "fp8" if fp8_q else "bf16" sparse_mla_fwd( q, kv_buffer, @@ -3493,7 +3531,7 @@ def triton_sparse_mla_fwd( kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, has_invalid=has_invalid, - dot_precision="fp8" if fp8_q else "bf16", + dot_precision=dot_precision, q_scale=q_scale if fp8_q else None, out=o, attn_sink=attn_sink, diff --git a/vllm/envs.py b/vllm/envs.py index dbbc49ff59a8..ac6ac02036ff 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1371,7 +1371,9 @@ def _resolve_rust_cli_path() -> str | None: # Whether sparse MLA prefill and decode run on aiter's Triton kernel, which # reads the KV cache as stored (bf16 or fp8, paged or flat). Used by the # ROCM_AITER_MLA_SPARSE backend (DeepSeek V3.2, GLM-5.x) and DeepSeek - # V4 / V4.1. gfx950 only. By default is disabled. + # V4 / V4.1. gfx950 only, except that the ROCM_AITER_MLA_SPARSE backend + # also uses it on gfx942 for rope-free sparse MLA (GLM-5.3-Flash) when the + # installed aiter has a gfx942 kernel. By default is disabled. "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "False").lower() in ("true", "1") diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index 037b5432ed03..20da2cf67c2a 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -780,6 +780,9 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor: class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]): supports_dcp = False use_aiter_sparse_mla = False + # gfx942's aiter kernel takes no fp8 q. With an fp8 cache it gets q in the + # model dtype and quantizes it itself, under fp8 dots. + aiter_sparse_mla_quantizes_q = False def __init__( self, @@ -845,10 +848,15 @@ def __init__( ) self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] - if rocm_aiter_ops.is_triton_sparse_mla_enabled(): + if rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True): reason = self._aiter_sparse_mla_unsupported_reason(vllm_config) if reason is None: + from vllm.platforms.rocm import on_gfx942 + self.use_aiter_sparse_mla = True + self.aiter_sparse_mla_quantizes_q = ( + on_gfx942() and kv_cache_dtype.startswith("fp8") + ) else: logger.warning_once( "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; " @@ -868,6 +876,19 @@ def _aiter_sparse_mla_unsupported_reason( # Context parallelism merges each rank's partials with a per-token LSE. if self.dcp_world_size > 1 or self.pcp_world_size > 1: return "context parallelism is not supported" + from vllm.platforms.rocm import on_gfx942 + + if on_gfx942(): + if self.qk_rope_head_dim != 0: + return "on gfx942 it covers rope-free sparse MLA only" + if ( + cache_dtype.startswith("fp8") + and not rocm_aiter_ops.triton_sparse_mla_reads_fp8_on_gfx942() + ): + return ( + "the installed aiter cannot read an fp8 KV cache on gfx942 " + "(ROCm/aiter#6199)" + ) return None def record_logical_topk_ready(self) -> None: @@ -1148,6 +1169,7 @@ def _forward_mla_aiter( q_scale=layer._q_scale, kv_scale=layer._k_scale, attn_sink=self.sinks, + dot_precision="fp8" if self.aiter_sparse_mla_quantizes_q else None, # triton_convert_req_index_to_global_index writes 0, never -1, # for an invalid top-k entry, so no slot in the stream is negative. has_invalid=False, @@ -1165,9 +1187,10 @@ def forward_mqa( # MQA 576/512 approach for both prefill and decode fp8_attention = self.kv_cache_dtype.startswith("fp8") + quantize_q = fp8_attention and not self.aiter_sparse_mla_quantizes_q if isinstance(q, tuple): ql_nope, q_pe = q - if fp8_attention: + if quantize_q: q = layer._decode_concat_quant_fp8_op( # type: ignore[attr-defined] ql_nope, q_pe, layer._q_scale ) @@ -1215,7 +1238,7 @@ def forward_mqa( # write the latent and rope to kv cache if fp8_attention: kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(current_platform.fp8_dtype()) - if q.dtype != current_platform.fp8_dtype(): + if quantize_q and q.dtype != current_platform.fp8_dtype(): original_q_shape = q.shape q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale) q = q.view(original_q_shape)