| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | |
| 11 | pytestmark = pytest.mark.skipif( |
| 12 | not current_platform.is_rocm(), reason="ROCm-only aiter path" |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]: |
| 17 | ends = indptr.tolist() |
| 18 | return [indices[s:e].tolist() for s, e in zip(ends, ends[1:])] |
| 19 | |
| 20 | |
| 21 | def test_triton_sparse_mla_gate(monkeypatch) -> None: |
| 22 | import vllm._aiter_ops as aiter_ops |
| 23 | from vllm._aiter_ops import rocm_aiter_ops |
| 24 | from vllm.platforms.rocm import on_gfx950 |
| 25 | |
| 26 | warnings: list[str] = [] |
| 27 | monkeypatch.setattr( |
| 28 | aiter_ops.logger, |
| 29 | "warning_once", |
| 30 | lambda msg, *args, **kwargs: warnings.append(msg % args), |
| 31 | ) |
| 32 | |
| 33 | def enabled(aiter: bool, flag: bool) -> bool: |
| 34 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", str(int(aiter))) |
| 35 | monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", str(int(flag))) |
| 36 | rocm_aiter_ops.refresh_env_variables() |
| 37 | return bool(rocm_aiter_ops.is_triton_sparse_mla_enabled()) |
| 38 | |
| 39 | try: |
| 40 | assert not enabled(aiter=True, flag=False) and not warnings |
| 41 | assert not enabled(aiter=False, flag=True) |
| 42 | assert "VLLM_ROCM_USE_AITER is off" in warnings.pop() |
| 43 | # On gfx950 the flag takes effect; elsewhere it warns and falls back. |
| 44 | assert enabled(aiter=True, flag=True) == on_gfx950() |
| 45 | if not on_gfx950(): |
| 46 | assert "gfx950" in warnings.pop() |
| 47 | finally: |
| 48 | monkeypatch.undo() |
| 49 | rocm_aiter_ops.refresh_env_variables() |
| 50 | |
| 51 | |
| 52 | @pytest.mark.parametrize( |
| 53 | "rope_dim", |
| 54 | [0, 64], |
| 55 | ids=["rope_free", "appended_rope"], # GLM-5.3-Flash; GLM-5.1/5.2, V3.2 |
| 56 | ) |
| 57 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 58 | @torch.inference_mode() |
| 59 | def test_forward_mqa_prepares_triton_sparse_mla_inputs( |
| 60 | monkeypatch, rope_dim: int, kv_cache_dtype: str |
| 61 | ) -> None: |
| 62 | """What forward_mqa hands aiter: global slots built from the request-local |
| 63 | top-k rows, q quantized with the layer's scale, and the cache as stored.""" |
| 64 | from vllm._aiter_ops import rocm_aiter_ops |
| 65 | from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( |
| 66 | ROCMAiterMLASparseImpl, |
| 67 | ) |
| 68 | |
| 69 | device = torch.device("cuda") |
| 70 | gen = torch.Generator().manual_seed(0) |
| 71 | kv_lora_rank, num_heads, topk, block_size = 512, 16, 128, 16 |
| 72 | head_size = kv_lora_rank + rope_dim |
| 73 | fp8 = kv_cache_dtype == "fp8" |
| 74 | |
| 75 | # A decode row, then a short prefill. |
| 76 | query_lens, seq_lens = [1, 5], [40, 20] |
| 77 | num_tokens = sum(query_lens) |
| 78 | block_table = torch.randperm(8, generator=gen).view(2, 4).to(device, torch.int32) |
| 79 | topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32) |
| 80 | req_id_per_token: list[int] = [] |
| 81 | rows: list[list[int]] = [] |
| 82 | for req, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)): |
| 83 | for pos in range(seq_len - query_len, seq_len): |
| 84 | local = torch.randperm(pos + 1, generator=gen) |
| 85 | topk_indices[len(rows), : pos + 1] = local.int() |
| 86 | slots = block_table[req, local // block_size].cpu() * block_size |
| 87 | rows.append((slots + local % block_size).tolist()) |
| 88 | req_id_per_token.append(req) |
| 89 | indptr = torch.tensor([0] + [len(row) for row in rows]).cumsum(0) |
| 90 | |
| 91 | impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl) |
| 92 | impl.num_heads = num_heads |
| 93 | impl.kv_lora_rank = kv_lora_rank |
| 94 | impl.qk_rope_head_dim = rope_dim |
| 95 | impl.scale = head_size**-0.5 |
| 96 | impl.sinks = torch.randn(num_heads, generator=gen).to(device) |
| 97 | impl.kv_cache_dtype = kv_cache_dtype |
| 98 | impl.topk_indices_buffer = topk_indices.to(device) |
| 99 | impl.use_aiter_sparse_mla = True |
| 100 | layer = SimpleNamespace( |
| 101 | _q_scale=torch.tensor([0.02], device=device), |
| 102 | _k_scale=torch.tensor([0.01], device=device), |
| 103 | ) |
| 104 | metadata = SimpleNamespace( |
| 105 | num_actual_tokens=num_tokens, |
| 106 | topk_tokens=topk, |
| 107 | req_id_per_token=torch.tensor(req_id_per_token, dtype=torch.int32).to(device), |
| 108 | block_table=block_table, |
| 109 | block_size=block_size, |
| 110 | paged_kv_indptr=indptr.to(device, torch.int32), |
| 111 | paged_kv_indices=torch.zeros(num_tokens * topk, dtype=torch.int32).to(device), |
| 112 | attn_out_dtype=torch.bfloat16, |
| 113 | ) |
| 114 | cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16 |
| 115 | kv_cache = torch.zeros(8, block_size, head_size, device=device).to(cache_dtype) |
| 116 | # Two rows of cudagraph padding past the real tokens. |
| 117 | q = torch.randn(num_tokens + 2, num_heads, head_size, generator=gen).to( |
| 118 | device, torch.bfloat16 |
| 119 | ) |
| 120 | calls = [] |
| 121 | monkeypatch.setattr( |
| 122 | rocm_aiter_ops, "triton_sparse_mla_fwd", lambda *a, **k: calls.append((a, k)) |
| 123 | ) |
| 124 | |
| 125 | out, lse = impl.forward_mqa(q, kv_cache, metadata, layer) |
| 126 | |
| 127 | assert len(calls) == 1 |
| 128 | (q_in, kv_in, o, sm_scale, kv_indptr, kv_indices), kwargs = calls[0] |
| 129 | assert _rows_from_ragged(kv_indices, kv_indptr) == rows |
| 130 | assert q_in.shape[0] == num_tokens |
| 131 | if fp8: |
| 132 | assert q_in.dtype == cache_dtype |
| 133 | torch.testing.assert_close( |
| 134 | q_in.float() * layer._q_scale, q[:num_tokens].float(), atol=1e-3, rtol=0.07 |
| 135 | ) |
| 136 | else: |
| 137 | assert q_in.data_ptr() == q.data_ptr() |
| 138 | assert kv_in.shape == (8 * block_size, 1, 1, head_size) |
| 139 | assert kv_in.data_ptr() == kv_cache.data_ptr() |
| 140 | assert o.data_ptr() == out.data_ptr() and o.shape[0] == num_tokens |
| 141 | assert sm_scale == impl.scale |
| 142 | assert kwargs["kv_lora_rank"] == kv_lora_rank |
| 143 | assert kwargs["qk_rope_head_dim"] == rope_dim |
| 144 | assert kwargs["q_scale"] is layer._q_scale |
| 145 | assert kwargs["kv_scale"] is layer._k_scale |
| 146 | assert kwargs["attn_sink"] is impl.sinks |
| 147 | assert lse is None |
| 148 | |
| 149 | |
| 150 | @pytest.mark.parametrize( |
| 151 | ("model", "compress_ratio"), |
| 152 | [("deepseek_v41", 1), ("deepseek_v41", 2), ("deepseek_v4", 4)], |
| 153 | ) |
| 154 | @torch.inference_mode() |
| 155 | def test_paged_prefill_indices_match_gathered_prefill( |
| 156 | monkeypatch, model: str, compress_ratio: int |
| 157 | ) -> None: |
| 158 | """The aiter prefill reads the paged caches through the SWA builder's rows |
| 159 | and build_prefill_topk_ragged_indices, and must attend exactly the keys the |
| 160 | dequantize-and-gather prefill does.""" |
| 161 | import importlib |
| 162 | |
| 163 | from tests.v1.attention.utils import create_vllm_config |
| 164 | from vllm._aiter_ops import rocm_aiter_ops |
| 165 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 166 | from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( |
| 167 | build_prefill_topk_ragged_indices, |
| 168 | ) |
| 169 | from vllm.v1.kv_cache_interface import SlidingWindowMLASpec |
| 170 | |
| 171 | rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") |
| 172 | device = torch.device("cuda") |
| 173 | gen = torch.Generator().manual_seed(compress_ratio) |
| 174 | window, width, block_size, max_model_len = 16, 32, 16, 128 |
| 175 | comp_block = block_size // compress_ratio |
| 176 | num_compressed = max_model_len // compress_ratio |
| 177 | |
| 178 | # A decode row, then prefills with and without a prefix. |
| 179 | prefix_lens, query_lens = [20, 0, 13, 40], [1, 7, 20, 9] |
| 180 | seq_lens = [p + q for p, q in zip(prefix_lens, query_lens)] |
| 181 | num_decodes = num_decode_tokens = 1 |
| 182 | num_tokens = sum(query_lens) |
| 183 | qsl = torch.tensor([0] + query_lens).cumsum(0).to(torch.int32) |
| 184 | positions = torch.cat([torch.arange(p, s) for p, s in zip(prefix_lens, seq_lens)]) |
| 185 | token_to_req = torch.repeat_interleave( |
| 186 | torch.arange(len(query_lens)), torch.tensor(query_lens) |
| 187 | ) |
| 188 | swa_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32) |
| 189 | comp_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32) |
| 190 | |
| 191 | vllm_config = create_vllm_config( |
| 192 | model_name="facebook/opt-125m", |
| 193 | max_model_len=max_model_len, |
| 194 | block_size=block_size, |
| 195 | max_num_seqs=4, |
| 196 | max_num_batched_tokens=64, |
| 197 | hf_config_override={ |
| 198 | "compress_ratios": [compress_ratio], |
| 199 | "index_topk": width, |
| 200 | "sliding_window": window, |
| 201 | }, |
| 202 | ) |
| 203 | swa_spec = SlidingWindowMLASpec( |
| 204 | block_size=block_size, |
| 205 | num_kv_heads=1, |
| 206 | head_size=512, |
| 207 | dtype=torch.bfloat16, |
| 208 | sliding_window=window, |
| 209 | cache_dtype_str="fp8_ds_mla", |
| 210 | model_version="deepseek_v4", |
| 211 | ) |
| 212 | slots = swa_bt.cpu()[token_to_req, positions // block_size] * block_size + ( |
| 213 | positions % block_size |
| 214 | ) |
| 215 | common = CommonAttentionMetadata( |
| 216 | query_start_loc=qsl.to(device), |
| 217 | query_start_loc_cpu=qsl, |
| 218 | seq_lens=torch.tensor(seq_lens, dtype=torch.int32).to(device), |
| 219 | seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), |
| 220 | num_reqs=len(query_lens), |
| 221 | num_actual_tokens=num_tokens, |
| 222 | max_query_len=max(query_lens), |
| 223 | max_seq_len=max(seq_lens), |
| 224 | block_table_tensor=swa_bt, |
| 225 | slot_mapping=slots.to(device, torch.int64), |
| 226 | positions=positions.to(device), |
| 227 | causal=True, |
| 228 | ) |
| 229 | |
| 230 | def build(enabled: bool): |
| 231 | monkeypatch.setattr( |
| 232 | rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled |
| 233 | ) |
| 234 | return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder( |
| 235 | swa_spec, ["swa"], vllm_config, device |
| 236 | ).build(0, common) |
| 237 | |
| 238 | assert build(enabled=False).prefill_swa_ragged_indices is None |
| 239 | md = build(enabled=True) |
| 240 | assert md.num_decode_tokens == num_decode_tokens |
| 241 | |
| 242 | # Local top-k rows. Under each row's causal cap sit a failed candidate and |
| 243 | # one past the pool; past the cap, stale values neither path may read. |
| 244 | num_prefill_tokens = num_tokens - num_decode_tokens |
| 245 | topk = torch.full((num_prefill_tokens, width), -1, dtype=torch.int32) |
| 246 | for row in range(num_prefill_tokens): |
| 247 | visible = (int(positions[row + num_decode_tokens]) + 1) // compress_ratio |
| 248 | cap = min(visible, width) |
| 249 | topk[row, :cap] = torch.randperm(visible, generator=gen)[:cap].int() |
| 250 | if cap >= 2: |
| 251 | topk[row, :2] = torch.tensor([-1, num_compressed + 3]) |
| 252 | topk[row, cap:] = torch.randint( |
| 253 | 0, num_compressed, (width - cap,), generator=gen |
| 254 | ).int() |
| 255 | topk = topk.to(device) |
| 256 | |
| 257 | topk_indices, topk_indptr = build_prefill_topk_ragged_indices( |
| 258 | topk, |
| 259 | md.token_to_req_indices, |
| 260 | md.query_start_loc, |
| 261 | md.seq_lens, |
| 262 | md.is_valid_token, |
| 263 | comp_bt, |
| 264 | block_size=comp_block, |
| 265 | compress_ratio=compress_ratio, |
| 266 | num_compressed=num_compressed, |
| 267 | token_offset=num_decode_tokens, |
| 268 | ) |
| 269 | paged_keys = [ |
| 270 | sorted([("s", x) for x in swa if x >= 0] + [("c", x) for x in top if x >= 0]) |
| 271 | for swa, top in zip( |
| 272 | _rows_from_ragged( |
| 273 | md.prefill_swa_ragged_indices, md.prefill_swa_ragged_indptr |
| 274 | ), |
| 275 | _rows_from_ragged(topk_indices, topk_indptr), |
| 276 | ) |
| 277 | ] |
| 278 | |
| 279 | # The gathered path's combine, mapped from workspace offsets to cache slots. |
| 280 | M = num_compressed + window + num_tokens |
| 281 | combined, combined_lens = rocm.combine_topk_swa_indices( |
| 282 | topk, |
| 283 | md.query_start_loc[num_decodes:], |
| 284 | md.prefill_seq_lens, |
| 285 | md.prefill_gather_lens, |
| 286 | window, |
| 287 | compress_ratio, |
| 288 | width, |
| 289 | M, |
| 290 | num_compressed, |
| 291 | ) |
| 292 | gather_start = (md.prefill_seq_lens - md.prefill_gather_lens).tolist() |
| 293 | swa_bt_cpu, comp_bt_cpu = swa_bt.cpu(), comp_bt.cpu() |
| 294 | |
| 295 | def slot(block_table, req, pos, size): |
| 296 | return int(block_table[req, pos // size]) * size + pos % size |
| 297 | |
| 298 | gathered_keys = [] |
| 299 | for row, (values, n) in enumerate(zip(combined.tolist(), combined_lens.tolist())): |
| 300 | req = int(token_to_req[row + num_decode_tokens]) |
| 301 | chunk_req = req - num_decodes |
| 302 | keys = [] |
| 303 | for value in values[:n]: |
| 304 | if value < 0: |
| 305 | continue |
| 306 | offset = value - chunk_req * M |
| 307 | if offset < num_compressed: |
| 308 | keys.append(("c", slot(comp_bt_cpu, req, offset, comp_block))) |
| 309 | else: |
| 310 | pos = gather_start[chunk_req] + offset - num_compressed |
| 311 | keys.append(("s", slot(swa_bt_cpu, req, pos, block_size))) |
| 312 | gathered_keys.append(sorted(keys)) |
| 313 | |
| 314 | assert paged_keys == gathered_keys |
| 315 | |
| 316 | |
| 317 | @pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"]) |
| 318 | def test_aiter_sparse_mla_stays_off_with_kv_connector(monkeypatch, model) -> None: |
| 319 | import importlib |
| 320 | |
| 321 | from vllm._aiter_ops import rocm_aiter_ops |
| 322 | |
| 323 | rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") |
| 324 | monkeypatch.setattr(rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: True) |
| 325 | assert rocm._aiter_sparse_mla_enabled(has_kv_transfer=False) |
| 326 | assert not rocm._aiter_sparse_mla_enabled(has_kv_transfer=True) |