[why1te]

[DSv4][P1-5] rmsnorm_residual

Summary

Implements P1-5 rmsnorm_residual from #383, targeting dsv4-p1-dev. Adds CUDA/Triton forward and backward, autograd, registry fallback, and P1 providers.

The operator returns RMS-normalized output plus a copy of the original input—not RMSNorm(x + residual). Both backends pass bitwise oracle alignment and P1 acceptance on H100.

Implementation

• BF16 activations/outputs; FP32 gamma, saved inverse RMS, and explicit dx/dgamma. Autograd returns BF16 x.grad and FP32 gamma.grad.
• D=4096 in production; D=128 for fixtures; eps=1e-6.
• Preserves rsqrt(mean(x²) + eps), (x * r) * gamma, ascending FP32 reductions, and separate rounding. No atomic reductions, Split-K, or Stream-K. Residual copying does not change the reduction layout.
• FP32 gamma is maintainer-confirmed. The CPU-oracle manifest updates 59 affected hashes; residual, Sinkhorn-edge, and gamma-gradient hashes remain unchanged.
• Registry order: CUDA → Triton → PyTorch reference. Providers replace P1-5 only; distributed P1-8 semantics are out of scope.

Operator documentation: docs/operators/rmsnorm-residual.md.

Validation environment

Item	Value
GPU	NVIDIA H100 PCIe, compute capability 9.0, 81,559 MiB, MIG disabled
Driver / CUDA toolkit	580.159.04 / 12.8 (nvcc V12.8.93)
PyTorch / Triton	2.8.0+cu128 / 3.4.0
Python / Host compiler	3.11.16 / GCC 13.3.0

Correctness / Tests

Coverage: T={1,7,16}, D={128,4096}; byte-equal forward/backward and saved state, edge values, strided inputs, repeatability, batch/padding invariance, both autograd branches, invalid inputs, and FP32 gamma rounding regressions.

uv run --no-sync python -m pytest \
  tests/test_rmsnorm_residual.py tests/test_rmsnorm_residual_triton.py -q -rs
# 28 passed in 7.89s

uv run --no-sync python -m pytest tests/test_rmsnorm_residual_gamma_precision.py -q
# 4 passed in 2.63s

uv run --no-sync python -m pytest \
  tests/test_p1_oracle.py tests/test_p1_provider.py \
  tests/test_operator_inputs.py tests/test_kernel_registry.py -q
# 77 passed in 2.16s, including the golden-manifest anchor

gtest

uv run --no-sync python scripts/check_operator.py \
  --op rmsnorm_residual --candidate cuda --device cuda \
  --dtype bf16 --batch 1 --seq 16 --normalized-dim 4096 --check-grad

uv run --no-sync python scripts/check_operator.py \
  --op rmsnorm_residual --candidate triton --device cuda \
  --dtype bf16 --batch 1 --seq 16 --normalized-dim 4096 --check-grad

Both report passed=True pass_rate=1.0000. Maximum absolute and relative errors are 0 for both outputs, BF16 x.grad, and FP32 gamma.grad.

P1 acceptance

uv run --no-sync python scripts/check_p1.py \
  --provider rl_engine.mhc.cuda_provider:CudaMHCProvider --device cuda

uv run --no-sync python scripts/check_p1.py \
  --provider rl_engine.mhc.triton_provider:TritonMHCProvider --device cuda

Both report RESULT: PASS (all boundaries byte-equal) for all five cases: one_row, packed_t16, packed_t7_odd, fused_pre_norm, and mixer_frozen.

Benchmarks

uv run --no-sync python benchmarks/benchmark_rmsnorm_residual.py \
  --tokens 32768 131072 262144 --warmup 50 --iterations 1000 \
  --output benchmark.json

H100 PCIe, D=4096, BF16 activations, FP32 gamma/explicit gradients; CUDA-event mean, eager execution, no CUDA Graphs. The script runs all three backends; both candidates pass pre-timing bitwise checks. CUDA versus torch-native forward + backward results:

T	torch-native ms	CUDA ms	Speedup	torch-native extra MiB	CUDA extra MiB
32768	12.908420	10.043916	1.29×	3584.3906	1024.1406
131072	51.176590	39.021008	1.31×	14337.5156	4096.5156
262144	102.286531	80.839070	1.27×	28675.0156	8193.0156

CUDA forward speedup is 1.62–1.63×; combined peak extra allocated memory is approximately 71.4% lower. Memory is the allocator peak above the live baseline, not total GPU usage. Torch-native uses vectorized reductions and is a performance baseline, not the fixed-order oracle. Results apply to these measured workloads.

Kernel PR checklist

Follow PR #204.

• CUDA and/or Triton kernel
• Binding + build integration
• Python/autograd wrapper + MHCProvider
• Hardware-gated registry entry with clean fallback
• check_p1.py passing result
• Forward/backward test sweep
• Bitwise batch-invariance test
• Benchmark: latency + peak extra memory
• Validation environment table
• Operator documentation

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: b170b67d-51bd-4352-99a2-a382d9c7c0e4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}

[zhangj1an]

LGTM. Thank you!

This PR implements rmsnorm_residual using CUDA and Triton backends.

Forward: y = (x · r) · gamma, whereby r = rsqrt(mean(x²) + eps).

Backward: since y = x · r · gamma, need to calculate dx and dgamma.

for dx: dx = r · (dy · gamma) − x · r³ · Σ_d(dy · gamma · x) / D + d_residual. (i asked against gpt to understand this formula derivation).

for dgamma: Σ_t (dy · x · r), whreby t means tokens, as every token uses the same gamma.

Both backends are byte-equal to the oracle (fixed-order reductions, mul/add separated, rsqrt not 1/sqrt).

We agreed via dm to use FP32 gamma. So the oracle was updated accordingly.