[Feat] Fuse MoE activation FP8 quantization kernels - #2089
Open
qiyueliuhuo23 wants to merge 5 commits into
Open
qiyueliuhuo23 wants to merge 5 commits into
qiyueliuhuo23 wants to merge 5 commits into
Conversation
heqs1013
reviewed
Sep 17, 2026
| def _assert_equal(actual: torch.Tensor, expected: torch.Tensor, name: str) -> None: | ||
| assert actual.shape == expected.shape, f"{name}: shape {actual.shape} != {expected.shape}" | ||
| assert actual.dtype == expected.dtype, f"{name}: dtype {actual.dtype} != {expected.dtype}" | ||
| assert torch.equal(actual, expected), f"{name}: values differ" |
Author
There was a problem hiding this comment.
已将两个测试文件统一改为 bitwise 比较，重新单测均通过
| ) | ||
| input_block = tl.load(input_offsets, mask=input_mask, other=0.0).to(tl.float32) | ||
|
|
||
| reciprocal_fp8_max = 1.0 / FP8_MAX |
There was a problem hiding this comment.
（仅讨论，不是修改建议）量化这一段在融合kernel和独立kernel里的共同部分，是否考虑抽出作为inline调用？重复代码看起来不容易维护
Author
There was a problem hiding this comment.
将公共量化计算抽取到了 quantization.py 的三个 @triton.jit helper，并且替换了xtuner/tests/float8/triton_kernels下的所有量化操作
| ([0, 1, 127, 128, 129], 6144), | ||
| ], | ||
| ) | ||
| def test_per_tile_quant_with_trans_per_block_matches_reference( |
Author
There was a problem hiding this comment.
目前检查了 XTuner 的 tests/，没有找到单测中记录 MFU/MBU 的现成例子，我理解单测还是主要做数值对齐。是不是可以单独提供 benchmark 脚本，记录MFU、MBU等指标呢
Reuse Triton row, column, and block quantization helpers across independent and fused kernels while preserving historical reduction and scale arithmetic. Require byte-level equality in activation fusion tests. Validated on H200 with 17 passing tests and 164 tensors bitwise equal to pre-refactor outputs.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR fuses activation processing in the tile-wise FP8 MoE grouped GEMM path.
The original implementation reads the same BF16 activation multiple times for:
per_tile_quanttrans_per_block_quant_expand_128xortrans_per_tile_quant_expand_128xnative_swigluThe new Triton kernels reduce redundant BF16 memory traffic and kernel launches while preserving the
existing FP8 layouts and quantization semantics.
Changes
per_tile_quant_with_trans_per_blockper_tile_quant+trans_per_block_quant_expand_128xper_tile_quant_with_trans_per_tileper_tile_quant+trans_per_tile_quant_expand_128xswiglu_per_tile_quant_with_trans_per_blocknative_swiglu+ independent dual-layout quantizationSiLU(gate) * upand W2 input quantization without materializing BF16 activationswiglu_backwardgrad_gate_upin one Triton kernelAdditional integration changes:
TileWiseFloat8GroupedLinearandMoEBlock.Float8Config.enable_fused_moe_activation.FUSED_MOE_ACTIVATIONsupport to the GLM-5.2 example configuration.Correctness
The fused outputs were compared with the original independent kernels on H200:
The latest 35-step GLM-5.2 candidate run was checked with the standard CI checker. The values below are
the actual baseline-to-candidate comparison at the largest drift position, except TGS, which follows the
checker’s step 11--34 one-sided P80 rule.
loss/local_lossloss/reduced_llm_lossgrad_normmemory/max_memory_GBruntime_info/tgsruntime_info/text_tokensloss/reduced_balancing_lossPerformance
Single-kernel benchmark
Measured on one H200, with 100 warm-up iterations and 300 measurement iterations. Values are p50;
the time column uses
baseline / candidatein milliseconds.per_tile_quant_with_trans_per_block[131072, 6144], 256 expertsper_tile_quant_with_trans_per_tile[131072, 6144], 256 expertsswiglu_per_tile_quant_with_trans_per_blockgate_up[32768, 4096]swiglu_backwardgate_up[32768, 4096]swiglu_per_tile_quant_with_trans_per_blockgate_up[131072, 8192], 256 expertsswiglu_backwardgate_up[131072, 8192], 256 expertsEnd-to-end A/B benchmark
GLM-5.2 EP4, 16K pack, 30-step warm-up plus 100-step measurement. Baseline disables
enable_fused_moe_activation; candidate enables it. Other configuration remains unchanged.The end-to-end gain is smaller than the single-kernel speedup because the fused activation kernels account
for only about 36.8 ms of a roughly 4.55 s training step. Even if the entire target kernel time were removed,
the theoretical upper bound for this workload would be approximately 0.81%.