[bjf-frz]

Required human contributor answers

The author's recorded answers in #160 describe the reusable CUDA-operator work that this PR continues:

1. What outcome are you trying to achieve?

Reuse the outcome in #160, question 1: a complete path for developing and integrating further native CUDA operators. This change applies that path to Ouro RMSNorm and its following residual addition.

2. What is the highest-leverage variable for that outcome?

Reuse #160, question 2: activation memory movement, kernel launches, and parallel computation while preserving the numerical contract.

3. How will this work move that variable rather than create busywork?

Reuse the fusion, aligned vector body, safe scalar boundaries, paired arithmetic, and matched numerical/performance checks described in #160, question 3. Here those techniques combine RMSNorm and its existing subsequent residual add. The Ouro integration uses a separately validated PyTorch-order reduction and falls back where that reduction has not been validated.

4. Who is the customer or user, and what do they actually need?

Reuse #160, question 4: the author developing CUDA operators for vllm-rlt needs an implementation, build, model-integration, and validation path for subsequent operators.

PR necessity and bugfix evidence

The author explicitly requested this RMSNorm PR. The existing implementation is a sequence of PyTorch operations; using a native kernel and fusing the subsequent add requires repository changes. This is an optimization, not a claimed pre-existing correctness bug.

Purpose

Add a native inference RMSNorm operator with optional residual addition and use it in Ouro. The operation preserves Ouro's order:

normalized = (input.float() * rsqrt(mean(input.float() ** 2) + epsilon)).to(input.dtype)
output = residual + weight * normalized  # residual is optional

• The weight is one-dimensional. The CUDA entry supports FP32/FP16/BF16, ranks 1–4, stride-1 last dimensions, independent tensor offsets, strided outer dimensions, empty inputs and the current CUDA stream. The output owns new contiguous storage. Unsupported layouts, mixed dtypes, autograd and Python-only installations use the Torch reference.
• The general kernel uses FP32 reduction, an aligned vector body with safe boundaries, register caching for selected widths and paired output arithmetic. Its reduction order can differ from Torch and is not the default Ouro execution path.
• Ouro's separately validated path matches the tested Torch reduction order for 2D width-2048 FP16/BF16 rows, with 8-byte-aligned input/weight/residual rows, on PyTorch 2.13.0 and SM103. Other versions, GPUs, shapes or alignments use the original Torch expression. This conservative gate matters because small reduction differences can change adaptive exits.
• The four layer norms and the norm after the layer loop opt in for Ouro. Other models keep their existing path. This does not depend on the SiLU alignment changes in open [Kernels] Vectorize misaligned SiLU inputs with a per-SM cutoff #163.

Test Plan

vllm-rlt Version / Commit: base 989921ac1d3897c5c43e27ec8cddbc91fa953194; runtime implementation 4f3ce8e1921af2224a98378a4dadb2be373b4bbb; final head 59f964cc741d38643fa22edaafb931da67291f8e adds documentation and the timeline only. Measurements and execution tests ran on the identical runtime source at 4f3ce8e.

• Python 3.12, PyTorch 2.13.0+cu130, CUDA toolkit 13.4, driver 610.43.02; single NVIDIA B300 SXM6 AC (SM103), physical GPU 7. The author authorized use of an idle GPU; GPU 7 was checked for other processes. Other GPUs had unrelated resident processes on the shared host. SM80 was cross-compiled but not run.
• pre-commit run --all-files: Ruff lint, format and the complete CPU regression suite passed. A separate Python-only import/fallback and CPU autograd check also passed.
• pytest tests/test_normalization.py tests/test_ouro.py tests/test_activation.py tests/test_native_build.py --run-gpu -q: 348 passed after rebuilding the extension for SM80 and SM103. Coverage includes independent input/weight/residual offsets, outer strides and singleton strides, odd/empty widths, special values, staged rounding, current streams, Graph replay, fallback and existing Ouro/SiLU behavior.
• Compute Sanitizer memcheck: 29 passed / 0 errors for stride and alignment cases; 7 passed / 0 errors for current-stream and Graph replay cases. The graph check used the normal CUDA allocator.
• A separate BF16 Graph model check compared actual greedy outputs and exit-depth sequences: all 8 fixed-four-loop natural requests (256 tokens) and all 16 adaptive-exit requests (512 tokens, threshold 0.7, depths 3 and 4 observed) matched. This is a small output-consistency check, not a full GSM8K quality result.
• The new RMSNorm microbenchmark was run with and without the optional source-built vLLM library. The latter is an optional performance reference, with its different intermediate rounding reported separately.

Test Result

Unprofiled, same-checkout three-way A/B after shape-specific warmup, three interleaved repetitions. Each backend has its own model worker resident on the same GPU, with only one submitting work at a time. Ouro-1.4B, BF16, Triton attention, fixed four loops, greedy, ignore EOS, KV 1536 × 16, max sequences 8, max batched tokens 1024, prefill chunk 256. Prompt IDs are constructed by repeating/truncating the first GSM8K questions to the indicated input lengths. Output throughput is successful output tokens divided by the local request window including prefill, scheduling and IPC; it excludes network time. All token IDs, exit depths, finish states and request/KV cleanup matched.

Concurrency / input / output / requests	Torch token/s	CUDA norm + separate add	CUDA norm + fused add
1 / 128 / 64 / 4	69.06	91.55	94.35
8 / 256 / 64 / 16	369.10	470.23	481.56

The columns are medians. The full replacement is +36.6% and +30.5%, respectively; the final residual fusion relative to the middle column is +3.1% and +2.4%. Per-run ranges were 68.69–69.11 / 90.79–91.89 / 94.15–94.92 token/s for concurrency 1, and 368.40–369.51 / 442.46–471.17 / 480.73–482.18 for concurrency 8. The latter separate-add run shows scheduling variation. For concurrency 1 the medians of per-run TPOT p50 were 14.151 / 10.650 / 10.331 ms; TTFT p50 was 35.334 / 26.932 / 24.732 ms. For concurrency 8 TPOT p50 was 19.837 / 15.398 / 15.006 ms and TTFT p50 was 96.387 / 87.039 / 88.473 ms. These are short closed-loop bursts, not saturation, tail-SLO, or task-quality measurements.

All three single-request workers captured one recurrent graph, replayed it 4032 times and had zero recurrent graph fallbacks; the concurrency-8 workers captured six, replayed 6060 times and also had zero fallbacks. Only the compatibility path is included in model speed claims; generic kernel timings are not substituted for it.

Source attribution and matched profiles

One fixed-depth decode batch invokes (24 layers × 4 RMSNorms + 1 final norm) × 4 loops = 388 norms. In this environment, the prior Torch expressions use 8 GPU kernels for a plain norm and 9 with residual, versus one each in the compatible implementation. For a warmed 128-input/2-output Graph request, matching separate ops-only and full CPU+CUDA profiles observed 11004 → 5188 total GPU kernels, a difference of 5816 across prefill plus one decode. Four recurrent graph launches were present in each decode; the first replay contained 1386 → 659 GPU kernels. The candidate trace contains 776 rms_norm_torch_order_kernel launches across the two forwards. Both profile variants completed with the same two output IDs and depths.

The full and zoomed panels use matching time scales; the first replay's profiled GPU span is 3.390 → 2.319 ms. Profiler timings are diagnostic, not the performance source. A separate unprofiled A/B with the same 128-input/2-output shape measured 41.61 → 57.03 output token/s, with 14.278 → 10.406 ms TPOT p50 medians; the four-request, two-output-token window is especially short. Raw ops-only/full archives, their shape/memory/stack artifacts, source hashes, request IDs and per-run results are retained locally under ignored artifacts/rmsnorm-submission/; they are not yet uploaded for external review. Peak memory and other GPU architectures have not been qualified. The PR remains draft while externally accessible raw traces are pending.

Reproduction

Use an allocated GPU and a prepared CUDA-enabled PyTorch/toolkit environment:

python -m pip install --no-build-isolation -e .
pre-commit run --all-files
CUDA_VISIBLE_DEVICES=<allocated_gpu> python -m pytest \
  tests/test_normalization.py tests/test_ouro.py tests/test_activation.py \
  tests/test_native_build.py --run-gpu -q
CUDA_VISIBLE_DEVICES=<allocated_gpu> python -m benchmarks.kernels.rms_norm \
  --tokens 1,8,16,128 --widths 2048 --dtypes bfloat16 \
  --output artifacts/rmsnorm/results.json --profile-dir artifacts/rmsnorm/profiles

The serving guide contains engine startup, a single request and batch client examples; those HTTP timing boundaries differ from the local worker results above. The exact same-checkout A/B, natural/adaptive, and profile drivers and results are retained in the local artifact directory above.

AI assistance was used for implementation, measurements and author self-review.