@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 166 lines 5.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """CPU-only tests for FlashInfer sparse MLA backend constraints."""
4
5 from types import SimpleNamespace
6 from unittest.mock import patch
7
8 import pytest
9 import torch
10
11 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
12 _DEFAULT_WORKSPACE_BUFFER_SIZE,
13 FlashInferMLASparseTRTLLMBackend,
14 _required_workspace_bytes,
15 compute_trtllm_sparse_mla_workspace_bytes,
16 )
17
18 # Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8,
19 # max_num_batched_tokens=16384, FlashInfer 0.6.14.
20 _REPORTER_HEADS_PER_RANK = 64 // 8
21 _REPORTER_DCP = 8
22 _REPORTER_MAX_TOKENS = 16384
23 # From the reported crash: trtllm-gen requested exactly this many bytes for
24 # trtllm_gen_softmax_workspace after an 8 MiB counter carve had left
25 # 404,750,336 of the 413,138,944-byte default buffer.
26 _REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312
27 _REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608
28
29
30 def test_reporter_softmax_carve_is_reproduced_byte_exactly():
31 # The crashing step scheduled 12288 tokens (a single 12K-token request);
32 # the softmax slab formula must reproduce the kernel's request exactly:
33 # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard.
34 softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes(
35 base_workspace_bytes=0,
36 dcp_world_size=_REPORTER_DCP,
37 num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
38 max_num_batched_tokens=12288,
39 )
40 assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES
41
42
43 def test_reporter_config_covers_observed_overflow():
44 computed = compute_trtllm_sparse_mla_workspace_bytes(
45 base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
46 dcp_world_size=_REPORTER_DCP,
47 num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
48 max_num_batched_tokens=_REPORTER_MAX_TOKENS,
49 )
50 assert computed >= (
51 _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES
52 )
53
54
55 def test_non_dcp_size_is_unchanged():
56 computed = compute_trtllm_sparse_mla_workspace_bytes(
57 base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE,
58 dcp_world_size=1,
59 num_heads_per_rank=128,
60 max_num_batched_tokens=65536,
61 )
62 assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE
63
64
65 def test_combined_pcp_dcp_rejects_flashinfer_sparse():
66 config = SimpleNamespace(
67 parallel_config=SimpleNamespace(
68 prefill_context_parallel_size=4,
69 decode_context_parallel_size=4,
70 ),
71 model_config=None,
72 )
73 with patch("vllm.config.get_current_vllm_config", return_value=config):
74 reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
75 head_size=576,
76 dtype=torch.bfloat16,
77 kv_cache_dtype="fp8",
78 block_size=64,
79 use_mla=True,
80 has_sink=False,
81 use_sparse=True,
82 use_mm_prefix=False,
83 device_capability=SimpleNamespace(major=10),
84 )
85
86 assert reason is not None
87 assert "use FLASHMLA_SPARSE" in reason
88
89
90 @pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)])
91 def test_single_context_parallel_mode_allows_flashinfer_sparse(
92 pcp_size: int, dcp_size: int
93 ):
94 config = SimpleNamespace(
95 parallel_config=SimpleNamespace(
96 prefill_context_parallel_size=pcp_size,
97 decode_context_parallel_size=dcp_size,
98 ),
99 model_config=None,
100 )
101 with patch("vllm.config.get_current_vllm_config", return_value=config):
102 reason = FlashInferMLASparseTRTLLMBackend.supports_combination(
103 head_size=576,
104 dtype=torch.bfloat16,
105 kv_cache_dtype="fp8",
106 block_size=64,
107 use_mla=True,
108 has_sink=False,
109 use_sparse=True,
110 use_mm_prefix=False,
111 device_capability=SimpleNamespace(major=10),
112 )
113
114 assert reason is None
115
116
117 def test_default_constant_matches_envs_default(monkeypatch):
118 from vllm import envs
119
120 monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
121 assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
122
123
124 def test_env_unset_returns_computed(monkeypatch):
125 monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False)
126 required = _required_workspace_bytes(
127 dcp_world_size=_REPORTER_DCP,
128 num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
129 max_num_batched_tokens=_REPORTER_MAX_TOKENS,
130 )
131 assert required == compute_trtllm_sparse_mla_workspace_bytes(
132 _DEFAULT_WORKSPACE_BUFFER_SIZE,
133 _REPORTER_DCP,
134 _REPORTER_HEADS_PER_RANK,
135 _REPORTER_MAX_TOKENS,
136 )
137
138
139 def test_env_override_below_computed_is_respected_with_warning(monkeypatch):
140 override = 100 * 1024 * 1024
141 monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
142 with patch(
143 "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
144 ) as mock_logger:
145 required = _required_workspace_bytes(
146 dcp_world_size=_REPORTER_DCP,
147 num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
148 max_num_batched_tokens=_REPORTER_MAX_TOKENS,
149 )
150 assert required == override
151 mock_logger.warning_once.assert_called_once()
152
153
154 def test_env_override_above_computed_is_respected_without_warning(monkeypatch):
155 override = 8 * 1024 * 1024 * 1024
156 monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override))
157 with patch(
158 "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger"
159 ) as mock_logger:
160 required = _required_workspace_bytes(
161 dcp_world_size=_REPORTER_DCP,
162 num_heads_per_rank=_REPORTER_HEADS_PER_RANK,
163 max_num_batched_tokens=_REPORTER_MAX_TOKENS,
164 )
165 assert required == override
166 mock_logger.warning_once.assert_not_called()