[BugFix] Support rms_norm for stacked 2D weights on XPU

[ccrhx4]

Fix #573

XPU's rms_norm kernel does not support stacked 2D weights. When callers pass a stacked 2D weight of shape [num_rows, hidden_size] (one row per layer) to normalize multiple layers in a single kernel launch, every row silently read weight row 0 instead of its own row, producing wrong results with no error.

How it is fixed:

• In the rms_norm() entry point, derive a weight_stride local variable: 0. The public rms_norm op schema is unchanged.
• Pass weight_stride through call_rms_norm_kernel into all three kernel paths that read weight.
• Each kernel reuses its existing batch_idx (already computed for input/out offsetting) to compute weight_row = weight + batch_idx * weight_stride.
• gemma_rms_norm's call site is updated to pass weight_stride=0 unchanged, since no vLLM caller stacks weights through it today.
• Adds tests/test_batched_weight_rms_norm.py covering batched-weight correctness against a per-row reference loop (5 shapes x 2 dtypes, covering the multi-row, vectorized, and scalar-fallback kernel paths), a direct regression test for the issue's zero-weight-row repro, and shape-validation error tests.

Test Result

• New tests/test_batched_weight_rms_norm.py (12 parametrized cases): batched-weight correctness vs. a per-row reference loop across 5 shapes x 2 dtypes covering all three kernel paths (multi-row, generic vectorized, scalar fallback for non-power-of-two hidden size), a direct regression test reproducing the issue's zero-weight-row repro, and a shape-validation test for the new TORCH_CHECK errors. All pass.
• Full existing tests/test_layernorm.py (1152 tests) plus tests/test_fused_qk_norm_rope.py, tests/test_fused_input_norm.py, and tests/test_fused_norm_quant.py (808 tests) -- zero regressions on the standard 1D-weight path.
• Cross-validated against vLLM's own tests/kernels/core/test_batched_weight_rms_norm.py via a locally patched (not committed) copy enabling XPU -- 13/13 pass, including float32.
• Verified the issue's exact repro script directly: previously out[1] read layer 0's weight (wrong, e.g. 3.79); now reads its own (zeroed) weight row, giving out[1] == 0.0 as expected.
• Benchmarked with benchmark/benchmark_rmsnorm.py (standard 1D-weight path, both with and without residual): correctness still matches HuggingFace-naive, and vLLM's kernel remains ~9-13x faster than HuggingFace-naive and ~1.2-1.6x faster than torch.compile across the full head_num/batch_size/seq_len sweep -- confirming no performance regression on the common path.

[Copilot]

🟡 Changes recommended

There are correctness issues to address (nullptr pointer arithmetic UB in the no-weight kernel path, missing weight dtype validation, and a likely-flaky exact-equality test).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR fixes incorrect behavior in the XPU rms_norm custom op when callers pass a stacked 2D weight ([num_rows, hidden_size]): previously all rows read weight row 0, silently producing wrong results. The change threads a weight_stride through the RMSNorm kernel implementations so the kernel can select the correct per-outer-row weight, and adds dedicated tests to prevent regression.

Changes:

• Add 2D (stacked) weight support to XPU rms_norm by passing weight_stride through the RMSNorm kernel launch and selecting weight_row based on the outer index.
• Add shape validation for 1D vs 2D weights in the rms_norm entry point.
• Add a new test module covering correctness and regression reproduction for batched-weight RMSNorm.

File summaries

File	Description
csrc/layernorm.cpp	Implements per-outer-row weight selection for 2D weights via weight_stride and adds weight shape validation.
tests/test_batched_weight_rms_norm.py	Adds regression + correctness tests for stacked 2D weight behavior in torch.ops._C.rms_norm.

Review details

Suppressed comments (2)

csrc/layernorm.cpp:215

• Same nullptr pointer-arithmetic issue as above: when HasWeight is false, weight can be nullptr, so computing weight_row = weight + ... is undefined behavior. Compute the offset only when HasWeight is true.

                  seq_idx * input_stride_d3 + head_idx * input_stride_d2;
    }
    const scalar_t* weight_row = weight + batch_idx * weight_stride;

csrc/layernorm.cpp:348

• Same nullptr pointer-arithmetic issue in the multi-row kernel: when HasWeight is false, weight may be nullptr, so weight + batch_idx * weight_stride is undefined behavior even if the pointer is never dereferenced. Guard the offset computation with if constexpr (HasWeight).

                  seq_idx * input_stride_d3 + head_idx * input_stride_d2;
    }
    const scalar_t* weight_row = weight + batch_idx * weight_stride;

• Files reviewed: 2/2 changed files
• Comments generated: 3
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jikunshang]

thanks for fixing.
cc @zufangzhu

[ccrhx4]

I have addressed the copilot review comments. Please help to review. @zufangzhu Thanks a lot.