From 70ae0b7435bcb1ceef051df28c50fb7ee7d1dfec Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Date: Mon, 28 Sep 2026 20:08:18 -0500 Subject: [PATCH] [ROCm][MLA] Enable sparse MLA Gluon kernel from Aiter (#53492) Signed-off-by: Mehmet Cagri Kaymak Signed-off-by: Mehmet Cagri --- .../test_rocm_aiter_triton_sparse_mla.py | 326 ++++++++++++++++++ vllm/_aiter_ops.py | 73 ++++ vllm/envs.py | 9 + vllm/models/deepseek_v4/amd/rocm.py | 184 +++++++++- vllm/models/deepseek_v41/amd/rocm.py | 250 +++++++++++++- .../backends/mla/rocm_aiter_mla_sparse.py | 73 +++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 96 ++++++ 7 files changed, 989 insertions(+), 22 deletions(-) create mode 100644 tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py new file mode 100644 index 0000000000..645eec2979 --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-only aiter path" +) + + +def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]: + ends = indptr.tolist() + return [indices[s:e].tolist() for s, e in zip(ends, ends[1:])] + + +def test_triton_sparse_mla_gate(monkeypatch) -> None: + import vllm._aiter_ops as aiter_ops + from vllm._aiter_ops import rocm_aiter_ops + from vllm.platforms.rocm import on_gfx950 + + warnings: list[str] = [] + monkeypatch.setattr( + aiter_ops.logger, + "warning_once", + lambda msg, *args, **kwargs: warnings.append(msg % args), + ) + + def enabled(aiter: bool, flag: bool) -> bool: + monkeypatch.setenv("VLLM_ROCM_USE_AITER", str(int(aiter))) + monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", str(int(flag))) + rocm_aiter_ops.refresh_env_variables() + return bool(rocm_aiter_ops.is_triton_sparse_mla_enabled()) + + try: + assert not enabled(aiter=True, flag=False) and not warnings + assert not enabled(aiter=False, flag=True) + assert "VLLM_ROCM_USE_AITER is off" in warnings.pop() + # On gfx950 the flag takes effect; elsewhere it warns and falls back. + assert enabled(aiter=True, flag=True) == on_gfx950() + if not on_gfx950(): + assert "gfx950" in warnings.pop() + finally: + monkeypatch.undo() + rocm_aiter_ops.refresh_env_variables() + + +@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"]) +@torch.inference_mode() +def test_forward_mqa_prepares_triton_sparse_mla_inputs( + monkeypatch, rope_dim: int, kv_cache_dtype: str +) -> 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.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseImpl, + ) + + device = torch.device("cuda") + gen = torch.Generator().manual_seed(0) + kv_lora_rank, num_heads, topk, block_size = 512, 16, 128, 16 + head_size = kv_lora_rank + rope_dim + fp8 = kv_cache_dtype == "fp8" + + # A decode row, then a short prefill. + query_lens, seq_lens = [1, 5], [40, 20] + num_tokens = sum(query_lens) + block_table = torch.randperm(8, generator=gen).view(2, 4).to(device, torch.int32) + topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32) + req_id_per_token: list[int] = [] + rows: list[list[int]] = [] + for req, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)): + for pos in range(seq_len - query_len, seq_len): + local = torch.randperm(pos + 1, generator=gen) + topk_indices[len(rows), : pos + 1] = local.int() + slots = block_table[req, local // block_size].cpu() * block_size + rows.append((slots + local % block_size).tolist()) + req_id_per_token.append(req) + indptr = torch.tensor([0] + [len(row) for row in rows]).cumsum(0) + + impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl) + impl.num_heads = num_heads + impl.kv_lora_rank = kv_lora_rank + impl.qk_rope_head_dim = rope_dim + impl.scale = head_size**-0.5 + impl.sinks = torch.randn(num_heads, generator=gen).to(device) + impl.kv_cache_dtype = kv_cache_dtype + impl.topk_indices_buffer = topk_indices.to(device) + impl.use_aiter_sparse_mla = True + layer = SimpleNamespace( + _q_scale=torch.tensor([0.02], device=device), + _k_scale=torch.tensor([0.01], device=device), + ) + metadata = SimpleNamespace( + num_actual_tokens=num_tokens, + topk_tokens=topk, + req_id_per_token=torch.tensor(req_id_per_token, dtype=torch.int32).to(device), + block_table=block_table, + block_size=block_size, + paged_kv_indptr=indptr.to(device, torch.int32), + paged_kv_indices=torch.zeros(num_tokens * topk, dtype=torch.int32).to(device), + attn_out_dtype=torch.bfloat16, + ) + cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16 + kv_cache = torch.zeros(8, block_size, head_size, device=device).to(cache_dtype) + # Two rows of cudagraph padding past the real tokens. + q = torch.randn(num_tokens + 2, num_heads, head_size, generator=gen).to( + device, torch.bfloat16 + ) + calls = [] + monkeypatch.setattr( + rocm_aiter_ops, "triton_sparse_mla_fwd", lambda *a, **k: calls.append((a, k)) + ) + + out, lse = impl.forward_mqa(q, kv_cache, metadata, layer) + + assert len(calls) == 1 + (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: + 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 + ) + else: + assert q_in.data_ptr() == q.data_ptr() + assert kv_in.shape == (8 * block_size, 1, 1, head_size) + assert kv_in.data_ptr() == kv_cache.data_ptr() + assert o.data_ptr() == out.data_ptr() and o.shape[0] == num_tokens + assert sm_scale == impl.scale + assert kwargs["kv_lora_rank"] == kv_lora_rank + assert kwargs["qk_rope_head_dim"] == rope_dim + assert kwargs["q_scale"] is layer._q_scale + assert kwargs["kv_scale"] is layer._k_scale + assert kwargs["attn_sink"] is impl.sinks + assert lse is None + + +@pytest.mark.parametrize( + ("model", "compress_ratio"), + [("deepseek_v41", 1), ("deepseek_v41", 2), ("deepseek_v4", 4)], +) +@torch.inference_mode() +def test_paged_prefill_indices_match_gathered_prefill( + monkeypatch, model: str, compress_ratio: int +) -> None: + """The aiter prefill reads the paged caches through the SWA builder's rows + and build_prefill_topk_ragged_indices, and must attend exactly the keys the + dequantize-and-gather prefill does.""" + import importlib + + from tests.v1.attention.utils import create_vllm_config + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_prefill_topk_ragged_indices, + ) + from vllm.v1.kv_cache_interface import SlidingWindowMLASpec + + rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") + device = torch.device("cuda") + gen = torch.Generator().manual_seed(compress_ratio) + window, width, block_size, max_model_len = 16, 32, 16, 128 + comp_block = block_size // compress_ratio + num_compressed = max_model_len // compress_ratio + + # A decode row, then prefills with and without a prefix. + prefix_lens, query_lens = [20, 0, 13, 40], [1, 7, 20, 9] + seq_lens = [p + q for p, q in zip(prefix_lens, query_lens)] + num_decodes = num_decode_tokens = 1 + num_tokens = sum(query_lens) + qsl = torch.tensor([0] + query_lens).cumsum(0).to(torch.int32) + positions = torch.cat([torch.arange(p, s) for p, s in zip(prefix_lens, seq_lens)]) + token_to_req = torch.repeat_interleave( + torch.arange(len(query_lens)), torch.tensor(query_lens) + ) + swa_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32) + comp_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32) + + vllm_config = create_vllm_config( + model_name="facebook/opt-125m", + max_model_len=max_model_len, + block_size=block_size, + max_num_seqs=4, + max_num_batched_tokens=64, + hf_config_override={ + "compress_ratios": [compress_ratio], + "index_topk": width, + "sliding_window": window, + }, + ) + swa_spec = SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + sliding_window=window, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + slots = swa_bt.cpu()[token_to_req, positions // block_size] * block_size + ( + positions % block_size + ) + common = CommonAttentionMetadata( + query_start_loc=qsl.to(device), + query_start_loc_cpu=qsl, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32).to(device), + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), + num_reqs=len(query_lens), + num_actual_tokens=num_tokens, + max_query_len=max(query_lens), + max_seq_len=max(seq_lens), + block_table_tensor=swa_bt, + slot_mapping=slots.to(device, torch.int64), + positions=positions.to(device), + causal=True, + ) + + def build(enabled: bool): + monkeypatch.setattr( + rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled + ) + return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder( + swa_spec, ["swa"], vllm_config, device + ).build(0, common) + + assert build(enabled=False).prefill_swa_ragged_indices is None + md = build(enabled=True) + assert md.num_decode_tokens == num_decode_tokens + + # Local top-k rows. Under each row's causal cap sit a failed candidate and + # one past the pool; past the cap, stale values neither path may read. + num_prefill_tokens = num_tokens - num_decode_tokens + topk = torch.full((num_prefill_tokens, width), -1, dtype=torch.int32) + for row in range(num_prefill_tokens): + visible = (int(positions[row + num_decode_tokens]) + 1) // compress_ratio + cap = min(visible, width) + topk[row, :cap] = torch.randperm(visible, generator=gen)[:cap].int() + if cap >= 2: + topk[row, :2] = torch.tensor([-1, num_compressed + 3]) + topk[row, cap:] = torch.randint( + 0, num_compressed, (width - cap,), generator=gen + ).int() + topk = topk.to(device) + + topk_indices, topk_indptr = build_prefill_topk_ragged_indices( + topk, + md.token_to_req_indices, + md.query_start_loc, + md.seq_lens, + md.is_valid_token, + comp_bt, + block_size=comp_block, + compress_ratio=compress_ratio, + num_compressed=num_compressed, + token_offset=num_decode_tokens, + ) + paged_keys = [ + sorted([("s", x) for x in swa if x >= 0] + [("c", x) for x in top if x >= 0]) + for swa, top in zip( + _rows_from_ragged( + md.prefill_swa_ragged_indices, md.prefill_swa_ragged_indptr + ), + _rows_from_ragged(topk_indices, topk_indptr), + ) + ] + + # The gathered path's combine, mapped from workspace offsets to cache slots. + M = num_compressed + window + num_tokens + combined, combined_lens = rocm.combine_topk_swa_indices( + topk, + md.query_start_loc[num_decodes:], + md.prefill_seq_lens, + md.prefill_gather_lens, + window, + compress_ratio, + width, + M, + num_compressed, + ) + gather_start = (md.prefill_seq_lens - md.prefill_gather_lens).tolist() + swa_bt_cpu, comp_bt_cpu = swa_bt.cpu(), comp_bt.cpu() + + def slot(block_table, req, pos, size): + return int(block_table[req, pos // size]) * size + pos % size + + gathered_keys = [] + for row, (values, n) in enumerate(zip(combined.tolist(), combined_lens.tolist())): + req = int(token_to_req[row + num_decode_tokens]) + chunk_req = req - num_decodes + keys = [] + for value in values[:n]: + if value < 0: + continue + offset = value - chunk_req * M + if offset < num_compressed: + keys.append(("c", slot(comp_bt_cpu, req, offset, comp_block))) + else: + pos = gather_start[chunk_req] + offset - num_compressed + keys.append(("s", slot(swa_bt_cpu, req, pos, block_size))) + gathered_keys.append(sorted(keys)) + + assert paged_keys == gathered_keys + + +@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"]) +def test_aiter_sparse_mla_stays_off_with_kv_connector(monkeypatch, model) -> None: + import importlib + + from vllm._aiter_ops import rocm_aiter_ops + + rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") + monkeypatch.setattr(rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: True) + assert rocm._aiter_sparse_mla_enabled(has_kv_transfer=False) + assert not rocm._aiter_sparse_mla_enabled(has_kv_transfer=True) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0e89b39818..ec249b1eca 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1927,6 +1927,7 @@ 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_FUSION_SHARED_EXPERTS: Controls shared expert fusion. VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation dtype (a16w4 | a8w4 | a4w4). @@ -1998,6 +1999,7 @@ class rocm_aiter_ops: _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE + _TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS _MOE_SITUV2 = _resolve_situv2_activation() # TODO: Consolidate under _LINEAR_ENABLED @@ -2027,6 +2029,7 @@ class rocm_aiter_ops: cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE + cls._TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS cls._MOE_SITUV2 = _resolve_situv2_activation() _sync_aiter_situv2_moe_env() @@ -2391,6 +2394,25 @@ class rocm_aiter_ops: def is_triton_rotary_embed_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED + @classmethod + @if_aiter_supported + def is_triton_sparse_mla_enabled(cls) -> bool: + if not cls._TRITON_SPARSE_MLA: + return False + from vllm.platforms.rocm import on_gfx950 + + if not cls._AITER_ENABLED: + reason = "VLLM_ROCM_USE_AITER is off" + elif not on_gfx950(): + reason = "the kernel is gfx950-only" + else: + return True + logger.warning_once( + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; ignoring it.", + reason, + ) + return False + @classmethod @if_aiter_supported def is_triton_gemm_enabled(cls) -> bool: @@ -3137,6 +3159,57 @@ class rocm_aiter_ops: ) -> 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_fwd( + q: torch.Tensor, + kv_buffer: torch.Tensor, + o: torch.Tensor, + sm_scale: float, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + kv_lora_rank: int = 512, + qk_rope_head_dim: int = 64, + q_scale: torch.Tensor | None = None, + kv_scale: torch.Tensor | None = None, + attn_sink: torch.Tensor | None = None, + extra_kv_buffer: torch.Tensor | None = None, + extra_kv_indptr: torch.Tensor | None = None, + extra_kv_indices: torch.Tensor | None = None, + has_invalid: bool = True, + ) -> None: + """Sparse MLA read straight from the KV cache, for prefill and decode. + + kv_indices are ragged global slot ids. With has_invalid they may hold + -1, which the kernel masks; a caller whose stream never holds a + 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 + compressed tokens here. + """ + from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd + + fp8_q = q.dtype == FP8_DTYPE + sparse_mla_fwd( + q, + kv_buffer, + kv_indptr, + kv_indices, + sm_scale, + kv_scale=kv_scale, + 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", + q_scale=q_scale if fp8_q else None, + out=o, + attn_sink=attn_sink, + extra_kv=extra_kv_buffer, + extra_indptr=extra_kv_indptr, + extra_indices=extra_kv_indices, + ) + @staticmethod def mla_decode_fwd( q: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index 5fee7e172e..c8e1a8b708 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -148,6 +148,7 @@ if TYPE_CHECKING: VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["asm", "segmented"] = "segmented" VLLM_ROCM_USE_AITER_MHA: bool = True VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False + VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False VLLM_ROCM_USE_AITER_FP8BMM: bool = True VLLM_ROCM_USE_AITER_FP4BMM: bool = True @@ -1343,6 +1344,14 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1") ), + # 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. + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "False").lower() + in ("true", "1") + ), # Whether to use aiter rope. # By default is disabled. "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: ( diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index d64eb3c9f3..292901dfe9 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -35,6 +35,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_prefill_topk_ragged_indices, build_ragged_indices_from_dense, rocm_inv_rope_einsum, rocm_inverse_rope_rows_, @@ -103,6 +104,22 @@ def _trust_dsv4_extra_cache_nan_free( ) +def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool: + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_triton_sparse_mla_enabled(): + return False + # The aiter kernel does not scrub NaNs, so like the trust check above it + # needs a cache only the canonical writer has touched. + if has_kv_transfer: + logger.warning_once( + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV " + "connector; DeepSeek V4 keeps the Triton sparse attention." + ) + return False + return True + + def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: lengths = lengths.to(dtype=torch.int32).contiguous() indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) @@ -446,6 +463,9 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata): class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None decode_swa_ragged_indptr: torch.Tensor | None = None + # Built only for the aiter prefill, which reads the SWA cache in place. + prefill_swa_ragged_indices: torch.Tensor | None = None + prefill_swa_ragged_indptr: torch.Tensor | None = None class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): @@ -565,6 +585,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild dtype=torch.int32, device=self.device, ) + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled( + self.vllm_config.kv_transfer_config is not None + ) def build( self, @@ -580,6 +603,23 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild replay_start=replay_start, ) + prefill_ragged_indices = None + prefill_ragged_indptr = None + if ( + self._use_aiter_sparse_mla + and base.num_prefill_tokens > 0 + and base.prefill_swa_indices is not None + and base.prefill_swa_lens is not None + ): + # Paged-coordinate rows (image visibility included), so every + # layer's prefill shares one build. + prefill_ragged_indices, prefill_ragged_indptr = ( + build_ragged_indices_from_dense( + base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1), + base.prefill_swa_lens, + ) + ) + ragged_indices = None ragged_indptr = None if ( @@ -606,6 +646,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild **vars(base), decode_swa_ragged_indices=ragged_indices, decode_swa_ragged_indptr=ragged_indptr, + prefill_swa_ragged_indices=prefill_ragged_indices, + prefill_swa_ragged_indptr=prefill_ragged_indptr, ) @@ -623,6 +665,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): """ROCm sparse MLA attention layer for DeepSeek V4.""" backend_cls = DeepseekV4ROCMAiterMLASparseBackend + _use_aiter_sparse_mla = False def __init__(self, *args, **kwargs): vllm_config = args[0] if args else kwargs["vllm_config"] @@ -648,6 +691,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): else: # Disable indexer inner overlap. self.indexer.aux_stream = None + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer) def _run_sequential_pipeline( self, @@ -1190,18 +1234,20 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): if attn_metadata is None: # Warmup dummy run: no real metadata. Reserve the same bf16 # gather workspace _forward_prefill would; the dequantize / topk - # / sparse_fwd kernels are skipped this step. - swa_only = self.compress_ratio <= 1 - N = ( - 0 - if swa_only - else (self.max_model_len + self.compress_ratio - 1) - // self.compress_ratio - ) - M = N + self.window_size + self.max_num_batched_tokens - current_workspace_manager().get_simultaneous( - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), - ) + # / sparse_fwd kernels are skipped this step. The aiter prefill + # reads the caches in place and never asks for it. + if not self._use_aiter_sparse_mla: + swa_only = self.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (self.max_model_len + self.compress_ratio - 1) + // self.compress_ratio + ) + M = N + self.window_size + self.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) output.zero_() return @@ -1312,6 +1358,22 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr + if self._use_aiter_sparse_mla: + assert swa_metadata.decode_swa_ragged_indices is not None + assert swa_metadata.decode_swa_ragged_indptr is not None + self._aiter_sparse_mla( + q, + output, + self.swa_cache_layer.kv_cache, + swa_metadata.decode_swa_ragged_indices, + swa_metadata.decode_swa_ragged_indptr, + kv_cache, + topk_ragged_indices, + topk_ragged_indptr, + ) + # The aiter kernel leaves the inverse RoPE to the caller. + return 0 + return rocm_sparse_attn_decode( q=q, kv_cache=kv_cache, @@ -1351,6 +1413,17 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, ) -> None: + if self._use_aiter_sparse_mla: + self._forward_prefill_aiter( + q=q, + compressed_k_cache=compressed_k_cache, + swa_k_cache=swa_k_cache, + output=output, + attn_metadata=attn_metadata, + swa_metadata=swa_metadata, + ) + return + swa_only = attn_metadata is None num_prefills = swa_metadata.num_prefills @@ -1476,3 +1549,90 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): attn_sink=self.attn_sink, output=output[query_start:query_end], ) + + def _forward_prefill_aiter( + self, + q: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + ) -> None: + """Prefill straight off the paged caches, with no bf16 gather.""" + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + assert swa_metadata.prefill_swa_ragged_indices is not None + assert swa_metadata.prefill_swa_ragged_indptr is not None + + topk_ragged_indices = None + topk_ragged_indptr = None + if attn_metadata is None: + compressed_k_cache = None + else: + assert compressed_k_cache is not None + assert swa_metadata.token_to_req_indices is not None + assert swa_metadata.query_start_loc is not None + assert swa_metadata.seq_lens is not None + assert swa_metadata.is_valid_token is not None + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + num_prefill_tokens + ] + else: + topk_indices = attn_metadata.c128a_prefill_topk_indices + assert topk_indices is not None + topk_ragged_indices, topk_ragged_indptr = build_prefill_topk_ragged_indices( + topk_indices, + swa_metadata.token_to_req_indices, + swa_metadata.query_start_loc, + swa_metadata.seq_lens, + swa_metadata.is_valid_token, + attn_metadata.block_table, + block_size=attn_metadata.block_size // self.compress_ratio, + compress_ratio=self.compress_ratio, + num_compressed=(self.max_model_len + self.compress_ratio - 1) + // self.compress_ratio, + token_offset=num_decode_tokens, + num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1], + ) + + self._aiter_sparse_mla( + q[:num_prefill_tokens], + output[:num_prefill_tokens], + swa_k_cache, + swa_metadata.prefill_swa_ragged_indices, + swa_metadata.prefill_swa_ragged_indptr, + compressed_k_cache, + topk_ragged_indices, + topk_ragged_indptr, + ) + + def _aiter_sparse_mla( + self, + q: torch.Tensor, + output: torch.Tensor, + swa_k_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_indptr: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + topk_indices: torch.Tensor | None, + topk_indptr: torch.Tensor | None, + ) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + # Both fp8 caches are read in place: the SWA window as the main segment, + # the top-k compressed tokens as the extra one. + rocm_aiter_ops.triton_sparse_mla_fwd( + q, + swa_k_cache, + output, + self.scale, + swa_indptr, + swa_indices, + attn_sink=self.attn_sink[: q.shape[1]], + extra_kv_buffer=compressed_k_cache, + extra_kv_indptr=topk_indptr, + extra_kv_indices=topk_indices, + ) diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index d38e412e31..897b0d12db 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -38,6 +38,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + build_prefill_topk_ragged_indices, build_ragged_indices_from_dense, rocm_inv_rope_einsum, rocm_inverse_rope_mxfp8_rows, @@ -64,6 +65,22 @@ def _trust_dsv4_extra_cache_nan_free( ) +def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool: + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_triton_sparse_mla_enabled(): + return False + # The aiter kernel does not scrub NaNs, so like the trust check above it + # needs a cache only the canonical writer has touched. + if has_kv_transfer: + logger.warning_once( + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV " + "connector; DeepSeek V4 keeps the Triton sparse attention." + ) + return False + return True + + def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: lengths = lengths.to(dtype=torch.int32).contiguous() indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) @@ -408,6 +425,9 @@ _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None decode_swa_ragged_indptr: torch.Tensor | None = None + # Built only for the aiter prefill, which reads the SWA cache in place. + prefill_swa_ragged_indices: torch.Tensor | None = None + prefill_swa_ragged_indptr: torch.Tensor | None = None class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): @@ -432,6 +452,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu dtype=torch.int32, device=self.device, ) + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled( + self.vllm_config.kv_transfer_config is not None + ) def build( self, @@ -447,6 +470,22 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu replay_start=replay_start, ) + prefill_ragged_indices = None + prefill_ragged_indptr = None + if ( + self._use_aiter_sparse_mla + and base.num_prefill_tokens > 0 + and base.prefill_swa_indices is not None + and base.prefill_swa_lens is not None + ): + # Paged-coordinate rows, so every layer's prefill shares one build. + prefill_ragged_indices, prefill_ragged_indptr = ( + build_ragged_indices_from_dense( + base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1), + base.prefill_swa_lens, + ) + ) + ragged_indices = None ragged_indptr = None if ( @@ -473,6 +512,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu **vars(base), decode_swa_ragged_indices=ragged_indices, decode_swa_ragged_indptr=ragged_indptr, + prefill_swa_ragged_indices=prefill_ragged_indices, + prefill_swa_ragged_indptr=prefill_ragged_indptr, ) @@ -497,6 +538,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend + _use_aiter_sparse_mla = False def __init__(self, *args, **kwargs): vllm_config = args[0] if args else kwargs["vllm_config"] @@ -518,6 +560,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): # same thing. The source owns one cache per compress ratio, refreshed # when it runs; consumers below it read through. self._topk_ragged_cache: dict[int, _TopkRagged] = {} + self._prefill_topk_ragged_cache: dict[ + str, tuple[torch.Tensor, torch.Tensor] + ] = {} self._index_source_prefix: str | None = None if self.compress_ratio > 0: assert self.index_source_layer_id is not None @@ -530,6 +575,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): "this rank; PP splits inside a v4.1 index-sharing group " "are not supported." ) + self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer) if self.compressor is None and self.indexer is None: # Dense layers have nothing to overlap; keep the base serial path. self.aux_stream_list = None @@ -915,9 +961,10 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): attn_metadata = forward_context.attn_metadata if attn_metadata is None: - # Warmup dummy run: no real metadata. Reserve the same bf16 - # gather workspace _forward_prefill would; the dequantize / topk - # / sparse_fwd kernels are skipped this step. + # Warmup dummy run: no real metadata. Reserve the workspace + # _forward_prefill would; the dequantize / topk / sparse_fwd kernels + # are skipped this step. The aiter prefill reads the caches in place, + # so it asks only for the bf16 rows of the MXFP8 output. swa_only = self.compress_ratio == 0 N = ( 0 @@ -926,9 +973,14 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): // self.compress_ratio ) M = N + self.window_size + self.max_num_batched_tokens - current_workspace_manager().get_simultaneous( - *self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q) + shapes = self._prefill_workspace_shapes( + M, + self.max_num_batched_tokens, + q, + gather=not self._use_aiter_sparse_mla, ) + if shapes: + current_workspace_manager().get_simultaneous(*shapes) if mxfp8_out is None: output.zero_() else: @@ -1087,6 +1139,38 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): num_decode_tokens=num_decode_tokens, ) + if self._use_aiter_sparse_mla: + assert swa_metadata.decode_swa_ragged_indices is not None + assert swa_metadata.decode_swa_ragged_indptr is not None + out = output + if output_mxfp8 is not None: + # The aiter kernel has no MXFP8 epilogue: write bf16 rows, then + # rotate and quantize them as the decode reduce would have. + out = self._mxfp8_bf16_rows(q.shape[0], q) + assert out is not None + self._aiter_sparse_mla( + q, + out, + self.swa_cache_layer.kv_cache, + swa_metadata.decode_swa_ragged_indices, + swa_metadata.decode_swa_ragged_indptr, + kv_cache, + topk_ragged_indices, + topk_ragged_indptr, + ) + if output_mxfp8 is not None: + rocm_inverse_rope_mxfp8_rows( + out, + positions, + self.rotary_emb.cos_sin_cache, + self.rope_head_dim, + output_mxfp8[0], + output_mxfp8[1], + ) + return q.shape[0] + # The aiter kernel leaves the inverse RoPE to the caller. + return 0 + return rocm_sparse_attn_decode( q=q, kv_cache=kv_cache, @@ -1132,6 +1216,19 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): With ``mxfp8_out`` the rows go to a bf16 workspace first and are then inverse-RoPE'd and quantized into it in one pass. """ + if self._use_aiter_sparse_mla: + self._forward_prefill_aiter( + q=q, + positions=positions, + compressed_k_cache=compressed_k_cache, + swa_k_cache=swa_k_cache, + output=output, + attn_metadata=attn_metadata, + swa_metadata=swa_metadata, + mxfp8_out=mxfp8_out, + ) + return + swa_only = attn_metadata is None num_prefills = swa_metadata.num_prefills @@ -1249,15 +1346,17 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): ) def _prefill_workspace_shapes( - self, M: int, num_prefill_tokens: int, q: torch.Tensor + self, M: int, num_prefill_tokens: int, q: torch.Tensor, gather: bool = True ) -> list[tuple[tuple[int, ...], torch.dtype]]: """The prefill gather buffer, plus bf16 rows for the MXFP8 output. One ``get_simultaneous`` call: separate calls alias the same memory. + ``gather=False`` leaves out the gather buffer, which the aiter kernel, + reading the caches in place, never needs. """ - shapes: list[tuple[tuple[int, ...], torch.dtype]] = [ - ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16) - ] + shapes: list[tuple[tuple[int, ...], torch.dtype]] = [] + if gather: + shapes.append(((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)) if _ON_GFX950: shapes.append( ( @@ -1266,3 +1365,136 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): ) ) return shapes + + def _prefill_topk_ragged( + self, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + attn_metadata: DeepseekV4FlashMLAMetadata, + compressed_k_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prefill counterpart of ``_decode_topk_ragged``: layers that share an + index source and a compressed cache reuse one build per step.""" + assert swa_metadata.token_to_req_indices is not None + assert swa_metadata.query_start_loc is not None + assert swa_metadata.seq_lens is not None + assert swa_metadata.is_valid_token is not None + assert self.topk_indices_buffer is not None + assert self._index_source_prefix is not None + assert self.compressed_cache_prefix is not None + + source = self._static_forward_context[self._index_source_prefix] + if source is self: + source._prefill_topk_ragged_cache = {} + cached = source._prefill_topk_ragged_cache.get(self.compressed_cache_prefix) + if cached is not None: + return cached + + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + built = build_prefill_topk_ragged_indices( + self.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + num_prefill_tokens + ], + swa_metadata.token_to_req_indices, + swa_metadata.query_start_loc, + swa_metadata.seq_lens, + swa_metadata.is_valid_token, + attn_metadata.block_table, + block_size=attn_metadata.block_size // self.compress_ratio, + compress_ratio=self.compress_ratio, + num_compressed=(self.max_model_len + self.compress_ratio - 1) + // self.compress_ratio, + token_offset=num_decode_tokens, + num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1], + ) + source._prefill_topk_ragged_cache[self.compressed_cache_prefix] = built + return built + + def _forward_prefill_aiter( + self, + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor | None, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + """Prefill straight off the paged caches, with no bf16 gather. + + With ``mxfp8_out`` the kernel writes bf16 rows to the workspace, which + are then inverse-RoPE'd and quantized into it as ``_forward_prefill`` + does. + """ + num_prefill_tokens = swa_metadata.num_prefill_tokens + assert swa_metadata.prefill_swa_ragged_indices is not None + assert swa_metadata.prefill_swa_ragged_indptr is not None + + topk_indices = None + topk_indptr = None + if attn_metadata is None: + compressed_k_cache = None + else: + assert compressed_k_cache is not None + topk_indices, topk_indptr = self._prefill_topk_ragged( + swa_metadata, attn_metadata, compressed_k_cache + ) + + if mxfp8_out is not None: + output = self._mxfp8_bf16_rows(num_prefill_tokens, q) + assert output is not None + self._aiter_sparse_mla( + q[:num_prefill_tokens], + output[:num_prefill_tokens], + swa_k_cache, + swa_metadata.prefill_swa_ragged_indices, + swa_metadata.prefill_swa_ragged_indptr, + compressed_k_cache, + topk_indices, + topk_indptr, + ) + if mxfp8_out is not None: + rocm_inverse_rope_mxfp8_rows( + output[:num_prefill_tokens], + positions[:num_prefill_tokens], + self.rotary_emb.cos_sin_cache, + self.rope_head_dim, + mxfp8_out[0][:num_prefill_tokens], + mxfp8_out[1][:num_prefill_tokens], + ) + + def _mxfp8_bf16_rows(self, num_tokens: int, q: torch.Tensor) -> torch.Tensor: + """bf16 attention rows for the aiter kernel to write when the layer's + output is MXFP8, taken from the workspace the warmup reserved.""" + shapes = self._prefill_workspace_shapes(0, num_tokens, q, gather=False) + assert shapes, "an MXFP8 attention output is gfx950-only" + return current_workspace_manager().get_simultaneous(*shapes)[0] + + def _aiter_sparse_mla( + self, + q: torch.Tensor, + output: torch.Tensor, + swa_k_cache: torch.Tensor, + swa_indices: torch.Tensor, + swa_indptr: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + topk_indices: torch.Tensor | None, + topk_indptr: torch.Tensor | None, + ) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + # Both fp8 caches are read in place: the SWA window as the main segment, + # the top-k compressed tokens as the extra one. + rocm_aiter_ops.triton_sparse_mla_fwd( + q, + swa_k_cache, + output, + self.scale, + swa_indptr, + swa_indices, + attn_sink=self.attn_sink[: q.shape[1]], + extra_kv_buffer=compressed_k_cache, + extra_kv_indptr=topk_indptr, + extra_kv_indices=topk_indices, + ) 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 f6258759d0..c5a672e5d4 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -396,9 +396,11 @@ class ROCMAiterMLASparseMetadataBuilder( attention_context = vllm_config.compilation_config.static_forward_context # Sink decode must use AITER's nonpersistent path. In particular, # gfx942 has no persistent+LSE kernel, and its metadata heuristic - # terminates for HY-V4's TP1 H64 shape. + # terminates for HY-V4's TP1 H64 shape. The Triton sparse MLA kernel + # never reads the persistent metadata either. self._use_persistent_metadata = all( getattr(attention_context[name].impl, "sinks", None) is None + and not getattr(attention_context[name].impl, "use_aiter_sparse_mla", False) for name in layer_names ) # Bounds the KV-split heuristic (see `_sparse_decode_max_split`). @@ -716,6 +718,7 @@ class ROCMAiterMLASparseImpl( is_sparse = True supports_dense_mha_prefill = False supports_dcp = False + use_aiter_sparse_mla = False def __init__( self, @@ -765,6 +768,32 @@ class ROCMAiterMLASparseImpl( (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( (q_concat_shape, vllm_config.model_config.dtype), ) + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + + if rocm_aiter_ops.is_triton_sparse_mla_enabled(): + reason = self._aiter_sparse_mla_unsupported_reason(vllm_config) + if reason is None: + self.use_aiter_sparse_mla = True + else: + logger.warning_once( + "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; " + "using the default sparse MLA path instead.", + reason, + ) + + def _aiter_sparse_mla_unsupported_reason( + self, vllm_config: VllmConfig + ) -> str | None: + model_dtype = vllm_config.model_config.dtype + if model_dtype != torch.bfloat16: + return f"the model dtype is {model_dtype}, not bfloat16" + cache_dtype = self.kv_cache_dtype + if cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"): + return f"kv-cache dtype {cache_dtype!r} is neither bfloat16 nor e4m3 fp8" + # 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" + return None def record_logical_topk_ready(self) -> None: # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer @@ -975,6 +1004,43 @@ class ROCMAiterMLASparseImpl( return output, lse + def _forward_mla_aiter( + self, + layer: AttentionLayer, + q: torch.Tensor, # [sq, heads, d_qk], not head-padded + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: ROCMAiterMLASparseMetadata, + ) -> torch.Tensor: + """_forward_mla on aiter's Triton sparse MLA kernel. + + Reads the same index stream, but needs no q head padding and no + persistent metadata, and its launch depends on shapes only, so it is + CUDA-graph capturable. + """ + num_actual_tokens = attn_metadata.num_actual_tokens + output = torch.empty( + [q.shape[0], self.num_heads, self.kv_lora_rank], + dtype=attn_metadata.attn_out_dtype, + device=q.device, + ) + rocm_aiter_ops.triton_sparse_mla_fwd( + q[:num_actual_tokens], + kv_c_and_k_pe_cache.view(-1, 1, 1, q.shape[-1]), + output[:num_actual_tokens], + self.scale, + attn_metadata.paged_kv_indptr, + attn_metadata.paged_kv_indices, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + q_scale=layer._q_scale, + kv_scale=layer._k_scale, + attn_sink=self.sinks, + # 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, + ) + return output + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -1026,6 +1092,11 @@ class ROCMAiterMLASparseImpl( 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) + if self.use_aiter_sparse_mla: + output = self._forward_mla_aiter( + layer, q, kv_c_and_k_pe_cache, attn_metadata + ) + return output, None mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q) return self._forward_mla( layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 005cbbb2ff..e676576fb7 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -4306,3 +4306,99 @@ def rocm_sparse_attn_decode( if direct_out is None: output.copy_(attn_out.to(output.dtype)) return output.shape[0] if inv_rope_positions is not None else 0 + + +@triton.jit(do_not_specialize=["token_offset"]) +def _prefill_topk_global_slots_kernel( + out_ptr, + lens_ptr, + topk_ptr, + topk_stride, + token_to_req_ptr, + query_start_loc_ptr, + seq_lens_ptr, + is_valid_token_ptr, + block_table_ptr, + block_table_stride, + token_offset, + num_compressed, + TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + BLOCK_W: tl.constexpr, +): + row = tl.program_id(0) + token = row + token_offset + req = tl.load(token_to_req_ptr + token) + query_start = tl.load(query_start_loc_ptr + req) + query_len = tl.load(query_start_loc_ptr + req + 1) - query_start + pos = tl.load(seq_lens_ptr + req) - query_len + token - query_start + # Only (pos + 1) // ratio compressed tokens exist yet; the rest of the row + # can hold stale indices. + row_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOPK) + row_len = tl.where(tl.load(is_valid_token_ptr + token), row_len, 0) + + for start in range(0, TOPK, BLOCK_W): + cols = start + tl.arange(0, BLOCK_W) + live = cols < row_len + local = tl.load(topk_ptr + row * topk_stride + cols, mask=live, other=-1) + valid = live & (local >= 0) & (local < num_compressed) + local = tl.where(valid, local, 0) + block = tl.load( + block_table_ptr + req * block_table_stride + local // BLOCK_SIZE, + mask=valid, + other=0, + ) + slot = tl.where(valid, block * BLOCK_SIZE + local % BLOCK_SIZE, -1) + tl.store(out_ptr + row * TOPK + cols, slot, mask=cols < TOPK) + tl.store(lens_ptr + row, row_len) + + +def build_prefill_topk_ragged_indices( + topk_indices: torch.Tensor, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + is_valid_token: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + compress_ratio: int, + num_compressed: int, + token_offset: int, + num_rows: int = -1, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map prefill top-k rows to a ragged stream of compressed-cache slots. + + ``topk_indices`` holds local compressed positions for the prefill tokens, + which sit at ``token_offset`` in the batch; ``token_to_req_indices``, + ``query_start_loc``, ``seq_lens`` and ``block_table`` are batch-wide. + ``block_size`` is the compressed cache's, i.e. already divided by the ratio. + """ + topk_indices = topk_indices.reshape(topk_indices.shape[0], -1) + num_tokens, width = topk_indices.shape + dense = torch.empty( + (num_tokens, width), dtype=torch.int32, device=topk_indices.device + ) + lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device) + if num_tokens > 0 and width > 0: + _prefill_topk_global_slots_kernel[(num_tokens,)]( + dense, + lens, + topk_indices, + topk_indices.stride(0), + token_to_req_indices, + query_start_loc, + seq_lens, + is_valid_token, + block_table, + block_table.stride(0), + token_offset, + num_compressed, + TOPK=width, + COMPRESS_RATIO=compress_ratio, + BLOCK_SIZE=block_size, + BLOCK_W=min(triton.next_power_of_2(width), 1024), + ) + else: + lens.zero_() + return build_ragged_indices_from_dense(dense, lens, num_rows=num_rows)