Skip to content

Add xe2 block_fp8 grouped_gemm kernel#436

Open
frost-intel wants to merge 7 commits into
vllm-project:mainfrom
frost-intel:moe_blockfp8
Open

Add xe2 block_fp8 grouped_gemm kernel#436
frost-intel wants to merge 7 commits into
vllm-project:mainfrom
frost-intel:moe_blockfp8

Conversation

@frost-intel

@frost-intel frost-intel commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

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

  • Correctness — block-FP8 output vs a dequantized (fp8 → x_dtype × scale) reference MoE:
    pytest tests/fused_moe/test_grouped_gemm.py
    Also exercised via the benchmark's calculate_diff against ref_fused_moe across model shapes, both activation dtypes (bf16/fp16), both weight dtypes (e4m3/e5m2), and with/without bias.
  • Performance — per-stage (remap / gemm1 / gemm2 / gather) latency + TFLOPS sweep over representative MoE shapes:
    python benchmark/benchmark_cutlass_fused_moe.py

Test Result

Tests were executed on PVC Max 1550.

  • Correctness: all configs pass torch.testing.assert_close within the per-dtype rtol/atol (block-FP8, per-tensor FP8, and BF16 baselines).
  • Performance: block-FP8 is ~10% slower than per-tensor FP8, with the worst gap at large M, and is within a few percent of per-tensor at small/decode M.
image

DeepSeek-V3

My objective is to improve DSv3 performance. For those kernels specifically, we achieve a mean of 1.04x


View with Codesmith Autofix with Codesmith
Need help on this PR? Tag /codesmith with what you need. Autofix is disabled.

Copilot AI review requested due to automatic review settings June 25, 2026 21:07

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_fp8block flag 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
@xinyu-intel

Copy link
Copy Markdown
Collaborator

thanks for the contribution. I tend to merge this first. #376

@jikunshang

Copy link
Copy Markdown
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>
Signed-off-by: frost-intel <frost.mitchell@intel.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants