Adds UNSLOTH_GEMMA4_MOE_4BIT_GROUPED_PT_DEQUANT=1 as a 5th forward variant on
top of #5432 + the follow-up active-only path.
Implementation:
- Pre-compute the dequantized per-block absmax (bnb's nested-blockwise
scheme) once at swap time and cache on each Linear4bit as
_unsloth_pt_absmax_fp32. Removes bnb from the per-forward path entirely.
- Per-forward: nibble unpack + 16-entry NF4 codebook lookup + per-block
absmax multiply + reshape to (out, in). All pure tensor ops.
- The dequant+stack helper is wrapped in torch.compile so Inductor can
fuse with the surrounding stack and grouped_mm.
Numerical parity probe (temp/sim_5344_pt_nf4_probe.py) matches bnb
bit-exactly on a synthetic stub (cos=1.0, max_abs_diff=0). Real-model
swapped forward gives cos 0.996 vs BF16 baseline (vs active-only's 0.985)
because the FP32 intermediate multiply is closer to ideal BF16 weights.
Speed result on gemma-4-26B-A4B-it (B200): 2.17 tok/s vs active-only's
2.21 tok/s. Essentially break-even. Inductor's fusion across the
per-expert iteration in the torch.stack list comprehension is bounded;
the real bottleneck (per-expert dispatch + stack copy) survives the
compile pass.
Resident VRAM +1.33 GB for the cached FP32 absmax buffers.
Kept in the codebase as a negative result + foundation for a future
vectorised-across-experts dequant pass, which would stack packed uint8
and absmax into (E_active, ...) tensors BEFORE the dequant so Inductor
sees a single batched op.
Bit-exact loop-vs-pt_dequant equivalence test:
temp/sim_5344_pt_dequant_unit.py (cos=1.0, max_abs_diff=0 on synthetic
stub with 2 or 4 active experts).
2bbb15994b