Add torch.compile support for qwen3-next on CPU (#12444)

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -384,6 +384,20 @@ class MambaAttnBackendBase(AttentionBackend):
             bs, req_pool_indices, forward_mode, spec_info, seq_lens_cpu
         )
 
+    def init_forward_metadata_capture_cpu_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        self.forward_metadata = self._capture_metadata(
+            bs, req_pool_indices, forward_mode, spec_info
+        )
+
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         assert (
             max_num_tokens % max_bs == 0
@@ -424,6 +438,23 @@ class MambaAttnBackendBase(AttentionBackend):
             device=self.device,
         )
 
+    def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
+        assert (
+            max_num_tokens % max_bs == 0
+        ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}"
+        for i in range(max_bs):
+            self.state_indices_list.append(
+                torch.full(
+                    (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device
+                )
+            )
+            self.query_start_loc_list.append(
+                torch.empty((i + 2,), dtype=torch.int32, device=self.device)
+            )
+        self.cached_cuda_graph_decode_query_start_loc = torch.arange(
+            0, max_bs + 1, dtype=torch.int32, device=self.device
+        )
+
     def _capture_metadata(
         self,
         bs: int,
@@ -530,6 +561,9 @@ class MambaAttnBackendBase(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return 1  # Mamba attn does not use seq lens to index kv cache
 
+    def get_cpu_graph_seq_len_fill_value(self):
+        return 1
+
     def _track_mamba_state_decode(
         self,
         forward_batch: ForwardBatch,
@@ -706,6 +740,10 @@ class HybridLinearAttnBackend(AttentionBackend):
         for attn_backend in self.attn_backend_list:
             attn_backend.init_cuda_graph_state(max_bs, max_num_tokens)
 
+    def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
+        for attn_backend in self.attn_backend_list:
+            attn_backend.init_cpu_graph_state(max_bs, max_num_tokens)
+
     def init_forward_metadata_capture_cuda_graph(
         self,
         bs: int,
@@ -727,6 +765,27 @@ class HybridLinearAttnBackend(AttentionBackend):
                 spec_info,
             )
 
+    def init_forward_metadata_capture_cpu_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        for attn_backend in self.attn_backend_list:
+            attn_backend.init_forward_metadata_capture_cpu_graph(
+                bs,
+                num_tokens,
+                req_pool_indices,
+                seq_lens,
+                encoder_lens,
+                forward_mode,
+                spec_info,
+            )
+
     def init_forward_metadata_replay_cuda_graph(
         self,
         bs: int,
@@ -753,6 +812,9 @@ class HybridLinearAttnBackend(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return self.full_attn_backend.get_cuda_graph_seq_len_fill_value()
 
+    def get_cpu_graph_seq_len_fill_value(self):
+        return self.full_attn_backend.get_cpu_graph_seq_len_fill_value()
+
     def forward_decode(
         self,
         layer: RadixAttention,

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -57,9 +57,35 @@ class IntelAMXAttnBackend(AttentionBackend):
             max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
         self.forward_metadata = (attn_logits, max_extend_len)
 
-    def get_graph_seq_len_fill_value(self):
+    def get_cpu_graph_seq_len_fill_value(self):
         return 1
 
+    def init_forward_metadata_capture_cpu_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens,
+        forward_mode,
+        spec_info,
+    ):
+        attn_logits = torch.zeros(
+            (
+                bs,
+                self.num_head,
+                8,  # self.num_kv_splits,
+                self.v_head_dim + 1,
+            ),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        max_extend_len = None
+        self.forward_metadata = (attn_logits, max_extend_len)
+
+    def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
+        pass
+
     def forward_extend(
         self,
         q,

python/sglang/srt/mem_cache/memory_pool.py

@@ -28,7 +28,7 @@ import abc
 import dataclasses
 import logging
 from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 import numpy as np
@@ -193,11 +193,15 @@ class MambaPool:
 
         def at_layer_idx(self, layer: int):
             kwargs = {}
-            for k, v in vars(self).items():
-                if k == "conv" or k == "intermediate_conv_window":
-                    kwargs[k] = [conv[layer] for conv in v]
+            # Use fields instead of vars to avoid torch.compile graph break
+            for f in fields(self):
+                name = f.name
+                v = getattr(self, name)
+                if name in ("conv", "intermediate_conv_window"):
+                    kwargs[name] = [conv[layer] for conv in v]
                 else:
-                    kwargs[k] = v[layer]
+                    kwargs[name] = v[layer]
+
             return type(self)(**kwargs)
 
         def mem_usage_bytes(self):

python/sglang/srt/model_executor/cpu_graph_runner.py

@@ -114,6 +114,9 @@ def register_fake_ops():
         "fused_add_rmsnorm_cpu",
         "decode_attention_cpu",
         "extend_attention_cpu",
+        "gemma_fused_add_rmsnorm_cpu",
+        "layernorm_cpu",
+        "fused_add_layernorm_cpu",
     ]
     for op in none_return_ops:
 
@@ -125,7 +128,12 @@ def register_fake_ops():
         "rmsnorm_cpu",
         "l2norm_cpu",
         "fused_experts_cpu",
+        "fused_rmsnorm_gated_cpu",
         "shared_expert_cpu",
+        "causal_conv1d_update_cpu",
+        "causal_conv1d_fwd_cpu",
+        "gemma_rmsnorm_cpu",
+        "gemma3_rmsnorm_cpu",
     ]:
 
         @torch.library.register_fake(f"sgl_kernel::{op}")
@@ -225,9 +233,19 @@ def register_fake_ops():
         v_input = k_input.narrow(-1, 0, kv_lora_rank)
         return q_input, k_input, v_input
 
+    def get_n_size(mat2, is_vnni):
+        tile_n = 16
+        if mat2.dtype == torch.float32:
+            return mat2.shape[1]
+        if not is_vnni and mat2.dim() == 2 and mat2.shape[0] < tile_n:
+            return mat2.shape[1]
+        return mat2.shape[0]
+
     @torch.library.register_fake("sgl_kernel::weight_packed_linear")
-    def _(x, weight, bias, is_vnni):
-        return x.new_empty(x.shape[0], weight.shape[0])
+    def _(mat1, mat2, bias, is_vnni):
+        M = mat1.shape[0]
+        N = get_n_size(mat2, is_vnni)
+        return mat1.new_empty(M, N)
 
     @torch.library.register_fake("sgl_kernel::per_token_quant_int8_cpu")
     def _(input):
@@ -306,9 +324,19 @@ def register_fake_ops():
             torch.empty(shape, device=hidden_states.device, dtype=torch.int),
         )
 
-    @torch.library.register_fake("sgl_kernel::silu_and_mul_cpu")
-    def _(input):
-        return input.new_empty(input.shape[0], input.shape[1] // 2)
+    for act_op in [
+        "silu_and_mul_cpu",
+        "gelu_tanh_and_mul_cpu",
+        "gelu_and_mul_cpu",
+    ]:
+
+        @torch.library.register_fake(f"sgl_kernel::{act_op}")
+        def _(input):
+            sizes = list(input.shape)
+            last_dim = input.dim() - 1
+            d = sizes[last_dim] // 2
+            sizes[last_dim] = d
+            return input.new_empty(sizes)
 
     @torch.library.register_fake("sgl_kernel::int8_scaled_mm_with_quant")
     def _(
@@ -337,6 +365,82 @@ def register_fake_ops():
         N = mat2.shape[0]
         return mat1.new_empty(M, N, dtype=out_dtype)
 
+    @torch.library.register_fake("sgl_kernel::fused_linear_sigmoid_mul")
+    def _(
+        mat1,
+        mat2,
+        bias,
+        is_vnni,
+        post_mul_mat,
+    ):
+        M = mat1.shape[0]
+        N = post_mul_mat.shape[1]
+        return mat1.new_empty(M, N)
+
+    @torch.library.register_fake("sgl_kernel::fused_qkvzba_split_reshape_cat_cpu")
+    def _(mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_qk, head_v):
+        batch = mixed_qkvz.shape[0]
+        qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v
+        mixed_qkv = mixed_qkvz.new_empty(batch, qkv_dim)
+        z = mixed_qkvz.new_empty(batch, num_heads_v, head_v)
+        b = mixed_ba.new_empty(batch, num_heads_v)
+        a = mixed_ba.new_empty(batch, num_heads_v)
+        return mixed_qkv, z, b, a
+
+    @torch.library.register_fake(
+        "sgl_kernel::fused_sigmoid_gating_delta_rule_update_cpu"
+    )
+    def _(
+        A_log,
+        dt_bias,
+        q,
+        k,
+        v,
+        a,
+        b,
+        initial_state_source,
+        initial_state_indices,
+        cu_seqlens,
+        use_qk_l2norm_in_kernel,
+        softplus_beta=1.0,
+        softplus_threshold=20.0,
+    ):
+        assert q.dim() == 4
+        assert v.dim() == 4
+        batch_size = q.shape[1]
+        seq_len = q.shape[0]
+        v_num_heads = v.shape[2]
+        v_head_dim = v.shape[3]
+        return q.new_empty(batch_size, seq_len, v_num_heads, v_head_dim)
+
+    @torch.library.register_fake("sgl_kernel::fused_gdn_gating_cpu")
+    def _(A_log, a, b, dt_bias):
+        batch = a.shape[0]
+        num_heads = a.shape[1]
+        out = a.new_empty(1, batch, num_heads, dtype=torch.float)
+        beta = b.new_empty(1, batch, num_heads)
+        return out, beta
+
+    @torch.library.register_fake("sgl_kernel::chunk_gated_delta_rule_cpu")
+    def _(
+        query,
+        key,
+        value,
+        g,
+        beta,
+        initial_state,
+        output_final_state,
+        cu_seqlens,
+        head_first,
+        use_qk_l2norm_in_kernel,
+        eps,
+    ):
+        output = torch.empty_like(value)
+        assert initial_state is not None
+        final_state = initial_state.to(torch.float32)
+
+        return output, final_state
+
 
 # TODO Remove unnecessary settings for CPUGraphRunner.
 # Re-abstract the graph runner and restructure CPUGraphRunner to reuse the same logic.
@@ -406,9 +510,12 @@ class CPUGraphRunner:
         # Attention backend
         self.max_bs = max(self.capture_bs)
         self.max_num_token = self.max_bs * self.num_tokens_per_bs
+        self.model_runner.attn_backend.init_cpu_graph_state(
+            self.max_bs, self.max_num_token
+        )
 
         self.seq_len_fill_value = (
-            self.model_runner.attn_backend.get_graph_seq_len_fill_value()
+            self.model_runner.attn_backend.get_cpu_graph_seq_len_fill_value()
         )
 
         if self.enable_torch_compile:
@@ -488,6 +595,29 @@ class CPUGraphRunner:
                 self.graphs[bs] = graph
                 self.output_buffers[bs] = output_buffers
 
+        # Re-init states for qwen3-next as
+        # torch.compile may change the states
+        self._reset_mamba_cache_if_needed()
+
+    def _reset_mamba_cache_if_needed(self) -> None:
+
+        mamba_pool = getattr(self.model_runner.req_to_token_pool, "mamba_pool", None)
+        if mamba_pool is None:
+            return
+        mamba_cache = getattr(mamba_pool, "mamba_cache", None)
+        if mamba_cache is None:
+            return
+
+        def _zero_nested(obj):
+            if isinstance(obj, torch.Tensor):
+                obj.zero_()
+            elif isinstance(obj, (list, tuple)):
+                for it in obj:
+                    _zero_nested(it)
+
+        for v in vars(mamba_cache).values():
+            _zero_nested(v)
+
     def capture_one_batch_size(self, bs: int, forward: Callable):
         num_tokens = bs * self.num_tokens_per_bs
 
@@ -497,7 +627,7 @@ class CPUGraphRunner:
         seq_lens = self.seq_lens[:bs]
         out_cache_loc = self.out_cache_loc[:num_tokens]
         positions = self.positions[:num_tokens]
-        mrope_positions = self.mrope_positions[:, :bs]
+        mrope_positions = self.mrope_positions[:, :num_tokens]
         self.num_token_non_padded[...] = num_tokens
 
         spec_info = self.get_spec_info(num_tokens)
@@ -528,21 +658,31 @@ class CPUGraphRunner:
         )
 
         # Attention backend
-        self.model_runner.attn_backend.init_forward_metadata(forward_batch)
+        self.model_runner.attn_backend.init_forward_metadata_capture_cpu_graph(
+            bs,
+            num_tokens,
+            req_pool_indices,
+            seq_lens,
+            None,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
+        )
         # Do infernence to avoid setting attr at runtime, e.g.,
         # self.attn_mha.kv_b_proj = self.kv_b_proj for full graph compile on CPU
-        self.model_runner.model.forward(
-            forward_batch.input_ids,
-            forward_batch.positions,
-            forward_batch,
-        )
+        with torch.no_grad():
+            self.model_runner.tp_group.barrier()
+            self.model_runner.model.forward(
+                forward_batch.input_ids,
+                forward_batch.positions,
+                forward_batch,
+            )
 
         # Run and capture
         def run_once():
             # Clean intermediate result cache for DP attention
             forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
             logits_output_or_pp_proxy_tensors = forward(
-                input_ids,
+                forward_batch.input_ids,
                 forward_batch.positions,
                 forward_batch,
             )