mjolnir-cli 0.2.0__py3-none-any.whl

mjolnir_cli/dockerctl.py

@@ -0,0 +1,533 @@
+"""Serve orchestration: drive the vLLM docker container from Python.
+
+The vLLM server runs in a container because the whole point of this repo is
+a *custom* image (the sm_110 patch stack + the Mjolnir decode kernel behind
+``--attention-backend MJOLNIR``) — ``mjolnir serve`` replaces the
+bash/docker-compose glue, not the container.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+import urllib.request
+from pathlib import Path
+
+from mjolnir_cli import cachectl
+from mjolnir_cli.config import CONTAINER_NAME, Settings, config_field, configs_root, user_configs_dir
+
+# Container-side env (the Thor-specific knobs from the old vllm.yaml service).
+_CONTAINER_ENV: list[str] = [
+    "DO_NOT_TRACK=1",
+    "VLLM_NO_USAGE_STATS=1",
+    "MAX_JOBS=4",
+    "NVCC_THREADS=2",
+    "FLASHINFER_NVCC_THREADS=2",
+    "VLLM_HTTP_TIMEOUT_KEEP_ALIVE=600",
+    "FLASHINFER_DISABLE_VERSION_CHECK=1",
+    "VLLM_CACHE_ROOT=/data/cache",
+    "VLLM_ASSETS_CACHE=/data/assets",
+    "VLLM_TUNED_CONFIG_FOLDER=/data/kernels-cache",
+    "FLASHINFER_JIT_DIR=/data/flashinfer-cache",
+    # FlashInfer 0.7 keeps its JIT modules, CuTe-DSL objects and downloaded
+    # cubins under $FLASHINFER_WORKSPACE_BASE/.cache/flashinfer (default:
+    # $HOME, i.e. the container's throwaway /root): point it at the mounted
+    # cache so fp4_gemm_cutlass & co. build once, not on every server start
+    "FLASHINFER_WORKSPACE_BASE=/data/flashinfer-cache",
+    "TRITON_CACHE_DIR=/data/triton-cache",
+    "TORCHINDUCTOR_CACHE_DIR=/data/cache/inductor",
+    "TORCH_COMPILE_CACHE_DIR=/data/cache/torch_compile_cache",
+    "XDG_CACHE_HOME=/data/cache/xdg",
+    "TIKTOKEN_ENCODINGS_BASE=/data/tiktoken",
+    "TIKTOKEN_RS_CACHE_DIR=/data/tiktoken",
+    "VLLM_LOGGING_LEVEL=INFO",
+    "NVIDIA_VISIBLE_DEVICES=all",
+    "NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics",
+    "CUDA_VISIBLE_DEVICES=0",
+    "VLLM_TRITON_USE_TD=1",
+    "PYTHONUNBUFFERED=1",
+    "SAFETENSORS_FAST_GPU=1",
+    "NVIDIA_DISABLE_REQUIRE=1",
+    "NVIDIA_FORWARD_COMPAT=1",
+    "NCCL_P2P_DISABLE=1",
+    "ENABLE_TRIATTENTION=0",
+    "CUDA_MANAGED_FORCE_DEVICE_ALLOC=1",
+    "TORCH_CUDA_ARCH_LIST=11.0a",
+    "FLASHINFER_CUDA_ARCH_LIST=11.0a",
+    "CUDA_TOOLKIT_PATH=/usr/local/cuda-13",
+    "CUTE_DSL_ARCH=sm_110a",
+    "TORCH_MATMUL_PRECISION=high",
+    "VLLM_DISABLED_KERNELS=CutlassFP8ScaledMMLinearKernel",
+    "VLLM_USAGE_SOURCE=production",
+    "TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas",
+    "VLLM_USE_BREAKABLE_CUDAGRAPH=0",
+    "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True",
+    # Thor-only host overrides from the old compose file:
+    "VLLM_ALLOW_LONG_MAX_MODEL_LEN=1",
+    "VLLM_MARLIN_USE_ATOMIC_ADD=1",
+]
+
+_HOST_MOUNTS = [
+    "/etc/nv_tegra_release",  # L4T id (Jetson)
+    "/tmp/nv_jetson_model",
+    "/etc/localtime",
+    "/etc/machine-id",
+]
+
+
+class DockerError(RuntimeError):
+    pass
+
+
+def _docker() -> str:
+    if shutil.which("docker") is None:
+        raise DockerError("docker not found on PATH")
+    return "docker"
+
+
+def _run(args: list[str], **kw) -> subprocess.CompletedProcess:
+    return subprocess.run(args, check=True, **kw)
+
+
+def _run_quiet(args: list[str], ok_codes: tuple[int, ...] = (0, 1)) -> tuple[int, str]:
+    p = subprocess.run(args, capture_output=True, text=True)
+    if p.returncode not in ok_codes:
+        raise DockerError(f"{' '.join(args[:3])}… failed ({p.returncode}): {(p.stderr or p.stdout).strip()[:500]}")
+    return p.returncode, p.stdout.strip()
+
+
+def container_running(name: str = CONTAINER_NAME) -> bool:
+    _, out = _run_quiet([_docker(), "inspect", "-f", "{{.State.Running}}", name])
+    return out == "true"
+
+
+# Mjolnir reads K/V fastest from head-major pages (each KV head's page is one
+# contiguous block): +11-17% decode at long context in `mjolnir verify
+# hndsweep`. vLLM allocates that layout (LBHNC) with VLLM_KV_CACHE_LAYOUT=HND;
+# both fallbacks (FlashAttention 4, FlashInfer) read it through strides, so
+# MJOLNIR configs get it by default.
+MJOLNIR_ENV_DEFAULTS = {"VLLM_KV_CACHE_LAYOUT": "HND"}
+
+# fp6 (e2m3) KV needs per-channel scales from `mjolnir calibrate`, stored next
+# to the model's configs (mounted at /configs in the container).
+KV_SCALES_FILE = "kv_scales.json"
+
+
+# What the image's installed mjolnir package is built from: `mjolnir image
+# build` stamps the checkout's HEAD into this label.
+REVISION_LABEL = "org.opencontainers.image.revision"
+KERNEL_PATHS = ("src/mjolnir", "docker/vllm-thor")
+
+
+def stale_image(image: str) -> str | None:
+    """Why ``image`` serves an older Mjolnir than this checkout, or None.
+
+    Serving uses the package installed in the image (only verify and
+    calibrate overlay the checkout's source), so a checkout that changed
+    ``src/mjolnir`` or the patch stack after the build serves old code. None
+    when it cannot tell: no checkout, no label (e.g. the official image), or
+    a revision git does not know."""
+    try:
+        from mjolnir_cli.config import find_repo_root
+
+        root = find_repo_root()
+    except FileNotFoundError:
+        return None
+    try:
+        rc, rev = _run_quiet(
+            [_docker(), "image", "inspect", "-f", '{{ index .Config.Labels "' + REVISION_LABEL + '" }}', image]
+        )
+    except (DockerError, OSError):  # no docker / no image: nothing to say
+        return None
+    rev = rev.strip()
+    if rc != 0 or not rev or rev in ("unknown", "<no value>"):
+        return None
+    r = subprocess.run(
+        ["git", "-C", str(root), "diff", "--quiet", rev, "HEAD", "--", *KERNEL_PATHS],
+        capture_output=True,
+    )
+    if r.returncode != 1:  # 0: same code; >1: revision unknown here
+        return None
+    log = (
+        subprocess.run(
+            ["git", "-C", str(root), "log", "--oneline", f"{rev}..HEAD", "--", *KERNEL_PATHS],
+            capture_output=True,
+            text=True,
+        )
+        .stdout.strip()
+        .splitlines()
+    )
+    head = f"{len(log)} commit(s)" if log else "changes"
+    return (
+        f"{image} was built from {rev}; {head} to {', '.join(KERNEL_PATHS)} since "
+        f"({'; '.join(log[:3])}{' ...' if len(log) > 3 else ''}). It serves the old code — "
+        "rebuild: mjolnir image build"
+    )
+
+
+def kv_scales_host_path(s: Settings) -> Path:
+    """Where `mjolnir calibrate` writes, and an fp6 config's scales are read
+    from (repo or user-local configs dir).
+
+    ``kv_scales.json`` is the model's — next to the config dirs, shared by
+    every config without its own drafter (the built-in MTP head is part of
+    the model). A config whose external drafter (DSpark, DFlash, ...) runs
+    on the Mjolnir backend has layers no other config has: it gets its own
+    ``kv_scales.json`` inside its config dir, target and drafter layers in
+    one file, written by ``mjolnir calibrate --config <config>``."""
+    from mjolnir_cli.config import config_drafter_on_mjolnir
+
+    cdir = s.config_path.parent
+    own = cdir / KV_SCALES_FILE
+    if own.exists() or config_drafter_on_mjolnir(s.config_path):
+        return own
+    return cdir.parent / KV_SCALES_FILE
+
+
+def missing_kv_scales(s: Settings) -> str | None:
+    """Why an fp6 config cannot start (its calibration file is missing), or None."""
+    if (s.attention_backend or "") not in ("MJOLNIR", "CUSTOM") or os.environ.get("MJOLNIR_ATTN_KV_SCALES"):
+        return None
+    if not (config_field(s.config_path, "kv-cache-dtype") or "").startswith("fp6"):
+        return None
+    path = kv_scales_host_path(s)
+    if path.exists():
+        return None
+    return f"{s.quant} serves fp6 KV but {path} is missing — calibrate it first: mjolnir calibrate --config {s.quant}"
+
+
+def backend_env(s: Settings) -> list[str]:
+    """Container env for the config's attention backend: the MJOLNIR
+    defaults (overridable from the host env) + every ``MJOLNIR_ATTN_*`` /
+    ``THOR_T5_*`` kernel knob set on the host. An fp6_e2m3 config also gets
+    MJOLNIR_ATTN_KV_SCALES pointed at its calibration file."""
+    out: list[str] = []
+    if (s.attention_backend or "") in ("MJOLNIR", "CUSTOM"):
+        out.append(f"MJOLNIR_CONFIG={s.quant}")  # the boot banner's `Model:` line
+        for k, v in MJOLNIR_ENV_DEFAULTS.items():
+            out.append(f"{k}={os.environ.get(k, v)}")
+        if (config_field(s.config_path, "kv-cache-dtype") or "").startswith("fp6") and not os.environ.get(
+            "MJOLNIR_ATTN_KV_SCALES"
+        ):
+            out.append(f"MJOLNIR_ATTN_KV_SCALES={_container_path(s, kv_scales_host_path(s))}")
+    out += config_env(s)
+    for k, v in sorted(os.environ.items()):
+        if k.startswith(("MJOLNIR_ATTN_", "THOR_T5_")):
+            out.append(f"{k}={v}")
+    return out
+
+
+def config_env(s: Settings) -> list[str]:
+    """``KEY=VALUE`` lines from the config's ``.env`` sidecar (inside its
+    config dir) — how an opt-in experiment config sets env vars vLLM's YAML
+    cannot carry (e.g. ``MJOLNIR_LM_HEAD=fp8``). ``{configs}`` expands to the
+    model's config dir as the container sees it; ``#`` starts a comment."""
+    side = s.config_path.parent / ".env"
+    if not side.exists():
+        return []
+    cdir = f"/configs/{s.model}"
+    out = []
+    for line in side.read_text().splitlines():
+        line = line.split("#", 1)[0].strip()
+        if "=" in line:
+            k, v = line.split("=", 1)
+            out.append(f"{k.strip()}={v.strip().replace('{configs}', cdir)}")
+    return out
+
+
+def _container_path(s: Settings, host: Path) -> str:
+    """A file under the configs dir, as the container sees it."""
+    try:
+        return "/configs/" + str(Path(host).relative_to(configs_root()))
+    except ValueError:  # user-local file: mirror its position under /configs/<model>/
+        try:
+            rel = str(Path(host).relative_to(user_configs_dir() / s.model))
+        except ValueError:
+            rel = Path(host).name
+        return f"/configs/{s.model}/{rel}"
+
+
+def build_run_args(
+    s: Settings,
+    image: str | None = None,
+    *,
+    name: str = CONTAINER_NAME,
+    detach: bool = True,
+    command: str | None = None,
+    extra_env: tuple[str, ...] = (),
+    extra_mounts: tuple[tuple[str, str], ...] = (),
+    serve_args: tuple[str, ...] = (),
+    chown: tuple[str, ...] = (),
+) -> list[str]:
+    """Assemble the full ``docker run`` argv for the vLLM container (the
+    server by default; ``command`` runs something else with the same mounts,
+    env and GPU access — ``mjolnir calibrate``). ``serve_args`` are appended
+    to ``vllm serve`` and override the config (e.g. ``--speculative-config``
+    for ``mjolnir verify spec``). ``chown`` lists container paths a
+    ``command`` writes into host mounts: they are handed to the invoking
+    user when it ends (the container runs as root)."""
+    image = image or s.image
+    cmd = [
+        _docker(),
+        "run",
+        *(["-d"] if detach else ["--rm"]),
+        "--name",
+        name,
+        "--gpus",
+        "all",
+        "--ipc",
+        "host",
+        "--privileged",
+        "--user",
+        "0:0",
+        # no --memory-swappiness: cgroup v2 (JetPack 7) has no per-container
+        # swappiness, docker only warns and discards it
+        "--ulimit",
+        "memlock=-1",
+        "--ulimit",
+        "stack=67108864",
+        "--ulimit",
+        "nofile=65536:65536",
+        "--restart",
+        "no",
+        # base image ENTRYPOINT is ["vllm","serve"] — pin ours, so the command
+        # below is the full argv (unpinned, docker would compose
+        # "vllm serve bash -lc ..." and argparse would reject it at startup)
+        "--entrypoint",
+        "bash",
+    ]
+
+    for k in _CONTAINER_ENV:
+        cmd += ["-e", k]
+    for kv in backend_env(s):
+        cmd += ["-e", kv]
+    for kv in extra_env:
+        cmd += ["-e", kv]
+    if os.environ.get("HF_TOKEN"):
+        cmd += ["-e", f"HF_TOKEN={os.environ['HF_TOKEN']}"]
+
+    d = s.data_dir
+    mounts = [
+        (str(d / "models"), "/data/models"),
+        (str(d / "cache"), "/data/cache"),
+        (str(d / "kernels-cache"), "/data/kernels-cache"),
+        (str(d / "flashinfer-cache"), "/data/flashinfer-cache"),
+        (str(d / "triton-cache"), "/data/triton-cache"),
+        (str(d / "tiktoken"), "/data/tiktoken"),
+        (str(configs_root()), "/configs:ro"),
+        ("/dev/shm", "/dev/shm"),
+        ("/etc/localtime", "/etc/localtime:ro"),
+        ("/etc/machine-id", "/etc/machine-id:ro"),
+    ]
+    # A user-local model config (outside the repo's configs/) is mounted as
+    # its config dir on top of the /configs mount (the more-specific mount
+    # wins): config.yaml + its .env sidecar, and its own scales when it has
+    # them.
+    local_cfg = s.config_path
+    repo_cfg = configs_root() / s.model / s.quant / "config.yaml"
+    if local_cfg != repo_cfg and local_cfg.exists():
+        mounts.append((str(local_cfg.parent), f"/configs/{s.model}/{s.quant}:ro"))
+        scales = kv_scales_host_path(s)
+        if scales.exists() and scales.parent != local_cfg.parent:
+            mounts.append((str(scales), f"/configs/{s.model}/{scales.name}:ro"))
+    mounts += list(extra_mounts)
+    for host, container in mounts:
+        if host.startswith("/data") is False and not Path(host).exists() and not host.startswith(("/dev/shm", "/etc/")):
+            continue  # optional host files (e.g. /etc/nv_tegra_release)
+        cmd += ["-v", f"{host}:{container}"]
+    if command is not None:
+        return [*cmd, image, "-lc", command + owner_suffix(chown)]
+    cmd += ["-p", f"{s.port}:8000"]
+
+    serve = (
+        "vllm serve "
+        f"--config /configs/{s.model}/{s.quant}/config.yaml "
+        "--host 0.0.0.0 "
+        "--cpu-offload-gb 0 "
+        "--pipeline-parallel-size 1 "
+        "--disable-fastapi-docs "
+        "--enable-force-include-usage "
+        "--enable-prompt-tokens-details "
+        "--download-dir /data/models/huggingface"
+    )
+    if serve_args:
+        import shlex
+
+        serve += " " + " ".join(shlex.quote(a) for a in serve_args)
+    cmd += [image, "-lc", serve]
+    return cmd
+
+
+def host_owner() -> str | None:
+    """``uid:gid`` of the user running the CLI (the sudo caller under sudo);
+    ``None`` off POSIX."""
+    if not hasattr(os, "getuid"):
+        return None
+    uid, gid = os.getuid(), os.getgid()
+    if uid == 0 and os.environ.get("SUDO_UID"):
+        uid, gid = int(os.environ["SUDO_UID"]), int(os.environ.get("SUDO_GID", gid))
+    return f"{uid}:{gid}"
+
+
+def owner_suffix(paths: tuple[str, ...] | list[str]) -> str:
+    """Shell tail for a root container command: hand ``paths`` (container
+    paths in host mounts) to the host user, keep the command's exit code.
+    Without it a file the container writes into the checkout is root-owned
+    and git hooks (end-of-file-fixer) cannot rewrite it."""
+    owner = host_owner()
+    if not paths or owner is None:
+        return ""
+    import shlex
+
+    quoted = " ".join(shlex.quote(p) for p in paths)
+    return f"; rc=$?; chown -R {owner} {quoted} 2>/dev/null; exit $rc"
+
+
+def run_once(argv: list[str], name: str) -> int:
+    """Run a one-shot container in the foreground (``docker run --rm``) and
+    always clean up after it: remove the container even when the CLI is
+    interrupted (stopping the docker client leaves the container running,
+    holding the GPU), then clear the host page cache — on Thor the cache the
+    container filled stays resident after it exits."""
+    cachectl.prepare_bench()
+    try:
+        return subprocess.run(argv).returncode
+    except KeyboardInterrupt:
+        print(f"\ninterrupted — removing container {name}", file=sys.stderr)
+        return 130
+    finally:
+        try:
+            _run_quiet([_docker(), "rm", "-f", name])
+        except DockerError:
+            pass
+        cachectl.prepare_bench()
+
+
+def serve_up(
+    s: Settings,
+    wait: bool = True,
+    wait_timeout_s: float = 1800.0,
+    dry_run: bool = False,
+    serve_args: tuple[str, ...] = (),
+) -> None:
+    cmd = build_run_args(s, serve_args=serve_args)
+    if dry_run:
+        print(" ".join(cmd))
+        return
+    why = missing_kv_scales(s)
+    if why:
+        raise DockerError(why)
+    # Always-before: clear the host page cache before starting vLLM (a clean,
+    # reproducible model-load). No-op + hint if the grant is missing; never
+    # prompts, never blocks.
+    cachectl.prepare_bench()
+    for sub in [
+        s.data_dir / p for p in ("models", "cache", "kernels-cache", "flashinfer-cache", "triton-cache", "tiktoken")
+    ]:
+        sub.mkdir(parents=True, exist_ok=True)
+    if container_running():
+        print("container already running — stopping the old one first")
+        serve_down(s)
+    _run(cmd)
+    if wait:
+        wait_healthy(s, wait_timeout_s)
+
+
+def serve_down(s: Settings, dry_run: bool = False) -> None:
+    if not container_running():
+        print("nothing to stop (container not running)")
+        return
+    if dry_run:
+        print(f"{_docker()} rm -f {CONTAINER_NAME}")
+        return
+    _run_quiet([_docker(), "rm", "-f", CONTAINER_NAME])
+    print("stopped")
+    # After stopping: clear the page cache (release the serving file activity
+    # that populated it). No-op + hint if the grant is missing.
+    cachectl.prepare_bench()
+
+
+def wait_healthy(s: Settings, timeout_s: float = 1800.0) -> None:
+    """Poll the OpenAI /health endpoint until the server answers (model load
+    + JIT on Thor takes minutes)."""
+    url = f"http://127.0.0.1:{s.port}/health"
+    t0 = time.time()
+    print(f"waiting for the server at {url} (up to {timeout_s:.0f}s)…")
+    while time.time() - t0 < timeout_s:
+        try:
+            with urllib.request.urlopen(url, timeout=5) as r:
+                if r.status == 200:
+                    print(f"healthy after {time.time() - t0:.0f}s")
+                    return
+        except Exception:  # noqa: BLE001
+            time.sleep(5)
+    raise TimeoutError(f"server not healthy after {timeout_s:.0f}s — check: mjolnir serve logs")
+
+
+def status(s: Settings) -> dict:
+    running = container_running()
+    state: dict = {"container": CONTAINER_NAME, "running": running}
+    if running:
+        _, out = _run_quiet(
+            [_docker(), "inspect", "-f", "{{.State.Status}} started {{.State.StartedAt}}", CONTAINER_NAME]
+        )
+        state["state"] = out
+    try:
+        with urllib.request.urlopen(f"http://127.0.0.1:{s.port}/health", timeout=5) as r:
+            state["health"] = "ok" if r.status == 200 else f"HTTP {r.status}"
+    except Exception:  # noqa: BLE001
+        state["health"] = "unreachable"
+    return state
+
+
+def logs(s: Settings, follow: bool = False, tail: int = 100) -> int:
+    if not container_running():
+        print("container not running", file=sys.stderr)
+        return 1
+    args = [_docker(), "logs", "--tail", str(tail)]
+    if follow:
+        args += ["-f"]
+    args += [CONTAINER_NAME]
+    try:
+        return subprocess.call(args)
+    except KeyboardInterrupt:
+        return 0
+
+
+# ANSI escape sequences: CSI (SGR colors, cursor moves) + OSC (e.g. hyperlinks).
+_ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)")
+
+
+def _strip_ansi(text: str) -> str:
+    return _ANSI_RE.sub("", text)
+
+
+def dump_server_log(since_epoch: float | None, out_path: Path) -> bool:
+    """Capture the vLLM container's logs to ``out_path``, with ANSI colors
+    stripped. ``since_epoch`` (unix seconds) bounds the capture from that
+    instant; ``None`` captures the whole container log since it started
+    (startup kernel-dispatch lines + everything after — what ``bench perf``
+    saves for kernel debugging). Best-effort: returns ``False`` (no raise)
+    when the container is absent/not running or docker fails — a log
+    capture failure never fails the bench."""
+    if not container_running():
+        return False
+    args = [_docker(), "logs"]
+    if since_epoch is not None:
+        ts = time.strftime("%Y-%m-%dT%H:%M:%S.000000000Z", time.gmtime(since_epoch))
+        args += ["--since", ts]
+    args += [CONTAINER_NAME]
+    try:
+        with out_path.open("w") as f:
+            subprocess.run(args, stdout=f, stderr=f, check=False)
+    except (OSError, DockerError):
+        return False
+    try:
+        out_path.write_text(_strip_ansi(out_path.read_text()))
+    except OSError:
+        return False
+    return True

mjolnir_cli/drivers/fa2_driver.py

@@ -0,0 +1,89 @@
+"""FlashAttention 2 paged decode driver for ncu — the stock-vLLM baseline.
+
+What vLLM 0.31.0 runs on sm_110 out of the box for a head_dim=256 model:
+``vllm.vllm_flash_attn.flash_attn_varlen_func(..., fa_version=2)`` over a
+paged KV cache in NHD order ``(num_blocks, block_size, h_k, d)``. FA2 has no
+fp8 KV path, so the cache is **bf16** — twice the bytes of the fp8 legs (the
+GB/s column counts the real bytes). One sequence of BENCH_L tokens, a
+single decode query, BENCH_ITERS launches, no warmup.
+
+Env: BENCH_L, BENCH_ITERS, BENCH_HQ (24), BENCH_HK (4), BENCH_D (256),
+     BENCH_PAGE (128), BENCH_CONTIG=1, BENCH_SEED, FA2_NUM_SPLITS (0 = FA2's
+     own split heuristic, the vLLM default).
+"""
+
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+
+
+def env_int(name: str, default: int) -> int:
+    return int(os.environ.get(f"BENCH_{name}", default))
+
+
+def main():
+    from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+    L, iters = env_int("L", 8192), env_int("ITERS", 50)
+    hq, hk, d, ps = env_int("HQ", 24), env_int("HK", 4), env_int("D", 256), env_int("PAGE", 128)
+    splits = int(os.environ.get("FA2_NUM_SPLITS", "0"))
+    torch.manual_seed(env_int("SEED", 1111))
+    dev = "cuda"
+
+    pages = (L + ps - 1) // ps
+    k_cache = torch.empty(pages, ps, hk, d, dtype=torch.bfloat16, device=dev)
+    v_cache = torch.empty_like(k_cache)
+    step = max(1, (64 << 20) // (ps * hk * d * 2))
+    for p0 in range(0, pages, step):
+        p1 = min(pages, p0 + step)
+        k_cache[p0:p1].normal_(0, 0.5)
+        v_cache[p0:p1].normal_(0, 0.5)
+    if os.environ.get("BENCH_CONTIG", "0") == "1":
+        table = torch.arange(pages, dtype=torch.int32, device=dev)
+    else:
+        table = torch.randperm(pages, device=dev).to(torch.int32)
+    block_table = table.view(1, pages)
+    q = torch.randn(1, hq, d, device=dev, dtype=torch.bfloat16)
+    cu_q = torch.tensor([0, 1], dtype=torch.int32, device=dev)
+    seqused_k = torch.tensor([L], dtype=torch.int32, device=dev)
+    out = torch.empty_like(q)
+    kw = dict(
+        q=q,
+        k=k_cache,
+        v=v_cache,
+        out=out,
+        cu_seqlens_q=cu_q,
+        max_seqlen_q=1,
+        seqused_k=seqused_k,
+        max_seqlen_k=L,
+        softmax_scale=1.0 / math.sqrt(d),
+        causal=True,
+        block_table=block_table,
+        fa_version=2,
+        num_splits=splits,
+    )
+
+    def run():
+        try:
+            return flash_attn_varlen_func(**kw)
+        except TypeError:  # older signature: no num_splits
+            kw.pop("num_splits", None)
+            return flash_attn_varlen_func(**kw)
+
+    torch.cuda.synchronize()
+    print(f"FA2: L={L} pages={pages}x{ps} h_q={hq} h_k={hk} d={d} bf16 KV iters={iters}", flush=True)
+    ev0, ev1 = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    ev0.record()
+    for _ in range(iters):
+        run()
+    ev1.record()
+    torch.cuda.synchronize()
+    print(f"FA2: mean {ev0.elapsed_time(ev1) * 1e3 / max(1, iters):.1f} us/iter (not ncu)", flush=True)
+    print("DONE")
+
+
+if __name__ == "__main__":
+    main()