fused_moe: fp8 block-scale experts, optional SwiGLU clamp by vasilevklart · Pull Request #270 · mstar-project/mstar

vasilevklart · 2026-09-17T23:06:47Z

The fp8 block-scale fused-MoE path, shared by both GLM models.

fused_experts_fp8 runs both grouped GEMMs fp8 x fp8: per-token per-group activation quant, fp32
accumulate, block scales in the checkpoint's convention, decode-shaped tiles for the two launches.
per_token_group_quant_fp8 is compiler-disabled because Inductor's recompile of it fails inside
Triton; stream capture still records the plain launch, so a compiled forward and CUDA graphs
coexist with it.

The M grid and the kernel's EM argument are min(padded slots, topk_ids.numel() * BLOCK_M)
rather than the allocator's worst case, since each (token, expert) slot opens at most one partial
tile. At a k=3 decode shape that is 512 rows instead of 3,872. Outputs are bit-identical: the valid
slots and the kernel are unchanged.

act_and_mul_kernel gains a constexpr-gated pre-activation clamp for GLM-5.3-Flash. With
swiglu_limit=None the branch compiles out, so existing callers are untouched.

Tested: CPU suites on this tree; the GPU numerical tests are present and skip without a device.

NSagan271

Some very initial comments; will fully review in the morning

NSagan271 · 2026-09-18T07:17:25Z

@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+"""ONE GPU, minutes: sweep Triton tile configs for M*'s fused fp8 MoE kernel


Nit: this comment is very long and hard to read

NSagan271 · 2026-09-18T07:22:56Z

+
+
+@triton.jit
+def per_token_group_quant_fp8_kernel(


This kernel (and the function below calling it) doesn't seem specific to the fused MoE; it should probably be in its own util file

per_token_group_quant_fp8 moves to mstar/utils/quant_fp8.py with its tests; the config sweep's module docstring drops the narrative.

NSagan271

Mainly looks good, some minor comments

NSagan271 · 2026-09-22T05:46:36Z

+    topk_weights, topk_ids = _route(x, seed=0)
+
+    ref = _reference(x, w1, w2, w1_s, w2_s, topk_weights, topk_ids)
+    got = fused_experts_fp8(x, w1, w2, w1_s, w2_s, topk_weights, topk_ids, block_size=BLOCK)


When I try to run this on a GPU machine, I get

NotImplementedError: Could not run '_mstar_moe_C::moe_align_block_size' with arguments from the 'CPU' backend

due to passing CPU tensors into the kernel. This is because the custom op is only gated in _cuda_op_available, not whether the inputs are on-device. moe_align_block_size should probably also check that its inputs (or one representative input) is on-device before dispatching torch.ops._mstar_moe_C.moe_align_block_size (and do the torch fallback if they are on CPU)

Thanks Naomi!

NSagan271 · 2026-09-22T06:11:51Z

+at a fixed config; this script fixes the grid and sweeps the compile-time
+tiles of ``fused_moe_kernel_fp8_w8a8`` -- ``BLOCK_SIZE_N``, ``GROUP_SIZE_M``,
+``num_warps``, ``num_stages`` -- for the gate/up and down GEMMs separately.
+``fused_experts_fp8`` tunes only gate/up today and reuses that config for


Nit: this comment is stale; mstar/utils/fused_moe/runner.py already tunes config_down separately

NSagan271 · 2026-09-22T06:12:34Z

+grid sized from the worst-case padded slot count (today) vs clamped.
+
+Today ``invoke_fused_moe_kernel_fp8_w8a8`` sizes its grid from
+``sorted_token_ids.shape[0]`` = ``tokens*top_k + E*(BLOCK_M-1)`` = 32 + 256*15 =


Nit: the grid in mstar/utils/fused_moe/runner.py sets a constant grid

NSagan271 · 2026-09-22T06:21:32Z

+    # the gate/up tiles; at decode shapes each launch gets its own
+    # sweep-picked tiles instead (benchmark/glm52/bench_fused_moe_config.py).
+    config_down = dict(config)
+    if num_tokens <= 16:


Question: would this config (and the existing MoE default config) have to be tuned for different devices (i.e., I would assume the optimal block size may be different on a B200 than H100)? I guess the existing get_default_config also has the same issue, and we don't have the resources to tune for many different GPUs, but just bringing it up.

Swept on the H200s gpus (coriander), since sm90 so H100 also works.. B200 would need another sweep, the override only happens at num_tokens <=16, so on another device would be slower but still correct

NSagan271 · 2026-09-22T06:25:57Z

+    raw_block_k = default_cfg["BLOCK_SIZE_K"]
+    default_cfg["BLOCK_SIZE_K"] = QBLOCK  # fused_experts_fp8's forced override
+    BLOCK_M = default_cfg["BLOCK_SIZE_M"]
+    assert BLOCK_M == 16, (


[non-blocking] get_default_config returns BLOCK_SIZE_M=64 when M > E (doesn't happen for any models we have so far), so it might also be good to tune for block size 64.

Fixed, the bench script now takes --block-m and runs a separate alignment pass per value (the is default [16, 64])

merceod

Ran this on an H100 against vLLM 0.30's Triton fused_experts and an unquantized fp32 reference, plus a stress test sweep I did (M from 1 to 4096, E=288 and 256, K=4096 and 6144, N=256 and 2048, uniform/skewed/one-expert/ round-robin routing, clamp on and off). The sweep I ran was on commit 9eadc88 though but the two commits since only change align.py and the benchmark script, so everything below still holds on your latest commit 65d12a0 I think.

TLDR;The kernel is correct and faster than vLLM's at decode sizes, but slower at prefill. A few things to fix:

1. The moe_align JIT lock race fix is in PR #263, not here. PR #263's align.py (lines 91 to 104 over there) retries the extension load once, because eight TP workers building the same .so can delete each other's FileBaton lock, and the loser drops to the bincount fallback. Since that fallback isn't capturable, that rank loses every decode graph and the whole TP group runs eager. Since #270 is the PR that ships the kernel, the retry belongs here (align.py lines 78-102), and the serve path should build the op once before spawning workers.

2. Prefill tiles are untuned. get_default_config (kernels.py lines 622-642) ignores N and K, and only M<=16 gets the sweep picked tiles (runner.py lines 253-255). Numbers from one rank of TP8 (E=288, K=4096, N=256), timing the whole fused_experts call.

M	mstar	vLLM Triton (default config, no tuned json)
1 to 16	164 us	184 to 204 us
32 / 64	250 / 337 us	216 / 301 us
512	534 us	440 us
2048	715 us	590 us
8192	2136 us	1602 us

So decode is 10-20% ahead and prefill 20-25% behind. vLLM picked its Triton MoE path for both GLM models in my serving runs on H100 as well. Now that the benchmark script sweeps BLOCK_SIZE_M=64 too, a per shape config for M>E, (BLOCK_SIZE_N 128, more warps and stages) is probably a day of running. On @NSagan271's per device question, the only data I have is H100, where the M<=16 tiles win and the M>=512 ones lose, which lines up with the override being decode only.

3. The activation quant launches one program per 128 element group (quant_fp8.py lines 75-85). An 8k x 6144 prefill is 393k programs. sglang and vLLM do several groups per program. This is not a correctness thing, but it shows up at prefill.

The Triton cast rounds ties toward zero where torch's .to(float8) rounds to even, on about 0.1% of elements with bf16 inputs. vLLM's CUDA quantizer does the same on the same elements, so this is fine, just don't write a torch reference and expect bitwise equality.

5. Nits. test/modular/test_quant_fp8.py and test/integration/test_quant_fp8.py share a basename, so pytest test/ fails at collection with an import file mismatch (still there on 65d12a0). The new device check has a stray space before the colon, ruff doesn't mind.

The A/B and sweep scripts are on my side if you want them.

Workers then find the .so cached instead of racing on the JIT build lock.

One program per 128-element group was 393k programs for an 8k x 6144 prefill.

vasilevklart · 2026-09-24T00:43:53Z

Retry ported from GLM-5.3-Flash (glm5_next) on the resource-pool engine #263 + the conductor now builds the op once before spawning workers
M > 16 uses the block-fp8 defaults you compared against (63a8d0c7). H200, whole fused_experts_fp8 call, E=288 / K=4096 / N=256, eager us:

M old tiles new tiles new tiles + quant fix

512 483 384 374

2048 654 537 489

8192 2067 1611 1425

M <= 64 is launch-bound in harness (~300 µs flat), the decode tiles are unchanged
3. The quant kernel now does 16 groups per program: 8192×4096 goes 162 to 32 us. Under CUDA graphs at M = 1…64 it is 1.3-2.9 -> 1.2-1.6 us per launch
4. Noted
5. Fixed

vasilevklart added 2 commits September 17, 2026 15:49

fused_moe: add block scaled fp8 experts

6711f28

fused_moe: add optional SwiGLU clamp

5577b82

vasilevklart requested review from NSagan271 and merceod September 17, 2026 23:06

This was referenced Sep 17, 2026

glm52: GLM-5.2 with MLA, fp8 MoE and MTP #272

Draft

GLM-5.2: MLA, fp8 MoE, MTP speculative decoding on the resource-pool engine #264

Closed

GLM-5.3-Flash (glm5_next) on the resource-pool engine #263

Draft

NSagan271 reviewed Sep 18, 2026

View reviewed changes

Nit: split fp8 quant module, trim bench docstring

9eadc88

per_token_group_quant_fp8 moves to mstar/utils/quant_fp8.py with its tests; the config sweep's module docstring drops the narrative.

NSagan271 requested changes Sep 22, 2026

View reviewed changes

vasilevklart added 2 commits September 22, 2026 15:40

Fused MoE benchmarking + Cuda check

3de0966

ruff

65d12a0

vasilevklart requested a review from NSagan271 September 22, 2026 23:11

merceod requested changes Sep 23, 2026

View reviewed changes

vasilevklart added 3 commits September 23, 2026 16:26

Improved CUDA handling, config update

63a8d0c

conductor: build fused-MoE op before spawn

103656e

Workers then find the .so cached instead of racing on the JIT build lock.

quant_fp8: quantize 16 groups per program

9a8980b

One program per 128-element group was 393k programs for an 8k x 6144 prefill.

vasilevklart requested a review from merceod September 24, 2026 00:44

		@@ -0,0 +1,361 @@
		#!/usr/bin/env python3
		"""ONE GPU, minutes: sweep Triton tile configs for M*'s fused fp8 MoE kernel



		@triton.jit
		def per_token_group_quant_fp8_kernel(

Conversation

vasilevklart commented Sep 17, 2026

Uh oh!

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

merceod left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasilevklart commented Sep 24, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

merceod left a comment •

edited

Loading