@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0d4ad47981

[Kernel] Add B12X causal paged attention backend (#52017)

Signed-off-by: Luke Alonso <lalonso@gmail.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Luke Alonso committed Sep 1, 2026 at 07:42 UTC 0d4ad47981b6133c8f835d78f7a19ca4eda113d5

8 files changed +1527 -5

.buildkite/test_areas/kernels.yaml

+1 -1

@@ -389,7 +389,7 @@ steps:

389 - tests/quantization/reference_mxfp4.py

390 - tests/quantization/test_auto_round.py

391 commands:

392 - - uv pip install --system b12x==1.2.6

392 + - uv pip install --system b12x==1.3.0

393 - pytest -v -s model_executor/kernels/test_b12x_linear.py

394 model_executor/test_b12x_warmup.py

395 kernels/quantization/test_block_fp8.py

docs/design/attention_backends.md

+10

@@ -136,6 +136,16 @@ Priority is **1 = highest** (tried first).

136 > `--block-size` that is not a multiple of 128 instead makes FlashAttention ineligible

137 > for such models, which is an error if the backend was requested explicitly.

138

139 +### b12x

140 +

141 +The optional [b12x](https://pypi.org/project/b12x/) backend supports causal

142 +decoder attention on NVIDIA SM120 and SM121 GPUs. Install and select it with:

143 +

144 +```bash

145 +uv pip install "vllm[b12x]"

146 +vllm serve <model> --attention-backend b12x

147 +```

148 +

149 ## MiniMax M3 Sparse Attention Backends

150

151 Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")

setup.py

+1 -1

@@ -1522,7 +1522,7 @@ setup(

1522 # only; also needs system GStreamer + libv4l (see docs).

1523 "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"],

1524 "flashinfer": [], # Kept for backwards compatibility

1525 - "b12x": ["b12x==1.2.6"],

1525 + "b12x": ["b12x==1.3.0"],

1526 # Optional deps for Helion kernel development

1527 # NOTE: When updating helion version, also update CI files:

1528 # - .buildkite/test_areas/kernels.yaml

tests/v1/attention/test_attention_backends.py

+53 -3

@@ -3,6 +3,7 @@

3 """Tests for v1 attention backends without GPUModelRunner dependency."""

5 from functools import partial

6 +from types import SimpleNamespace

8 import pytest

9 import torch

@@ -282,6 +283,7 @@ def run_attention_backend(

283 sliding_window: int | None = None,

284 kv_cache_dtype: str = "auto",

285 sinks: torch.Tensor | None = None,

286 + use_cuda_graph: bool = False,

287 ) -> torch.Tensor:

288 """Run attention computation using the specified backend's AttentionImpl."""

289

@@ -367,9 +369,26 @@ def run_attention_backend(

369 impl.do_kv_cache_update(

370 mock_layer, key, value, kv_cache, attn_metadata.slot_mapping

371 )

370 - output = impl.forward(

371 - mock_layer, query, key, value, kv_cache, attn_metadata, output=output

372 - )

372 + if use_cuda_graph:

373 + impl.forward(

374 + mock_layer, query, key, value, kv_cache, attn_metadata, output=output

375 + )

376 + graph = torch.cuda.CUDAGraph()

377 + with torch.cuda.graph(graph):

378 + impl.forward(

379 + mock_layer,

380 + query,

381 + key,

382 + value,

383 + kv_cache,

384 + attn_metadata,

385 + output=output,

386 + )

387 + graph.replay()

388 + else:

389 + output = impl.forward(

390 + mock_layer, query, key, value, kv_cache, attn_metadata, output=output

391 + )

392

393 return output

394

@@ -389,6 +408,11 @@ def _test_backend_correctness(

408 kv_cache_dtype: str = "auto",

409 use_sinks: bool = False,

410 layout: KVCacheLayout | None = None,

411 + use_cuda_graph: bool = False,

412 + num_speculative_tokens: int = 0,

413 + model_dtype: torch.dtype | None = None,

414 + max_num_seqs: int | None = None,

415 + max_num_batched_tokens: int | None = None,

416 ):

417 """

418 Test that all backends produce similar outputs to a reference implementation

@@ -433,10 +457,19 @@ def _test_backend_correctness(

457 model_name=model,

458 tensor_parallel_size=1, # Always use TP=1 to avoid multi-GPU requirements

459 max_model_len=max(batch_spec.seq_lens),

460 + dtype=model_dtype or "auto",

461 block_size=block_size,

462 num_gpu_blocks=8192,

463 hf_config_override=hf_config_override,

464 )

465 + if max_num_seqs is not None:

466 + vllm_config.scheduler_config.max_num_seqs = max_num_seqs

467 + if max_num_batched_tokens is not None:

468 + vllm_config.scheduler_config.max_num_batched_tokens = max_num_batched_tokens

469 + if num_speculative_tokens > 0:

470 + vllm_config.speculative_config = SimpleNamespace(

471 + num_speculative_tokens=num_speculative_tokens

472 + )

473 vllm_config.cache_config.cache_dtype = kv_cache_dtype

474 device = torch.device(f"{DEVICE_TYPE}:0")

475

@@ -619,6 +652,22 @@ def _test_backend_correctness(

652 backend_layout = backend_supported[0]

653 kv_cache_for_backend = _clone_kv_cache_in_layout(kv_cache, backend_layout)

654

655 + if backend_name == AttentionBackendEnum.B12X:

656 + cache_dtype = (

657 + FP8_KV_CACHE_DTYPES[kv_cache_dtype]

658 + if is_quantized_kv_cache(kv_cache_dtype)

659 + else kv_cache.dtype

660 + )

661 + typed_cache = kv_cache.view(cache_dtype)

662 + key_cache = typed_cache[..., :head_size].permute(0, 2, 1, 3)

663 + value_cache = typed_cache[..., head_size:].permute(0, 2, 1, 3)

664 + packed_cache = torch.stack((key_cache, value_cache), dim=1).flatten(-2)

665 + if is_quantized_kv_cache(kv_cache_dtype):

666 + packed_cache = packed_cache.view(torch.uint8)

667 + kv_cache_for_backend = _clone_kv_cache_in_layout(

668 + packed_cache, backend_layout

669 + )

670 +

671 # FlashInfer reads the layout at plan time; set it to match

672 # the physical order of the test cache.

673 vllm_config.cache_config.kv_cache_layout = backend_layout.name

@@ -638,6 +687,7 @@ def _test_backend_correctness(

687 attn_type=attn_type,

688 kv_cache_dtype=kv_cache_dtype,

689 sinks=sinks,

690 + use_cuda_graph=use_cuda_graph,

691 )

692

693 # Check shape and dtype consistency

tests/v1/attention/test_b12x.py new

+360

@@ -0,0 +1,360 @@

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 + )

vllm/utils/b12x.py

@@ -34,6 +34,7 @@ def _import_submodule(module_name: str) -> ModuleType | None:

34 _B12X_SUBMODULES = {

35 module_name: _import_submodule(module_name)

36 for module_name in (

37 + "b12x.attention.paged",

38 "b12x.gemm.blockscaled",

39 # TODO: Remove once B12X exposes the scale-swizzle API publicly.

40 "b12x._lib.intrinsics",

@@ -73,6 +74,10 @@ def get_b12x_fused_moe() -> ModuleType | None:

74 return _get_submodule("b12x.moe.fused_moe")

77 +def get_b12x_paged_attention() -> ModuleType | None:

78 + return _get_submodule("b12x.attention.paged")

79 +

80 +

81 def b12x_warmup_token_counts(

82 *,

83 max_tokens: int,

vllm/v1/attention/backends/b12x.py new

+1096

@@ -0,0 +1,1096 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""b12x paged causal attention backend for SM12x."""

4 +

5 +from __future__ import annotations

6 +

7 +import copy

8 +import math

9 +from dataclasses import dataclass, replace

10 +from functools import partial

11 +from typing import Any, ClassVar

12 +

13 +import torch

14 +

15 +from vllm.config import VllmConfig, get_current_vllm_config

16 +from vllm.config.cache import CacheDType

17 +from vllm.logger import init_logger

18 +from vllm.model_executor.warmup.cutedsl_warmup import (

19 + CuTeDSLCompileUnit,

20 + register_cutedsl_warmup_provider,

21 +)

22 +from vllm.platforms import current_platform

23 +from vllm.platforms.interface import DeviceCapability

24 +from vllm.utils.b12x import get_b12x_paged_attention

25 +from vllm.utils.math_utils import cdiv

26 +from vllm.utils.torch_utils import (

27 + canonicalize_singleton_dim_strides,

28 + get_dtype_size,

29 + is_quantized_kv_cache,

30 +)

31 +from vllm.v1.attention.backend import (

32 + AttentionBackend,

33 + AttentionCGSupport,

34 + AttentionImpl,

35 + AttentionLayer,

36 + AttentionMetadata,

37 + AttentionMetadataBuilder,

38 + AttentionType,

39 + CommonAttentionMetadata,

40 + MultipleOf,

41 +)

42 +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout, KVCacheSpec

43 +from vllm.v1.worker.workspace import current_workspace_manager

44 +

45 +logger = init_logger(__name__)

46 +

47 +_B12X_SUPPORTED_PAGE_SIZES = (64, 128)

48 +_B12X_PREFERRED_PAGE_SIZE = 128

49 +_MIN_PAGED_TILE_Q = 16

50 +_B12X_FP8_KV_CACHE_DTYPES: tuple[CacheDType, ...] = ("fp8", "fp8_e4m3")

51 +_B12X_SUPPORTED_KV_CACHE_DTYPES: tuple[CacheDType, ...] = (

52 + "auto",

53 + "bfloat16",

54 + *_B12X_FP8_KV_CACHE_DTYPES,

55 +)

56 +

57 +

58 +def _max_page_table_width(

59 + max_model_len: int,

60 + block_size: int,

61 + max_num_batched_tokens: int,

62 + is_hybrid: bool,

63 +) -> int:

64 + width = max(cdiv(max(max_model_len, 1), block_size), 1)

65 + if is_hybrid:

66 + # Hybrid cache setup can enlarge the storage block after attention

67 + # layers are initialized. Its expansion into kernel-sized blocks adds

68 + # at most one storage block of trailing page-table capacity.

69 + width += cdiv(max_num_batched_tokens, block_size)

70 + return width

71 +

72 +

73 +def _kv_page_size(key_cache: torch.Tensor, value_cache: torch.Tensor) -> int:

74 + """Return the static kernel page geometry negotiated by vLLM.

75 +

76 + The KV manager can split the configured storage block into a smaller

77 + kernel page when another backend shares its cache group. Cache shapes are

78 + fixed before graph capture, so this is not a live-length policy decision.

79 + """

80 + if key_cache.ndim < 2 or value_cache.ndim < 2:

81 + raise ValueError(

82 + "b12x expects paged K/V caches with a page dimension, got "

83 + f"{tuple(key_cache.shape)} and {tuple(value_cache.shape)}."

84 + )

85 + key_page_size = int(key_cache.shape[1])

86 + value_page_size = int(value_cache.shape[1])

87 + if key_page_size != value_page_size:

88 + raise ValueError(

89 + "b12x requires matching K/V page sizes, got "

90 + f"{key_page_size} and {value_page_size}."

91 + )

92 + if key_page_size not in _B12X_SUPPORTED_PAGE_SIZES:

93 + raise ValueError(

94 + "b12x requires runtime page size in "

95 + f"{_B12X_SUPPORTED_PAGE_SIZES}, got {key_page_size}."

96 + )

97 + return key_page_size

98 +

99 +

100 +def _capture_alloc_forbidden() -> bool:

101 + if not torch.cuda.is_available():

102 + return False

103 + try:

104 + return bool(torch.cuda.is_current_stream_capturing())

105 + except RuntimeError:

106 + return False

107 +

108 +

109 +def _ensure_i32_contiguous(tensor: torch.Tensor, name: str) -> torch.Tensor:

110 + if tensor.dtype != torch.int32:

111 + if _capture_alloc_forbidden():

112 + raise RuntimeError(

113 + f"b12x would convert {name} to int32 during CUDA graph "

114 + "capture. Prepare int32 metadata before capture."

115 + )

116 + tensor = tensor.to(torch.int32)

117 + if not tensor.is_contiguous():

118 + if _capture_alloc_forbidden():

119 + raise RuntimeError(

120 + f"b12x would make {name} contiguous during CUDA graph "

121 + "capture. Prepare contiguous metadata before capture."

122 + )

123 + tensor = tensor.contiguous()

124 + return tensor

125 +

126 +

127 +def _dtype_from_cache_config(

128 + kv_cache_dtype: str,

129 + vllm_config: VllmConfig,

130 +) -> torch.dtype:

131 + if kv_cache_dtype == "bfloat16":

132 + return torch.bfloat16

133 + if kv_cache_dtype in _B12X_FP8_KV_CACHE_DTYPES:

134 + return current_platform.fp8_dtype()

135 + if kv_cache_dtype != "auto":

136 + raise NotImplementedError(

137 + "b12x currently supports only auto, bfloat16, "

138 + "fp8, and fp8_e4m3 "

139 + f"KV cache dtypes; got {kv_cache_dtype!r}."

140 + )

141 + return vllm_config.model_config.dtype

142 +

143 +

144 +def _is_b12x_fp8_kv_cache(kv_cache_dtype: str) -> bool:

145 + return kv_cache_dtype in _B12X_FP8_KV_CACHE_DTYPES

146 +

147 +

148 +class B12xPagedAttentionBackend(AttentionBackend):

149 + """b12x paged attention backend for regular/GQA decoder layers."""

150 +

151 + @classmethod

152 + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:

153 + if spec.state_content_bytes is not None:

154 + return spec

155 + assert spec.head_size == spec.head_size_v, (

156 + "Separate K/V planes require symmetric K/V head sizes."

157 + )

158 + return replace(

159 + spec,

160 + num_head_slots=2,

161 + state_content_bytes=spec.num_kv_heads

162 + * spec.head_size

163 + * get_dtype_size(spec.dtype),

164 + )

165 +

166 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

167 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = list(

168 + _B12X_SUPPORTED_KV_CACHE_DTYPES

169 + )

170 +

171 + forward_includes_kv_cache_update: bool = False

172 +

173 + @staticmethod

174 + def get_name() -> str:

175 + return "B12X"

176 +

177 + @classmethod

178 + def get_impl_cls(cls) -> type[B12xPagedAttentionImpl]:

179 + return B12xPagedAttentionImpl

180 +

181 + @staticmethod

182 + def get_builder_cls() -> type[B12xPagedMetadataBuilder]:

183 + return B12xPagedMetadataBuilder

184 +

185 + @staticmethod

186 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

187 + return list(_B12X_SUPPORTED_PAGE_SIZES)

188 +

189 + @classmethod

190 + def supports_block_size(cls, block_size: int | None) -> bool:

191 + return block_size is None or int(block_size) in _B12X_SUPPORTED_PAGE_SIZES

192 +

193 + @classmethod

194 + def get_preferred_block_size(cls, default_block_size: int) -> int:

195 + if int(default_block_size) in _B12X_SUPPORTED_PAGE_SIZES:

196 + return int(default_block_size)

197 + return _B12X_PREFERRED_PAGE_SIZE

198 +

199 + @classmethod

200 + def get_supported_head_sizes(cls) -> list[int]:

201 + return [64, 128, 192, 256]

202 +

203 + @classmethod

204 + def supports_sink(cls) -> bool:

205 + return True

206 +

207 + @classmethod

208 + def supports_sliding_window(cls) -> bool:

209 + return True

210 +

211 + @classmethod

212 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:

213 + # Consumer Blackwell SM120 / SM121. The b12x paged kernels also gate

214 + # internally, but keep vLLM selection fail-fast and explicit.

215 + return (capability.major, capability.minor) in ((12, 0), (12, 1))

216 +

217 + @classmethod

218 + def supports_combination(

219 + cls,

220 + head_size: int,

221 + dtype: torch.dtype,

222 + kv_cache_dtype: CacheDType | None,

223 + block_size: int | None,

224 + use_mla: bool,

225 + has_sink: bool,

226 + use_sparse: bool,

227 + use_mm_prefix: bool,

228 + device_capability: DeviceCapability,

229 + ) -> str | None:

230 + if dtype != torch.bfloat16:

231 + return "b12x currently requires bfloat16 queries"

232 + if kv_cache_dtype == "float16":

233 + return "b12x does not support float16 KV cache"

234 + if (

235 + kv_cache_dtype is not None

236 + and is_quantized_kv_cache(kv_cache_dtype)

237 + and not _is_b12x_fp8_kv_cache(kv_cache_dtype)

238 + ):

239 + return "b12x currently supports only fp8/fp8_e4m3 quantized KV cache dtypes"

240 + paged_attention = get_b12x_paged_attention()

241 + if paged_attention is None:

242 + return "Install the b12x backend with `pip install vllm[b12x]`"

243 + if not paged_attention.is_supported():

244 + return "b12x paged attention is not supported on the current device"

245 + return None

246 +

247 + @classmethod

248 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:

249 + return (KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)

250 +

251 +

252 +@dataclass

253 +class B12xPagedMetadata(AttentionMetadata):

254 + num_actual_tokens: int

255 + max_query_len: int

256 + query_start_loc: torch.Tensor

257 + max_seq_len: int

258 + seq_lens: torch.Tensor

259 + block_table: torch.Tensor

260 + slot_mapping: torch.Tensor

261 + causal: bool = True

262 +

263 +

264 +class B12xPagedMetadataBuilder(AttentionMetadataBuilder[B12xPagedMetadata]):

265 + """Metadata builder for b12x.

266 +

267 + Decode and uniform speculative-verifier batches use preplanned graph

268 + buckets. Extend/prefill remains eager and does not affect uniform decode

269 + graph eligibility.

270 + """

271 +

272 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

273 + supports_update_block_table: bool = True

274 +

275 + @classmethod

276 + def get_cudagraph_support(

277 + cls,

278 + vllm_config: VllmConfig,

279 + kv_cache_spec: KVCacheSpec,

280 + ) -> AttentionCGSupport:

281 + del vllm_config, kv_cache_spec

282 + return cls._cudagraph_support

283 +

284 + def __init__(

285 + self,

286 + kv_cache_spec: AttentionSpec,

287 + layer_names: list[str],

288 + vllm_config: VllmConfig,

289 + device: torch.device,

290 + ) -> None:

291 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

292 +

293 + def build(

294 + self,

295 + common_prefix_len: int,

296 + common_attn_metadata: CommonAttentionMetadata,

297 + fast_build: bool = False,

298 + ) -> B12xPagedMetadata:

299 + del common_prefix_len, fast_build

300 + cm = common_attn_metadata

301 + return B12xPagedMetadata(

302 + num_actual_tokens=cm.num_actual_tokens,

303 + max_query_len=cm.max_query_len,

304 + query_start_loc=cm.query_start_loc,

305 + max_seq_len=cm.max_seq_len,

306 + seq_lens=cm.seq_lens,

307 + block_table=cm.block_table_tensor,

308 + slot_mapping=cm.slot_mapping,

309 + causal=cm.causal,

310 + )

311 +

312 + def update_block_table(

313 + self,

314 + metadata: B12xPagedMetadata,

315 + blk_table: torch.Tensor,

316 + slot_mapping: torch.Tensor,

317 + ) -> B12xPagedMetadata:

318 + new_metadata = copy.copy(metadata)

319 + new_metadata.block_table = blk_table

320 + new_metadata.slot_mapping = slot_mapping

321 + return new_metadata

322 +

323 +

324 +class B12xPagedAttentionImpl(AttentionImpl[B12xPagedMetadata]):

325 + """b12x paged GQA attention implementation."""

326 +

327 + can_return_lse_for_decode: bool = False

328 + supports_dcp: bool = False

329 +

330 + def __init__(

331 + self,

332 + num_heads: int,

333 + head_size: int,

334 + scale: float,

335 + num_kv_heads: int,

336 + alibi_slopes: list[float] | None,

337 + sliding_window: int | None,

338 + kv_cache_dtype: str,

339 + logits_soft_cap: float | None = None,

340 + attn_type: AttentionType = AttentionType.DECODER,

341 + kv_sharing_target_layer_name: str | None = None,

342 + sinks: torch.Tensor | None = None,

343 + ) -> None:

344 + if alibi_slopes is not None:

345 + raise NotImplementedError("b12x does not support ALiBi.")

346 + if logits_soft_cap not in (None, 0):

347 + raise NotImplementedError("b12x does not support logits soft cap.")

348 + if attn_type != AttentionType.DECODER:

349 + raise NotImplementedError(

350 + "b12x currently supports decoder self-attention only."

351 + )

352 + if is_quantized_kv_cache(kv_cache_dtype) and not _is_b12x_fp8_kv_cache(

353 + kv_cache_dtype

354 + ):

355 + raise NotImplementedError(

356 + "b12x currently supports only fp8/fp8_e4m3 quantized KV cache dtypes."

357 + )

358 + if num_heads % num_kv_heads != 0:

359 + raise ValueError("b12x requires q heads divisible by kv heads.")

360 +

361 + expected_scale = head_size**-0.5

362 + if not math.isclose(float(scale), expected_scale, rel_tol=1e-5, abs_tol=1e-7):

363 + raise NotImplementedError(

364 + "b12x currently requires canonical softmax scale "

365 + f"head_dim**-0.5={expected_scale}, got {scale}."

366 + )

367 + if self.total_cp_world_size > 1:

368 + raise NotImplementedError(

369 + "b12x does not yet support decode/prefill context parallelism."

370 + )

371 +

372 + self.num_heads = int(num_heads)

373 + self.head_size = int(head_size)

374 + self.output_head_size = self.head_size

375 + self.scale = float(scale)

376 + self.num_kv_heads = int(num_kv_heads)

377 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads

378 + self.kv_cache_dtype = kv_cache_dtype

379 + self.attn_type = attn_type

380 + self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

381 + self.window_left = -1 if sliding_window is None else int(sliding_window) - 1

382 +

383 + self._sinks_source = sinks

384 + if sinks is not None and (

385 + sinks.ndim != 1 or int(sinks.shape[0]) != self.num_heads

386 + ):

387 + raise ValueError(

388 + "b12x sinks must have shape "

389 + f"[{self.num_heads}], got {tuple(sinks.shape)}."

390 + )

391 + self.sinks = sinks if sinks is None or sinks.dtype == torch.float32 else None

392 +

393 + vllm_config = get_current_vllm_config()

394 + scheduler_config = vllm_config.scheduler_config

395 + model_config = vllm_config.model_config

396 + cache_config = vllm_config.cache_config

397 + spec_config = vllm_config.speculative_config

398 + default_block_size = int(cache_config.block_size)

399 + if default_block_size not in _B12X_SUPPORTED_PAGE_SIZES:

400 + raise ValueError(

401 + "b12x requires --block-size in "

402 + f"{_B12X_SUPPORTED_PAGE_SIZES}, got "

403 + f"{cache_config.block_size}."

404 + )

405 +

406 + self.device = torch.device("cuda", torch.accelerator.current_device_index())

407 + self.dtype = model_config.dtype

408 + self.kv_torch_dtype = _dtype_from_cache_config(kv_cache_dtype, vllm_config)

409 + if self.dtype != torch.bfloat16:

410 + raise NotImplementedError("b12x currently requires bfloat16 queries.")

411 + max_batched = int(scheduler_config.max_num_batched_tokens)

412 + max_num_seqs = int(scheduler_config.max_num_seqs)

413 + max_model_len = int(model_config.max_model_len)

414 + self._max_num_seqs = max_num_seqs

415 + max_page_table_widths = {

416 + page_size: _max_page_table_width(

417 + max_model_len,

418 + page_size,

419 + max_batched,

420 + model_config.is_hybrid,

421 + )

422 + for page_size in _B12X_SUPPORTED_PAGE_SIZES

423 + }

424 +

425 + # Extend dispatch may depend on the static Q tensor capacity, but never

426 + # on live per-request lengths. Keep a small set of capacity buckets so

427 + # short/tail prefills do not replay the maximum 8K CTA grid.

428 + self._extend_q_capacities = tuple(

429 + sorted(

430 + {

431 + min(max_batched, q_capacity)

432 + for q_capacity in (128, 512, 1024, 2048, 4096, max_batched)

433 + if q_capacity > 0

434 + }

435 + )

436 + )

437 +

438 + paged_attention = get_b12x_paged_attention()

439 + assert paged_attention is not None and paged_attention.is_supported()

440 +

441 + def _extend_work_items(

442 + page_size: int,

443 + q_capacity: int,

444 + batch_size: int,

445 + ) -> int:

446 + capacity = paged_attention.extend_graph_capacity(

447 + device=self.device,

448 + q_dtype=self.dtype,

449 + kv_dtype=self.kv_torch_dtype,

450 + num_q_heads=self.num_heads,

451 + num_kv_heads=self.num_kv_heads,

452 + head_dim_qk=self.head_size,

453 + head_dim_vo=self.output_head_size,

454 + page_size=page_size,

455 + batch=batch_size,

456 + total_q_capacity=q_capacity,

457 + max_cache_page_count=max_page_table_widths[page_size],

458 + window_left=self.window_left,

459 + )

460 + return capacity.max_work_items

461 +

462 + self._paged_attention = paged_attention

463 +

464 + def _make_plan(

465 + page_size: int,

466 + mode: str,

467 + max_total_q: int,

468 + max_batch: int,

469 + max_work_items: int,

470 + max_partial_rows: int,

471 + use_cuda_graph: bool,

472 + num_cache_pages: int,

473 + copy_runtime_metadata: bool,

474 + ) -> Any:

475 + return paged_attention.plan(

476 + paged_attention.Caps(

477 + device=self.device,

478 + mode=mode,

479 + dtype=self.dtype,

480 + kv_dtype=self.kv_torch_dtype,

481 + num_q_heads=self.num_heads,

482 + num_kv_heads=self.num_kv_heads,

483 + head_dim_qk=self.head_size,

484 + head_dim_vo=self.output_head_size,

485 + page_size=page_size,

486 + max_total_q=max_total_q,

487 + max_batch=max_batch,

488 + max_page_table_width=max_page_table_widths[page_size],

489 + max_work_items=max_work_items,

490 + max_partial_rows=max_partial_rows,

491 + # Shape-only planning tensor; runtime cache shape is

492 + # validated by head/page geometry, not page count.

493 + num_cache_pages=num_cache_pages,

494 + use_cuda_graph=use_cuda_graph,

495 + copy_runtime_metadata=copy_runtime_metadata,

496 + )

497 + )

498 +

499 + capture_sizes = vllm_config.compilation_config.cudagraph_capture_sizes or []

500 + decode_plan_sizes = {

501 + int(size) for size in capture_sizes if 0 < int(size) <= max_num_seqs

502 + }

503 + decode_plan_sizes.add(max_num_seqs)

504 +

505 + def _create_decode_plan(page_size: int, batch_size: int) -> Any:

506 + max_page_table_width = max_page_table_widths[page_size]

507 + capacity = paged_attention.decode_graph_capacity(

508 + device=self.device,

509 + q_dtype=self.dtype,

510 + kv_dtype=self.kv_torch_dtype,

511 + num_q_heads=self.num_heads,

512 + num_kv_heads=self.num_kv_heads,

513 + head_dim_qk=self.head_size,

514 + head_dim_vo=self.output_head_size,

515 + page_size=page_size,

516 + batch=batch_size,

517 + max_cache_page_count=max_page_table_width,

518 + window_left=self.window_left,

519 + )

520 + plan = _make_plan(

521 + page_size,

522 + "decode",

523 + batch_size,

524 + batch_size,

525 + capacity.max_work_items,

526 + capacity.max_partial_rows,

527 + True,

528 + max_page_table_width,

529 + True,

530 + )

531 + plan.prepare_decode_graph_replay_state(

532 + batch=batch_size,

533 + total_q_capacity=batch_size,

534 + max_page_table_width=max_page_table_width,

535 + max_cache_page_count=max_page_table_width,

536 + window_left=self.window_left,

537 + )

538 + return plan

539 +

540 + self._create_decode_plan = _create_decode_plan

541 + self._verify_q_per_req = 0

542 + if spec_config is not None:

543 + self._verify_q_per_req = 1 + int(

544 + getattr(spec_config, "num_speculative_tokens", None) or 0

545 + )

546 + if self._verify_q_per_req <= 1:

547 + self._verify_q_per_req = 0

548 +

549 + def _create_verify_plan(page_size: int, batch_size: int) -> Any:

550 + if self._verify_q_per_req <= 1:

551 + raise RuntimeError("b12x verifier plan requested without speculation")

552 + max_page_table_width = max_page_table_widths[page_size]

553 + total_q = batch_size * self._verify_q_per_req

554 + capacity = paged_attention.verify_graph_capacity(

555 + device=self.device,

556 + q_dtype=self.dtype,

557 + kv_dtype=self.kv_torch_dtype,

558 + num_q_heads=self.num_heads,

559 + num_kv_heads=self.num_kv_heads,

560 + head_dim_qk=self.head_size,

561 + head_dim_vo=self.output_head_size,

562 + page_size=page_size,

563 + batch=batch_size,

564 + query_len=self._verify_q_per_req,

565 + max_cache_page_count=max_page_table_width,

566 + window_left=self.window_left,

567 + )

568 + plan = _make_plan(

569 + page_size,

570 + "verify",

571 + total_q,

572 + batch_size,

573 + capacity.max_work_items,

574 + capacity.max_partial_rows,

575 + True,

576 + max_page_table_width,

577 + True,

578 + )

579 + page_ids = torch.arange(

580 + max_page_table_width,

581 + dtype=torch.int32,

582 + device=self.device,

583 + )

584 + max_page_table = page_ids.unsqueeze(0).expand(batch_size, -1).contiguous()

585 + max_cache_seqlens = torch.full(

586 + (batch_size,),

587 + capacity.representative_cache_seqlen,

588 + dtype=torch.int32,

589 + device=self.device,

590 + )

591 + max_cu_seqlens_q = torch.arange(

592 + 0,

593 + total_q + 1,

594 + self._verify_q_per_req,

595 + dtype=torch.int32,

596 + device=self.device,

597 + )

598 + plan.prepare_graph_replay_state(

599 + page_table=max_page_table,

600 + cache_seqlens=max_cache_seqlens,

601 + cu_seqlens_q=max_cu_seqlens_q,

602 + active_total_q=total_q,

603 + window_left=self.window_left,

604 + )

605 + return plan

606 +

607 + self._create_verify_plan = _create_verify_plan

608 +

609 + def _create_extend_plan(

610 + page_size: int,

611 + batch_size: int,

612 + q_capacity: int,

613 + ) -> Any:

614 + """Prepare a fixed-capacity extend plan without reading live lengths."""

615 + max_page_table_width = max_page_table_widths[page_size]

616 + plan = _make_plan(

617 + page_size,

618 + "extend",

619 + q_capacity,

620 + batch_size,

621 + _extend_work_items(page_size, q_capacity, batch_size),

622 + 0,

623 + True,

624 + max_page_table_width,

625 + False,

626 + )

627 + page_ids = torch.arange(

628 + max_page_table_width,

629 + dtype=torch.int32,

630 + device=self.device,

631 + )

632 + max_page_table = page_ids.unsqueeze(0).expand(batch_size, -1).contiguous()

633 + max_cache_seqlens = torch.full(

634 + (batch_size,),

635 + min(max_model_len, max_page_table_width * page_size),

636 + dtype=torch.int32,

637 + device=self.device,

638 + )

639 + # Put one row in every request except the last, which owns the

640 + # remainder. This represents the full total-Q capacity while the

641 + # replay kernel remains responsible for packing arbitrary live

642 + # per-request lengths from device cu_seqlens_q.

643 + max_cu_seqlens_q = torch.arange(

644 + 0,

645 + batch_size + 1,

646 + dtype=torch.int32,

647 + device=self.device,

648 + )

649 + max_cu_seqlens_q[-1] = q_capacity

650 + plan.prepare_graph_replay_state(

651 + page_table=max_page_table,

652 + cache_seqlens=max_cache_seqlens,

653 + cu_seqlens_q=max_cu_seqlens_q,

654 + active_total_q=q_capacity,

655 + window_left=self.window_left,

656 + )

657 + return plan

658 +

659 + self._create_extend_plan = _create_extend_plan

660 + decode_scratch_envelopes = {

661 + page_size: paged_attention.decode_graph_scratch_envelope(

662 + device=self.device,

663 + q_dtype=self.dtype,

664 + kv_dtype=self.kv_torch_dtype,

665 + num_q_heads=self.num_heads,

666 + num_kv_heads=self.num_kv_heads,

667 + head_dim_qk=self.head_size,

668 + head_dim_vo=self.output_head_size,

669 + page_size=page_size,

670 + max_batch=max_num_seqs,

671 + max_page_table_width=max_page_table_widths[page_size],

672 + max_cache_page_count=max_page_table_widths[page_size],

673 + window_left=self.window_left,

674 + copy_runtime_metadata=True,

675 + )

676 + for page_size in _B12X_SUPPORTED_PAGE_SIZES

677 + }

678 + self._decode_plans: dict[tuple[int, int], Any] = {}

679 + self._verify_plans: dict[tuple[int, int], Any] = {}

680 + self._extend_plans: dict[tuple[int, int, int], Any] = {}

681 + for page_size in _B12X_SUPPORTED_PAGE_SIZES:

682 + for batch_size in sorted(decode_plan_sizes):

683 + self._decode_plans[page_size, batch_size] = self._create_decode_plan(

684 + page_size, batch_size

685 + )

686 + if self._verify_q_per_req > 1:

687 + for batch_size in range(1, max_num_seqs + 1):

688 + self._verify_plans[page_size, batch_size] = (

689 + self._create_verify_plan(page_size, batch_size)

690 + )

691 + for batch_size in range(1, max_num_seqs + 1):

692 + for q_capacity in self._extend_q_capacities:

693 + # Equal capacity is necessarily one query token per

694 + # request, which is handled by the decode plan.

695 + if batch_size >= q_capacity:

696 + continue

697 + self._extend_plans[page_size, batch_size, q_capacity] = (

698 + self._create_extend_plan(

699 + page_size,

700 + batch_size,

701 + q_capacity,

702 + )

703 + )

704 + self._scratch_nbytes = max(

705 + *(int(envelope.nbytes) for envelope in decode_scratch_envelopes.values()),

706 + *(int(plan.layout.nbytes) for plan in self._verify_plans.values()),

707 + *(int(plan.layout.nbytes) for plan in self._extend_plans.values()),

708 + )

709 +

710 + current_workspace_manager().get_simultaneous(

711 + ((self._scratch_nbytes,), torch.uint8),

712 + )

713 +

714 + self.supports_quant_query_input = False

715 + register_cutedsl_warmup_provider(self)

716 +

717 + logger.info_once(

718 + "Using b12x with q_heads=%d kv_heads=%d head_dim_qk=%d "

719 + "head_dim_vo=%d window_left=%d planned_page_sizes=%s "

720 + "verify_q_per_req=%d extend_q_capacities=%s scratch=%d bytes.",

721 + self.num_heads,

722 + self.num_kv_heads,

723 + self.head_size,

724 + self.output_head_size,

725 + self.window_left,

726 + _B12X_SUPPORTED_PAGE_SIZES,

727 + self._verify_q_per_req,

728 + self._extend_q_capacities,

729 + self._scratch_nbytes,

730 + )

731 +

732 + def _compile_paged_extend_entry(self, page_size: int) -> None:

733 + """Compile fixed-capacity paged-prefill entries without a live plan."""

734 + warmup_plans: list[tuple[int, int, Any, bool]] = []

735 + for batch_size in range(1, self._max_num_seqs + 1):

736 + candidates = sorted(

737 + (q_capacity, plan)

738 + for (plan_page_size, plan_batch, q_capacity), plan in (

739 + self._extend_plans.items()

740 + )

741 + if plan_page_size == page_size and plan_batch == batch_size

742 + )

743 + for index, (q_capacity, plan) in enumerate(candidates):

744 + warmup_plans.append((batch_size, q_capacity, plan, index == 0))

745 + if not warmup_plans:

746 + return

747 +

748 + max_q_rows = max(

749 + min(q_capacity, max(64, batch_size + 1))

750 + for batch_size, q_capacity, _, _ in warmup_plans

751 + )

752 + q = torch.zeros(

753 + (max_q_rows, self.num_heads, self.head_size),

754 + dtype=self.dtype,

755 + device=self.device,

756 + )

757 + output = torch.zeros(

758 + (max_q_rows, self.num_heads, self.output_head_size),

759 + dtype=self.dtype,

760 + device=self.device,

761 + )

762 + kv_cache = torch.zeros(

763 + (1, 2, page_size, self.num_kv_heads * self.head_size),

764 + dtype=self.kv_torch_dtype,

765 + device=self.device,

766 + )

767 + key_cache, value_cache = self._kv_cache_views(kv_cache)

768 + (scratch_storage,) = current_workspace_manager().get_simultaneous(

769 + ((self._scratch_nbytes,), torch.uint8),

770 + )

771 + for batch_size, q_capacity, plan, execute in warmup_plans:

772 + q_rows = min(q_capacity, max(64, batch_size + 1))

773 + page_table = torch.zeros(

774 + (batch_size, plan.caps.max_page_table_width),

775 + dtype=torch.int32,

776 + device=self.device,

777 + )

778 + cache_seqlens = torch.full(

779 + (batch_size,), page_size, dtype=torch.int32, device=self.device

780 + )

781 + cu_seqlens_q = torch.arange(

782 + 0,

783 + batch_size + 1,

784 + dtype=torch.int32,

785 + device=self.device,

786 + )

787 + cu_seqlens_q[-1] = q_rows

788 + k_descale = None

789 + v_descale = None

790 + if _is_b12x_fp8_kv_cache(self.kv_cache_dtype):

791 + k_descale = torch.ones(

792 + (), dtype=torch.float32, device=self.device

793 + ).expand(batch_size)

794 + v_descale = torch.ones(

795 + (), dtype=torch.float32, device=self.device

796 + ).expand(batch_size)

797 + binding = plan.bind(

798 + scratch=scratch_storage,

799 + q=q[:q_rows],

800 + k_cache=key_cache,

801 + v_cache=value_cache,

802 + output=output[:q_rows],

803 + page_table=page_table,

804 + cache_seqlens=cache_seqlens,

805 + cu_seqlens_q=cu_seqlens_q,

806 + window_left=self.window_left,

807 + attention_sink_bias=self.sinks,

808 + k_descale=k_descale,

809 + v_descale=v_descale,

810 + )

811 + self._paged_attention.compile(binding=binding)

812 + if execute:

813 + # Compile-only warmup does not launch the device-side compact

814 + # scheduler. One execution per batch covers the capture-static

815 + # metadata variant shared by its Q-capacity plans.

816 + self._paged_attention.run(binding=binding)

817 +

818 + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:

819 + common_key = (

820 + "b12x_paged_extend",

821 + str(self.device),

822 + str(self.dtype),

823 + str(self.kv_torch_dtype),

824 + self.num_heads,

825 + self.num_kv_heads,

826 + self.head_size,

827 + self.output_head_size,

828 + self.window_left,

829 + self.sinks is not None,

830 + )

831 + return tuple(

832 + CuTeDSLCompileUnit(

833 + name="b12x_paged_extend",

834 + key=(*common_key, page_size),

835 + compile=partial(self._compile_paged_extend_entry, page_size),

836 + )

837 + for page_size in _B12X_SUPPORTED_PAGE_SIZES

838 + )

839 +

840 + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:

841 + del act_dtype

842 + source_sinks = self._sinks_source

843 + if source_sinks is None:

844 + return

845 + if source_sinks.dtype == torch.float32:

846 + self.sinks = source_sinks

847 + elif self.sinks is None or self.sinks.dtype != torch.float32:

848 + self.sinks = source_sinks.to(torch.float32)

849 + else:

850 + self.sinks.copy_(source_sinks)

851 +

852 + def _prepare_fp8_descales(

853 + self,

854 + layer: AttentionLayer,

855 + num_reqs: int,

856 + device: torch.device,

857 + ) -> tuple[torch.Tensor | None, torch.Tensor | None]:

858 + if not _is_b12x_fp8_kv_cache(self.kv_cache_dtype):

859 + return None, None

860 + if num_reqs <= 0:

861 + raise ValueError("b12x fp8 KV descale request count must be positive.")

862 +

863 + def _prepare(scale: torch.Tensor, name: str) -> torch.Tensor:

864 + if scale.device != device:

865 + raise RuntimeError(f"b12x {name} must be on the query device.")

866 + if scale.dtype != torch.float32:

867 + raise RuntimeError(f"b12x {name} must be float32.")

868 + if scale.ndim == 0:

869 + return scale.expand(num_reqs)

870 + if scale.ndim == 1:

871 + if int(scale.shape[0]) == 1:

872 + return scale.expand(num_reqs)

873 + if int(scale.shape[0]) >= num_reqs:

874 + return scale[:num_reqs]

875 + raise ValueError(

876 + f"b12x {name} must be scalar or rank-1 with at least "

877 + f"{num_reqs} values; got shape {tuple(scale.shape)}."

878 + )

879 +

880 + return _prepare(layer._k_scale, "k_scale"), _prepare(layer._v_scale, "v_scale")

881 +

882 + def _kv_cache_views(

883 + self,

884 + kv_cache: torch.Tensor,

885 + ) -> tuple[torch.Tensor, torch.Tensor]:

886 + key_cache, value_cache = kv_cache.unbind(1)

887 + key_cache = key_cache.unflatten(-1, (self.num_kv_heads, self.head_size))

888 + value_cache = value_cache.unflatten(-1, (self.num_kv_heads, self.head_size))

889 + key_cache = canonicalize_singleton_dim_strides(key_cache)

890 + value_cache = canonicalize_singleton_dim_strides(value_cache)

891 + if _is_b12x_fp8_kv_cache(self.kv_cache_dtype):

892 + fp8_dtype = current_platform.fp8_dtype()

893 + if key_cache.dtype == torch.uint8:

894 + key_cache = key_cache.view(fp8_dtype)

895 + if value_cache.dtype == torch.uint8:

896 + value_cache = value_cache.view(fp8_dtype)

897 + if (

898 + key_cache.dtype != self.kv_torch_dtype

899 + or value_cache.dtype != self.kv_torch_dtype

900 + ):

901 + raise TypeError(

902 + f"b12x plan expects KV dtype {self.kv_torch_dtype}, got "

903 + f"{key_cache.dtype}/{value_cache.dtype}."

904 + )

905 + return key_cache, value_cache

906 +

907 + def _select_plan(

908 + self,

909 + attn_metadata: B12xPagedMetadata,

910 + total_q: int,

911 + q_capacity: int,

912 + num_reqs: int,

913 + page_size: int,

914 + ) -> Any:

915 + if attn_metadata.max_query_len <= 1 and int(total_q) == int(num_reqs):

916 + batch_size = int(total_q)

917 + plan_key = (page_size, batch_size)

918 + plan = self._decode_plans.get(plan_key)

919 + if plan is None:

920 + if _capture_alloc_forbidden():

921 + raise RuntimeError(

922 + "b12x decode plan was not prepared before CUDA graph "

923 + f"capture for page size {page_size}, batch size "

924 + f"{batch_size}."

925 + )

926 + plan = self._create_decode_plan(page_size, batch_size)

927 + if int(plan.layout.nbytes) > self._scratch_nbytes:

928 + raise RuntimeError(

929 + "b12x lazily created decode plan exceeds reserved "

930 + f"scratch: {int(plan.layout.nbytes)} > "

931 + f"{self._scratch_nbytes} bytes."

932 + )

933 + self._decode_plans[plan_key] = plan

934 + return plan

935 + elif (

936 + self._verify_q_per_req > 1

937 + and attn_metadata.max_query_len == self._verify_q_per_req

938 + and int(total_q) == int(num_reqs) * self._verify_q_per_req

939 + ):

940 + plan_key = (page_size, int(num_reqs))

941 + plan = self._verify_plans.get(plan_key)

942 + if plan is None:

943 + if _capture_alloc_forbidden():

944 + raise RuntimeError(

945 + "b12x verifier plan was not prepared before CUDA "

946 + f"graph capture for page size {page_size}, batch size "

947 + f"{num_reqs}."

948 + )

949 + plan = self._create_verify_plan(page_size, int(num_reqs))

950 + if int(plan.layout.nbytes) > self._scratch_nbytes:

951 + raise RuntimeError(

952 + "b12x lazily created verifier plan exceeds reserved "

953 + f"scratch: {int(plan.layout.nbytes)} > "

954 + f"{self._scratch_nbytes} bytes."

955 + )

956 + self._verify_plans[plan_key] = plan

957 + return plan

958 + extend_q_capacity = next(

959 + (

960 + capacity

961 + for capacity in self._extend_q_capacities

962 + if q_capacity <= capacity

963 + ),

964 + None,

965 + )

966 + if extend_q_capacity is None:

967 + raise ValueError(

968 + f"b12x extend Q capacity {q_capacity} exceeds prepared "

969 + f"maximum {self._extend_q_capacities[-1]}."

970 + )

971 + return self._extend_plans[

972 + page_size,

973 + int(num_reqs),

974 + extend_q_capacity,

975 + ]

976 +

977 + def forward(

978 + self,

979 + layer: AttentionLayer,

980 + query: torch.Tensor,

981 + key: torch.Tensor,

982 + value: torch.Tensor,

983 + kv_cache: torch.Tensor,

984 + attn_metadata: B12xPagedMetadata,

985 + output: torch.Tensor,

986 + output_scale: torch.Tensor | None = None,

987 + output_block_scale: torch.Tensor | None = None,

988 + ) -> torch.Tensor:

989 + del key, value

990 + if output_scale is not None or output_block_scale is not None:

991 + raise NotImplementedError(

992 + "b12x does not support fused output quantization."

993 + )

994 + if attn_metadata is None:

995 + return output.fill_(0)

996 + if output.shape[-1] != self.output_head_size:

997 + raise ValueError(

998 + f"b12x expected output head dim {self.output_head_size}, got "

999 + f"{output.shape[-1]}."

1000 + )

1001 + if kv_cache.numel() == 0:

1002 + return output.fill_(0)

1003 +

1004 + # In FULL cudagraph mode vLLM may pad attention metadata to the graph

1005 + # bucket while still passing per-layer Q/output tensors with only the

1006 + # real rows. Use tensor capacity as the launch contract and avoid

1007 + # selecting decode graph replay for padded virtual requests.

1008 + q_capacity = min(

1009 + int(query.shape[0]),

1010 + int(output.shape[0]),

1011 + )

1012 + num_actual_tokens = min(

1013 + int(attn_metadata.num_actual_tokens),

1014 + q_capacity,

1015 + )

1016 + if num_actual_tokens <= 0:

1017 + return output

1018 + q = query[:num_actual_tokens]

1019 + out = output[:num_actual_tokens]

1020 + if q.dtype != self.dtype or out.dtype != self.dtype:

1021 + raise TypeError(

1022 + f"b12x plan expects dtype {self.dtype}, got "

1023 + f"q={q.dtype}, output={out.dtype}."

1024 + )

1025 +

1026 + key_cache, value_cache = self._kv_cache_views(kv_cache)

1027 + page_size = _kv_page_size(key_cache, value_cache)

1028 + if not attn_metadata.causal:

1029 + raise NotImplementedError("b12x supports causal attention only.")

1030 +

1031 + page_table = _ensure_i32_contiguous(attn_metadata.block_table, "block_table")

1032 + cache_seqlens = _ensure_i32_contiguous(attn_metadata.seq_lens, "seq_lens")

1033 + cu_seqlens_q = _ensure_i32_contiguous(

1034 + attn_metadata.query_start_loc,

1035 + "query_start_loc",

1036 + )

1037 + num_reqs = int(cache_seqlens.shape[0])

1038 + if attn_metadata.max_query_len <= 1 and num_actual_tokens < num_reqs:

1039 + num_reqs = num_actual_tokens

1040 + page_table = page_table[:num_reqs]

1041 + cache_seqlens = cache_seqlens[:num_reqs]

1042 + cu_seqlens_q = cu_seqlens_q[: num_reqs + 1]

1043 + k_descale, v_descale = self._prepare_fp8_descales(

1044 + layer,

1045 + num_reqs,

1046 + q.device,

1047 + )

1048 + plan = self._select_plan(

1049 + attn_metadata,

1050 + num_actual_tokens,

1051 + q_capacity,

1052 + num_reqs,

1053 + page_size,

1054 + )

1055 + (scratch_storage,) = current_workspace_manager().get_simultaneous(

1056 + ((self._scratch_nbytes,), torch.uint8),

1057 + )

1058 + binding = plan.bind(

1059 + scratch=scratch_storage,

1060 + q=q,

1061 + k_cache=key_cache,

1062 + v_cache=value_cache,

1063 + output=out,

1064 + page_table=page_table,

1065 + cache_seqlens=cache_seqlens,

1066 + cu_seqlens_q=cu_seqlens_q,

1067 + window_left=self.window_left,

1068 + active_total_q=(None if plan.caps.mode == "extend" else num_actual_tokens),

1069 + attention_sink_bias=self.sinks,

1070 + k_descale=k_descale,

1071 + v_descale=v_descale,

1072 + )

1073 + self._paged_attention.run(binding=binding)

1074 + return output

1075 +

1076 + def do_kv_cache_update(

1077 + self,

1078 + layer: AttentionLayer,

1079 + key: torch.Tensor,

1080 + value: torch.Tensor,

1081 + kv_cache: torch.Tensor,

1082 + slot_mapping: torch.Tensor,

1083 + ) -> None:

1084 + if kv_cache.numel() == 0:

1085 + return

1086 + key_cache, value_cache = self._kv_cache_views(kv_cache)

1087 + torch.ops._C_cache_ops.reshape_and_cache_flash(

1088 + key,

1089 + value,

1090 + key_cache,

1091 + value_cache,

1092 + slot_mapping,

1093 + self.kv_cache_dtype,

1094 + layer._k_scale,

1095 + layer._v_scale,

1096 + )

vllm/v1/attention/backends/registry.py

@@ -94,6 +94,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):

94 ROCM_FLASHMLA_SPARSE_DSV4 = (

95 "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend"

96 )

97 + B12X = "vllm.v1.attention.backends.b12x.B12xPagedAttentionBackend"

98 FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"

99 FLASH_ATTN_MLA_SPARSE = (

100 "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend"