From cac32693056f8efd094e31680e2fec1485f4aaf4 Mon Sep 17 00:00:00 2001 From: Juan Muneton <102537701+jmunetong@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:23:23 -0700 Subject: [PATCH] [XPU] Fix NemotronH (hybrid mamba2) launch on --device xpu (#32227) Co-authored-by: Ma Mingfei --- .../srt/layers/attention/triton_backend.py | 16 ++++++---------- python/sglang/srt/mem_cache/memory_pool.py | 7 +++++++ python/sglang/srt/server_args.py | 4 ++-- 3 files changed, 15 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 54a6c2885..c6b9cb92b 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -10,11 +10,7 @@ from sglang.kernels.ops.attention.metadata import get_num_kv_splits_triton from sglang.kernels.ops.kvcache.kv_indices import ( create_flashinfer_kv_indices_triton, ) -from sglang.srt.configs.hybrid_arch import ( - hybrid_gdn_config, - kimi_linear_config, - linear_attn_model_spec, -) +from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.model_config import AttentionArch, is_kimi_k3 from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, @@ -213,12 +209,12 @@ class TritonAttnBackend(AttentionBackend): if self.sliding_window_size is not None and swa_v_head_dim != full_v_head_dim: self.v_head_dim = full_v_head_dim self.swa_v_head_dim = swa_v_head_dim - elif ( - hybrid_gdn_config(model_runner.model_config) is not None - 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: # For hybrid linear models, layer_id = 0 may not be full attention + # (e.g. NemotronH's full-attn layers are [5,12,19,...]). mambaish_config + # unions mamba2 (NemotronH/FalconH1/...), hybrid-GDN, kimi-linear, and + # linear-attn specs, so we ask get_v_head_dim() instead of indexing + # layer 0, which is not guaranteed to be a full-attention layer. self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim() self.swa_v_head_dim = None else: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 55b1f7269..cecd7b940 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -2303,6 +2303,13 @@ class MHATokenToKVPool(KVCache): self.layer_transfer_counter.wait_until(layer_id - self.start_layer) return self._get_value_buffer(layer_id) + def get_v_head_dim(self): + # Every layer in this pool is full-attention, so the value head dim is + # uniform and known at construction. Mirrors HybridLinearKVPool's + # get_v_head_dim() so the TritonAttnBackend mambaish branch works when a + # mamba2 config is served by a plain MHA pool (no per-linear-layer split). + return self.v_head_dim + def get_kv_buffer(self, layer_id: int): return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ad34ea51b..22782138b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5707,8 +5707,8 @@ class ServerArgs: view, model_arch ), f"extra_buffer is not supported for {model_arch}; use no_buffer." assert ( - is_cuda() or is_musa() or is_npu() or is_hip() - ), "extra_buffer needs CUDA/MUSA/NPU/ROCm (FLA)." + is_cuda() or is_musa() or is_npu() or is_hip() or is_xpu() + ), "extra_buffer needs CUDA/MUSA/NPU/ROCm/XPU (FLA)." if view.mamba_radix_cache_strategy == "extra_buffer_lazy": # The PD-disagg decode pool is not wired for lazy slots. assert view.disaggregation_mode == "null", (