[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu by jmunetong · Pull Request #32227 · sgl-project/sglang

jmunetong · 2026-07-23T16:55:51Z

Motivation

Bringing up NVIDIA-Nemotron-3-Nano-30B-A3B (a hybrid Mamba-2 / attention model,
NemotronHForCausalLM) on --device xpu fails during server init with two separate,
sequential errors. Both are device/arch-coverage gaps rather than kernel bugs.

Error 1 — `extra_buffer` mamba radix cache asserts out on XPU

With the default --mamba-radix-cache-strategy auto, the resolver selects
extra_buffer for supported hybrid archs (Nemotron qualifies). A device assertion
then rejects XPU:

File ".../server_args.py", line 5251, in _validate_mamba_extra_buffer
    is_cuda() or is_musa() or is_npu() or is_hip()
AssertionError: extra_buffer needs CUDA/MUSA/NPU/ROCm (FLA).

The extra_buffer path is FLA/Triton-backed and Triton has an XPU backend, but
is_xpu() was simply missing from the guard, so the strategy could never be used
on XPU even where it works.

Error 2 — Triton backend indexes layer 0 as full-attention for a hybrid model

After clearing Error 1, init crashes selecting v_head_dim:

File ".../layers/attention/triton_backend.py", line 198, in __init__
    self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
File ".../mem_cache/memory_pool.py", line 3649, in _transfer_full_attention_id
    raise ValueError(
ValueError: layer_id=0 not in full attention layers: dict_keys([5, 12, 19, 26, 33, 42])

TritonAttnBackend.__init__ has a hybrid-model branch that avoids indexing layer 0
(which may be a Mamba layer, not full-attention), but it only tested
hybrid_gdn_config / kimi_linear_config / linear_attn_model_spec. NemotronH is
detected by mamba2_config, which is not in that set, so it fell through to the
else branch that assumes layer 0 is full-attention. Nemotron's full-attention
layers are [5, 12, 19, 26, 33, 42], so get_value_buffer(0) raises. This affects
other mamba2_config archs too (FalconH1, Lfm2, ...), not just XPU.

Modifications

server_args.py — add is_xpu() to the extra_buffer device guard (and to
the assertion message).
triton_backend.py — replace the three-way hybrid check in the v_head_dim
selection with the umbrella mambaish_config(), which unions mamba2_config
(NemotronH/FalconH1/Lfm2/...) + hybrid_gdn_config + kimi_linear_config +
linear-attn specs, so every hybrid-linear model takes the get_v_head_dim() path
instead of indexing layer 0. Unused imports removed.

Fix 2 is device-agnostic (fixes the crash on any backend for these archs); fix 1 is
XPU-specific.

Why the fix is in the Triton backend, not the `intel_xpu` backend

On XPU this model runs its full-attention layers through TritonAttnBackend, not the
intel_xpu backend. That is by design of the hybrid path, not an oversight:
NemotronHForCausalLM is a hybrid Mamba-2 / attention model, so it is wrapped in
HybridLinearAttnBackend, whose full-attention sub-backend is selected by device as
trtllm_mha (CUDA SM100) / fa3 (CUDA SM≥80) / triton (everything else,
including XPU). The intel_xpu attention backend is not currently wired into the
hybrid-linear-attention path, so XPU falls to the Triton full-attn sub-backend (Triton
has an XPU backend, so it executes on-device). Hence the crash — and the fix — live in
TritonAttnBackend. Integrating intel_xpu as a hybrid full-attn sub-backend is a
separate, larger enablement effort and out of scope here.

Accuracy / correctness

Draft: fix 1 removes a launch guard for a path whose FLA extra-buffer kernels are
well-exercised on CUDA/ROCm/NPU but new to XPU — reviewers should weigh whether XPU
needs the same validation bar as the other gated backends before this lands.
Fix 2 changes only which detector gates an existing code path; behavior for the
already-covered archs is unchanged.

Checklist / how to reproduce & validate

Repro (before this PR, either error above depending on stage):

SGLANG_USE_SGL_XPU=1 python3 -m sglang.launch_server \
  --model-path nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
  --device xpu --tp 4 --kv-cache-dtype fp8_e4m3

With this PR the server initializes past both errors on XPU.

CI States

Latest PR Test (Base): ❌ Run #31546362673
Latest PR Test (Extra): ❌ Run #31546362268

gemini-code-assist · 2026-07-23T16:55:54Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

gemini-code-assist · 2026-07-23T17:11:12Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

airMeng · 2026-07-23T17:26:08Z

-            or kimi_linear_config(model_runner.model_config) is not None
-            or linear_attn_model_spec(model_runner.model_config) is not None
-        ):
+        elif mambaish_config(model_runner.model_config) is not None:


will it break other models?

should be OK:

defined here:

https://github.com/sgl-project/sglang/blob/f00a13ae48e16c796bc317fb11f4478441e0e2a3/python/sglang/srt/configs/hybrid_arch.py#L113-L123

mingfeima · 2026-07-29T07:12:54Z

@jmunetong

ci has failure at

https://github.com/sgl-project/sglang/actions/runs/30398898975/job/90411477370?pr=32227

ERROR: test_layout_robustness_cases (__main__.TestTritonMamba2BackendCorrectness) (case='mamba2_extend_zero_prefix_exact_page', layout='interleaved_pages')
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/actions-runner/_work/sglang/sglang/test/registered/attention/unittests/mamba/test_mamba2.py", line 145, in test_layout_robustness_cases
    run_mamba2_attention_case(self, case, loc_layout=layout)
  File "/actions-runner/_work/sglang/sglang/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py", line 853, in run_mamba2_attention_case
    fixture = build_mamba2_attention_fixture(
  File "/actions-runner/_work/sglang/sglang/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py", line 692, in build_mamba2_attention_fixture
    ATTENTION_BACKENDS[case.backend](runner)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/attention_registry.py", line 172, in create_triton_backend
    return TritonAttnBackend(runner)
  File "/actions-runner/_work/sglang/sglang/python/sglang/srt/layers/attention/triton_backend.py", line 196, in __init__
    self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim()
AttributeError: 'MHATokenToKVPool' object has no attribute 'get_v_head_dim'. Did you mean: 'v_head_dim'?

jmunetong · 2026-07-30T14:14:58Z

@mingfeima . Thanks for letting me know. I have addressed this and test should be passing now.

mingfeima · 2026-08-03T02:19:22Z

@mingfeima . Thanks for letting me know. I have addressed this and test should be passing now.

sglang project CIs are kind of tricky as you can see a lot of REDs and you need to identify whether your failures are real or they are just failed on other PRs as well. You can also let AI do the job for you, for example https://github.com/mingfeima/sglang/tree/main/.claude/skills/intel-xpu-cpu-review-pr, this saves me a lot of work.

anyway still need to be careful :)

The mamba extra_buffer radix-cache strategy is gated to CUDA/MUSA/NPU/ROCm by a device assertion, so hybrid-Mamba models (e.g. NVIDIA-Nemotron-3-Nano) launched on --device xpu with the default (auto) strategy hit: AssertionError: extra_buffer needs CUDA/MUSA/NPU/ROCm (FLA). The extra_buffer path is FLA/Triton-backed and Triton has an XPU backend, so add is_xpu() to the guard to let the path run on XPU. This unblocks launch; runtime correctness of the FLA extra-buffer kernels on XPU still needs validation (tracked separately).

TritonAttnBackend.__init__ picks v_head_dim via a hybrid-model branch that only checked hybrid_gdn_config / kimi_linear_config / linear_attn_model_spec. NemotronH (and other mamba2 archs: FalconH1, Lfm2, ...) are detected by mamba2_config, not those three, so they fell through to the else branch which indexes layer 0 as full-attention: ValueError: layer_id=0 not in full attention layers: dict_keys([5, 12, 19, 26, 33, 42]) NemotronH's full-attention layers are [5,12,19,26,33,42], so layer 0 is a mamba layer and get_value_buffer(0) raises. Switch the branch to the umbrella mambaish_config(), which unions mamba2 + hybrid-GDN + kimi-linear + linear-attn, so all hybrid linear models take the get_v_head_dim() path. Surfaced bringing up NVIDIA-Nemotron-3-Nano-30B on --device xpu.

TritonAttnBackend.__init__ routes mambaish (mamba2 / hybrid-GDN / kimi-linear / linear-attn) configs through token_to_kv_pool.get_v_head_dim(), but that method only existed on HybridLinearKVPool. When a mamba2 config is served by a plain MHATokenToKVPool (as in the mamba2 attention unit-test fixture), the call raised: AttributeError: 'MHATokenToKVPool' object has no attribute 'get_v_head_dim' Add get_v_head_dim() to MHATokenToKVPool returning self.v_head_dim (uniform across its all-full-attention layers), mirroring HybridLinearKVPool. Inherited by NoOp/PageMajor/MXFP8/FP4 subclasses. Fixes CI failure in TestTritonMamba2BackendCorrectness.test_layout_robustness_cases.

mingfeima · 2026-08-05T02:19:03Z

failure on base-c-test-8-gpu-h20 seems flaky, check again.

mingfeima · 2026-08-06T00:26:31Z

@arathi-hlab could you please help check xpu ci failure?

Resolve triton_backend.py import conflict: keep mambaish_config (this branch's umbrella hybrid detector) alongside upstream's new is_kimi_k3 import for the use_verify_mla gate.

airMeng · 2026-08-11T19:15:43Z

/rerun-failed-ci

mingfeima · 2026-08-12T05:23:09Z

haha at last, xpu ci passed...

…ct#32227) Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

jmunetong marked this pull request as ready for review July 23, 2026 17:11

jmunetong requested review from Fridge003, HaiShaw, Qiaolin-Yu, hebiao064, ispobock and merrymercy as code owners July 23, 2026 17:11

airMeng reviewed Jul 23, 2026

View reviewed changes

mingfeima approved these changes Jul 24, 2026

View reviewed changes

mingfeima added intel ci continue integration related xpu intel gpu with device `torch.xpu` run-ci CI: run the baseline test suite on this PR and removed ci continue integration related labels Jul 24, 2026

jmunetong requested review from Ying1123, alphabetc1, hanming-lu, hnyls2002, hzh0425, xiezhq-hermann and yizhang2077 as code owners July 29, 2026 21:58

jmunetong added 4 commits August 3, 2026 23:02

ci: trigger re-run

90a8f6f

jmunetong force-pushed the xpu-mamba-extra-buffer branch from d88f83f to 90a8f6f Compare August 3, 2026 23:42

jmunetong added 3 commits August 4, 2026 09:16

Merge branch 'main' into xpu-mamba-extra-buffer

b25eea8

Merge branch 'main' into xpu-mamba-extra-buffer

03986ac

Merge branch 'main' into xpu-mamba-extra-buffer

424086e

mingfeima and others added 2 commits August 5, 2026 10:19

Merge branch 'main' into xpu-mamba-extra-buffer

7a7e8cb

Merge branch 'main' into xpu-mamba-extra-buffer

9979c03

jmunetong and others added 5 commits August 6, 2026 16:11

Merge branch 'main' into xpu-mamba-extra-buffer

8eed890

Merge upstream/main into xpu-mamba-extra-buffer

ad03df7

Resolve triton_backend.py import conflict: keep mambaish_config (this branch's umbrella hybrid detector) alongside upstream's new is_kimi_k3 import for the use_verify_mla gate.

Merge branch 'main' into xpu-mamba-extra-buffer

e2bcb25

Merge branch 'main' into xpu-mamba-extra-buffer

0dbd389

Merge branch 'main' into xpu-mamba-extra-buffer

ee70b5f

Merge branch 'main' into xpu-mamba-extra-buffer

78f39b1

mingfeima merged commit cac3269 into sgl-project:main Aug 12, 2026
150 of 171 checks passed

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026

[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu (sgl-proje…

43e646a

…ct#32227) Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026

[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu (sgl-proje…

5fc0cf8

…ct#32227) Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

Xia-Weiwen mentioned this pull request Sep 17, 2026

[XPU] Qwen3.8-flash-next enablement #37213

Merged

3 of 5 tasks

jmunetong mentioned this pull request Oct 6, 2026

[XPU] Let Inkling use the Mamba extra_buffer strategy on XPU #42814

Draft

5 tasks done

Navigation Menu

Repository navigation

[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu - #32227

[XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu#32227
mingfeima merged 15 commits into
sgl-project:mainfrom
jmunetong:xpu-mamba-extra-buffer

jmunetong commented Jul 23, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 23, 2026

Uh oh!

gemini-code-assist Bot commented Jul 23, 2026

Uh oh!

airMeng Jul 23, 2026

Uh oh!

mingfeima Jul 24, 2026

Uh oh!

mingfeima commented Jul 29, 2026

Uh oh!

jmunetong commented Jul 30, 2026

Uh oh!

mingfeima commented Aug 3, 2026

Uh oh!

mingfeima commented Aug 5, 2026

Uh oh!

mingfeima commented Aug 6, 2026

Uh oh!

airMeng commented Aug 11, 2026

Uh oh!

mingfeima commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

jmunetong commented Jul 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Error 1 — extra_buffer mamba radix cache asserts out on XPU

Error 2 — Triton backend indexes layer 0 as full-attention for a hybrid model

Modifications

Why the fix is in the Triton backend, not the intel_xpu backend

Accuracy / correctness

Checklist / how to reproduce & validate

CI States

Uh oh!

gemini-code-assist Bot commented Jul 23, 2026

Uh oh!

gemini-code-assist Bot commented Jul 23, 2026

Uh oh!

airMeng Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

mingfeima Jul 24, 2026

Choose a reason for hiding this comment

Uh oh!

mingfeima commented Jul 29, 2026

Uh oh!

jmunetong commented Jul 30, 2026

Uh oh!

mingfeima commented Aug 3, 2026

Uh oh!

mingfeima commented Aug 5, 2026

Uh oh!

mingfeima commented Aug 6, 2026

Uh oh!

airMeng commented Aug 11, 2026

Uh oh!

mingfeima commented Aug 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jmunetong commented Jul 23, 2026 •

edited by github-actions Bot

Loading

Error 1 — `extra_buffer` mamba radix cache asserts out on XPU

Why the fix is in the Triton backend, not the `intel_xpu` backend