[MatsueYu]

refactoring add_rmsnorm_bias, rmsnorm_without_weight, rmsnorm_split for new kernel sturctures
related to Refactor: Ops Architecture - #869

Test

Test design

1. add_rmsnorm_bias — 120 cases

3 dtypes × 5 shapes × 4 modes × 2 eps

Dimension	Values	Purpose
dtype	fp32, fp16, bf16	Algorithm check, plus low-precision reduction/cast paths
eps	1e-6, 1e-5	Stability term; zero-variance finiteness
shape	(1, 4096)	Qwen3-8B single-token decode (model width borrowed)
	(16, 7168)	DeepSeek-V3 16-token decode
	(3, 6144)	Legacy width; multiple of the 2048 quant tile
	(1, 1)	Minimum non-empty reduction
	(7, 2049)	Odd rows + one-column tail of the 2048 quant tile
mode	normal	Residual add + affine RMSNorm main path
	zero	eps handling; output should equal the bias
	offset_only	Offset without scale must not enable quantization
	quantized	int8 truncation, both saturation endpoints, masked tail

quantized uses per-column x/residual = ±0.25/±0.75 (sum ±1), weight=1, bias=0, scale ∈ {0.5, 2.0, 200.0, 200.0}, offset=0.25, so the normal, truncation and +127/-128 saturation paths are all deterministic.

2. rmsnorm_without_weight — 90 cases

3 dtypes × 5 shapes × 3 modes × 2 eps

Dimension	Values	Purpose
shape	(1, 130, 2048)	LTX-2 audio; previous test shape
	(2, 16, 2048)	Audio width, multiple tokens
	(2, 1, 4096)	Video width
	(1, 1, 2048)	Minimum token count
	(1, 65, 2048)	One token past block_l=64
mode	normal	General random reduction
	zero	Zero input must produce zero and stay finite
	constant	Uniform rows verify the normalization scale

All C values are powers of two because the kernel reduces with tl.arange(0, C).

3. fused_variance + fused_rsqrt_mul — 75 cases

Test	Matrix	Cases
fused_variance	1 dtype (fp32) × 5 shapes × 3 modes	15
fused_rsqrt_mul	1 dtype (fp32) × 5 shapes × 2 variance layouts × 3 modes × 2 eps	60

Dimension	Values	Purpose
shape (variance)	(1,1024,512), (1,8190,512), (2,4096,2048), (1,1,512), (1,65,1024)	Legacy shapes, TP shard widths, minimum token, block tail
shape (rsqrt_mul)	(1,8190,2560), (1,1024,512), (2,4096,640), (1,1,128), (3,17,100)	Production-like widths, TP-sharded width, unaligned width C=100
variance layout	[B,L,1], [B*L]	Real caller layout and legacy flat layout
mode	normal, zero, constant	Random positive variance, eps handling, uniform scale
dtype	fp32 only	sglang casts x and sharded weight to fp32 before this path

No multi-rank all_reduce is exercised: the operators are single-rank phases and the collective stays in sglang.

Accuracy and performance policy

Item	Setting
Tested / reference impl	auto (= Triton) vs torch_native
fp32 tolerance	atol=rtol=1e-5
fp16 tolerance	atol=rtol=1e-3
bf16 tolerance	atol=rtol=5e-2
int8 quantized output	exact match (atol=rtol=0)
Accuracy failure	Stops before performance collection
Performance cases	2 per operator, selected by perf_shape + dtype + mode

Accuracy Test

add_rmsnorm_bias

python -m pytest tests/python/operators/norm/test_add_rmsnorm_bias.py --impl=triton --impl-ref=torch_native --skip-acc-check=false -v -s -x -rs

rmsnorm_without_weight

python -m pytest tests/python/operators/norm/test_rmsnorm_without_weight.py \
  --impl=triton --impl-ref=torch_native --skip-acc-check=false -v -s -x -rs

rmsnorm_split

Performance Test

python -m pytest tests/python/operators/norm/test_add_rmsnorm_bias.py \
  --enable-perf=true --prof-dir=results/profilings \
  --csv-path=results/perf_add_rmsnorm_bias.csv -v -s

python -m pytest tests/python/operators/norm/test_rmsnorm_without_weight.py   --enable-perf=true --prof-dir=results/profilings   --csv-path=results/perf_rmsnorm_without_weight.csv -v -s

python -m pytest tests/python/operators/norm/test_rmsnorm_split.py \ --enable-perf=true --prof-dir=results/profilings \ --csv-path=results/perf_rmsnorm_split.csv -v -s

[Todobe]

Optimize the description, delete the chat content.

[MatsueYu]

Optimize the description, delete the chat content.

[MatsueYu]

Optimize the description, delete the chat content.

modified