fix(cpu-moe): honor padded fp8 scale strides, add avx512f tier and float64 parity (builds on #36)

[ChenyuHeee]

Relationship to #36

This builds directly on @gdevenyi's #36 and includes his commit unchanged (4a04c9b, +360 lines — the format definition, bank plumbing, ABI probe and benchbw integration). My three commits on top are +262 lines. If you would rather land #36 first, these three cherry-pick cleanly onto it and I am happy to send them there instead — I opened this separately only because the stride bug below is a correctness issue that would otherwise ship silently.

Credit for the feature is his. What follows is a bug fix and test hardening.

1. The scale bank's trailing dimension is padded (correctness)

kernel/aot_models.py pads the scale bank's last dimension so per-expert rows land on a 16-byte boundary for the fused copy:

def fp8_block_scale_pad(rows: int, cols: int) -> int:
    """Trailing scale-bank dim padded so per-expert row bytes are 16B-aligned (fused copy)."""
    while (rows * cols * 2) % 16:
        cols += 1
    return cols

So the row stride is not ceil(K/128). #36 recomputes it from the unpadded block count in two places (cpu_moe_ext.cpp and the shape assert in cpu_executor.py), which reads scales at the wrong offset whenever padding applies.

Qwen3.8 is the case already documented in the tree — models/qwen3_5_moe/weight.py records down_scale as 20x5 bf16 = 200 B, and 200 % 16 != 0 pads cols 5 → 6. The assert demands (E, 20, 5) and fails; with asserts off the kernel silently produces wrong numbers.

Neither the benchbw fp8 workload (H=2048, I=512) nor DSV4-Flash (H=7168, I=2048) trips it — both happen to give rows*cols divisible by 8, which is why #36's own tests are green.

The fix reads the stride off the tensor and threads it through to the kernel. Reverting just the C++ half turns 7 tests red (4 on a CPU-only host).

Worth noting: the Triton kernel was already stride-aware. Feeding it padded banks keeps the GPU cross-check green, so the CPU path was the only one getting this wrong.

2. The parity test could not see accumulator precision loss

The existing check compares two fp32 accumulations (CPU vs Triton) at rel < 2e-2, so common-mode error cancels and a kernel with ~1e-3 error passes unnoticed.

Added a float64 reference — torch.float8_e4m3fn dequant with the bf16 block scales, GEMV in float64 — over mixed-sign full-range e4m3 weights, zero-mean activations, and ragged K. Measured error is 2.674e-3 against a 5e-3 bound (the executor stores bf16, so the floor is not tighter than that).

No existing assertion was weakened; the 2e-2 check is untouched.

Also addressed the two Copilot comments on #36: the ISA-tier test now actually sets FREETOKEN_CPU_MOE_ISA, and test_subnormal_and_sign_decode_exactly's docstring no longer claims to validate the C++ decoder when it only exercises a PyTorch round-trip.

3. An avx512f tier for CPUs without AVX512-BF16

#36's fast path is _mm512_dpbf16_ps, so a CPU with AVX-512 but no BF16 extension falls all the way to the AVX2 tier — 8-wide fp32, 512-bit registers unused. That is Skylake-SP, Cascade Lake and Ice Lake-SP.

Added an avx512f tier between them, following the existing per-function __attribute__((target(...))) + runtime dispatch structure. All three reachable tiers are bit-identical across all 256 e4m3 codes.

Measured on Xeon Platinum 8368 (Ice Lake-SP, no avx512_bf16), three repetitions each:

	avx2	avx512f
single-thread, cache-resident	3.31 GB/s	3.94 GB/s	+19%, intervals disjoint
ft bench bw --dtype fp8, 75 threads	60.4 GB/s	62.7 GB/s	+4%, intervals overlap

At 75 threads the kernel is memory-bound and the tier is largely washed out — the project's own benchmark cannot resolve it. The single-thread number is the honest figure.

4. Tests now cover the code they ship

Before this, deleting the entire avx512f branch from select_fp8dot left every test passing, because isa_name() reports select_dot() (the bf16 chain) rather than the fp8 dispatch — and no test ran numerics under a forced tier. Injecting a ragged-tail bug into dot_fp8_block_avx2 was invisible for the same reason.

Each reachable tier now runs the float64 parity GEMV, and the fp8 dispatch is observable. Verified by mutation:

mutation	before	after
drop the avx512f dispatch branch	all green	assert 'avx2' == 'avx512f'
ragged-tail scale bug in dot_fp8_block_avx2	all green	rel = 0.257
kernel ignores the passed stride	—	7 failed (4 CPU-only)
break only scalar	—	caught, that tier only
break only avx512f	—	caught, that tier only

The module-level CUDA guard also moved onto the tests that need a GPU — all ten were skipping on a CPU-only host, which is exactly where a CPU MoE kernel matters most.

Known limits, stated up front

• The avx512bf16 fp8 tier has never executed here. This host lacks the ISA and I have no SDE, so that path is unvalidated on hardware — and the code comment says it is the tier that matters. A mutation that halves its output goes undetected on this machine. Someone with Emerald Rapids or Sapphire Rapids should confirm it; @gdevenyi's 6526Y would do.
• Tier verification is name-mediated. A tier that is numerically correct but silently does not use AVX-512 would pass. Correctness tests cannot distinguish that; only the throughput number argues the tier earns its place.
• CPU-only runs exercise num_threads=1 only. Row splitting is format-independent, so the risk is low, but multi-threaded fp8 is covered solely by the GPU cross-check.

Test results

CUDA:      12 passed, 1 skipped     (avx512bf16 tier skipped — host lacks the ISA)
CPU-only:   9 passed, 4 skipped
tests/moe: 130 passed, 7 skipped, 1 failed

The one failure is test_cpu_moe_q4_0.py::test_cpu_decode_q4_0_matches_ggml_mmvq, a GGUF JIT compile error under gcc 15 (need 'typename' before ... in torch headers). It fails identically on main — unrelated to this PR.

Built with gcc 15.2.1 / binutils 2.44. Function-level AVX-512 target attributes need binutils >= 2.36; the gcc 11 / binutils 2.35 on RHEL 9 cannot assemble them.