[XPU] Dispatch nn.LayerNorm to fused SYCL kernel via CustomOp

[mganczarenko]

Purpose

vllm-xpu-kernels exposes fused SYCL kernels for standard (non-RMS) LayerNorm (vllm-xpu-kernels#577) - layer_norm, fused_add_layer_norm, nemotron_layer_norm, fused_add_nemotron_layer_norm - but nothing in vLLM dispatched to them yet. Every model with a bare
nn.LayerNorm fell back to torch-xpu-ops' generic kernel on XPU.

• Adds StandardLayerNorm, a CustomOp (registered "layer_norm") - drop-in for a bare nn.LayerNorm, dispatching via forward_native/forward_cuda/forward_xpu, matching RMSNorm/GemmaRMSNorm in the same file.
• Reworks NemotronLayerNorm1P (Nemotron's "+1 to weights" variant) from an nn.LayerNorm subclass into a CustomOp sibling ("nemotron_layer_norm"); its XPU path also dispatches the fused-residual-add kernel when a residual is present.
• Wires both into GPT-2, GPT-J, OPT, SigLIP (PaliGemma's vision tower), and Nemotron.
• Adds replace_layernorm_class (transformers/utils.py), matching replace_linear_class/replace_conv_class/replace_embedding_class, so the Transformers modeling backend's generic recursive module-replace picks this up too - covers StarCoder2-style models with no dedicated vLLM model file.
• Kernel calls are wrapped as real torch.ops.vllm.xpu_* custom ops (_xpu_ops.py, via direct_register_custom_op), the same way GemmaRMSNorm wraps its kernel - compile-safe by construction, gated only by the standard CustomOp +/-custom_ops toggle.
• Every dispatch path is hasattr(torch.ops._C, ...)-guarded and falls back to native F.layer_norm when the op isn't present - currently a no-op everywhere, since no vllm-xpu-kernels release yet contains Incorrect output when using hf-internal-testing/llama-tokenizer #577 (just merged) and requirements/xpu.txt's pin predates it.

Test Plan

• ruff check, ruff format, full pre-commit (incl. mypy) on every touched file.
• Kernel correctness on Intel B70 hardware: parametrized cases across hidden sizes (incl. non-power-of-2), bf16/fp16, weight/bias/weight_bias combinations, plus fused-add cases, for both layer_norm and nemotron_layer_norm/fused_add_nemotron_layer_norm, checked against a plain-PyTorch fp32 reference.
• End-to-end serve (--enforce-eager) on the same hardware, per affected model.
• End-to-end performance via the official vllm bench serve/throughput/latency, baseline vs patched, N=3 median.
• Not yet added: a correctness test in the repo (tests/kernels/core/test_layernorm.py covers RMSNorm/GemmaRMSNorm, nothing exercised plain LayerNorm before this PR). Not yet done: lm_eval accuracy check for any affected model. A reproducible kernel-level benchmark (benchmark/benchmark_layer_norm.py) is already
  committed on the vllm-xpu-kernels PR itself.

Test Result

Correctness: 64 parametrized cases + 2 fused-add cases, all pass vs plain-PyTorch fp32 reference.

Kernel microbenchmark (native F.layer_norm vs SYCL, bf16, hidden_size=4096):

tokens	native	SYCL	speedup
1	8.35us	3.60us	2.32x
128	9.63us	3.70us	2.60x
512	28.20us	5.90us	4.78x
1024	51.73us	10.83us	4.78x (peak)
2048	103.91us	56.61us	1.84x

End-to-end serve: GPT-J, GPT-2, OPT-125m, Nemotron/Minitron-8B, PaliGemma (SigLIP), and StarCoder2-3B all served and generated correctly with --enforce-eager; SYCL dispatch confirmed firing inside the real engine, not just via a direct standalone op call.

End-to-end performance:

Model	Path	Throughput	Latency	Serve
GPT-2-124M	direct swap	+9.1%	-9.5% (faster)	+5.3%
OPT-125M	direct swap	+9.2%	-1.2% (faster)	+9.4%
GPT-J-405M	direct swap	+6.1%	-4.4% (faster)	+9.2%
StarCoder2-3B	generic swap	+1.0%	-1.2% (faster)	+1.3%
Minitron-8B (Nemotron)	direct swap	+0.8%	-0.9% (faster)	+0.5%

Caveat: the above was measured against an earlier revision of this dispatch code (direct torch.ops._C.* calls, gated by torch.compiler.is_compiling()). This PR wraps those same calls through torch.ops.vllm.xpu_* instead, for torch.compile safety - a thin pass-through to the identical kernel call, so behavior should
be unchanged, but hasn't itself been re-run on hardware.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hmellor]

Why not just call it LayerNorm?

Suggested change

	class StandardLayerNorm(CustomOp):
	class LayerNorm(CustomOp):

[mganczarenko]

There is already LayerNorm class being nn.Module at the end of the file, but I am open for alternatives for the name

[hmellor]

Could we port the other LayerNorm to be a custom op instead of adding a new one with a different name?

[mganczarenko]

I think we could. Please, take a look and let me know if it looks good to you - ff208db

[mganczarenko]

CC: @jikunshang, @mayuyuace

[hmellor]

Given that this is directly swapping nn.LayerNorm, why not put the fused kernel in torch-xpu-ops?

[mganczarenko]

Given that this is directly swapping nn.LayerNorm, why not put the fused kernel in torch-xpu-ops?

I've tried to match the precedents already set by RMSNorm/GemmaRMSNorm in the same file (both dispatch fused kernels living in vllm-xpu-kernels), so I've chosen vllm-xpu-kernels at first. This is the place where vLLM's XPU-specific fast paths for attention/MoE/quant/norm ops already live.

Nothing rules out upstreaming this to torch-xpu-ops, but this is another process.

[hmellor]

IIRC the reason that the RMSNorm/GemmaRNSNorm classes exist was so that they could be picked up by the custom op machinery in vLLM and used in vLLM's fusion passes.

Since there are no users of this new layer there are no fusion passes for it to be captured by.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mganczarenko.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mganczarenko]

IIRC the reason that the RMSNorm/GemmaRNSNorm classes exist was so that they could be picked up by the custom op machinery in vLLM and used in vLLM's fusion passes.

Since there are no users of this new layer there are no fusion passes for it to be captured by.

Thanks for the clarification. Following up on the question about why this couldn't just be a kernel contribution to torch-xpu-ops instead of a vLLM-specific CustomOp + vllm-xpu-kernels kernel.

Why this kernel is faster than native?

torch-xpu-ops already has a LayerNorm kernel for XPU (at::native::xpu::VectorizedLayerNormKernelFunctor) - the pre-patch path in this PR already uses it, via aten::native_layer_norm. So this isn't about filling a gap in torch-xpu-ops. I profiled both paths end-to-end (GPT-J-405M and OPT-125m, B60, torch profiler traces down to the individual event level) to see where the improvement comes from, and it splits into two pieces:

1. The new SYCL kernel itself is modestly faster than torch-xpu-ops's native kernel at this scale, but small in
   absolute terms for these model sizes.
2. The bigger effect comes from dispatch, not the kernel. aten::native_layer_norm is a full ATen "native" op - it's registered with autograd and functionalization support, because torch-xpu-ops has to serve every PyTorch consumer, including training. Every call walks that dispatcher-key machinery before reaching the backend kernel, and none of it shows up as a named event in a profile (it's just unlabeled C++ time between the last visible op and the driver call) - but it's real, unavoidable cost tied to the op's schema, not to which kernel runs underneath it. vllm.model_executor.layers.layernorm.StandardLayerNorm's CustomOp registers xpu_layer_norm/_C::layer_norm as an inference-only op with no autograd/functionalization layering -there's structurally less dispatcher work to do, and it's Python-visible (torch.ops.* frames) rather than a black-box gap.

Why it does not meet requirements to be native?

That second part is the actual reason this lives outside torch-xpu-ops: it's not a limitation of torch-xpu-ops's kernel quality, it's that any backend registered for aten::native_layer_norm inherits the full differentiable-op dispatch contract, which vLLM's inference-only serving path never needs. Bypassing it requires an op outside that contract - which is what CustomOp + vllm-xpu-kernels gives us, and what a general-purpose PyTorch backend project can't do without breaking autograd for everyone else using aten::native_layer_norm. Same pattern already used here for GemmaRMSNorm/RMSNorm, so this isn't a new precedent, just extending an existing one to LayerNorm.

[hmellor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92117 for commit e9f2f1c3702d.