flashinfer-ai/flashinfer

[Bug] sm12x MoE CuTe-DSL kernels are never disk-cached — cute.compile hardcodes no_cache=True, so every fresh engine re-pays 15–40 s MLIR compiles per kernel shape

Offen

#4.317 geöffnet am 02.08.2026

 (7 Kommentare) (0 Reaktionen) (1 zugewiesene Person)Python (1.031 Forks)github user discovery
good first issueneeds-triageop: moe

Repository-Metriken

Stars
 (5.756 Sterne)
PR-Merge-Metriken
 (Durchschn. Merge 11T 23h) (186 gemergte PRs in 30 T)

Beschreibung

Summary

The Blackwell sm12x fused-MoE path compiles all of its CuTe-DSL kernels through explicit cute.compile(...), which unconditionally disables the DSL's file cache. The only caching layer is three in-process Python dicts, so every engine start re-pays the full MLIR JIT cost for every kernel specialization it touches — measured at ~40 s for the dynamic kernel's first hit and ~15–18 s per small-m micro-kernel specialization. In a serving context (vLLM --moe-backend=flashinfer_b12x) these compiles land mid-request and freeze the whole engine (vllm#47458: ~17.5 s stop-the-world stalls, all in-flight requests frozen including mid-decode).

CUTE_DSL_CACHE_DIR / CUTE_DSL_DISABLE_FILE_CACHING have no effect on this path — the gate that would consult them is short-circuited before they are read.

Environment

  • flashinfer 0.6.12, 0.6.13, and 0.6.15.post1 (all verified; moe_dispatch.py cache structure unchanged across them)
  • nvidia-cutlass-dsl 4.5.0 and 4.6.0 (both verified)
  • SM120: RTX 5070 Ti (×2, TP/PP), RTX 5090 (unsharded), RTX PRO 6000 Blackwell — reproduced on all three in the vllm#47458 thread
  • Model: nvidia/Qwen3.6-35B-A3B-NVFP4 via vLLM --moe-backend=flashinfer_b12x

Mechanism (exact code paths)

  1. The DSL hardcodes no_cache=True for explicit compiles. CompileCallable._compile sets kwargs["no_cache"] = True unconditionally:

    • nvidia-cutlass-dsl 4.5.0: base_dsl/compiler.py:604-605
    • nvidia-cutlass-dsl 4.6.0: base_dsl/compiler.py:1319
  2. The file cache is gated behind if not no_cache — both the load and the dump — so the env vars are never even consulted for explicit compiles:

    • 4.5.0: base_dsl/dsl.py:1645 (load), :1736 (dump)
    • 4.6.0: base_dsl/dsl.py:2004 (load), :2105 (dump); in-memory jit_cache lookups likewise skipped under no_cache at :2012, :2293
  3. The sm12x MoE dispatch only ever compiles through cute.compile, with plain in-process dicts on top — nothing reaches disk:

    • 0.6.12/0.6.13: blackwell_sm12x/moe_dispatch.py:676/912/1675 (compile sites); _STATIC/_MICRO/_DYNAMIC_KERNEL_CACHE
    • 0.6.15.post1: compile sites at :697/:942/:1723; caches at :486/:733/:1524
  4. Compounding factor (tracked separately as #3836): the micro-kernel cache_key includes both m and workspace.max_rows (0.6.15.post1 moe_dispatch.py:770-791, fields at :774/:778), so workspace growth invalidates already-compiled small-batch entries even within a process lifetime. With no disk layer underneath, each invalidation is a full 15 s+ recompile, not a reload.

Verification

Behavioral check from the vllm#47458 thread (credit @waynehacking8, on 0.6.12 / DSL 4.5.0):

  • cute.compile of a trivial @cute.jit function leaves CUTE_DSL_CACHE_DIR empty.
  • Calling the same jit function directly (implicit compile path) writes cute_dsl_<hash>.mlir, and a fresh process then logs JIT cache hit IN-FILE.
  • Only the implicit direct-call path caches; flashinfer's MoE dispatch never uses it.

Source-level re-verification on 0.6.15.post1 / DSL 4.6.0 (2026-08-02, inside the vllm/vllm-openai nightly image): every line cite in the Mechanism section above was confirmed against the installed packages — the hardcoded no_cache=True, the if not no_cache gates around both file-cache load and dump, the three cute.compile sites, the three in-process dicts, and m/max_rows still present in the micro-kernel cache key.

Measured first-hit compile costs (RTX PRO 6000, synthetic MoE E=32/topk=4 driving B12xMoEWrapper.run exactly as vLLM's FlashInferB12xExperts does):

kernel path first hit subsequent (same key)
dynamic (m-independent key) ~40 s 1–3 ms
micro m=1 14.9 s ms
micro m=2 15.2 s ms
micro m=64 18.3 s ms

Production symptom: vllm#47458 — py-spy captured 16 consecutive samples over 8.35 s pinned in b12x_fused_moe (flashinfer/fused_moe/cute_dsl/b12x_moe.py:166) during a live ~17.5 s engine-wide stall; perf showed _cutlass_ir.cpython-312.so frames present only during stalls.

Why this deserves a disk layer even after #3836

#3836 stops rekeying (recompiles within one process). But serving engines restart — deploys, config changes, crash recovery, autoscaling — and each fresh process re-pays the entire compile set: ~40 s dynamic + 15–18 s per small-shape class. For a model like Qwen3.6-35B-A3B this is 1–2 minutes of mid-serving stalls per engine lifetime that a disk cache would reduce to milliseconds. The artifacts are deterministic products of (kernel source, shape/config tuple, DSL version, arch) — exactly what file caches are for.

Suggested fix directions

Either would resolve it; the first is contained entirely in flashinfer:

  1. flashinfer grows its own disk layer for sm12x MoE kernels: serialize the compiled artifact keyed by the existing cache_key tuples + flashinfer version + nvidia-cutlass-dsl version + SM arch, under FLASHINFER_WORKSPACE_DIR (or a new env), loading into _STATIC/_MICRO/_DYNAMIC_KERNEL_CACHE on miss. The DSL's own compile_and_cache machinery demonstrates the artifact is serializable (the implicit path already writes/reads .mlir + compiled objects).
  2. Upstream: nvidia-cutlass-dsl honors its file cache under explicit cute.compile (stop hardcoding no_cache=True, or add an opt-in kwarg flashinfer can pass).

A third, cheaper-but-partial option is an official warm-up API (the attempt in vllm#47599 was closed unmerged): it moves the cost to boot but still re-pays it every process start, and needs the largest-workspace-first ordering to avoid #3836 rekeying — a disk cache composes with, rather than replaces, that.

Cross-references

  • vllm#47458 — serving-stall report this splits from (mechanism established in-thread by @waynehacking8; py-spy confirmation by me)
  • flashinfer#3836 — the workspace-rekeying half (open)
  • vllm#47599 — warm-up mitigation PR (closed, unmerged)
  • flashinfer#3807 — earlier sm12x smem-overflow issue in the same dispatch (fixed)

Contributor Guide