expXa

Research

Warp sharing in expert and attention kernels

Can threads in a warp share expert routing, weight metadata or attention scalars instead of each recomputing them, and does that make the kernels faster?

Finding

No. The existing expert kernel already shares the useful things. The new variants were either slower than a matched control or only slightly faster, by amounts too small to attribute. In attention, broadcasting rescaling values made the 4×8 schedule we actually use about 19% slower. Nothing was adopted.

Expert kernel

The two-token Q6_K down projection already broadcasts expert IDs with warp shuffles and groups matching experts. We tested four additions across every expert overlap count from 0 to 8, including padded strides and invalid-ID tails.

Median change in operator time. Negative means faster.
Changevs current kernelvs matched control
Paired Q6 fragment hoist−0.76%+0.31%
Lane-zero expert reduction−0.76%+0.31%
Block-shared expert routing+2.45%+3.73%
Warp-shared Q6 metadata+1.51%+2.69%

The matched control rewrites routing and launch structure, and on its own ran 1.06% faster than the current kernel. Measured against it, every candidate was slower. So the small wins against the current kernel point to the routing rewrite, not to warp sharing.

Disassembly showed identical global-load instruction counts for the current kernel, the control and the paired hoist. Shared metadata added 16 indexed shuffles. Instruction counts don’t measure memory traffic.

4×8 attention

Row maxima and sums feed several output fragments, so their broadcasts are required. We tested broadcasting two other values instead: the common rescaling exponentials and the final reciprocal denominators.

Median change in operator time. Negative means faster.
Change4×8 schedule (in use)8×8 schedule
Rescaling exponential broadcast+18.87%−1.62%
Reciprocal broadcast+1.42%+0.65%
Both+20.35%−0.93%

Fixtures covered 1, 2 and 4 queries, 256 to 32,768 KV rows, masked and unmasked, each checked against an independent dense reference. The probe calls the real tile implementation but leaves out serving’s Q8-to-F16 conversion and stream-K scheduling.

Q8_0 vector attention

Head size 128, native Q8_0 K/V. Median change in operator time.
ChangeFP16 accumulateFP32 accumulate
Root-only reductions, lane-zero consumers+0.39%0.00%
Common exponential broadcast+0.71%−0.32%
Both−0.24%+0.16%

Results depend on shape and don’t add up to a general improvement.

Method

  • 15,408 scored timing batches: 12 balanced repetitions per configuration, each averaging 64 expert or 16 attention launches.
  • Every final candidate matched its reference bit for bit: 1,152 expert, 216 MMA and 288 vector comparisons across both GPUs. Memory and synchronization sanitizers reported no errors.
  • Figures are medians of per-fixture paired time ratios across both GPUs, with 2,000-resample bootstrap intervals in the download. They describe these fixtures, not serving throughput.

Excluded attempts

  • Two early MMA harness bugs were found and fixed: a wrong output-head offset, and weak template symbols that made two arms run the same kernel. The final build checks for eight distinct kernel symbols.
  • Earlier expert and vector sweeps that lacked binary hashes were rerun and left out of scoring. All of these attempts are kept in the archive.