onnxruntime

a926b837 - Fix CUDA LayerNorm/RMSNorm silent wrong output for rank-reduced scale (#28971)

Fix CUDA LayerNorm/RMSNorm silent wrong output for rank-reduced scale (#28971) ### Description The CUDA `LayerNormalization` / `RMSNormalization` kernels only support broadcasting scale/bias through the encoded `broadcast_param` fast-path; they do **not** implement the generic NumPy-broadcasting path the CPU kernels use. `LayerNormHelper::GetBroadcastParam` required `scale.rank == X.rank` to emit a fast-path value, so a rank-reduced but NumPy-broadcastable scale — e.g. scale `(S, W)` for input `(B, S, W)` with `axis=2` — returned `kLayerNormInvalidInput`. On CPU this falls back to the correct generic path, but the CUDA kernel silently used `broadcast_param = -1`, computing offset `norm_size * (x_row % 1) = 0` and applying `scale[0:W]` to **every** row → wrong results (no error raised). ### Impact This breaks Mamba2 / hybrid SSM models (e.g. **NemotronH**), whose gated `RMSNorm` uses a per-group scale `(n_groups, head_dim)` broadcast over the batch dimension. Such models produce correct logits on CPU but a wrong argmax on CUDA. Minimal repro (fp16): ```python # X:(5,8,960), scale:(8,960), axis=-1 # CPU matches numpy reference; CUDA was off by ~13 (max-abs) before this fix. ``` ### Fix Right-align the scale/bias shapes to `X`'s rank inside `GetBroadcastParam` (NumPy-style, prepending 1s) before matching the `(B,S,...)`, `(B,1,...)`, `(1,S,...)` fast-path layouts. A rank-reduced `(S, W)` scale now maps onto the existing `(1, S, ...)` fast-path and yields the correct per-row offset on CUDA — identical to an equal-rank `(1, S, W)` scale. - Equal-rank shapes are unchanged (right-alignment is the identity). - The CPU generic path is unaffected. ### Testing Adds a CPU+CUDA regression test `RMSNorm_Scale_RankReduced_Sx3_Axis2`. Verified end-to-end: the NemotronH prefill argmax now matches between CPU and CUDA, and a broadcast sweep (RMSNorm + LayerNorm, multiple scale shapes, CPU vs CUDA vs numpy) passes at float epsilon with no regressions. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>