@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 152 lines 5.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Routing/selection tests for the AITER ASM cprr DCP-verify route.
4
5 Companion to the upstream segmented-verify tests. Everything here is pure CPU:
6 it pins the *decisions* (which route a KV group takes, which head count the
7 kernel runs at, which configurations are refused) rather than kernel numerics.
8 The numerics parity test against the segmented route needs a GPU and both
9 routes built; it is marked and skipped when unavailable.
10 """
11
12 from __future__ import annotations
13
14 import pytest
15
16 rocm_aiter_mla = pytest.importorskip(
17 "vllm.v1.attention.backends.mla.rocm_aiter_mla",
18 reason="ROCm AITER MLA backend not available",
19 )
20
21 _heads = rocm_aiter_mla._asm_dcp_verify_heads
22 _configured = rocm_aiter_mla._asm_dcp_verify_configured
23 _select_route = rocm_aiter_mla._select_dcp_decode_route
24 Route = rocm_aiter_mla._DCPDecodeRoute
25 NATIVE = rocm_aiter_mla._NATIVE_CPRR_HEADS
26 MIN_QLEN = rocm_aiter_mla._MIN_CPRR_QLEN
27
28
29 def test_dcp_verify_env_defaults_to_segmented(monkeypatch):
30 monkeypatch.delenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", raising=False)
31 assert rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY == "segmented"
32
33
34 def test_dcp_verify_env_rejects_unknown_route(monkeypatch):
35 monkeypatch.setenv("VLLM_ROCM_AITER_MLA_DCP_VERIFY", "unknown")
36 with pytest.raises(ValueError, match="Invalid value"):
37 _ = rocm_aiter_mla.envs.VLLM_ROCM_AITER_MLA_DCP_VERIFY
38
39
40 # --------------------------------------------------------------------------
41 # _asm_dcp_verify_heads -- which head count the cprr kernel runs at
42 # --------------------------------------------------------------------------
43
44
45 @pytest.mark.parametrize("n", NATIVE)
46 def test_native_head_counts_are_used_as_is(n):
47 assert _heads(n) == n
48
49
50 def test_non_native_head_count_pads_up_to_the_next_native():
51 """K3 at TP8/DCP8 gathers 96 heads for the target; the kernel has no 96
52 variant, so it must run at 128, not silently fall out of the asm path."""
53 assert _heads(96) == 128
54 assert _heads(1) == 16
55 assert _heads(17) == 32
56 assert _heads(65) == 128
57
58
59 def test_head_count_above_the_largest_native_is_unservable():
60 """0 means 'cannot serve'; anything else would route a shape the kernel
61 has no build for."""
62 assert _heads(max(NATIVE) + 1) == 0
63 assert _heads(256) == 0
64
65
66 # --------------------------------------------------------------------------
67 # _asm_dcp_verify_configured -- reachability
68 # --------------------------------------------------------------------------
69
70
71 @pytest.fixture
72 def on_gfx950(monkeypatch):
73 """The cprr kernels are built only for gfx950, so the gate checks the arch.
74 Patched rather than detected, so both branches run on any runner."""
75 import vllm.platforms.rocm as rocm
76
77 monkeypatch.setattr(rocm, "on_gfx950", lambda: True)
78
79
80 def test_route_needs_dcp(on_gfx950):
81 assert _configured(dcp_world_size=1, cp_interleave=1, multi_token_decode=True) is (
82 False
83 )
84 assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=True) is (
85 True
86 )
87
88
89 def test_round_robin_interleave_other_than_one_is_excluded(on_gfx950):
90 """The kernel's global-position causal window assumes interleave 1, the
91 same restriction the segmented gate carries."""
92 assert _configured(dcp_world_size=8, cp_interleave=2, multi_token_decode=True) is (
93 False
94 )
95 assert _configured(dcp_world_size=8, cp_interleave=16, multi_token_decode=True) is (
96 False
97 )
98
99
100 def test_route_needs_multi_token_decode(on_gfx950):
101 """Single-token decode is already served. Enabling the route for it would
102 pad the head count on every step and buy nothing."""
103 assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=False) is (
104 False
105 )
106
107
108 def test_route_needs_gfx950(monkeypatch):
109 """hsa/gfx942/mla and hsa/gfx1250/mla carry no cprr rows, so off gfx950 the
110 kernel lookup fails at the first verify step rather than falling back."""
111 import vllm.platforms.rocm as rocm
112
113 monkeypatch.setattr(rocm, "on_gfx950", lambda: False)
114 assert _configured(dcp_world_size=8, cp_interleave=1, multi_token_decode=True) is (
115 False
116 )
117
118
119 # --------------------------------------------------------------------------
120 # qlen floor
121 # --------------------------------------------------------------------------
122
123
124 def test_min_cprr_qlen_is_above_two():
125 """Qlen 2 is below CPRR's floor and must use segmented verification."""
126 assert MIN_QLEN > 2
127
128
129 @pytest.mark.parametrize(
130 "supports,causal,qlen,asm,expected",
131 [
132 (True, True, 1, True, Route.PLAIN),
133 (True, True, 2, True, Route.SEGMENTED),
134 (True, True, 3, True, Route.CPRR),
135 (True, True, 4, True, Route.CPRR),
136 (True, True, 4, False, Route.SEGMENTED),
137 (True, False, 2, True, Route.PLAIN),
138 (True, False, 4, False, Route.PLAIN),
139 (False, True, 3, True, Route.CPRR),
140 ],
141 )
142 def test_dcp_decode_route(supports, causal, qlen, asm, expected):
143 assert _select_route(supports, causal, qlen, asm) is expected
144
145
146 def test_causal_multi_token_batch_without_a_valid_route_fails():
147 with pytest.raises(RuntimeError, match="requires either segmented MLA"):
148 _select_route(False, True, 2, True)
149
150
151 # Kernel numerics live in test_rocm_aiter_mla_dcp_cprr_numerics.py, which runs
152 # the real builder and the real asm kernel on a gfx950 GPU.