ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

flex/moe: compile_walker option — 2x decode throughput on Qwen3 MoE

Browse source 
Optional ``FlexMoEInference(compile_walker=True)`` / env var
``UNSLOTH_FLEX_COMPILE_WALKER=1`` wraps the decode walker
(``call_moe_model_with_flex_kwargs``) with
``torch.compile(fullgraph=False, dynamic=False)`` before the CUDA
graph capture kicks in. Inductor fuses the layernorm + residual +
router pointwise ops, and the compiled kernels end up recorded
inside the captured graph. Net: ~2x decode tok/s on the grouped_mm
path with no change in VRAM, no correctness regression, and no
user-facing API change unless the flag is set.

Numbers on Qwen3-30B-A3B-Instruct-2507, B200, 128 new tokens,
bs sweep, median of 2 timed rounds after 1 warmup:

| precision | bs | baseline (v2) | + compile_walker | speedup |
|-----------|---:|--------------:|-----------------:|--------:|
| 4bit      | 16 |           699 |           1347.8 |   1.93x |
| 4bit      | 32 |          1243 |           2423.9 |   1.95x |
| 4bit      | 48 |          1735 |         3383.2   |   1.95x |
| 4bit      | 64 |          1523 |           2981.9 |   1.96x |
| bf16      | 16 |             — |           1401.1 |       — |
| bf16      | 32 |             — |           2864.0 |       — |
| bf16      | 48 |             — |       **3911.1** |       — |

Peak throughput: 3911 tok/s at bf16 bs=48 — 39x the pure-HF naive
baseline on the same workload (101.1 tok/s with
``AutoModelForCausalLM`` + eager attn, left-padded, no unsloth).
At 4bit bs=48, 51x the pure-HF naive baseline (66.7 tok/s).

GRPO 5-step validation (Qwen3_MoE_GRPO.py --backend flex
--max_steps 5 on DAPO-Math-17k):

| precision | baseline (v2) | + compile_walker | speedup |
|-----------|--------------:|-----------------:|--------:|
| 4bit      |        548.3s |           434.5s |   1.26x |
| bf16      |        451.8s |       **407.4s** |   1.10x |

Peak VRAM unchanged (130-133 GB). Loss / KL stable on both, no
NaN, rewards pegged at -7.5 (base-model artifact; orthogonal).

Parity (greedy 32 tokens × 3 prompts at bf16 and 4bit via
``FLEX_MOE_COMPILE_WALKER=1 tests/flex_moe_parity.py``):
flex-captured with the compile wrap matches flex-captured without
the compile wrap on 6/6 prompts with no gibberish, and matches pure
``transformers.AutoModelForCausalLM`` 32/32 on 5 of 6 (prompt ×
precision) pairs (the one divergence is a tie-break logit boundary
on an open-ended continuation — both coherent English).

Bisection of a few torch.compile flag sets against the default at
bs=32 4bit (max_batch_size=32):

| config                                                     |  tok/s |
|------------------------------------------------------------|-------:|
| default (``torch.compile(fullgraph=False, dynamic=False)``)| 1581.5 |
| + max_autotune + coord_descent + aggressive_fusion         | 1704.9 |
| + ``freezing=True``                                        |  935.7 |

``freezing=True`` is a regression on this path; shipping with the
default config only. The other flags are +7.8% at this size but
at large bs (48+) the max_autotune variant timed out during
compile (>40 min) so the default stays the ship-target for now.

Other attention backends don't help on B200 today:
- pure HF with ``attn_implementation="sdpa"``: cuDNN Frontend error
  ("No valid execution plans built") on sm_100 + torch 2.11.
- ``flash_attention_2`` 2.8.3: works, but kernels compiled for
  sm_80/sm_90 only — slower than eager on B200 (46.7 / 67.9 tok/s
  vs eager 66.7 / 101.1 at 4bit / bf16).
- ``flash_attention_3``: ``no kernel image for execution on the
  device`` — sm_100 kernels not yet in flash_attn_interface.
- FA4 / ``flash_attention_4``: works standalone but transformers'
  integration hard-codes ``flash_attn_with_kvcache = None`` for it,
  so it can't service decode. Prefill-only, out of scope here.

New tests:
- ``tests/flex_moe_micro_bench.py``: tight probe that loads the
  model once, sweeps batch sizes, prints a sample completion per
  bucket (catches gibberish early). Supports ``--compile_mode
  {off, walker, walker_fullgraph}`` and ``--compile_opts
  {stock, unsloth_O3, inference_freeze}``.
- ``tests/flex_moe_bench.py``: add ``--backend hf_naive`` which
  imports pure ``transformers`` (no ``import unsloth``) for the
  reference HF baseline, with ``HF_ATTN_IMPL`` env var to switch
  between eager / sdpa / flash_attention_{2,3,4}.
- ``tests/flex_moe_parity.py``: add ``FLEX_MOE_COMPILE_WALKER=1``
  env var to exercise the compile wrap through the parity harness.
This commit is contained in:
danielhanchen 2026-04-22 17:27:41 +00:00
commit cc4832b605
4 changed files with 229 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tests/flex_moe_bench.py
View file

@ -74,12 +74,14 @@ def main():
if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>":
tokenizer.pad_token = "<|vision_pad|>"
tokenizer.padding_side = "left"
attn_impl = os.environ.get("HF_ATTN_IMPL", "eager")
print(f"[bench] HF attn_implementation={attn_impl}")
model = AutoModelForCausalLM.from_pretrained(
args.model,
dtype=dtype,
quantization_config=quant_cfg,
device_map="cuda",
attn_implementation="eager",
attn_implementation=attn_impl,
)
model.eval()
print(f"[bench] pure transformers (no unsloth patches)")

184
tests/flex_moe_micro_bench.py Normal file
View file

@ -0,0 +1,184 @@
# SPDX-License-Identifier: GNU Affero General Public License v3.0
# Copyright 2023-present the Unsloth team. All rights reserved.
"""Tight, single-load decode throughput probe for Qwen3 MoE.
Loads the model once, captures CUDA graphs, then sweeps batch sizes.
Avoids the 30s cold-load tax of the full bench so optimization
iterations can run in under a minute per config.
Usage:
CUDA_VISIBLE_DEVICES=5 UNSLOTH_FAST_INFERENCE=1 \\
UNSLOTH_MOE_BACKEND=grouped_mm python -u \\
tests/flex_moe_micro_bench.py --load_in_4bit --bs 1,4,8,16,32
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
sys.path.insert(0, str(_REPO_ROOT))
def main():
p = argparse.ArgumentParser()
p.add_argument("--model", default="unsloth/Qwen3-30B-A3B-Instruct-2507")
p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
p.add_argument("--load_in_4bit", action="store_true")
p.add_argument("--bs", default="1,4,8,16,32",
help="comma-separated batch sizes to sweep")
p.add_argument("--max_new_tokens", type=int, default=128)
p.add_argument("--max_seq_length", type=int, default=1024)
p.add_argument("--max_batch_size", type=int, default=32)
p.add_argument("--warmup_rounds", type=int, default=1)
p.add_argument("--timed_rounds", type=int, default=2)
p.add_argument("--tag", default="baseline",
help="label for this config in the output JSON")
p.add_argument("--compile_mode", choices=["off", "walker", "walker_fullgraph"],
default="off",
help="wrap call_moe_model_with_flex_kwargs in torch.compile")
p.add_argument("--compile_opts", choices=["stock", "unsloth_O3", "inference_freeze"],
default="stock",
help="which inductor / dynamo options profile to apply before compile")
p.add_argument("--explain", action="store_true",
help="run torch._dynamo.explain on the walker first to list breaks")
p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2")
args = p.parse_args()
bs_list = [int(x) for x in args.bs.split(",") if x.strip()]
os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
import torch
import unsloth # noqa: F401
from unsloth import FastLanguageModel
from unsloth.inference import flex_moe as _flex_moe_mod
# Apply inductor / dynamo config BEFORE wrapping with torch.compile.
if args.compile_opts == "unsloth_O3":
# Aggressive autotune + coord descent + aggressive_fusion. Matches
# unsloth_zoo.patching_utils.patch_torch_compile(O3=True).
import torch._inductor.config as _ic
import torch._dynamo.config as _dc
_ic.max_autotune = True
_ic.max_autotune_pointwise = True
_ic.coordinate_descent_tuning = True
_ic.aggressive_fusion = True
_ic.cuda.use_fast_math = True
_dc.cache_size_limit = 1024
_dc.recompile_limit = 1024
_dc.capture_scalar_outputs = True
_dc.capture_dynamic_output_shape_ops = True
print("[micro] inductor/dynamo options: unsloth_O3")
elif args.compile_opts == "inference_freeze":
# Inference-friendly: constant-fold weights via freezing=True.
# Only safe when the model weights won't be updated after compile
# (true here — we capture graphs post-load and never refresh
# during bench).
import torch._inductor.config as _ic
import torch._dynamo.config as _dc
_ic.freezing = True
_ic.max_autotune = True
_ic.coordinate_descent_tuning = True
_ic.cuda.use_fast_math = True
_dc.cache_size_limit = 1024
_dc.capture_scalar_outputs = True
print("[micro] inductor/dynamo options: inference_freeze")
# Apply torch.compile to the decode walker BEFORE the engine is built
# / graphs are captured, so the compiled kernels get recorded into the
# CUDA graph.
if args.compile_mode != "off":
fullgraph = args.compile_mode == "walker_fullgraph"
orig_walker = _flex_moe_mod.call_moe_model_with_flex_kwargs
compile_kwargs = dict(fullgraph=fullgraph, dynamic=False)
tmode = os.environ.get("FLEX_COMPILE_MODE", "")
if tmode:
compile_kwargs["mode"] = tmode
compiled = torch.compile(orig_walker, **compile_kwargs)
_flex_moe_mod.call_moe_model_with_flex_kwargs = compiled
print(f"[micro] wrapped call_moe_model_with_flex_kwargs with "
f"torch.compile(fullgraph={fullgraph}, mode={tmode or 'default'})")
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
torch.cuda.reset_peak_memory_stats()
t0 = time.perf_counter()
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=args.model,
max_seq_length=args.max_seq_length,
dtype=dtype,
load_in_4bit=args.load_in_4bit,
fast_inference=True,
max_batch_size=args.max_batch_size,
)
print(f"[micro] loaded in {time.perf_counter() - t0:.1f}s")
class _SP:
max_tokens = args.max_new_tokens
temperature = 0.0
results = []
for bs in bs_list:
prompts = [f"The quick brown fox jumps over fence {i}, then"
for i in range(bs)]
# Warmup (first call captures the graphs for all buckets).
for _ in range(args.warmup_rounds):
_ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
# Timed.
wall = []
n_tok = []
for _ in range(args.timed_rounds):
torch.cuda.synchronize()
t0 = time.perf_counter()
outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
torch.cuda.synchronize()
wall.append(time.perf_counter() - t0)
n_tok.append(sum(len(o.outputs[0].token_ids) for o in outs))
med_wall = sorted(wall)[len(wall) // 2]
med_tok = n_tok[len(wall) // 2]
tps = med_tok / med_wall if med_wall > 0 else 0.0
# Sanity: print the first completion so we can eyeball for
# gibberish. A compile bug or bad capture shows up here first.
sample_text = outs[0].outputs[0].text if outs else ""
sample_preview = sample_text.replace("\n", "\\n")[:120]
print(f"[micro] bs={bs:>3} tok={med_tok:>5} "
f"wall={med_wall:.3f}s tok/s={tps:.1f}")
print(f"[micro] bs={bs:>3} completion[0]: {sample_preview!r}")
results.append({
"bs": bs,
"max_new_tokens": args.max_new_tokens,
"median_wall_s": round(med_wall, 3),
"median_tok": med_tok,
"tok_per_s": round(tps, 1),
"wall_times_s": wall,
"sample_completion": sample_text[:400],
})
peak = torch.cuda.max_memory_reserved() / 1024**3
precision = "4bit" if args.load_in_4bit else args.dtype
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / f"micro_bench_{args.tag}_{precision}.json"
with open(out_path, "w") as f:
json.dump({
"tag": args.tag,
"precision": precision,
"peak_vram_gb": round(peak, 2),
"results": results,
}, f, indent=2)
print(f"[micro] peak VRAM: {peak:.1f} GB")
print(f"[micro] wrote {out_path}")
if __name__ == "__main__":
main()

11
tests/flex_moe_parity.py
View file

@ -41,6 +41,17 @@ def _run_flex(prompts, args, dtype, *, capture: bool):
from unsloth.inference.flex_moe import FlexMoEInference
FlexMoEInference.capture_decode_cudagraph = lambda self: None
# Opt-in torch.compile wrap of the decode walker for parity check
# — enabled via env var to avoid cluttering the CLI further.
if os.environ.get("FLEX_MOE_COMPILE_WALKER") == "1":
import torch as _torch
from unsloth.inference import flex_moe as _flex_moe_mod
_orig = _flex_moe_mod.call_moe_model_with_flex_kwargs
_flex_moe_mod.call_moe_model_with_flex_kwargs = _torch.compile(
_orig, fullgraph=False, dynamic=False
)
print("[parity] torch.compile(call_moe_model_with_flex_kwargs) enabled")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=args.model,
max_seq_length=args.max_seq_length,

36
unsloth/inference/flex_moe.py
View file

@ -25,6 +25,7 @@ LoRA double-copy refresh — is shared verbatim with the dense path.
from __future__ import annotations
import os
import types
from collections import deque
from typing import Optional
@ -140,6 +141,7 @@ class FlexMoEInference:
base_model = None,
peft_model = None,
cumem_allocator = None,
compile_walker = None,
):
# FastQwen3MoeModel.pre_patch (unsloth/models/qwen3_moe.py) installs
# a legacy Qwen3MoeSparseMoeBlock_fast_forward that expects
@ -255,13 +257,37 @@ class FlexMoEInference:
L = max_seq_length,
)
# MoE decode is eager-only for this first cut. Set the captured
# flag to False permanently so ``generate(capture_cudagraph=True)``
# still runs the eager fallback.
self.cudagraph_captured = False
self.graphs = {}
self.graph_vars = {}
# Optional: wrap ``call_moe_model_with_flex_kwargs`` with
# ``torch.compile(fullgraph=False, dynamic=False)`` BEFORE CUDA
# graph capture. On Qwen3-30B-A3B 4bit this gives ~2x decode
# throughput (378 → 753 tok/s at bs=8, 1735 → 3383 tok/s at
# bs=48) because Inductor fuses the layernorm + residual +
# router pointwise ops and the compiled kernels get recorded
# into the captured graph. Opt-in for now: either pass
# ``compile_walker=True`` explicitly or set
# ``UNSLOTH_FLEX_COMPILE_WALKER=1``.
if compile_walker is None:
compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1"
self._moe_walker = call_moe_model_with_flex_kwargs
if compile_walker:
try:
self._moe_walker = torch.compile(
call_moe_model_with_flex_kwargs,
fullgraph = False,
dynamic = False,
)
print(
"[flex-moe] wrapped call_moe_model_with_flex_kwargs "
"with torch.compile(fullgraph=False, dynamic=False)"
)
except Exception as e:
print(f"[flex-moe] torch.compile wrap failed, falling back: {e}")
self._moe_walker = call_moe_model_with_flex_kwargs
# --- tokenize / prefill / decode ---------------------------------------
# Near-verbatim from FlexInference. Only difference is the walker.
@ -318,7 +344,7 @@ class FlexMoEInference:
flex_kernel_options = self.prefill_kernel_options,
)
position_ids = input_pos
hidden = call_moe_model_with_flex_kwargs(
hidden = self._moe_walker(
self.model, input_ids, position_ids, flex_kwargs
)
return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0)
@ -372,7 +398,7 @@ class FlexMoEInference:
flex_batch_idx = batch_idx,
flex_kernel_options = self.decode_kernel_options,
)
hidden = call_moe_model_with_flex_kwargs(
hidden = self._moe_walker(
self.model, input_ids.view(B, 1), position_ids, flex_kwargs
)
return self.model.lm_head(hidden[:, -1, :])