diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 3428625e96..b25656cf6e 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -9,6 +9,7 @@ causality), ckv/kpe cache splitting, and the backend's model-shape gates. """ from types import SimpleNamespace +from typing import Any import pytest import torch @@ -50,6 +51,8 @@ class FakeWrapper: def __init__(self): self.plan_args = None self.run_args = None + self._plan_info: Any = None + self._planned_backend: Any = None def plan(self, *args, **kwargs): self.plan_args = (args, kwargs) @@ -69,8 +72,11 @@ class FakeState: self.kv_indptr = torch.zeros(max_tokens + 1, dtype=torch.int32) self.wrapper = FakeWrapper() self.plan_calls = [] + self.index_topk = 2048 + self.index_kpool = 4 + self.max_valid = 2051 - def plan(self, num_tokens, kv_lens): + def plan(self, num_tokens, kv_lens, cam, req_id_per_token): self.plan_calls.append((num_tokens, kv_lens)) def pack_indices(self, slots): @@ -168,13 +174,12 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes): """MHA prefill rows must not make the MQA kernel read beyond its query.""" builder = object.__new__(FlashInferMLASparseSM90Builder) - builder._index_topk = 2048 - builder._index_kpool = 4 - builder._async_scheduling = False + builder._adaptive_verification = False builder.state = FakeState(TOPK) builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha) metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata) metadata.state = None + metadata.req_id_per_token = None metadata.num_prefills = 1 metadata.num_decode_tokens = num_decodes monkeypatch.setattr( @@ -206,12 +211,17 @@ def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_dec def test_plan_uses_state_params(monkeypatch): """The NoPE/rope dims and scale live on the builder state, not the layer. - plan() takes exact per-row KV lengths; the schedule is rebuilt on every - call (contexts grow between steps) and the indptrs are always full-size + Without a device batch layout plan() takes exact per-row KV lengths and + rebuilds the schedule on every call; the indptrs are always full-size with zero-query padding rows past num_tokens. """ - impl, rows = make_impl(64, "auto") wrapper = FakeWrapper() + plan_info = [0] * sm90_mod._PLAN_INFO_LEN + plan_info[sm90_mod._PI_KV_INDPTR] = 64 * 4 + wrapper._plan_info = torch.tensor(plan_info) + wrapper._planned_backend = SimpleNamespace( + _int_workspace_buffer=torch.zeros(64, dtype=torch.int32) + ) state = sm90_mod._SM90State.__new__(sm90_mod._SM90State) state.device = torch.device("cpu") state.wrapper = wrapper @@ -227,8 +237,9 @@ def test_plan_uses_state_params(monkeypatch): state._qo_cpu = torch.empty(5, dtype=torch.int32) state._kv_cpu = torch.empty(5, dtype=torch.int32) state._lens_cpu = torch.full((4,), TOPK, dtype=torch.int32) + state._saved_kv_end = torch.empty(0, dtype=torch.int32) - state.plan(3, torch.tensor([2, 5, 7], dtype=torch.int32)) + state.plan(3, torch.tensor([2, 5, 7], dtype=torch.int32), None, None) assert wrapper.plan_args is not None args, kwargs = wrapper.plan_args (qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args @@ -241,10 +252,10 @@ def test_plan_uses_state_params(monkeypatch): assert kwargs["kv_data_type"] == torch.bfloat16 # Replanning a smaller batch must clear the previous rows' lengths. - state.plan(1, torch.tensor([TOPK], dtype=torch.int32)) + state.plan(1, torch.tensor([TOPK], dtype=torch.int32), None, None) assert state._kv_cpu.tolist() == [0, TOPK, TOPK, TOPK, TOPK] assert state._lens_cpu.tolist() == [TOPK, 0, 0, 0] - state.plan(0, torch.empty(0, dtype=torch.int32)) + state.plan(0, torch.empty(0, dtype=torch.int32), None, None) assert state._kv_cpu.tolist() == [0, 0, 0, 0, 0] assert state._lens_cpu.tolist() == [0, 0, 0, 0] @@ -279,38 +290,49 @@ def test_pack_indices_replays_with_updated_offsets(width): assert (state.kv_indices[expected.numel() :] == -99).all() -def test_kv_lens_host_formula(): +@pytest.mark.parametrize("adaptive", [False, True]) +def test_kv_lens_host_formula(adaptive): """Per-row host lengths: context == position + 1; capped at - index_topk + trailing-pool remainder past the sparse threshold.""" + index_topk + trailing-pool remainder past the sparse threshold. + + Exact (synced) without a host upper bound; with one, a sync-free bound + capped at the largest valid count, which the plan clamps on device. + """ builder = object.__new__(FlashInferMLASparseSM90Builder) - builder._index_topk = 2048 - builder._index_kpool = 4 - builder._async_scheduling = False + builder.state = FakeState(TOPK) + builder._adaptive_verification = adaptive + seq_lens = torch.tensor([100, 9, 3000], dtype=torch.int32) cam = SimpleNamespace( num_reqs=3, query_start_loc_cpu=torch.tensor([0, 5, 7, 10], dtype=torch.int32), - seq_lens=torch.tensor([100, 9, 3000], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([100, 9, 3000], dtype=torch.int32), - positions=None, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=None, ) - num_rows, lens = builder._kv_lens_host(cam) - assert num_rows == 10 + num_rows, lens, exact = builder._kv_lens_host(cam) + assert num_rows == 10 and exact # req0: positions 95..99 -> ctx 96..100 (all <= 2048: full context) # req1: positions 7,8 -> ctx 8,9 # req2: positions 2997..2999 -> ctx 2998..3000 (> 2048: topk + ctx%4) assert lens.tolist() == [96, 97, 98, 99, 100, 8, 9, 2050, 2051, 2048] + cam.seq_lens_cpu_upper_bound = seq_lens + num_rows, lens, exact = builder._kv_lens_host(cam) + assert num_rows == 10 and not exact + if adaptive: + # Host query split may differ from the device one: batch-wide bound. + assert lens.tolist() == [2051] * 10 + else: + assert lens.tolist() == [96, 97, 98, 99, 100, 8, 9, 2051, 2051, 2051] + def test_kv_lens_host_empty(): builder = object.__new__(FlashInferMLASparseSM90Builder) - builder._index_topk = 2048 - builder._index_kpool = 4 cam = SimpleNamespace( num_reqs=0, query_start_loc_cpu=torch.tensor([0], dtype=torch.int32), seq_lens=torch.zeros(0, dtype=torch.int32), ) - num_rows, lens = builder._kv_lens_host(cam) + num_rows, lens, _ = builder._kv_lens_host(cam) assert num_rows == 0 and lens.numel() == 0 @@ -334,7 +356,13 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens) "__init__", lambda *_args, **_kwargs: None, ) - monkeypatch.setattr(sm90_mod, "_SM90State", lambda *_args, **_kwargs: None) + + class _RecordingState: + def __init__(self, *_args, **kwargs): + self.index_topk = kwargs["index_topk"] + self.index_kpool = kwargs["index_kpool"] + + monkeypatch.setattr(sm90_mod, "_SM90State", _RecordingState) impl, _ = make_impl(64) spec = MLAAttentionSpec( 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) max_num_batched_tokens=64, async_scheduling=False ), model_config=SimpleNamespace(hf_text_config=hf_config), + speculative_config=None, ) builder = FlashInferMLASparseSM90Builder( spec, ["attn"], vllm_config, torch.device("cpu") @@ -359,14 +388,138 @@ def test_builder_kpool_from_model_config(monkeypatch, index_kpool, prefill_lens) cam = SimpleNamespace( num_reqs=2, query_start_loc_cpu=torch.tensor([0, 4, 6], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([42295, 100], dtype=torch.int32), - positions=None, + seq_lens=torch.tensor([42295, 100], dtype=torch.int32), + seq_lens_cpu_upper_bound=None, ) - num_rows, lens = builder._kv_lens_host(cam) - assert num_rows == 6 + num_rows, lens, exact = builder._kv_lens_host(cam) + assert num_rows == 6 and exact assert lens.tolist() == prefill_lens + [99, 100] +def _sm90_flashinfer_available(): + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + and sm90_mod.has_flashinfer_sm90_nope_mla() + ) + + +SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH = 16, 2048, 4, 2176 + + +def _real_state(rows): + return sm90_mod._SM90State( + torch.device("cuda"), + SM90_HEADS, + torch.bfloat16, + rows, + SM90_WIDTH, + HEAD, + 0, + HEAD**-0.5, + SM90_TOPK, + SM90_KPOOL, + ) + + +def _clamp_batch(num_reqs, q_per_req, lo, hi, topk, kpool, width): + """Random batch: per-row contexts, device layout, and top-k slots.""" + last = torch.randint(lo + q_per_req, hi, (num_reqs,)) + ctx = (last[:, None] + torch.arange(1 - q_per_req, 1)).reshape(-1) + valid = torch.where(ctx <= topk, ctx, topk + ctx % kpool) + cam = SimpleNamespace( + seq_lens=last.to(torch.int32).cuda(), + query_start_loc=(torch.arange(num_reqs + 1) * q_per_req).to(torch.int32).cuda(), + ) + req_id = torch.arange(num_reqs, dtype=torch.int32).repeat_interleave(q_per_req) + slots = torch.randint(0, 1 << 16, (ctx.numel(), width), dtype=torch.int32) + slots[torch.arange(width) >= valid[:, None]] = -1 + return ctx, valid.to(torch.int32), cam, req_id.cuda(), slots.cuda() + + +@pytest.mark.skipif( + not _sm90_flashinfer_available(), reason="Requires SM90 FlashInfer MLA" +) +@pytest.mark.parametrize( + "num_reqs,q_per_req,lo,hi", + [ + (64, 1, 1, 3000), + (16, 6, 2030, 2070), # MTP-style rows straddling index_topk + (4096, 1, 1, 3000), # clamp spans several program blocks + ], +) +def test_upper_bound_plan_matches_exact(num_reqs, q_per_req, lo, hi): + """Sync-free planning from optimistic host bounds (async spec decode) + must match exact planning: the device clamp trims every work item back + to its row's valid count, so no -1 slot is read.""" + torch.manual_seed(0) + heads, topk, kpool, width = SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH + ctx, valid, cam, req_id, slots = _clamp_batch( + num_reqs, q_per_req, lo, hi, topk, kpool, width + ) + rows = ctx.numel() + + kv = torch.randn(1 << 16, 1, HEAD, device="cuda", dtype=torch.bfloat16) + q = torch.randn(rows, heads, HEAD, device="cuda", dtype=torch.bfloat16) + q_pe = q.new_empty(rows, heads, 0) + + def run(state): + state.pack_indices(slots) + return state.wrapper.run(q, q_pe, kv, kv[..., :0]) + + exact = _real_state(rows) + exact.plan(rows, valid, None, None) + ub = _real_state(rows) + host_bound = torch.clamp(ctx + torch.randint(0, 7, (rows,)), max=ub.max_valid) + ub.plan(rows, host_bound, cam, req_id) + out = run(ub) + assert not out.isnan().any() + torch.testing.assert_close(out, run(exact), atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not _sm90_flashinfer_available(), reason="Requires SM90 FlashInfer MLA" +) +def test_upper_bound_plan_reused_across_steps(monkeypatch): + """Growing contexts (draft/decode steps) reuse the padded plan and only + re-clamp on device, matching an exact replan every step.""" + torch.manual_seed(0) + heads, topk, kpool, width = SM90_HEADS, SM90_TOPK, SM90_KPOOL, SM90_WIDTH + rows = 48 + ctx, _, _, req_id, _ = _clamp_batch(rows, 1, 1, 3000, topk, kpool, width) + + exact, ub = _real_state(rows), _real_state(rows) + replans: list[int] = [] + plan_fn = ub._plan + + def counting_plan(num_tokens, kv_lens): + replans.append(num_tokens) + plan_fn(num_tokens, kv_lens) + + monkeypatch.setattr(ub, "_plan", counting_plan) + kv = torch.randn(1 << 16, 1, HEAD, device="cuda", dtype=torch.bfloat16) + for step in range(5): + step_ctx = ctx + step + valid = torch.where(step_ctx <= topk, step_ctx, topk + step_ctx % kpool) + slots = torch.randint(0, 1 << 16, (rows, width), dtype=torch.int32) + slots[torch.arange(width) >= valid[:, None]] = -1 + slots = slots.cuda() + cam = SimpleNamespace( + seq_lens=step_ctx.to(torch.int32).cuda(), + query_start_loc=torch.arange(rows + 1, dtype=torch.int32).cuda(), + ) + num_plans = len(replans) + ub.plan(rows, torch.clamp(step_ctx + 5, max=ub.max_valid), cam, req_id) + assert len(replans) - num_plans == (1 if step == 0 else 0) + exact.plan(rows, valid.to(torch.int32), None, None) + q = torch.randn(rows, heads, HEAD, device="cuda", dtype=torch.bfloat16) + outs = [] + for state in (ub, exact): + state.pack_indices(slots) + outs.append(state.wrapper.run(q, q[..., :0], kv, kv[..., :0])) + torch.testing.assert_close(outs[0], outs[1], atol=1e-2, rtol=1e-2) + + def test_supports_combination_gates(monkeypatch, default_vllm_config): monkeypatch.setattr(sm90_mod, "has_flashinfer_sm90_nope_mla", lambda: True) call = lambda **kw: FlashInferMLASparseSM90Backend.supports_combination( diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 9dc899c903..a5f7757d38 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -17,14 +17,14 @@ encoded by the indexer's selection, so ``causal=False``. CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally, so it must stay outside graph capture. Each metadata builder owns a wrapper, reserved capture-stable device buffers, and the plan parameters. The wrapper -bakes the per-row ``kv_len`` into its int schedule -at plan() time — ``run()`` never reads the device-side buffer — so the -metadata builder replans every step (outside capture) with exact host-side -lengths derived from the batch's sequence lengths; a full-width schedule -would send the kernel past each row's valid count into the -1 tail of the -converted index buffer (illegal address). Per-step content (top-k slots) -is written into the reserved buffers by kernels inside the captured -forward, and captured runs read the refreshed plan buffers on replay. +bakes the per-row ``kv_len`` into its int schedule at plan() time — ``run()`` +never reads the device-side buffer. The builder plans (outside capture) from +sync-free host upper bounds and clamps each scheduled work item's ``kv_end`` +on device to the exact valid count, so the kernel never reads past a row's +valid prefix (the -1 tail of the converted index buffer) and async scheduling +needs no D2H sync. Per-step content (top-k slots) is written into the +reserved buffers by kernels inside the captured forward, and captured runs +read the refreshed plan buffers on replay. KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]`` (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 _FP8_KV_DTYPES = ("fp8", "fp8_e4m3") _WORKSPACE_BYTES = 128 * 1024 * 1024 +# TODO: FlashInfer bakes per-row kv_len into the plan on the host, which +# forces a D2H sync under async scheduling. The plan-info layout constants +# below and the int-workspace clamp in _SM90State work around this by +# planning from sync-free upper bounds; delete them once FlashInfer accepts +# an upper-bound plan() that clamps kv_end internally, or exposes a +# GPU-side plan API (upstream ask: +# https://github.com/flashinfer-ai/flashinfer/issues/5559, see +# vllm-project/vllm#58684). +# FlashInfer MLAPlanInfo layout (scheduler.cuh) field indices. +_PLAN_INFO_LEN = 18 +_PI_NUM_BLKS_Y = 1 +_PI_Q_INDPTR = 2 +_PI_KV_INDPTR = 3 +_PI_KV_START = 13 +_PI_KV_END = 14 +_PI_WORK_INDPTR = 15 +# Planned lengths exceed the host upper bound by up to this many tokens so a +# plan stays reusable across draft steps and consecutive decode steps; the +# device-side clamp restores the exact lengths. +_PLAN_SLACK = 32 @triton.jit @@ -82,6 +102,40 @@ def _pack_topk_indices( tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask) +@triton.jit +def _clamp_work_kv_end( + int_ws, + saved_kv_end, + seq_lens, + query_start_loc, + req_id_per_token, + num_works_idx, + q_indptr_off, + kv_start_off, + kv_end_off, + TOPK: tl.constexpr, + KPOOL: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # Clamp each planned work item to its row's exact valid count, derived + # from the device batch layout exactly as the convert kernel counts it. + works = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = works < tl.load(int_ws + num_works_idx) + row = tl.load(int_ws + q_indptr_off + works, mask=mask, other=0) + req = tl.load(req_id_per_token + row, mask=mask, other=0) + seq_len = tl.load(seq_lens + req, mask=mask, other=0) + query_end = tl.load(query_start_loc + req + 1, mask=mask, other=0) + ctx = seq_len - query_end + row + 1 + valid = tl.where(ctx <= TOPK, ctx, TOPK + ctx % KPOOL) + start = tl.load(int_ws + kv_start_off + works, mask=mask, other=0) + end = tl.load(saved_kv_end + works, mask=mask, other=0) + tl.store( + int_ws + kv_end_off + works, + tl.maximum(tl.minimum(end, valid), start), + mask=mask, + ) + + class FlashInferMLASparseSM90Backend(AttentionBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -190,6 +244,8 @@ class _SM90State: kv_lora_rank: int, qk_rope_head_dim: int, sm_scale: float, + index_topk: int, + index_kpool: int, ) -> None: from flashinfer.mla import BatchMLAPagedAttentionWrapper @@ -202,6 +258,10 @@ class _SM90State: self.kv_lora_rank = kv_lora_rank self.qk_rope_head_dim = qk_rope_head_dim self.sm_scale = sm_scale + self.index_topk = index_topk + self.index_kpool = max(index_kpool, 1) + # Largest valid count the convert kernel can produce for any context. + self.max_valid = min(self.index_topk + self.index_kpool - 1, topk_width) # User-reserved buffers: with use_cuda_graph=True plan() refreshes # these in place, so run()'s captured kernels always read them. self.kv_indices = torch.zeros( @@ -226,24 +286,54 @@ class _SM90State: self._lens_cpu = torch.full( (self.max_tokens,), self.topk_width, dtype=torch.int32 ) + # Planned per-work kv_end of the current plan, restored by each clamp; + # sized on first plan to the planner's work capacity. + self._saved_kv_end = torch.empty(0, dtype=torch.int32, device=device) + self._planned_num_tokens = -1 + self._int_ws: torch.Tensor | None = None + self._plan_offsets: tuple[int, int, int, int] | None = None - def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None: - """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``). + def plan( + self, + num_tokens: int, + kv_lens: torch.Tensor, + cam: CommonAttentionMetadata | None, + req_id_per_token: torch.Tensor | None, + ) -> None: + """Plan per-row KV lengths (CPU int32, ``[num_tokens]``). - The wrapper bakes kv_len into its int schedule from host values; - run() never reads the device kv_len_arr buffer. Scheduling with the - full buffer width while kv_indices rows carry a -1 tail past each - row's valid count makes the kernel compute ``-1 * ckv_stride_page`` - (illegal address), so the lengths must be exact at plan time and - replanned every step as contexts grow. Must run outside CUDA graph - capture: the in-place refreshed plan_info/indptr buffers are what - captured runs read. + The wrapper bakes kv_len into its int schedule from host values, so + rows past their valid count would send the kernel into the -1 tail of + the converted index buffer. With ``cam`` and ``req_id_per_token`` + (device batch layout), ``kv_lens`` only needs to upper-bound the + valid counts: the plan is padded by ``_PLAN_SLACK`` and reused while + the bound still fits, and every call clamps each work item's kv_end + on device to the exact count, so no D2H sync is needed. Without them + the lengths must be exact. Must run outside CUDA graph capture: the + in-place refreshed plan_info/indptr buffers are what captured runs + read. """ if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "FlashInferMLASparseSM90 plan() called inside CUDA graph " "capture; lengths must be planned host-side before capture." ) + kv_lens = kv_lens.to(torch.int32) + if cam is None or req_id_per_token is None: + self._plan(num_tokens, kv_lens) + self._planned_num_tokens = -1 + return + planned = self._lens_cpu[:num_tokens] + if num_tokens != self._planned_num_tokens or not bool( + ((planned >= kv_lens) & (planned <= kv_lens + _PLAN_SLACK)).all() + ): + self._plan( + num_tokens, torch.clamp(kv_lens + _PLAN_SLACK, max=self.max_valid) + ) + self._planned_num_tokens = num_tokens + self._clamp_kv_end(cam, req_id_per_token) + + def _plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None: # CPU staging buffers are filled in place: plan() runs per step # (once per draft/verify metadata build), so per-call allocations and # device round trips are on the hot path. Passing CPU tensors lets @@ -257,13 +347,16 @@ class _SM90State: # matching the exact lengths, including zero-length padded rows. torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu) self._lens_cpu.zero_() - self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32) + self._lens_cpu[:num_tokens] = kv_lens self._kv_cpu[0] = 0 torch.cumsum(self._lens_cpu, dim=0, out=self._kv_cpu[1:]) + # Only the live prefix of kv_indices: the wrapper snapshots and + # copies whatever is passed. + num_indices = max(int(self._kv_cpu[-1]), 1) self.wrapper.plan( self._qo_cpu, self._kv_cpu, - self.kv_indices, + self.kv_indices[:num_indices], self._lens_cpu, self.num_heads, self.kv_lora_rank, # head_dim_ckv @@ -274,6 +367,42 @@ class _SM90State: q_data_type=torch.bfloat16, kv_data_type=self.kv_dtype, ) + plan_info = [int(x) for x in self.wrapper._plan_info] + assert len(plan_info) == _PLAN_INFO_LEN, plan_info + # Per-work arrays are laid out back to back at the planner's capacity. + max_works = (plan_info[_PI_KV_INDPTR] - plan_info[_PI_Q_INDPTR]) // 4 + if self._saved_kv_end.numel() != max_works: + self._saved_kv_end = torch.empty( + max_works, dtype=torch.int32, device=self.device + ) + self._int_ws = self.wrapper._planned_backend._int_workspace_buffer.view( + torch.int32 + ) + self._plan_offsets = ( + plan_info[_PI_WORK_INDPTR] // 4 + plan_info[_PI_NUM_BLKS_Y], + plan_info[_PI_Q_INDPTR] // 4, + plan_info[_PI_KV_START] // 4, + plan_info[_PI_KV_END] // 4, + ) + kv_end_off = self._plan_offsets[3] + self._saved_kv_end.copy_(self._int_ws[kv_end_off : kv_end_off + max_works]) + + def _clamp_kv_end( + self, cam: CommonAttentionMetadata, req_id_per_token: torch.Tensor + ) -> None: + assert self._int_ws is not None and self._plan_offsets is not None + block = 1024 + _clamp_work_kv_end[(triton.cdiv(self._saved_kv_end.numel(), block),)]( + self._int_ws, + self._saved_kv_end, + cam.seq_lens, + cam.query_start_loc, + req_id_per_token, + *self._plan_offsets, + TOPK=self.index_topk, + KPOOL=self.index_kpool, + BLOCK_SIZE=block, + ) def pack_indices(self, topk_slots: torch.Tensor) -> None: num_tokens, width = topk_slots.shape @@ -312,6 +441,8 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): ) topk_indices_buffer = impl.topk_indices_buffer assert topk_indices_buffer is not None + hf_config = vllm_config.model_config.hf_text_config + assert hf_config.index_topk is not None self.state = _SM90State( device, impl.num_heads, @@ -321,20 +452,18 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): kv_lora_rank=impl.kv_lora_rank, qk_rope_head_dim=impl.qk_rope_head_dim, sm_scale=impl.scale, + index_topk=int(hf_config.index_topk), + index_kpool=int(getattr(hf_config, "index_kpool", 1) or 1), + ) + spec_config = vllm_config.speculative_config + self._adaptive_verification = bool( + spec_config is not None and spec_config.enable_adaptive_verification ) - # seq_lens_cpu_upper_bound is optimistic on decode rows under async - # spec decode, so the fast sync-free path is only safe without it; - # under async scheduling the exact (device) positions are used at - # the cost of one D2H sync per metadata build. - self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling) - hf_config = vllm_config.model_config.hf_text_config - assert hf_config.index_topk is not None - self._index_topk = int(hf_config.index_topk) - self._index_kpool = int(getattr(hf_config, "index_kpool", 1) or 1) - def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]: - """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes - them into its schedule at plan time; there is no device-side path). + def _kv_lens_host( + self, cam: CommonAttentionMetadata + ) -> tuple[int, torch.Tensor, bool]: + """Host per-row KV lengths for the plan and whether they are exact. A row for the j-th query token of request i attends ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection @@ -343,53 +472,41 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): ``index_topk`` pool-expanded tokens plus the trailing incomplete pool (valid == ``index_topk + context % index_kpool``). Both match the count of non -1 entries the convert kernel produces. + + With ``seq_lens_cpu_upper_bound`` (optimistic under async spec + decode) returns sync-free upper bounds, which the plan clamps on + device. Otherwise returns exact lengths via a D2H copy. """ num_reqs = cam.num_reqs qsl = cam.query_start_loc_cpu[: num_reqs + 1] num_rows = int(qsl[-1]) if num_rows == 0: - return 0, torch.zeros(0, dtype=torch.int32) - # Row context == position + 1. Without async scheduling the - # maintained host upper bound is exact, giving a sync-free path; - # under async scheduling it is optimistic on decode rows, so fall - # back to the exact device positions (one D2H sync per build). The - # seq_lens derivation equals position + 1 because positions are - # contiguous per request. + return 0, torch.zeros(0, dtype=torch.int32), True + state = self.state sl_host = cam.seq_lens_cpu_upper_bound - positions = cam.positions - if not self._async_scheduling and sl_host is not None: - seq_lens = sl_host[:num_reqs].to(torch.int32) - q_lens = qsl[1:] - qsl[:-1] - first_pos = seq_lens - q_lens - req_of_row = torch.repeat_interleave( - torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64) - ) - rows = torch.arange(num_rows, dtype=torch.int32) - ctx = ( - first_pos.to(torch.int64)[req_of_row] - + rows.to(torch.int64) - - qsl.to(torch.int64)[req_of_row] - + 1 - ) - elif positions is not None and num_rows <= positions.shape[0]: - ctx = positions[:num_rows].cpu().to(torch.int64) + 1 + if sl_host is None: + ctx = self._row_contexts(qsl, cam.seq_lens[:num_reqs].cpu()) + topk, kpool = state.index_topk, state.index_kpool + lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool) + return num_rows, lens.to(torch.int32), True + if self._adaptive_verification: + # The host query split may not match the device one, so bound + # every row by the longest request. + ctx = torch.full((num_rows,), int(sl_host[:num_reqs].max())) else: - seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32) - q_lens = qsl[1:] - qsl[:-1] - first_pos = seq_lens - q_lens - req_of_row = torch.repeat_interleave( - torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64) - ) - rows = torch.arange(num_rows, dtype=torch.int32) - ctx = ( - first_pos.to(torch.int64)[req_of_row] - + rows.to(torch.int64) - - qsl.to(torch.int64)[req_of_row] - + 1 - ) - topk = self._index_topk - lens = torch.where(ctx <= topk, ctx, topk + ctx % self._index_kpool) - return num_rows, lens.to(torch.int32) + ctx = self._row_contexts(qsl, sl_host[:num_reqs]) + return num_rows, torch.clamp(ctx, max=state.max_valid), False + + @staticmethod + def _row_contexts(qsl: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor: + num_reqs = qsl.shape[0] - 1 + q_lens = (qsl[1:] - qsl[:-1]).to(torch.int64) + first_ctx = seq_lens.to(torch.int64) - q_lens + 1 + req_of_row = torch.repeat_interleave( + torch.arange(num_reqs, dtype=torch.int64), q_lens + ) + rows = torch.arange(int(qsl[-1]), dtype=torch.int64) + return first_ctx[req_of_row] + rows - qsl.to(torch.int64)[req_of_row] def build( self, @@ -399,9 +516,9 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): ) -> FlashInferMLASparseSM90Metadata: metadata = super().build(common_prefix_len, common_attn_metadata, fast_build) assert isinstance(metadata, FlashInferMLASparseSM90Metadata) - # Replan every step outside any CUDA graph capture with this step's - # exact per-row lengths; captured runs read the refreshed buffers. - num_rows, kv_lens = self._kv_lens_host(common_attn_metadata) + # Plan outside any CUDA graph capture; captured runs read the + # refreshed buffers. + num_rows, kv_lens, exact = self._kv_lens_host(common_attn_metadata) # MHA prefills are removed from q before forward_mqa. Plan only the # decode prefix in that case, using the same routing decision as MLA. if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( @@ -409,7 +526,12 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder): ): num_rows = metadata.num_decode_tokens kv_lens = kv_lens[:num_rows] - self.state.plan(num_rows, kv_lens) + if exact: + self.state.plan(num_rows, kv_lens, None, None) + else: + self.state.plan( + num_rows, kv_lens, common_attn_metadata, metadata.req_id_per_token + ) metadata.state = self.state return metadata