amd-agi/geak/apply-flydsl-moe-to-vllm

When to use

A validated FlyDSL rewrite of an int4 W4A16 (GPTQ/AWQ, no-zp) fused-MoE grouped GEMM exists (see the sibling kernel skill flydsl_rewrite_quantized_moe, isolated ~3.6x) and must be landed end-to-end in a live vLLM server on MI300X (gfx942) TP4 — i.e. the isolated win has to survive serving. Use when the apply hits HIP OOM / `Engine core initialization failed` during `determine_available_memory`, or when an A/B shows the new kernel slower than Triton e2e.

Mechanism

A FlyDSL grouped-GEMM that is fast in isolation fails its first e2e apply in two ways, each with a fix:

• HIP OOM at startup — the converted weights/scales are kept beside the originals, or re-materialized

per forward ([E,N,K] int16 unpack), doubling/spiking MoE memory. Caused by a runtime-path conversion whose cache key includes the activation `data_ptr()` (reallocated every step → cache misses forever → re-converts every forward). Fix: convert once at load time, in place, same-byte, keyed by the new weight `data_ptr()` only. GOTCHA (verified 2026-06-26, the dominant real hog): re-home BOTH the weight AND the scale param — layer.w*_weight_packed.data = w_flat AND layer.w*_weight_scale.data = s_flat. Re-homing only the weight (a common miss) leaves the original [E,N,G] bf16 scale alive on the layer while the converted scale is also cached → scales DUPLICATED ≈ +246 MiB/layer × N layers ≈ +14.5 GiB → Available KV collapses (e.g. 5.96 GiB vs 17.16 needed for 262144) → Engine core initialization failed. This — NOT the repeat_interleave A2 pre-gather — was the binding KV constraint on Kimi-K2.6 (measured: repeat_interleave cost only ~0.23 GiB; the scale leak cost ~14.5 GiB).

• Eager A/B is slower (−22…−34%) despite isolated 3.6x — the shim issues more launches/layer than

Triton's single fused op; eager per-launch host latency dominates. Fix: run with HIP graphs (drop --enforce-eager); capture replays the sequence at zero launch cost and the faster GEMMs surface.

• **Graph capture DEADLOCKS under TP>1 — server never becomes healthy (e2e_delta=null, hung on a capture

batch, 0 forwards served) — this is the #1 reason a kernel wins isolated yet scores nothing e2e. The shim JIT-compiles its FlyDSL exes lazily on first call, keyed by `(weight.data_ptr, stage, M)`. vLLM captures the decode path into CUDA graphs over ~86 distinct M-bucket descriptors (PIECEWISE + FULL); ANY `(layer, M)` not already compiled compiles inside the captured region → a blocking op under capture → TP ranks desync → c10d collective timeout, capture frozen (observed hung at PIECEWISE 2/51 and at FULL 25/35 on Kimi-K2.6 TP4). The isolated unittest CANNOT catch this (single process, no capture) — only the e2e gate does, so a shim can pass isolated at high speedup and still be undeployable. Fix: precompile EVERY capture M-bucket BEFORE capture** via the vendored flydsl_capture_precompile.py (Procedure step 2b) — an exhaustive eager warmup that JITs all exes OUTSIDE the capture stream, on all TP ranks in lockstep. VERIFIED 2026-07-20 same-box controlled A/B: a 2.51x isolated shim WITHOUT the seam deadlocked (FULL 25/35, never healthy); the SAME shim WITH the seam precompiled 86/86 outside capture, captured 51/51 + 35/35, reached Application startup complete, and served requests. The seam is env-gated (VLLM_USE_FLYDSL_MOE=1) so it only affects THIS candidate's server — zero impact on any other kernel.

The golden rule: convert weights ONCE at load time; key all caches by the NEW weight `data_ptr()`, never the activation pointer; and PRECOMPILE every capture M-bucket BEFORE capture (never rely on lazy first-call JIT inside the captured region). The runtime shim then only LAUNCHES (capture-safe: routing blocks = sorted_expert_ids.numel() is a host shape read, aiter GPU sort + FlyDSL cexe are capturable, all exes already compiled+cached by the step-2b warmup before capture). The graph/inductor path needs NO custom op.

Procedure

The integration shim is vendored in this skill dir: flydsl_moe_shim.py (functions convert_layer_inplace, flydsl_fused_experts_impl). Apply it via two env-gated edits to the installed vLLM (off by default ⇒ byte-identical stock; back up the two files first so revert = restore).

1. Load-time conversion — vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

→ CompressedTensorsWNA16MoEMethod.process_weights_after_loading, AFTER the existing repack to [E,N,K//2] uint8 + scale transposes, append: ``python import os as _os if _os.environ.get("VLLM_USE_FLYDSL_MOE") == "1": import sys as _sys _d = _os.environ["FLYDSL_SHIM_DIR"] if _d not in _sys.path: _sys.path.insert(0, _d) import flydsl_moe_shim flydsl_moe_shim.convert_layer_inplace(layer) # in-place same-byte; frees originals ``

1. Runtime route — vllm/model_executor/layers/fused_moe/fused_moe.py → TOP of fused_experts_impl

(before the constraint checks), insert: ``python import os as _os if (_os.environ.get("VLLM_USE_FLYDSL_MOE") == "1" and use_int4_w4a16 and w1_zp is None and w2_zp is None): import sys as _sys _d = _os.environ["FLYDSL_SHIM_DIR"] if _d not in _sys.path: _sys.path.insert(0, _d) import flydsl_moe_shim return flydsl_moe_shim.flydsl_fused_experts_impl( hidden_states, w1, w2, topk_weights, topk_ids, inplace, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, global_num_experts=global_num_experts, expert_map=expert_map, w1_scale=w1_scale, w2_scale=w2_scale) # shim absorbs extras via **_ignored ` 2b. **Precompile-before-capture seam — MANDATORY for TP>1 graph capture (skip it and the server deadlocks at capture; see Mechanism).** The vendored flydsl_capture_precompile.py (in THIS skill dir) wraps GPUModelRunner.capture_model to run an exhaustive eager warmup over every capture descriptor BEFORE the real capture, so all FlyDSL exes JIT+cache OUTSIDE the capture stream. Add ONE env-gated edit at the END of vllm/v1/worker/gpu_model_runner.py (same FLYDSL_SHIM_DIR as steps 1–2): `python import os as _os if _os.environ.get("VLLM_USE_FLYDSL_MOE") == "1": import sys as _sys _d = _os.environ["FLYDSL_SHIM_DIR"] if _d not in _sys.path: _sys.path.insert(0, _d) import flydsl_capture_precompile # vendored in this skill dir flydsl_capture_precompile.install_capture_precompile(_sys.modules[__name__]) ` The seam is env-gated (VLLM_USE_FLYDSL_MOE=1) and self-disables otherwise, so it is scoped to THIS candidate's server ONLY and cannot affect any other kernel/candidate. It skips FULL-mode descriptors on purpose (the same M set is covered by PIECEWISE warmup; a raw FULL _dummy_run` dirties vLLM's overlapped shared-experts buffer) and resets those buffers before capture — do not remove that logic.

1. Launch (graph mode — NO `--enforce-eager`). Point FLYDSL_SHIM_DIR at THIS skill dir (the vendored

shim) and FLYDSL_ROOT at a FlyDSL checkout whose kernels AND build-fly bindings are the SAME tree: ``bash export FLYDSL_SHIM_DIR=<this_skill_dir> # contains the vendored flydsl_moe_shim.py export FLYDSL_ROOT=<flydsl_checkout> # kernels + build-fly MUST be one tree (see pitfalls) export PYTHONPATH=$FLYDSL_ROOT:$FLYDSL_ROOT/build-fly/python_packages export VLLM_USE_FLYDSL_MOE=1 export VLLM_ROCM_USE_AITER=1 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 VLLM_ROCM_USE_AITER_RMSNORM=0 # full config, NO --max-model-len cap (verified no-OOM 2026-07-23 at mem 0.9, full 262144): # scale re-home + stage-2 accumulate=True make the convert memory-neutral, so the old # --max-model-len 32768 cap is NOT needed (see Knobs). Add it back only to squeeze concurrency. vllm serve <model> -tp 4 --gpu-memory-utilization 0.9 \ --trust-remote-code --tool-call-parser kimi_k2 --enable-auto-tool-choice \ --reasoning-parser kimi_k2 --no-enable-prefix-caching --mm-encoder-tp-mode data ``

1. Verify startup (in order, no error): weight load → (silent in-place convert) → Available KV cache

(no OOM) → [flydsl-precompile] BEGIN exhaustive eager warmup ... → [flydsl-precompile] DONE eager warmup (N/N); all FlyDSL exes compiled+cached outside capture (from step 2b — one per TP rank) → Capturing CUDA graphs ... 100% (must reach 100%, not hang) → Application startup complete. Smoke a completion; confirm 0 shim failed / weights not converted lines. If capture hangs (no DONE eager warmup line, or the capture bar freezes with repeated No available shared memory broadcast block found in 60 seconds), step 2b was not wired — do NOT report a number; fix the seam, do not reject.

1. Gate: same-session A/B baseline (Triton, VLLM_USE_FLYDSL_MOE unset) vs FlyDSL, BOTH graph, identical

vllm bench serve (ISL/OSL/conc), plus GSM8K within noise. Quote same-session ratios only.

• **GATE ON THE LIVE SERVING REGIME — pick the acceptance metric from the actual ISL/OSL; do NOT assume

decode (regime-adaptive, verified 2026-08-04).** The isolated op-bench M-weighting MUST mirror the served shapes, and the same-session A/B must improve the metric the dominant regime is bound by:

• decode-dominated serving (short ISL, e.g. 1024/1024/64) → primary signal is a TPOT drop (best

run 07-26: 158.9 -> 64.7 ms, decode 2.45x). Here a candidate whose isolated bake-off is weighted toward large-M prefill (observed: M8192 = 72% of the metric) can be REAL 1.83x isolated yet REGRESS decode serving by -15.95% e2e, because its tuned tiles lose to vLLM's default at the live decode M-buckets → rejected.

• prefill-heavy serving (long ISL, e.g. 8192/1024/64) → primary signal is a **TTFT drop and/or

output-throughput gain** (TPOT is secondary here); the prefill M-buckets dominate e2e. Validated on ISL8192: 300.51 -> 398.35 tok/s (+32.6%) and the +79.95% ISL8192 run. Do NOT reject a real prefill win merely because TPOT did not move — that would be the wrong metric for this regime.

• balanced / unknown → require a same-session throughput gain AND no regression in EITHER TTFT or TPOT.

Universal rule: the isolated bake-off M-weighting must match the served ISL/OSL, and never stack an e2e regression. The shim already buckets tiles by M (prefill M>=512 swept / mid M>64 / decode M<=64 tile_n=64), so the decode tile_n=64 refinement is scoped to decode and never perturbs the prefill bucket.

• Sit the flydsl MoE on a decode-FULL cuda-graph config stack. The big decode gain only surfaces when

the decode path is a FULL cuda-graph. On Kimi-K2.6 that upgrade came from the aiter MLA backend (enabled by VLLM_ROCM_USE_AITER=1), which flips decode capture from PIECEWISE-only to FULL; pair with VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4. The +143.2% best run = this cfg0 stack THEN the flydsl kernel.

• Note on ceiling vs AMD's published +162%: AMD's blog (+162% / TPOT -69% / TTFT -65%, Kimi-K2.5, gfx942)

came from an INSTRUCTION-LEVEL hand-authored FLIR kernel (hand-placed mfma sched / direct-to-LDS / LDS swizzle / software pipeline). This skill only LANDS a rewrite (apply mechanics); reproducing the full ceiling needs the sibling flydsl_rewrite_quantized_moe to author a kernel that wins DECODE, not just a per-M tile reuse. gfx942 caveat: async-copy/direct-to-LDS is gfx950-default (off on gfx942), LDS budget 65536 B — so gfx942 needs arch-specific hand-tuning, not a stock-primitive swap.