expXa

Research

Reordering expert work for weight reuse

If jobs that use the same expert run close together, they should reuse the same weight tiles while those tiles are still in cache. Is the saving larger than the cost of sorting?

Finding

Only for large batches. With 255 active jobs, grouping by expert and walking 32-row bands cut operator time by 3.8–4.2%, including the cost of the sort. With 63 jobs it was about 0.7% slower, and it slowed the existing paired kernel by 11–21%. That isn’t enough to adopt it. It might be worth revisiting as a policy that only applies to large batches.

Confirmation run: 32-row bands against job-major order, both GPUs. Times include the GPU sort. Negative means faster.
GPUActive jobsControl µsGrouped µsChange95% interval
06397.2897.98+0.66%−0.60 to +0.79%
0255353.92339.10−4.21%−4.24 to −4.17%
16397.4798.08+0.66%+0.59 to +0.76%
1255356.58343.10−3.76%−3.81 to −3.66%

The large fixture uses 32 expert IDs, 2,048 output rows and 512 input columns. These are standalone operator timings, not model throughput, and they don’t show an improvement over the engine’s existing batched MoE path.

What changed

The research kernel keeps expXa’s Q6_K × Q8_1 dot product and output indexing. A parallel GPU rank sort groups job indices by expert ID. Work then sweeps a small band of rows across those groups before moving on, so consumers of the same weight tiles run closer together in time. This changes how work is assigned, not how the hardware schedules warps.

Six schedules were compared: job-major control, full-row expert sort, and sorted bands of 4, 32, 128 and 512 rows. A seventh arm calls the existing paired kernel directly. Sorting and banding were always combined, so their individual effects can’t be separated.

Full sweep for the 32-row schedule
Change in operator time including the sort. Pool bound is the range expert IDs are drawn from.
GPUKernelJob slotsPool boundChange
0Paired168+11.04%
0Paired1632+20.49%
0Independent168+9.40%
0Independent1632+20.29%
0Independent16256+13.93%
0Independent648+2.01%
0Independent6432−1.82%
0Independent64256+4.60%
0Independent2568+0.55%
0Independent25632−4.87%
0Independent256256+5.38%
1Paired168+11.24%
1Paired1632+21.01%
1Independent168+9.13%
1Independent1632+13.56%
1Independent16256+12.89%
1Independent648+2.20%
1Independent6432−1.58%
1Independent64256+4.09%
1Independent2568+1.04%
1Independent25632−4.14%
1Independent256256+5.24%

High overlap between experts doesn’t guarantee a win. An earlier prototype with a serial sort lost badly once the sort was counted; its results are in the download too.

Method

  • Main sweep: 3,264 timing batches, each averaging 16 launches after three warmups, 12 repetitions with candidate order rotated and reversed.
  • Because the sweep order wasn’t fully balanced, a separate confirmation fixed the 32-row arm in advance and alternated A/B and B/A for 12 fresh repetitions per shape and GPU.
  • Every fixture includes one invalid-ID sentinel, so 16, 64 and 256 job slots mean 15, 63 and 255 active jobs.
  • All 136 sweep and 24 confirmation comparisons matched their references bit for bit. Three rows per job were also checked against CPU dequantization. Memory and synchronization sanitizers reported no errors on small fixtures.
  • Intervals are 2,000-resample bootstrap intervals over paired ratios. They describe these fixtures only.

Limits

GPU performance counters weren’t available on the test host (Nsight Compute returned ERR_NVGPUCTRPERM), so cache hit rates and DRAM traffic weren’t measured. The time saving is real for these fixtures, but cache reuse as its cause is still a hypothesis. CPU routing, graph execution and whole-model latency were out of scope.