id
papers/hina-moe-amx-general-layout
Authors
Weiling YangJunwen ZhangDezun DongJianbin FangEnda YuZhe BaiXiaopeng Deng
Affiliations
National University of Defense Technology
Problem
When routed experts live in host RAM and only a few tokens hit each one, copying an expert to the GPU costs more than computing it, so the expert FFN stays on the CPU and sits on the critical path. The CPU kernels that are actually fast, KTransformers and IPEX, get there by storing weights in an AMX-packed layout and by sharding those weights across NUMA nodes with a custom allocator. Both break what a hybrid stack needs: one weight buffer a GPU kernel can also read, and an allocator the framework already owns.
Mechanism
Keep the expert weights in the plain row-major layout the GPU library already expects, and do the layout work on the small gathered-token matrix instead. Token gathering already copies the routed tokens into an intermediate matrix I; HiNa-MoE fuses a register-level transform into that copy so I lands in the AMX-specific operand (tmm1) while the weight stays the general operand (tmm0). The same idea repeats at the two stores: the gate and up projections are written out in the layout the next AMX multiply wants, and the down projection is transposed back to row-major in registers before the store. The tile size is chosen to fill the eight 1 KB AMX tiles: k_r = 32 BF16, and m_r = n_r = 32 when there are at least 32 rows, otherwise m_r equals the row count, so AMX is used even for a single token instead of falling back to AVX-512 below four rows the way KTransformers does. When the row count is 32 or less, weights are loaded with tileloadt1, a non-temporal load that skips the hardware prefetcher and avoids evicting data that has no reuse. For NUMA, weights are allocated with ordinary page interleaving (numactl --interleave=all, 4 KB pages holding 2048 BF16 elements) and the kernel, not the allocator, recovers locality. Because expert dimensions are multiples of 128, the page-to-node mapping is periodic with a short period T_x, which is 1 or 2 rows on Mixtral-8x7B at two nodes and up to 16 on Qwen3-30B-A3B, so each node is handed only the segments that physically sit on it and partial products are reduced across nodes. Decode attention gets the same treatment: under GQA the query heads that share one KV head form a small block (8 by head dimension on Qwen3-30B-A3B), which turns the decode GEMV into a small GEMM that AMX can run after one cached layout transform of that block.
Speedup reported
Authors own measurements. FFN kernel, geometric mean of 100 runs, BF16, expert shapes from Mixtral-8x7B (4096x14336, top-2), Mixtral-8x22B (6144x16384, top-2), DeepSeek-V3 (7168x2048, top-8) and Qwen3-235B (4096x1536, top-8), 1 to 16 tokens per expert. Dual-socket Xeon Gold 6430 (32 cores x2, 2.1 GHz, 134.4 BF16 TFLOPS, 2 NUMA nodes, STREAM local 234 GB/s, cross-socket 95 GB/s) and quad-socket Xeon Gold 6448H (32 cores x4, 2.4 GHz, 300.4 BF16 TFLOPS, 4 NUMA nodes, local 158 GB/s because only six DIMMs per socket, cross-socket 32 GB/s). Average 1.73x over IPEX and 1.68x over KTransformers, peaks 3.37x and 2.16x; on the 6430 the kernel reaches 96% of measured local peak bandwidth. INT8 with the same symmetric group-wise scheme as KTransformers averages 2.18x over it on the 6430 (Figure 11, no per-point readout). Dual-socket Xeon Max 9462 in HBM-only mode, 64 GB per socket, measured peak 853 GB/s, local 425 GB/s, cross-socket 89 GB/s: FFN kernel only, since HBM cannot hold the models, averages 9.65x over oneDNN and 3.41x over IPEX. End to end, batch 1, 128 input and 128 output tokens, geometric mean of 5 runs, dual 6430 plus an RTX A6000: decode averages 1.99x, 1.70x and 1.22x over vLLM CPU-only, IPEX and KTransformers, with peaks of 2.87x, 2.09x and 1.25x, on Mixtral-8x7B, Qwen3-30B-A3B and Qwen2-57B-A14B. The IPEX and KTransformers numbers are against their own builds with the HiNa-MoE operators swapped in. Expert FFNs are about 80% of decode time, and the micro-kernel alone is +27% over the baseline in the ablation. Prefill, Qwen2-57B-A14B, batch 1, SGLang, GQA attention on the A6000 and experts on the CPU, tokens/s: with GPU expert scheduling allowed, 1526.66, 2065.21, 1977.82, 2134.25, 2085.47 at 2K, 4K, 8K, 16K, 32K, against 1187.43, 1317.39, 1377.77, 1343.52, 1252.54 for SGLang plus KTransformers, so 1.29x at 2K rising to 1.67x at 32K. Without GPU scheduling the two stay close (HiNa-MoE 1293.36 down to 1187.70 across the same lengths). On the KTransformers prefill path, attention goes from 1.536 ms and 97.88 TFLOPS at 2K to 85.466 ms and 112.57 TFLOPS at 32K, while the CPU expert FFN goes from 67.124 ms and 13.44 TFLOPS to 919.989 ms and 15.69 TFLOPS.
notes
VENUE: the schema has no pact value. The arXiv comment and the v1 HTML header both say accepted at PACT 2026 (Chicago, October 19-22, DOI 10.1145/3838684.3846846, ISBN 979-8-4007-2915-7/2026/10), but a Crossref query for that DOI on 2026-10-18 returned Resource not found, so venue is other with the real venue recorded here and venue_track main. Year is the conference year, which matches the arXiv submission year. WHAT TRANSFERS to a hand-written stack that keeps attention and shared experts resident and streams routed experts from RAM. (1) The layout decision. Packing expert weights for the CPU matrix engine (AMX here; the same pressure exists for XMX joint_matrix tiles and for any packed low-bit kernel) makes the buffer unreadable by the GPU kernel, so every expert that moves has to be reconverted or stored twice. Transforming the gathered token matrix instead costs a few register ops on a buffer that is tiny next to the expert, and it leaves one canonical weight the GPU path can consume directly. This is the same dequant-into-the-tile idea already recorded for the SYCL XMX path, applied one step earlier: the layout conversion rides on a copy the runtime already performs. (2) The tile-size rule. With eight 1 KB tiles, BF16 and a 32-wide reduction, the compute-per-byte ratio m_r*n_r/(m_r+n_r) peaks at 32 by 32, and below 32 rows you shrink m_r rather than drop to the vector unit. KTransformers only uses AMX above four tokens per expert and dips there; a hand-written kernel should not have that cliff, because decode is exactly the one-token case. (3) Non-temporal loads for the expert. At one to 32 rows the weight is streamed once, so a prefetch-triggering cached load only evicts the KV cache and the shared expert that you are trying to keep hot. tileloadt1 is the AMX spelling of that; the equivalent on a GPU is a streaming load for the expert and a cached load for everything resident. (4) NUMA without a custom allocator. Page-interleaved allocation plus a periodic segment table, T_x = 2048*N_nodes / gcd(K, 2048*N_nodes) with 4 KB pages holding 2048 BF16, gets the remote-DRAM fraction to about zero, matching KTransformers hand-sharded placement. The remote penalty it removes is large: cross-socket bandwidth is 95 GB/s against 234 GB/s local on the 6430, 32 against 158 on the 6448H, and 89 against 425 on the 9462, which is why the 9462 gap over IPEX opens up to 3.41x. On a single-socket desktop this part is worth nothing; on a dual-socket box holding the experts in RAM it is most of the win. (5) The crossover the paper plots but does not tabulate: on the 6430 plus A6000, CPU-side expert execution stays competitive at small and moderate batch sizes and the GPU wins once the transfer is amortized, which is the decision a hand-written scheduler has to make per expert per step. WHAT DOES NOT TRANSFER. Every number is the authors own, and the FFN curves are figures, so 1.73x and 3.37x are the only kernel-level figures stated in prose. The end-to-end IPEX and KTransformers comparisons are against builds that already contain the HiNa-MoE operators, which flatters the baseline and makes the 1.22x over KTransformers a lower bound on the kernel difference rather than a system comparison. The 2.09x peak is over IPEX, not over a GPU-offload system, and no absolute tok/s is printed for decode. The NUMA scheme assumes dimensions divisible by 128 and a page holding 2048 elements; a low-bit expert changes both the element count per page and the period, so the segment table has to be recomputed rather than copied. There is no code link, and the paper deliberately does not change the offload policy, so it says nothing about overlapping the RAM fetch with compute.