cuda: preserve weighted RMS reduction order with vector loads

[GenerelSchwerz]

Overview

Use four original logical RMS threads per physical thread to load/store float4 values with a 256-thread block while preserving the original 1024-thread accumulation and reduction tree. Weighted RMS keeps normalized-before-gamma multiplication and its existing parameter broadcast rules. No temporary gamma-first stores or changed sum association are needed.

The new path requires plain F32 weighted RMS, no ADD/scale, physical warp size 32, width >= 1024 divisible by 4, and aligned source/destination rows and strides. Small/odd/unaligned cases, unweighted RMS, live-output and explicit in-place fallbacks, HC pending/custom reads and quantized producer emission retain their existing paths. No model names, fixed hidden widths, new graph matcher or allocator policy.

This is a one-file leaf, 41 additions and 1 deletion, atop PR156 prerequisite commit 08b1d9bb1d1fa59996052591f3197782871030b7. The PR base remains reference/upstream-kernels-20261001, official dcd387a412ca54e172a8d60eb71ef6753850c8ca. The full diff includes earlier prerequisite commits; review the one-file leaf commit for the RMS change.

Measurements

RTX 5070 Ti, SM120a, matching PR156 control, normal upstream scheduler/allocation/dispatch and GPU events around 2000 captured graph replays per measurement.

Component	Control us	Candidate us	Less time
Single-row BF16 HC, no pending residual, six pairs	22.521208	20.501712	8.967%
Independent repeat, three pairs	22.521696	20.499537	8.979%
Three-row BF16 HC	28.662304	28.662895	Flat
Pending residual BF16 HC	22.542255	22.541616	Flat
Standalone width 16384 weighted RMS	6.143632	6.148304	0.076% more time

Other standalone norms, F16/F32 multirow components and live-intermediate controls are flat. These are isolated component measurements, not model-serving, MTP, pristine-upstream or additive series gains. Older-GPU speed is unmeasured.

Validation

• 25 paired cases and four datasets per case preserve 2,946,800 original GPU output elements byte-for-byte, plus inputs and independent CPU-reference bytes. CPU NMSE bound remains 5e-4. Coverage includes vector tails, 4D gamma broadcast/repeat, padded and unaligned views, finite extremes/subnormals, zero/epsilon, ADD/live/small/odd fallbacks and full BF16/F16/F32 components.
• Normal allocation actually reuses a dead ADD producer as the weighted destination. A live producer remains separate. All 12 targeted memcheck/racecheck/synccheck checks pass; alias and ownership checks remain enabled.
• Both modules pass 97/97 existing filtered test-backend-ops CUDA cases. All 818 SM120a shader resource records match control; 817 executable shaders are unchanged after nvcc internal source-token normalization. Only the intended weighted 256-thread RMS shader changes, remaining REG40/STACK0/LOCAL0. No new shader entry.
• SM75/80/89/90 compilation passes with REG58/32/40/32 and zero stack/local memory for the changed shader. SM90 reports 1024 bytes static shared memory. Compilation does not establish older-GPU performance or HIP/MUSA validation.
• Whole ordered build/GPU leases, verified 4 GiB host cap, swap0, finite timeouts and complete process-tree teardown. Successful work was reused after recorded fixture bookkeeping corrections. No full-CI, perplexity or model-server claim.

Evidence and exact commands/source/module/raw-array/profile hashes: /home/gencoolpc/moe-cache-tests/results/upstream-kernel-series-20261001/HYPERCONNECTION/HC-ORDER-ASSESSMENT-20261008, especially FINAL-NORM-REGRESSION-RESULTS-V1.json, NORM-AUDIT-REPORT-V1.json, ORDER-COMPARISON-RESULTS-V1.md and NORM-ROOT-SOURCE-REVIEW-V1.md.

Related work and scope

Upstream PR20520 and PR29720 already propose vectorized RMS with a different sum tree. This fork alternative preserves the original tree and existing callback/shader interface. No speed comparison against those branches is claimed. Future upstream integration should reconcile these approaches with the existing authors.

Pending HC reads, quantized emission, shared activation staging, prefetch and paired down/injection are separate work and are not covered by this leaf.

Requirements

• Maintained-fork publication is explicitly authorized by the repository owner. No upstream submission.
• AI usage disclosure: YES. Codex implemented, reviewed and qualified this owner-requested change. The owner retains review and maintenance responsibility.