|
|
@@ -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
|
|
|
@@ -64,6 +64,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
|
|
|
|
|
|
# GB10: a Triton sparse MLA kernel in place of FlashInfer's FA2 path, which
|
|
|
# runs at ~10 TFLOPS on SM121 (gb10_sparse_mla.py; about 3.3x faster).
|
|
|
@@ -87,6 +107,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]] = [
|
|
|
@@ -195,6 +249,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
|
|
|
|
|
|
@@ -207,6 +263,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(
|
|
|
@@ -231,24 +291,54 @@ class _SM90State:
|
|
|
self._lens_cpu = torch.full(
|
|
|
(self.max_tokens,), self.topk_width, dtype=torch.int32
|
|
|
)
|
|
|
-
|
|
|
- def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
|
|
|
- """Replan with exact 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.
|
|
|
+ # 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,
|
|
|
+ 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, 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
|
|
|
@@ -262,13 +352,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
|
|
|
@@ -279,6 +372,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
|
|
|
@@ -317,6 +446,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,
|
|
|
@@ -327,20 +458,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
|
|
|
@@ -349,53 +478,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,
|
|
|
@@ -411,9 +528,9 @@ class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
|
|
|
# scheduling (it stalls the CPU until the previous step ends).
|
|
|
metadata.state = self.state
|
|
|
return metadata
|
|
|
- # 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(
|
|
|
@@ -421,7 +538,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
|
|
|
|
