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,)

sgl-kernel/csrc/cpu/extend.cpp

@@ -9,6 +9,17 @@ namespace {
 //   2. can handle non-contiguous k_extend and v_extend
 //   3. computes attention for prefix and extend separately
 //   4. TODO: apply head dimension blocking to optimize GQA
+//   5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1):
+//      `tree_mask` is a flat [batches * qlen * qlen] bool tensor in
+//      TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] ==
+//      max_len_extend (uniform across the batch, equal to draft_token_num).
+//      Row i = query draft token, column j = key draft token; true means query i
+//      may attend key j (each row marks self + ancestors + root). The committed
+//      prefix (stage 1) is implicitly fully visible to every draft token, which
+//      is why the mask only covers the qlen x qlen new-token block; the GPU
+//      FULL_MASK layout carries the prefix columns explicitly but they are
+//      all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the
+//      plain causal mask (correct for non-spec extend and topk == 1 chains).
 //
 
 template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N>
@@ -27,6 +38,7 @@ void extend_attention_kernel_impl(
     const index_t* __restrict__ extend_start_loc,
     const void* __restrict__ buffer,
     const scalar_t* __restrict__ sinks,
+    const bool* __restrict__ tree_mask,
     int batches,
     int num_heads,
     int num_heads_kv,
@@ -109,6 +121,16 @@ void extend_attention_kernel_impl(
         TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix);
       }
 
+      if (tree_mask != nullptr) {
+        // QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY)
+        TORCH_CHECK(
+            seq_len_extend == max_len_extend,
+            "extend attention: tree_mask requires uniform extend_seq_lens, got ",
+            seq_len_extend,
+            " vs ",
+            max_len_extend);
+      }
+
       // offset and size in MB
       int m = mb * BLOCK_M;
       int m_size = std::min(BLOCK_M, seq_len_extend - m);
@@ -223,18 +245,38 @@ void extend_attention_kernel_impl(
               /* B     */ Btmp,
               /* C     */ s_i);
 
-          // apply causal mask
-          // [Note] condition to apply causal mask.
-          // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
-          // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
-          // when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.  With
-          // BLOCK_M=512, BLOCK_N=768:
-          //   BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
-          //   Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
-          //   Query row=0 is at position 512, so keys 513..767 are future and must be
-          //   masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
-          //   producing wrong (non-causal) attention for rows 0..254 of this m-block.
-          if (n + n_size - 1 > m) {
+          // apply tree mask (speculative TARGET_VERIFY) or causal mask
+          if (tree_mask != nullptr) {
+            // [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY).
+            // mask[bs][m + row][n + col] == false -> query draft token (m + row)
+            // may not attend key draft token (n + col); set the score to -inf
+            // before softmax. The tree mask subsumes the causal constraint:
+            // ancestors always precede descendants in the draft token ordering,
+            // so permitted keys satisfy j <= i and the causal `num_keys` bound
+            // above remains valid.
+            const bool* __restrict__ mask_base =
+                tree_mask + (static_cast<int64_t>(bs) * seq_len_extend + m) * seq_len_extend + n;
+            for (int row = 0; row < m_size; ++row) {
+              float* __restrict__ row_ptr = s_i + row * BLOCK_N;
+              const bool* __restrict__ mask_ptr = mask_base + static_cast<int64_t>(row) * seq_len_extend;
+              for (int col = 0; col < n_size; ++col) {
+                if (!mask_ptr[col]) {
+                  row_ptr[col] = -std::numeric_limits<float>::infinity();
+                }
+              }
+            }
+          } else if (n + n_size - 1 > m) {
+            // apply causal mask
+            // [Note] condition to apply causal mask.
+            // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
+            // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was
+            // correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.
+            // With BLOCK_M=512, BLOCK_N=768:
+            //   BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
+            //   Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
+            //   Query row=0 is at position 512, so keys 513..767 are future and must be
+            //   masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
+            //   producing wrong (non-causal) attention for rows 0..254 of this m-block.
             for (int row = 0; row < m_size; ++row) {
               int last_col = m + row - n;
               // [Note] mask the entire row if last_col < 0.
@@ -333,6 +375,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
         extend_start_loc.data_ptr<index_t>(),                                              \
         buffer.data_ptr(),                                                                 \
         sinks_tensor.data_ptr<scalar_t>(),                                                 \
+        tree_mask_ptr,                                                                     \
         num_seqs,                                                                          \
         num_heads,                                                                         \
         num_heads_kv,                                                                      \
@@ -377,6 +420,8 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
 // extend_start_loc: [num_seqs]
 // encoder_lens: [num_seqs] int64 or None
 // sinks: [num_heads] or None
+// tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None
+//   TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above.
 void extend_attention_cpu(
     at::Tensor& q_extend,
     const std::optional<at::Tensor>& k_extend_opt,
@@ -395,7 +440,8 @@ void extend_attention_cpu(
     bool is_cross_attn,
     int64_t sliding_window_size,
     std::optional<at::Tensor> encoder_lens,
-    std::optional<at::Tensor> sinks) {
+    std::optional<at::Tensor> sinks,
+    std::optional<at::Tensor> tree_mask) {
   if (!is_cross_attn) {
     TORCH_CHECK(
         k_extend_opt.has_value() && v_extend_opt.has_value(),
@@ -481,6 +527,26 @@ void extend_attention_cpu(
   CHECK_DIM(1, sinks_tensor);
   CHECK_EQ(sinks_tensor.size(0), num_heads);
 
+  const bool* tree_mask_ptr = nullptr;
+  if (tree_mask.has_value()) {
+    const at::Tensor& tree_mask_t = tree_mask.value();
+    CHECK_INPUT(tree_mask_t);
+    TORCH_CHECK(
+        tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type());
+    TORCH_CHECK(
+        tree_mask_t.numel() == static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
+        "extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ",
+        static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend,
+        ", got ",
+        tree_mask_t.numel());
+    TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention");
+    // The window mask derives query positions from the row index
+    // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth,
+    // which is <= the row index; combining the two would over-mask the prefix.
+    TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention");
+    tree_mask_ptr = tree_mask_t.data_ptr<bool>();
+  }
+
   AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
     AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
       if (max_len_extend <= 256) {

sgl-kernel/csrc/cpu/kvcache.cpp

@@ -128,3 +128,57 @@ void store_cache_cpu(
     });
   });
 }
+
+// CPU counterpart of the Triton kernel `copy_all_layer_kv_cache_tiled`:
+// for every K/V buffer b, copy the slot rows `src_loc` to `tgt_loc`:
+//   buf_b[tgt_loc[i], :] = buf_b[src_loc[i], :]  for i in [0, num_locs)
+//
+//   data_ptrs : [2 * layer_num] uint64; base address of each K/V buffer
+//   strides   : [2 * layer_num] int64; bytes per slot row of each buffer
+//   tgt_loc   : [num_locs] int64/int32 slot indices
+//   src_loc   : [num_locs] int64/int32 slot indices
+//
+// Like the Triton kernel, the copy is safe when tgt_loc and src_loc overlap
+// arbitrarily: all source rows of a buffer are staged before any target row
+// of that buffer is written (gather then scatter).
+void copy_all_layer_kv_cache_cpu(
+    const at::Tensor& data_ptrs, const at::Tensor& strides, const at::Tensor& tgt_loc, const at::Tensor& src_loc) {
+  CHECK_INPUT(data_ptrs);
+  CHECK_INPUT(strides);
+  CHECK_INPUT(tgt_loc);
+  CHECK_INPUT(src_loc);
+  CHECK_EQ(data_ptrs.scalar_type(), at::kUInt64);
+  CHECK_EQ(strides.scalar_type(), at::kLong);
+  CHECK_EQ(tgt_loc.scalar_type(), src_loc.scalar_type());
+
+  int64_t num_bufs = data_ptrs.numel();
+  CHECK_EQ(strides.numel(), num_bufs);
+  int64_t num_locs = tgt_loc.numel();
+  CHECK_EQ(src_loc.numel(), num_locs);
+  if (num_bufs == 0 || num_locs == 0) {
+    return;
+  }
+
+  const uint64_t* __restrict__ ptrs = reinterpret_cast<const uint64_t*>(data_ptrs.data_ptr());
+  const int64_t* __restrict__ stride_ptr = strides.data_ptr<int64_t>();
+
+  AT_DISPATCH_INDEX_TYPES(tgt_loc.scalar_type(), "copy_all_layer_kv_cache_cpu", [&] {
+    const index_t* __restrict__ tgt_ptr = tgt_loc.data_ptr<index_t>();
+    const index_t* __restrict__ src_ptr = src_loc.data_ptr<index_t>();
+
+    at::parallel_for(0, num_bufs, 0, [&](int64_t begin, int64_t end) {
+      std::vector<uint8_t> staging;
+      for (int64_t b = begin; b < end; ++b) {
+        uint8_t* base = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ptrs[b]));
+        const int64_t stride = stride_ptr[b];
+        staging.resize(num_locs * stride);
+        for (int64_t i = 0; i < num_locs; ++i) {
+          std::memcpy(staging.data() + i * stride, base + src_ptr[i] * stride, stride);
+        }
+        for (int64_t i = 0; i < num_locs; ++i) {
+          std::memcpy(base + tgt_ptr[i] * stride, staging.data() + i * stride, stride);
+        }
+      }
+    });
+  });
+}

sgl-kernel/csrc/cpu/spec.cpp

@@ -0,0 +1,854 @@
+#include "common.h"
+
+namespace {
+
+// Contract shared by every kernel in this file: all tensors are dense,
+// contiguous CPU tensors (checked below), so strides are the canonical
+// row-major ones; per-function comments list shapes and dtypes only.
+// `index_t` params accept int32 or int64 via AT_DISPATCH_INDEX_TYPES so
+// callers never pay a dtype-conversion copy.
+
+template <typename rpi_t, typename off_t>
+void assign_req_to_token_pool_kernel_impl(
+    const rpi_t* __restrict__ req_pool_indices,
+    int32_t* __restrict__ req_to_token,
+    const off_t* __restrict__ start_offset,
+    const off_t* __restrict__ end_offset,
+    const int64_t* __restrict__ out_cache_loc,
+    int64_t num_cache_locs,
+    int64_t batch_size,
+    int64_t pool_len) {
+  // Pre-compute exclusive prefix sum of (end - start) to avoid O(N^2) work.
+  std::vector<int64_t> prefix(batch_size + 1, 0);
+  for (int64_t i = 0; i < batch_size; ++i) {
+    prefix[i + 1] = prefix[i] + (end_offset[i] - start_offset[i]);
+  }
+  TORCH_CHECK(
+      prefix[batch_size] <= num_cache_locs,
+      "assign_req_to_token_pool: out_cache_loc has ",
+      num_cache_locs,
+      " entries but offsets require ",
+      prefix[batch_size]);
+
+  at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t pid = begin; pid < end; ++pid) {
+      int64_t kv_start = start_offset[pid];
+      int64_t kv_end = end_offset[pid];
+      int32_t* token_pool = req_to_token + req_pool_indices[pid] * pool_len;
+      int64_t out_offset = prefix[pid];
+
+      for (int64_t j = kv_start; j < kv_end; ++j) {
+        token_pool[j] = static_cast<int32_t>(out_cache_loc[out_offset + (j - kv_start)]);
+      }
+    }
+  });
+}
+
+template <typename index_t>
+void verify_tree_greedy_kernel_impl(
+    int32_t* __restrict__ predicts,
+    int32_t* __restrict__ accept_index,
+    int32_t* __restrict__ accept_token_num,
+    const index_t* __restrict__ candidates,
+    const index_t* __restrict__ retrive_index,
+    const index_t* __restrict__ retrive_next_token,
+    const index_t* __restrict__ retrive_next_sibling,
+    const index_t* __restrict__ target_predict,
+    int64_t batch_size,
+    int64_t num_spec_step,
+    int64_t num_draft_tokens) {
+  at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t bx = begin; bx < end; ++bx) {
+      int64_t off = bx * num_draft_tokens;
+      int64_t ai_off = bx * num_spec_step;
+
+      int64_t last_accept_index = retrive_index[off];  // retrive_index[bx, 0]
+      accept_index[ai_off] = static_cast<int32_t>(last_accept_index);
+
+      int32_t num_correct_drafts = 0;
+      int64_t cur = 0;
+
+      for (int64_t j = 1; j < num_spec_step; ++j) {
+        cur = retrive_next_token[off + cur];  // move to next token
+        while (cur != -1) {
+          int64_t draft_idx = retrive_index[off + cur];
+          int64_t draft_tok = candidates[off + cur];
+          int64_t target_tok = target_predict[last_accept_index];
+          if (draft_tok == target_tok) {
+            predicts[last_accept_index] = static_cast<int32_t>(target_tok);
+            ++num_correct_drafts;
+            accept_index[ai_off + num_correct_drafts] = static_cast<int32_t>(draft_idx);
+            last_accept_index = draft_idx;
+            break;
+          }
+          cur = retrive_next_sibling[off + cur];  // try sibling
+        }
+        if (cur == -1) break;
+      }
+      accept_token_num[bx] = num_correct_drafts;
+      predicts[last_accept_index] = static_cast<int32_t>(target_predict[last_accept_index]);
+    }
+  });
+}
+
+// Find the node index in `selected_index[bid]` holding `token_idx`; -1 when the
+// tree is malformed and the parent is absent (callers warn and stop the walk,
+// mirroring the CUDA kernel's "invalid eagle tree" printf).
+template <typename index_t>
+int64_t
+find_parent_node(const index_t* __restrict__ selected_index, int64_t row_off, int64_t sel_stride, int64_t token_idx) {
+  for (int64_t i = 0; i < sel_stride; ++i) {
+    if (selected_index[row_off + i] == token_idx) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+template <typename index_t>
+void build_tree_kernel_efficient_impl(
+    const index_t* __restrict__ parent_list,
+    const index_t* __restrict__ selected_index,
+    const index_t* __restrict__ verified_seq_len,
+    bool* __restrict__ tree_mask,
+    index_t* __restrict__ positions,
+    index_t* __restrict__ retrive_index,
+    index_t* __restrict__ retrive_next_token,
+    index_t* __restrict__ retrive_next_sibling,
+    int64_t bs,
+    int64_t topk,
+    int64_t depth,
+    int64_t draft_token_num,
+    int64_t tree_mask_mode) {
+  int64_t parent_stride = topk * (depth - 1) + 1;
+  int64_t sel_stride = draft_token_num - 1;
+
+  // FULL_MASK row offsets depend on a prefix sum over verified_seq_len;
+  // precompute it so the batch loop can run in parallel.
+  std::vector<int64_t> mask_offsets(bs, 0);
+  if (tree_mask_mode == 0) {  // FULL_MASK
+    int64_t acc = 0;
+    for (int64_t i = 0; i < bs; ++i) {
+      mask_offsets[i] = i * draft_token_num * draft_token_num + acc;
+      acc += static_cast<int64_t>(verified_seq_len[i]) * draft_token_num;
+    }
+  }
+
+  at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t bid = begin; bid < end; ++bid) {
+      int64_t off = bid * draft_token_num;
+      int64_t sel_off = bid * sel_stride;
+      int64_t seq_len = verified_seq_len[bid];
+
+      // tid == 0 logic: build retrive_index, retrive_next_token, retrive_next_sibling
+      positions[off] = seq_len;
+      retrive_index[off] = off;  // retrive_index[bid, 0] = bid * draft_token_num
+
+      for (int64_t i = draft_token_num - 1; i > 0; --i) {
+        retrive_index[off + i] = off + i;
+        int64_t parent_tb_idx = selected_index[sel_off + i - 1] / topk;
+        int64_t parent_position = 0;
+        if (parent_tb_idx > 0) {
+          int64_t parent_token_idx = parent_list[bid * parent_stride + parent_tb_idx];
+          int64_t found = find_parent_node(selected_index, sel_off, sel_stride, parent_token_idx);
+          if (found < 0) {
+            TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, parent of node ", i, " not found");
+            continue;  // skip invalid
+          }
+          parent_position = found + 1;
+        }
+        if (retrive_next_token[off + parent_position] == -1) {
+          retrive_next_token[off + parent_position] = i;
+        } else {
+          int64_t origin = retrive_next_token[off + parent_position];
+          retrive_next_token[off + parent_position] = i;
+          retrive_next_sibling[off + i] = origin;
+        }
+      }
+
+      // Build tree_mask and positions for tid > 0
+      if (tree_mask_mode == 1) {  // QLEN_ONLY
+        int64_t mask_stride = draft_token_num;
+        for (int64_t tid = 0; tid < draft_token_num; ++tid) {
+          int64_t row_start = (off + tid) * mask_stride;
+          tree_mask[row_start] = true;  // attend to the root token (column 0)
+          for (int64_t j = 1; j < draft_token_num; ++j) {
+            tree_mask[row_start + j] = false;
+          }
+          if (tid == 0) {
+            continue;
+          }
+          int64_t position = 0;
+          int64_t cur = tid - 1;
+          // A valid root-ward walk has at most `depth` steps; the bound turns a
+          // malformed (cyclic) tree into a warning instead of a scheduler hang.
+          while (position < depth) {
+            position++;
+            tree_mask[row_start + cur + 1] = true;
+            int64_t ptb = selected_index[sel_off + cur] / topk;
+            if (ptb == 0) break;
+            int64_t tok_idx = parent_list[bid * parent_stride + ptb];
+            cur = find_parent_node(selected_index, sel_off, sel_stride, tok_idx);
+            if (cur < 0) {
+              TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, ancestor of node ", tid, " not found");
+              break;  // stop the walk on a malformed tree
+            }
+          }
+          positions[off + tid] = position + seq_len;
+        }
+      } else {  // FULL_MASK (mode 0)
+        // Full mask includes the seq_len prefix
+        int64_t seq_tree_idx = mask_offsets[bid];
+        for (int64_t tid = 0; tid < draft_token_num; ++tid) {
+          int64_t row_start = seq_tree_idx + (seq_len + draft_token_num) * tid + seq_len;
+          tree_mask[row_start] = true;  // attend to the root token (column 0)
+          for (int64_t j = 1; j < draft_token_num; ++j) {
+            tree_mask[row_start + j] = false;
+          }
+          if (tid == 0) {
+            continue;
+          }
+          int64_t position = 0;
+          int64_t cur = tid - 1;
+          // Same depth bound as the QLEN_ONLY branch above.
+          while (position < depth) {
+            position++;
+            tree_mask[row_start + cur + 1] = true;
+            int64_t ptb = selected_index[sel_off + cur] / topk;
+            if (ptb == 0) {
+              break;
+            }
+            int64_t tok_idx = parent_list[bid * parent_stride + ptb];
+            cur = find_parent_node(selected_index, sel_off, sel_stride, tok_idx);
+            if (cur < 0) {
+              TORCH_WARN("build_tree_kernel_efficient_cpu: invalid eagle tree, ancestor of node ", tid, " not found");
+              break;  // stop the walk on a malformed tree
+            }
+          }
+          positions[off + tid] = position + seq_len;
+        }
+      }
+    }
+  });
+}
+
+}  // anonymous namespace
+
+// Greedy tree verification: walk each request's draft tree, accepting the
+// longest root path whose draft tokens match the target model's argmax.
+//
+// predicts:            [bs * num_draft_tokens] int32; out, verified tokens by flat draft index
+// accept_index:        [bs, num_spec_step] int32; out, flat indices of accepted
+//                      tokens; caller pre-fills with -1 (rejected slots keep it)
+// accept_token_num:    [bs] int32; out, accepted drafts per request (bonus excluded)
+// candidates:          [bs, num_draft_tokens] int32 or int64; draft tokens
+// retrive_index:       [bs, num_draft_tokens] int32 or int64; flat index of each tree node
+// retrive_next_token:  [bs, num_draft_tokens] int32 or int64; first child, -1 = none
+// retrive_next_sibling:[bs, num_draft_tokens] int32 or int64; next sibling, -1 = none
+// target_predict:      [bs, num_draft_tokens] int32 or int64; target argmax per draft slot
+void verify_tree_greedy_cpu(
+    at::Tensor predicts,
+    at::Tensor accept_index,
+    at::Tensor accept_token_num,
+    const at::Tensor& candidates,
+    const at::Tensor& retrive_index,
+    const at::Tensor& retrive_next_token,
+    const at::Tensor& retrive_next_sibling,
+    const at::Tensor& target_predict) {
+  CHECK_INPUT(candidates);
+  CHECK_DIM(2, candidates);
+  CHECK_DIM(2, accept_index);
+
+  const auto index_dtype = retrive_index.scalar_type();
+  int64_t batch_size = candidates.size(0);
+  int64_t num_draft_tokens = candidates.size(1);
+  int64_t num_spec_step = accept_index.size(1);
+
+  CHECK_EQ(candidates.scalar_type(), index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(predicts, {batch_size * num_draft_tokens}, at::kInt);
+  CHECK_INPUT_SHAPE_DTYPE<false>(accept_index, {batch_size, num_spec_step}, at::kInt);
+  CHECK_INPUT_SHAPE_DTYPE<false>(accept_token_num, {batch_size}, at::kInt);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_index, {batch_size, num_draft_tokens}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_token, {batch_size, num_draft_tokens}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_sibling, {batch_size, num_draft_tokens}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(target_predict, {batch_size, num_draft_tokens}, index_dtype);
+
+  AT_DISPATCH_INDEX_TYPES(index_dtype, "verify_tree_greedy_indices", [&] {
+    verify_tree_greedy_kernel_impl<index_t>(
+        predicts.data_ptr<int32_t>(),
+        accept_index.data_ptr<int32_t>(),
+        accept_token_num.data_ptr<int32_t>(),
+        candidates.data_ptr<index_t>(),
+        retrive_index.data_ptr<index_t>(),
+        retrive_next_token.data_ptr<index_t>(),
+        retrive_next_sibling.data_ptr<index_t>(),
+        target_predict.data_ptr<index_t>(),
+        batch_size,
+        num_spec_step,
+        num_draft_tokens);
+  });
+}
+
+// Build the draft token tree consumed by target verify: tree attention mask,
+// per-token positions, and the retrieval linkage (index / first child /
+// next sibling) used by verify_tree_greedy.
+//
+// parent_list:         [bs, topk * (depth - 1) + 1] int32 or int64
+//                      (empty [bs, 0] when depth == 1, e.g. MTP steps=1)
+// selected_index:      [bs, draft_token_num - 1] int32 or int64
+// verified_seq_len:    [bs] int32 or int64; committed prefix length per request
+// tree_mask:           out, bool.
+//                      QLEN_ONLY: [bs * draft_token_num * draft_token_num]; rows
+//                      are fully overwritten here.
+//                      FULL_MASK: [sum_i(seq_len_i * draft_token_num) + bs * draft_token_num^2];
+//                      only each row's qlen block is written -- the caller must
+//                      pre-fill the seq_len prefix columns with true.
+// positions:           [bs * draft_token_num]; out, same dtype as parent_list
+// retrive_index:       [bs, draft_token_num]; out
+// retrive_next_token:  [bs, draft_token_num]; out, pre-filled with -1
+// retrive_next_sibling:[bs, draft_token_num]; out, pre-filled with -1
+// tree_mask_mode:      0 = FULL_MASK, 1 = QLEN_ONLY (2 = QLEN_ONLY_BITPACKING is rejected)
+void build_tree_kernel_efficient_cpu(
+    const at::Tensor& parent_list,
+    const at::Tensor& selected_index,
+    const at::Tensor& verified_seq_len,
+    at::Tensor tree_mask,
+    at::Tensor positions,
+    at::Tensor retrive_index,
+    at::Tensor retrive_next_token,
+    at::Tensor retrive_next_sibling,
+    int64_t topk,
+    int64_t depth,
+    int64_t draft_token_num,
+    int64_t tree_mask_mode) {
+  CHECK_INPUT(parent_list);
+  CHECK_DIM(2, parent_list);
+
+  // CPU workers always use FULL_MASK (0) or QLEN_ONLY (1); QLEN_ONLY_BITPACKING
+  // (2) has no CPU producer and any other value is a caller bug.
+  TORCH_CHECK(
+      tree_mask_mode == 0 || tree_mask_mode == 1,
+      "build_tree_kernel_efficient_cpu: only FULL_MASK (0) and QLEN_ONLY (1) are supported, got ",
+      tree_mask_mode);
+
+  const auto index_dtype = parent_list.scalar_type();
+  int64_t bs = parent_list.size(0);
+
+  // depth == 1 (e.g. MTP steps=1) has no non-root parents, so
+  // organize_draft_results emits an empty (bs, 0) parent_list that the kernel
+  // never indexes; only the multi-step layout is width topk*(depth-1)+1.
+  if (depth > 1) {
+    CHECK_EQ(parent_list.size(1), topk * (depth - 1) + 1);
+  }
+  CHECK_INPUT_SHAPE_DTYPE<false>(selected_index, {bs, draft_token_num - 1}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(verified_seq_len, {bs}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(positions, {bs * draft_token_num}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_index, {bs, draft_token_num}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_token, {bs, draft_token_num}, index_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(retrive_next_sibling, {bs, draft_token_num}, index_dtype);
+
+  CHECK_INPUT(tree_mask);
+  CHECK_EQ(tree_mask.scalar_type(), at::kBool);
+  if (tree_mask_mode == 1) {
+    CHECK_EQ(tree_mask.numel(), bs * draft_token_num * draft_token_num);
+  } else {
+    int64_t seq_len_sum = verified_seq_len.sum().item<int64_t>();
+    CHECK_EQ(tree_mask.numel(), (seq_len_sum + bs * draft_token_num) * draft_token_num);
+  }
+
+  AT_DISPATCH_INDEX_TYPES(index_dtype, "build_tree_kernel_efficient_indices", [&] {
+    build_tree_kernel_efficient_impl<index_t>(
+        parent_list.data_ptr<index_t>(),
+        selected_index.data_ptr<index_t>(),
+        verified_seq_len.data_ptr<index_t>(),
+        tree_mask.data_ptr<bool>(),
+        positions.data_ptr<index_t>(),
+        retrive_index.data_ptr<index_t>(),
+        retrive_next_token.data_ptr<index_t>(),
+        retrive_next_sibling.data_ptr<index_t>(),
+        bs,
+        topk,
+        depth,
+        draft_token_num,
+        tree_mask_mode);
+  });
+}
+
+// Scatter freshly allocated KV slots into the request-to-token map:
+// req_to_token[req_pool_indices[i], start_offset[i]:end_offset[i]] =
+// out_cache_loc[prefix[i]:prefix[i+1]].
+//
+// req_pool_indices: [bs] int32 or int64
+// req_to_token:     [max_num_reqs, pool_len] int32; out
+// start_offset:     [bs] int32 or int64 (independent of req_pool_indices;
+//                   eagle_prepare_for_decode passes int64 indices with int32 kv lens)
+// end_offset:       [bs] same dtype as start_offset
+// out_cache_loc:    [sum_i(end_offset[i] - start_offset[i])] int64
+void assign_req_to_token_pool_cpu(
+    const at::Tensor& req_pool_indices,
+    at::Tensor req_to_token,
+    const at::Tensor& start_offset,
+    const at::Tensor& end_offset,
+    const at::Tensor& out_cache_loc,
+    int64_t pool_len) {
+  CHECK_INPUT(req_pool_indices);
+  CHECK_INPUT(req_to_token);
+  CHECK_INPUT(start_offset);
+  CHECK_INPUT(end_offset);
+  CHECK_INPUT(out_cache_loc);
+  CHECK_DIM(2, req_to_token);
+  CHECK_EQ(req_to_token.scalar_type(), at::kInt);
+  CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
+  CHECK_EQ(end_offset.scalar_type(), start_offset.scalar_type());
+  CHECK_EQ(req_to_token.size(1), pool_len);
+
+  int64_t batch_size = req_pool_indices.size(0);
+  CHECK_EQ(start_offset.numel(), batch_size);
+  CHECK_EQ(end_offset.numel(), batch_size);
+
+  AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_req_to_token_pool_rpi", [&] {
+    using rpi_t = index_t;
+    const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
+    AT_DISPATCH_INDEX_TYPES(start_offset.scalar_type(), "assign_req_to_token_pool_offsets", [&] {
+      assign_req_to_token_pool_kernel_impl<rpi_t, index_t>(
+          rpi_ptr,
+          req_to_token.data_ptr<int32_t>(),
+          start_offset.data_ptr<index_t>(),
+          end_offset.data_ptr<index_t>(),
+          out_cache_loc.data_ptr<int64_t>(),
+          out_cache_loc.numel(),
+          batch_size,
+          pool_len);
+    });
+  });
+}
+
+// Expand req_to_token for multi-step draft decode: row b*topk+tk holds the
+// committed prefix of request b followed by candidate tk's draft slots
+// (which assign_draft_cache_locs_contiguous laid out at sl + tk*num_steps).
+//
+// req_to_token:     [max_num_reqs, pool_len] int32
+// req_pool_indices: [num_seqs] int32 or int64
+// seq_lens:         [num_seqs] int32 or int64 (independent of req_pool_indices)
+// returns:          [num_seqs * topk, pool_len] int32; only the first
+//                   seq_lens[b] + num_steps entries of each row are defined
+at::Tensor build_draft_decode_metadata_cpu(
+    const at::Tensor& req_to_token,
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& seq_lens,
+    int64_t topk,
+    int64_t num_steps,
+    int64_t pool_len) {
+  CHECK_INPUT(req_to_token);
+  CHECK_INPUT(req_pool_indices);
+  CHECK_INPUT(seq_lens);
+  CHECK_DIM(2, req_to_token);
+  CHECK_EQ(req_to_token.scalar_type(), at::kInt);
+  CHECK_EQ(req_to_token.size(1), pool_len);
+
+  int64_t num_seqs = req_pool_indices.size(0);
+  int64_t bs = num_seqs * topk;
+  CHECK_EQ(seq_lens.numel(), num_seqs);
+
+  auto req_to_token_draft = at::empty({bs, pool_len}, req_to_token.options());
+
+  auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
+  auto* draft_ptr = req_to_token_draft.data_ptr<int32_t>();
+
+  AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "build_draft_decode_metadata_rpi", [&] {
+    using rpi_t = index_t;
+    const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
+    AT_DISPATCH_INDEX_TYPES(seq_lens.scalar_type(), "build_draft_decode_metadata_lens", [&] {
+      const index_t* sl_ptr = seq_lens.data_ptr<index_t>();
+
+      at::parallel_for(0, num_seqs, 0, [&](int64_t begin, int64_t end) {
+        for (int64_t b = begin; b < end; ++b) {
+          int64_t idx = rpi_ptr[b];
+          int64_t sl = sl_ptr[b];
+          const int32_t* src_row = rtt_ptr + idx * pool_len;
+
+          for (int64_t tk = 0; tk < topk; ++tk) {
+            int64_t flat = b * topk + tk;
+            int32_t* dst_row = draft_ptr + flat * pool_len;
+
+            // Copy prefix
+            std::memcpy(dst_row, src_row, sl * sizeof(int32_t));
+
+            // Copy draft tokens for this candidate
+            int64_t draft_start = sl + tk * num_steps;
+            for (int64_t s = 0; s < num_steps; ++s) {
+              dst_row[sl + s] = src_row[draft_start + s];
+            }
+          }
+        }
+      });
+    });
+  });
+
+  return req_to_token_draft;
+}
+
+// Pick the last accepted token of each request as its bonus token.
+//
+// accept_tokens: [bs, accept_stride] int32; row-major, accept_stride = accept_index.shape[1]
+// accept_lens:   [bs] int32; number of accepted tokens per request (bonus included)
+// bonus_tokens:  [bs] int32; out
+void fill_bonus_tokens_cpu(
+    const at::Tensor& accept_tokens, const at::Tensor& accept_lens, at::Tensor bonus_tokens, int64_t accept_stride) {
+  CHECK_INPUT(accept_tokens);
+  CHECK_INPUT(accept_lens);
+  CHECK_INPUT(bonus_tokens);
+  CHECK_EQ(accept_tokens.scalar_type(), at::kInt);
+  CHECK_EQ(accept_lens.scalar_type(), at::kInt);
+  CHECK_EQ(bonus_tokens.scalar_type(), at::kInt);
+
+  int64_t bs = accept_lens.size(0);
+  CHECK_EQ(accept_tokens.numel(), bs * accept_stride);
+  CHECK_EQ(bonus_tokens.numel(), bs);
+  auto* accept_ptr = accept_tokens.data_ptr<int32_t>();
+  auto* al_ptr = accept_lens.data_ptr<int32_t>();
+  auto* out_ptr = bonus_tokens.data_ptr<int32_t>();
+
+  at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t pid = begin; pid < end; ++pid) {
+      int64_t idx = accept_stride * pid + al_ptr[pid] - 1;
+      out_ptr[pid] = accept_ptr[idx];
+    }
+  });
+}
+
+// Compact the accepted tokens' KV slots: gather out_cache_loc at the accepted
+// indices, skipping -1 (rejected) entries. Sequential by design: the output
+// write position depends on how many prior entries were accepted.
+//
+// accept_index:         [bs * num_spec_step] int32 or int64; flat, -1 = rejected
+// out_cache_loc:        [bs * num_draft_tokens] int64
+// accept_out_cache_loc: [>= num_accept] int64; out, only the first num_accept
+//                       entries are written
+void fill_accept_out_cache_loc_cpu(
+    const at::Tensor& accept_index, const at::Tensor& out_cache_loc, at::Tensor accept_out_cache_loc) {
+  CHECK_INPUT(accept_index);
+  CHECK_INPUT(out_cache_loc);
+  CHECK_INPUT(accept_out_cache_loc);
+  CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
+  CHECK_EQ(accept_out_cache_loc.scalar_type(), at::kLong);
+  // num_accept <= accept_index.numel(), so this bounds every write below.
+  CHECK_GE(accept_out_cache_loc.numel(), accept_index.numel());
+
+  int64_t num_indices = accept_index.numel();
+  int64_t num_cache_locs = out_cache_loc.numel();
+  auto* ocl_ptr = out_cache_loc.data_ptr<int64_t>();
+  auto* out_ptr = accept_out_cache_loc.data_ptr<int64_t>();
+
+  AT_DISPATCH_INDEX_TYPES(accept_index.scalar_type(), "fill_accept_out_cache_loc_indices", [&] {
+    const index_t* ai_ptr = accept_index.data_ptr<index_t>();
+    int64_t dst = 0;
+    for (int64_t i = 0; i < num_indices; ++i) {
+      int64_t src = static_cast<int64_t>(ai_ptr[i]);
+      if (src > -1) {
+        TORCH_CHECK(src < num_cache_locs, "fill_accept_out_cache_loc: accept_index ", src, " out of range");
+        out_ptr[dst++] = ocl_ptr[src];
+      }
+    }
+  });
+}
+
+// Read back the draft KV slots reserved by the allocator: for each request,
+// copy the topk*num_steps slots starting at seq_lens[pid] out of req_to_token.
+//
+// req_pool_indices: [bs] int32 or int64
+// req_to_token:     [max_num_reqs, pool_len] int32
+// seq_lens:         [bs] int32 or int64 (independent of req_pool_indices)
+// out_cache_loc:    [bs * topk * num_steps] int64; out
+void assign_draft_cache_locs_contiguous_cpu(
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& req_to_token,
+    const at::Tensor& seq_lens,
+    at::Tensor out_cache_loc,
+    int64_t pool_len,
+    int64_t topk,
+    int64_t num_steps) {
+  // Contiguous slot layout: requires page_size == 1 or topk == 1 (see prepare_for_v2_draft guard).
+  CHECK_INPUT(req_pool_indices);
+  CHECK_INPUT(req_to_token);
+  CHECK_INPUT(seq_lens);
+  CHECK_INPUT(out_cache_loc);
+  CHECK_DIM(2, req_to_token);
+  CHECK_EQ(req_to_token.scalar_type(), at::kInt);
+  CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
+  CHECK_EQ(req_to_token.size(1), pool_len);
+  CHECK_EQ(out_cache_loc.numel(), req_pool_indices.numel() * topk * num_steps);
+
+  int64_t bs = req_pool_indices.size(0);
+  int64_t copy_len = topk * num_steps;
+  CHECK_EQ(seq_lens.numel(), bs);
+
+  auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
+  auto* out_ptr = out_cache_loc.data_ptr<int64_t>();
+
+  AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_draft_cache_locs_contiguous_rpi", [&] {
+    using rpi_t = index_t;
+    const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
+    AT_DISPATCH_INDEX_TYPES(seq_lens.scalar_type(), "assign_draft_cache_locs_contiguous_lens", [&] {
+      const index_t* sl_ptr = seq_lens.data_ptr<index_t>();
+
+      at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
+        for (int64_t pid = begin; pid < end; ++pid) {
+          int64_t kv_start = sl_ptr[pid];
+          int64_t req_idx = rpi_ptr[pid];
+          const int32_t* src = rtt_ptr + req_idx * pool_len + kv_start;
+          int64_t* dst = out_ptr + pid * copy_len;
+          for (int64_t j = 0; j < copy_len; ++j) {
+            dst[j] = static_cast<int64_t>(src[j]);
+          }
+        }
+      });
+    });
+  });
+}
+
+// Gather each request's KV slots in [start_offset, end_offset) out of
+// req_to_token into a dense int64 vector (verify/extend cache locations).
+//
+// req_pool_indices: [bs] int32 or int64
+// req_to_token:     [max_num_reqs, pool_len] int32
+// start_offset:     [bs] int32 or int64 (independent of req_pool_indices)
+// end_offset:       [bs] same dtype as start_offset
+// out_cache_loc:    [sum_i(end_offset[i] - start_offset[i])] int64; out
+void assign_extend_cache_locs_cpu(
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& req_to_token,
+    const at::Tensor& start_offset,
+    const at::Tensor& end_offset,
+    at::Tensor out_cache_loc,
+    int64_t pool_len) {
+  CHECK_INPUT(req_pool_indices);
+  CHECK_INPUT(req_to_token);
+  CHECK_INPUT(start_offset);
+  CHECK_INPUT(end_offset);
+  CHECK_INPUT(out_cache_loc);
+  CHECK_DIM(2, req_to_token);
+  CHECK_EQ(req_to_token.scalar_type(), at::kInt);
+  CHECK_EQ(out_cache_loc.scalar_type(), at::kLong);
+  CHECK_EQ(end_offset.scalar_type(), start_offset.scalar_type());
+  CHECK_EQ(req_to_token.size(1), pool_len);
+
+  int64_t bs = req_pool_indices.size(0);
+  CHECK_EQ(start_offset.numel(), bs);
+  CHECK_EQ(end_offset.numel(), bs);
+  auto* rtt_ptr = req_to_token.data_ptr<int32_t>();
+  auto* out_ptr = out_cache_loc.data_ptr<int64_t>();
+
+  AT_DISPATCH_INDEX_TYPES(req_pool_indices.scalar_type(), "assign_extend_cache_locs_rpi", [&] {
+    using rpi_t = index_t;
+    const rpi_t* rpi_ptr = req_pool_indices.data_ptr<rpi_t>();
+    AT_DISPATCH_INDEX_TYPES(start_offset.scalar_type(), "assign_extend_cache_locs_offsets", [&] {
+      const index_t* start_ptr = start_offset.data_ptr<index_t>();
+      const index_t* end_ptr = end_offset.data_ptr<index_t>();
+
+      // Compute prefix sum for output offsets (sequential)
+      std::vector<int64_t> out_offsets(bs + 1, 0);
+      for (int64_t i = 0; i < bs; ++i) {
+        out_offsets[i + 1] = out_offsets[i] + (end_ptr[i] - start_ptr[i]);
+      }
+      // Callers may size out_cache_loc at max capacity (e.g. bs * num_spec_step
+      // in move_accept_tokens) and leave the tail untouched, hence <= not ==.
+      TORCH_CHECK(
+          out_offsets[bs] <= out_cache_loc.numel(),
+          "assign_extend_cache_locs: out_cache_loc has ",
+          out_cache_loc.numel(),
+          " entries but offsets require ",
+          out_offsets[bs]);
+
+      at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
+        for (int64_t pid = begin; pid < end; ++pid) {
+          int64_t kv_start = start_ptr[pid];
+          int64_t kv_end = end_ptr[pid];
+          int64_t req_idx = rpi_ptr[pid];
+          int64_t length = kv_end - kv_start;
+          const int32_t* src = rtt_ptr + req_idx * pool_len + kv_start;
+          int64_t* dst = out_ptr + out_offsets[pid];
+          for (int64_t j = 0; j < length; ++j) {
+            dst[j] = static_cast<int64_t>(src[j]);
+          }
+        }
+      });
+    });
+  });
+}
+
+// Recover tree linkage from a QLEN-layout boolean tree mask (NGRAM path):
+// depth/position, retrieval index, first child and next sibling per node.
+//
+// tree_mask:           [bs * draft_token_num * draft_token_num] bool
+// verified_seq_len:    [bs] int32 or int64
+// positions:           [bs * draft_token_num]; out, same dtype as verified_seq_len
+// retrive_index:       [bs, draft_token_num]; out
+// retrive_next_token:  [bs, draft_token_num]; out
+// retrive_next_sibling:[bs, draft_token_num]; out
+void reconstruct_indices_from_tree_mask_cpu(
+    const at::Tensor& tree_mask,
+    const at::Tensor& verified_seq_len,
+    at::Tensor positions,
+    at::Tensor retrive_index,
+    at::Tensor retrive_next_token,
+    at::Tensor retrive_next_sibling,
+    int64_t batch_size,
+    int64_t draft_token_num) {
+  CHECK_INPUT(tree_mask);
+  CHECK_INPUT(verified_seq_len);
+  CHECK_INPUT(positions);
+  CHECK_INPUT(retrive_index);
+  CHECK_INPUT(retrive_next_token);
+  CHECK_INPUT(retrive_next_sibling);
+  CHECK_EQ(tree_mask.scalar_type(), at::kBool);
+  CHECK_EQ(tree_mask.numel(), batch_size * draft_token_num * draft_token_num);
+  CHECK_EQ(verified_seq_len.numel(), batch_size);
+  CHECK_EQ(positions.numel(), batch_size * draft_token_num);
+  CHECK_EQ(retrive_index.numel(), batch_size * draft_token_num);
+  CHECK_EQ(retrive_next_token.numel(), batch_size * draft_token_num);
+  CHECK_EQ(retrive_next_sibling.numel(), batch_size * draft_token_num);
+  const auto index_dtype = verified_seq_len.scalar_type();
+  CHECK_EQ(positions.scalar_type(), index_dtype);
+  CHECK_EQ(retrive_index.scalar_type(), index_dtype);
+  CHECK_EQ(retrive_next_token.scalar_type(), index_dtype);
+  CHECK_EQ(retrive_next_sibling.scalar_type(), index_dtype);
+
+  const bool* mask_ptr = tree_mask.data_ptr<bool>();
+  int64_t base_offset = draft_token_num * draft_token_num;
+
+  AT_DISPATCH_INDEX_TYPES(index_dtype, "reconstruct_indices_from_tree_mask_indices", [&] {
+    const index_t* seq_len_ptr = verified_seq_len.data_ptr<index_t>();
+    index_t* pos_ptr = positions.data_ptr<index_t>();
+    index_t* ri_ptr = retrive_index.data_ptr<index_t>();
+    index_t* rnt_ptr = retrive_next_token.data_ptr<index_t>();
+    index_t* rns_ptr = retrive_next_sibling.data_ptr<index_t>();
+
+    at::parallel_for(0, batch_size * draft_token_num, 0, [&](int64_t begin, int64_t end) {
+      for (int64_t idx = begin; idx < end; ++idx) {
+        int64_t bid = idx / draft_token_num;
+        int64_t tid = idx % draft_token_num;
+
+        int64_t token_idx = bid * draft_token_num;
+        int64_t tree_mask_offset = bid * base_offset;
+
+        // Step 1: depth and parent via backward scan
+        int64_t depth = 0;
+        int64_t parent_idx = -1;
+        for (int64_t i = tid - 1, start_idx = tree_mask_offset + tid * draft_token_num; i >= 0; --i) {
+          if (mask_ptr[start_idx + i]) {
+            depth++;
+            if (parent_idx == -1) {
+              parent_idx = i;
+            }
+          }
+        }
+
+        // Step 2: retrive_index (identity)
+        ri_ptr[token_idx + tid] = token_idx + tid;
+
+        // Step 3: position = depth + verified_seq_len
+        pos_ptr[token_idx + tid] = depth + seq_len_ptr[bid];
+
+        // Step 4: first child (next_token)
+        int64_t next_token_idx = -1;
+        for (int64_t i = tid + 1; i < draft_token_num; ++i) {
+          if (mask_ptr[tree_mask_offset + i * draft_token_num + tid]) {
+            next_token_idx = i;
+            break;
+          }
+        }
+        rnt_ptr[token_idx + tid] = next_token_idx;
+
+        // Step 5: next sibling (shares parent, no intervening ancestors)
+        int64_t next_sibling_idx = -1;
+        if (parent_idx != -1) {
+          for (int64_t i = tid + 1; i < draft_token_num; ++i) {
+            int64_t si = tree_mask_offset + i * draft_token_num + parent_idx;
+            if (mask_ptr[si]) {
+              bool is_sibling = true;
+              int64_t ei = tree_mask_offset + i * draft_token_num + i;
+              for (int64_t j = si + 1; j < ei; ++j) {
+                if (mask_ptr[j]) {
+                  is_sibling = false;
+                  break;
+                }
+              }
+              if (is_sibling) {
+                next_sibling_idx = i;
+                break;
+              }
+            }
+          }
+        }
+        rns_ptr[token_idx + tid] = next_sibling_idx;
+      }
+    });
+  });
+}
+
+// Shift each request's extend segment left by one token and write the new
+// draft token at the end (or at select_index when given). Mutates input_ids
+// in place; callers rely on this.
+//
+// input_ids:        [num_extend_tokens] int64; in/out
+// extend_start_loc: [bs] int32 or int64
+// extend_seq_lens:  [bs] int32 or int64 (independent of extend_start_loc; the
+//                   spec decode-extend batch pairs int64 lens with int32 locs)
+// topk_index:       [bs] int64; new draft token per request
+// select_index:     [bs] int64 or None; global slot for the new token
+void rotate_input_ids_cpu(
+    at::Tensor input_ids,
+    const at::Tensor& extend_start_loc,
+    const at::Tensor& extend_seq_lens,
+    const at::Tensor& topk_index,
+    const std::optional<at::Tensor>& select_index_opt) {
+  CHECK_INPUT(input_ids);
+  CHECK_INPUT(extend_start_loc);
+  CHECK_INPUT(extend_seq_lens);
+  CHECK_INPUT(topk_index);
+  CHECK_EQ(input_ids.scalar_type(), at::kLong);
+  CHECK_EQ(topk_index.scalar_type(), at::kLong);
+
+  int64_t bs = extend_seq_lens.size(0);
+  CHECK_EQ(extend_start_loc.numel(), bs);
+  CHECK_EQ(topk_index.numel(), bs);
+  if (select_index_opt.has_value()) {
+    CHECK_INPUT(select_index_opt.value());
+    CHECK_EQ(select_index_opt.value().scalar_type(), at::kLong);
+    CHECK_EQ(select_index_opt.value().numel(), bs);
+  }
+
+  auto* ids_ptr = input_ids.data_ptr<int64_t>();
+  auto* topk_ptr = topk_index.data_ptr<int64_t>();
+  const int64_t* select_ptr = conditional_data_ptr<int64_t>(select_index_opt);
+
+  AT_DISPATCH_INDEX_TYPES(extend_start_loc.scalar_type(), "rotate_input_ids_start", [&] {
+    using start_t = index_t;
+    const start_t* start_ptr = extend_start_loc.data_ptr<start_t>();
+    AT_DISPATCH_INDEX_TYPES(extend_seq_lens.scalar_type(), "rotate_input_ids_lens", [&] {
+      const index_t* lens_ptr = extend_seq_lens.data_ptr<index_t>();
+
+      at::parallel_for(0, bs, 0, [&](int64_t begin, int64_t end) {
+        for (int64_t pid = begin; pid < end; ++pid) {
+          int64_t start = start_ptr[pid];
+          int64_t seq_len = lens_ptr[pid];
+          int64_t new_token = topk_ptr[pid];
+
+          // Shift left by 1
+          if (seq_len > 1) {
+            std::memmove(ids_ptr + start, ids_ptr + start + 1, (seq_len - 1) * sizeof(int64_t));
+          }
+          // Write new token
+          if (seq_len > 0) {
+            if (select_ptr != nullptr) {
+              ids_ptr[select_ptr[pid]] = new_token;
+            } else {
+              ids_ptr[start + seq_len - 1] = new_token;
+            }
+          }
+        }
+      });
+    });
+  });
+}

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -75,6 +75,87 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qk_gemma_rmsnorm_with_gate_
     int64_t head_dim,
     int64_t num_head);
 
+// speculative decoding
+void verify_tree_greedy_cpu(
+    at::Tensor predicts,
+    at::Tensor accept_index,
+    at::Tensor accept_token_num,
+    const at::Tensor& candidates,
+    const at::Tensor& retrive_index,
+    const at::Tensor& retrive_next_token,
+    const at::Tensor& retrive_next_sibling,
+    const at::Tensor& target_predict);
+
+void build_tree_kernel_efficient_cpu(
+    const at::Tensor& parent_list,
+    const at::Tensor& selected_index,
+    const at::Tensor& verified_seq_len,
+    at::Tensor tree_mask,
+    at::Tensor positions,
+    at::Tensor retrive_index,
+    at::Tensor retrive_next_token,
+    at::Tensor retrive_next_sibling,
+    int64_t topk,
+    int64_t depth,
+    int64_t draft_token_num,
+    int64_t tree_mask_mode);
+
+void assign_req_to_token_pool_cpu(
+    const at::Tensor& req_pool_indices,
+    at::Tensor req_to_token,
+    const at::Tensor& start_offset,
+    const at::Tensor& end_offset,
+    const at::Tensor& out_cache_loc,
+    int64_t pool_len);
+
+at::Tensor build_draft_decode_metadata_cpu(
+    const at::Tensor& req_to_token,
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& seq_lens,
+    int64_t topk,
+    int64_t num_steps,
+    int64_t pool_len);
+
+void fill_bonus_tokens_cpu(
+    const at::Tensor& accept_tokens, const at::Tensor& accept_lens, at::Tensor bonus_tokens, int64_t accept_stride);
+
+void fill_accept_out_cache_loc_cpu(
+    const at::Tensor& accept_index, const at::Tensor& out_cache_loc, at::Tensor accept_out_cache_loc);
+
+void assign_draft_cache_locs_contiguous_cpu(
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& req_to_token,
+    const at::Tensor& seq_lens,
+    at::Tensor out_cache_loc,
+    int64_t pool_len,
+    int64_t topk,
+    int64_t num_steps);
+
+void assign_extend_cache_locs_cpu(
+    const at::Tensor& req_pool_indices,
+    const at::Tensor& req_to_token,
+    const at::Tensor& start_offset,
+    const at::Tensor& end_offset,
+    at::Tensor out_cache_loc,
+    int64_t pool_len);
+
+void reconstruct_indices_from_tree_mask_cpu(
+    const at::Tensor& tree_mask,
+    const at::Tensor& verified_seq_len,
+    at::Tensor positions,
+    at::Tensor retrive_index,
+    at::Tensor retrive_next_token,
+    at::Tensor retrive_next_sibling,
+    int64_t batch_size,
+    int64_t draft_token_num);
+
+void rotate_input_ids_cpu(
+    at::Tensor input_ids,
+    const at::Tensor& extend_start_loc,
+    const at::Tensor& extend_seq_lens,
+    const at::Tensor& topk_index,
+    const std::optional<at::Tensor>& select_index_opt);
+
 // topk
 std::tuple<at::Tensor, at::Tensor>
 topk_sigmoid_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t topk, bool renormalize);
@@ -142,7 +223,8 @@ void extend_attention_cpu(
     bool is_cross_attn,
     int64_t sliding_window_size,
     std::optional<at::Tensor> encoder_lens,
-    std::optional<at::Tensor> sinks);
+    std::optional<at::Tensor> sinks,
+    std::optional<at::Tensor> tree_mask);
 
 // flash attention
 at::Tensor flash_attn_varlen_func(
@@ -449,6 +531,9 @@ void store_cache_cpu(
     const at::Tensor& indices,
     std::optional<int64_t> row_dim);
 
+void copy_all_layer_kv_cache_cpu(
+    const at::Tensor& data_ptrs, const at::Tensor& strides, const at::Tensor& tgt_loc, const at::Tensor& src_loc);
+
 // [NOTE] When registering kernels, we should accurately describe the in-place information.
 // Taking fused_add_rmsnorm_cpu as an example, add `Tensor(a!)` modifier to all tensors that
 // will be modified in-place to avoid incorrect fusing and execution order on graph mode.
@@ -496,6 +581,64 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "(Tensor, Tensor, Tensor)");
   m.impl("fused_qk_gemma_rmsnorm_with_gate_cpu", torch::kCPU, &fused_qk_gemma_rmsnorm_with_gate_cpu);
 
+  // speculative decoding
+  m.def(
+      "verify_tree_greedy_cpu(Tensor(a!) predicts, Tensor(a!) accept_index, "
+      "Tensor(a!) accept_token_num, Tensor candidates, Tensor retrive_index, "
+      "Tensor retrive_next_token, Tensor retrive_next_sibling, Tensor target_predict) -> ()");
+  m.impl("verify_tree_greedy_cpu", torch::kCPU, &verify_tree_greedy_cpu);
+
+  m.def(
+      "build_tree_kernel_efficient_cpu(Tensor parent_list, Tensor selected_index, "
+      "Tensor verified_seq_len, Tensor(a!) tree_mask, Tensor(a!) positions, "
+      "Tensor(a!) retrive_index, Tensor(a!) retrive_next_token, "
+      "Tensor(a!) retrive_next_sibling, int topk, int depth, "
+      "int draft_token_num, int tree_mask_mode) -> ()");
+  m.impl("build_tree_kernel_efficient_cpu", torch::kCPU, &build_tree_kernel_efficient_cpu);
+
+  m.def(
+      "assign_req_to_token_pool_cpu(Tensor req_pool_indices, Tensor(a!) req_to_token, "
+      "Tensor start_offset, Tensor end_offset, Tensor out_cache_loc, "
+      "int pool_len) -> ()");
+  m.impl("assign_req_to_token_pool_cpu", torch::kCPU, &assign_req_to_token_pool_cpu);
+
+  m.def(
+      "build_draft_decode_metadata_cpu(Tensor req_to_token, Tensor req_pool_indices, "
+      "Tensor seq_lens, int topk, int num_steps, int pool_len) -> Tensor");
+  m.impl("build_draft_decode_metadata_cpu", torch::kCPU, &build_draft_decode_metadata_cpu);
+
+  m.def(
+      "fill_bonus_tokens_cpu(Tensor accept_tokens, Tensor accept_lens, "
+      "Tensor(a!) bonus_tokens, int accept_stride) -> ()");
+  m.impl("fill_bonus_tokens_cpu", torch::kCPU, &fill_bonus_tokens_cpu);
+
+  m.def(
+      "fill_accept_out_cache_loc_cpu(Tensor accept_index, Tensor out_cache_loc, "
+      "Tensor(a!) accept_out_cache_loc) -> ()");
+  m.impl("fill_accept_out_cache_loc_cpu", torch::kCPU, &fill_accept_out_cache_loc_cpu);
+
+  m.def(
+      "assign_draft_cache_locs_contiguous_cpu(Tensor req_pool_indices, Tensor req_to_token, "
+      "Tensor seq_lens, Tensor(a!) out_cache_loc, int pool_len, int topk, int num_steps) -> ()");
+  m.impl("assign_draft_cache_locs_contiguous_cpu", torch::kCPU, &assign_draft_cache_locs_contiguous_cpu);
+
+  m.def(
+      "assign_extend_cache_locs_cpu(Tensor req_pool_indices, Tensor req_to_token, "
+      "Tensor start_offset, Tensor end_offset, Tensor(a!) out_cache_loc, int pool_len) -> ()");
+  m.impl("assign_extend_cache_locs_cpu", torch::kCPU, &assign_extend_cache_locs_cpu);
+
+  m.def(
+      "rotate_input_ids_cpu(Tensor(a!) input_ids, Tensor extend_start_loc, "
+      "Tensor extend_seq_lens, Tensor topk_index, Tensor? select_index=None) -> ()");
+  m.impl("rotate_input_ids_cpu", torch::kCPU, &rotate_input_ids_cpu);
+
+  m.def(
+      "reconstruct_indices_from_tree_mask_cpu(Tensor tree_mask, Tensor verified_seq_len, "
+      "Tensor(a!) positions, Tensor(a!) retrive_index, "
+      "Tensor(a!) retrive_next_token, Tensor(a!) retrive_next_sibling, "
+      "int batch_size, int draft_token_num) -> ()");
+  m.impl("reconstruct_indices_from_tree_mask_cpu", torch::kCPU, &reconstruct_indices_from_tree_mask_cpu);
+
   // topk
   m.def("topk_sigmoid_cpu(Tensor hidden_states, Tensor gating_output, int topk, bool renormalize) -> (Tensor, Tensor)");
   m.impl("topk_sigmoid_cpu", torch::kCPU, &topk_sigmoid_cpu);
@@ -528,7 +671,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "Tensor v_buffer, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor extend_seq_lens, Tensor "
       "extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, int "
       "sliding_window_size, Tensor? "
-      "encoder_lens, Tensor? sinks) -> ()");
+      "encoder_lens, Tensor? sinks, Tensor? tree_mask=None) -> ()");
   m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
 
   // flash attn
@@ -716,6 +859,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "store_cache_cpu(Tensor k, Tensor v, Tensor(a!) k_cache, Tensor(a!) v_cache, Tensor indices, int? row_dim) -> "
       "()");
   m.impl("store_cache_cpu", torch::kCPU, &store_cache_cpu);
+
+  // The copy mutates the K/V buffers addressed via `data_ptrs` (a table of
+  // raw base pointers), which schema-level alias annotations cannot express.
+  m.def("copy_all_layer_kv_cache_cpu(Tensor data_ptrs, Tensor strides, Tensor tgt_loc, Tensor src_loc) -> ()");
+  m.impl("copy_all_layer_kv_cache_cpu", torch::kCPU, &copy_all_layer_kv_cache_cpu);
 }
 
 TORCH_LIBRARY_IMPL(sgl_kernel, CatchAll, m) {

sgl-kernel/python/sgl_kernel/__init__.py

@@ -76,6 +76,7 @@ else:
         max_pooling_1d_varlen,
     )
     from sgl_kernel.kvcacheio import (
+        copy_all_layer_kv_cache_cpu,
         transfer_kv_all_layer,
         transfer_kv_all_layer_mla,
         transfer_kv_per_layer,
@@ -113,11 +114,20 @@ else:
         top_p_renorm_prob,
     )
     from sgl_kernel.speculative import (
+        assign_draft_cache_locs_contiguous_cpu,
+        assign_extend_cache_locs_cpu,
+        assign_req_to_token_pool_cpu,
+        build_draft_decode_metadata_cpu,
         build_tree_kernel_efficient,
+        build_tree_kernel_efficient_cpu,
+        fill_accept_out_cache_loc_cpu,
+        fill_bonus_tokens_cpu,
         reconstruct_indices_from_tree_mask,
+        rotate_input_ids_cpu,
         segment_packbits,
         tree_speculative_sampling_target_only,
         verify_tree_greedy,
+        verify_tree_greedy_cpu,
     )
     from sgl_kernel.top_k import (
         fast_topk,

sgl-kernel/python/sgl_kernel/kvcacheio.py

@@ -305,3 +305,17 @@ def transfer_kv_all_layer_mla_lf_pf(
         block_quota,
         num_warps_per_block,
     )
+
+
+def copy_all_layer_kv_cache_cpu(
+    data_ptrs: torch.Tensor,
+    strides: torch.Tensor,
+    tgt_loc: torch.Tensor,
+    src_loc: torch.Tensor,
+):
+    torch.ops.sgl_kernel.copy_all_layer_kv_cache_cpu(
+        data_ptrs,
+        strides,
+        tgt_loc,
+        src_loc,
+    )

sgl-kernel/python/sgl_kernel/speculative.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 
 
@@ -97,16 +99,28 @@ def reconstruct_indices_from_tree_mask(
     batch_size: int,
     draft_token_num: int,
 ) -> None:
-    torch.ops.sgl_kernel.reconstruct_indices_from_tree_mask.default(
-        tree_mask,
-        verified_seq_len,
-        positions,
-        retrive_index,
-        retrive_next_token,
-        retrive_next_sibling,
-        batch_size,
-        draft_token_num,
-    )
+    if tree_mask.is_cpu:
+        torch.ops.sgl_kernel.reconstruct_indices_from_tree_mask_cpu(
+            tree_mask,
+            verified_seq_len,
+            positions,
+            retrive_index,
+            retrive_next_token,
+            retrive_next_sibling,
+            batch_size,
+            draft_token_num,
+        )
+    else:
+        torch.ops.sgl_kernel.reconstruct_indices_from_tree_mask.default(
+            tree_mask,
+            verified_seq_len,
+            positions,
+            retrive_index,
+            retrive_next_token,
+            retrive_next_sibling,
+            batch_size,
+            draft_token_num,
+        )
 
 
 def segment_packbits(
@@ -124,3 +138,171 @@ def segment_packbits(
         batch_size,
         torch.cuda.current_stream().cuda_stream,
     )
+
+
+def verify_tree_greedy_cpu(
+    predicts: torch.Tensor,  # mutable
+    accept_index: torch.Tensor,  # mutable
+    accept_token_num: torch.Tensor,  # mutable
+    candidates: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    target_predict: torch.Tensor,
+) -> None:
+    torch.ops.sgl_kernel.verify_tree_greedy_cpu(
+        predicts,
+        accept_index,
+        accept_token_num,
+        candidates,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        target_predict,
+    )
+
+
+def build_tree_kernel_efficient_cpu(
+    parent_list: torch.Tensor,
+    selected_index: torch.Tensor,
+    verified_seq_len: torch.Tensor,
+    tree_mask: torch.Tensor,
+    positions: torch.Tensor,
+    retrive_index: torch.Tensor,
+    retrive_next_token: torch.Tensor,
+    retrive_next_sibling: torch.Tensor,
+    topk: int,
+    depth: int,
+    draft_token_num: int,
+    tree_mask_mode: int,
+) -> None:
+    torch.ops.sgl_kernel.build_tree_kernel_efficient_cpu(
+        parent_list,
+        selected_index,
+        verified_seq_len,
+        tree_mask,
+        positions,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        topk,
+        depth,
+        draft_token_num,
+        tree_mask_mode,
+    )
+
+
+def assign_req_to_token_pool_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_req_to_token_pool_cpu(
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        pool_len,
+    )
+
+
+def build_draft_decode_metadata_cpu(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    topk: int,
+    num_steps: int,
+    pool_len: int,
+) -> torch.Tensor:
+    return torch.ops.sgl_kernel.build_draft_decode_metadata_cpu(
+        req_to_token,
+        req_pool_indices,
+        seq_lens,
+        topk,
+        num_steps,
+        pool_len,
+    )
+
+
+def fill_bonus_tokens_cpu(
+    accept_tokens: torch.Tensor,
+    accept_lens: torch.Tensor,
+    bonus_tokens: torch.Tensor,
+    accept_stride: int,
+) -> None:
+    torch.ops.sgl_kernel.fill_bonus_tokens_cpu(
+        accept_tokens,
+        accept_lens,
+        bonus_tokens,
+        accept_stride,
+    )
+
+
+def fill_accept_out_cache_loc_cpu(
+    accept_index: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    accept_out_cache_loc: torch.Tensor,  # mutable
+) -> None:
+    torch.ops.sgl_kernel.fill_accept_out_cache_loc_cpu(
+        accept_index,
+        out_cache_loc,
+        accept_out_cache_loc,
+    )
+
+
+def assign_draft_cache_locs_contiguous_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+    topk: int,
+    num_steps: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_draft_cache_locs_contiguous_cpu(
+        req_pool_indices,
+        req_to_token,
+        seq_lens,
+        out_cache_loc,
+        pool_len,
+        topk,
+        num_steps,
+    )
+
+
+def assign_extend_cache_locs_cpu(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    pool_len: int,
+) -> None:
+    torch.ops.sgl_kernel.assign_extend_cache_locs_cpu(
+        req_pool_indices,
+        req_to_token,
+        start_offset,
+        end_offset,
+        out_cache_loc,
+        pool_len,
+    )
+
+
+def rotate_input_ids_cpu(
+    input_ids: torch.Tensor,
+    extend_start_loc: torch.Tensor,
+    extend_seq_lens: torch.Tensor,
+    topk_index: torch.Tensor,
+    select_index: Optional[torch.Tensor] = None,
+) -> None:
+    torch.ops.sgl_kernel.rotate_input_ids_cpu(
+        input_ids,
+        extend_start_loc,
+        extend_seq_lens,
+        topk_index,
+        select_index,
+    )

test/registered/cpu/test_spec_eagle_cpu.py

@@ -0,0 +1,65 @@
+"""EAGLE spec-decoding core on CPU: the standard config (topk=1, page_size=1)
+on the synchronous (non-overlap) path. topk > 1 tree drafting is covered in
+test_spec_eagle_topk_cpu.py (split to stay under the per-file CI timeout).
+"""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.matched_stop_kit import MatchedStopMixin
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecCorrectnessKit,
+    SpecFeatureKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
+
+# Measured 780s all-green on a 40-core GNR socket (1 launch + 18 methods).
+register_cpu_ci(est_time=800, suite="base-b-test-cpu")
+
+_KITS = (
+    SpecCorrectnessKit,
+    SpecAccuracyKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+    SpecFeatureKit,
+    MatchedStopMixin,
+)
+
+
+class _Core(EagleLlama2Base):
+    """EAGLE (Llama-2) preset on CPU."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    gsm8k_num_examples = 64
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2NoOverlap(_Core, *_KITS):
+    """Spec v1 (overlap scheduler off) -- the only mode reachable on CPU."""
+
+    # Standard chain config (topk=1, page_size=1), same shape as the CUDA core.
+    spec_steps = 5
+    spec_topk = 1
+    spec_tokens = 6
+    # EAGLE/Llama-2 topk=1 accepts modestly; tune against CI if needed.
+    acc_length_thres = 1.6
+    batch_accept_len_thres = 1.3
+    gsm8k_accept_len_thres = 1.3
+
+    @unittest.skip(
+        "constrained decoding on CPU needs a vocab-mask CPU branch in the "
+        "xgrammar backend (upstream gap, not spec-specific); the other grammar "
+        "backends lack the rollback spec verification requires"
+    )
+    def test_constrained_decoding(self):
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_spec_eagle_parity_cpu.py

@@ -0,0 +1,23 @@
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.spec_server_kits import SpecParityKit
+from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
+
+# Estimated: 2 sequential 8B server launches + one 4-prompt greedy method
+# (CUDA sibling: 360); tune from CI TIMINGS once it has run there.
+register_cpu_ci(est_time=480, suite="base-b-test-cpu")
+
+
+class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):
+    """EAGLE3 spec (intel_amx) greedy output == non-spec reference."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/cpu/test_spec_eagle_topk_cpu.py

@@ -0,0 +1,62 @@
+"""EAGLE topk > 1 tree drafting on CPU (Llama-2 topk=4, synchronous path).
+
+Split from test_spec_eagle_cpu.py, mirroring the CUDA test_spec_eagle.py /
+test_spec_eagle_topk.py layout, so each file stays under the per-file CI
+timeout.
+"""
+
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.kits.spec_server_kits import (
+    SpecAccuracyKit,
+    SpecCorrectnessKit,
+    SpecFeatureKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+)
+from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
+
+# Measured 830s all-green on a 40-core GNR socket (1 launch + 14 methods).
+register_cpu_ci(est_time=850, suite="base-b-test-cpu")
+
+
+class _Core(EagleLlama2Base):
+    """EAGLE (Llama-2) preset on CPU."""
+
+    attention_backend = "intel_amx"
+    disable_overlap = True
+    mem_fraction_static = 0.3
+    gsm8k_num_examples = 64
+    env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
+
+
+class TestEagleLlama2Topk4(
+    _Core,
+    SpecCorrectnessKit,
+    SpecAccuracyKit,
+    SpecLogprobKit,
+    SpecPenaltyKit,
+    SpecFeatureKit,
+):
+    """EAGLE/Llama-2 topk=4 tree coverage (kits listed in bases)."""
+
+    spec_steps = 3
+    spec_topk = 4
+    spec_tokens = 8
+    acc_length_thres = 2.4
+    batch_accept_len_thres = 1.6
+    gsm8k_accept_len_thres = 2.0
+
+    @unittest.skip(
+        "constrained decoding on CPU needs a vocab-mask CPU branch in the "
+        "xgrammar backend (upstream gap, not spec-specific); the other grammar "
+        "backends lack the rollback spec verification requires"
+    )
+    def test_constrained_decoding(self):
+        pass
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -15,13 +15,13 @@ import torch
 from sglang.srt.speculative.adaptive_runtime_state import SpecRuntimeState
 from sglang.srt.speculative.eagle_utils import organize_draft_results
 from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
-from sglang.srt.utils import get_device
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
 register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
 
-DEVICE = get_device()
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def _fake_server_args(**fields):

test/registered/unit/spec/test_spec_cpu_overlap_constraint.py

@@ -0,0 +1,75 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+
+
+def _make_spec_args(device: str, algorithm: str = "EAGLE", **overrides) -> ServerArgs:
+    # model_path="dummy" short-circuits ServerArgs.__post_init__; invoke the
+    # speculative hook directly (same pattern as the unit/server_args tests).
+    args = ServerArgs(model_path="dummy")
+    args.speculative_algorithm = algorithm
+    args.device = device
+    # Fully specify the chain config so the hook doesn't auto-choose params.
+    args.speculative_num_steps = 3
+    args.speculative_eagle_topk = 1
+    args.speculative_num_draft_tokens = 4
+    args.get_model_config = lambda: SimpleNamespace(
+        hf_config=SimpleNamespace(
+            architectures=["LlamaForCausalLM"],
+            get_text_config=lambda: SimpleNamespace(),
+        )
+    )
+    for key, value in overrides.items():
+        setattr(args, key, value)
+    return args
+
+
+class TestSpecCPUOverlapConstraint(CustomTestCase):
+    def test_cpu_eagle_forces_disable_overlap_schedule(self):
+        args = _make_spec_args(device="cpu")
+        self.assertFalse(args.disable_overlap_schedule)
+
+        handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+
+    def test_cpu_eagle3_forces_disable_overlap_schedule(self):
+        args = _make_spec_args(device="cpu", algorithm="EAGLE3")
+
+        handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+
+    def test_cpu_explicit_disable_overlap_is_preserved(self):
+        args = _make_spec_args(device="cpu", disable_overlap_schedule=True)
+
+        # Already disabled: the hook must not flip the flag, and (unlike the
+        # forced-disable cases) must not warn about overriding it.
+        with self.assertLogs(
+            "sglang.srt.arg_groups.speculative_hook", "WARNING"
+        ) as logs:
+            handle_speculative_decoding(args)
+
+        self.assertTrue(args.disable_overlap_schedule)
+        self.assertFalse(
+            any("Overlap schedule" in message for message in logs.output),
+            f"hook warned about overriding an already-disabled overlap: {logs.output}",
+        )
+
+    def test_cuda_eagle_keeps_overlap_schedule(self):
+        # Guard the constraint's scope: the hook must not touch non-CPU devices.
+        args = _make_spec_args(device="cuda")
+
+        handle_speculative_decoding(args)
+
+        self.assertFalse(args.disable_overlap_schedule)
+
+
+if __name__ == "__main__":
+    unittest.main()