From 339f3f10a33972421aba32778dac8adca7fcb96a Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Wed, 1 Apr 2026 19:58:12 +0100 Subject: [PATCH] Use HND layout --- bench/prefill_decode_bench.py | 399 ++++++++++++++++++ justfile | 21 + pyproject.toml | 4 +- src/exo/disaggregated/batch_connector.py | 10 +- src/exo/disaggregated/streaming_connector.py | 15 +- src/exo/worker/engines/vllm/growable_cache.py | 201 +++++++++ src/exo/worker/engines/vllm/vllm_generator.py | 5 +- src/exo/worker/runner/bootstrap.py | 5 +- 8 files changed, 648 insertions(+), 12 deletions(-) create mode 100644 bench/prefill_decode_bench.py diff --git a/bench/prefill_decode_bench.py b/bench/prefill_decode_bench.py new file mode 100644 index 00000000..902b9d50 --- /dev/null +++ b/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 --pp 2048,8192 --tg 128 + uv run python prefill_decode_bench.py --model --prefill-model --pp 4096 --tg 128 + uv run python prefill_decode_bench.py --model --pp 2048 --tg 128 --no-overlapping + uv run python prefill_decode_bench.py --model --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()) diff --git a/justfile b/justfile index e21453d0..5864ff36 100644 --- a/justfile +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 6d7cd6ea..b799d465 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/exo/disaggregated/batch_connector.py b/src/exo/disaggregated/batch_connector.py index bbbd9a3c..6846fc66 100644 --- a/src/exo/disaggregated/batch_connector.py +++ b/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] diff --git a/src/exo/disaggregated/streaming_connector.py b/src/exo/disaggregated/streaming_connector.py index 85ec4c64..24b21236 100644 --- a/src/exo/disaggregated/streaming_connector.py +++ b/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] diff --git a/src/exo/worker/engines/vllm/growable_cache.py b/src/exo/worker/engines/vllm/growable_cache.py index 48ce1dd5..2b703d31 100644 --- a/src/exo/worker/engines/vllm/growable_cache.py +++ b/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}") diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py index d75b2d6c..ab5bb192 100644 --- a/src/exo/worker/engines/vllm/vllm_generator.py +++ b/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) diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index c3c60f53..b111354d 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/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()