[fix](kt-kernel): AVX-VNNI-256 RAWINT4 per-expert weight loading

[blazingphoenix7]

What does this PR do?

Native RAWINT4 checkpoints (the Kimi K2 layout) cannot load on CPUs where the AVX-VNNI-256 backend is selected. NativeMoEWrapper always hands the C++ side per-expert weight pointers (moe_config.gate_projs = [[...]], one pointer per expert), but the backend's TP wrapper only accepted flat buffers:

if (config.gate_proj == nullptr || config.gate_scale == nullptr) {
  throw std::runtime_error("AVX-VNNI RAWINT4 MoE only supports flat packed INT4 with KGroup bf16 scales");
}

The wrapper never sets the flat pointers, so on a build without AMX support running on a CPU with avx_vnni, where _select_rawint4_backend prefers this backend for group sizes up to 256, weight loading always throws. The exception is raised on the CPUInfer worker thread, which has no handler, so the process dies during startup. This is the consumer-CPU lane that #1897 asked about.

The fix teaches TP_MOE<AVXVNNI256_RAW_INT4_MOE_TP>::load_weights to gather TP slices from per-expert pointers into the same per-partition staging buffers the flat path already builds, then hand off to the unchanged per-part loader. The slicing arithmetic is the same as the per-expert branch of the K2 K-Group loader in k2-moe.hpp, which serves the AMX/AVX512 builds today: gate/up take a contiguous row block per partition, down takes a per-row column gather. Because the gather works for any partition count, RAWINT4 serving with --kt-threadpool-count greater than 1 also works on these hosts now. The flat path is untouched, and the transient staging cost is identical to flat mode and to the K2 loader.

Found by equivalence testing the RAWINT4 load paths against each other (same weight bytes loaded two ways must produce bitwise identical forwards) while looking into the reports in #2076.

Changes

• kt-kernel/operators/avx2/rawint4_avxvnni-moe.hpp: per-expert pointer support in the TP wrapper's load_weights, mirroring the K2 loader. Flat mode unchanged.
• kt-kernel/test/per_commit/test_moe_rawint4_load_equivalence.py: new CPU per-commit test, registered with register_cpu_ci. Each check runs two backend instances holding the same weight bytes loaded through different paths and compares full forward outputs with torch.equal, since identical kernels on identical bytes must agree bitwise:
  • per-expert vs flat load on AVX-VNNI-256, with one and with two worker subpools (the two-subpool case exercises the new down-projection column gather)
  • per-expert vs flat load on AVX2 (single subpool, that backend's supported mode)
  • permuted vs identity physical_to_logical_map through per-expert load, on both backends
  • int32 pack-quantized storage vs uint8 byte-packed storage, repacked independently through the logical nibble order (pins the layout identity behind the loader normalization from fix: normalize compressed RAWINT4 weights #2075)
  • the one non-bitwise check: one vs two subpools with the same flat weights, compared at 5e-3 mean relative tolerance, since the split down projection sums its partials in a different order

This PR is correctness and coverage only; no performance behavior changes and none is claimed.

Not addressed here, left as known limitations of the AVX2 sibling: AVX2RawInt4_MOE still requires --kt-threadpool-count 1 in per-expert mode because it serves weights zero-copy from the source tensors and a direct pointer cannot express the down projection's column slice. Lifting that needs a strided weight view or a copy fallback, a separate change. Hosts with avx_vnni stop hitting that limit with this fix since selection prefers the VNNI backend there. write_weights_to_buffer on the VNNI backend (GPU expert staging) also remains unimplemented, as before.

Testing

cd kt-kernel/test
python3 per_commit/test_moe_rawint4_load_equivalence.py
python3 per_commit/test_moe_rawint4_accuracy.py

On an Intel Core Ultra 5 225 (AVX2 + AVX-VNNI, no AVX512, no AMX), Linux, gcc 15, CPUINFER_CPU_INSTRUCT=NATIVE CPUINFER_ENABLE_AMX=OFF:

• test_moe_rawint4_load_equivalence.py: all six tests run (this host has avx_vnni, so nothing skips), PASSED.
• test_moe_rawint4_accuracy.py: both backends at qlen 1 and 16, three iterations each, every diff under its threshold (about 0.003 for AVX2, about 0.014 for AVX-VNNI against the torch reference), PASSED.
• Rest of test/per_commit on this host: test_basic_cpu, test_load_experts_count_guard, test_port_checker, test_moe_avx2_accuracy_bf16, test_moe_avx2_accuracy_fp8, test_moe_gptq_int4_accuracy and the placeholders pass; the sglang adapter tests pass under pytest (12 passed, 5 integration cases skip without model files). The eight test_moe_amx_* files stop at torch calls that require CUDA before reaching any kt-kernel code (CPU-only torch here); they cannot run on this machine with or without this change.

Rebuilding with the fix reverted, the first per-expert AVX-VNNI check aborts the process with the exception quoted above escaping on the CPUInfer worker thread, while the flat-mode loads in the same run keep working. That is the exact failure the serving path hits today, and it shows the flat path is not affected by this change.

Before submitting

• Read the contributor guideline.
• Added regression tests for the fixed path.

[gemini-code-assist]

Code Review

This pull request adds support for loading weights via per-expert pointers in the AVX-VNNI RAWINT4 MoE operator, alongside a comprehensive Python test suite to verify load-path equivalence across different configurations. Feedback suggests adding defensive validation to ensure all per-expert pointer vectors are properly populated and contain non-null pointers to prevent potential segmentation faults.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.