Support speculative decoding on CPU (#27862)

Co-authored-by: Valentine233 <xuan.liao@intel.com>

python/sglang/srt/arg_groups/speculative_hook.py

@@ -11,6 +11,16 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _disable_overlap_schedule_for_cpu(server_args: ServerArgs) -> None:
+    if server_args.device != "cpu" or server_args.disable_overlap_schedule:
+        return
+
+    server_args.disable_overlap_schedule = True
+    logger.warning(
+        "Overlap schedule is not implemented for speculative decoding on CPU."
+    )
+
+
 def _resolve_speculative_algorithm_alias(
     speculative_algorithm: Optional[str],
     speculative_draft_model_path: Optional[str],
@@ -332,6 +342,8 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
             "Max running requests is reset to 48 for speculative decoding. You can override this by explicitly setting --max-running-requests."
         )
 
+    _disable_overlap_schedule_for_cpu(server_args)
+
     if resolved_view(server_args).disable_overlap_schedule:
         logger.warning(
             "Non-overlap (synchronous) spec v2 is used for eagle/eagle3/standalone "
@@ -469,8 +481,12 @@ def _handle_eagle_family(server_args: ServerArgs) -> None:
 
 
 def _handle_ngram(server_args: ServerArgs) -> None:
-    if not server_args.device.startswith("cuda"):
-        raise ValueError("Ngram speculative decoding only supports CUDA device.")
+    if server_args.device not in ("cuda", "cpu"):
+        raise ValueError(
+            "Ngram speculative decoding only supports CUDA or CPU devices."
+        )
+
+    _disable_overlap_schedule_for_cpu(server_args)
 
     if server_args.max_running_requests is None:
         server_args.max_running_requests = 48

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

@@ -20,11 +20,14 @@ class IntelAMXAttnBackend(AttentionBackend):
 
         super().__init__()
         self.forward_metadata = None
+        self.extend_metadata = None
+        self.draft_decode_metadata = None
         self.device = model_runner.device
         # Pool refs — captured at construction so they survive deletion of the
         # corresponding ForwardBatch fields.
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
+        self.max_context_len = model_runner.model_config.context_len
 
         # full->SWA translated out_cache_loc, computed once per forward (the only
         # set_kv_buffer is in eager forward_extend; decode writes KV in-kernel).
@@ -51,6 +54,68 @@ class IntelAMXAttnBackend(AttentionBackend):
         self.decode_attention_fwd = torch.ops.sgl_kernel.decode_attention_cpu
         self.extend_attention_fwd = torch.ops.sgl_kernel.extend_attention_cpu
 
+        # Number of KV splits used by decode_attention_cpu; attn_logits is
+        # sized [bs, num_head, num_kv_splits, v_head_dim + 1] to match.
+        self.num_kv_splits = 8
+
+        # speculative decoding params
+        self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
+
+    def _build_extend_metadata(self, forward_batch: ForwardBatch):
+        """Resolve (seq_lens, extend_seq_lens, extend_start_loc, tree_mask) for
+        forward_extend, once per forward pass.
+
+        In TARGET_VERIFY mode the batch carries no extend_* fields, so they are
+        derived from spec_info (mirrors the CUDA unified path in
+        triton_backend.py); each request extends by exactly num_draft_tokens
+        tokens. Outside spec decoding the fields are passed through.
+        """
+        bs = forward_batch.batch_size
+        seq_lens = forward_batch.seq_lens
+        tree_mask = None
+
+        if forward_batch.forward_mode.is_target_verify():
+            spec_info = forward_batch.spec_info
+            if spec_info is None:
+                raise RuntimeError(
+                    "spec_info is unset in TARGET_VERIFY mode; the extend_* "
+                    "metadata can only be derived from spec_info for "
+                    "speculative verify batches."
+                )
+            num_draft_tokens = spec_info.draft_token_num
+            extend_seq_lens = torch.full(
+                (bs,), num_draft_tokens, dtype=torch.int32, device=self.device
+            )
+            # Uniform extend lengths: start locations form a plain range.
+            extend_start_loc = torch.arange(
+                0,
+                bs * num_draft_tokens,
+                num_draft_tokens,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            seq_lens = forward_batch.seq_lens + num_draft_tokens
+            # Speculative verify with a token tree: each draft token may only
+            # attend to its ancestors among the draft tokens (the committed
+            # prefix stays fully visible).
+            #
+            # NOTE: unlike triton_backend.py, which forwards spec_info.custom_mask
+            # unconditionally, the mask is gated on tree_topk here. tree_topk == 1
+            # means the draft tokens form a simple chain whose visibility is
+            # exactly the kernel's built-in causal masking, and skipping the explicit
+            # mask lets extend_attention_cpu take its faster mask-free path. EAGLE
+            # has tree_topk == topk (> 1 for real trees); NGRAM has tree_topk == -1
+            # (irregular tree); both need the mask.
+            if spec_info.tree_topk != 1:
+                custom_mask = spec_info.custom_mask
+                if custom_mask is not None and custom_mask.numel() > 0:
+                    tree_mask = custom_mask
+        else:
+            extend_seq_lens = forward_batch.extend_seq_lens
+            extend_start_loc = forward_batch.extend_start_loc
+
+        return seq_lens, extend_seq_lens, extend_start_loc, tree_mask
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Init the metadata for a forward pass."""
 
@@ -59,7 +124,7 @@ class IntelAMXAttnBackend(AttentionBackend):
             (
                 bs,
                 self.num_head,
-                8,  # self.num_kv_splits,
+                self.num_kv_splits,
                 self.v_head_dim + 1,
             ),
             dtype=torch.float32,
@@ -67,8 +132,13 @@ class IntelAMXAttnBackend(AttentionBackend):
         )
         if forward_batch.forward_mode.is_decode_or_idle():
             max_extend_len = None
+            self.extend_metadata = None
+        elif forward_batch.forward_mode.is_target_verify():
+            max_extend_len = self.num_draft_tokens
+            self.extend_metadata = self._build_extend_metadata(forward_batch)
         else:
             max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
+            self.extend_metadata = self._build_extend_metadata(forward_batch)
         self.forward_metadata = (attn_logits, max_extend_len)
 
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
@@ -97,7 +167,7 @@ class IntelAMXAttnBackend(AttentionBackend):
             (
                 bs,
                 self.num_head,
-                8,  # self.num_kv_splits,
+                self.num_kv_splits,
                 self.v_head_dim + 1,
             ),
             dtype=torch.float32,
@@ -105,6 +175,7 @@ class IntelAMXAttnBackend(AttentionBackend):
         )
         max_extend_len = None
         self.forward_metadata = (attn_logits, max_extend_len)
+        self.extend_metadata = None
 
     def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
         pass
@@ -136,6 +207,10 @@ class IntelAMXAttnBackend(AttentionBackend):
                 layer, KVWriteLoc(cache_loc, swa_loc), k, v
             )
 
+        # Precomputed once per forward pass in init_forward_metadata (spec
+        # verify batches carry no extend_* fields; see _build_extend_metadata).
+        seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata
+
         _, max_extend_len = self.forward_metadata
         self.extend_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
@@ -146,9 +221,9 @@ class IntelAMXAttnBackend(AttentionBackend):
             self.token_to_kv_pool.get_value_buffer(layer.layer_id),
             self.req_to_token_pool.req_to_token,
             forward_batch.req_pool_indices,
-            forward_batch.seq_lens,
-            forward_batch.extend_seq_lens,
-            forward_batch.extend_start_loc,
+            seq_lens,
+            extend_seq_lens,
+            extend_start_loc,
             max_extend_len,
             layer.scaling,
             layer.logit_cap,
@@ -156,6 +231,7 @@ class IntelAMXAttnBackend(AttentionBackend):
             layer.sliding_window_size + 1,
             forward_batch.encoder_lens,
             sinks,
+            tree_mask,
         )
         return o
 
@@ -171,6 +247,13 @@ class IntelAMXAttnBackend(AttentionBackend):
     ):
         attn_logits, _ = self.forward_metadata
 
+        if self.draft_decode_metadata is not None:
+            req_to_token, seq_lens, req_pool_indices = self.draft_decode_metadata
+        else:
+            req_to_token = self.req_to_token_pool.req_to_token
+            req_pool_indices = forward_batch.req_pool_indices
+            seq_lens = forward_batch.seq_lens
+
         q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
 
         if layer.qk_head_dim != layer.v_head_dim:
@@ -191,9 +274,9 @@ class IntelAMXAttnBackend(AttentionBackend):
             v,
             cache_loc,
             attn_logits,
-            self.req_to_token_pool.req_to_token,
-            forward_batch.req_pool_indices,
-            forward_batch.seq_lens,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
             layer.scaling,
             layer.logit_cap,
             layer.is_cross_attention,
@@ -205,3 +288,67 @@ class IntelAMXAttnBackend(AttentionBackend):
 
     def support_triton(self):
         return False
+
+
+class IntelAMXMultiStepDraftBackend:
+    """
+    Wrap multiple intel amx attention backends as one for multiple consecutive
+    draft decoding steps.
+    """
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        from sgl_kernel import build_draft_decode_metadata_cpu
+
+        self.build_draft_decode_metadata = build_draft_decode_metadata_cpu
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.attn_backends: list[IntelAMXAttnBackend] = []
+        for _ in range(self.speculative_num_steps - 1):
+            self.attn_backends.append(IntelAMXAttnBackend(model_runner))
+        self.device = model_runner.device
+        self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        num_seqs = forward_batch.batch_size
+        topk = self.topk
+        bs = num_seqs * topk
+        num_steps = self.speculative_num_steps
+        req_to_token = self.attn_backends[0].req_to_token_pool.req_to_token
+        seq_lens = forward_batch.seq_lens
+        pool_len = self.pool_len
+        num_head = self.attn_backends[0].num_head
+        v_head_dim = self.attn_backends[0].v_head_dim
+        device = self.device
+
+        # Build expanded req_to_token via C++ kernel
+        req_to_token_draft = self.build_draft_decode_metadata(
+            req_to_token,
+            forward_batch.req_pool_indices,
+            seq_lens,
+            topk,
+            num_steps,
+            pool_len,
+        )
+
+        req_pool_indices_expanded = torch.arange(bs, dtype=torch.int64, device=device)
+
+        num_kv_splits = self.attn_backends[0].num_kv_splits
+        for step in range(num_steps - 1):
+            # Each candidate sees prefix + (step + 1) draft tokens.
+            seq_lens_expanded = seq_lens.repeat_interleave(topk) + step + 1
+            attn_logits = torch.zeros(
+                (bs, num_head, num_kv_splits, v_head_dim + 1),
+                dtype=torch.float32,
+                device=device,
+            )
+            self.attn_backends[step].forward_metadata = (attn_logits, None)
+            self.attn_backends[step].draft_decode_metadata = (
+                req_to_token_draft,
+                seq_lens_expanded,
+                req_pool_indices_expanded,
+            )

python/sglang/srt/layers/logits_processor.py

@@ -936,7 +936,7 @@ class LogitsProcessor(nn.Module):
         logits = self._copy_logits_to_buffer(logits, logits_metadata)
 
         if self.final_logit_softcapping:
-            if not _is_npu:
+            if not (_is_npu or _is_cpu):
                 fused_softcap(logits, self.final_logit_softcapping)
             else:
                 logits = self.final_logit_softcapping * torch.tanh(

python/sglang/srt/mem_cache/memory_pool.py

@@ -61,7 +61,7 @@ from sglang.srt.mem_cache.layout.page_major import (
     mha_entry_bytes,
 )
 from sglang.srt.mem_cache.triton_ops.cache_move import (
-    copy_all_layer_kv_cache_tiled,
+    copy_all_layer_kv_cache_func,
     set_kv_buffer_prefix_valid_tiled,
     store_cache_4d,
 )
@@ -1479,18 +1479,14 @@ class MHATokenToKVPool(KVCache):
         }
 
         dummy_loc = torch.zeros(chunk_upper, dtype=torch.int64, device=self.device)
-        grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"])
-
-        copy_all_layer_kv_cache_tiled[grid](
+        copy_all_layer_kv_cache_func(
             self.data_ptrs,
             self.data_strides,
             dummy_loc,
             dummy_loc,
             1,
             chunk_upper,
-            BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"],
-            num_warps=self._kv_copy_config["num_warps"],
-            num_stages=2,
+            self._kv_copy_config,
         )
 
     def _create_buffers(self):
@@ -1998,20 +1994,16 @@ class MHATokenToKVPool(KVCache):
 
         cfg = self._kv_copy_config
         cap = int(cfg.get("num_locs_upper", 256))
-        grid = (self.data_ptrs.numel(), cfg["byte_tiles"])
 
         if N <= cap:
-            upper = next_power_of_2(N)
-            copy_all_layer_kv_cache_tiled[grid](
+            copy_all_layer_kv_cache_func(
                 self.data_ptrs,
                 self.data_strides,
                 tgt_loc,
                 src_loc,
                 N,
-                upper,
-                BYTES_PER_TILE=cfg["bytes_per_tile"],
-                num_warps=cfg["num_warps"],
-                num_stages=2,
+                next_power_of_2(N),
+                cfg,
             )
             return
 
@@ -2019,17 +2011,14 @@ class MHATokenToKVPool(KVCache):
         for start in range(0, N, cap):
             end = min(start + cap, N)
             chunk_len = end - start
-            upper = next_power_of_2(chunk_len)
-            copy_all_layer_kv_cache_tiled[grid](
+            copy_all_layer_kv_cache_func(
                 self.data_ptrs,
                 self.data_strides,
                 tgt_loc[start:end],
                 src_loc[start:end],
                 chunk_len,
-                upper,
-                BYTES_PER_TILE=cfg["bytes_per_tile"],
-                num_warps=cfg["num_warps"],
-                num_stages=2,
+                next_power_of_2(chunk_len),
+                cfg,
             )
 
 

python/sglang/srt/mem_cache/triton_ops/cache_move.py

@@ -2,6 +2,13 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import copy_all_layer_kv_cache_cpu
+
 
 @triton.jit
 def set_kv_buffer_prefix_valid_tiled(
@@ -86,6 +93,37 @@ def copy_all_layer_kv_cache_tiled(
     tl.store(tgt_ptr, vals, mask=mask)
 
 
+def copy_all_layer_kv_cache_func(
+    data_ptrs: torch.Tensor,
+    strides: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+    num_locs: int,
+    num_locs_upper: int,
+    kv_copy_config: dict,
+):
+    if _is_cpu:
+        copy_all_layer_kv_cache_cpu(
+            data_ptrs,
+            strides,
+            tgt_loc[:num_locs],
+            src_loc[:num_locs],
+        )
+        return
+    grid = (data_ptrs.numel(), kv_copy_config["byte_tiles"])
+    copy_all_layer_kv_cache_tiled[grid](
+        data_ptrs,
+        strides,
+        tgt_loc,
+        src_loc,
+        num_locs,
+        num_locs_upper,
+        BYTES_PER_TILE=kv_copy_config["bytes_per_tile"],
+        num_warps=kv_copy_config["num_warps"],
+        num_stages=2,
+    )
+
+
 # ---------------------------------------------------------------------------
 # store_cache_4d — single-launch Triton write into the 4-D page-major envelope
 # K/V views. At `PAGE_SIZE = 1` the kernel constexpr-folds to byte-identical

python/sglang/srt/models/llama.py

@@ -52,6 +52,7 @@ from sglang.srt.model_loader.weight_utils import (
     kv_cache_scales_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.utils import add_prefix, is_cuda, is_npu, is_xpu, make_layers
 from sglang.utils import get_exception_traceback
@@ -783,8 +784,8 @@ class LlamaForCausalLM(nn.Module):
             torch.xpu.empty_cache()
             torch.xpu.synchronize()
         else:
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
+            current_platform.empty_cache()
+            current_platform.synchronize()
 
     def get_embed(self):
         return self.model.embed_tokens.weight
@@ -802,8 +803,8 @@ class LlamaForCausalLM(nn.Module):
             torch.xpu.empty_cache()
             torch.xpu.synchronize()
         else:
-            torch.cuda.empty_cache()
-            torch.cuda.synchronize()
+            current_platform.empty_cache()
+            current_platform.synchronize()
 
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
         self.model.load_kv_cache_scales(quantization_param_path)

python/sglang/srt/models/mimo_mtp.py

@@ -18,6 +18,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.qwen2 import Qwen2DecoderLayer
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_parallel
 
 
@@ -198,8 +199,8 @@ class MiMoMTP(nn.Module):
         del self.lm_head.weight
         self.model.embed_tokens.weight = embed
         self.lm_head.weight = head
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        current_platform.empty_cache()
+        current_platform.synchronize()
 
 
 EntryClass = MiMoMTP

python/sglang/srt/models/qwen2.py

@@ -49,6 +49,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     kv_cache_scales_loader,
 )
+from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
@@ -646,8 +647,8 @@ class Qwen2ForCausalLM(nn.Module):
         del self.lm_head.weight
         self.model.embed_tokens.weight = embed
         self.lm_head.weight = head
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        current_platform.empty_cache()
+        current_platform.synchronize()
 
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
         self.model.load_kv_cache_scales(quantization_param_path)

python/sglang/srt/models/qwen2_eagle.py

@@ -135,6 +135,7 @@ class Qwen2ForCausalLMEagle(Qwen2ForCausalLM):
                 prefix=add_prefix("lm_head", prefix),
             )
         self.logits_processor = LogitsProcessor(config)
+        self.capture_aux_hidden_states = False
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         for name, loaded_weight in weights:

python/sglang/srt/speculative/base_spec_worker.py

@@ -5,6 +5,13 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import assign_draft_cache_locs_contiguous_cpu
+
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
     from sglang.srt.managers.tp_worker import TpModelWorker
@@ -202,16 +209,27 @@ class EagleDraftWorkerBase(ABC):
                     dtype=torch.int64,
                     device=batch.device,
                 )
-                # FIXME(lsyin): align with the default code path
-                assign_draft_cache_locs_contiguous[(bs,)](
-                    batch.req_pool_indices,
-                    req_to_token_pool.req_to_token,
-                    batch.seq_lens,
-                    batch.out_cache_loc,
-                    req_to_token_pool.req_to_token.shape[1],
-                    topk,
-                    num_steps,
-                )
+                if _is_cpu:
+                    assign_draft_cache_locs_contiguous_cpu(
+                        batch.req_pool_indices,
+                        req_to_token_pool.req_to_token,
+                        batch.seq_lens,
+                        batch.out_cache_loc,
+                        req_to_token_pool.req_to_token.shape[1],
+                        topk,
+                        num_steps,
+                    )
+                else:
+                    # FIXME(lsyin): align with the default code path
+                    assign_draft_cache_locs_contiguous[(bs,)](
+                        batch.req_pool_indices,
+                        req_to_token_pool.req_to_token,
+                        batch.seq_lens,
+                        batch.out_cache_loc,
+                        req_to_token_pool.req_to_token.shape[1],
+                        topk,
+                        num_steps,
+                    )
             else:
                 # page_size > 1 + topk > 1: per-branch page-aligned draft pages.
                 # Reduce out_cache_loc from the page-aligned tree region down to the

python/sglang/srt/speculative/draft_utils.py

@@ -1,7 +1,14 @@
 import logging
 
 from sglang.srt.server_args import ServerArgs
-from sglang.srt.utils.common import is_blackwell, is_hip, is_musa, is_npu
+from sglang.srt.utils.common import (
+    cpu_has_amx_support,
+    is_blackwell,
+    is_cpu,
+    is_hip,
+    is_musa,
+    is_npu,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -44,13 +51,10 @@ class DraftBackendFactory:
         backend_map = {
             "flashinfer": self._create_flashinfer_decode_backend,
             "triton": self._create_triton_decode_backend,
+            "intel_amx": self._create_intel_amx_decode_backend,
             "aiter": self._create_aiter_decode_backend,
             "fa3": self._create_fa3_decode_backend,
-            "hybrid_linear_attn": (
-                self._create_fa3_decode_backend
-                if not is_blackwell()
-                else self._create_triton_decode_backend
-            ),
+            "hybrid_linear_attn": self._create_hybrid_linear_attn_decode_backend,
             "flashmla": self._create_flashmla_decode_backend,
             "trtllm_mha": self._create_trtllm_mha_decode_backend,
             "trtllm_mla": self._create_trtllm_mla_decode_backend,
@@ -73,13 +77,10 @@ class DraftBackendFactory:
         backend_map = {
             "flashinfer": self._create_flashinfer_prefill_backend,
             "triton": self._create_triton_prefill_backend,
+            "intel_amx": self._create_intel_amx_prefill_backend,
             "aiter": self._create_aiter_prefill_backend,
             "fa3": self._create_fa3_prefill_backend,
-            "hybrid_linear_attn": (
-                self._create_fa3_prefill_backend
-                if not is_blackwell()
-                else self._create_triton_prefill_backend
-            ),
+            "hybrid_linear_attn": self._create_hybrid_linear_attn_prefill_backend,
             "flashmla": self._create_flashmla_prefill_backend,
             "trtllm_mha": self._create_trtllm_mha_prefill_backend,
             "trtllm_mla": self._create_trtllm_mla_prefill_backend,
@@ -144,6 +145,29 @@ class DraftBackendFactory:
             self.draft_model_runner, self.topk, self.speculative_num_steps
         )
 
+    def _create_intel_amx_decode_backend(self):
+        from sglang.srt.layers.attention.intel_amx_backend import (
+            IntelAMXMultiStepDraftBackend,
+        )
+
+        return IntelAMXMultiStepDraftBackend(
+            self.draft_model_runner, self.topk, self.speculative_num_steps
+        )
+
+    def _create_hybrid_linear_attn_decode_backend(self):
+        if is_cpu() and cpu_has_amx_support():
+            return self._create_intel_amx_decode_backend()
+        if is_blackwell():
+            return self._create_triton_decode_backend()
+        return self._create_fa3_decode_backend()
+
+    def _create_hybrid_linear_attn_prefill_backend(self):
+        if is_cpu() and cpu_has_amx_support():
+            return self._create_intel_amx_prefill_backend()
+        if is_blackwell():
+            return self._create_triton_prefill_backend()
+        return self._create_fa3_prefill_backend()
+
     def _create_aiter_decode_backend(self):
         from sglang.srt.layers.attention.aiter_backend import AiterMultiStepDraftBackend
 
@@ -277,6 +301,11 @@ class DraftBackendFactory:
 
         return TritonAttnBackend(self.draft_model_runner, skip_prefill=False)
 
+    def _create_intel_amx_prefill_backend(self):
+        from sglang.srt.layers.attention.intel_amx_backend import IntelAMXAttnBackend
+
+        return IntelAMXAttnBackend(self.draft_model_runner)
+
     def _create_aiter_prefill_backend(self):
         from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
 

python/sglang/srt/speculative/eagle_info.py

@@ -59,19 +59,21 @@ class EagleVerifyInput(SpecInput):
         return self.draft_token_num, self.draft_token_num
 
     @classmethod
-    def create_idle_input(cls, topk: int, spec_steps: int, num_verify_tokens: int):
+    def create_idle_input(
+        cls, topk: int, spec_steps: int, num_verify_tokens: int, device: str
+    ):
         return cls(
-            draft_token=torch.empty((0,), dtype=torch.long, device="cuda"),
-            custom_mask=torch.full((0,), True, dtype=torch.bool, device="cuda"),
-            positions=torch.empty((0,), dtype=torch.int64, device="cuda"),
+            draft_token=torch.empty((0,), dtype=torch.long, device=device),
+            custom_mask=torch.full((0,), True, dtype=torch.bool, device=device),
+            positions=torch.empty((0,), dtype=torch.int64, device=device),
             retrieve_index=torch.full(
-                (0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
             ),
             retrieve_next_token=torch.full(
-                (0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
             ),
             retrieve_next_sibling=torch.full(
-                (0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
+                (0, num_verify_tokens), -1, dtype=torch.long, device=device
             ),
             retrieve_cum_len=None,
             topk=topk,

python/sglang/srt/speculative/eagle_utils.py

@@ -25,6 +25,7 @@ from sglang.srt.speculative.triton_ops.spec_tree import (
     verify_tree_greedy_kernel_triton,
 )
 from sglang.srt.utils import (
+    is_cpu,
     is_cuda,
     is_hip,
     is_musa,
@@ -46,6 +47,7 @@ _is_hip = is_hip()
 _is_npu = is_npu()
 _is_musa = is_musa()
 _is_xpu = is_xpu()
+_is_cpu = is_cpu()
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +55,11 @@ if _is_cuda or _is_hip or _is_musa:
     from sgl_kernel import (
         build_tree_kernel_efficient as sgl_build_tree_kernel_efficient,
     )
+elif _is_cpu:
+    from sgl_kernel import (
+        build_tree_kernel_efficient_cpu as sgl_build_tree_kernel_efficient_cpu,
+    )
+    from sgl_kernel import verify_tree_greedy_cpu as sgl_verify_tree_greedy_cpu
 
 
 ALLOC_EXTEND_FUNCS = defaultdict(
@@ -139,6 +146,12 @@ class TreeMaskMode(IntEnum):
     QLEN_ONLY_BITPACKING = 2
 
 
+def default_tree_mask_mode() -> TreeMaskMode:
+    # The CPU verify attention kernel (intel_amx) consumes the qlen x qlen
+    # QLEN_ONLY tree mask directly; FULL_MASK is for the GPU kernels.
+    return TreeMaskMode.QLEN_ONLY if _is_cpu else TreeMaskMode.FULL_MASK
+
+
 def build_tree_kernel_efficient(
     bonus_tokens: torch.Tensor,
     parent_list: List[torch.Tensor],
@@ -243,6 +256,21 @@ def build_tree_kernel_efficient(
             num_verify_tokens,
             tree_mask_mode,
         )
+    elif _is_cpu:
+        sgl_build_tree_kernel_efficient_cpu(
+            parent_list,
+            top_scores_index,
+            seq_lens,
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            topk,
+            spec_steps,
+            num_verify_tokens,
+            tree_mask_mode,
+        )
     else:
         sgl_build_tree_kernel_efficient(
             parent_list,
@@ -376,6 +404,20 @@ def verify_tree_greedy_func(
             target_predict=target_predict,
         )
 
+    elif _is_cpu:
+        sgl_verify_tree_greedy_cpu(
+            predicts=predicts,  # mutable
+            accept_index=accept_index,  # mutable
+            accept_token_num=accept_token_num,  # mutable
+            candidates=candidates,
+            # kwarg LHS retained as `retrive_*` to match the CUDA op schema, so
+            # the CPU/CUDA call sites stay grep-symmetric.
+            retrive_index=retrieve_index,
+            retrive_next_token=retrieve_next_token,
+            retrive_next_sibling=retrieve_next_sibling,
+            target_predict=target_predict,
+        )
+
     elif _is_npu:
         from sgl_kernel_npu.sample.verify_tree_greedy import verify_tree_greedy
 
@@ -617,7 +659,7 @@ def eagle_sample(
 
     # Sample tokens
     target_predict = None
-    if sampling_info.is_all_greedy or _is_npu or _is_hip or _is_xpu:
+    if sampling_info.is_all_greedy or _is_cpu or _is_npu or _is_hip or _is_xpu:
         target_predict = torch.argmax(next_token_logits, dim=-1)
         target_predict = target_predict.reshape(bs, verify_input.draft_token_num)
         predict, accept_index, num_correct_drafts = verify_tree_greedy_func(

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -66,9 +66,9 @@ from sglang.srt.speculative.eagle_info import (
     EagleVerifyInput,
 )
 from sglang.srt.speculative.eagle_utils import (
-    TreeMaskMode,
     _eagle_prefill_tail_tokens,
     build_tree_kernel_efficient,
+    default_tree_mask_mode,
     eagle_prepare_for_verify,
     eagle_sample,
     get_draft_recurrent_hidden_state_spec,
@@ -90,7 +90,7 @@ from sglang.srt.speculative.spec_utils import (
     select_top_k_tokens,
     spec_stage_span,
 )
-from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens
+from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens_func
 from sglang.srt.utils.async_probe import (
     maybe_detect_inf,
     maybe_detect_nan,
@@ -101,6 +101,7 @@ from sglang.srt.utils.common import (
     empty_context,
     fast_topk,
     get_available_gpu_memory,
+    is_cpu,
     is_cuda,
     is_hip,
     is_musa,
@@ -110,12 +111,14 @@ from sglang.srt.utils.common import (
 )
 from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 
+_is_cpu = is_cpu()
 _is_npu = is_npu()
 _is_cuda = is_cuda()
 _is_musa = is_musa()
 _is_hip = is_hip()
 _is_xpu = is_xpu()
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -199,7 +202,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        self.tree_mask_mode = TreeMaskMode.FULL_MASK
+        self.tree_mask_mode = default_tree_mask_mode()
 
         self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
 
@@ -393,14 +396,14 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         self.draft_runner.draft_attn_backend = self.draft_attn_backend
         if self.draft_extend_attn_backend is not None:
             self.draft_runner.attn_backend = self.draft_extend_attn_backend
-        self.tree_mask_mode = TreeMaskMode.FULL_MASK
+        self.tree_mask_mode = default_tree_mask_mode()
 
     def _capture_cuda_graphs(self):
         """Capture the draft worker's own cuda graphs (decode + draft-extend)."""
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
+        if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
             return
 
         if self.server_args.model_impl == "mindspore":
@@ -552,6 +555,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
                 self.topk,
                 self.speculative_num_steps,
                 self.speculative_num_draft_tokens,
+                self.device,
             )
 
         # Build tree mask
@@ -1275,7 +1279,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         """
         if batch.forward_mode.is_idle():
             return EagleVerifyInput.create_idle_input(
-                topk=self.topk, spec_steps=0, num_verify_tokens=1
+                topk=self.topk, spec_steps=0, num_verify_tokens=1, device=self.device
             )
 
         draft_input: EagleDraftInput = batch.spec_info
@@ -1673,11 +1677,12 @@ class EAGLEWorkerV2(BaseSpecWorker):
             bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
             # stride = accept_tokens per-req width = accept_index.shape[1]
             # (spec_steps + 1); NOT num_draft_tokens, wrong for topk > 1 trees.
-            fill_bonus_tokens[(bs,)](
+            fill_bonus_tokens_func(
                 accept_tokens,
                 accept_lens,
                 bonus_tokens,
                 accept_index.shape[1],
+                bs,
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -435,7 +435,7 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
         for step in range(num_steps):
             _, topk_p, topk_index = runner.replay(step)
             if step < num_steps - 1:
-                rotate_input_ids_triton(...)  # advance the draft chain
+                rotate_input_ids(...)  # advance the draft chain
 
     Not itself a DecodeCudaGraphRunner -- it only routes work to the per-step
     runners.

python/sglang/srt/speculative/multi_layer_eagle_utils.py

@@ -13,11 +13,11 @@
 # ==============================================================================
 
 from sglang.srt.speculative.triton_ops.multi_layer_eagle import (
+    rotate_input_ids,
     rotate_input_ids_kernel,
-    rotate_input_ids_triton,
 )
 
 __all__ = [
+    "rotate_input_ids",
     "rotate_input_ids_kernel",
-    "rotate_input_ids_triton",
 ]

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -49,8 +49,8 @@ from sglang.srt.speculative.eagle_info import (
     EagleVerifyInput,
 )
 from sglang.srt.speculative.eagle_utils import (
-    TreeMaskMode,
     build_tree_kernel_efficient,
+    default_tree_mask_mode,
     eagle_prepare_for_verify,
     eagle_sample,
     get_draft_recurrent_hidden_state_spec,
@@ -58,7 +58,7 @@ from sglang.srt.speculative.eagle_utils import (
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
 )
-from sglang.srt.speculative.multi_layer_eagle_utils import rotate_input_ids_triton
+from sglang.srt.speculative.multi_layer_eagle_utils import rotate_input_ids
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import (
     draft_tp_context,
@@ -67,8 +67,8 @@ from sglang.srt.speculative.spec_utils import (
     sample_draft_proposal,
     select_top_k_tokens,
 )
-from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens
-from sglang.srt.utils import is_npu
+from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens_func
+from sglang.srt.utils import is_cpu, is_npu
 from sglang.srt.utils.async_probe import (
     maybe_detect_inf,
     maybe_detect_nan,
@@ -77,6 +77,8 @@ from sglang.srt.utils.async_probe import (
 from sglang.srt.utils.common import empty_context, fast_topk
 
 _is_npu = is_npu()
+_is_cpu = is_cpu()
+
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput
@@ -168,7 +170,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        self.tree_mask_mode = TreeMaskMode.FULL_MASK
+        self.tree_mask_mode = default_tree_mask_mode()
         self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
 
     def alloc_memory_pool(
@@ -232,7 +234,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
+        if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
             return
 
         if not _is_npu:
@@ -264,6 +266,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 self.topk,
                 self.speculative_num_steps,
                 self.speculative_num_draft_tokens,
+                self.device,
             )
 
         # Build tree mask
@@ -348,7 +351,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                             (tree_info[2].size(0), 1),
                             i,
                             dtype=torch.long,
-                            device="cuda",
+                            device=tree_info[2].device,
                         )
                     )
 
@@ -435,7 +438,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         # Construct input_ids
         # TODO: same chunked-prefill chain divergence as PR #26329.
         if not batch.forward_mode.is_idle():
-            rotate_input_ids_triton(
+            rotate_input_ids(
                 forward_batch.input_ids,
                 forward_batch.extend_start_loc,
                 forward_batch.extend_seq_lens,
@@ -479,7 +482,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                     output.logits_output.hidden_states
                 )
             if forward_batch.extend_seq_lens is not None:
-                rotate_input_ids_triton(
+                rotate_input_ids(
                     forward_batch.input_ids,
                     forward_batch.extend_start_loc,
                     forward_batch.extend_seq_lens,
@@ -568,7 +571,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 # Advance the draft chain by rotating the shared input_ids window
                 # in place; step N+1's graph then reads the rotated values.
                 if step < self.speculative_num_steps - 1:
-                    rotate_input_ids_triton(
+                    rotate_input_ids(
                         cgr.buffers.input_ids[: cgr.raw_num_tokens],
                         cgr.buffers.extend_start_loc[: cgr.raw_bs],
                         cgr.buffers.extend_seq_lens[: cgr.raw_bs],
@@ -621,7 +624,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                         draft_logits_output.logits_output.hidden_states
                     )
                 if forward_batch.extend_seq_lens is not None:
-                    rotate_input_ids_triton(
+                    rotate_input_ids(
                         forward_batch.input_ids,
                         forward_batch.extend_start_loc,
                         forward_batch.extend_seq_lens,
@@ -869,11 +872,12 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
             accept_tokens = predict[accept_index]
             bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32)
             # stride = accept_tokens per-req width = accept_index.shape[1].
-            fill_bonus_tokens[(bs,)](
+            fill_bonus_tokens_func(
                 accept_tokens,
                 accept_lens,
                 bonus_tokens,
                 accept_index.shape[1],
+                bs,
             )
         else:
             bonus_tokens = torch.empty((0,), device=self.device, dtype=torch.int32)

python/sglang/srt/speculative/ngram_worker.py

@@ -26,8 +26,11 @@ from sglang.srt.speculative.spec_utils import (
 from sglang.srt.speculative.triton_ops.cache_locs import (
     assign_extend_cache_locs_func as assign_extend_cache_locs_func,
 )
+from sglang.srt.utils import is_cpu
 from sglang.srt.utils.async_probe import maybe_detect_inf, maybe_detect_nan
 
+_is_cpu = is_cpu()
+
 logger = logging.getLogger(__name__)
 
 
@@ -68,7 +71,7 @@ class NGRAMWorker(BaseSpecWorker):
         self.speculative_num_steps = server_args.speculative_num_steps
         # req_to_token_pool / token_to_kv_pool_allocator are set in
         # alloc_memory_pool(), after the target pools are allocated.
-        self.device = f"cuda:{gpu_id}" if gpu_id >= 0 else "cuda"
+        self.device = server_args.device
 
         self.adaptive_controller = None
         # rids of the last decode batch; used to erase corpus match state for
@@ -298,7 +301,7 @@ class NGRAMWorker(BaseSpecWorker):
 
         # NOTE: QLEN_MASK is faster than FULL_MASK, but requires corresponding changes in flashinfer.
         # Testing shows about 8% performance improvement (the effect is roughly proportional to batch size).
-        if USE_FULL_MASK:
+        if USE_FULL_MASK and not _is_cpu:
             tree_mask = []
             mask = mask.reshape(bs, self.draft_token_num, self.draft_token_num)
             # TODO(siyuan): the for loop here leads to significant overhead in large batch size. Can be written into a kernel.

python/sglang/srt/speculative/spec_utils.py

@@ -41,9 +41,17 @@ from sglang.srt.speculative.triton_ops.cache_locs import (
     get_target_cache_loc as get_target_cache_loc,
 )
 from sglang.srt.speculative.triton_ops.eagle import (
-    fill_accept_out_cache_loc as fill_accept_out_cache_loc,
+    fill_accept_out_cache_loc_func as fill_accept_out_cache_loc_func,
+)
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+    next_power_of_2,
 )
-from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, is_xpu, next_power_of_2
 from sglang.srt.utils.async_probe import maybe_detect_oob
 from sglang.srt.utils.nvtx_utils import profile_range
 
@@ -52,6 +60,7 @@ _is_hip = is_hip()
 _is_npu = is_npu()
 _is_musa = is_musa()
 _is_xpu = is_xpu()
+_is_cpu = is_cpu()
 
 if TYPE_CHECKING:
     from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
@@ -69,6 +78,9 @@ elif _is_hip:
 else:
     from sglang.srt.utils.common import fast_topk
 
+if _is_cpu:
+    from sgl_kernel import assign_extend_cache_locs_cpu
+
 
 logger = logging.getLogger(__name__)
 
@@ -347,7 +359,7 @@ def generate_simulated_accept_index(
 
     accept_indx_first_col = accept_index[:, 0].view(-1, 1)
     sim_accept_index = torch.full(
-        (bs, spec_steps + 1), -1, dtype=torch.int32, device="cuda"
+        (bs, spec_steps + 1), -1, dtype=torch.int32, device=accept_index.device
     )
     sim_accept_index[:, :simulate_acc_len] = accept_indx_first_col + torch.arange(
         simulate_acc_len, device=accept_index.device
@@ -571,20 +583,30 @@ def move_accept_tokens_to_target_kvcache(
         device=device,
     )
     accept_out_cache_loc = torch.zeros(size, dtype=torch.int64, device=device)
-    assign_extend_cache_locs[(bs,)](
-        batch.req_pool_indices,
-        batch.req_to_token_pool.req_to_token,
-        batch.seq_lens,
-        batch.seq_lens + num_correct_drafts + 1,
-        tgt_cache_loc,
-        batch.req_to_token_pool.req_to_token.shape[1],
-        next_power_of_2(bs),
-    )
-    fill_accept_out_cache_loc[(size,)](
+    if _is_cpu:
+        assign_extend_cache_locs_cpu(
+            batch.req_pool_indices,
+            batch.req_to_token_pool.req_to_token,
+            batch.seq_lens,
+            batch.seq_lens + num_correct_drafts + 1,
+            tgt_cache_loc,
+            batch.req_to_token_pool.req_to_token.shape[1],
+        )
+    else:
+        assign_extend_cache_locs[(bs,)](
+            batch.req_pool_indices,
+            batch.req_to_token_pool.req_to_token,
+            batch.seq_lens,
+            batch.seq_lens + num_correct_drafts + 1,
+            tgt_cache_loc,
+            batch.req_to_token_pool.req_to_token.shape[1],
+            next_power_of_2(bs),
+        )
+    fill_accept_out_cache_loc_func(
         accept_index,
         batch.out_cache_loc,
         accept_out_cache_loc,
-        next_power_of_2(size),
+        size,
     )
     token_to_kv_pool_allocator.get_kvcache().move_kv_cache(
         tgt_cache_loc, accept_out_cache_loc

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -11,7 +11,7 @@ from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.adaptive_runtime_state import (
     AdaptiveController,
 )
-from sglang.srt.speculative.eagle_utils import TreeMaskMode
+from sglang.srt.speculative.eagle_utils import default_tree_mask_mode
 from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import draft_tp_context
@@ -102,7 +102,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        self.tree_mask_mode = TreeMaskMode.FULL_MASK
+        self.tree_mask_mode = default_tree_mask_mode()
         self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
         # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
         self.index_share_for_mtp_iteration = (

python/sglang/srt/speculative/triton_ops/cache_locs.py

@@ -4,14 +4,26 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu, is_xpu, next_power_of_2
+from sglang.srt.utils import (
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+    next_power_of_2,
+)
 
+_is_cpu = is_cpu()
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_npu = is_npu()
 _is_musa = is_musa()
 _is_xpu = is_xpu()
 
+if _is_cpu:
+    from sgl_kernel import assign_extend_cache_locs_cpu, assign_req_to_token_pool_cpu
+
 
 @triton.jit
 def assign_req_to_token_pool(
@@ -56,6 +68,16 @@ def assign_req_to_token_pool_func(
     out_cache_loc: torch.Tensor,
     batch_size: int,
 ):
+    if _is_cpu:
+        assign_req_to_token_pool_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+        return
     assign_req_to_token_pool[(batch_size,)](
         req_pool_indices,
         req_to_token,
@@ -377,3 +399,20 @@ def assign_extend_cache_locs_func(
         )
 
         return out_cache_loc
+
+    elif _is_cpu:
+        out_cache_loc = torch.empty(
+            (batch_size * draft_token_num,),
+            dtype=torch.int64,
+            device=device,
+        )
+        assign_extend_cache_locs_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+        )
+
+        return out_cache_loc

python/sglang/srt/speculative/triton_ops/eagle.py

@@ -1,6 +1,14 @@
+import torch
 import triton
 import triton.language as tl
 
+from sglang.srt.utils import is_cpu, next_power_of_2
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import fill_accept_out_cache_loc_cpu, fill_bonus_tokens_cpu
+
 
 @triton.jit
 def fill_bonus_tokens(
@@ -21,6 +29,29 @@ def fill_bonus_tokens(
     tl.store(bonus_tokens_ptr + pid, bonus_token)
 
 
+def fill_bonus_tokens_func(
+    accept_tokens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    bonus_tokens: torch.Tensor,  # mutable
+    accept_stride: int,
+    batch_size: int,
+):
+    if _is_cpu:
+        fill_bonus_tokens_cpu(
+            accept_tokens,
+            accept_lens,
+            bonus_tokens,
+            accept_stride,
+        )
+        return
+    fill_bonus_tokens[(batch_size,)](
+        accept_tokens,
+        accept_lens,
+        bonus_tokens,
+        accept_stride,
+    )
+
+
 @triton.jit
 def fill_accept_out_cache_loc(
     accept_index,
@@ -37,3 +68,24 @@ def fill_accept_out_cache_loc(
     if src > -1:
         value = tl.load(out_cache_loc + src)
         tl.store(accept_out_cache_loc + dst, value)
+
+
+def fill_accept_out_cache_loc_func(
+    accept_index: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_out_cache_loc: torch.Tensor,  # mutable
+    size: int,
+):
+    if _is_cpu:
+        fill_accept_out_cache_loc_cpu(
+            accept_index,
+            out_cache_loc,
+            accept_out_cache_loc,
+        )
+        return
+    fill_accept_out_cache_loc[(size,)](
+        accept_index,
+        out_cache_loc,
+        accept_out_cache_loc,
+        next_power_of_2(size),
+    )

python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py

@@ -15,6 +15,13 @@
 import triton
 import triton.language as tl
 
+from sglang.srt.utils import is_cpu
+
+_is_cpu = is_cpu()
+
+if _is_cpu:
+    from sgl_kernel import rotate_input_ids_cpu
+
 
 @triton.jit
 def rotate_input_ids_kernel(
@@ -53,9 +60,19 @@ def rotate_input_ids_kernel(
         tl.store(last_pos_ptr, new_token)
 
 
-def rotate_input_ids_triton(
+def rotate_input_ids(
     input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None
 ):
+    if _is_cpu:
+        rotate_input_ids_cpu(
+            input_ids,
+            extend_start_loc,
+            extend_seq_lens,
+            topk_index,
+            select_index,
+        )
+        return input_ids
+
     batch_size = extend_seq_lens.shape[0]
     BLOCK_SIZE = 4096 if select_index is not None else 8
     grid = (batch_size,)