fix(distill): pre-warm non-LoRA backward kernels (PMAT-698g)

[noahgift]

Summary

Phase 3 dispatch v8 on gx10 reached the training loop and the first backward step started JIT-compiling backward kernels on demand mid-training, then failed with:

forward_backward_with_grad returned None (CUDA stream poisoned or gradient shape mismatch)

This is the documented Blackwell sm_121 bug (trueno#200, CLAUDE.md "Backward kernels: Crash because they compile on-demand when GPU is already active").

Root cause

pre_warm_lora_backward_kernels had a short-circuit:

if lora_rank == 0 {
    eprintln!("[BWD-PREWARM] Skipping (lora_rank=0)");
    return Ok(());
}

The function name implies LoRA-only, but its BODY contains pre-warms for silu_backward, batched_softmax_backward, and batched_rms_norm_backward — which distillation training also needs. The early-out left those to JIT during training, exactly when Blackwell can't.

Fix

Restructure — only LoRA-specific GEMM-backward warm-ups are gated on lora_rank > 0. Activation/norm/standard-FP32-GEMM backward kernels always pre-warm.

Test plan

• cargo check --features cuda clean build
• 18 cuda_backward lib tests pass
• Live gx10 dispatch reaches stepping (post-merge)

Defect cascade so far

Stage 4 of the Phase 3 cuda dispatch cascade:

Stage	PR	What
1	#1804 PMAT-700-B	Skip PTX GEMM pre-warm when cuBLAS active (JIT cache pressure)
2	#1808 PMAT-698e	Cap max_seq_len at 2048 (workspace sizing)
3	#1809 PMAT-698f	Accept APR magic in load_safetensors_weights (format compat)
4	THIS PMAT-698g	Pre-warm non-LoRA backward kernels (stream poisoning)

Each surfaced the next defect once unblocked.

🤖 Generated with Claude Code