Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds several new Xe-Fuse kernel-fusion patterns (and supporting epilogue-builder aliases) targeting BMG, along with new correctness tests and micro-benchmarks to validate/measure the new fused paths.
Changes:
- Introduces new fused kernels for (1) QK RMSNorm+RoPE, (2) gate+residual+gamma epilogue, and (3) FP8 GEMM dequant (+ SwiGLU) epilogues.
- Extends standalone/vLLM-equivalent baselines and builder aliases to support the new fusion patterns.
- Adds new tests and benchmarks and wires them into
tests/CMakeLists.txt.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_qk_norm_rope.cpp | New correctness test for fused per-head RMSNorm + RoPE |
| tests/test_k0g.cpp | New correctness test for GEMM + gate + residual + gamma epilogue |
| tests/test_fp8_k2.cpp | New correctness/perf test for FP8 GEMM dequant + SwiGLU |
| tests/CMakeLists.txt | Registers new tests and benchmarks |
| tests/bench_swiglu_requant.cpp | New benchmark comparing SwiGLU+INT8 requant strategies |
| tests/bench_qk_norm_rope.cpp | New benchmark comparing QK norm+RoPE fusion strategies |
| tests/bench_norm_quantize.cpp | New benchmark comparing norm+INT8 quantization strategies |
| include/xe-fuse/standalone/vllm_ops.hpp | Adds per-head RMSNorm and interleaved RoPE baseline kernels |
| include/xe-fuse/standalone/ops.hpp | Adds quantize_bf16_to_int8 standalone quantization helper |
| include/xe-fuse/kernels/swiglu_requant.hpp | New standalone SwiGLU(+optional) + INT8 requant kernel |
| include/xe-fuse/kernels/qk_norm_rope.hpp | New fused per-head RMSNorm + RoPE kernel |
| include/xe-fuse/kernels/gemm_gate_residual_norm.hpp | New K0g epilogue fusion pattern (gate + residual + gamma) |
| include/xe-fuse/kernels/gemm_fp8_dequant.hpp | New FP8 dequant (+ SwiGLU) GEMM epilogue fusion kernels |
| include/xe-fuse/kernels/compute_rstd.hpp | Adds dual-output norm+quant kernel (BF16 + INT8) |
| include/xe-fuse/builder/epilogue_builder.hpp | Adds FP8 dequant-related epilogue-builder aliases |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+152
to
+163
| // ── Benchmark NAIVE ─────────────────────────────────────────────────────── | ||
| timer.start(); | ||
| for (int i = 0; i < iters; ++i) { | ||
| // swiglu: reads+writes [M, N_ffn] in-place | ||
| xe_fuse::standalone::swiglu(q, swiglu_work_naive.get(), M * L, N_ffn, 1); | ||
| // rstd: reads [M, d] post-swiglu | ||
| xe_fuse::standalone::compute_rstd(q, rstd_naive.get(), | ||
| block_post_swiglu.get(), M, d, L); | ||
| // quantize: reads [M, d] twice + writes [M, d] int8 | ||
| xe_fuse::standalone::quantize_activations(q, block_post_swiglu.get(), | ||
| rstd_naive.get(), quant_naive.get(), scale_naive.get(), M, d, L); | ||
| } |
Comment on lines
+158
to
+160
| // GemmFP8DequantSwiGLU — K2_FP8 | ||
| // D[m,n] = SwiGLU( acc[m,n] * scale_a[m] * scale_b[n] ) | ||
| // For FFN gate+up projections; SwiGLU contracts N→N/2 at output. |
Comment on lines
+102
to
+121
| std::vector<ElementQ> h_ref(input_size); | ||
| constexpr float eps = 1e-6f; | ||
|
|
||
| for (int tok = 0; tok < M; ++tok) { | ||
| for (int h = 0; h < num_heads; ++h) { | ||
| int row_base = tok * num_heads * head_dim + h * head_dim; | ||
| int cs_base = tok * head_dim; | ||
|
|
||
| // RMSNorm | ||
| float sum_sq = 0.f; | ||
| for (int d = 0; d < head_dim; ++d) { | ||
| float v = static_cast<float>(h_in[row_base + d]); | ||
| sum_sq += v * v; | ||
| } | ||
| float rstd = 1.f / std::sqrt(sum_sq / head_dim + eps); | ||
|
|
||
| // normalize + gamma | ||
| std::vector<float> normed(head_dim); | ||
| for (int d = 0; d < head_dim; ++d) | ||
| normed[d] = static_cast<float>(h_in[row_base + d]) * rstd * h_gamma[d]; |
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.
Three new kernel-fusion patterns targeting BMG:
QK normalization + RoPE (
kernels/qk_norm_rope.hpp)Per-head RMSNorm fused with rotary positional embedding in a single two-pass
subgroup kernel. Used in Gemma 3, Qwen2.5-VL, FLUX.2.
Gate + residual + gamma (
kernels/gemm_gate_residual_norm.hpp)EVT:
D[m,n] = gamma[n] * (gate[m] * acc[m,n] + residual[m,n]).Covers the FLUX.2 single-stream block output path.
FP8 GEMM epilogue fusion (
kernels/gemm_fp8_dequant.hpp)GemmFP8DequantandGemmFP8DequantSwiGLUusingfloat_e4m3_toperands,manually constructed mainloop (
MainloopIntelW8A8+XE_8x16x16_F32F16F16F32_TTupcast path) and per-token × per-channel dequant EVT identical in structure to
the W8A8 tree.
Epilogue builder additions (
builder/epilogue_builder.hpp)DequantFP8,DequantFP8SwiGLU,DequantFP8GeGLUtype aliases.Tests and benchmarks
tests/test_qk_norm_rope.cpp— correctness across LLaMA/Qwen/Gemma head sizestests/test_k0g.cpp— correctness for FLUX.2 and LLM dimensionstests/test_fp8_k2.cpp— correctness + throughput for FP8 SwiGLU GEMMtests/bench_{qk_norm_rope,norm_quantize,swiglu_requant}.cpp— throughput benchmarks