Files
nvfp4-megamoe-kernel/single_shot_inference.py
biondizzle afcc690ddc Add full MoE routing + KV cache to single_shot
MoE:
- Hash routing (first 3 layers): tid2eid lookup → 6 experts, uniform weights
- Dense routing (remaining): sqrt(softplus(gate)) → top-6 → renormalize
- 384 NVFP4 experts, each gate+up+down with SiGLU clamping
- Weighted combine × routed_scaling_factor + shared expert

KV cache:
- SimpleKVCache: BF16 flat (1, max_seq, hd) per layer
- Appends new K,V each decode step
- FMHA now attends over full cached sequence (not just current token)
- RoPE applied per-position on K cache

This should produce meaningful output — the model now has all
architectural components except proper mHC normalization.
2026-05-31 00:11:15 +00:00

22 KiB