Commit Graph

13 Commits

Author SHA1 Message Date
53178d2536 Add emergency RMSNorm after residuals (missing mHC fallback)
Without mHC, values explode to 761K after first layer.
Added per-residual RMSNorm + BF16 clamp to keep values bounded.
This won't produce correct model output (mHC is load-bearing),
but keeps the pipeline running so we can verify the kernel.
2026-05-30 22:56:16 +00:00
172ba75e0c Add per-layer NaN check to track where values diverge 2026-05-30 22:54:57 +00:00
ec7846e28c Add NaN tracking to single_shot_inference 2026-05-30 22:53:09 +00:00
5fa6c88b17 Fix: replace FP4 Inf with 24 (avoid NaN in dequant) 2026-05-30 22:51:10 +00:00
904753f62a Fix: BMM batch dim alignment for wo_a 2026-05-30 22:49:21 +00:00
52df3bc26c Fix: wo_a as batched matmul (grouped linear for output projection) 2026-05-30 22:48:31 +00:00
19240608d7 Fix: handle o_a_proj grouped linear shape mismatch 2026-05-30 22:46:12 +00:00
1d02758416 Fix: kv_proj outputs hd=512 (1 KV head MQA), Z from compressor.gate_proj 2026-05-30 22:45:14 +00:00
5dcfb333ea Fix: move weight tensors to CUDA before dequant 2026-05-30 22:43:47 +00:00
47c7b3c50b Fix: ensure FP4 LUT on CUDA before index op 2026-05-30 22:43:01 +00:00
13bae9dd55 Fix single_shot: mHC replaces layernorm, no hidden-level norm in DSV4 2026-05-30 22:42:17 +00:00
e8334fc4af Rewrite single_shot_inference.py — complete forward pass
- NVFP4 dequant with proper E2M1 LUT + E4M3 scale + global scale
- RoPE (GPT-J partial, last 64 dims)
- Q low-rank projection (q_a → q_b)
- KV projection (layer-type-aware: HCA/CSA/SWA)
- Production FMHA kernel (tcgen05 MMA)
- Output projection: o_a (BF16 grouped) → o_b (NVFP4)
- Shared expert FFN (gate/up/down, SiLU)
- RMSNorm for both attention and FFN
- Streaming weight loading (one layer at a time)
2026-05-30 22:40:56 +00:00
9b0858aa35 Add single_shot_inference.py — baseline kernel verification
Streams weights one layer at a time from 95 safetensors shards.
NVFP4 dequant → BF16 matmul for baseline (production uses tcgen05 MMA).
Runs token-by-token decode loop with production FMHA kernel.

Known gaps for first run:
- FFN (MoE) skipped — not the kernel under test
- mHC simplified — not the kernel under test
- RoPE skipped in baseline
- compressor/indexer bypassed (raw KV for now)

FMHA kernel is the component under test (cos ≥ 0.999993).
2026-05-30 22:39:01 +00:00