| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Which next_n the DSA indexer decode path may hand to DeepGEMM unflattened. |
| 4 | |
| 5 | Getting this wrong is not a slow path but a crash: `fp8_fp4_paged_mqa_logits` |
| 6 | asserts both that the architecture implements the requested `next_n` and that |
| 7 | the schedule metadata was sized for the matching slot count. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.deep_gemm import _paged_mqa_logits_schedule_slots |
| 14 | from vllm.v1.attention.backends.mla import indexer |
| 15 | |
| 16 | NUM_SMS = 114 # H100 PCIe |
| 17 | |
| 18 | |
| 19 | def _set_arch(monkeypatch, family: int, *, cuda: bool = True, deep_gemm: bool = True): |
| 20 | monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda) |
| 21 | monkeypatch.setattr( |
| 22 | current_platform, |
| 23 | "is_device_capability_family", |
| 24 | lambda capability, device_id=0: capability // 10 == family, |
| 25 | ) |
| 26 | monkeypatch.setattr(indexer, "has_deep_gemm", lambda: deep_gemm) |
| 27 | |
| 28 | |
| 29 | @pytest.mark.parametrize( |
| 30 | "family,expected_native", |
| 31 | [ |
| 32 | # SM90 gained next_n=4 (MTP=3) via 2-CTA multicast, but never 3. |
| 33 | (9, {1, 2, 4}), |
| 34 | # SM100 schedules any next_n with multi-atom tiles. |
| 35 | (10, {1, 2, 3, 4, 5, 8}), |
| 36 | # SM120 advertises multi-atom too but is unvalidated on hardware, so |
| 37 | # it stays on the conservative gate. Loosen it only with measurements. |
| 38 | (12, {1, 2}), |
| 39 | ], |
| 40 | ) |
| 41 | def test_native_decode_gate_per_architecture(monkeypatch, family, expected_native): |
| 42 | _set_arch(monkeypatch, family) |
| 43 | for next_n in (1, 2, 3, 4, 5, 8): |
| 44 | assert indexer._supports_native_decode(next_n) == (next_n in expected_native), ( |
| 45 | f"family={family} next_n={next_n}" |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | @pytest.mark.parametrize( |
| 50 | "cuda,deep_gemm", [(False, True), (True, False), (False, False)] |
| 51 | ) |
| 52 | def test_native_decode_gate_without_deepgemm(monkeypatch, cuda, deep_gemm): |
| 53 | """Without the DeepGEMM kernels only the shapes every backend handles.""" |
| 54 | _set_arch(monkeypatch, 9, cuda=cuda, deep_gemm=deep_gemm) |
| 55 | assert [indexer._supports_native_decode(n) for n in (1, 2, 3, 4)] == [ |
| 56 | True, |
| 57 | True, |
| 58 | False, |
| 59 | False, |
| 60 | ] |
| 61 | |
| 62 | |
| 63 | def test_sm90_next_n_4_halves_the_schedule_slots(monkeypatch): |
| 64 | """SM90 next_n=4 runs one scheduler task per 2-CTA cluster, not per SM.""" |
| 65 | _set_arch(monkeypatch, 9) |
| 66 | assert _paged_mqa_logits_schedule_slots(NUM_SMS, 4) == NUM_SMS // 2 |
| 67 | for next_n in (1, 2, 3): |
| 68 | assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS |
| 69 | |
| 70 | |
| 71 | @pytest.mark.parametrize("family", [10, 12]) |
| 72 | def test_multicast_is_sm90_only(monkeypatch, family): |
| 73 | _set_arch(monkeypatch, family) |
| 74 | for next_n in (1, 2, 3, 4): |
| 75 | assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS |