Add xe2 block_fp8 grouped_gemm kernel#436
Open
frost-intel wants to merge 7 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds block-wise FP8 (W8A16) support for the Xe2 grouped GEMM used by the fused MoE op, enabling 128×128 (K×N) block scale grids to match block-scaled FP8 checkpoints.
Changes:
- Extends the CUTLASS grouped GEMM interface end-to-end with an
is_B_fp8blockflag and block-scale validation. - Implements a new Xe2 device kernel path (
xe_gemm_blockfp8) that applies per-128×128 block FP32 scales for FP8 weights. - Adds correctness/perf benchmarking plumbing and a new unit test covering block-FP8 grouped GEMM.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| vllm_xpu_kernels/fused_moe_interface.py | Plumbs is_B_fp8block through Python fused-MoE calls into the XPU op. |
| tests/fused_moe/test_grouped_gemm.py | Adds a new block-FP8 grouped GEMM correctness test. |
| csrc/xpu/torch_bindings.cpp | Updates the Torch library schema to include is_B_fp8block. |
| csrc/xpu/ops.h | Updates the C++ op declaration with the new boolean parameter. |
| csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp | Adds compile-time IsBlockScale handling and dispatch to the block-FP8 GEMM path. |
| csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2.h | Extends Xe2 grouped GEMM API to accept is_B_fp8block. |
| csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2.cpp | Forwards is_B_fp8block into the Xe2 implementation. |
| csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2_interface.hpp | Adds block-FP8 argument validation and launch plumbing (IsBlockScale=true). |
| csrc/xpu/grouped_gemm/xe_2/gemm_xe2.hpp | Implements xe_gemm_blockfp8 device kernel for block-scaled FP8 weights. |
| csrc/xpu/grouped_gemm/grouped_gemm_interface.h | Extends the top-level grouped GEMM interface signature with is_B_fp8block. |
| csrc/xpu/grouped_gemm/grouped_gemm_interface.cpp | Wires is_B_fp8block into the Xe2 path and rejects it for XE default kernel. |
| benchmark/src/get_model_config.py | Adds block-FP8 variations to correctness/perf config generation with validity filtering. |
| benchmark/src/fused_moe_interface_.py | Adds is_block_fp8 flag and forwards it to grouped GEMM calls. |
| benchmark/benchmark_cutlass_fused_moe.py | Adds block-FP8 data generation, diff checking, and benchmark parameterization. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+325
to
+330
| TORCH_CHECK( | ||
| is_weight_fp8, "block fp8 grouped gemm must have fp8 weights"); | ||
| TORCH_CHECK( | ||
| ptr_scales.has_value(), "block fp8 grouped gemm must have scales"); | ||
| TORCH_CHECK(ptr_scales->is_contiguous(), "ptr_scales must be contiguous"); | ||
| TORCH_CHECK(ptr_scales->dtype() == at::kFloat, "ptr_scales must be float"); |
Comment on lines
+281
to
+283
| scale_full = scale_B[i].repeat_interleave( | ||
| block_size, dim=0).repeat_interleave(block_size, dim=1) | ||
| weight = input_B_fp8[i].to(torch.float32) * scale_full |
Collaborator
|
thanks for the contribution. I tend to merge this first. #376 |
Member
|
please rebase and fix conflicts. |
Signed-off-by: frost-intel <frost.mitchell@intel.com>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
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
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.
Purpose
Add block-wise FP8 (W8A16) weight quantization support to the Xe2 grouped GEMM that backs the fused-MoE op. Weights are stored in FP8 (
e4m3/e5m2) with a 128×128 block scale grid and dequantized on the fly against BF16/FP16 activations, complementing the existing per-tensor FP8 path. This matches the block-scaled FP8 checkpoints in DeepSeek-V2/V3, Qwen3-MoE, etc.Test Plan
fp8 → x_dtype × scale) reference MoE:calculate_diffagainstref_fused_moeacross model shapes, both activation dtypes (bf16/fp16), both weight dtypes (e4m3/e5m2), and with/without bias.Test Result
Tests were executed on PVC Max 1550.
torch.testing.assert_closewithin the per-dtypertol/atol(block-FP8, per-tensor FP8, and BF16 baselines).DeepSeek-V3
My objective is to improve DSv3 performance. For those kernels specifically, we achieve a mean of 1.04x
Need help on this PR? Tag
/codesmithwith what you need. Autofix is disabled.