@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c8d7a7dd13

[Perf][Attention] Remove D2H sync from FlashInfer SM90 sparse MLA plan under async scheduling (#58684)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Sep 27, 2026 at 13:06 UTC c8d7a7dd13e2b40c013fb6d46be800937e4335fb

2 files changed +381 -106

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

+181 -28

@@ -9,6 +9,7 @@ causality), ckv/kpe cache splitting, and the backend's model-shape gates.

9 """

11 from types import SimpleNamespace

12 +from typing import Any

14 import pytest

15 import torch

@@ -50,6 +51,8 @@ class FakeWrapper:

51 def __init__(self):

52 self.plan_args = None

53 self.run_args = None

54 + self._plan_info: Any = None

55 + self._planned_backend: Any = None

57 def plan(self, *args, **kwargs):

58 self.plan_args = (args, kwargs)

@@ -69,8 +72,11 @@ class FakeState:

72 self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32)

73 self.wrapper = FakeWrapper()

74 self.plan_calls = []

75 + self.index_topk = 2048

76 + self.index_kpool = 4

77 + self.max_valid = 2051

73 - def plan(self, num_tokens, kv_lens):

79 + def plan(self, num_tokens, kv_lens, cam, req_id_per_token):

80 self.plan_calls.append((num_tokens, kv_lens))

82 def pack_indices(self, slots):

@@ -168,13 +174,12 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):

174 def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes):

175 """MHA prefill rows must not make the MQA kernel read beyond its query."""

176 builder = object.__new__(FlashInferMLASparseSM90Builder)

171 - builder._index_topk = 2048

172 - builder._index_kpool = 4

173 - builder._async_scheduling = False

177 + builder._adaptive_verification = False

178 builder.state = FakeState(TOPK)

179 builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha)

180 metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata)

181 metadata.state = None

182 + metadata.req_id_per_token = None

183 metadata.num_prefills = 1

184 metadata.num_decode_tokens = num_decodes

185 monkeypatch.setattr(

@@ -206,12 +211,17 @@ def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_dec

211 def test_plan_uses_state_params(monkeypatch):

212 """The NoPE/rope dims and scale live on the builder state, not the layer.

213

209 - plan() takes exact per-row KV lengths; the schedule is rebuilt on every

210 - call (contexts grow between steps) and the indptrs are always full-size

214 + Without a device batch layout plan() takes exact per-row KV lengths and

215 + rebuilds the schedule on every call; the indptrs are always full-size

216 with zero-query padding rows past num_tokens.

217 """

213 - impl, rows = make_impl(64, "auto")

218 wrapper = FakeWrapper()

219 + plan_info = [0] * sm90_mod._PLAN_INFO_LEN

220 + plan_info[sm90_mod._PI_KV_INDPTR] = 64 * 4

221 + wrapper._plan_info = torch.tensor(plan_info)

222 + wrapper._planned_backend = SimpleNamespace(

223 + _int_workspace_buffer=torch.zeros(64, dtype=torch.int32)

224 + )

225 state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)

226 state.device = torch.device("cpu")

227 state.wrapper = wrapper

@@ -227,8 +237,9 @@ def test_plan_uses_state_params(monkeypatch):

237 state._qo_cpu = torch.empty(5, dtype=torch.int32)

238 state._kv_cpu = torch.empty(5, dtype=torch.int32)

239 state._lens_cpu = torch.full((4,), TOPK, dtype=torch.int32)

240 + state._saved_kv_end = torch.empty(0, dtype=torch.int32)

241

231 - state.plan(3, torch.tensor([2, 5, 7], dtype=torch.int32))

242 + state.plan(3, torch.tensor([2, 5, 7], dtype=torch.int32), None, None)

243 assert wrapper.plan_args is not None

244 args, kwargs = wrapper.plan_args

245 (qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args

@@ -241,10 +252,10 @@ def test_plan_uses_state_params(monkeypatch):

252 assert kwargs["kv_data_type"] == torch.bfloat16

253

254 # Replanning a smaller batch must clear the previous rows' lengths.

244 - state.plan(1, torch.tensor([TOPK], dtype=torch.int32))

255 + state.plan(1, torch.tensor([TOPK], dtype=torch.int32), None, None)

256 assert state._kv_cpu.tolist() == [0, TOPK, TOPK, TOPK, TOPK]

257 assert state._lens_cpu.tolist() == [TOPK, 0, 0, 0]

247 - state.plan(0, torch.empty(0, dtype=torch.int32))

258 + state.plan(0, torch.empty(0, dtype=torch.int32), None, None)

259 assert state._kv_cpu.tolist() == [0, 0, 0, 0, 0]

260 assert state._lens_cpu.tolist() == [0, 0, 0, 0]

261

@@ -279,38 +290,49 @@ def test_pack_indices_replays_with_updated_offsets(width):

290 assert (state.kv_indices[expected.numel() :] == -99).all()

291

292

282 -def test_kv_lens_host_formula():

293 +@pytest.mark.parametrize("adaptive", [False, True])

294 +def test_kv_lens_host_formula(adaptive):

295 """Per-row host lengths: context == position + 1; capped at

284 - index_topk + trailing-pool remainder past the sparse threshold."""

296 + index_topk + trailing-pool remainder past the sparse threshold.

297 +

298 + Exact (synced) without a host upper bound; with one, a sync-free bound

299 + capped at the largest valid count, which the plan clamps on device.

300 + """

301 builder = object.__new__(FlashInferMLASparseSM90Builder)

286 - builder._index_topk = 2048

287 - builder._index_kpool = 4

288 - builder._async_scheduling = False

302 + builder.state = FakeState(TOPK)

303 + builder._adaptive_verification = adaptive

304 + seq_lens = torch.tensor([100, 9, 3000], dtype=torch.int32)

305 cam = SimpleNamespace(

306 num_reqs=3,

307 query_start_loc_cpu=torch.tensor([0, 5, 7, 10], dtype=torch.int32),

292 - seq_lens=torch.tensor([100, 9, 3000], dtype=torch.int32),

293 - seq_lens_cpu_upper_bound=torch.tensor([100, 9, 3000], dtype=torch.int32),

294 - positions=None,

308 + seq_lens=seq_lens,

309 + seq_lens_cpu_upper_bound=None,

310 )

296 - num_rows, lens = builder._kv_lens_host(cam)

297 - assert num_rows == 10

311 + num_rows, lens, exact = builder._kv_lens_host(cam)

312 + assert num_rows == 10 and exact

313 # req0: positions 95..99 -> ctx 96..100 (all <= 2048: full context)

314 # req1: positions 7,8 -> ctx 8,9

315 # req2: positions 2997..2999 -> ctx 2998..3000 (> 2048: topk + ctx%4)

316 assert lens.tolist() == [96, 97, 98, 99, 100, 8, 9, 2050, 2051, 2048]

317

318 + cam.seq_lens_cpu_upper_bound = seq_lens

319 + num_rows, lens, exact = builder._kv_lens_host(cam)

320 + assert num_rows == 10 and not exact

321 + if adaptive:

322 + # Host query split may differ from the device one: batch-wide bound.

323 + assert lens.tolist() == [2051] * 10

324 + else:

325 + assert lens.tolist() == [96, 97, 98, 99, 100, 8, 9, 2051, 2051, 2051]

326 +

327

328 def test_kv_lens_host_empty():

329 builder = object.__new__(FlashInferMLASparseSM90Builder)

306 - builder._index_topk = 2048

307 - builder._index_kpool = 4

330 cam = SimpleNamespace(

331 num_reqs=0,

332 query_start_loc_cpu=torch.tensor([0], dtype=torch.int32),

333 seq_lens=torch.zeros(0, dtype=torch.int32),

334 )

313 - num_rows, lens = builder._kv_lens_host(cam)

335 + num_rows, lens, _ = builder._kv_lens_host(cam)

336 assert num_rows == 0 and lens.numel() == 0

337

338

@@ -334,7 +356,13 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens)

356 "__init__",

357 lambda *_args, **_kwargs: None,

358 )

337 - monkeypatch.setattr(sm90_mod, "_SM90State", lambda *_args, **_kwargs: None)

359 +

360 + class _RecordingState:

361 + def __init__(self, *_args, **kwargs):

362 + self.index_topk = kwargs["index_topk"]

363 + self.index_kpool = kwargs["index_kpool"]

364 +

365 + monkeypatch.setattr(sm90_mod, "_SM90State", _RecordingState)

366 impl, _ = make_impl(64)

367 spec = MLAAttentionSpec(

368 block_size=BLOCK_SIZE, num_kv_heads=1, head_size=576, dtype=torch.bfloat16

@@ -351,6 +379,7 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens)

379 max_num_batched_tokens=64, async_scheduling=False

380 ),

381 model_config=SimpleNamespace(hf_text_config=hf_config),

382 + speculative_config=None,

383 )

384 builder = FlashInferMLASparseSM90Builder(

385 spec, ["attn"], vllm_config, torch.device("cpu")

@@ -359,14 +388,138 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens)

388 cam = SimpleNamespace(

389 num_reqs=2,

390 query_start_loc_cpu=torch.tensor([0, 4, 6], dtype=torch.int32),

362 - seq_lens_cpu_upper_bound=torch.tensor([42295, 100], dtype=torch.int32),

363 - positions=None,

391 + seq_lens=torch.tensor([42295, 100], dtype=torch.int32),

392 + seq_lens_cpu_upper_bound=None,

393 )

365 - num_rows, lens = builder._kv_lens_host(cam)

366 - assert num_rows == 6

394 + num_rows, lens, exact = builder._kv_lens_host(cam)

395 + assert num_rows == 6 and exact

396 assert lens.tolist() == prefill_lens + [99, 100]

397

398

399 +def _sm90_flashinfer_available():

400 + return (

401 + torch.cuda.is_available()

402 + and torch.cuda.get_device_capability()[0] == 9

403 + and sm90_mod.has_flashinfer_sm90_nope_mla()

404 + )

405 +

406 +

407 +SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH = 16, 2048, 4, 2176

408 +

409 +

410 +def _real_state(rows):

411 + return sm90_mod._SM90State(

412 + torch.device("cuda"),

413 + SM90_HEADS,

414 + torch.bfloat16,

415 + rows,

416 + SM90_WIDTH,

417 + HEAD,

418 + 0,

419 + HEAD**-0.5,

420 + SM90_TOPK,

421 + SM90_KPOOL,

422 + )

423 +

424 +

425 +def _clamp_batch(num_reqs, q_per_req, lo, hi, topk, kpool, width):

426 + """Random batch: per-row contexts, device layout, and top-k slots."""

427 + last = torch.randint(lo + q_per_req, hi, (num_reqs,))

428 + ctx = (last[:, None] + torch.arange(1 - q_per_req, 1)).reshape(-1)

429 + valid = torch.where(ctx <= topk, ctx, topk + ctx % kpool)

430 + cam = SimpleNamespace(

431 + seq_lens=last.to(torch.int32).cuda(),

432 + query_start_loc=(torch.arange(num_reqs + 1) * q_per_req).to(torch.int32).cuda(),

433 + )

434 + req_id = torch.arange(num_reqs, dtype=torch.int32).repeat_interleave(q_per_req)

435 + slots = torch.randint(0, 1 << 16, (ctx.numel(), width), dtype=torch.int32)

436 + slots[torch.arange(width) >= valid[:, None]] = -1

437 + return ctx, valid.to(torch.int32), cam, req_id.cuda(), slots.cuda()

438 +

439 +

440 +@pytest.mark.skipif(

441 + not _sm90_flashinfer_available(), reason="Requires SM90 FlashInfer MLA"

442 +)

443 +@pytest.mark.parametrize(

444 + "num_reqs,q_per_req,lo,hi",

445 + [

446 + (64, 1, 1, 3000),

447 + (16, 6, 2030, 2070), # MTP-style rows straddling index_topk

448 + (4096, 1, 1, 3000), # clamp spans several program blocks

449 + ],

450 +)

451 +def test_upper_bound_plan_matches_exact(num_reqs, q_per_req, lo, hi):

452 + """Sync-free planning from optimistic host bounds (async spec decode)

453 + must match exact planning: the device clamp trims every work item back

454 + to its row's valid count, so no -1 slot is read."""

455 + torch.manual_seed(0)

456 + heads, topk, kpool, width = SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH

457 + ctx, valid, cam, req_id, slots = _clamp_batch(

458 + num_reqs, q_per_req, lo, hi, topk, kpool, width

459 + )

460 + rows = ctx.numel()

461 +

462 + kv = torch.randn(1 << 16, 1, HEAD, device="cuda", dtype=torch.bfloat16)

463 + q = torch.randn(rows, heads, HEAD, device="cuda", dtype=torch.bfloat16)

464 + q_pe = q.new_empty(rows, heads, 0)

465 +

466 + def run(state):

467 + state.pack_indices(slots)

468 + return state.wrapper.run(q, q_pe, kv, kv[..., :0])

469 +

470 + exact = _real_state(rows)

471 + exact.plan(rows, valid, None, None)

472 + ub = _real_state(rows)

473 + host_bound = torch.clamp(ctx + torch.randint(0, 7, (rows,)), max=ub.max_valid)

474 + ub.plan(rows, host_bound, cam, req_id)

475 + out = run(ub)

476 + assert not out.isnan().any()

477 + torch.testing.assert_close(out, run(exact), atol=1e-2, rtol=1e-2)

478 +

479 +

480 +@pytest.mark.skipif(

481 + not _sm90_flashinfer_available(), reason="Requires SM90 FlashInfer MLA"

482 +)

483 +def test_upper_bound_plan_reused_across_steps(monkeypatch):

484 + """Growing contexts (draft/decode steps) reuse the padded plan and only

485 + re-clamp on device, matching an exact replan every step."""

486 + torch.manual_seed(0)

487 + heads, topk, kpool, width = SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH

488 + rows = 48

489 + ctx, _, _, req_id, _ = _clamp_batch(rows, 1, 1, 3000, topk, kpool, width)

490 +

491 + exact, ub = _real_state(rows), _real_state(rows)

492 + replans: list[int] = []

493 + plan_fn = ub._plan

494 +

495 + def counting_plan(num_tokens, kv_lens):

496 + replans.append(num_tokens)

497 + plan_fn(num_tokens, kv_lens)

498 +

499 + monkeypatch.setattr(ub, "_plan", counting_plan)

500 + kv = torch.randn(1 << 16, 1, HEAD, device="cuda", dtype=torch.bfloat16)

501 + for step in range(5):

502 + step_ctx = ctx + step

503 + valid = torch.where(step_ctx <= topk, step_ctx, topk + step_ctx % kpool)

504 + slots = torch.randint(0, 1 << 16, (rows, width), dtype=torch.int32)

505 + slots[torch.arange(width) >= valid[:, None]] = -1

506 + slots = slots.cuda()

507 + cam = SimpleNamespace(

508 + seq_lens=step_ctx.to(torch.int32).cuda(),

509 + query_start_loc=torch.arange(rows + 1, dtype=torch.int32).cuda(),

510 + )

511 + num_plans = len(replans)

512 + ub.plan(rows, torch.clamp(step_ctx + 5, max=ub.max_valid), cam, req_id)

513 + assert len(replans) - num_plans == (1 if step == 0 else 0)

514 + exact.plan(rows, valid.to(torch.int32), None, None)

515 + q = torch.randn(rows, heads, HEAD, device="cuda", dtype=torch.bfloat16)

516 + outs = []

517 + for state in (ub, exact):

518 + state.pack_indices(slots)

519 + outs.append(state.wrapper.run(q, q[..., :0], kv, kv[..., :0]))

520 + torch.testing.assert_close(outs[0], outs[1], atol=1e-2, rtol=1e-2)

521 +

522 +

523 def test_supports_combination_gates(monkeypatch, default_vllm_config):

524 monkeypatch.setattr(sm90_mod, "has_flashinfer_sm90_nope_mla", lambda: True)

525 call = lambda **kw: FlashInferMLASparseSM90Backend.supports_combination(

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

+200 -78

@@ -17,14 +17,14 @@ encoded by the indexer's selection, so ``causal=False``.

17 CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally,

18 so it must stay outside graph capture. Each metadata builder owns a wrapper,

19 reserved capture-stable device buffers, and the plan parameters. The wrapper

20 -bakes the per-row ``kv_len`` into its int schedule

21 -at plan() time — ``run()`` never reads the device-side buffer — so the

22 -metadata builder replans every step (outside capture) with exact host-side

23 -lengths derived from the batch's sequence lengths; a full-width schedule

24 -would send the kernel past each row's valid count into the -1 tail of the

25 -converted index buffer (illegal address). Per-step content (top-k slots)

26 -is written into the reserved buffers by kernels inside the captured

27 -forward, and captured runs read the refreshed plan buffers on replay.

20 +bakes the per-row ``kv_len`` into its int schedule at plan() time — ``run()``

21 +never reads the device-side buffer. The builder plans (outside capture) from

22 +sync-free host upper bounds and clamps each scheduled work item's ``kv_end``

23 +on device to the exact valid count, so the kernel never reads past a row's

24 +valid prefix (the -1 tail of the converted index buffer) and async scheduling

25 +needs no D2H sync. Per-step content (top-k slots) is written into the

26 +reserved buffers by kernels inside the captured forward, and captured runs

27 +read the refreshed plan buffers on replay.

29 KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]``

30 (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The

@@ -63,6 +63,26 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout

64 _FP8_KV_DTYPES = ("fp8", "fp8_e4m3")

65 _WORKSPACE_BYTES = 128 * 1024 * 1024

66 +# TODO: FlashInfer bakes per-row kv_len into the plan on the host, which

67 +# forces a D2H sync under async scheduling. The plan-info layout constants

68 +# below and the int-workspace clamp in _SM90State work around this by

69 +# planning from sync-free upper bounds; delete them once FlashInfer accepts

70 +# an upper-bound plan() that clamps kv_end internally, or exposes a

71 +# GPU-side plan API (upstream ask:

72 +# https://github.com/flashinfer-ai/flashinfer/issues/5559, see

73 +# vllm-project/vllm#58684).

74 +# FlashInfer MLAPlanInfo layout (scheduler.cuh) field indices.

75 +_PLAN_INFO_LEN = 18

76 +_PI_NUM_BLKS_Y = 1

77 +_PI_Q_INDPTR = 2

78 +_PI_KV_INDPTR = 3

79 +_PI_KV_START = 13

80 +_PI_KV_END = 14

81 +_PI_WORK_INDPTR = 15

82 +# Planned lengths exceed the host upper bound by up to this many tokens so a

83 +# plan stays reusable across draft steps and consecutive decode steps; the

84 +# device-side clamp restores the exact lengths.

85 +_PLAN_SLACK = 32

88 @triton.jit

@@ -82,6 +102,40 @@ def _pack_topk_indices(

102 tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)

103

104

105 +@triton.jit

106 +def _clamp_work_kv_end(

107 + int_ws,

108 + saved_kv_end,

109 + seq_lens,

110 + query_start_loc,

111 + req_id_per_token,

112 + num_works_idx,

113 + q_indptr_off,

114 + kv_start_off,

115 + kv_end_off,

116 + TOPK: tl.constexpr,

117 + KPOOL: tl.constexpr,

118 + BLOCK_SIZE: tl.constexpr,

119 +):

120 + # Clamp each planned work item to its row's exact valid count, derived

121 + # from the device batch layout exactly as the convert kernel counts it.

122 + works = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

123 + mask = works < tl.load(int_ws + num_works_idx)

124 + row = tl.load(int_ws + q_indptr_off + works, mask=mask, other=0)

125 + req = tl.load(req_id_per_token + row, mask=mask, other=0)

126 + seq_len = tl.load(seq_lens + req, mask=mask, other=0)

127 + query_end = tl.load(query_start_loc + req + 1, mask=mask, other=0)

128 + ctx = seq_len - query_end + row + 1

129 + valid = tl.where(ctx <= TOPK, ctx, TOPK + ctx % KPOOL)

130 + start = tl.load(int_ws + kv_start_off + works, mask=mask, other=0)

131 + end = tl.load(saved_kv_end + works, mask=mask, other=0)

132 + tl.store(

133 + int_ws + kv_end_off + works,

134 + tl.maximum(tl.minimum(end, valid), start),

135 + mask=mask,

136 + )

137 +

138 +

139 class FlashInferMLASparseSM90Backend(AttentionBackend):

140 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

141 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

@@ -190,6 +244,8 @@ class _SM90State:

244 kv_lora_rank: int,

245 qk_rope_head_dim: int,

246 sm_scale: float,

247 + index_topk: int,

248 + index_kpool: int,

249 ) -> None:

250 from flashinfer.mla import BatchMLAPagedAttentionWrapper

251

@@ -202,6 +258,10 @@ class _SM90State:

258 self.kv_lora_rank = kv_lora_rank

259 self.qk_rope_head_dim = qk_rope_head_dim

260 self.sm_scale = sm_scale

261 + self.index_topk = index_topk

262 + self.index_kpool = max(index_kpool, 1)

263 + # Largest valid count the convert kernel can produce for any context.

264 + self.max_valid = min(self.index_topk + self.index_kpool - 1, topk_width)

265 # User-reserved buffers: with use_cuda_graph=True plan() refreshes

266 # these in place, so run()'s captured kernels always read them.

267 self.kv_indices = torch.zeros(

@@ -226,24 +286,54 @@ class _SM90State:

286 self._lens_cpu = torch.full(

287 (self.max_tokens,), self.topk_width, dtype=torch.int32

288 )

229 -

230 - def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:

231 - """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``).

232 -

233 - The wrapper bakes kv_len into its int schedule from host values;

234 - run() never reads the device kv_len_arr buffer. Scheduling with the

235 - full buffer width while kv_indices rows carry a -1 tail past each

236 - row's valid count makes the kernel compute ``-1 * ckv_stride_page``

237 - (illegal address), so the lengths must be exact at plan time and

238 - replanned every step as contexts grow. Must run outside CUDA graph

239 - capture: the in-place refreshed plan_info/indptr buffers are what

240 - captured runs read.

289 + # Planned per-work kv_end of the current plan, restored by each clamp;

290 + # sized on first plan to the planner's work capacity.

291 + self._saved_kv_end = torch.empty(0, dtype=torch.int32, device=device)

292 + self._planned_num_tokens = -1

293 + self._int_ws: torch.Tensor | None = None

294 + self._plan_offsets: tuple[int, int, int, int] | None = None

295 +

296 + def plan(

297 + self,

298 + num_tokens: int,

299 + kv_lens: torch.Tensor,

300 + cam: CommonAttentionMetadata | None,

301 + req_id_per_token: torch.Tensor | None,

302 + ) -> None:

303 + """Plan per-row KV lengths (CPU int32, ``[num_tokens]``).

304 +

305 + The wrapper bakes kv_len into its int schedule from host values, so

306 + rows past their valid count would send the kernel into the -1 tail of

307 + the converted index buffer. With ``cam`` and ``req_id_per_token``

308 + (device batch layout), ``kv_lens`` only needs to upper-bound the

309 + valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while

310 + the bound still fits, and every call clamps each work item's kv_end

311 + on device to the exact count, so no D2H sync is needed. Without them

312 + the lengths must be exact. Must run outside CUDA graph capture: the

313 + in-place refreshed plan_info/indptr buffers are what captured runs

314 + read.

315 """

316 if torch.cuda.is_current_stream_capturing():

317 raise RuntimeError(

318 "FlashInferMLASparseSM90 plan() called inside CUDA graph "

319 "capture; lengths must be planned host-side before capture."

320 )

321 + kv_lens = kv_lens.to(torch.int32)

322 + if cam is None or req_id_per_token is None:

323 + self._plan(num_tokens, kv_lens)

324 + self._planned_num_tokens = -1

325 + return

326 + planned = self._lens_cpu[:num_tokens]

327 + if num_tokens != self._planned_num_tokens or not bool(

328 + ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all()

329 + ):

330 + self._plan(

331 + num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid)

332 + )

333 + self._planned_num_tokens = num_tokens

334 + self._clamp_kv_end(cam, req_id_per_token)

335 +

336 + def _plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:

337 # CPU staging buffers are filled in place: plan() runs per step

338 # (once per draft/verify metadata build), so per-call allocations and

339 # device round trips are on the hot path. Passing CPU tensors lets

@@ -257,13 +347,16 @@ class _SM90State:

347 # matching the exact lengths, including zero-length padded rows.

348 torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)

349 self._lens_cpu.zero_()

260 - self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)

350 + self._lens_cpu[:num_tokens] = kv_lens

351 self._kv_cpu[0] = 0

352 torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:])

353 + # Only the live prefix of kv_indices: the wrapper snapshots and

354 + # copies whatever is passed.

355 + num_indices = max(int(self._kv_cpu[-1]), 1)

356 self.wrapper.plan(

357 self._qo_cpu,

358 self._kv_cpu,

266 - self.kv_indices,

359 + self.kv_indices[:num_indices],

360 self._lens_cpu,

361 self.num_heads,

362 self.kv_lora_rank, # head_dim_ckv

@@ -274,6 +367,42 @@ class _SM90State:

367 q_data_type=torch.bfloat16,

368 kv_data_type=self.kv_dtype,

369 )

370 + plan_info = [int(x) for x in self.wrapper._plan_info]

371 + assert len(plan_info) == _PLAN_INFO_LEN, plan_info

372 + # Per-work arrays are laid out back to back at the planner's capacity.

373 + max_works = (plan_info[_PI_KV_INDPTR] - plan_info[_PI_Q_INDPTR]) // 4

374 + if self._saved_kv_end.numel() != max_works:

375 + self._saved_kv_end = torch.empty(

376 + max_works, dtype=torch.int32, device=self.device

377 + )

378 + self._int_ws = self.wrapper._planned_backend._int_workspace_buffer.view(

379 + torch.int32

380 + )

381 + self._plan_offsets = (

382 + plan_info[_PI_WORK_INDPTR] // 4 + plan_info[_PI_NUM_BLKS_Y],

383 + plan_info[_PI_Q_INDPTR] // 4,

384 + plan_info[_PI_KV_START] // 4,

385 + plan_info[_PI_KV_END] // 4,

386 + )

387 + kv_end_off = self._plan_offsets[3]

388 + self._saved_kv_end.copy_(self._int_ws[kv_end_off : kv_end_off + max_works])

389 +

390 + def _clamp_kv_end(

391 + self, cam: CommonAttentionMetadata, req_id_per_token: torch.Tensor

392 + ) -> None:

393 + assert self._int_ws is not None and self._plan_offsets is not None

394 + block = 1024

395 + _clamp_work_kv_end[(triton.cdiv(self._saved_kv_end.numel(), block),)](

396 + self._int_ws,

397 + self._saved_kv_end,

398 + cam.seq_lens,

399 + cam.query_start_loc,

400 + req_id_per_token,

401 + *self._plan_offsets,

402 + TOPK=self.index_topk,

403 + KPOOL=self.index_kpool,

404 + BLOCK_SIZE=block,

405 + )

406

407 def pack_indices(self, topk_slots: torch.Tensor) -> None:

408 num_tokens, width = topk_slots.shape

@@ -312,6 +441,8 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

441 )

442 topk_indices_buffer = impl.topk_indices_buffer

443 assert topk_indices_buffer is not None

444 + hf_config = vllm_config.model_config.hf_text_config

445 + assert hf_config.index_topk is not None

446 self.state = _SM90State(

447 device,

448 impl.num_heads,

@@ -321,20 +452,18 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

452 kv_lora_rank=impl.kv_lora_rank,

453 qk_rope_head_dim=impl.qk_rope_head_dim,

454 sm_scale=impl.scale,

455 + index_topk=int(hf_config.index_topk),

456 + index_kpool=int(getattr(hf_config, "index_kpool", 1) or 1),

457 + )

458 + spec_config = vllm_config.speculative_config

459 + self._adaptive_verification = bool(

460 + spec_config is not None and spec_config.enable_adaptive_verification

461 )

325 - # seq_lens_cpu_upper_bound is optimistic on decode rows under async

326 - # spec decode, so the fast sync-free path is only safe without it;

327 - # under async scheduling the exact (device) positions are used at

328 - # the cost of one D2H sync per metadata build.

329 - self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling)

330 - hf_config = vllm_config.model_config.hf_text_config

331 - assert hf_config.index_topk is not None

332 - self._index_topk = int(hf_config.index_topk)

333 - self._index_kpool = int(getattr(hf_config, "index_kpool", 1) or 1)

462

335 - def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:

336 - """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes

337 - them into its schedule at plan time; there is no device-side path).

463 + def _kv_lens_host(

464 + self, cam: CommonAttentionMetadata

465 + ) -> tuple[int, torch.Tensor, bool]:

466 + """Host per-row KV lengths for the plan and whether they are exact.

467

468 A row for the j-th query token of request i attends

469 ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection

@@ -343,53 +472,41 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

472 ``index_topk`` pool-expanded tokens plus the trailing incomplete

473 pool (valid == ``index_topk + context % index_kpool``). Both match

474 the count of non -1 entries the convert kernel produces.

475 +

476 + With ``seq_lens_cpu_upper_bound`` (optimistic under async spec

477 + decode) returns sync-free upper bounds, which the plan clamps on

478 + device. Otherwise returns exact lengths via a D2H copy.

479 """

480 num_reqs = cam.num_reqs

481 qsl = cam.query_start_loc_cpu[: num_reqs + 1]

482 num_rows = int(qsl[-1])

483 if num_rows == 0:

351 - return 0, torch.zeros(0, dtype=torch.int32)

352 - # Row context == position + 1. Without async scheduling the

353 - # maintained host upper bound is exact, giving a sync-free path;

354 - # under async scheduling it is optimistic on decode rows, so fall

355 - # back to the exact device positions (one D2H sync per build). The

356 - # seq_lens derivation equals position + 1 because positions are

357 - # contiguous per request.

484 + return 0, torch.zeros(0, dtype=torch.int32), True

485 + state = self.state

486 sl_host = cam.seq_lens_cpu_upper_bound

359 - positions = cam.positions

360 - if not self._async_scheduling and sl_host is not None:

361 - seq_lens = sl_host[:num_reqs].to(torch.int32)

362 - q_lens = qsl[1:] - qsl[:-1]

363 - first_pos = seq_lens - q_lens

364 - req_of_row = torch.repeat_interleave(

365 - torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)

366 - )

367 - rows = torch.arange(num_rows, dtype=torch.int32)

368 - ctx = (

369 - first_pos.to(torch.int64)[req_of_row]

370 - + rows.to(torch.int64)

371 - - qsl.to(torch.int64)[req_of_row]

372 - + 1

373 - )

374 - elif positions is not None and num_rows <= positions.shape[0]:

375 - ctx = positions[:num_rows].cpu().to(torch.int64) + 1

487 + if sl_host is None:

488 + ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu())

489 + topk, kpool = state.index_topk, state.index_kpool

490 + lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)

491 + return num_rows, lens.to(torch.int32), True

492 + if self._adaptive_verification:

493 + # The host query split may not match the device one, so bound

494 + # every row by the longest request.

495 + ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max()))

496 else:

377 - seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32)

378 - q_lens = qsl[1:] - qsl[:-1]

379 - first_pos = seq_lens - q_lens

380 - req_of_row = torch.repeat_interleave(

381 - torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)

382 - )

383 - rows = torch.arange(num_rows, dtype=torch.int32)

384 - ctx = (

385 - first_pos.to(torch.int64)[req_of_row]

386 - + rows.to(torch.int64)

387 - - qsl.to(torch.int64)[req_of_row]

388 - + 1

389 - )

390 - topk = self._index_topk

391 - lens = torch.where(ctx <= topk, ctx, topk + ctx % self._index_kpool)

392 - return num_rows, lens.to(torch.int32)

497 + ctx = self._row_contexts(qsl, sl_host[:num_reqs])

498 + return num_rows, torch.clamp(ctx, max=state.max_valid), False

499 +

500 + @staticmethod

501 + def _row_contexts(qsl: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor:

502 + num_reqs = qsl.shape[0] - 1

503 + q_lens = (qsl[1:] - qsl[:-1]).to(torch.int64)

504 + first_ctx = seq_lens.to(torch.int64) - q_lens + 1

505 + req_of_row = torch.repeat_interleave(

506 + torch.arange(num_reqs, dtype=torch.int64), q_lens

507 + )

508 + rows = torch.arange(int(qsl[-1]), dtype=torch.int64)

509 + return first_ctx[req_of_row] + rows - qsl.to(torch.int64)[req_of_row]

510

511 def build(

512 self,

@@ -399,9 +516,9 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

516 ) -> FlashInferMLASparseSM90Metadata:

517 metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)

518 assert isinstance(metadata, FlashInferMLASparseSM90Metadata)

402 - # Replan every step outside any CUDA graph capture with this step's

403 - # exact per-row lengths; captured runs read the refreshed buffers.

404 - num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)

519 + # Plan outside any CUDA graph capture; captured runs read the

520 + # refreshed buffers.

521 + num_rows, kv_lens, exact = self._kv_lens_host(common_attn_metadata)

522 # MHA prefills are removed from q before forward_mqa. Plan only the

523 # decode prefix in that case, using the same routing decision as MLA.

524 if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha(

@@ -409,7 +526,12 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):

526 ):

527 num_rows = metadata.num_decode_tokens

528 kv_lens = kv_lens[:num_rows]

412 - self.state.plan(num_rows, kv_lens)

529 + if exact:

530 + self.state.plan(num_rows, kv_lens, None, None)

531 + else:

532 + self.state.plan(

533 + num_rows, kv_lens, common_attn_metadata, metadata.req_id_per_token

534 + )

535 metadata.state = self.state

536 return metadata

537