| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | import sys |
| 6 | from types import SimpleNamespace |
| 7 | from unittest import mock |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | if not current_platform.is_rocm(): |
| 15 | pytest.skip("ROCm AITER MLA tests", allow_module_level=True) |
| 16 | |
| 17 | from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402 |
| 18 | from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402 |
| 19 | AiterMLAHelper, |
| 20 | AiterMLAImpl, |
| 21 | AiterMLAMetadataBuilder, |
| 22 | ) |
| 23 | from vllm.v1.attention.ops.rocm_aiter_mla_merge import ( # noqa: E402 |
| 24 | merge_mla_segments_triton, |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | class _NoOpTritonKernel: |
| 29 | def __getitem__(self, grid): |
| 30 | self.grid = grid |
| 31 | return self |
| 32 | |
| 33 | def __call__(self, *args, **kwargs): |
| 34 | pass |
| 35 | |
| 36 | |
| 37 | def _lse_combine_natural( |
| 38 | out_a: torch.Tensor, |
| 39 | lse_a: torch.Tensor, |
| 40 | out_b: torch.Tensor, |
| 41 | lse_b: torch.Tensor, |
| 42 | out_dtype: torch.dtype, |
| 43 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 44 | normalizer = torch.logaddexp(lse_a, lse_b) |
| 45 | out = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) |
| 46 | out += out_b * torch.exp(lse_b - normalizer).unsqueeze(-1) |
| 47 | return out.to(out_dtype), normalizer |
| 48 | |
| 49 | |
| 50 | class _ExpandPageIndicesKernel: |
| 51 | def __getitem__(self, grid): |
| 52 | self.grid = grid |
| 53 | return self |
| 54 | |
| 55 | def __call__( |
| 56 | self, |
| 57 | page_indices, |
| 58 | block_table_tensor, |
| 59 | stride_0, |
| 60 | stride_1, |
| 61 | paged_kv_indptr, |
| 62 | *, |
| 63 | KERNEL_BLOCK_SIZE, |
| 64 | BLOCK_SIZE, |
| 65 | ): |
| 66 | # Both block-table strides are passed explicitly: the kernel no longer |
| 67 | # assumes the column stride is 1. This stand-in indexes the tensor |
| 68 | # directly so torch applies the strides for it, but the values are |
| 69 | # checked against the tensor actually handed in -- a caller passing the |
| 70 | # wrong stride is exactly the bug the parameter exists to prevent, and |
| 71 | # silently ignoring them here would hide it. |
| 72 | assert stride_0 == block_table_tensor.stride(0) |
| 73 | assert stride_1 == block_table_tensor.stride(1) |
| 74 | self.kernel_block_size = KERNEL_BLOCK_SIZE |
| 75 | for req_idx in range(self.grid[0]): |
| 76 | out_start = int(paged_kv_indptr[req_idx].item()) |
| 77 | seq_len = int( |
| 78 | (paged_kv_indptr[req_idx + 1] - paged_kv_indptr[req_idx]).item() |
| 79 | ) |
| 80 | for token_idx in range(seq_len): |
| 81 | block_id = int( |
| 82 | block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item() |
| 83 | ) |
| 84 | page_indices[out_start + token_idx] = ( |
| 85 | block_id * KERNEL_BLOCK_SIZE + token_idx % KERNEL_BLOCK_SIZE |
| 86 | ) |
| 87 | |
| 88 | |
| 89 | def _builder( |
| 90 | *, |
| 91 | mtp_decode_qlen: int, |
| 92 | has_full_cudagraphs: bool = False, |
| 93 | kernel_block_size: int = 1, |
| 94 | max_decode_rows: int = 32, |
| 95 | num_heads: int = 16, |
| 96 | kv_cache_dtype: str = "auto", |
| 97 | dcp_world_size: int = 1, |
| 98 | dcp_rank: int = 0, |
| 99 | asm_dcp_verify: bool = False, |
| 100 | decode_causal: bool = True, |
| 101 | supports_segmented_dcp_verify: bool | None = None, |
| 102 | ): |
| 103 | decode_num_heads = num_heads * dcp_world_size |
| 104 | asm_heads = ( |
| 105 | rocm_aiter_mla._asm_dcp_verify_heads(decode_num_heads) if asm_dcp_verify else 0 |
| 106 | ) |
| 107 | if supports_segmented_dcp_verify is None: |
| 108 | supports_segmented_dcp_verify = rocm_aiter_mla._segmented_dcp_verify_supported( |
| 109 | dcp_world_size, 1 |
| 110 | ) |
| 111 | stub = SimpleNamespace( |
| 112 | device=torch.device("cpu"), |
| 113 | num_heads=num_heads, |
| 114 | # DCP gathers every rank's head shard before decode, and the routing |
| 115 | # predicates read the gathered count. |
| 116 | _decode_num_heads=decode_num_heads, |
| 117 | dcp_world_size=dcp_world_size, |
| 118 | dcp_rank=dcp_rank, |
| 119 | cp_kv_cache_interleave_size=1, |
| 120 | # Mirrors the production constructor: configuration decides whether the |
| 121 | # segmented route is available, query length decides per batch. CPRR |
| 122 | # preferred still keeps this True so qlen < _MIN_CPRR_QLEN can fall back. |
| 123 | _supports_segmented_dcp_verify=supports_segmented_dcp_verify, |
| 124 | # Derived once in the real constructor, so derive it once here too. |
| 125 | _segmented_page_size=rocm_aiter_mla._segmented_mla_page_size(kernel_block_size), |
| 126 | _dcp_verify_buffers=None, |
| 127 | _graph_seq_lens=None, |
| 128 | _kv_cache_dtype_str=kv_cache_dtype, |
| 129 | # Unsized as during profiling, so the Gluon KV bound stays out of the way. |
| 130 | _kv_cache_bytes=None, |
| 131 | paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32), |
| 132 | paged_kv_indices=torch.empty(1024, dtype=torch.int32), |
| 133 | paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), |
| 134 | qo_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32), |
| 135 | compilation_config=SimpleNamespace( |
| 136 | cudagraph_mode=SimpleNamespace( |
| 137 | has_full_cudagraphs=lambda: has_full_cudagraphs |
| 138 | ) |
| 139 | ), |
| 140 | _mtp_decode_qlen=mtp_decode_qlen, |
| 141 | _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), |
| 142 | _use_persistent_metadata=False, |
| 143 | kernel_block_size=kernel_block_size, |
| 144 | _num_attention_heads=( |
| 145 | asm_heads |
| 146 | if asm_dcp_verify |
| 147 | else AiterMLAHelper.get_actual_mla_num_heads(num_heads) |
| 148 | ), |
| 149 | _mla_work_meta_data=torch.empty(1, dtype=torch.int32), |
| 150 | _mla_work_info_set=torch.empty(1, dtype=torch.int32), |
| 151 | _mla_work_indptr=torch.empty(1, dtype=torch.int32), |
| 152 | _mla_reduce_indptr=torch.empty(1, dtype=torch.int32), |
| 153 | _mla_reduce_final_map=torch.empty(1, dtype=torch.int32), |
| 154 | _mla_reduce_partial_map=torch.empty(1, dtype=torch.int32), |
| 155 | _mla_q_dtype=torch.bfloat16, |
| 156 | _mla_kv_dtype=torch.bfloat16, |
| 157 | decode_attn_out_dtype=torch.bfloat16, |
| 158 | _decode_causal=decode_causal, |
| 159 | # Default off so existing cases keep exercising segmented/persistent. |
| 160 | # Pass asm_dcp_verify=True for the CPRR route and its qlen-floor fallback. |
| 161 | _asm_dcp_verify=asm_dcp_verify, |
| 162 | _asm_dcp_verify_heads=asm_heads, |
| 163 | _g_kv_indptr_buf=( |
| 164 | torch.zeros(max_decode_rows + 1, dtype=torch.int32) |
| 165 | if asm_dcp_verify and dcp_world_size > 1 |
| 166 | else None |
| 167 | ), |
| 168 | # Only read when the asm route is on; the real builder derives it from |
| 169 | # the device CU count. |
| 170 | _mla_max_split_per_batch=256, |
| 171 | ) |
| 172 | # Bound to the stub rather than faked: the verify flatten's per-row view is |
| 173 | # part of what _build_decode is being tested for. |
| 174 | stub._build_dcp_verify_row_view = ( |
| 175 | AiterMLAMetadataBuilder._build_dcp_verify_row_view.__get__(stub) |
| 176 | ) |
| 177 | stub._fill_dcp_verify_page_table = ( |
| 178 | AiterMLAMetadataBuilder._fill_dcp_verify_page_table.__get__(stub) |
| 179 | ) |
| 180 | return stub |
| 181 | |
| 182 | |
| 183 | def test_backend_declares_uniform_batch_support(): |
| 184 | # UNIFORM/UNIFORM_BATCH is unconditional: MTP yields uniform qlen>1 and |
| 185 | # non-MTP yields qlen==1, both uniform batches. |
| 186 | assert ( |
| 187 | AiterMLAMetadataBuilder.query_len_support |
| 188 | == rocm_aiter_mla.QueryLenSupport.UNIFORM |
| 189 | ) |
| 190 | assert ( |
| 191 | AiterMLAMetadataBuilder._cudagraph_support |
| 192 | == rocm_aiter_mla.AttentionCGSupport.UNIFORM_BATCH |
| 193 | ) |
| 194 | |
| 195 | |
| 196 | @pytest.mark.parametrize("causal", [True, False]) |
| 197 | def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal): |
| 198 | seen = [] |
| 199 | |
| 200 | def fake_common_build(self, *args): |
| 201 | seen.append(self._decode_causal) |
| 202 | return SimpleNamespace(decode=None, prefill=None) |
| 203 | |
| 204 | monkeypatch.setattr( |
| 205 | rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build |
| 206 | ) |
| 207 | builder = object.__new__(AiterMLAMetadataBuilder) |
| 208 | builder._fp8_prefill_enabled = False |
| 209 | |
| 210 | builder.build(0, SimpleNamespace(causal=causal)) |
| 211 | |
| 212 | assert seen == [causal] |
| 213 | |
| 214 | |
| 215 | def test_dcp_verify_row_view_is_causal_per_row(): |
| 216 | """Row t must cover this rank's share of positions [0, seq_len - qlen + t]. |
| 217 | |
| 218 | Regression guard: giving every row of a request the committed prefix only |
| 219 | drops the whole verify block, including each row's own token, since nothing |
| 220 | else adds it back on the target path. |
| 221 | """ |
| 222 | qlen = 4 |
| 223 | builder = _builder(mtp_decode_qlen=qlen, dcp_world_size=2, kernel_block_size=2) |
| 224 | block_table = torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32) |
| 225 | |
| 226 | view = builder._build_dcp_verify_row_view( |
| 227 | qlen, |
| 228 | block_table, |
| 229 | torch.tensor([10, 12], dtype=torch.int32), |
| 230 | ) |
| 231 | |
| 232 | # seq_len 10 -> rows see 7, 8, 9, 10 global positions, of which rank 0 of 2 |
| 233 | # holds ceil(n / 2); seq_len 12 -> 9, 10, 11, 12. |
| 234 | assert view.row_lens.tolist() == [4, 4, 5, 5, 5, 5, 6, 6] |
| 235 | assert view.qo_indptr.tolist() == list(range(2 * qlen + 1)) |
| 236 | # One query per row, so the page table is the request's ascending shard and |
| 237 | # only the row length selects the causal prefix. |
| 238 | assert view.block_table.tolist() == [[7, 8, 9]] * qlen + [[10, 11, 12]] * qlen |
| 239 | assert view.max_kv_seq_len == 6 |
| 240 | |
| 241 | |
| 242 | def test_dcp_verify_rows_sum_to_the_global_window_across_ranks(): |
| 243 | """Every rank's rows together must cover each row's window exactly once.""" |
| 244 | qlen = 4 |
| 245 | dcp_world_size = 2 |
| 246 | seq_lens = torch.tensor([10, 12], dtype=torch.int32) |
| 247 | per_rank = [] |
| 248 | for dcp_rank in range(dcp_world_size): |
| 249 | builder = _builder( |
| 250 | mtp_decode_qlen=qlen, |
| 251 | dcp_world_size=dcp_world_size, |
| 252 | dcp_rank=dcp_rank, |
| 253 | kernel_block_size=2, |
| 254 | ) |
| 255 | view = builder._build_dcp_verify_row_view( |
| 256 | qlen, |
| 257 | torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32), |
| 258 | seq_lens, |
| 259 | ) |
| 260 | per_rank.append(view.row_lens.tolist()) |
| 261 | |
| 262 | expected = [ |
| 263 | int(seq_len) - qlen + t + 1 |
| 264 | for seq_len in seq_lens.tolist() |
| 265 | for t in range(qlen) |
| 266 | ] |
| 267 | assert [sum(lens) for lens in zip(*per_rank)] == expected |
| 268 | |
| 269 | |
| 270 | def test_dcp_verify_row_view_uses_static_graph_bound(): |
| 271 | """Under full graphs the bound comes from the buffers, not the batch.""" |
| 272 | qlen = 3 |
| 273 | builder = _builder( |
| 274 | mtp_decode_qlen=qlen, |
| 275 | dcp_world_size=8, |
| 276 | has_full_cudagraphs=True, |
| 277 | kernel_block_size=1536, |
| 278 | ) |
| 279 | builder._dcp_verify_buffers = rocm_aiter_mla.AiterMLADCPVerifyMetadata( |
| 280 | row_lens=torch.zeros(qlen, dtype=torch.int32), |
| 281 | block_table=torch.zeros((qlen, 24), dtype=torch.int32), |
| 282 | qo_indptr=torch.zeros(qlen + 1, dtype=torch.int32), |
| 283 | page_size=128, |
| 284 | max_kv_seq_len=3072, |
| 285 | ) |
| 286 | |
| 287 | view = builder._build_dcp_verify_row_view( |
| 288 | qlen, |
| 289 | torch.arange(2, dtype=torch.int32).view(1, 2), |
| 290 | torch.tensor([1000], dtype=torch.int32), |
| 291 | ) |
| 292 | |
| 293 | buffers = builder._dcp_verify_buffers |
| 294 | assert view.row_lens.tolist() == [125] * qlen |
| 295 | assert view.block_table.data_ptr() == buffers.block_table.data_ptr() |
| 296 | assert view.qo_indptr.data_ptr() == buffers.qo_indptr.data_ptr() |
| 297 | assert view.block_table[0].tolist() == list(range(24)) |
| 298 | assert view.max_kv_seq_len == 3072 |
| 299 | |
| 300 | |
| 301 | def test_dcp_fp8_verify_build_uses_segmented(monkeypatch): |
| 302 | qlen = 4 |
| 303 | monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) |
| 304 | |
| 305 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 306 | _builder( |
| 307 | mtp_decode_qlen=qlen, |
| 308 | dcp_world_size=2, |
| 309 | kv_cache_dtype="fp8", |
| 310 | kernel_block_size=2, |
| 311 | ), |
| 312 | block_table_tensor=torch.tensor([[0, 1, 2], [10, 11, 12]], dtype=torch.int32), |
| 313 | seq_lens_device=torch.tensor([5, 6], dtype=torch.int32), |
| 314 | max_seq_len=6, |
| 315 | query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), |
| 316 | query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), |
| 317 | num_decode_tokens=2 * qlen, |
| 318 | max_query_len=qlen, |
| 319 | dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32), |
| 320 | ) |
| 321 | |
| 322 | assert metadata.dcp_verify is not None |
| 323 | |
| 324 | |
| 325 | def _cprr_decode_batch(builder, *, qlen: int, num_reqs: int = 1): |
| 326 | seq_lens = torch.full((num_reqs,), 16, dtype=torch.int32) |
| 327 | query_start_loc = torch.arange(0, (num_reqs + 1) * qlen, qlen, dtype=torch.int32) |
| 328 | return AiterMLAMetadataBuilder._build_decode( |
| 329 | builder, |
| 330 | block_table_tensor=torch.zeros(num_reqs, 8, dtype=torch.int32), |
| 331 | seq_lens_device=seq_lens, |
| 332 | max_seq_len=16, |
| 333 | query_start_loc_cpu=query_start_loc, |
| 334 | query_start_loc_device=query_start_loc, |
| 335 | num_decode_tokens=num_reqs * qlen, |
| 336 | max_query_len=qlen, |
| 337 | dcp_tot_seq_lens_device=seq_lens * 8, |
| 338 | ) |
| 339 | |
| 340 | |
| 341 | def test_asm_preferred_qlen2_falls_back_to_segmented(monkeypatch): |
| 342 | """K=4 can clamp to qlen 2; CPRR has no kernel, so use segmented MLA.""" |
| 343 | monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) |
| 344 | get_mla_metadata_v1 = mock.MagicMock() |
| 345 | monkeypatch.setitem( |
| 346 | sys.modules, |
| 347 | "aiter", |
| 348 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 349 | ) |
| 350 | |
| 351 | metadata = _cprr_decode_batch( |
| 352 | _builder( |
| 353 | mtp_decode_qlen=5, |
| 354 | dcp_world_size=8, |
| 355 | num_heads=12, |
| 356 | asm_dcp_verify=True, |
| 357 | kernel_block_size=2, |
| 358 | ), |
| 359 | qlen=2, |
| 360 | ) |
| 361 | |
| 362 | assert metadata.dcp_verify is not None |
| 363 | assert metadata.g_kv_indptr is None |
| 364 | assert not metadata.has_persistent_metadata |
| 365 | get_mla_metadata_v1.assert_not_called() |
| 366 | |
| 367 | |
| 368 | def test_asm_preferred_qlen_at_cprr_floor_stays_on_cprr(monkeypatch): |
| 369 | monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) |
| 370 | get_mla_metadata_v1 = mock.MagicMock() |
| 371 | monkeypatch.setitem( |
| 372 | sys.modules, |
| 373 | "aiter", |
| 374 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 375 | ) |
| 376 | monkeypatch.setattr( |
| 377 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 378 | ) |
| 379 | |
| 380 | metadata = _cprr_decode_batch( |
| 381 | _builder( |
| 382 | mtp_decode_qlen=5, |
| 383 | dcp_world_size=8, |
| 384 | num_heads=12, |
| 385 | asm_dcp_verify=True, |
| 386 | ), |
| 387 | qlen=rocm_aiter_mla._MIN_CPRR_QLEN, |
| 388 | ) |
| 389 | |
| 390 | assert metadata.dcp_verify is None |
| 391 | assert metadata.g_kv_indptr is not None |
| 392 | assert metadata.has_persistent_metadata |
| 393 | assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is True |
| 394 | |
| 395 | |
| 396 | def test_cprr_capable_single_token_decode_keeps_persistent_metadata(monkeypatch): |
| 397 | """The CPRR qlen floor must not suppress the ordinary qlen-1 schedule.""" |
| 398 | get_mla_metadata_v1 = mock.MagicMock() |
| 399 | monkeypatch.setitem( |
| 400 | sys.modules, |
| 401 | "aiter", |
| 402 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 403 | ) |
| 404 | monkeypatch.setattr( |
| 405 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 406 | ) |
| 407 | |
| 408 | metadata = _cprr_decode_batch( |
| 409 | _builder( |
| 410 | mtp_decode_qlen=5, |
| 411 | dcp_world_size=8, |
| 412 | num_heads=12, |
| 413 | asm_dcp_verify=True, |
| 414 | ), |
| 415 | qlen=1, |
| 416 | ) |
| 417 | |
| 418 | assert metadata.g_kv_indptr is None |
| 419 | assert metadata.has_persistent_metadata |
| 420 | assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is None |
| 421 | |
| 422 | |
| 423 | def test_non_causal_dcp_block_uses_plain_mask0_metadata(monkeypatch): |
| 424 | """A non-causal block needs no CPRR global-position window or head pad.""" |
| 425 | get_mla_metadata_v1 = mock.MagicMock() |
| 426 | monkeypatch.setitem( |
| 427 | sys.modules, |
| 428 | "aiter", |
| 429 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 430 | ) |
| 431 | monkeypatch.setattr( |
| 432 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 433 | ) |
| 434 | builder = _builder( |
| 435 | mtp_decode_qlen=5, |
| 436 | dcp_world_size=8, |
| 437 | num_heads=12, |
| 438 | asm_dcp_verify=True, |
| 439 | ) |
| 440 | builder._decode_causal = False |
| 441 | |
| 442 | metadata = _cprr_decode_batch(builder, qlen=4) |
| 443 | |
| 444 | assert metadata.dcp_verify is None |
| 445 | assert metadata.g_kv_indptr is None |
| 446 | assert metadata.has_persistent_metadata |
| 447 | assert get_mla_metadata_v1.call_args.args[3] == 96 |
| 448 | assert get_mla_metadata_v1.call_args.args[5] is False |
| 449 | assert get_mla_metadata_v1.call_args.kwargs.get("is_cp_round_robin") is None |
| 450 | |
| 451 | |
| 452 | def test_asm_qlen2_without_segmented_fails_during_build(monkeypatch): |
| 453 | monkeypatch.setitem( |
| 454 | sys.modules, |
| 455 | "aiter", |
| 456 | SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()), |
| 457 | ) |
| 458 | monkeypatch.setattr( |
| 459 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 460 | ) |
| 461 | |
| 462 | with pytest.raises(RuntimeError, match="requires either segmented MLA"): |
| 463 | _cprr_decode_batch( |
| 464 | _builder( |
| 465 | mtp_decode_qlen=5, |
| 466 | dcp_world_size=8, |
| 467 | num_heads=12, |
| 468 | asm_dcp_verify=True, |
| 469 | supports_segmented_dcp_verify=False, |
| 470 | ), |
| 471 | qlen=2, |
| 472 | ) |
| 473 | |
| 474 | |
| 475 | def test_non_causal_dcp_block_bypasses_cprr_in_forward(monkeypatch): |
| 476 | """Process-level CPRR capability must not pad or route a non-causal batch.""" |
| 477 | captured = {} |
| 478 | |
| 479 | def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs): |
| 480 | captured["q_heads"] = q.shape[1] |
| 481 | captured["kwargs"] = kwargs |
| 482 | return None, torch.zeros(q.shape[0], q.shape[1]) |
| 483 | |
| 484 | monkeypatch.setattr( |
| 485 | rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode |
| 486 | ) |
| 487 | |
| 488 | impl = object.__new__(AiterMLAImpl) |
| 489 | impl.num_heads = 12 |
| 490 | impl.dcp_world_size = 8 |
| 491 | impl.kv_cache_dtype = "auto" |
| 492 | impl.kv_lora_rank = 512 |
| 493 | impl.qk_rope_head_dim = 64 |
| 494 | impl.scale = 576**-0.5 |
| 495 | decode = SimpleNamespace( |
| 496 | max_qo_len=4, |
| 497 | qo_indptr=torch.tensor([0, 4], dtype=torch.int32), |
| 498 | paged_kv_indptr=torch.tensor([0, 1], dtype=torch.int32), |
| 499 | paged_kv_indices=torch.zeros(1, dtype=torch.int32), |
| 500 | paged_kv_last_page_len=torch.ones(1, dtype=torch.int32), |
| 501 | use_gluon_decode=False, |
| 502 | use_gluon_verify=False, |
| 503 | dcp_route=rocm_aiter_mla._DCPDecodeRoute.PLAIN, |
| 504 | dcp_verify=None, |
| 505 | g_kv_indptr=None, |
| 506 | has_persistent_metadata=False, |
| 507 | attn_out_dtype=torch.bfloat16, |
| 508 | asm_decode_num_heads=0, |
| 509 | mla_num_kv_splits=0, |
| 510 | cp_world_size=8, |
| 511 | cp_rank=0, |
| 512 | min_kv_seq_len=1, |
| 513 | ) |
| 514 | attn_metadata = SimpleNamespace(decode=decode, causal=False, work_meta_data=None) |
| 515 | layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0)) |
| 516 | q = torch.zeros(4, 96, 576, dtype=torch.bfloat16) |
| 517 | |
| 518 | output, lse = impl.forward_mqa(q, torch.zeros(1, 1, 576), attn_metadata, layer) |
| 519 | |
| 520 | assert captured["q_heads"] == 96 |
| 521 | assert captured["kwargs"]["causal"] is False |
| 522 | assert "g_kv_indptr" not in captured["kwargs"] |
| 523 | assert "num_kv_splits" not in captured["kwargs"] |
| 524 | assert output.shape[1] == 96 |
| 525 | assert lse is not None and lse.shape == (4, 96) |
| 526 | |
| 527 | |
| 528 | def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch): |
| 529 | """Single-token DCP decode must come back with an LSE the merge can use. |
| 530 | |
| 531 | This is the path every DCP step takes outside verification. The vLLM custom |
| 532 | op drops aiter's LSE, so the impl calls aiter directly; without an LSE the |
| 533 | MLA layer's cross-rank combine asserts. The head count is deliberately a |
| 534 | non-divisor of 16 so both the output and the LSE have to be unpadded back |
| 535 | to the gathered head count. |
| 536 | """ |
| 537 | num_heads, dcp_world_size = 6, 2 |
| 538 | decode_heads = num_heads * dcp_world_size # 12, padded to 16 by aiter |
| 539 | num_tokens, head_dim = 3, 576 |
| 540 | captured = {} |
| 541 | |
| 542 | def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs): |
| 543 | captured["q_heads"] = q.shape[1] |
| 544 | captured["return_lse"] = kwargs.get("return_lse") |
| 545 | return None, torch.zeros(num_tokens, q.shape[1]) |
| 546 | |
| 547 | monkeypatch.setattr( |
| 548 | rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode |
| 549 | ) |
| 550 | |
| 551 | impl = object.__new__(AiterMLAImpl) |
| 552 | impl.num_heads = num_heads |
| 553 | impl.dcp_world_size = dcp_world_size |
| 554 | impl.kv_cache_dtype = "auto" |
| 555 | impl.kv_lora_rank = 512 |
| 556 | impl.qk_rope_head_dim = 64 |
| 557 | impl.scale = head_dim**-0.5 |
| 558 | decode = SimpleNamespace( |
| 559 | max_qo_len=1, |
| 560 | qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32), |
| 561 | paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32), |
| 562 | paged_kv_indices=torch.zeros(1, dtype=torch.int32), |
| 563 | paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32), |
| 564 | use_gluon_decode=False, |
| 565 | use_gluon_verify=False, |
| 566 | dcp_route=rocm_aiter_mla._DCPDecodeRoute.PLAIN, |
| 567 | dcp_verify=None, |
| 568 | has_persistent_metadata=False, |
| 569 | attn_out_dtype=torch.bfloat16, |
| 570 | # 0 == the round-robin (cprr) asm route is off, this rank's decode takes |
| 571 | # the existing path. Mirrors the AiterMLADecodeMetadata default. |
| 572 | asm_decode_num_heads=0, |
| 573 | g_kv_indptr=None, |
| 574 | ) |
| 575 | attn_metadata = SimpleNamespace(decode=decode, causal=True, work_meta_data=None) |
| 576 | layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0)) |
| 577 | |
| 578 | q = torch.zeros(num_tokens, decode_heads, head_dim, dtype=torch.bfloat16) |
| 579 | output, lse = impl.forward_mqa(q, torch.zeros(1, 1, head_dim), attn_metadata, layer) |
| 580 | |
| 581 | assert captured["return_lse"] is True |
| 582 | # aiter is handed the padded head count, but both results come back at the |
| 583 | # gathered count the cross-rank merge expects. |
| 584 | assert captured["q_heads"] == 16 |
| 585 | assert lse is not None |
| 586 | assert output.shape[1] == decode_heads |
| 587 | assert lse.shape == (num_tokens, decode_heads) |
| 588 | |
| 589 | |
| 590 | def test_verify_partial_attention_merge(): |
| 591 | device = torch.device("cuda") |
| 592 | out_a = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]], device=device) |
| 593 | out_b = torch.tensor([[[5.0, 6.0]], [[7.0, 8.0]]], device=device) |
| 594 | lse_a = torch.tensor([[0.0], [1.0]], device=device) |
| 595 | lse_b = torch.tensor([[1.0], [0.0]], device=device) |
| 596 | |
| 597 | out, lse = _lse_combine_natural(out_a, lse_a, out_b, lse_b, torch.float32) |
| 598 | normalizer = torch.logaddexp(lse_a, lse_b) |
| 599 | expected = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) + out_b * torch.exp( |
| 600 | lse_b - normalizer |
| 601 | ).unsqueeze(-1) |
| 602 | |
| 603 | torch.testing.assert_close(out, expected) |
| 604 | torch.testing.assert_close(lse, normalizer) |
| 605 | |
| 606 | |
| 607 | def test_segmented_verify_reduce_returns_natural_lse_and_masks_empty_rows(): |
| 608 | segment_output = torch.tensor( |
| 609 | [1.0, 3.0, 99.0, 99.0], |
| 610 | device="cuda", |
| 611 | ).reshape(2, 1, 2, 1) |
| 612 | segment_max = torch.tensor( |
| 613 | [[[0.0, 1.0]], [[99.0, 99.0]]], |
| 614 | device="cuda", |
| 615 | ) |
| 616 | segment_expsum = torch.ones_like(segment_max) |
| 617 | |
| 618 | output, lse = merge_mla_segments_triton( |
| 619 | segment_output, |
| 620 | segment_max, |
| 621 | segment_expsum, |
| 622 | torch.tensor([2, 0], dtype=torch.int32, device="cuda"), |
| 623 | tile_size=1, |
| 624 | out_dtype=torch.float32, |
| 625 | ) |
| 626 | |
| 627 | torch.testing.assert_close(output[0], torch.tensor([[7.0 / 3]], device="cuda")) |
| 628 | assert output[1].item() == 0 |
| 629 | torch.testing.assert_close(lse[0], torch.tensor([math.log(3.0)], device="cuda")) |
| 630 | assert lse[1].item() == float("-inf") |
| 631 | |
| 632 | |
| 633 | def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): |
| 634 | """Exercise the target qlen>1 path and merge two rank-local results.""" |
| 635 | monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True) |
| 636 | |
| 637 | torch.manual_seed(0) |
| 638 | device = torch.device("cuda") |
| 639 | dcp_world_size = 2 |
| 640 | qlen = 3 |
| 641 | num_heads = 64 |
| 642 | kv_lora_rank = 512 |
| 643 | rope_dim = 64 |
| 644 | head_dim = kv_lora_rank + rope_dim |
| 645 | global_seq_len = 259 |
| 646 | block_size = 128 |
| 647 | kv_scale = 0.02 |
| 648 | sm_scale = head_dim**-0.5 |
| 649 | |
| 650 | q_nope = torch.randn( |
| 651 | qlen, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device |
| 652 | ) |
| 653 | q_pe = torch.randn(qlen, num_heads, rope_dim, dtype=torch.bfloat16, device=device) |
| 654 | kv_source = torch.randn( |
| 655 | global_seq_len, head_dim, dtype=torch.float32, device=device |
| 656 | ) |
| 657 | kv_fp8 = (kv_source / kv_scale).to(torch.float8_e4m3fn) |
| 658 | kv_dequant = kv_fp8.float() * kv_scale |
| 659 | |
| 660 | partials = [] |
| 661 | for dcp_rank in range(dcp_world_size): |
| 662 | local_kv = kv_fp8[dcp_rank::dcp_world_size] |
| 663 | num_blocks = math.ceil(local_kv.shape[0] / block_size) |
| 664 | kv_cache = torch.zeros( |
| 665 | num_blocks, |
| 666 | block_size, |
| 667 | head_dim, |
| 668 | dtype=torch.float8_e4m3fn, |
| 669 | device=device, |
| 670 | ) |
| 671 | kv_cache.view(-1, head_dim)[: local_kv.shape[0]].copy_(local_kv) |
| 672 | |
| 673 | builder = _builder( |
| 674 | mtp_decode_qlen=qlen, |
| 675 | kernel_block_size=block_size, |
| 676 | num_heads=num_heads // dcp_world_size, |
| 677 | kv_cache_dtype="fp8", |
| 678 | dcp_world_size=dcp_world_size, |
| 679 | dcp_rank=dcp_rank, |
| 680 | ) |
| 681 | block_table = torch.arange( |
| 682 | num_blocks, dtype=torch.int32, device=device |
| 683 | ).unsqueeze(0) |
| 684 | view = builder._build_dcp_verify_row_view( |
| 685 | qlen, |
| 686 | block_table, |
| 687 | torch.tensor([global_seq_len], dtype=torch.int32, device=device), |
| 688 | ) |
| 689 | decode = SimpleNamespace( |
| 690 | max_qo_len=qlen, |
| 691 | paged_kv_indptr=torch.tensor( |
| 692 | [0, local_kv.shape[0]], dtype=torch.int32, device=device |
| 693 | ), |
| 694 | paged_kv_indices=torch.arange( |
| 695 | local_kv.shape[0], dtype=torch.int32, device=device |
| 696 | ), |
| 697 | use_gluon_decode=False, |
| 698 | use_gluon_verify=False, |
| 699 | dcp_route=rocm_aiter_mla._DCPDecodeRoute.SEGMENTED, |
| 700 | dcp_verify=view, |
| 701 | attn_out_dtype=torch.bfloat16, |
| 702 | ) |
| 703 | attn_metadata = SimpleNamespace(decode=decode, causal=True) |
| 704 | impl = object.__new__(AiterMLAImpl) |
| 705 | impl.num_heads = num_heads // dcp_world_size |
| 706 | impl.dcp_world_size = dcp_world_size |
| 707 | impl.kv_cache_dtype = "fp8" |
| 708 | impl.kv_lora_rank = kv_lora_rank |
| 709 | impl.qk_rope_head_dim = rope_dim |
| 710 | impl.scale = sm_scale |
| 711 | layer = SimpleNamespace(_k_scale=torch.tensor(kv_scale, device=device)) |
| 712 | |
| 713 | partials.append( |
| 714 | impl.forward_mqa((q_nope, q_pe), kv_cache, attn_metadata, layer) |
| 715 | ) |
| 716 | |
| 717 | output, lse = partials[0] |
| 718 | assert lse is not None |
| 719 | for rank_output, rank_lse in partials[1:]: |
| 720 | assert rank_lse is not None |
| 721 | output, lse = _lse_combine_natural( |
| 722 | output.float(), lse, rank_output.float(), rank_lse, torch.float32 |
| 723 | ) |
| 724 | |
| 725 | reference = torch.empty_like(output) |
| 726 | q_nope_fp32 = q_nope.float() |
| 727 | q_pe_fp32 = q_pe.float() |
| 728 | for query_pos in range(qlen): |
| 729 | visible = global_seq_len - qlen + query_pos + 1 |
| 730 | keys = kv_dequant[:visible] |
| 731 | scores = torch.einsum( |
| 732 | "hd,nd->hn", q_nope_fp32[query_pos], keys[:, :kv_lora_rank] |
| 733 | ) |
| 734 | scores += torch.einsum( |
| 735 | "hd,nd->hn", q_pe_fp32[query_pos], keys[:, kv_lora_rank:] |
| 736 | ) |
| 737 | probs = torch.softmax(scores * sm_scale, dim=-1) |
| 738 | reference[query_pos] = probs @ keys[:, :kv_lora_rank] |
| 739 | |
| 740 | torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2) |
| 741 | |
| 742 | |
| 743 | @pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128]) |
| 744 | @pytest.mark.parametrize( |
| 745 | "spec_method, parallel_drafting", |
| 746 | [ |
| 747 | ("deepseek_mtp", False), |
| 748 | # A drafter that is not one of the historically recognized MTP methods, |
| 749 | # and a parallel one, so the threshold is 1 + 2 * num_spec rather than |
| 750 | # 1 + num_spec. Sizing the metadata off a method name instead leaves |
| 751 | # these at qlen=1 while the router still admits the full range. |
| 752 | ("custom", True), |
| 753 | ("eagle", False), |
| 754 | ], |
| 755 | ) |
| 756 | def test_mtp_builder_init_sizes_native_fp8_metadata( |
| 757 | monkeypatch, num_heads, spec_method, parallel_drafting |
| 758 | ): |
| 759 | """Aiter init sizes the metadata for every query length decode can be handed. |
| 760 | |
| 761 | Sweeping num_heads asserts metadata is sized for the padded decode shape, |
| 762 | covering Kimi-K3 TP4's 24 -> 32 head path and native fp8 nhead=32 folding. |
| 763 | """ |
| 764 | dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16") |
| 765 | info_calls = [] |
| 766 | |
| 767 | def get_mla_metadata_info_v1( |
| 768 | max_batch_size, |
| 769 | max_qo_len, |
| 770 | num_attention_heads, |
| 771 | q_dtype, |
| 772 | kv_dtype, |
| 773 | *, |
| 774 | is_sparse, |
| 775 | fast_mode, |
| 776 | ): |
| 777 | info_calls.append( |
| 778 | { |
| 779 | "max_batch_size": max_batch_size, |
| 780 | "max_qo_len": max_qo_len, |
| 781 | "num_attention_heads": num_attention_heads, |
| 782 | "q_dtype": q_dtype, |
| 783 | "kv_dtype": kv_dtype, |
| 784 | "is_sparse": is_sparse, |
| 785 | "fast_mode": fast_mode, |
| 786 | } |
| 787 | ) |
| 788 | return tuple((1, torch.int32) for _ in range(6)) |
| 789 | |
| 790 | def init_common_builder(self, *args, **kwargs): |
| 791 | self.num_heads = num_heads |
| 792 | self.dcp_world_size = 1 |
| 793 | self.cp_kv_cache_interleave_size = 1 |
| 794 | # Mirror what _init_reorder_batch_threshold would have left behind: the |
| 795 | # metadata is sized from the routing threshold, so a stub that skips it |
| 796 | # would size for qlen=1 and hide the very mismatch this test covers. |
| 797 | spec = config.speculative_config |
| 798 | self.reorder_batch_threshold = ( |
| 799 | 1 + (2 if spec.parallel_drafting else 1) * spec.num_speculative_tokens |
| 800 | ) |
| 801 | |
| 802 | monkeypatch.setitem( |
| 803 | sys.modules, |
| 804 | "aiter", |
| 805 | SimpleNamespace( |
| 806 | dtypes=dtypes, |
| 807 | get_mla_metadata_info_v1=get_mla_metadata_info_v1, |
| 808 | ), |
| 809 | ) |
| 810 | monkeypatch.setattr( |
| 811 | rocm_aiter_mla.MLACommonMetadataBuilder, |
| 812 | "__init__", |
| 813 | init_common_builder, |
| 814 | ) |
| 815 | monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False) |
| 816 | |
| 817 | config = SimpleNamespace( |
| 818 | speculative_config=SimpleNamespace( |
| 819 | method=spec_method, |
| 820 | num_speculative_tokens=3, |
| 821 | parallel_drafting=parallel_drafting, |
| 822 | ), |
| 823 | parallel_config=SimpleNamespace( |
| 824 | tensor_parallel_size=8, |
| 825 | decode_context_parallel_size=1, |
| 826 | cp_kv_cache_interleave_size=1, |
| 827 | ), |
| 828 | model_config=SimpleNamespace( |
| 829 | max_model_len=16, |
| 830 | dtype=torch.bfloat16, |
| 831 | get_num_attention_heads=lambda parallel_config: num_heads, |
| 832 | ), |
| 833 | scheduler_config=SimpleNamespace(max_num_seqs=2), |
| 834 | cache_config=SimpleNamespace(cache_dtype="fp8_e4m3", num_gpu_blocks=None), |
| 835 | compilation_config=SimpleNamespace( |
| 836 | cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False), |
| 837 | # Empty: the per-layer head-count probe finds no attention layer and |
| 838 | # falls back to the model config above. |
| 839 | static_forward_context={}, |
| 840 | ), |
| 841 | ) |
| 842 | builder = AiterMLAMetadataBuilder( |
| 843 | kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16), |
| 844 | layer_names=["layer.0"], |
| 845 | vllm_config=config, |
| 846 | device=torch.device("cpu"), |
| 847 | ) |
| 848 | |
| 849 | assert info_calls == [ |
| 850 | { |
| 851 | "max_batch_size": config.scheduler_config.max_num_seqs, |
| 852 | "max_qo_len": builder.reorder_batch_threshold, |
| 853 | "num_attention_heads": AiterMLAHelper.get_actual_mla_num_heads(num_heads), |
| 854 | "q_dtype": dtypes.fp8, |
| 855 | "kv_dtype": dtypes.fp8, |
| 856 | "is_sparse": False, |
| 857 | "fast_mode": True, |
| 858 | } |
| 859 | ] |
| 860 | assert builder._mla_q_dtype == dtypes.fp8 |
| 861 | assert builder._mla_kv_dtype == dtypes.fp8 |
| 862 | |
| 863 | |
| 864 | def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch): |
| 865 | get_mla_metadata_v1 = mock.MagicMock() |
| 866 | monkeypatch.setitem( |
| 867 | sys.modules, |
| 868 | "aiter", |
| 869 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 870 | ) |
| 871 | monkeypatch.setattr( |
| 872 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 873 | ) |
| 874 | |
| 875 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 876 | _builder(mtp_decode_qlen=4), |
| 877 | block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), |
| 878 | seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), |
| 879 | max_seq_len=7, |
| 880 | query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32), |
| 881 | query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32), |
| 882 | num_decode_tokens=8, |
| 883 | max_query_len=4, |
| 884 | dcp_tot_seq_lens_device=None, |
| 885 | ) |
| 886 | |
| 887 | assert metadata.max_qo_len == 4 |
| 888 | assert torch.equal(metadata.seq_lens, torch.tensor([7, 5], dtype=torch.int32)) |
| 889 | assert torch.equal(metadata.qo_indptr, torch.tensor([0, 4, 8], dtype=torch.int32)) |
| 890 | assert metadata.has_persistent_metadata |
| 891 | assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == 4 |
| 892 | assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4 |
| 893 | |
| 894 | |
| 895 | def test_min_kv_seq_len_ignores_cudagraph_padding_rows(monkeypatch): |
| 896 | """min_kv_seq_len must not be driven by padded dummy requests. |
| 897 | |
| 898 | Full-CG verify pins zero-qo rows to max_qo_len for paged_kv metadata; taking |
| 899 | per_req_len.min() over all rows then reports qlen instead of a real KV length |
| 900 | and collapses Gluon KV-split parallelism on single-request decode. |
| 901 | """ |
| 902 | monkeypatch.setitem( |
| 903 | sys.modules, |
| 904 | "aiter", |
| 905 | SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()), |
| 906 | ) |
| 907 | monkeypatch.setattr( |
| 908 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 909 | ) |
| 910 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) |
| 911 | monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") |
| 912 | |
| 913 | mtp_qlen = 8 |
| 914 | num_reqs = 8 |
| 915 | active_seq_len = 1032 |
| 916 | seq_lens = torch.tensor([active_seq_len] + [0] * (num_reqs - 1), dtype=torch.int32) |
| 917 | query_start_loc = torch.tensor( |
| 918 | [0, mtp_qlen] + [mtp_qlen] * (num_reqs - 1), dtype=torch.int32 |
| 919 | ) |
| 920 | |
| 921 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 922 | _builder( |
| 923 | mtp_decode_qlen=mtp_qlen, |
| 924 | has_full_cudagraphs=True, |
| 925 | max_decode_rows=num_reqs, |
| 926 | num_heads=12, |
| 927 | ), |
| 928 | block_table_tensor=torch.zeros(num_reqs, active_seq_len, dtype=torch.int32), |
| 929 | seq_lens_device=seq_lens, |
| 930 | max_seq_len=active_seq_len, |
| 931 | query_start_loc_cpu=query_start_loc, |
| 932 | query_start_loc_device=query_start_loc, |
| 933 | num_decode_tokens=num_reqs * mtp_qlen, |
| 934 | max_query_len=mtp_qlen, |
| 935 | dcp_tot_seq_lens_device=None, |
| 936 | ) |
| 937 | |
| 938 | assert metadata.min_kv_seq_len == active_seq_len |
| 939 | |
| 940 | |
| 941 | def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr( |
| 942 | monkeypatch, |
| 943 | ): |
| 944 | """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759.""" |
| 945 | get_mla_metadata_v1 = mock.MagicMock() |
| 946 | monkeypatch.setitem( |
| 947 | sys.modules, |
| 948 | "aiter", |
| 949 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 950 | ) |
| 951 | monkeypatch.setattr( |
| 952 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 953 | ) |
| 954 | |
| 955 | mtp_qlen = 4 |
| 956 | seq_lens = torch.tensor([7, 0], dtype=torch.int32) |
| 957 | qo_lens = torch.tensor([mtp_qlen, 0], dtype=torch.int32) |
| 958 | expected_seq_lens = torch.where(qo_lens > 0, seq_lens, mtp_qlen) |
| 959 | expected_paged_kv_indptr = torch.cat( |
| 960 | [ |
| 961 | torch.zeros(1, dtype=torch.int32), |
| 962 | expected_seq_lens.cumsum(dim=0, dtype=torch.int32), |
| 963 | ] |
| 964 | ) |
| 965 | expected_qo_indptr = torch.arange( |
| 966 | 0, |
| 967 | (seq_lens.numel() + 1) * mtp_qlen, |
| 968 | step=mtp_qlen, |
| 969 | dtype=torch.int32, |
| 970 | ) |
| 971 | |
| 972 | builder = _builder( |
| 973 | mtp_decode_qlen=mtp_qlen, |
| 974 | has_full_cudagraphs=True, |
| 975 | max_decode_rows=4, |
| 976 | ) |
| 977 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 978 | builder, |
| 979 | block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), |
| 980 | seq_lens_device=seq_lens, |
| 981 | max_seq_len=int(seq_lens.max().item()), |
| 982 | query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), |
| 983 | query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32), |
| 984 | num_decode_tokens=seq_lens.numel() * mtp_qlen, |
| 985 | max_query_len=mtp_qlen, |
| 986 | dcp_tot_seq_lens_device=None, |
| 987 | ) |
| 988 | |
| 989 | assert metadata.max_qo_len == mtp_qlen |
| 990 | assert torch.equal(metadata.seq_lens, expected_seq_lens) |
| 991 | assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) |
| 992 | assert torch.equal(metadata.qo_indptr, expected_qo_indptr) |
| 993 | assert torch.all( |
| 994 | builder.paged_kv_indptr[expected_paged_kv_indptr.numel() :] |
| 995 | == expected_paged_kv_indptr[-1] |
| 996 | ) |
| 997 | assert torch.all( |
| 998 | builder.qo_indptr[expected_qo_indptr.numel() :] == expected_qo_indptr[-1] |
| 999 | ) |
| 1000 | assert metadata.has_persistent_metadata |
| 1001 | assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == mtp_qlen |
| 1002 | assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == mtp_qlen |
| 1003 | |
| 1004 | |
| 1005 | def test_decode_expands_kernel_block_page_indices(monkeypatch): |
| 1006 | """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704.""" |
| 1007 | expand_kernel = _ExpandPageIndicesKernel() |
| 1008 | monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel) |
| 1009 | # qlen==1 now takes the persistent-metadata path, which imports |
| 1010 | # get_mla_metadata_v1 from aiter; mock it so this test needs no real kernel. |
| 1011 | monkeypatch.setitem( |
| 1012 | sys.modules, |
| 1013 | "aiter", |
| 1014 | SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()), |
| 1015 | ) |
| 1016 | |
| 1017 | kernel_block_size = 2 |
| 1018 | block_table = torch.tensor( |
| 1019 | [ |
| 1020 | [10, 11, 99], |
| 1021 | [20, 21, 22], |
| 1022 | ], |
| 1023 | dtype=torch.int32, |
| 1024 | ) |
| 1025 | seq_lens = torch.tensor([3, 5], dtype=torch.int32) |
| 1026 | expected_paged_kv_indptr = torch.cat( |
| 1027 | [ |
| 1028 | torch.zeros(1, dtype=torch.int32), |
| 1029 | seq_lens.cumsum(dim=0, dtype=torch.int32), |
| 1030 | ] |
| 1031 | ) |
| 1032 | expected_indices = torch.tensor( |
| 1033 | [ |
| 1034 | int(block_table[req_idx, token_idx // kernel_block_size].item()) |
| 1035 | * kernel_block_size |
| 1036 | + token_idx % kernel_block_size |
| 1037 | for req_idx, seq_len in enumerate(seq_lens.tolist()) |
| 1038 | for token_idx in range(seq_len) |
| 1039 | ], |
| 1040 | dtype=torch.int32, |
| 1041 | ) |
| 1042 | |
| 1043 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 1044 | _builder( |
| 1045 | mtp_decode_qlen=1, |
| 1046 | kernel_block_size=kernel_block_size, |
| 1047 | ), |
| 1048 | block_table_tensor=block_table, |
| 1049 | seq_lens_device=seq_lens, |
| 1050 | max_seq_len=int(seq_lens.max().item()), |
| 1051 | query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32), |
| 1052 | query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32), |
| 1053 | num_decode_tokens=seq_lens.numel(), |
| 1054 | max_query_len=1, |
| 1055 | dcp_tot_seq_lens_device=None, |
| 1056 | ) |
| 1057 | |
| 1058 | assert metadata.max_qo_len == 1 |
| 1059 | assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr) |
| 1060 | assert torch.equal( |
| 1061 | metadata.paged_kv_indices[: expected_indices.numel()], |
| 1062 | expected_indices, |
| 1063 | ) |
| 1064 | # The launch now carries a second grid dimension over token chunks, so |
| 1065 | # the expansion is spread across the context instead of one workgroup |
| 1066 | # walking it serially. The width comes from the block-table upper bound |
| 1067 | # on tokens per request, which is available host-side and so costs no |
| 1068 | # device sync; derived here rather than hard-coded so the assertion |
| 1069 | # still describes the launch if the chunk size changes. |
| 1070 | expected_chunks = max(1, -(-block_table.shape[1] * kernel_block_size // 1024)) |
| 1071 | assert expand_kernel.grid == (seq_lens.numel(), expected_chunks) |
| 1072 | assert expand_kernel.kernel_block_size == kernel_block_size |
| 1073 | |
| 1074 | |
| 1075 | @pytest.mark.parametrize( |
| 1076 | "mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent", |
| 1077 | [ |
| 1078 | (1, 1, 16, "auto", True), # non-MTP decode |
| 1079 | (4, 2, 16, "auto", True), # MTP deployment, in-range step |
| 1080 | (4, 4, 16, "auto", True), # MTP deployment, full-qlen verification step |
| 1081 | (1, 1, 24, "auto", True), # unaligned H24 pads to H32 persistent decode |
| 1082 | (2, 4, 16, "auto", False), # step demand exceeds provisioned K -> fallback |
| 1083 | (1, 1, 8, "auto", False), # divisor head count -> Gluon decode owns qlen==1 |
| 1084 | (4, 4, 8, "auto", False), # small head count -> Gluon flatten owns qlen>1 |
| 1085 | # A non-divisor head count is padded to 16 and runs the asm decode, so it |
| 1086 | # needs the schedule even though num_heads < 16. Its verify still |
| 1087 | # flattens onto Gluon though -- divisibility only steers the qlen==1 |
| 1088 | # decode, and bf16 has no small-head multi-token asm kernel either way. |
| 1089 | (1, 1, 12, "auto", True), |
| 1090 | (8, 8, 12, "auto", False), |
| 1091 | # An fp8 cache never reaches Gluon, at any head count or qlen, so it |
| 1092 | # always needs the schedule. (1, 1, 8) is the case the head-count gate |
| 1093 | # got wrong: divisor head count, so it read as Gluon-owned. |
| 1094 | (1, 1, 8, "fp8", True), |
| 1095 | (1, 1, 12, "fp8", True), |
| 1096 | (1, 1, 16, "fp8", True), |
| 1097 | # DSpark verify shape. The fp8 fold path rejects non-persistent outright |
| 1098 | # once qlen > 4, so these must not fall through. |
| 1099 | (8, 8, 8, "fp8", True), |
| 1100 | (8, 8, 12, "fp8", True), |
| 1101 | ], |
| 1102 | ) |
| 1103 | def test_persistent_metadata_gate( |
| 1104 | monkeypatch, mtp_decode_qlen, qo_len, num_heads, kv_cache_dtype, expect_persistent |
| 1105 | ): |
| 1106 | """Persistent metadata is passed iff the asm decode runs and 1 <= qlen <= K. |
| 1107 | |
| 1108 | K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets |
| 1109 | the pre-built schedule only when its qlen fits those buffers, otherwise it |
| 1110 | falls back to the kernel computing its own. qlen==1 (non-MTP) must stay |
| 1111 | in-range -- dropping it is the regression this guards. This includes |
| 1112 | unaligned H24, whose padded H32 decode uses persistent metadata. |
| 1113 | |
| 1114 | Only the Gluon paths ignore the schedule, so the gate follows the routing |
| 1115 | predicates rather than the raw head count. Reading `num_heads >= 16` instead |
| 1116 | denies the schedule to two sets of shapes that do run the asm kernels: a |
| 1117 | non-divisor head count padded up to 16, and any fp8 cache below 16 heads. |
| 1118 | """ |
| 1119 | get_mla_metadata_v1 = mock.MagicMock() |
| 1120 | monkeypatch.setitem( |
| 1121 | sys.modules, |
| 1122 | "aiter", |
| 1123 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 1124 | ) |
| 1125 | monkeypatch.setattr( |
| 1126 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 1127 | ) |
| 1128 | # The gate follows the routing now, so the small-head bf16 rows above only |
| 1129 | # hold where a Gluon build exists. Pin the arch they describe: on gfx942 |
| 1130 | # those shapes run the asm decode and do get the schedule, which is a |
| 1131 | # different assertion, made arch-free in the fp8 routing test. |
| 1132 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True) |
| 1133 | monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") |
| 1134 | |
| 1135 | # Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so |
| 1136 | # num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in. |
| 1137 | num_reqs = 2 |
| 1138 | query_start_loc = torch.arange( |
| 1139 | 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 |
| 1140 | ) |
| 1141 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 1142 | _builder( |
| 1143 | mtp_decode_qlen=mtp_decode_qlen, |
| 1144 | num_heads=num_heads, |
| 1145 | kv_cache_dtype=kv_cache_dtype, |
| 1146 | ), |
| 1147 | block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), |
| 1148 | seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), |
| 1149 | max_seq_len=8, |
| 1150 | query_start_loc_cpu=query_start_loc, |
| 1151 | query_start_loc_device=query_start_loc, |
| 1152 | num_decode_tokens=num_reqs * qo_len, |
| 1153 | max_query_len=qo_len, |
| 1154 | dcp_tot_seq_lens_device=None, |
| 1155 | ) |
| 1156 | |
| 1157 | assert metadata.max_qo_len == qo_len |
| 1158 | assert metadata.has_persistent_metadata is expect_persistent |
| 1159 | assert get_mla_metadata_v1.called is expect_persistent |
| 1160 | if expect_persistent: |
| 1161 | assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len |
| 1162 | assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len |
| 1163 | |
| 1164 | |
| 1165 | @pytest.mark.parametrize( |
| 1166 | "qo_len, num_heads, kv_cache_dtype, expect_persistent", |
| 1167 | [ |
| 1168 | # A padded rank keeps the schedule wherever a persistent kernel exists: |
| 1169 | # bf16 up to qseqlen 4, fp8 at every qlen (its fold requires it). |
| 1170 | (1, 12, "auto", True), |
| 1171 | (4, 12, "auto", True), |
| 1172 | (1, 12, "fp8", True), |
| 1173 | (8, 12, "fp8", True), |
| 1174 | # Past qseqlen 4 bf16 has only the non-persistent entry, and the fold |
| 1175 | # that would reach a persistent one is gfx950-only, so the padded rank |
| 1176 | # must fall back rather than ask for a kernel that is not built. |
| 1177 | (8, 12, "auto", False), |
| 1178 | # A native 16-head rank is unchanged: it took the schedule before this |
| 1179 | # backend keyed the gate off the routing, and still does. |
| 1180 | (8, 16, "auto", True), |
| 1181 | ], |
| 1182 | ) |
| 1183 | def test_persistent_metadata_gate_without_gluon_build( |
| 1184 | monkeypatch, qo_len, num_heads, kv_cache_dtype, expect_persistent |
| 1185 | ): |
| 1186 | """The gate on an arch with no Gluon build, i.e. gfx942. |
| 1187 | |
| 1188 | Both Gluon predicates read False there, so a small-head shape reaches the |
| 1189 | padded asm decode instead of being flattened. The schedule then has to |
| 1190 | follow what the asm dispatch actually ships: bf16 has no gqa=16 persistent |
| 1191 | kernel above qseqlen 4, and the q-row fold that reaches one on gfx950 is |
| 1192 | guarded on that arch, so asking for the schedule there selects a kernel |
| 1193 | that does not exist. |
| 1194 | """ |
| 1195 | get_mla_metadata_v1 = mock.MagicMock() |
| 1196 | monkeypatch.setitem( |
| 1197 | sys.modules, |
| 1198 | "aiter", |
| 1199 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 1200 | ) |
| 1201 | monkeypatch.setattr( |
| 1202 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 1203 | ) |
| 1204 | monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False) |
| 1205 | monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto") |
| 1206 | |
| 1207 | num_reqs = 2 |
| 1208 | query_start_loc = torch.arange( |
| 1209 | 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32 |
| 1210 | ) |
| 1211 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 1212 | _builder( |
| 1213 | mtp_decode_qlen=qo_len, |
| 1214 | num_heads=num_heads, |
| 1215 | kv_cache_dtype=kv_cache_dtype, |
| 1216 | ), |
| 1217 | block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), |
| 1218 | seq_lens_device=torch.tensor([8, 8], dtype=torch.int32), |
| 1219 | max_seq_len=8, |
| 1220 | query_start_loc_cpu=query_start_loc, |
| 1221 | query_start_loc_device=query_start_loc, |
| 1222 | num_decode_tokens=num_reqs * qo_len, |
| 1223 | max_query_len=qo_len, |
| 1224 | dcp_tot_seq_lens_device=None, |
| 1225 | ) |
| 1226 | |
| 1227 | assert metadata.has_persistent_metadata is expect_persistent |
| 1228 | assert get_mla_metadata_v1.called is expect_persistent |
| 1229 | |
| 1230 | |
| 1231 | def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen): |
| 1232 | """Drive _build_decode with causal=False and hand back the aiter mock.""" |
| 1233 | get_mla_metadata_v1 = mock.MagicMock() |
| 1234 | monkeypatch.setitem( |
| 1235 | sys.modules, |
| 1236 | "aiter", |
| 1237 | SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), |
| 1238 | ) |
| 1239 | monkeypatch.setattr( |
| 1240 | rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() |
| 1241 | ) |
| 1242 | builder = _builder( |
| 1243 | mtp_decode_qlen=mtp_qlen, |
| 1244 | kv_cache_dtype=kv_cache_dtype, |
| 1245 | num_heads=num_heads, |
| 1246 | ) |
| 1247 | builder._decode_causal = False |
| 1248 | metadata = AiterMLAMetadataBuilder._build_decode( |
| 1249 | builder, |
| 1250 | block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), |
| 1251 | seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), |
| 1252 | max_seq_len=7, |
| 1253 | query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), |
| 1254 | query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), |
| 1255 | num_decode_tokens=2 * qlen, |
| 1256 | max_query_len=qlen, |
| 1257 | dcp_tot_seq_lens_device=None, |
| 1258 | ) |
| 1259 | return metadata, get_mla_metadata_v1 |
| 1260 | |
| 1261 | |
| 1262 | def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch): |
| 1263 | """The schedule carries the mask, so is_causal must follow the block.""" |
| 1264 | metadata, get_mla_metadata_v1 = _build_non_causal( |
| 1265 | monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8 |
| 1266 | ) |
| 1267 | assert metadata.has_persistent_metadata |
| 1268 | # 6th positional arg of get_mla_metadata_v1 is is_causal. |
| 1269 | assert get_mla_metadata_v1.call_args.args[5] is False |
| 1270 | |
| 1271 | |
| 1272 | @pytest.mark.parametrize("num_heads", [12, 16, 48]) |
| 1273 | def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads): |
| 1274 | """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel. |
| 1275 | |
| 1276 | 48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they |
| 1277 | must be refused the same way as native 16-head qlen 2. |
| 1278 | """ |
| 1279 | with pytest.raises(ValueError, match="2-token"): |
| 1280 | _build_non_causal( |
| 1281 | monkeypatch, |
| 1282 | num_heads=num_heads, |
| 1283 | kv_cache_dtype="fp8", |
| 1284 | qlen=2, |
| 1285 | mtp_qlen=8, |
| 1286 | ) |
| 1287 | |
| 1288 | |
| 1289 | def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch): |
| 1290 | """32 heads at qlen 2 folds onto the 16-head / 4-token kernel.""" |
| 1291 | metadata, get_mla_metadata_v1 = _build_non_causal( |
| 1292 | monkeypatch, |
| 1293 | num_heads=32, |
| 1294 | kv_cache_dtype="fp8", |
| 1295 | qlen=2, |
| 1296 | mtp_qlen=8, |
| 1297 | ) |
| 1298 | assert metadata.has_persistent_metadata |
| 1299 | assert get_mla_metadata_v1.called |
| 1300 | |
| 1301 | |
| 1302 | def test_a_two_token_bf16_block_is_allowed(monkeypatch): |
| 1303 | """The bf16 fold has no such hole, so the guard must not fire there.""" |
| 1304 | metadata, _ = _build_non_causal( |
| 1305 | monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 |
| 1306 | ) |
| 1307 | assert metadata.has_persistent_metadata |