Use HND layout

bench/prefill_decode_bench.py

@@ -0,0 +1,399 @@
+# type: ignore
+#!/usr/bin/env python3
+"""Disaggregated prefill-decode benchmark for exo.
+
+Measures throughput when a vLLM node handles prefill (KV generation + transfer)
+and an MLX node handles decode (token generation).
+
+Requires a cluster with at least one CUDA/vLLM node and one Apple Silicon/MLX node.
+
+Usage:
+    uv run python prefill_decode_bench.py --model <decode-model> --pp 2048,8192 --tg 128
+    uv run python prefill_decode_bench.py --model <decode-model> --prefill-model <vllm-model> --pp 4096 --tg 128
+    uv run python prefill_decode_bench.py --model <model-id> --pp 2048 --tg 128 --no-overlapping
+    uv run python prefill_decode_bench.py --model <model-id> --pp 2048 --tg 128 --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import itertools
+import json
+import sys
+import time
+from statistics import mean
+from typing import Any
+
+from exo_bench import (
+    PromptSizer,
+    format_peak_memory,
+    load_tokenizer_for_bench,
+    parse_int_list,
+)
+from harness import (
+    ExoClient,
+    ExoHttpError,
+    add_common_instance_args,
+    instance_id_from_instance,
+    nodes_used_in_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+    settle_and_fetch_placements,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
+from loguru import logger
+
+
+def find_vllm_placement(placements: list[dict[str, Any]]) -> dict[str, Any] | None:
+    for p in placements:
+        if "vllm" in str(p.get("instance_meta", "")).lower():
+            return p
+    return None
+
+
+def find_mlx_placement(
+    placements: list[dict[str, Any]], decode_meta: str
+) -> dict[str, Any] | None:
+    for p in placements:
+        meta = str(p.get("instance_meta", "")).lower()
+        if decode_meta == "both":
+            if "ring" in meta or "jaccl" in meta:
+                return p
+        elif decode_meta in meta:
+            return p
+    return None
+
+
+def run_one_disaggregated(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+) -> tuple[dict[str, Any], int]:
+    content, pp_tokens = prompt_sizer.build(pp_hint)
+    payload: dict[str, Any] = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": content}],
+        "stream": False,
+        "max_tokens": tg,
+    }
+
+    t0 = time.perf_counter()
+    out = client.post_bench_disaggregated(payload)
+    elapsed = time.perf_counter() - t0
+
+    stats = out.get("generation_stats")
+
+    choices = out.get("choices") or [{}]
+    message = choices[0].get("message", {}) if choices else {}
+    text = message.get("content") or ""
+    preview = text[:200] if text else ""
+
+    return {
+        "elapsed_s": elapsed,
+        "output_text_preview": preview,
+        "stats": stats,
+    }, pp_tokens
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(
+        prog="prefill-decode-bench",
+        description="Benchmark disaggregated prefill-decode (vLLM prefill → MLX decode).",
+    )
+    add_common_instance_args(ap)
+    ap.add_argument(
+        "--pp",
+        nargs="+",
+        required=True,
+        help="Prompt-size hints (ints, must be >1000). Accepts commas.",
+    )
+    ap.add_argument(
+        "--tg",
+        nargs="+",
+        required=True,
+        help="Generation lengths (ints). Accepts commas.",
+    )
+    ap.add_argument(
+        "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
+    )
+    ap.add_argument(
+        "--warmup",
+        type=int,
+        default=0,
+        help="Warmup runs per placement (uses first pp/tg).",
+    )
+    ap.add_argument(
+        "--prefill-model",
+        default=None,
+        help="Model for the vLLM prefill node (defaults to --model if not set).",
+    )
+    ap.add_argument(
+        "--no-overlapping",
+        action="store_true",
+        help="Use batch KV transfer instead of streaming (overlapping).",
+    )
+    ap.add_argument(
+        "--decode-instance-meta",
+        choices=["ring", "jaccl", "both"],
+        default="ring",
+        help="Instance meta for the decode (MLX) node.",
+    )
+    ap.add_argument(
+        "--json-out",
+        default="bench/prefill_decode_results.json",
+        help="Write raw per-run results JSON to this path.",
+    )
+    ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
+    ap.add_argument(
+        "--dry-run", action="store_true", help="List selected placements and exit."
+    )
+    ap.add_argument(
+        "--all-combinations",
+        action="store_true",
+        help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
+    )
+    args = ap.parse_args()
+
+    pp_list = parse_int_list(args.pp)
+    tg_list = parse_int_list(args.tg)
+    if not pp_list or not tg_list:
+        logger.error("pp and tg lists must be non-empty")
+        return 2
+    for pp in pp_list:
+        if pp <= 1000:
+            logger.error(f"pp={pp} must be >1000 (remote prefill requires >1000 uncached tokens)")
+            return 2
+    if args.repeat <= 0:
+        logger.error("--repeat must be >= 1")
+        return 2
+
+    use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
+    if use_combinations:
+        logger.info(f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs")
+    else:
+        logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
+
+    overlapping = not args.no_overlapping
+    logger.info(f"KV transfer mode: {'overlapping (streaming)' if overlapping else 'batch'}")
+
+    client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+
+    decode_short_id, decode_full_id = resolve_model_short_id(
+        client, args.model, force_download=args.force_download
+    )
+
+    prefill_model_arg = args.prefill_model or args.model
+    if args.prefill_model:
+        prefill_short_id, prefill_full_id = resolve_model_short_id(
+            client, args.prefill_model, force_download=args.force_download
+        )
+    else:
+        prefill_short_id, prefill_full_id = decode_short_id, decode_full_id
+
+    tokenizer = load_tokenizer_for_bench(decode_full_id)
+    if tokenizer is None:
+        raise RuntimeError("[prefill-decode-bench] tokenizer load failed")
+
+    try:
+        prompt_sizer = PromptSizer(tokenizer)
+    except Exception:
+        logger.error("[prefill-decode-bench] tokenizer usable but prompt sizing failed")
+        raise
+
+    original_instance_meta = args.instance_meta
+    original_max_nodes = args.max_nodes
+
+    args.instance_meta = "vllm"
+    args.min_nodes = 1
+    args.max_nodes = 1
+    vllm_placements = settle_and_fetch_placements(
+        client, prefill_full_id, args, settle_timeout=args.settle_timeout
+    )
+
+    args.instance_meta = args.decode_instance_meta
+    args.min_nodes = 1
+    args.max_nodes = original_max_nodes
+    mlx_placements = settle_and_fetch_placements(
+        client, decode_full_id, args, settle_timeout=args.settle_timeout
+    )
+
+    args.instance_meta = original_instance_meta
+
+    vllm_placement = find_vllm_placement(vllm_placements)
+    mlx_placement = find_mlx_placement(mlx_placements, args.decode_instance_meta)
+
+    if not vllm_placement:
+        logger.error("No vLLM placement found. Need a CUDA node for prefill.")
+        return 1
+    if not mlx_placement:
+        logger.error(f"No MLX ({args.decode_instance_meta}) placement found for decode.")
+        return 1
+
+    vllm_instance = vllm_placement["instance"]
+    mlx_instance = mlx_placement["instance"]
+    vllm_instance_id = instance_id_from_instance(vllm_instance)
+    mlx_instance_id = instance_id_from_instance(mlx_instance)
+    vllm_meta = str(vllm_placement.get("instance_meta", ""))
+    mlx_meta = str(mlx_placement.get("instance_meta", ""))
+    mlx_sharding = str(mlx_placement.get("sharding", ""))
+    mlx_nodes = nodes_used_in_instance(mlx_instance)
+
+    logger.info("=" * 80)
+    logger.info(f"PREFILL (vLLM):  {vllm_meta} / {prefill_short_id} ({prefill_full_id}) / instance_id={vllm_instance_id}")
+    logger.info(f"DECODE  (MLX):   {mlx_meta} / {mlx_sharding} / nodes={mlx_nodes} / {decode_short_id} ({decode_full_id}) / instance_id={mlx_instance_id}")
+    logger.info(f"Overlapping: {overlapping}")
+
+    if args.dry_run:
+        return 0
+
+    settle_deadline = (
+        time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+    )
+
+    logger.info("Planning phase: checking downloads...")
+    download_duration_s = run_planning_phase(
+        client,
+        prefill_full_id,
+        vllm_placement,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    download_duration_mlx = run_planning_phase(
+        client,
+        decode_full_id,
+        mlx_placement,
+        args.danger_delete_downloads,
+        args.timeout,
+        settle_deadline,
+    )
+    if download_duration_s is not None:
+        logger.info(f"Download (vLLM): {download_duration_s:.1f}s")
+    if download_duration_mlx is not None:
+        logger.info(f"Download (MLX): {download_duration_mlx:.1f}s")
+
+    # Create vLLM instance first (starts prefill server)
+    logger.info("Creating vLLM (prefill) instance...")
+    client.request_json("POST", "/instance", body={"instance": vllm_instance})
+    try:
+        wait_for_instance_ready(client, vllm_instance_id)
+    except (RuntimeError, TimeoutError) as e:
+        logger.error(f"Failed to initialize vLLM instance: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        return 1
+    logger.info("vLLM (prefill) instance ready")
+
+    # Create MLX instance (decode)
+    logger.info("Creating MLX (decode) instance...")
+    client.request_json("POST", "/instance", body={"instance": mlx_instance})
+    try:
+        wait_for_instance_ready(client, mlx_instance_id)
+    except (RuntimeError, TimeoutError) as e:
+        logger.error(f"Failed to initialize MLX instance: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{mlx_instance_id}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        return 1
+    logger.info("MLX (decode) instance ready")
+
+    time.sleep(2)
+
+    all_rows: list[dict[str, Any]] = []
+
+    try:
+        for i in range(args.warmup):
+            run_one_disaggregated(
+                client, decode_full_id, pp_list[0], tg_list[0], prompt_sizer
+            )
+            logger.debug(f"  warmup {i + 1}/{args.warmup} done")
+
+        if use_combinations:
+            pp_tg_pairs = list(itertools.product(pp_list, tg_list))
+        else:
+            pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
+
+        for pp, tg in pp_tg_pairs:
+            logger.info(f"--- pp={pp} tg={tg} ---")
+            runs: list[dict[str, Any]] = []
+            for r in range(args.repeat):
+                time.sleep(3)
+
+                try:
+                    row, actual_pp_tokens = run_one_disaggregated(
+                        client, decode_full_id, pp, tg, prompt_sizer
+                    )
+                except Exception as e:
+                    logger.error(e)
+                    continue
+
+                row.update(
+                    {
+                        "prefill_model_short_id": prefill_short_id,
+                        "prefill_model_id": prefill_full_id,
+                        "prefill_instance_meta": vllm_meta,
+                        "prefill_instance_id": vllm_instance_id,
+                        "decode_model_short_id": decode_short_id,
+                        "decode_model_id": decode_full_id,
+                        "decode_instance_meta": mlx_meta,
+                        "decode_sharding": mlx_sharding,
+                        "decode_nodes": mlx_nodes,
+                        "decode_instance_id": mlx_instance_id,
+                        "overlapping": overlapping,
+                        "pp_tokens": actual_pp_tokens,
+                        "tg": tg,
+                        "repeat_index": r,
+                    }
+                )
+                runs.append(row)
+                all_rows.append(row)
+
+            if runs:
+                prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
+                gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
+                gtok = mean(x["stats"]["generation_tokens"] for x in runs)
+                peak = mean(
+                    x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
+                )
+                avg_elapsed = mean(x["elapsed_s"] for x in runs)
+
+                logger.info(
+                    f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
+                    f"prompt_tokens={ptok} gen_tokens={gtok}    "
+                    f"peak_memory={format_peak_memory(peak)}    "
+                    f"avg_elapsed={avg_elapsed:.2f}s\n"
+                )
+            time.sleep(2)
+    finally:
+        try:
+            client.request_json("DELETE", f"/instance/{mlx_instance_id}")
+        except ExoHttpError as e:
+            if e.status != 404:
+                raise
+        try:
+            client.request_json("DELETE", f"/instance/{vllm_instance_id}")
+        except ExoHttpError as e:
+            if e.status != 404:
+                raise
+        wait_for_instance_gone(client, mlx_instance_id)
+        wait_for_instance_gone(client, vllm_instance_id)
+        logger.debug("Deleted both instances")
+
+    if args.stdout:
+        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+    elif args.json_out:
+        with open(args.json_out, "w", encoding="utf-8") as f:
+            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+        logger.debug(f"\nWrote results JSON: {args.json_out}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

justfile

@@ -21,12 +21,33 @@ sync-cuda:
     if command -v nvidia-smi &>/dev/null; then
         arch=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d ' ')
         export TORCH_CUDA_ARCH_LIST="$arch"
+        export FLASHINFER_CUDA_ARCH_LIST="${arch}a"
         export VLLM_TARGET_DEVICE=cuda
     fi
     uv pip install "cmake>=3.26.1" ninja "packaging>=24.2" "setuptools>=77.0.3,<81.0.0" "setuptools-scm>=8.0" wheel jinja2
     find ~/.cache/uv/git-v0 -name CMakeCache.txt -delete 2>/dev/null || true
     uv sync --extra cuda
 
+build-flashinfer-sm121:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    FLASHINFER_VERSION="v0.6.6"
+    BUILD_DIR="/tmp/flashinfer-build"
+    EXO_VENV="$(pwd)/.venv"
+    rm -rf "$BUILD_DIR"
+    git clone --depth 1 --branch "$FLASHINFER_VERSION" --recurse-submodules --shallow-submodules https://github.com/flashinfer-ai/flashinfer.git "$BUILD_DIR"
+    cd "$BUILD_DIR"
+    export FLASHINFER_CUDA_ARCH_LIST="12.1a"
+    export TORCH_CUDA_ARCH_LIST="12.1"
+    export FLASHINFER_ENABLE_AOT=1
+    UV="$(command -v uv || echo "$HOME/.local/bin/uv")"
+    VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation
+    echo "Building AOT cubin wheel..."
+    cd flashinfer-cubin && VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation && cd ..
+    echo "Building JIT cache wheel..."
+    cd flashinfer-jit-cache && VIRTUAL_ENV="$EXO_VENV" "$UV" pip install . --no-build-isolation && cd ..
+    echo "FlashInfer built from source with SM121 AOT kernels"
+
 sync-clean:
     uv sync --all-packages --force-reinstall --no-cache
 

pyproject.toml

@@ -134,8 +134,8 @@ environments = [
   "sys_platform == 'darwin'",
   "sys_platform == 'linux' and platform_machine == 'aarch64'",
 ]
-no-binary-package = ["vllm"]
-no-build-isolation-package = ["vllm"]
+no-binary-package = ["vllm", "flashinfer-python"]
+no-build-isolation-package = ["vllm", "flashinfer-python"]
 extra-build-dependencies = { vllm = [
   "cmake>=3.26.1",
   "ninja",

src/exo/disaggregated/batch_connector.py

@@ -66,12 +66,14 @@ class BatchConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[report
             return
 
         if slot_mapping is not None:
+            from exo.disaggregated.streaming_connector import _to_nhd
+
             if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
-                k_all = kv_layer[0]  # pyright: ignore[reportAny]
-                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[1])  # pyright: ignore[reportAny]
             else:
-                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
-                v_all = kv_layer[:, 1]  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[:, 0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[:, 1])  # pyright: ignore[reportAny]
             k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
             v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
             valid = slot_mapping >= 0  # pyright: ignore[reportAny]

src/exo/disaggregated/streaming_connector.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import queue
 import re
 from dataclasses import dataclass
@@ -16,6 +17,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (  # pyright: igno
 _LAYER_RE = re.compile(r"layers\.(\d+)\.")
 
 
+def _to_nhd(t: torch.Tensor) -> torch.Tensor:
+    if os.environ.get("VLLM_KV_CACHE_LAYOUT", "HND") == "HND" and t.dim() == 3:
+        return t.permute(1, 0, 2)
+    return t
+
+
 def _to_bf16(t: torch.Tensor) -> torch.Tensor:
     if t.dtype == torch.uint8:
         t = t.view(torch.float8_e4m3fn)  # type: ignore
@@ -94,11 +101,11 @@ class StreamingConnector(KVConnectorBase_V1, SupportsHMA):  # pyright: ignore[re
 
         if slot_mapping is not None:
             if kv_layer.shape[0] == 2:  # pyright: ignore[reportAny]
-                k_all = kv_layer[0]  # pyright: ignore[reportAny]
-                v_all = kv_layer[1]  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[1])  # pyright: ignore[reportAny]
             else:
-                k_all = kv_layer[:, 0]  # pyright: ignore[reportAny]
-                v_all = kv_layer[:, 1]  # pyright: ignore[reportAny]
+                k_all = _to_nhd(kv_layer[:, 0])  # pyright: ignore[reportAny]
+                v_all = _to_nhd(kv_layer[:, 1])  # pyright: ignore[reportAny]
             k_flat = k_all.reshape(-1, *k_all.shape[-2:])  # pyright: ignore[reportAny]
             v_flat = v_all.reshape(-1, *v_all.shape[-2:])  # pyright: ignore[reportAny]
             valid = slot_mapping >= 0  # pyright: ignore[reportAny]

src/exo/worker/engines/vllm/growable_cache.py

@@ -1,3 +1,6 @@
+import os
+import sys
+
 import torch
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
@@ -46,6 +49,7 @@ def patch_vllm() -> None:
     _patch_get_computed_blocks()
     _patch_moe_sum()
     _patch_marlin_w2_thread_config()
+    _patch_flashinfer_sinks()
     logger.info("vLLM growable KV cache patch applied")
 
 
@@ -123,8 +127,17 @@ def _patch_initialize_kv_cache_tensors() -> None:
 
 def _patch_initialize_from_config() -> None:
     from vllm.v1.worker.gpu_worker import Worker
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
     original = Worker.initialize_from_config
+    original_init_attn = GPUModelRunner.initialize_attn_backend
+
+    def idempotent_init_attn(self: "GPUModelRunner", *args: object, **kwargs: object) -> None:
+        if len(self.attn_groups) > 0:
+            return
+        original_init_attn(self, *args, **kwargs)
+
+    GPUModelRunner.initialize_attn_backend = idempotent_init_attn  # type: ignore[reportAttributeAccessIssue]
 
     def patched(self: "Worker", kv_cache_config: "object") -> None:
         original(self, kv_cache_config)
@@ -442,3 +455,191 @@ def _patch_get_computed_blocks() -> None:
         return self.empty_kv_cache_blocks, num_matched
 
     KVCacheManager.get_computed_blocks = patched  # type: ignore[reportAttributeAccessIssue]
+
+
+def _patch_flashinfer_sinks() -> None:
+    cap = torch.cuda.get_device_capability()
+    if cap[0] < 12:
+        return
+
+    _patch_is_blackwell_class()
+    _patch_flashinfer_compilation_context()
+    _patch_flashinfer_fmha_sm12x()
+    _patch_triton_ptxas_path()
+    logger.info("FlashInfer SM12x Blackwell-class patches applied")
+
+
+def _patch_is_blackwell_class() -> None:
+    from vllm.platforms.cuda import CudaPlatform
+
+    original_is_family = CudaPlatform.is_device_capability_family.__func__  # type: ignore[reportAttributeAccessIssue]
+
+    @classmethod  # type: ignore[reportArgumentType]
+    def _patched_is_family(cls: type, capability: int, device_id: int = 0) -> bool:
+        if capability == 100:
+            current = cls.get_device_capability(device_id=device_id)
+            if current is not None and current.major in (10, 11, 12):
+                return True
+        return original_is_family(cls, capability, device_id)
+
+    CudaPlatform.is_device_capability_family = _patched_is_family  # type: ignore[reportAttributeAccessIssue]
+
+    try:
+        from vllm.v1.attention.backends.flashinfer import FlashInferBackend, FlashInferMetadataBuilder
+    except ImportError:
+        return
+
+    import vllm.utils.flashinfer as fi_mod
+
+    def _supports_sink_sm12x(cls: type) -> bool:
+        return True
+    FlashInferBackend.supports_sink = classmethod(_supports_sink_sm12x)  # type: ignore[reportAttributeAccessIssue]
+
+    original_init = FlashInferMetadataBuilder.__init__
+
+    def _patched_init(self: object, *args: object, **kwargs: object) -> None:
+        try:
+            original_init(self, *args, **kwargs)
+        except NotImplementedError as e:
+            if "attention sinks" not in str(e):
+                raise
+            logger.info("Bypassing FlashInfer sinks check for SM12x (native compiled attention)")
+            _orig = fi_mod.can_use_trtllm_attention
+            fi_mod.can_use_trtllm_attention = lambda *a, **k: True  # type: ignore[reportAttributeAccessIssue]
+            for mod in list(sys.modules.values()):
+                if mod is not None and mod is not fi_mod:
+                    for attr in list(vars(mod)):
+                        if vars(mod)[attr] is _orig:
+                            setattr(mod, attr, fi_mod.can_use_trtllm_attention)
+            original_init(self, *args, **kwargs)
+            fi_mod.can_use_trtllm_attention = _orig  # type: ignore[reportAttributeAccessIssue]
+            for mod in list(sys.modules.values()):
+                if mod is not None and mod is not fi_mod:
+                    for attr in list(vars(mod)):
+                        if vars(mod)[attr] is fi_mod.can_use_trtllm_attention:
+                            setattr(mod, attr, _orig)
+            setattr(self, 'use_trtllm_decode_attention', False)
+
+    FlashInferMetadataBuilder.__init__ = _patched_init  # type: ignore[reportAttributeAccessIssue]
+
+    original_build = FlashInferMetadataBuilder.build
+
+    def _patched_build(self: object, *args: object, **kwargs: object) -> object:
+        real_sinks = getattr(self, 'has_sinks', False)
+        if real_sinks:
+            setattr(self, 'has_sinks', False)
+        try:
+            return original_build(self, *args, **kwargs)
+        finally:
+            if real_sinks:
+                setattr(self, 'has_sinks', True)
+
+    FlashInferMetadataBuilder.build = _patched_build  # type: ignore[reportAttributeAccessIssue]
+
+
+def _patch_flashinfer_compilation_context() -> None:
+    try:
+        from flashinfer.compilation_context import CompilationContext
+    except ImportError:
+        return
+
+    original_get_flags = CompilationContext.get_nvcc_flags_list
+
+    def _patched_get_flags(self: object, supported_major_versions: list[int] | None = None) -> list[str]:
+        if supported_major_versions and 10 in supported_major_versions and 12 not in supported_major_versions:
+            supported_major_versions = list(supported_major_versions) + [12]
+        return original_get_flags(self, supported_major_versions)
+
+    CompilationContext.get_nvcc_flags_list = _patched_get_flags  # type: ignore[reportAttributeAccessIssue]
+    logger.info("Patched FlashInfer CompilationContext: SM12x added to SM10x version filters")
+
+
+def _patch_flashinfer_fmha_sm12x() -> None:
+    try:
+        import flashinfer
+        fi_include = os.path.join(os.path.dirname(flashinfer.__file__), "data", "include", "flashinfer", "trtllm")
+    except ImportError:
+        return
+
+    common_h = os.path.join(fi_include, "common.h")
+    runner_cuh = os.path.join(fi_include, "fmha", "fmhaRunner.cuh")
+    kernels_cuh = os.path.join(fi_include, "fmha", "fmhaKernels.cuh")
+
+    if not os.path.exists(common_h):
+        return
+
+    with open(common_h) as f:
+        content = f.read()
+    needs_patch = False
+    if "kSM_120" not in content:
+        content = content.replace(
+            "constexpr int32_t kSM_103 = 103;",
+            "constexpr int32_t kSM_103 = 103;\nconstexpr int32_t kSM_120 = 120;",
+        )
+        needs_patch = True
+    if "sm > 120 && sm < 130" not in content:
+        content = content.replace(
+            "return sm_major * 10 + sm_minor;",
+            "int sm = sm_major * 10 + sm_minor;\n  if (sm > 120 && sm < 130) sm = 120;\n  return sm;",
+        )
+        needs_patch = True
+    if needs_patch:
+        with open(common_h, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer common.h: kSM_120 + SM121->120 normalization")
+
+    with open(runner_cuh) as f:
+        content = f.read()
+    if "kSM_120" not in content:
+        content = content.replace(
+            "mSM == kSM_100 || mSM == kSM_103",
+            "mSM == kSM_100 || mSM == kSM_103 || mSM == kSM_120",
+        )
+        with open(runner_cuh, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer fmhaRunner.cuh: added kSM_120")
+
+    with open(kernels_cuh) as f:
+        content = f.read()
+    if "kSM_120" not in content:
+        content = content.replace(
+            "if (gpuSM == kSM_103) {\n    return kernelSM == kSM_100f || kernelSM == kSM_103;",
+            "if (gpuSM == kSM_103 || gpuSM == kSM_120) {\n    return kernelSM == kSM_100f || kernelSM == kSM_103 || kernelSM == kSM_120;",
+        )
+        with open(kernels_cuh, "w") as f:
+            f.write(content)
+        logger.info("Patched flashinfer fmhaKernels.cuh: added kSM_120 kernel dispatch")
+
+    fi_csrc = os.path.join(os.path.dirname(flashinfer.__file__), "data", "csrc")
+    moe_launcher = os.path.join(fi_csrc, "trtllm_fused_moe_kernel_launcher.cu")
+    if os.path.exists(moe_launcher):
+        with open(moe_launcher) as f:
+            content = f.read()
+        if "major, 10) <<" in content and "major, 10) || std::get<0>(device_props) == 12)" not in content:
+            content = content.replace(
+                'TVM_FFI_ICHECK_EQ(major, 10) << "MoE kernel requires 10.x',
+                'TVM_FFI_ICHECK(major == 10 || major == 12) << "MoE kernel requires 10.x/12.x',
+            )
+            content = content.replace(
+                'std::get<0>(device_props) == 10)',
+                'std::get<0>(device_props) == 10 || std::get<0>(device_props) == 12)',
+            )
+            with open(moe_launcher, "w") as f:
+                f.write(content)
+            cache_dir = os.path.expanduser("~/.cache/flashinfer/0.6.6/121a/cached_ops/fused_moe_trtllm_sm100")
+            if os.path.isdir(cache_dir):
+                import shutil
+                shutil.rmtree(cache_dir)
+            logger.info("Patched flashinfer trtllm_fused_moe_kernel_launcher.cu: added SM12x support")
+
+
+def _patch_triton_ptxas_path() -> None:
+    import os
+    import shutil
+
+    if os.environ.get("TRITON_PTXAS_PATH"):
+        return
+    ptxas = shutil.which("ptxas")
+    if ptxas:
+        os.environ["TRITON_PTXAS_PATH"] = ptxas
+        logger.info(f"Set TRITON_PTXAS_PATH={ptxas}")

src/exo/worker/engines/vllm/vllm_generator.py

@@ -599,6 +599,7 @@ def load_vllm_engine(
     if has_mamba:
         backends = ["FLASH_ATTN", "TRITON_ATTN"]
     else:
+        # backends = ["FLASH_ATTN", "TRITON_ATTN"]
         backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"]
 
     engine: LLMEngine | None = None
@@ -612,11 +613,13 @@ def load_vllm_engine(
                 load_format="fastsafetensors",
                 enable_prefix_caching=False,
                 attention_backend=backend,
-                compilation_config={"cudagraph_mode": "none"},
+                compilation_config={},
                 disable_log_stats=True,
                 max_num_batched_tokens=4096,
                 kv_transfer_config=kv_transfer_config,  # type: ignore
                 disable_hybrid_kv_cache_manager=False,
+                kv_cache_dtype="auto",
+                **({"moe_backend": "flashinfer_cutlass"} if is_mxfp4 else {}),
             )
 
             set_weight_loading_callback(on_layer_loaded)

src/exo/worker/runner/bootstrap.py

@@ -103,7 +103,10 @@ def entrypoint(
     try:
         if isinstance(bound_instance.instance, VllmInstance):
             os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
+            import torch
+            cap = torch.cuda.get_device_capability()
+            layout = "HND" if cap[0] >= 12 else "NHD"
+            os.environ["VLLM_KV_CACHE_LAYOUT"] = layout
             # os.environ["VLLM_BATCH_INVARIANT"] = "1"
             _ensure_cuda_libs()
             _ensure_tiktoken_encodings()