fix(distill): warm! macro used hardcoded "silu_forward" key — THE Blackwell root cause (PMAT-698j)

[noahgift]

THE root-cause bug

The warm! macro in pre_warm_for_model had self.get_or_compile("silu_forward", &ptx)? hardcoded — every pre-warm call stored its compiled module under the SAME hashmap key. Only the first kernel actually got compiled; all subsequent warm!() invocations hit the Occupied path and silently discarded their PTX.

At runtime, every kernel looks up its real key (e.g. batched_rmsnorm_fwd_896_eps358637bd) — cache miss → JIT mid-forward-pass → Blackwell sm_121 stream poisoning → forward_backward_with_grad returned None.

Discovered via PMAT-698i diagnostic logging

The diagnostic PR's [FWD-CACHE] log surfaced every "pre-warmed" kernel actually JIT-compiling at first runtime call. Proves the cache held exactly one entry (under silu_forward), no matter how many warm!() calls ran.

Fix

 macro_rules! warm {
     ($key:expr, $kernel:expr) => {{
+        let key = $key;
         let ptx = $kernel.emit_ptx_for_target(&target);
-        self.get_or_compile("silu_forward", &ptx)?;
+        self.get_or_compile(&key, &ptx)?;
         count += 1;
     }};
 }

One-character behavioral change. The previous PMAT-698g/h fixes were defense-in-depth on the backward cache (whose warm! macro was correct); this fixes the dominant forward-cache failure that all 5 prior iterations were chasing.

Cascade summary

#	PR	Status
1	#1804 PMAT-700-B	✅ legit independent (JIT cache pressure on cuBLAS-active hosts)
2	#1808 PMAT-698e	✅ legit independent (workspace sizing)
3	#1809 PMAT-698f	✅ legit independent (APR format compat)
4	#1810 PMAT-698g	✅ legit defense-in-depth (backward cache gap)
5	#1813 PMAT-698h	✅ legit defense-in-depth (rms_norm_gamma_reduce gap)
6	#1815 PMAT-698i	✅ diagnostic (surfaced THE bug)
7	THIS PMAT-698j	🎯 THE root cause

Test plan

• cargo check --features cuda clean
• 366 autograd lib tests pass
• Live gx10 dispatch: zero [FWD-CACHE] Compiling events post-pre-warm
• Forward pass + backward pass + optimizer step actually run
• Smoke completes 50 steps with final_loss < initial_loss (F-DISTILL-SMOKE-001)

🤖 Generated with Claude Code