From 4f0f6cd9d0259fb5081ecc565dad2bb47643e8fa Mon Sep 17 00:00:00 2001 From: Cao E Date: Fri, 27 Feb 2026 15:28:03 +0800 Subject: [PATCH] Add torch.compile support for qwen3-next on CPU (#12444) --- .../attention/hybrid_linear_attn_backend.py | 62 +++++++ .../srt/layers/attention/intel_amx_backend.py | 28 ++- python/sglang/srt/mem_cache/memory_pool.py | 14 +- .../srt/model_executor/cpu_graph_runner.py | 168 ++++++++++++++++-- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 8 +- 5 files changed, 256 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 4fd4d48ac..91194c494 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py index 7ab275374..e64bb45c0 100644 --- a/python/sglang/srt/layers/attention/intel_amx_backend.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5312b2edb..c51e004fb 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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): diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 20579a3c2..f71e83ec4 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/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, ) diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index da7d09cc3..6b29da857 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -359,11 +359,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.impl("l2norm_cpu", torch::kCPU, &l2norm_cpu); m.def("fused_rmsnorm_gated_cpu(Tensor input, Tensor weight, Tensor gate, float eps) -> Tensor"); m.impl("fused_rmsnorm_gated_cpu", torch::kCPU, &fused_rmsnorm_gated_cpu); - m.def("fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor residual, Tensor weight, float eps) -> ()"); + m.def("fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()"); m.impl("fused_add_rmsnorm_cpu", torch::kCPU, &fused_add_rmsnorm_cpu); - m.def("gemma_fused_add_rmsnorm_cpu(Tensor input, Tensor residual, Tensor weight, float eps) -> ()"); + m.def("gemma_fused_add_rmsnorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()"); m.impl("gemma_fused_add_rmsnorm_cpu", torch::kCPU, &gemma_fused_add_rmsnorm_cpu); - m.def("fused_add_layernorm_cpu(Tensor(a!) input, Tensor residual, Tensor weight, float eps) -> ()"); + m.def("fused_add_layernorm_cpu(Tensor(a!) input, Tensor(a!) residual, Tensor weight, float eps) -> ()"); m.impl("fused_add_layernorm_cpu", torch::kCPU, &fused_add_layernorm_cpu); // topk @@ -502,7 +502,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); m.def( - "causal_conv1d_update_cpu(Tensor x, Tensor conv_states, Tensor weight, Tensor? bias, bool silu_activation," + "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor weight, Tensor? bias, bool silu_activation," "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor"); m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);