[Perf][Sampler] Add opt-in FlashInfer Cake top-k/top-p sampling

benchmarks/kernels/benchmark_cake_sampling.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compare Cake with the existing top-k/top-p paths on captured decode logits.
+
+Example:
+    VLLM_USE_FLASHINFER_CAKE_SAMPLER=1 \
+    python benchmarks/kernels/benchmark_cake_sampling.py \
+        --logits flash_next_logits.pt --output microbench.json
+
+Masking mutates its input, so graph timings restore logits before each timed
+replay. Direct sampling uses FlashInfer's CUPTI helper. Both flush L2 before
+timing; compilation, allocations, input restoration and transfers are excluded.
+Requires FlashInfer >= 0.7.1 and cupti-python with CUDA 13 or newer.
+
+"""
+
+import argparse
+import json
+import statistics
+import warnings
+from functools import partial
+from pathlib import Path
+
+import flashinfer
+import torch
+from flashinfer.testing import bench_gpu_time_with_cupti
+
+from vllm.v1.sample.ops.topk_topp_cake import (
+    CAKE_MASK_MAX_ROWS,
+    apply_top_k_top_p_cake,
+    cake_eligible,
+    cake_sample,
+)
+from vllm.v1.sample.ops.topk_topp_sampler import (
+    apply_top_k_top_p_pytorch,
+    flashinfer_sample,
+)
+from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
+
+
+def mask_time_us(fn, logits, flush, repeats):
+    work = logits.clone()
+    for _ in range(5):
+        work.copy_(logits)
+        fn(work)
+    torch.accelerator.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        work.copy_(logits)
+        fn(work)
+    torch.cuda.current_stream().wait_stream(stream)
+    with torch.cuda.graph(graph):
+        fn(work)
+    events = []
+    for _ in range(repeats):
+        work.copy_(logits)
+        flush.zero_()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        graph.replay()
+        end.record()
+        events.append((start, end))
+    torch.accelerator.synchronize()
+    return statistics.median(s.elapsed_time(e) * 1000 for s, e in events)
+
+
+def sample_time_us(fn, logits, k, p, repeats):
+    for _ in range(5):
+        fn(logits, k, p)
+    torch.accelerator.synchronize()
+    with warnings.catch_warnings():
+        warnings.filterwarnings("error", message="CUPTI is not installed.*")
+        warnings.filterwarnings("error", message="cold_l2_cache=True but no GPU.*")
+        times = bench_gpu_time_with_cupti(
+            fn,
+            use_cuda_graph=True,
+            cold_l2_cache=True,
+            dry_run_time_ms=250,
+            repeat_iters=repeats,
+            input_args=(logits, k, p),
+        )
+    return statistics.median(times) * 1000
+
+
+def correctness(logits, k, p, top_k):
+    reference = apply_top_k_top_p_pytorch(logits.clone(), k, p)
+    baseline = apply_top_k_top_p_triton(logits.clone(), k, p)
+    candidate = apply_top_k_top_p_cake(logits.clone(), k, p, top_k)
+    expected = reference.double().softmax(-1)
+    result = {}
+    for name, actual in (("triton", baseline), ("cake", candidate)):
+        tv = (actual.double().softmax(-1) - expected).abs().sum(-1) / 2
+        result[name] = {
+            "mask_diff_rows": int(
+                ((actual > -torch.inf) != (reference > -torch.inf)).any(-1).sum()
+            ),
+            "tv_mean": float(tv.mean()),
+            "tv_max": float(tv.max()),
+        }
+        assert torch.isfinite(actual).any(-1).all(), "empty retained support"
+    # The reference retains ties at the top-k boundary; Cake keeps exactly k.
+    # Report these differences instead of claiming token identity across backends.
+    tokens = cake_sample(logits, k, p, top_k).long()
+    assert (candidate > -torch.inf).gather(1, tokens[:, None]).all()
+    # Independently check Cake's exact-k, lower-index tie contract.
+    probs = flashinfer.sampling.softmax(logits)
+    values, ids = probs.sort(descending=True, stable=True)
+    values, ids = values[:, :top_k].double(), ids[:, :top_k]
+    values /= values.sum(-1, keepdim=True)
+    keep = values.cumsum(-1) - values < p[:, None]
+    values = values.masked_fill(~keep, 0)
+    values /= values.sum(-1, keepdim=True)
+    exact = torch.zeros_like(logits, dtype=torch.float64).scatter_(1, ids, values)
+    tv = (candidate.double().softmax(-1) - exact).abs().sum(-1) / 2
+    result["cake_exact_k_tv_max"] = float(tv.max())
+    assert tv.max() < 2e-5, "Cake disagrees with its exact-k/top-p contract"
+    return result
+
+
+@torch.inference_mode()
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--logits", type=Path)
+    parser.add_argument("--vocab-size", type=int, default=248320)
+    parser.add_argument(
+        "--rows", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64, 128, 256]
+    )
+    parser.add_argument("--top-k", type=int, default=20)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--rounds", type=int, default=3)
+    parser.add_argument("--repeats", type=int, default=50)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    torch.manual_seed(123)
+    if args.logits:
+        captured = torch.load(args.logits, map_location="cuda", weights_only=True)
+        source = str(args.logits)
+    else:
+        captured = torch.randn(max(args.rows), args.vocab_size, device="cuda") * 3
+        source = "synthetic; shape validation only"
+    captured = captured.float().contiguous()
+    vocab = captured.shape[-1]
+    assert cake_eligible(captured, args.top_k), "Cake pipeline is unavailable"
+    flush = torch.empty(128 * 1024**2, device="cuda", dtype=torch.int8)
+    result = {
+        "metadata": {
+            "source": source,
+            "gpu": torch.cuda.get_device_name(),
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "flashinfer": flashinfer.__version__,
+            "dtype": str(captured.dtype),
+            "vocab_size": vocab,
+            "top_k": args.top_k,
+            "top_p": args.top_p,
+            "rounds": args.rounds,
+            "repeats": args.repeats,
+            "mask_timing": "CUDA graph events; restore and 128 MiB L2 flush excluded",
+            "sample_timing": "FlashInfer CUPTI helper; CUDA graph; cold L2",
+        },
+        "cases": [],
+    }
+    for rows in args.rows:
+        logits = captured[torch.arange(rows, device="cuda") % len(captured)].clone()
+        k = torch.full((rows,), args.top_k, device="cuda", dtype=torch.int32)
+        p = torch.full((rows,), args.top_p, device="cuda")
+        checks = correctness(logits, k, p, args.top_k)
+        masks = {
+            "triton": partial(apply_top_k_top_p_triton, k=k, p=p),
+            "cake": partial(apply_top_k_top_p_cake, k=k, p=p, k_max=args.top_k),
+        }
+        samples = {
+            "flashinfer": flashinfer_sample,
+            "cake": partial(cake_sample, k_max=args.top_k),
+        }
+        timings = {
+            "mask_triton": [],
+            "mask_cake": [],
+            "sample_flashinfer": [],
+            "sample_cake": [],
+        }
+        for round_idx in range(args.rounds):
+            for name in list(masks)[:: 1 if round_idx % 2 == 0 else -1]:
+                timings[f"mask_{name}"].append(
+                    mask_time_us(masks[name], logits, flush, args.repeats)
+                )
+            for name in list(samples)[:: 1 if round_idx % 2 == 0 else -1]:
+                timings[f"sample_{name}"].append(
+                    sample_time_us(samples[name], logits, k, p, args.repeats)
+                )
+        case = {
+            "rows": rows,
+            "mask_dispatch": "cake" if rows <= CAKE_MASK_MAX_ROWS else "triton",
+            "correctness": checks,
+            "timings_us": timings,
+            "median_us": {
+                name: statistics.median(values) for name, values in timings.items()
+            },
+        }
+        result["cases"].append(case)
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(result, indent=2) + "\n")
+        print(json.dumps(case), flush=True)
+
+
+if __name__ == "__main__":
+    main()

tests/v1/sample/test_topk_topp_cake.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer_cake_sampling
+from vllm.v1.sample.ops.topk_topp_cake import (
+    apply_top_k_top_p_cake,
+    cake_eligible,
+    cake_sample,
+)
+from vllm.v1.sample.ops.topk_topp_sampler import (
+    apply_top_k_top_p,
+    apply_top_k_top_p_pytorch,
+)
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda()
+    or not current_platform.has_device_capability(90)
+    or not has_flashinfer_cake_sampling(),
+    reason="FlashInfer Cake sampling needs FlashInfer >= 0.7.1 on SM90+",
+)
+
+DEVICE = "cuda"
+
+
+def _params(batch_size, top_k, top_p):
+    k = torch.full((batch_size,), top_k, device=DEVICE, dtype=torch.int32)
+    p = None if top_p is None else torch.full((batch_size,), top_p, device=DEVICE)
+    return k, p
+
+
+@pytest.mark.parametrize("batch_size", [1, 8, 64])
+@pytest.mark.parametrize("vocab_size", [32000, 151936])
+@pytest.mark.parametrize("top_k", [1, 20, 1024])
+@pytest.mark.parametrize("top_p", [None, 0.9])
+def test_mask_matches_reference(batch_size, vocab_size, top_k, top_p):
+    torch.manual_seed(0)
+    logits = torch.randn(batch_size, vocab_size, device=DEVICE) * 3
+    k, p = _params(batch_size, top_k, top_p)
+
+    expected = apply_top_k_top_p_pytorch(logits.clone(), k, p)
+    actual = apply_top_k_top_p_cake(logits.clone(), k, p, top_k)
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_mask_per_row_params():
+    torch.manual_seed(1)
+    batch_size, vocab_size = 13, 50000
+    logits = torch.randn(batch_size, vocab_size, device=DEVICE) * 2
+    k = torch.randint(1, 300, (batch_size,), device=DEVICE, dtype=torch.int32)
+    p = torch.rand(batch_size, device=DEVICE) * 0.8 + 0.1
+
+    expected = apply_top_k_top_p_pytorch(logits.clone(), k, p)
+    actual = apply_top_k_top_p_cake(logits.clone(), k, p, int(k.max()))
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_mask_ties_keep_lowest_index():
+    logits = torch.zeros(2, 1000, device=DEVICE)
+    logits[:, 10:20] = 5.0
+    k, _ = _params(2, 4, None)
+
+    kept = apply_top_k_top_p_cake(logits.clone(), k, None, 4) > float("-inf")
+
+    assert kept.sum(dim=1).tolist() == [4, 4]
+    assert kept[:, 10:14].all()
+
+
+@pytest.mark.parametrize("top_p", [None, 0.8])
+def test_sample_within_kept_set(top_p):
+    torch.manual_seed(2)
+    batch_size, vocab_size, top_k = 24, 32000, 50
+    logits = torch.randn(batch_size, vocab_size, device=DEVICE) * 2
+    k, p = _params(batch_size, top_k, top_p)
+    kept = apply_top_k_top_p_pytorch(logits.clone(), k, p) > float("-inf")
+
+    for _ in range(20):
+        tokens = cake_sample(logits, k, p, top_k).long()
+        assert kept.gather(1, tokens[:, None]).all()
+
+
+def test_sample_preserves_nonuniform_probabilities():
+    """Catch biased draws that a support-only check misses."""
+    torch.manual_seed(4)
+    batch_size = 128
+    logits = torch.full((batch_size, 32000), -torch.inf, device=DEVICE)
+    probabilities = torch.tensor([0.4, 0.3, 0.2, 0.1], device=DEVICE)
+    logits[:, :4] = probabilities.log()
+    k, p = _params(batch_size, 4, 1.0)
+    draws = torch.cat([cake_sample(logits, k, p, 4).clone() for _ in range(128)])
+    observed = draws.long().bincount(minlength=32000).float() / draws.numel()
+    assert observed[4:].sum() == 0
+    torch.testing.assert_close(observed[:4], probabilities, atol=0.03, rtol=0)
+
+
+@pytest.mark.parametrize("batch_size", [8, 128])
+def test_dispatch_matches_reference(batch_size, monkeypatch):
+    """Small batches take Cake, large ones Triton; both match the reference."""
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_CAKE_SAMPLER", "1")
+    torch.manual_seed(3)
+    logits = torch.randn(batch_size, 32000, device=DEVICE) * 3
+    k, p = _params(batch_size, 20, 0.9)
+
+    expected = apply_top_k_top_p_pytorch(logits.clone(), k, p)
+    actual = apply_top_k_top_p(logits.clone(), k, p, k_max=20)
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_eligibility(monkeypatch):
+    logits = torch.zeros(4, 32000, device=DEVICE)
+    monkeypatch.delenv("VLLM_USE_FLASHINFER_CAKE_SAMPLER", raising=False)
+    assert not cake_eligible(logits, 20)
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_CAKE_SAMPLER", "1")
+    assert cake_eligible(logits, 20)
+    assert cake_eligible(logits, 1024)
+    assert not cake_eligible(logits, None)
+    assert not cake_eligible(logits, 1025)
+    # A request without top-k carries k == vocab_size.
+    assert not cake_eligible(logits, 32000)
+
+    monkeypatch.setenv("VLLM_USE_FLASHINFER_CAKE_SAMPLER", "0")
+    assert not cake_eligible(logits, 20)

vllm/envs.py

@@ -48,6 +48,7 @@
     VLLM_LOG_STATS_INTERVAL: float = 10.0
     VLLM_TRACE_FUNCTION: int = 0
     VLLM_USE_FLASHINFER_SAMPLER: bool = True
+    VLLM_USE_FLASHINFER_CAKE_SAMPLER: bool = False
     VLLM_PP_LAYER_PARTITION: str | None = None
     VLLM_CPU_KVCACHE_SPACE: int | None = 0
     VLLM_CPU_OMP_THREADS_BIND: str = "auto"
@@ -876,6 +877,14 @@ def _resolve_rust_cli_path() -> str | None:
         if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ
         else True
     ),
+    # Whether to use FlashInfer's Cake radix top-k/top-p kernels (FlashInfer
+    # >= 0.7.1) when every request in the batch sets 1 <= top_k <= 1024.
+    # Opt in with 1; defaults to the Triton masking and FlashInfer sampling paths.
+    "VLLM_USE_FLASHINFER_CAKE_SAMPLER": lambda: (
+        bool(int(os.environ["VLLM_USE_FLASHINFER_CAKE_SAMPLER"]))
+        if "VLLM_USE_FLASHINFER_CAKE_SAMPLER" in os.environ
+        else False
+    ),
     # Pipeline stage partition strategy
     "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None),
     # (CPU backend only) CPU key-value cache space.

vllm/utils/flashinfer.py

@@ -364,6 +364,19 @@ def has_flashinfer_comm() -> bool:
     return has_flashinfer() and importlib.util.find_spec("flashinfer.comm") is not None
 
 
+@functools.cache
+def has_flashinfer_cake_sampling() -> bool:
+    """Return whether FlashInfer's Cake radix top-k/top-p sampler is available."""
+    if not has_flashinfer():
+        return False
+    module = _get_submodule("flashinfer.cake_sampling")
+    return bool(
+        module
+        and callable(getattr(module, "top_k_top_p_sampling_from_probs", None))
+        and callable(getattr(module, "cake_sampling_route", None))
+    )
+
+
 @functools.cache
 def has_flashinfer_packed_fused_kda_decode() -> bool:
     """Return whether FlashInfer's packed fused KDA decode API is available."""