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

flex/moe: CUDA graph decode for Qwen3 MoE on grouped_mm backend

Browse source 
Qwen3-30B-A3B decode on bs=8/64-new-tokens jumps from 55.6 tok/s
(eager) to 374.3 tok/s at 4bit / 421.6 tok/s at bf16 — 4-7x over pure
Hugging Face generate (66.7 / 101.1 tok/s) on B200. Output is bit-exact
against eager decode on all 32 greedy tokens across 3 prompts.

FlexMoEInference.capture_decode_cudagraph now mirrors the dense
FlexInference capture pattern: reserve a scratch page per batch_idx
slot, allocate static input_ids / batch_idx / outputs buffers, capture
one graph per bucket in [1, 2, 4, 8] + range(16, max_bs+1, 16), share
the CUDA memory pool across buckets, erase scratch pages at the end.

_decode_step replays on subsequent calls: zero graph buffers (except
outputs), copy live input_ids / batch_idx into the pinned buffers,
graph.replay(), slice outputs. Matches flex_qwen3_llama.py:725-801
verbatim.

Capture is gated on select_moe_backend() == "grouped_mm". Other
backends (unsloth_triton, native_torch) print a warning and leave
cudagraph_captured = False so generate() silently falls back to eager.
The legacy NotImplementedError rationale — torch.where + Python
for-loop over experts — only applied to native_torch; grouped_mm's
bincount + cumsum + argsort + torch._grouped_mm + index_add_ path is
fully capture-friendly on H100/B200.

FlexEngine.__init__ drops the blanket capture_cudagraph=False force
for qwen3_moe arch since capture correctness is now backend-aware.

tests/flex_moe_bench.py gains a hf_naive backend that imports pure
transformers (no unsloth patches) for a clean "fast inference vs naive
HF" comparison: flex 374.3 vs 66.7 tok/s at 4bit = 5.6x, flex 421.6
vs 101.1 tok/s at bf16 = 4.2x.

tests/flex_moe_parity.py adds greedy-decode parity across three
prompts with flex-captured / flex-eager / HF-pure. Results at bf16:
flex-captured matches flex-eager 32/32 on all three prompts (CUDA
graph is numerically identical to eager), and matches HF pure
transformers 32/32 on 2/3 prompts with the remaining divergence at a
tie-break logit boundary.

Pairs with an unsloth-zoo PR dropping @torch.compiler.disable on the
sparse MoE block and replacing torch.bincount with a capture-safe
scatter_add_ — without those, capture trips on a CPU→CUDA scalar copy
inside bincount(minlength=python_int).
This commit is contained in:
danielhanchen 2026-04-22 14:58:19 +00:00
commit 80aaf1121b
4 changed files with 323 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

59
tests/flex_moe_bench.py
View file

@ -31,7 +31,7 @@ if str(_REPO_ROOT) not in sys.path:
def main():
p = argparse.ArgumentParser()
p.add_argument("--backend", choices = ["flex", "hf"], default = "flex")
p.add_argument("--backend", choices = ["flex", "hf", "hf_naive"], default = "flex")
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")
@ -49,21 +49,52 @@ def main():
os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
import unsloth
print(f"[bench] unsloth={unsloth.__file__}")
from unsloth import FastLanguageModel
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
torch.cuda.reset_peak_memory_stats()
t_load0 = 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 = args.backend == "flex",
)
if args.backend == "hf_naive":
# Pure transformers path: NO ``import unsloth`` so none of
# Unsloth's Qwen3 MoE attention / MLP patches run. This is the
# fair naive reference to compare flex fast-inference against.
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
)
quant_cfg = None
if args.load_in_4bit:
quant_cfg = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=dtype,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
tokenizer = AutoTokenizer.from_pretrained(args.model)
if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>":
tokenizer.pad_token = "<|vision_pad|>"
tokenizer.padding_side = "left"
model = AutoModelForCausalLM.from_pretrained(
args.model,
dtype=dtype,
quantization_config=quant_cfg,
device_map="cuda",
attn_implementation="eager",
)
model.eval()
print(f"[bench] pure transformers (no unsloth patches)")
else:
import unsloth
print(f"[bench] unsloth={unsloth.__file__}")
from unsloth import FastLanguageModel
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 = args.backend == "flex",
)
t_load = time.perf_counter() - t_load0
peak_load = torch.cuda.max_memory_reserved() / 1024**3
print(f"[bench] loaded in {t_load:.1f}s peak {peak_load:.1f} GB")
@ -91,7 +122,7 @@ def main():
sum(len(o.outputs[0].token_ids) for o in outs)
)
else:
# HF generate
# HF generate (shared for "hf" unsloth-patched and "hf_naive" pure).
inputs = tokenizer(prompts, return_tensors = "pt", padding = True).to("cuda")
gen_kwargs = dict(
max_new_tokens = args.max_new_tokens,

173
tests/flex_moe_parity.py Normal file
View file

@ -0,0 +1,173 @@
# SPDX-License-Identifier: GNU Affero General Public License v3.0
# Copyright 2023-present the Unsloth team. All rights reserved.
"""Token-level parity: FlexMoEInference (CUDA-graph capture) vs HF generate.
Same prompt set, temperature=0, max_new_tokens fixed. Reports per-prompt
token-id match rate + first divergence index. Serves as the correctness
check for the v2 grouped_mm + CUDA-graph-capture changes.
Usage:
CUDA_VISIBLE_DEVICES=5 UNSLOTH_FAST_INFERENCE=1 \
UNSLOTH_MOE_BACKEND=grouped_mm python -u \
tests/flex_moe_parity.py --load_in_4bit
"""
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 _run_flex(prompts, args, dtype, *, capture: bool):
import torch
os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
import unsloth # noqa: F401
from unsloth import FastLanguageModel
if not capture:
# Monkey-patch ``capture_decode_cudagraph`` to a no-op BEFORE the
# engine is built so ``generate`` takes the eager branch. The
# engine's ``self.graphs`` stays empty, ``cudagraph_captured``
# stays False, and every step goes through ``_decode_step_eager``.
from unsloth.inference.flex_moe import FlexMoEInference
FlexMoEInference.capture_decode_cudagraph = lambda self: None
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,
)
class _SP:
max_tokens = args.max_new_tokens
temperature = 0.0
# First call warms / captures; second call is the measurement.
_ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
token_ids = [list(o.outputs[0].token_ids) for o in outputs]
texts = [o.outputs[0].text for o in outputs]
return token_ids, texts, tokenizer
def _run_hf(prompts, args, dtype):
import torch
# Pure Hugging Face: NO ``import unsloth`` — we want the unpatched
# reference forward to compare flex against. Quantization via
# transformers' ``BitsAndBytesConfig`` matches what unsloth loads
# under the hood for ``load_in_4bit=True``.
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,
)
# Translate the unsloth-flavoured model id (unsloth/Qwen3-30B-A3B-Instruct-2507)
# to the 4bit variant if load_in_4bit was requested (FastLanguageModel
# does this implicitly; do it explicitly here for the naive path).
model_id = args.model
quant_cfg = None
if args.load_in_4bit:
quant_cfg = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=dtype,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
dtype=dtype,
quantization_config=quant_cfg,
device_map="cuda",
attn_implementation="eager",
)
model.eval()
# Qwen3-30B-A3B-Instruct-2507 uses <|vision_pad|> as its pad token.
# Unsloth's loader may swap it to a sentinel; reset to the HF default
# so batched left-padded generation matches the authoritative config.
if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>":
tokenizer.pad_token = "<|vision_pad|>"
tokenizer.padding_side = "left"
gen_kwargs = dict(
max_new_tokens=args.max_new_tokens,
do_sample=False,
temperature=1.0,
pad_token_id=tokenizer.pad_token_id,
)
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda")
out = model.generate(**inputs, **gen_kwargs)
prompt_len = inputs["input_ids"].shape[1]
eos = tokenizer.eos_token_id
pad = tokenizer.pad_token_id
token_ids = []
texts = []
for row in out:
ids = row[prompt_len:].tolist()
while ids and ids[-1] in (eos, pad):
ids.pop()
token_ids.append(ids)
texts.append(tokenizer.decode(ids, skip_special_tokens=True))
return token_ids, texts, tokenizer
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("--max_new_tokens", type=int, default=32)
p.add_argument("--max_seq_length", type=int, default=1024)
p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True)
p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2")
args = p.parse_args()
import torch
prompts = [
"The quick brown fox jumps over",
"Q: What is 23 + 19?\nA:",
"Paris is the capital of",
]
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
if args.backend == "flex":
token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=True)
elif args.backend == "flex_eager":
token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=False)
else:
token_ids, texts, tok = _run_hf(prompts, args, dtype)
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"parity_{args.backend}_{precision}.json"
with open(out_path, "w") as f:
json.dump(
{
"backend": args.backend,
"precision": precision,
"prompts": prompts,
"token_ids": token_ids,
"texts": texts,
},
f,
indent=2,
)
print(f"[parity-{args.backend}] wrote {out_path}")
for i, (p_, t_) in enumerate(zip(prompts, texts)):
print(f"[parity-{args.backend}] prompt {i}: {p_!r}")
print(f"[parity-{args.backend}] completion {i}: {t_!r}")
if __name__ == "__main__":
main()

11
unsloth/inference/flex_engine.py
View file

@ -397,12 +397,11 @@ class FlexEngine:
Impl = FlexGemma4Inference
elif arch == "qwen3_moe":
Impl = FlexMoEInference
# MoE decode uses bincount + Python expert loops inside
# ``forward_moe_backend`` (unsloth_zoo moe_utils), which is
# not CUDA-graph capturable. Force eager decode so a stray
# ``capture_cudagraph=True`` does not fail inside a captured
# graph on the first token.
self.capture_cudagraph = False
# CUDA graph capture is supported on the ``grouped_mm`` MoE
# backend only. ``FlexMoEInference.capture_decode_cudagraph``
# re-checks the active backend at capture time and skips
# capture on any other backend, leaving ``capture_cudagraph``
# alone here.
else:
Impl = FlexInference
# Pass the cuMem allocator through so the impl can wrap ONLY

115
unsloth/inference/flex_moe.py
View file

@ -380,26 +380,100 @@ class FlexMoEInference:
def _decode_step(
self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor
):
# MoE path is always eager — no CUDA graph replay. See capture
# docstring below.
self.input_pos_buffer.zero_()
self.input_pos_buffer[batch_idx] = input_pos
return self._decode_step_eager(batch_idx, input_ids)
if not self.cudagraph_captured:
return self._decode_step_eager(batch_idx, input_ids)
bs = input_ids.size(0)
key = next(x for x in self.graph_bs if x >= bs)
graph = self.graphs[key]
gv = self.graph_vars
# batch_idx=0 is the reserved no-op slot. Zero out the unused part
# of each capture-shape buffer so padded entries don't write into
# real KV pages.
for k, v in gv.items():
if k != "outputs":
v.zero_()
gv["input_ids"][:bs] = input_ids
gv["batch_idx"][:bs] = batch_idx
graph.replay()
return gv["outputs"][:bs]
def capture_decode_cudagraph(self):
"""Not supported for MoE. ``Qwen3MoeExperts.forward`` uses
``torch.where`` + a data-dependent Python for-loop over experts
(shapes depend on routing), which cannot be captured. Raising
here so a stray ``capture_cudagraph=True`` fails loudly.
"""Capture one CUDA graph per batch-size bucket for MoE decode.
Future: a padded-fixed-shape dispatch can be gated behind
``UNSLOTH_MOE_STATIC_DISPATCH=1`` to make capture viable — out
of scope for the first cut.
Supported on the ``grouped_mm`` MoE backend only. On that backend
the decode path is fixed-shape:
``bincount(minlength=num_experts) → cumsum → argsort →
torch._grouped_mm × 2 → index_add_``. Python control flow in
``sparse_moe_block_forward`` runs once at capture time; only the
recorded CUDA kernels replay.
For any other backend (``unsloth_triton``, ``native_torch``) this
method logs a warning and returns without enabling replay, so
``generate(capture_cudagraph=True)`` silently falls back to eager
decode instead of failing inside the captured graph.
Pre-reserves a page for every ``batch_idx`` slot so the paged-KV
``index_put_`` during capture hits valid physical addresses. The
reservations are erased after capture — replay reads / writes the
same physical pages regardless of the logical batch state,
because ``batch_idx = 0`` is reserved as a padding slot.
"""
raise NotImplementedError(
"FlexMoEInference does not support CUDA graph capture: MoE "
"expert routing has data-dependent shapes. Run with "
"capture_cudagraph=False."
try:
from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend
backend = select_moe_backend()
except Exception:
backend = None
if backend != "grouped_mm":
print(
f"[flex] MoE CUDA graph capture requires the 'grouped_mm' "
f"backend (got {backend!r}); skipping capture, decode stays "
f"eager."
)
return
max_bs = self.max_batch_size
reserved_batches = []
for bi in range(1, max_bs):
try:
allocated = self.page_table.allocate()
self.page_table.reserve(
allocated,
torch.tensor([allocated], device = self.device, dtype = torch.long),
self.page_size,
)
reserved_batches.append(allocated)
except Exception:
break
input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device)
batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device)
outputs = torch.zeros(
(max_bs, self.model.config.vocab_size),
dtype = self.model.dtype,
device = self.device,
)
self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
pool = None
for bs in reversed(self.graph_bs):
if bs > max_bs:
continue
print(f"[flex-moe] capturing CUDA graph for bs={bs}")
torch.cuda.synchronize()
_ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph, pool):
outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
if pool is None:
pool = graph.pool()
self.graphs[bs] = graph
torch.cuda.synchronize()
for bi in reserved_batches:
self.page_table.erase(bi)
self.graph_vars = dict(
input_ids = input_ids, batch_idx = batch_idx, outputs = outputs
)
def refresh_inference_from_base(self):
@ -430,13 +504,22 @@ class FlexMoEInference:
@torch.inference_mode()
def generate(self, sequences: list[Sequence], capture_cudagraph = False):
"""Decode loop. ``capture_cudagraph`` is ignored for MoE —
always runs eager."""
"""Decode loop. Captures one CUDA graph per bucket on first call
when ``capture_cudagraph=True`` and the ``grouped_mm`` MoE
backend is active; otherwise falls back to eager decode."""
self.tokenize(sequences)
waiting = deque(sequences)
running = deque()
done = []
if capture_cudagraph and not self.cudagraph_captured:
self.capture_decode_cudagraph()
# ``capture_decode_cudagraph`` leaves ``cudagraph_captured``
# alone when it skips (non-grouped_mm backend), so only flip
# the flag when at least one bucket was actually captured.
if self.graphs:
self.cudagraph_captured = True
while waiting or running:
batch = []
while waiting and self.page_table.can_reserve(waiting[0].total_length):