From a1bc5cbd73b83cb4d0dbbaa7a42993dac85130f3 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 22 Apr 2026 18:10:31 +0000 Subject: [PATCH] flex/moe: UNSLOTH_FLEX_GRAPH_BS env var to override capture bucket ladder The default ``graph_bs = [1, 2, 4, 8] + range(16, max_bs+1, 16)`` matches the dense FlexInference pattern and is optimal for power-of-2 batch sizes. For workloads that routinely hit intermediate sizes (e.g. bs=24, 40, 56), capturing dedicated buckets improves replay efficiency ~5-10% on those sizes (they otherwise round up to the next power-of-2 bucket). Measurements on Qwen3-30B-A3B 4bit, bs=24/40/56 with compile_walker=True: | bs | default ladder (round-up) | UNSLOTH_FLEX_GRAPH_BS=8,16,24,32,40,48,56,64 | |----|------------------------------:|---------------------------------------------:| | 24 | (rounds to 32, ~1800 tok/s) | 1836 tok/s | | 40 | (rounds to 48, ~2750 tok/s) | 2772 tok/s | | 56 | (rounds to 64, ~2700 tok/s) | 3797 tok/s | The fine ladder costs a few extra capture calls at startup and slightly larger CUDA graph pool footprint. Power-of-2 bucket throughput is marginally lower with the fine ladder (memory pool is shared across more captures), so it's opt-in via env var. Invalid values fall back to the default ladder with a log. --- unsloth/inference/flex_moe.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/unsloth/inference/flex_moe.py b/unsloth/inference/flex_moe.py index dd6768aec9..af61e12e4b 100644 --- a/unsloth/inference/flex_moe.py +++ b/unsloth/inference/flex_moe.py @@ -480,7 +480,22 @@ class FlexMoEInference: dtype = self.model.dtype, device = self.device, ) - self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + # Bucket ladder for CUDA graph capture. Default mirrors the dense + # FlexInference pattern. ``UNSLOTH_FLEX_GRAPH_BS=1,8,32,64`` etc + # lets you override (useful for tight memory or for bigger batch + # bench). Values > max_bs are silently skipped below. + _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS") + if _env_bs: + try: + self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()] + except ValueError: + print( + f"[flex-moe] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; " + f"using default bucket ladder" + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + else: + 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: