Xpu fixes for finegrained-kernels - #1104
Open
kaixuanliu wants to merge 9 commits into
Open
Conversation
github-actions
Bot
added
chore
kaixuanliu
marked this pull request as ready for review
August 28, 2026 07:46
The suite already had device-agnostic helpers (TEST_DEVICE, accelerator_module, SUPPORTS_FP8) and conftest already admitted XPU, but 21 per-test gates still read skipif(TEST_DEVICE != "cuda", reason="CUDA required") and several tests hardcoded device="cuda" / torch.cuda.synchronize(). On XPU that skipped all 254 tests, and forcing them on then failed with "Torch not compiled with CUDA enabled". - gates: TEST_DEVICE != "cuda" -> TEST_DEVICE is None (accelerator required) - device="cuda" -> TEST_DEVICE, torch.cuda.sync -> accelerator_module().sync - conftest: _visible_gpu_pool() takes (env_var, list_cmd); worker pinning probes ZE_AFFINITY_MASK/xpu-smi before CUDA_VISIBLE_DEVICES/nvidia-smi - bench_moe: derive DEV/ACCEL/DEV_MASK_ENV instead of assuming cuda 414 tests now collect on XPU (was 254, all skipped).
The XPU arm pruned bn_span/bk_span to a single value (128,). The autotune grid is the only source of tiles, so a single-valued span makes every N (or K) that is not a multiple of 128 unschedulable -- the launch raises ValueError: N=320 is not a multiple of any BLOCK_SIZE_N in the autotune grid before a single config is ever benched. This took out every 320/1024-shaped scenario on XPU (10 failures across test_ops and test_moe). Add 64 to both spans; 256/512 stay out since large Xe-core tiles spill. CUDA is untouched. The 10 previously-failing cases now pass, with no regression in the 10 already-passing cases that share those shapes (20 passed, 0 failed).
…dgrad
WARP_SPEC is emitted into the autotune grid only under `get_active_device_type()
== "cuda"` (compat.py), but both dgrad kernels declared it as a constexpr with no
default. Off-CUDA the tuner therefore hands the kernel a config that never carries
the key and every launch dies in the launcher:
TypeError: dynamic_func() missing 1 required positional argument: 'WARP_SPEC'
165 configs failed this way on XPU. The autotuner scores a failed config as inf
and moves on, so it surfaces as "no config succeeded" rather than as this error.
Giving the axis its CUDA-off value as the default leaves the CUDA grid unchanged
(it always passes the key explicitly) and lets every other backend launch.
make_weights drew the whole (E, out, in) stack in fp32 and then materialized a
second scaled copy of it. At the production MoE shape the suite already exercises
(E=128, N=4096, K=6144) each of those is 12GiB, so the fixture needs 24GiB before
a single kernel launches:
tests/utils.py:159: torch.OutOfMemoryError: XPU out of memory.
Tried to allocate 12.00 GiB. GPU 0 has a total capacity of 31.89 GiB
of which 7.67 GiB is free. Of the allocated memory 24.00 GiB is
allocated by PyTorch
Building one expert at a time keeps the peak at 2 * N * K * 4 (~200MB here) and
concatenating at the end gives a bit-identical result: the block-scale reduction
is already per-expert, so nothing crosses the expert axis.
test_moe.py's two large-shape fp8 cases go from OOM to `2 passed in 395.35s`.
…ll/ rows
Two independent fixes to bench_moe.py.
1. _Experts fed the transformers-integration forwards our interleaved gate|up
weight while leaving them to split it with chunk(2), which pairs the wrong
halves. The result is not a crash but a silently scrambled forward, and it
showed up as the transformers arms reporting parity 1.2e+00 with cosine ~0
against every other arm -- i.e. the whole transformers baseline column was
measuring the wrong computation.
Determined empirically (7 candidate splits scored against moe_fused_batched
and moe_fused_grouped, identical results for both):
stacked gate=[0,I) up=[I,2I) rel 1.1991 cos -0.0193
stacked-swap up=[0,I) gate=[I,2I) rel 1.8155 cos -0.0008
row-interleave gate=0::2 up=1::2 rel 0.0035 cos 1.0000 <--
row-interleave-swap gate=1::2 up=0::2 rel 0.7877 cos 0.7847
block-interleave BN=32/64/128 rel 1.5/1.9/1.5, cos ~0
Overriding _apply_gate to split on the stride is the fix; de-interleaving the
weight itself is not possible for the fp8_128x128 recipe because a 128-row
block scale spans 64 gate and 64 up rows. Setting is_concatenated=False has no
effect, since _Experts overrides _apply_gate and the flag is never consulted.
Parity for the transformers arms goes from 1.2e+00 to 4.7e-03 ~ 6.4e-02.
2. Three small/ problems (E32 H2048 I1024, /8 experts and /2 dims of the DeepSeek
-V4 and MiniMax-M3 geometries). The full-size rows need ~16GB for a single
pre-quant gate_up grid, so they cannot run on a 32GB part at all.
a9e4828 made the element-level row interleave the kernel layout, but several docstrings kept describing the pre-interleave stacked ordering. They are the only statements in the tree that assert WHICH rows hold gate and which hold up, and they assert it wrongly: batched.py "``B`` is the (E, 2N, K) stack (gate rows [0,N), interleaved rows)" mma.py "with its STACKED gate|up tile (gate rows first, ...)" split_gate_up is the ground truth and reads the accumulator the other way: "de-interleave the doubled N extent ... The accumulator's N axis ALTERNATES gate/up PER COLUMN, so this is a trailing-axis split" g, u = tl.split(tl.reshape(flat, (rows, BLOCK_SIZE_N, 2))) Comment-only in the kernels. "gate|up stack" as the name of the (2N, K) slab is left alone -- it names the tensor, not an ordering. bench_moe.py also acted on the stale description: the triton_kernels arm remapped our rows to interleaved before handing them to GPT-OSS, on the stated grounds that "our layout stacks them". Both sides interleave, so the remap was the thing that paired the wrong halves. Dropped -- the slab transfers as-is. (Not exercised here: triton_kernels has no XPU build.)
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
kaixuanliu
force-pushed
the
xpu-fixes
branch
2 times, most recently
from
September 7, 2026 08:01
44d7997 to
eeeb5e0
Compare
Contributor
Author
|
@IlyasMoutawwakil pls help review, thx! |
Contributor
Author
|
After optimization, current status w/ this PR: MoESpeedup = baseline latency / finegrained-kernels latency. Measured on Intel Arc Pro B70 (XPU).
attn FP8 (dense gemm)
|
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
1 participant
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.
No description provided.