[XinyuanQiao]

Overview

FLASH_ATTN_MLA_SPARSE_FA4 sends sparse MQA prefill to FlashInfer's trtllm-gen kernel, one query at a time. This PR runs it on FA4 instead, with 2, 4 or 8 consecutive queries sharing each 128-head kernel row (the grouped mode of vllm-project/flash-attention#205) and a Triton kernel building each group's top-k union. It covers both prefill layouts:

• ShardQ steps ([Perf][DSA] Token-sharded sparse-MLA prefill under TP; remove indexer row sharding #63956, on by default under TP): each rank gets a slice of the tokens with every head, token-major, so 128 / heads consecutive tokens are already one packed row. GLM-5.x (64 heads) groups 2 per row at any TP size, including TP8.
• Steps without ShardQ (ShardQ off, TP1, FULL cudagraph mixed mode): 128 / heads per rank consecutive queries of a request per row, as before.

Depends on vllm-project/flash-attention#205 (approved). cmake/external_projects/vllm_flash_attn.cmake pins e7352c9, which is #205 merged onto d8dca8f (main's current pin) and differs from d8dca8f only in #205's four files. flash-attention main has since taken the Dao-AILab sync (#213), whose _flash_attn_fwd returns five values where vLLM main's wrapper unpacks four. Once #205 is merged and vLLM moves to a flash-attention commit with it, the pin goes back to vllm-project/flash-attention.

Related: #48726 (LiteDSA, under review) groups neighbouring queries in the same way for the FP8 path in the FlashInfer sparse backend; the two PRs touch no common files. The same idea is merged in SGLang for SM90 and SM120 (sgl-project/sglang#32779). #63328 adds an FP8 KV cache on top of this PR.

Claims

• Against main's default configuration (ShardQ on), prefill is 3.5-4.4% faster for pure prefills of 8K-46K tokens and 1.7-3.1% faster for 8K new tokens on 16-24K cached ones (nvidia/GLM-5.2-NVFP4, TP4, B300, BF16 KV). Steps that ShardQ does not shard (fewer than 1024 prefill tokens) take main's path and are unchanged.
• With ShardQ off, sparse MQA prefill is 1.33-1.46x faster end to end (46K tokens: 2.55 → 1.75 s).
• No accuracy change: perplexity over 46K teacher-forced tokens moves within the run-to-run noise of either server.
• Decode, and the configurations grouping does not cover (DCP, PCP, HiSparse, or heads that do not make 128 / heads queries per row), keep their current kernels.

Validation

All runs: B300. This branch (merged with main 86fdabe) is installed editable over main's precompiled wheel for d60c440, which has no csrc/ change up to 86fdabe, with the flash_attn/cute files of the pinned flash-attention commit copied in as the cmake build does. The baseline is main 86fdabe installed the same way, with flash-attention d8dca8f. I did not compile this branch from source.

Both servers ran with --kernel-config '{"enable_flashinfer_autotune": false}'. With FlashInfer autotuning on, main's server crashed at startup in my environment (an illegal memory access during the autotune dummy runs, whenever their prefill went through trtllm-gen, e.g. under ShardQ), which I have not investigated further.

Tests

pytest tests/v1/attention/test_sparse_mla_backends.py
# 551 passed, 1432 skipped (split over 8 GPUs)

• test_build_query_group_row_union_matches_reference[2, 4]: the ShardQ union builder against a Python reference, with an empty group, rows past each query's count, negative rows and rows up to the 24-bit limit.
• test_fa4_shardq_prefill_groups_rank_tokens[32, 64 heads]: every rank of a 4-way ShardQ split runs the real forward_mqa with narrowed metadata, in calls of 2 rows, with tails; together the ranks match an fp32 reference of the whole step.
• test_fa4_grouped_prefill_warmup_covers_runtime_keys: after the compile-only warmup, grouped calls on both sides of FA4's 256-row scheduler threshold, with and without ShardQ, compile nothing (cute.compile raises inside them).
• test_build_query_group_union_matches_reference, test_fa4_sparse_grouped_prefill_correctness[16, 32, 64] and test_fa4_sparse_grouped_prefill_sub_chunks[16, 64] cover the path without ShardQ, as before.

Dropping the second query's membership bit, or the tail's output, fails the new tests; I checked both. The grouped tests need SM100 and an FA build with flash-attention#205, so CI covers them once the pin includes it.

End-to-end prefill latency. Baseline main on GPUs 0-3 and this PR on GPUs 4-7 of one machine, with:

vllm serve nvidia/GLM-5.2-NVFP4 --served-model-name glm52 --tensor-parallel-size 4 --max-model-len 65536 \
  --max-num-batched-tokens 65536 --max-num-seqs 16 --gpu-memory-utilization 0.85 \
  --attention-backend FLASH_ATTN_MLA_SPARSE_FA4 --kv-cache-dtype bfloat16 --enable-prefix-caching \
  --kernel-config '{"enable_flashinfer_autotune": false}'
# right column: both servers with --attention-config '{"sparse_mla_shard_q": false}'
python3 e2e_prefill.py --urls http://127.0.0.1:8400 http://127.0.0.1:8500 --trials 5 --out r.json

Each request asks for one token on public-domain text (Project Gutenberg #1342 and #2701). Each cell is the median of 5 trials, with the servers alternating, from the second of two runs: in the first, both servers were up to 3x slower on prompt lengths they had not seen yet, including on paths this PR does not touch. Every prompt starts with its own salt, so no request hits an earlier one's prefix cache; for the cached rows the prefix is sent once, untimed, first.

Prompt (tokens)	Default (ShardQ on): main → this PR	ShardQ off: main → this PR
7.7K new	0.328 → 0.316 s (1.04x)	0.267 → 0.266 s (masked MHA, unchanged)
15.4K new	0.597 → 0.575 s (1.04x)	0.544 → 0.539 s (masked MHA, unchanged)
33.6K new	1.254 → 1.211 s (1.04x)	at the masked-MHA threshold
46.3K new	1.860 → 1.782 s (1.04x)	2.552 → 1.750 s (1.46x)
16K cached + 8K new	0.363 → 0.352 s (1.03x)	0.468 → 0.337 s (1.39x)
24K cached + 8K new	0.395 → 0.389 s (1.02x)	0.512 → 0.385 s (1.33x)
32K cached + 256 new (15 trials)	0.109 → 0.109 s (main's path)
32K cached + 900 new (15 trials)	0.141 → 0.145 s (main's path; the interquartile ranges overlap)

The ShardQ-off column ran with FlashInfer autotuning on, on the commit before this PR's ShardQ path; the path it measures is unchanged since. Under ShardQ, main routes every sharded step to MQA, so short prompts there do not use masked MHA.

Accuracy. Default configuration, so the ~8K-token prompts go through ShardQ and the grouped path. Teacher-forced prompt_logprobs over six public-text prompts, two runs per server:

python3 prompt_logprob_probe.py --base-url http://127.0.0.1:8400 --cache-salt r1 --out base_r1.json  # also r2, and :8500
python3 prompt_logprob_probe.py --compare base_r1.json new_r1.json

	main	This PR
Perplexity on 46,070 positions, runs 1 / 2	1.1548 / 1.1527	1.1561 / 1.1570
Top-1 agreement, run 1 vs run 2 on the same server	98.02%	97.78%
Top-1 agreement with the other server (runs 1, runs 2)	97.11%, 97.10%

Agreement between the servers is a little below each server's run-to-run agreement, as expected from kernels that round differently. Before the main merge, the path without ShardQ also kept GSM8K 5-shot unchanged with every prefill forced through it (0.9363 ± 0.0067 → 0.9386 ± 0.0066).

Where a ShardQ rank's time goes. One layer on one rank of a 4-way split of a ~15K-token prompt (real GLM-5.2 top-k, ~3.7K tokens per rank, L2 flushed, bench_shardq_rank.py):

	Early and late chunk (rank 0)	Middle chunks (rank 3)
trtllm-gen, 64 heads (main)	1,102-1,151 µs	1,425-1,588 µs
This PR: union + FA4	863-970 µs (1.2-1.3x)	1,053-1,166 µs (1.3-1.4x)
of which the union	96-99 µs	110-112 µs

At 64 heads FA4's grouped kernel alone is 1.4-1.6x faster than trtllm-gen (at 16 heads, without ShardQ, it is 4x), which is why the gain under ShardQ is small.

Scripts: e2e_prefill.py, prompt_logprob_probe.py, bench_shardq_rank.py

"""Prefill latency of two servers on public-domain text (Project Gutenberg #1342 and #2701).

Each request asks for one output token, so its latency is the prefill. Every trial uses a
different slice of text, and every prompt starts with a per-run salt, so neither trials nor
reruns hit each other's prefix cache. The servers take turns on identical prompts.
  python e2e_prefill.py --urls http://127.0.0.1:8400 http://127.0.0.1:8500 --out r.json
"""

import argparse
import json
import secrets
import statistics
import time
import urllib.request

BOOKS = ["pg1342.txt", "pg2701.txt"]  # https://www.gutenberg.org/cache/epub/{1342,2701}/pg{1342,2701}.txt
CHARS_PER_TOKEN = 4.0


def load_text() -> str:
    parts = []
    for path in BOOKS:
        t = open(path, encoding="utf-8").read()
        parts.append(t[t.index("*** START OF THE PROJECT GUTENBERG EBOOK") + 200 :])
    return "\n\n".join(parts)


def complete(url: str, prompt: str) -> tuple[float, int]:
    body = {"model": "glm52", "prompt": prompt, "max_tokens": 1, "temperature": 0}
    req = urllib.request.Request(f"{url}/v1/completions", data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    t0 = time.perf_counter()
    with urllib.request.urlopen(req, timeout=1800) as r:
        usage = json.load(r)["usage"]
    return time.perf_counter() - t0, usage["prompt_tokens"]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--urls", nargs="+", required=True)
    ap.add_argument("--pure", type=int, nargs="*", default=[8192, 16384, 32768, 49152])
    ap.add_argument("--context", type=int, nargs=2, action="append", default=None,
                    metavar=("CACHED", "NEW"), help="cached prefix tokens and new tokens")
    ap.add_argument("--trials", type=int, default=3)
    ap.add_argument("--salt", default=secrets.token_hex(8))
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    contexts = args.context if args.context is not None else [[16384, 8192], [24576, 8192]]
    text, cursor, takes = load_text(), 0, 0
    print(f"salt {args.salt}", flush=True)

    def take(tokens: int, salted: bool) -> str:
        # each salted take gets its own salt: the text wraps around, and a repeated
        # salted slice would hit the prefix cache
        nonlocal cursor, takes
        n = int(tokens * CHARS_PER_TOKEN)
        if cursor + n > len(text):
            cursor = 0
        s = text[cursor : cursor + n]
        cursor += n
        takes += 1
        return f"[{args.salt}-{takes}]\n{s}" if salted else s

    results = []
    cases = [("pure", n, 0) for n in args.pure] + [("context", new, cached) for cached, new in contexts]
    for kind, new, cached in cases:
        runs = {u: [] for u in args.urls}
        tokens = {}
        for trial in range(args.trials):
            prefix = take(cached, salted=True) if cached else ""
            body = take(new, salted=not cached)
            order = args.urls if trial % 2 == 0 else args.urls[::-1]
            for u in order:
                if cached:
                    complete(u, prefix)  # warm the prefix cache, not timed
                dt, n_tok = complete(u, prefix + body)
                runs[u].append(dt)
                tokens[u] = n_tok
        row = {"kind": kind, "cached_tokens": cached, "new_tokens": new,
               "prompt_tokens": tokens, "latency_s": {u: [round(x, 4) for x in v] for u, v in runs.items()},
               "median_s": {u: round(statistics.median(v), 4) for u, v in runs.items()}}
        base, other = args.urls[0], args.urls[-1]
        row["speedup"] = round(row["median_s"][base] / row["median_s"][other], 3)
        print(json.dumps(row), flush=True)
        results.append(row)
    json.dump(results, open(args.out, "w"), indent=1)


if __name__ == "__main__":
    main()

"""Teacher-forced accuracy of the prefill: per-position logprob of each prompt token and the model's
top-1 next token, read through prompt_logprobs. Every value comes out of the prefill forward, so
two servers can be compared position by position over many tokens, and two runs on one server give
the run-to-run noise floor. Prompts are ~8K-token slices of public text (Project Gutenberg #1342,
#2701 and the GPL-3 / GFDL-1.3 license texts).
  python prompt_logprob_probe.py --base-url http://127.0.0.1:8400 --cache-salt r1 --out a.json
  python prompt_logprob_probe.py --compare a.json b.json
"""

import argparse
import json
import math
from pathlib import Path

import httpx

CHARS = 32_000


def prompts() -> list[str]:
    texts = []
    for path in ["pg1342.txt", "pg2701.txt"]:  # Project Gutenberg #1342 and #2701
        t = open(path, encoding="utf-8").read()
        texts.append(t[t.index("*** START OF THE PROJECT GUTENBERG EBOOK") + 200 :])
    lic = Path("/usr/share/common-licenses")
    licenses = (lic / "GPL-3").read_text() + "\n\n" + (lic / "GFDL-1.3").read_text()
    austen, melville = texts
    return [austen[200_000 : 200_000 + CHARS], austen[500_000 : 500_000 + CHARS],
            melville[300_000 : 300_000 + CHARS], melville[700_000 : 700_000 + CHARS],
            melville[1_000_000 : 1_000_000 + CHARS], licenses[:CHARS]]


def probe(base_url: str, model: str, cache_salt: str) -> list[dict]:
    out = []
    with httpx.Client(timeout=1800) as c:
        for i, p in enumerate(prompts()):
            ids = c.post(f"{base_url}/tokenize", json={"model": model, "prompt": p}).json()["tokens"]
            r = c.post(f"{base_url}/v1/completions", json={
                "model": model, "prompt": ids, "max_tokens": 1, "temperature": 0, "prompt_logprobs": 1,
                "cache_salt": cache_salt}).json()
            # position j maps token id -> {logprob, rank} for prompt token j and the top-1 token
            rows = r["choices"][0]["prompt_logprobs"][1:]
            lp = [row[str(t)]["logprob"] for row, t in zip(rows, ids[1:], strict=True)]
            top = [int(next(k for k, v in row.items() if v["rank"] == 1)) for row in rows]
            out.append({"prompt_tokens": len(ids), "logprobs": lp, "top1": top})
            print(f"prompt {i}: {len(ids)} tokens, mean logprob {sum(lp) / len(lp):.4f}", flush=True)
    return out


def compare(a: list[dict], b: list[dict]) -> None:
    da, agree, n, sa, sb = [], 0, 0, 0.0, 0.0
    for x, y in zip(a, b):
        assert len(x["logprobs"]) == len(y["logprobs"])
        da += [abs(p - q) for p, q in zip(x["logprobs"], y["logprobs"])]
        agree += sum(p == q for p, q in zip(x["top1"], y["top1"]))
        n += len(x["logprobs"])
        sa += sum(x["logprobs"]); sb += sum(y["logprobs"])
    da.sort()
    print(f"{n} positions | mean |dlogprob| {sum(da) / n:.5f}, p99 {da[int(0.99 * n)]:.4f}, max {da[-1]:.4f} | "
          f"top-1 agreement {agree / n:.4%} | perplexity {math.exp(-sa / n):.4f} vs {math.exp(-sb / n):.4f}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-url")
    ap.add_argument("--model", default="glm52")
    ap.add_argument("--cache-salt", default="r1")
    ap.add_argument("--out")
    ap.add_argument("--compare", nargs=2)
    a = ap.parse_args()
    if a.compare:
        compare(*(json.load(open(f)) for f in a.compare))
    else:
        json.dump(probe(a.base_url, a.model, a.cache_salt), open(a.out, "w"))

"""One layer of a ShardQ rank's sparse MQA prefill on real GLM-5.2 top-k: main's
trtllm-gen with every head against the grouped path, and where the grouped path's
time goes (union build, FA4 calls, calls of 1,024 rows against one call).

The rank's tokens are its zigzag chunks of one ~15K-token prompt (world 4).
L2 is flushed before every timed run, as in bench_grouped_prefill.py.
  python3 bench_shardq_rank.py /d/topk_0{12,20,28,36}.pt --out x.json
"""

import argparse
import json
import os
import statistics

import torch

import vllm.v1.attention.backends.mla.flashattn_mla_sparse as be
from vllm.v1.attention.backends.mla.sparse_utils import build_query_group_row_union
from vllm.v1.attention.ops.shardq import zigzag_split
from vllm.v1.worker.workspace import init_workspace_manager

LORA, ROPE, HEADS, WORLD, GROUP = 512, 64, 64, 4, 2


def timed(fn, flush, iters=20):
    for _ in range(3):
        fn()
    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
    for s, e in zip(starts, ends):
        flush.zero_()
        s.record()
        fn()
        e.record()
    torch.cuda.synchronize()
    return statistics.median(s.elapsed_time(e) * 1000 for s, e in zip(starts, ends))


def make_impl(topk_buf):
    impl = object.__new__(be.FlashAttnMLASparseFA4Impl)
    impl.__dict__.update(kv_lora_rank=LORA, qk_rope_head_dim=ROPE, scale=(LORA + ROPE) ** -0.5)
    impl.topk_indices_buffer = topk_buf
    return impl


def main():
    from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla

    ap = argparse.ArgumentParser()
    ap.add_argument("files", nargs="+")
    ap.add_argument("--out", required=True)
    args = ap.parse_args()
    init_workspace_manager(torch.device("cuda"))
    flush = torch.empty(256 << 20, dtype=torch.uint8, device="cuda")
    workspace = torch.zeros(512 << 20, dtype=torch.uint8, device="cuda")
    gen = torch.Generator(device="cuda").manual_seed(0)
    results = []
    for path in args.files:
        d = torch.load(path)
        topk_all = d["topk"].to(torch.int32).cuda()
        valid_all = d["valid"].to(torch.int32).cuda()
        n = topk_all.shape[0]
        kv_rows = -(-(int(topk_all.max()) + 1) // 64) * 64
        kv = (torch.rand(kv_rows, LORA + ROPE, device="cuda", generator=gen) - 0.5).to(torch.bfloat16)
        kv_cache = kv.view(-1, 64, LORA + ROPE)
        token_ids, per_rank = zigzag_split(torch.tensor([n]).numpy(), WORLD)
        start = 0
        for rank, count in enumerate(per_rank.tolist()):
            local = torch.from_numpy(token_ids[start : start + count]).cuda()
            start += count
            if rank not in (0, WORLD - 1):
                continue
            topk, valid = topk_all[local].contiguous(), valid_all[local].contiguous()
            q = (torch.rand(count, HEADS, LORA + ROPE, device="cuda", generator=gen) - 0.5).to(torch.bfloat16)
            impl = make_impl(topk)

            def trtllm():
                trtllm_batch_decode_with_kv_cache_mla(
                    query=q.unsqueeze(1), kv_cache=kv_cache.unsqueeze(1), workspace_buffer=workspace,
                    qk_nope_head_dim=192, kv_lora_rank=LORA, qk_rope_head_dim=ROPE,
                    block_tables=topk.unsqueeze(1), seq_lens=valid, max_seq_len=topk.shape[1],
                    bmm1_scale=impl.scale, bmm2_scale=1.0, sparse_mla_top_k=topk.shape[1])

            def grouped(rows):
                impl._shardq_rows = rows
                impl._forward_row_groups(q, kv_cache, topk, valid, 64, GROUP)

            groups = count // GROUP
            enc = torch.empty(groups, GROUP * topk.shape[1], dtype=torch.int32, device="cuda")
            lens = torch.empty(groups, dtype=torch.int32, device="cuda")

            def union_once():
                build_query_group_row_union(topk[: groups * GROUP], valid[: groups * GROUP],
                                            group_size=GROUP, enc=enc, lens=lens)

            union_once()
            kv_flat = kv_cache.view(-1, LORA + ROPE)
            rows_q = q[: groups * GROUP].view(groups, 128, -1)
            out = torch.empty(groups, 128, LORA, dtype=torch.bfloat16, device="cuda")

            def fa4_once():
                impl._fa4_sparse_attention(rows_q[..., :LORA], rows_q[..., LORA:], kv_flat, enc, lens,
                                           return_lse=False, out=out, query_group_size=GROUP)

            row = dict(case=f"{os.path.basename(path)}:layer{d.get('layer')}", rank=rank, tokens=count,
                       union_lists=round(float(lens.float().mean()) / float(valid.float().mean()), 3),
                       us_trtllm=timed(trtllm, flush), us_grouped_1024=timed(lambda: grouped(1024), flush),
                       us_grouped_one_call=timed(lambda: grouped(groups + 1), flush),
                       us_union_once=timed(union_once, flush), us_fa4_once=timed(fa4_once, flush))
            row["speedup_1024"] = round(row["us_trtllm"] / row["us_grouped_1024"], 3)
            row["speedup_one_call"] = round(row["us_trtllm"] / row["us_grouped_one_call"], 3)
            print(json.dumps({k: round(v, 1) if isinstance(v, float) else v for k, v in row.items()}), flush=True)
            results.append(row)
    json.dump(results, open(args.out, "w"), indent=1)


if __name__ == "__main__":
    main()

Details

Why. FA4's sparse MLA kernel needs 128 query heads per KV head. GLM-5.x has 64 heads and DeepSeek-V3.2 128, so under TP a rank holds 8-64; this backend therefore runs prefill on trtllm-gen, one query at a time, each streaming its own 2,048 KV rows. Consecutive queries select mostly the same rows: on GLM-5.2 the union of 2 neighbours' top-k is 1.15x one list, and of 8 neighbours 1.2-1.7x. flash-attention#205 lets each 128-head row hold 128 / heads queries that gather their union once, with a membership mask per query.

Changes

• ShardQ steps (_forward_row_groups): the rank's query is token-major with every head, so groups are views of it, and a group may span requests, since each query keeps to its own top-k through its membership bit. sparse_utils.build_query_group_row_union builds the unions from the converted KV rows: groups of two sort one list and look the other's rows up in it, larger groups sort all their rows, so the cost does not depend on the sequence length. A step takes one call, its unions in workspace that __init__ reserves from ShardQ's per-rank token bound (134 MB at 65,536 batched tokens and TP4). The last group repeats the rank's last token to fill its row.
• Layers that use ShardQ group only its steps. The steps it leaves are below sparse_mla_shard_q_min_tokens, where grouping cost more than it saved (16-32K cached + 256 new tokens: 0.87-0.88x with grouping).
• Steps without ShardQ (_forward_query_groups): the first n // G * G queries of each request run grouped, the rest on the one-query-per-row FA4 kernel this backend uses for decode. The unions come from sparse_utils.build_query_group_union (a position bitmap per group), in calls of at most 1,024 packed rows whose scratch (272 MB at 16 heads) __init__ reserves. The DSA model's absorbed q_nope is head-major, so each call copies its rows token-major.
• Prefill stays on trtllm-gen under DCP or PCP, with HiSparse, when decode needs an LSE, when the cache has 2^24 - 1 or more KV rows, or when the heads do not make 2, 4 or 8 queries per row (DeepSeek-V3.2 under ShardQ has 128 heads per rank).
• _forward_mqa_kernel doesn't receive the metadata, so forward_mqa keeps it on the impl for the duration of the call.
• flash_attn_varlen_func passes gather_kv_query_group_size through (default 1, no change), and compile_flash_attn_varlen_func_from_specs takes qv and the top-k gather specs, so FA4GroupedPrefillKernel registers the grouped kernel's two compile keys with the JIT warmup (compile-only). ShardQ's strided query views reuse them.

Tradeoffs and limitations

• Gain under ShardQ. ShardQ already lets each gathered KV row serve 64 heads; pairing queries only doubles that, and FA4's grouped kernel is 1.4-1.6x faster than trtllm-gen at 64 heads, of which the union build takes about a tenth. When a rank's token count is odd, the tail group costs about 90 µs per layer, as much as trtllm-gen spends on that one query.
• Union build without ShardQ. Each group scans a bitmap as long as the request, so its cost grows with prefill length × sequence length (1.5 ms per layer at 48K). The ShardQ path's builder avoids that and could replace it.
• Low overlap. When consecutive queries' top-k barely overlap, grouping is slower than one query per row: at 16 heads, break-even is near a union of 5.8 lists per group of 8, while GLM-5.2's unions are 1.2-1.7 lists (flash-attention#205).
• Startup warmup. Both grouped compile keys compile at startup, 5-9 s per process. The union builders' Triton kernels compile on first use (~0.15 s).
• KV dtype. BF16 KV only here; [Perf][Attention] FA4 sparse MLA: grouped prefill with an FP8 KV cache #63328 adds FP8.

AI assistance (Claude) was used to write this change, its tests and the benchmarks, which ran on my machine; I reviewed every changed line.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @XinyuanQiao.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork