[XPU] Norm Inkling's q in place on the packed QKVR row via sgl-kernel-xpu

[jmunetong]

Motivation

This reworks #37323 so the q-norm optimization is XPU-only: the kernel-side trick moves to sgl-kernel-xpu (sgl-project/sgl-kernel-xpu#529), and this PR only calls it on XPU. Two fixes that stop Inkling on XPU at default settings are also carried over.

• In-place q norm. Inkling's q is a strided slice of the packed [q | k | v | r] projection row. apply_qk_norm reshapes it to (T*H, D), which forces a compaction copy and an extra launch. On XPU, q is now normed in place on the row with sgl_kernel.rmsnorm_heads_inplace.
• inkling_common.attn doesn't import on XPU. is_batch_invariant's module eagerly imports cutlass, which is CUDA-only (absent from pyproject_xpu.toml), so import fails with No module named 'cutlass'.
• _get_inkling_ar_resources raises on XPU. It calls torch.cuda.is_current_stream_capturing() before checking is_cuda(), and that raises Tried to instantiate dummy base class _cuda_isCurrentStreamCapturing.

Modifications

3 files, +240/−6.

• inkling_common/attn.py:
  • rmsnorm_heads_inplace / rmsnorm_heads_inplace_supported are imported only under is_xpu(), inside try/except ImportError. Against a wheel without them, the fast path is simply off.
  • __init__ decides once whether the row geometry is admitted, through the kernel repo's _supported helper (TP=16/32 are refused since their tail isn't whole heads).
  • The forward gets one XPU-only branch: elif self._xpu_q_norm_ok and self._xpu_norm_q_inplace(q). It recovers the packed row from q's own storage and declines (falling through to apply_qk_norm) if q isn't the leading slice of a row of the expected width, or isn't bf16/fp16. _project_qkvr and every CUDA/ROCm line are unchanged.
  • The is_batch_invariant import falls back to False on ImportError. _batch_invariant is only ever set by the fa4 path, which needs cutlass itself, so False is what the real module returns wherever this fallback can be reached. On CUDA the import succeeds as before.
• inkling_common/kernels/comm.py: check is_cuda() before torch.cuda.is_current_stream_capturing(). On CUDA the first check passes, so behavior is unchanged; on ROCm both orders already returned None.

Compared to #37323, this drops the _project_qkvr signature change, the inline pad arithmetic (now rmsnorm_heads_inplace_supported), the position-id cache (now in the op), and the attention-backend assert-message edit.

Accuracy Tests

test/registered/xpu/test_inkling_xpu_q_norm.py, registered to stage-b-test-1-gpu-xpu, 5 tests:

case	guards	runs on today's pinned wheel?
attn imports, is_batch_invariant() is False	the cutlass import failure	yes
AR resources return None	the is_current_stream_capturing raise	yes
Inkling TP=8/TP=4 geometries admitted	xpu_qkvr_width and the init-time gate	skips until the op ships
q matches F.rms_norm; k, v, r bit-identical	row recovery from q in _xpu_norm_q_inplace, bf16/fp16, T ∈ {1, 8, 512}	skips until the op ships
compacted q and fp32 q declined, row untouched	the fall-through branch	skips until the op ships

Both regression cases fail on current main with the exact errors above, and pass here. With the op injected from sgl-project/sgl-kernel-xpu#529, all 5 pass. The op's own contract (rotary off at nonzero positions, the tail never written, refusals) is tested in sgl-kernel-xpu. Intel Arc Pro B60, torch 2.15.0.dev20260912+xpu.

Not yet verified: a full Inkling server run on this branch. #37323's end-to-end A/B used the same in-place call.

Speed Tests and Profiling

The native call is the same as #37323's, so these are that PR's numbers, not re-measured here. q+k norm device time (profiler self-time, min of 5 interleaved reps), Arc Pro B60, torch 2.12.0+xpu, Inkling TP=8 per-rank rows:

base / this	T=1	T=8	T=64	T=512	T=4096	T=8192
SWA row 1664	1.18×	1.62×	1.50×	1.82×	3.28×	1.82×
full row 1408	1.19×	1.60×	1.50×	1.88×	3.08×	2.50×

Kernel count per layer goes from 4 to 3. End to end, #37323's tp=4 A/B on a 6-layer reduced checkpoint found the win real but immaterial at that geometry: 0.010% of a 682 ms TTFT, well inside run-to-run noise.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. N/A, no user-facing change.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #37543020377
Latest PR Test (Extra): ❌ Run #37543020094
Latest PR Test (AMD ROCm 10): ❌ Run #37543020471