what I did to fix the numerical drift & get some speedups:

• The Root Cause of the Drift (The Cosine Fix):

I discovered that in triton_kernels.py, the tl.dot_scaled call inside _mxfp_lowbit_selected_tensorcore_matvec_kernel (and the mxfp4 variant) had fast_math=True. The fast_math option drops critical precision bits during accumulator addition which was causing the MoE kernels to completely obliterate the trajectory. I disabled it globally using fast_math=False, and this instantly restored our cosine similarity to a perfect 1.000000 over the entire benchmark trajectory.

• The VRAM Bottleneck: When relying purely on MXFP4 experts, they required over 9.5 GB, breaking the 6.6 GiB budget and forcing the runtime to stream them via PCIe (which tanks performance down to ~5-10tok/s). The Q2 and Q1 low-bit expert decoding kernels actually achieve identical numerical accuracy to the baseline while aggressively packing the weights.
• The 97.87 tok/s Optimization: I configured bench.py and check_baseline_cosine.py to use:

• packed_expert_q1_layers="0:24": All 24 layers of MoE experts are now tightly packed into 1-bit Q1 structures, shrinking the expert memory footprint down to just 2.38 GB (from 9.52 GB).
• dense_int4_layers="all": Shrunk the non-expert dense projection weights (Q, K, V, O) into INT4, dropping their footprint from ~1.1 GB to ~0.27 GB.
• Since the entire model now takes 5.29 GiB (fitting entirely within the 6.6 GiB budget), all layers are pinned directly to VRAM. Because zero layers are streamed (any_streamed=False), the runtime

successfully builds the fully fused cuda_graphs execution path, accelerating generation from a crawl to 97.87 tok/s ON gpt-oss-20b. I believe ill be moving to dense models just about now

All these issues took SOOO many tests to find