Skip to content

triton-kernels: fix Intel opt-flag heuristics copied from the NVIDIA path - #1075

Draft
jiqing-feng wants to merge 1 commit into
huggingface:mainfrom
jiqing-feng:triton-kernels-intel-opt-flags
Draft

triton-kernels: fix Intel opt-flag heuristics copied from the NVIDIA path#1075
jiqing-feng wants to merge 1 commit into
huggingface:mainfrom
jiqing-feng:triton-kernels-intel-opt-flags

Conversation

@jiqing-feng

Copy link
Copy Markdown
Contributor

Summary

opt_flags_intel.py was derived from opt_flags_nvidia.py but kept two
constants that encode NVIDIA hardware properties. Both are wrong for Xe and
together cost roughly an order of magnitude on Battlemage.

block_k. The NVIDIA helper derives block_k from the cacheline size; the
Intel copy hardcoded min(128, ...) instead. An Xe cacheline is half as wide,
so operand tiles were twice the size the hardware wants. Deriving it the same
way NVIDIA does needs the operand dtypes, hence the signature change.

num_warps. block_m * block_n // 4096 assumes a 32-lane warp. An Xe
sub-group is SIMD16, so a warp covers half as many elements and the same tile
needs twice as many to keep per-lane register pressure in range. A standalone
tl.dot sweep puts the optimum consistently at // 2048.

Also allows block_n as a tuning constraint, so the Intel path can be swept
from the outside like block_m and block_k already can.

Only the backend == "xpu" branch is touched; CUDA and HIP go through
make_default_opt_flags_nvidia / _amd and are unaffected.

Validation

bf16 on Intel Arc Pro B60, torch 2.13.0+xpu, one card:

shape speedup
MoE, 4096 tokens / 32 experts / top4 ~10x
MoE, 2048 tokens / 8 experts / top2 ~9x
dense 4096³ ~5x

Numerics are unchanged: max relative error against a float32 reference stays at
the same order as before, for both the dense and the MoE path.

For context, a hand-written Triton tl.dot matmul tops out around 3x the
post-fix throughput on this device and oneDNN around 9x, so a gap remains — but
it is in the Triton XPU backend, not in these heuristics.

Note

Benchmarking the MoE path against matmul_ogs_torch also needs the reference
implementation fix in the companion PR; the matmul_ogs path itself does not
depend on it.

@github-actions github-actions Bot added the chore Version bumps, releases, misc maintenance label Aug 13, 2026
@jiqing-feng
jiqing-feng force-pushed the triton-kernels-intel-opt-flags branch from d76e5e5 to c2f9584 Compare August 13, 2026 07:26
…path

`opt_flags_intel.py` was derived from `opt_flags_nvidia.py` but kept two
constants that encode NVIDIA-specific hardware properties. Both are badly
wrong for Xe and together cost ~10x on Battlemage (Arc Pro B60).

1. block_k. The NVIDIA helper derives block_k from the cacheline size
   (1024 bits), while the Intel copy hardcoded `min(128, ...)`. An Xe
   cacheline is 512 bits, so the operand tiles were twice as large as the
   hardware wants. Measured on a real MoE shape (4096 tokens, 32 experts,
   top4, 2048x2048): block_k=128 -> 1.11 TF/s, 64 -> 2.45, 32 -> 11.00.
   Derive it from the cacheline like the NVIDIA path does, which needs the
   operand dtypes, hence the signature change.

2. num_warps. `block_m * block_n // 4096` assumes a 32-lane warp. An Xe
   sub-group is SIMD16, so the same tile needs twice as many warps to keep
   per-lane register pressure in range. A standalone `tl.dot` sweep shows
   the optimum is consistently `block_m * block_n // 2048`, e.g. for a
   128x128x32 tile 4 warps gives 4.39 TF/s and 8 warps gives 25.64 TF/s.

Also allow `block_n` as a tuning constraint so the Intel path can be swept
from the outside like block_m/block_k already can.

Measured end to end (bf16, Arc Pro B60, ZE_AFFINITY_MASK=0):

  MoE 4096 tok / 32 experts / top4 : 1.09 -> 11.20 TF/s  (10.3x)
  MoE 2048 tok /  8 experts / top2 : 1.05 ->  9.50 TF/s   (9.0x)
  dense 4096^3                     : 1.98 -> 10.33 TF/s   (5.2x)

Numerics are unchanged: max relative error vs a float32 reference stays at
~2e-3 for dense and ~3e-3 vs `matmul_ogs_torch` for the MoE path.

For reference, a hand-written Triton `tl.dot` matmul tops out at ~25 TF/s
on this device while oneDNN reaches ~90 TF/s, so the remaining gap is in
the Triton XPU backend rather than in these heuristics.

Only the `backend == "xpu"` branch is touched; CUDA and HIP go through
`make_default_opt_flags_nvidia` / `_amd` and are unaffected.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

chore Version bumps, releases, misc maintenance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant