| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from functools import partial |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.v1.attention.test_attention_backends import ( |
| 11 | BATCH_SPECS, |
| 12 | _test_backend_correctness, |
| 13 | ) |
| 14 | from tests.v1.attention.utils import BatchSpec |
| 15 | from vllm.config import ModelConfig |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.platforms.interface import DeviceCapability |
| 18 | from vllm.utils.b12x import get_b12x_paged_attention |
| 19 | from vllm.v1.attention.backends import b12x |
| 20 | from vllm.v1.attention.backends.b12x import ( |
| 21 | B12xPagedAttentionBackend, |
| 22 | B12xPagedAttentionImpl, |
| 23 | _kv_page_size, |
| 24 | _max_page_table_width, |
| 25 | ) |
| 26 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 27 | from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout |
| 28 | |
| 29 | |
| 30 | def _require_b12x_paged_attention() -> None: |
| 31 | capability = current_platform.get_device_capability() |
| 32 | if ( |
| 33 | not current_platform.is_cuda() |
| 34 | or capability is None |
| 35 | or not B12xPagedAttentionBackend.supports_compute_capability(capability) |
| 36 | ): |
| 37 | pytest.skip("b12x paged attention requires SM120 or SM121.") |
| 38 | |
| 39 | paged_attention = get_b12x_paged_attention() |
| 40 | if paged_attention is None or not paged_attention.is_supported(): |
| 41 | pytest.skip("b12x paged attention is not available.") |
| 42 | |
| 43 | |
| 44 | def _causal_mask( |
| 45 | b: torch.Tensor, |
| 46 | h: torch.Tensor, |
| 47 | q_idx: torch.Tensor, |
| 48 | kv_idx: torch.Tensor, |
| 49 | *, |
| 50 | context_len: int, |
| 51 | ): |
| 52 | return q_idx + context_len >= kv_idx |
| 53 | |
| 54 | |
| 55 | def _causal_sliding_window_mask( |
| 56 | b: torch.Tensor, |
| 57 | h: torch.Tensor, |
| 58 | q_idx: torch.Tensor, |
| 59 | kv_idx: torch.Tensor, |
| 60 | *, |
| 61 | context_len: int, |
| 62 | sliding_window: int, |
| 63 | ): |
| 64 | causal_mask = q_idx + context_len >= kv_idx |
| 65 | window_mask = q_idx + context_len - kv_idx < sliding_window |
| 66 | return causal_mask & window_mask |
| 67 | |
| 68 | |
| 69 | @pytest.mark.parametrize( |
| 70 | ("dtype", "kv_cache_dtype", "paged_attention", "expected_reason"), |
| 71 | [ |
| 72 | pytest.param( |
| 73 | torch.float16, |
| 74 | "fp8_e4m3", |
| 75 | None, |
| 76 | "b12x currently requires bfloat16 queries", |
| 77 | id="query-dtype", |
| 78 | ), |
| 79 | pytest.param( |
| 80 | torch.bfloat16, |
| 81 | "float16", |
| 82 | None, |
| 83 | "b12x does not support float16 KV cache", |
| 84 | id="kv-cache-dtype", |
| 85 | ), |
| 86 | pytest.param( |
| 87 | torch.bfloat16, |
| 88 | "auto", |
| 89 | None, |
| 90 | "Install the b12x backend with `pip install vllm[b12x]`", |
| 91 | id="package-not-installed", |
| 92 | ), |
| 93 | pytest.param( |
| 94 | torch.bfloat16, |
| 95 | "auto", |
| 96 | SimpleNamespace(is_supported=lambda: False), |
| 97 | "b12x paged attention is not supported on the current device", |
| 98 | id="device-api", |
| 99 | ), |
| 100 | pytest.param( |
| 101 | torch.bfloat16, |
| 102 | "auto", |
| 103 | SimpleNamespace(is_supported=lambda: True), |
| 104 | None, |
| 105 | id="supported", |
| 106 | ), |
| 107 | ], |
| 108 | ) |
| 109 | def test_b12x_attention_config_support( |
| 110 | monkeypatch: pytest.MonkeyPatch, |
| 111 | dtype: torch.dtype, |
| 112 | kv_cache_dtype: str, |
| 113 | paged_attention, |
| 114 | expected_reason: str | None, |
| 115 | ) -> None: |
| 116 | monkeypatch.setattr( |
| 117 | b12x, |
| 118 | "get_b12x_paged_attention", |
| 119 | lambda: paged_attention, |
| 120 | ) |
| 121 | |
| 122 | reason = B12xPagedAttentionBackend.supports_combination( |
| 123 | head_size=128, |
| 124 | dtype=dtype, |
| 125 | kv_cache_dtype=kv_cache_dtype, |
| 126 | block_size=128, |
| 127 | use_mla=False, |
| 128 | has_sink=False, |
| 129 | use_sparse=False, |
| 130 | use_mm_prefix=False, |
| 131 | device_capability=DeviceCapability(12, 0), |
| 132 | ) |
| 133 | |
| 134 | assert reason == expected_reason |
| 135 | |
| 136 | |
| 137 | def test_b12x_attention_uses_two_plane_nhd_cache() -> None: |
| 138 | spec = B12xPagedAttentionBackend.customize_spec( |
| 139 | FullAttentionSpec( |
| 140 | block_size=128, |
| 141 | num_kv_heads=4, |
| 142 | head_size=128, |
| 143 | dtype=torch.bfloat16, |
| 144 | ) |
| 145 | ) |
| 146 | |
| 147 | assert spec.num_head_slots == 2 |
| 148 | assert spec.state_content_bytes == 4 * 128 * 2 |
| 149 | assert B12xPagedAttentionBackend.supported_kv_cache_layouts() == ( |
| 150 | KVCacheLayout.LBHNC, |
| 151 | KVCacheLayout.BLHNC, |
| 152 | ) |
| 153 | assert B12xPagedAttentionBackend.supports_block_size(128) |
| 154 | assert not B12xPagedAttentionBackend.supports_block_size(32) |
| 155 | |
| 156 | |
| 157 | def test_b12x_attention_hybrid_cache_capacity_includes_expansion() -> None: |
| 158 | assert _max_page_table_width(4096, 128, 4096, False) == 32 |
| 159 | assert _max_page_table_width(4096, 128, 4096, True) == 64 |
| 160 | |
| 161 | |
| 162 | def test_b12x_attention_runtime_page_size_comes_from_cache() -> None: |
| 163 | key_cache = torch.empty((3, 64, 4, 128), device="meta") |
| 164 | value_cache = torch.empty_like(key_cache) |
| 165 | |
| 166 | assert _kv_page_size(key_cache, value_cache) == 64 |
| 167 | with pytest.raises(ValueError, match="matching K/V page sizes"): |
| 168 | _kv_page_size(key_cache, torch.empty((3, 128, 4, 128), device="meta")) |
| 169 | |
| 170 | |
| 171 | def test_b12x_attention_lazily_prepares_decode_bucket(monkeypatch) -> None: |
| 172 | impl = object.__new__(B12xPagedAttentionImpl) |
| 173 | plan = SimpleNamespace(layout=SimpleNamespace(nbytes=96)) |
| 174 | created: list[tuple[int, int]] = [] |
| 175 | |
| 176 | def create_plan(page_size: int, batch_size: int) -> SimpleNamespace: |
| 177 | created.append((page_size, batch_size)) |
| 178 | return plan |
| 179 | |
| 180 | impl._decode_plans = {} |
| 181 | impl._create_decode_plan = create_plan |
| 182 | impl._scratch_nbytes = 128 |
| 183 | impl._extend_plans = {} |
| 184 | impl._verify_q_per_req = 0 |
| 185 | metadata = SimpleNamespace(max_query_len=1) |
| 186 | monkeypatch.setattr(b12x, "_capture_alloc_forbidden", lambda: False) |
| 187 | |
| 188 | assert impl._select_plan(metadata, 7, 7, 7, 64) is plan |
| 189 | assert impl._select_plan(metadata, 7, 7, 7, 64) is plan |
| 190 | assert created == [(64, 7)] |
| 191 | |
| 192 | |
| 193 | def test_b12x_attention_fp8_descales_follow_request_batch() -> None: |
| 194 | impl = object.__new__(B12xPagedAttentionImpl) |
| 195 | impl.kv_cache_dtype = "fp8_e4m3" |
| 196 | layer = SimpleNamespace( |
| 197 | _k_scale=torch.tensor(2.0), |
| 198 | _v_scale=torch.tensor([3.0, 4.0, 5.0]), |
| 199 | ) |
| 200 | |
| 201 | k_descale, v_descale = impl._prepare_fp8_descales( |
| 202 | layer, num_reqs=2, device=torch.device("cpu") |
| 203 | ) |
| 204 | |
| 205 | torch.testing.assert_close(k_descale, torch.tensor([2.0, 2.0])) |
| 206 | torch.testing.assert_close(v_descale, torch.tensor([3.0, 4.0])) |
| 207 | assert k_descale.stride() == (0,) |
| 208 | assert v_descale.stride() == (1,) |
| 209 | |
| 210 | |
| 211 | def test_b12x_attention_sinks_refresh_in_place_after_reload() -> None: |
| 212 | impl = object.__new__(B12xPagedAttentionImpl) |
| 213 | source = torch.tensor([1.0, 2.0], dtype=torch.bfloat16) |
| 214 | impl._sinks_source = source |
| 215 | impl.sinks = None |
| 216 | |
| 217 | impl.process_weights_after_loading(torch.bfloat16) |
| 218 | assert impl.sinks is not None |
| 219 | sinks_ptr = impl.sinks.data_ptr() |
| 220 | source.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16)) |
| 221 | impl.process_weights_after_loading(torch.bfloat16) |
| 222 | |
| 223 | assert impl.sinks.data_ptr() == sinks_ptr |
| 224 | torch.testing.assert_close(impl.sinks, source.float()) |
| 225 | |
| 226 | |
| 227 | @pytest.mark.parametrize( |
| 228 | "batch_spec_name", |
| 229 | ["small_decode", "small_prefill", "mixed_small", "medium_decode"], |
| 230 | ) |
| 231 | @pytest.mark.parametrize( |
| 232 | ("kv_cache_dtype", "model_dtype"), |
| 233 | [ |
| 234 | ("auto", None), |
| 235 | ("bfloat16", torch.bfloat16), |
| 236 | ("fp8_e4m3", torch.bfloat16), |
| 237 | ], |
| 238 | ) |
| 239 | @pytest.mark.parametrize("block_size", [64, 128]) |
| 240 | def test_b12x_causal_backend_correctness( |
| 241 | default_vllm_config, |
| 242 | workspace_init, |
| 243 | batch_spec_name: str, |
| 244 | kv_cache_dtype: str, |
| 245 | model_dtype: torch.dtype | None, |
| 246 | block_size: int, |
| 247 | ) -> None: |
| 248 | """b12x causal paged attention matches the shared SDPA reference.""" |
| 249 | _require_b12x_paged_attention() |
| 250 | batch_spec = BATCH_SPECS[batch_spec_name] |
| 251 | |
| 252 | _test_backend_correctness( |
| 253 | batch_spec, |
| 254 | "Qwen/Qwen3-0.6B", |
| 255 | [AttentionBackendEnum.B12X], |
| 256 | _causal_mask, |
| 257 | block_size=block_size, |
| 258 | kv_cache_dtype=kv_cache_dtype, |
| 259 | model_dtype=model_dtype, |
| 260 | max_num_seqs=batch_spec.batch_size, |
| 261 | max_num_batched_tokens=max(sum(batch_spec.query_lens), 64), |
| 262 | ) |
| 263 | |
| 264 | |
| 265 | @pytest.mark.parametrize( |
| 266 | "batch_spec", |
| 267 | [ |
| 268 | pytest.param(BatchSpec(seq_lens=[2080, 2200], query_lens=[1, 1]), id="decode"), |
| 269 | pytest.param(BatchSpec(seq_lens=[2080, 2200], query_lens=[8, 8]), id="prefill"), |
| 270 | ], |
| 271 | ) |
| 272 | def test_b12x_causal_sliding_window( |
| 273 | default_vllm_config, |
| 274 | workspace_init, |
| 275 | batch_spec: BatchSpec, |
| 276 | ) -> None: |
| 277 | """b12x causal sliding-window attention matches the shared reference.""" |
| 278 | _require_b12x_paged_attention() |
| 279 | |
| 280 | model = "microsoft/Phi-tiny-MoE-instruct" |
| 281 | sliding_window = ModelConfig( |
| 282 | model=model, max_model_len=max(batch_spec.seq_lens) |
| 283 | ).get_sliding_window() |
| 284 | assert sliding_window is not None |
| 285 | mask = partial(_causal_sliding_window_mask, sliding_window=sliding_window) |
| 286 | |
| 287 | _test_backend_correctness( |
| 288 | batch_spec, |
| 289 | model, |
| 290 | [AttentionBackendEnum.B12X], |
| 291 | mask, |
| 292 | block_size=64, |
| 293 | atol=3e-2, |
| 294 | rtol=3e-2, |
| 295 | max_num_seqs=batch_spec.batch_size, |
| 296 | max_num_batched_tokens=max(sum(batch_spec.query_lens), 64), |
| 297 | ) |
| 298 | |
| 299 | |
| 300 | def test_b12x_attention_sinks( |
| 301 | default_vllm_config, |
| 302 | workspace_init, |
| 303 | ) -> None: |
| 304 | """b12x attention sinks match the explicit sink reference.""" |
| 305 | _require_b12x_paged_attention() |
| 306 | batch_spec = BATCH_SPECS["small_prefill"] |
| 307 | |
| 308 | _test_backend_correctness( |
| 309 | batch_spec, |
| 310 | "Qwen/Qwen3-0.6B", |
| 311 | [AttentionBackendEnum.B12X], |
| 312 | _causal_mask, |
| 313 | block_size=64, |
| 314 | atol=3e-2, |
| 315 | rtol=3e-2, |
| 316 | use_sinks=True, |
| 317 | max_num_seqs=batch_spec.batch_size, |
| 318 | max_num_batched_tokens=max(sum(batch_spec.query_lens), 64), |
| 319 | ) |
| 320 | |
| 321 | |
| 322 | def test_b12x_decode_cuda_graph_replay( |
| 323 | default_vllm_config, |
| 324 | workspace_init, |
| 325 | ) -> None: |
| 326 | """b12x decode output remains correct after CUDA graph replay.""" |
| 327 | _require_b12x_paged_attention() |
| 328 | batch_spec = BATCH_SPECS["small_decode"] |
| 329 | |
| 330 | _test_backend_correctness( |
| 331 | batch_spec, |
| 332 | "Qwen/Qwen3-0.6B", |
| 333 | [AttentionBackendEnum.B12X], |
| 334 | _causal_mask, |
| 335 | block_size=64, |
| 336 | use_cuda_graph=True, |
| 337 | max_num_seqs=batch_spec.batch_size, |
| 338 | max_num_batched_tokens=max(sum(batch_spec.query_lens), 64), |
| 339 | ) |
| 340 | |
| 341 | |
| 342 | def test_b12x_speculative_verification_uses_cuda_graph_plan( |
| 343 | default_vllm_config, |
| 344 | workspace_init, |
| 345 | ) -> None: |
| 346 | """Exercise the verifier plan rather than the general extend plan.""" |
| 347 | _require_b12x_paged_attention() |
| 348 | batch_spec = BatchSpec(seq_lens=[32, 40], query_lens=[4, 4]) |
| 349 | |
| 350 | _test_backend_correctness( |
| 351 | batch_spec, |
| 352 | "Qwen/Qwen3-0.6B", |
| 353 | [AttentionBackendEnum.B12X], |
| 354 | _causal_mask, |
| 355 | block_size=128, |
| 356 | num_speculative_tokens=3, |
| 357 | use_cuda_graph=True, |
| 358 | max_num_seqs=batch_spec.batch_size, |
| 359 | max_num_batched_tokens=max(sum(batch_spec.query_lens), 64), |
| 360 | ) |