[WS1 Closeout][Backward] Reusable gradient-invariance harness and operator adapters #270

Parent: #266
Class: Must (blocks Full WS1 EXIT; required by #150 / C10)
Depends on: C1 (#267), C2 (#268)
Source: #153 Backward-pass consistency

Goal

Ensure ops and the full chain share one comparison semantic for training-style gradients under batch/chunk, so tests do not invent their own upstream grads or denominators.

Scope

• Cross-config gradient-invariance API
• Separate gradient accuracy (vs reference VJP) from invariance (across configs)
• Fix across configs: upstream grad, loss reduction, active-token denominator
• Compare on C2’s same logical sample/token multiset: N× B=1 singleton fixed-order aggregate vs one B=N
• Compare all required trainable parameter gradients; give stable logical names for required activation/input VJPs
• Register adapters for differentiable WS1 ops (OP_SPECS or equivalent; name-only mention in a chain report does not count)
• Audit: no atomicAdd / no shape-dependent bwd accumulation on BI candidates
• Gradient thresholds only from C1
• Register and validate required gradient adapters separately for cuda_bf16 and triton_cuda_bf16; a missing Triton backward for a required differentiable node is red

Suggested API

assert_gradient_batch_invariant(op, configs, contract) -> GradientInvarianceReport

Required adapters

• RMSNorm dx/dweight; GEMM dX/dW; Attention dQ/dK/dV
• Embedding / LM-head; selected-logprob
• RoPE / SiLU / SwiGLU
• Pack/KV when declared supported and differentiable

Acceptance criteria

• Every WS1 op with backward is enumerable by the harness
• Accuracy / invariance separate; thresholds via C1
• Required Batch/Chunk gradient invariance is bitwise (atol=0, rtol=0) after logical aggregation; only gradient accuracy uses FP32-reference tolerance
• B1/BN use same sample set, upstream grad, active-token denominator, loss reduction, fixed aggregation order
• Norm, embedding/LM-head, QKV/out/MLP weight grads and required input/VJPs have stable names and verdicts
• Report max abs/rel, grad tensor name, first failing op
• Status matrix has no untracked red (else open Blocker)
• No private gradient thresholds
• CUDA and Triton required gradient adapters are complete and green; neither profile may borrow the other's implementation as fallback

Relation to EXIT

• C10 Must: full-chain training-style gradients invariant on Batch=1/N and Chunked-Prefill on/off
• C4 is a hard prerequisite of C10; C4 green alone ≠ EXIT