iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
This commit is contained in:
Nicolò Lucchesi
2026-09-27 13:06:23 +02:00
committed by GitHub
parent 24c9772d19
commit c8d7a7dd13
2 changed files with 379 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
+181 -28
View File
@@ -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(
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
+198 -76
View File
@@ -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