onnxruntime

1e3c937d - [CUDA] Optimize GatedDeltaNetChunkedKernel with ldmatrix/SyncBeforeStore and add GatedRMSNormWarp128Kernel (#33601)

[CUDA] Optimize GatedDeltaNetChunkedKernel with ldmatrix/SyncBeforeStore and add GatedRMSNormWarp128Kernel (#33601) ## Description Optimizes `GatedDeltaNetChunkedKernel` (and shared `gated_delta_net_mma.cuh` helpers) and `GatedRMSNorm` for Qwen3.5 / Qwen3.8 linear-attention prefill and speculative decode on SM80+/SM90 GPUs: 1. **Hardware `ldmatrix` & Direct Shared-Memory Epilogues in `gated_delta_net_mma.cuh`**: - Replaces scalar 16-bit shared-memory loads + `PackHalf2` in `LoadFragA<Transposed>` and `LoadFragB<Transposed>` with `ldmatrix.sync.aligned.m8n8.{x4,x2}{.trans}.shared.b16` on `__CUDA_ARCH__ >= 800`. - Adds 64-bit/32-bit `StoreC2` (`float2` / `PackHalf2`), `SyncBeforeStore`, and `IdentityMinus` to `SmemGemm`, and vectorizes 3-slice `S -> s_h` staging in `SmemGemmFp32State` with `float2` loads and `PackHalf2` stores. - Stages the 16-element forward-substitution column in `BuildTriInverse` in unused `tf` shared memory instead of a thread-local `float col[16]` array (`STACK:64 -> STACK:0`) and writes Horner steps (`I - nb @ ti`) and `ti @ db` directly to `ti`. 2. **Barrier & Shared-Memory Round-Trip Elimination + `uint4` Vectorized Global I/O in `GatedDeltaNetChunkedKernel` (`gated_delta_net_impl.cu`)**: - Vectorizes global `key`/`query`/`value` loads and `output` stores with 16-byte `uint4` transactions. The vector paths are taken only when the bound buffers are 16-byte aligned (checked at run time, so offset IOBinding views use the existing scalar paths), and the dynamic shared-memory base is `__align__(16)`. - Passes `c_scale = __expf(g_total)` directly to `SmemGemm<DK, kDVB, BT, true, false, true>(s.s_f, ...)`, eliminating the 8,192-element `s.s_f *= decay_all` shared-memory loop and barrier. - Writes `s.u_h = s.ti_h @ s.u_h` directly via `SyncBeforeStore=true`, caches `qk_l2_norm` inputs in registers, overlaps `InclusiveScanBT` before the barrier, and vectorizes `s.v_h -> s.u_h` and `s.k_h`/`s.q_h` decay scaling in pairs, eliminating 6 `__syncthreads()` barriers and 5 shared-memory matrix round-trips per chunk. 3. **4-Row/CTA Warp-Shuffle `GatedRMSNormWarp128Kernel` (`linear_attention_gates_impl.cu`)**: - For `sizeof(T) == 2 && norm_size == 128`, replaces the 1-row/CTA scalar `cub::BlockReduce` kernel (which read `input` twice and used `__syncthreads()`) with a 4-warp/CTA `uint2`-vectorized warp-shuffle kernel (`REG:26 STACK:0 SHARED:0`, 0 barriers, 4x fewer CTAs). It is selected only when the output, input, scale and gate pointers are 8-byte aligned; otherwise the generic kernel runs. ## Testing - `GatedDeltaNetTest.ChunkedMisalignedFloat16IoBuffers` / `ChunkedMisalignedBFloat16IoBuffers`: bind query/key/value/output through IOBinding at 2-byte and 8-byte offsets into a larger allocation, with `ORT_GDN_PLAN=chunked` pinned, `chunk_size` 64 and 32, and the BF16 `fp32`/`fp16`/`bf16` precision modes. They fault if either alignment guard is removed. - `ContribOpGatedRMSNormTest.Float16_PartialFinalCta` / `BFloat16_PartialFinalCta`: 1, 3 and 5 rows (a partially filled last CTA) with SiLU and sigmoid gates. - Verified on an H200 (SM90): `GatedDeltaNet*:ContribOpGatedRMSNormTest.*:ContribOpLinearAttentionGateTest.*:*LinearAttention*` passes (142 passed, 19 WebGPU tests skipped). ## Performance (NVIDIA H200, Qwen3.8-Flash-Next 36 GDN layers) - **`GatedDeltaNetChunkedKernel` (=2048$, 36 layers)**: - Executed instructions cut by **`-45.5%`** (`264.53M -> 144.05M` per launch), stack frame reduced from `64 B -> 0 B`. - Latency reduced from **`34.74 ms` (`964.96 us/call`) -> `22.18 ms` (`616.11 us/call`) (`1.57x` faster, `-12.56 ms` saved per 2K chunk)**. - **`GatedRMSNormWarp128Kernel` (=2048$, 36 layers)**: - Latency reduced from **`3.08 ms` (`85.48 us/call`) -> `1.05 ms` (`29.08 us/call`) (`2.94x` faster, `-2.03 ms` saved per 2K chunk)**, achieving `85.36%` Compute SOL and `89.73%` occupancy. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>