diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py
index 776ff0f72..b42c0b621 100644
--- a/python/sglang/srt/arg_groups/speculative_hook.py
+++ b/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
diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py
index 808733d3b..80660d808 100644
--- a/python/sglang/srt/layers/attention/intel_amx_backend.py
+++ b/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,
+            )
diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py
index e3dde55cc..2ade16095 100644
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/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(
diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index 2a535af1c..613492668 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/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,
             )
 
 
diff --git a/python/sglang/srt/mem_cache/triton_ops/cache_move.py b/python/sglang/srt/mem_cache/triton_ops/cache_move.py
index 2f2389e89..e4e408f9b 100644
--- a/python/sglang/srt/mem_cache/triton_ops/cache_move.py
+++ b/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
diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py
index 105a96441..91f89ef14 100644
--- a/python/sglang/srt/models/llama.py
+++ b/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)
diff --git a/python/sglang/srt/models/mimo_mtp.py b/python/sglang/srt/models/mimo_mtp.py
index da6f2558d..d25b87c59 100644
--- a/python/sglang/srt/models/mimo_mtp.py
+++ b/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
diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py
index 46a0c34fa..ad19bdaa7 100644
--- a/python/sglang/srt/models/qwen2.py
+++ b/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)
diff --git a/python/sglang/srt/models/qwen2_eagle.py b/python/sglang/srt/models/qwen2_eagle.py
index b5944d183..017e75bc8 100644
--- a/python/sglang/srt/models/qwen2_eagle.py
+++ b/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:
diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py
index 23deb34ba..909bdbdf1 100644
--- a/python/sglang/srt/speculative/base_spec_worker.py
+++ b/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
diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py
index 2e55c5fd0..b831da232 100644
--- a/python/sglang/srt/speculative/draft_utils.py
+++ b/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
 
diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py
index 078581ee2..11406f6d8 100644
--- a/python/sglang/srt/speculative/eagle_info.py
+++ b/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,
diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py
index daf9fa82e..c3fe5f7ba 100644
--- a/python/sglang/srt/speculative/eagle_utils.py
+++ b/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(
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index 591260089..e1a5f3064 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/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)
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
index 7ad662cba..4d9776acd 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+++ b/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.
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_utils.py b/python/sglang/srt/speculative/multi_layer_eagle_utils.py
index d9ae4cb77..bc54df1ca 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_utils.py
+++ b/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",
 ]
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
index e6bd85f2f..873848a46 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+++ b/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)
diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py
index 4cbcb7f0e..90524c74c 100644
--- a/python/sglang/srt/speculative/ngram_worker.py
+++ b/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.
diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py
index 240678864..b4268a09b 100644
--- a/python/sglang/srt/speculative/spec_utils.py
+++ b/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
diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py
index bf91afbcc..10dc6f7af 100644
--- a/python/sglang/srt/speculative/standalone_worker_v2.py
+++ b/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 = (
diff --git a/python/sglang/srt/speculative/triton_ops/cache_locs.py b/python/sglang/srt/speculative/triton_ops/cache_locs.py
index 663c23872..19dc00f11 100644
--- a/python/sglang/srt/speculative/triton_ops/cache_locs.py
+++ b/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
diff --git a/python/sglang/srt/speculative/triton_ops/eagle.py b/python/sglang/srt/speculative/triton_ops/eagle.py
index de9b2ec4c..10342899c 100644
--- a/python/sglang/srt/speculative/triton_ops/eagle.py
+++ b/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),
+    )
diff --git a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py b/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py
index a4716c9a9..a41c9f1d7 100644
--- a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py
+++ b/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,)
diff --git a/sgl-kernel/csrc/cpu/extend.cpp b/sgl-kernel/csrc/cpu/extend.cpp
index 01458f5fb..4592b1425 100644
--- a/sgl-kernel/csrc/cpu/extend.cpp
+++ b/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) {
diff --git a/sgl-kernel/csrc/cpu/kvcache.cpp b/sgl-kernel/csrc/cpu/kvcache.cpp
index ca2dc81b8..a0508f967 100644
--- a/sgl-kernel/csrc/cpu/kvcache.cpp
+++ b/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);
+        }
+      }
+    });
+  });
+}
diff --git a/sgl-kernel/csrc/cpu/spec.cpp b/sgl-kernel/csrc/cpu/spec.cpp
new file mode 100644
index 000000000..859a38c05
--- /dev/null
+++ b/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;
+            }
+          }
+        }
+      });
+    });
+  });
+}
diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
index 80a910e7f..eea0fe06a 100644
--- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+++ b/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) {
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index 4cbe4b355..199c1f922 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/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,
diff --git a/sgl-kernel/python/sgl_kernel/kvcacheio.py b/sgl-kernel/python/sgl_kernel/kvcacheio.py
index 908a3efa9..090142a98 100644
--- a/sgl-kernel/python/sgl_kernel/kvcacheio.py
+++ b/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,
+    )
diff --git a/sgl-kernel/python/sgl_kernel/speculative.py b/sgl-kernel/python/sgl_kernel/speculative.py
index 17e424972..6df9d8521 100644
--- a/sgl-kernel/python/sgl_kernel/speculative.py
+++ b/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,
+    )
diff --git a/test/registered/cpu/test_spec_eagle_cpu.py b/test/registered/cpu/test_spec_eagle_cpu.py
new file mode 100644
index 000000000..5ef074384
--- /dev/null
+++ b/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()
diff --git a/test/registered/cpu/test_spec_eagle_parity_cpu.py b/test/registered/cpu/test_spec_eagle_parity_cpu.py
new file mode 100644
index 000000000..7694dcb5b
--- /dev/null
+++ b/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()
diff --git a/test/registered/cpu/test_spec_eagle_topk_cpu.py b/test/registered/cpu/test_spec_eagle_topk_cpu.py
new file mode 100644
index 000000000..2bda20df0
--- /dev/null
+++ b/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()
diff --git a/test/registered/cpu/test_spec_kernels.py b/test/registered/cpu/test_spec_kernels.py
new file mode 100644
index 000000000..73d7fe076
--- /dev/null
+++ b/test/registered/cpu/test_spec_kernels.py
@@ -0,0 +1,1381 @@
+import unittest
+
+import sgl_kernel  # noqa: F401
+import torch
+import torch.nn.functional as F
+from utils import precision
+
+from sglang.srt.speculative.eagle_utils import TreeMaskMode, organize_draft_results
+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-b-test-cpu")
+
+
+def _topk1_chain_inputs(bs, num_steps):
+    parent_width = num_steps if num_steps > 1 else 0
+    parent_list = torch.arange(-1, parent_width - 1, dtype=torch.int64).repeat(bs, 1)
+    selected_index = torch.arange(num_steps, dtype=torch.int64).repeat(bs, 1)
+    return parent_list, selected_index
+
+
+def _gen_draft_tree(bs, topk, num_steps, draft_token_num):
+    """Simulate the EAGLE draft loop (select_top_k_tokens + organize_draft_results)
+    with random probabilities to obtain a valid random (parent_list, selected_index).
+    """
+    scores = torch.rand(bs, topk, dtype=torch.float32)
+    score_chunks = [scores]
+    parents_chunks = [torch.arange(-1, topk, dtype=torch.int64).expand(bs, -1)]
+    cum_scores = scores
+    for i in range(1, num_steps):
+        # Probabilities in (0, 1): a child cumulative score is strictly smaller
+        # than its parent's, so the global topk always keeps full ancestor chains.
+        step_p = torch.rand(bs, topk, topk, dtype=torch.float32)
+        expand_scores = cum_scores.unsqueeze(2) * step_p
+        cum_scores, topk_cs_index = torch.topk(
+            expand_scores.flatten(start_dim=1), topk, dim=-1
+        )
+        score_chunks.append(expand_scores.flatten(start_dim=1))
+        parents_chunks.append(topk_cs_index + (topk * topk * (i - 1) + topk))
+    score_flat = torch.cat(score_chunks, dim=1)
+    selected_index = torch.sort(
+        torch.topk(score_flat, draft_token_num - 1, dim=-1).indices, dim=-1
+    ).values
+    parent_list = torch.cat(parents_chunks[:-1], dim=1)
+    return parent_list, selected_index
+
+
+def _ref_build_tree(parent_list, selected_index, seq_lens, topk, draft_token_num):
+    bs = selected_index.shape[0]
+    retrieve_index = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+    retrieve_next_token = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+    retrieve_next_sibling = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+    positions = torch.zeros(bs * draft_token_num, dtype=torch.int64)
+    tree_mask = torch.zeros(bs, draft_token_num, draft_token_num, dtype=torch.bool)
+
+    for bid in range(bs):
+        off = bid * draft_token_num
+        sel = selected_index[bid].tolist()
+        parents = parent_list[bid].tolist()
+        retrieve_index[bid] = torch.arange(off, off + draft_token_num)
+
+        # parent position (in tree-node numbering) of each node i >= 1
+        parent_pos = [0] * draft_token_num
+        for i in range(1, draft_token_num):
+            parent_tb_idx = sel[i - 1] // topk
+            if parent_tb_idx == 0:
+                parent_pos[i] = 0
+            else:
+                parent_token_idx = parents[parent_tb_idx]
+                parent_pos[i] = sel.index(parent_token_idx) + 1
+
+        # head-insertion linking, iterating from the last node (kernel order)
+        next_token = [-1] * draft_token_num
+        next_sibling = [-1] * draft_token_num
+        for i in range(draft_token_num - 1, 0, -1):
+            p = parent_pos[i]
+            if next_token[p] == -1:
+                next_token[p] = i
+            else:
+                next_sibling[i] = next_token[p]
+                next_token[p] = i
+        retrieve_next_token[bid] = torch.tensor(next_token, dtype=torch.int64)
+        retrieve_next_sibling[bid] = torch.tensor(next_sibling, dtype=torch.int64)
+
+        seq_len = int(seq_lens[bid])
+        positions[off] = seq_len
+        tree_mask[bid, :, 0] = True
+        for i in range(1, draft_token_num):
+            ancestors = []
+            j = i
+            while j != 0:
+                ancestors.append(j)
+                j = parent_pos[j]
+            positions[off + i] = seq_len + len(ancestors)
+            for j in ancestors:
+                tree_mask[bid, i, j] = True
+
+    return (
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        positions,
+        tree_mask,
+    )
+
+
+def _run_build_tree_kernel(
+    parent_list, selected_index, seq_lens, topk, num_steps, draft_token_num, mode
+):
+    bs = seq_lens.numel()
+    if mode == TreeMaskMode.QLEN_ONLY:
+        tree_mask = torch.full(
+            (bs * draft_token_num * draft_token_num,), True, dtype=torch.bool
+        )
+    else:  # FULL_MASK carries the (all-true) prefix columns explicitly
+        seq_lens_sum = int(seq_lens.sum())
+        tree_mask = torch.full(
+            (seq_lens_sum * draft_token_num + bs * draft_token_num * draft_token_num,),
+            True,
+            dtype=torch.bool,
+        )
+    positions = torch.zeros(bs * draft_token_num, dtype=torch.int64)
+    retrieve_buf = torch.full((3, bs, draft_token_num), -1, dtype=torch.int64)
+    retrieve_index, retrieve_next_token, retrieve_next_sibling = retrieve_buf
+    torch.ops.sgl_kernel.build_tree_kernel_efficient_cpu(
+        parent_list,
+        selected_index,
+        seq_lens,
+        tree_mask,
+        positions,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        topk,
+        num_steps,
+        draft_token_num,
+        mode,
+    )
+    return (
+        tree_mask,
+        positions,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+    )
+
+
+def _ref_verify_tree_greedy(
+    candidates,
+    retrieve_index,
+    retrieve_next_token,
+    retrieve_next_sibling,
+    target_predict,
+    num_spec_step,
+):
+    bs, draft_token_num = candidates.shape
+    predicts = torch.full((bs * draft_token_num,), -1, dtype=torch.int32)
+    accept_indices = torch.full((bs, num_spec_step), -1, dtype=torch.int32)
+    num_correct_drafts = torch.zeros(bs, dtype=torch.int32)
+    target_flat = target_predict.reshape(-1)
+
+    for bx in range(bs):
+        last_accept_idx = int(retrieve_index[bx, 0])
+        accept_indices[bx, 0] = last_accept_idx
+        num_correct = 0
+        cur = 0
+        for _ in range(1, num_spec_step):
+            cur = int(retrieve_next_token[bx, cur])
+            while cur != -1:
+                draft_idx = int(retrieve_index[bx, cur])
+                draft_token = int(candidates[bx, cur])
+                target_token = int(target_flat[last_accept_idx])
+                if draft_token == target_token:
+                    predicts[last_accept_idx] = target_token
+                    num_correct += 1
+                    accept_indices[bx, num_correct] = draft_idx
+                    last_accept_idx = draft_idx
+                    break
+                cur = int(retrieve_next_sibling[bx, cur])
+            if cur == -1:
+                break
+        num_correct_drafts[bx] = num_correct
+        predicts[last_accept_idx] = int(target_flat[last_accept_idx])
+    return predicts, accept_indices, num_correct_drafts
+
+
+class TestVerifyTreeGreedy(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def _run_and_check(
+        self,
+        candidates,
+        retrieve_index,
+        retrieve_next_token,
+        retrieve_next_sibling,
+        target_predict,
+        num_spec_step,
+    ):
+        bs, draft_token_num = candidates.shape
+        predicts = torch.full((bs * draft_token_num,), -1, dtype=torch.int32)
+        accept_indices = torch.full((bs, num_spec_step), -1, dtype=torch.int32)
+        num_correct_drafts = torch.empty((bs,), dtype=torch.int32)
+        torch.ops.sgl_kernel.verify_tree_greedy_cpu(
+            predicts,
+            accept_indices,
+            num_correct_drafts,
+            candidates,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            target_predict,
+        )
+        predicts_ref, accept_indices_ref, num_correct_drafts_ref = (
+            _ref_verify_tree_greedy(
+                candidates,
+                retrieve_index,
+                retrieve_next_token,
+                retrieve_next_sibling,
+                target_predict,
+                num_spec_step,
+            )
+        )
+        torch.testing.assert_close(predicts, predicts_ref, atol=0, rtol=0)
+        torch.testing.assert_close(accept_indices, accept_indices_ref, atol=0, rtol=0)
+        torch.testing.assert_close(
+            num_correct_drafts, num_correct_drafts_ref, atol=0, rtol=0
+        )
+        return predicts, accept_indices, num_correct_drafts
+
+    def test_verify_tree_greedy_chain_topk1(self):
+        # MTP config: topk=1, steps=3, draft_token_num=4 -> linear chain.
+        # Per row, force exactly `m` leading drafts to match the target so every
+        # accept count in [0, num_steps] is exercised.
+        num_steps, draft_token_num = 3, 4
+        vocab = 100
+        for bs in [1, 3]:
+            with self.subTest(bs=bs):
+                offs = (
+                    torch.arange(bs, dtype=torch.int64).unsqueeze(1) * draft_token_num
+                )
+                retrieve_index = (
+                    torch.arange(draft_token_num, dtype=torch.int64).repeat(bs, 1)
+                    + offs
+                )
+                retrieve_next_token = torch.tensor(
+                    [[1, 2, 3, -1]], dtype=torch.int64
+                ).repeat(bs, 1)
+                retrieve_next_sibling = torch.full(
+                    (bs, draft_token_num), -1, dtype=torch.int64
+                )
+                target_predict = torch.randint(
+                    0, vocab, (bs, draft_token_num), dtype=torch.int64
+                )
+                candidates = torch.randint(
+                    0, vocab, (bs, draft_token_num), dtype=torch.int64
+                )
+                forced_num_correct = [
+                    bx % (num_steps + 1) if bs > 1 else num_steps for bx in range(bs)
+                ]
+                for bx, m in enumerate(forced_num_correct):
+                    for i in range(m):
+                        candidates[bx, i + 1] = target_predict[bx, i]
+                    for i in range(m, draft_token_num - 1):
+                        candidates[bx, i + 1] = (target_predict[bx, i] + 1) % vocab
+                _, _, num_correct_drafts = self._run_and_check(
+                    candidates,
+                    retrieve_index,
+                    retrieve_next_token,
+                    retrieve_next_sibling,
+                    target_predict,
+                    num_steps + 1,
+                )
+                self.assertEqual(num_correct_drafts.tolist(), forced_num_correct)
+
+    def test_verify_tree_greedy_chain_hand_case(self):
+        # Deterministic chain: drafts 3, 4 match the target, draft 5 does not.
+        retrieve_index = torch.arange(4, dtype=torch.int64).unsqueeze(0)
+        retrieve_next_token = torch.tensor([[1, 2, 3, -1]], dtype=torch.int64)
+        retrieve_next_sibling = torch.full((1, 4), -1, dtype=torch.int64)
+        candidates = torch.tensor([[10, 3, 4, 5]], dtype=torch.int64)
+        target_predict = torch.tensor([[3, 4, 9, 9]], dtype=torch.int64)
+        predicts, accept_indices, num_correct_drafts = self._run_and_check(
+            candidates,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            target_predict,
+            4,
+        )
+        self.assertEqual(predicts.tolist(), [3, 4, 9, -1])
+        self.assertEqual(accept_indices.tolist(), [[0, 1, 2, -1]])
+        self.assertEqual(num_correct_drafts.tolist(), [2])
+
+    def test_verify_tree_greedy_upstream_golden(self):
+        # Golden fixture ported from the CUDA kernel UT
+        # sgl-kernel/tests/speculative/test_eagle_utils.py::test_verify_tree_greedy
+        # (device swapped to CPU); expected outputs are the CUDA kernel's.
+        candidates = torch.tensor(
+            [
+                [0, 1, 2, 3, 4, 5],
+                [7, 8, 9, 10, 11, 12],
+            ],
+            dtype=torch.int64,
+        )
+        retrieve_index = torch.tensor(
+            [
+                [0, 1, 2, 3, 4, 5],
+                [6, 7, 8, 9, 10, 11],
+            ],
+            dtype=torch.int64,
+        )
+        retrieve_next_token = torch.tensor(
+            [
+                [1, 2, -1, 4, 5, -1],
+                [4, 2, 3, -1, 5, -1],
+            ],
+            dtype=torch.int64,
+        )
+        retrieve_next_sibling = torch.tensor(
+            [
+                [-1, 3, -1, -1, -1, -1],
+                [-1, -1, -1, -1, 1, -1],
+            ],
+            dtype=torch.int64,
+        )
+
+        target_logits = torch.full((2, 6, 20), 1, dtype=torch.float32)
+        target_logits[0, 0, 3] = 10
+        target_logits[0, 3, 4] = 10
+        target_logits[0, 4, 5] = 10
+        target_logits[1, 0, 11] = 10
+        target_logits[1, 4, 12] = 10
+        for i in range(target_logits.shape[0]):
+            for j in range(target_logits.shape[1]):
+                if torch.max(target_logits[i][j]) < 10:
+                    target_logits[i][j][18] = 10
+        target_predict = torch.argmax(target_logits, dim=-1)
+
+        predicts, accept_indices, num_correct_drafts = self._run_and_check(
+            candidates,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+            target_predict,
+            4,  # num_spec_step
+        )
+        self.assertEqual(
+            predicts.tolist(), [3, -1, -1, 4, 5, 18, 11, -1, -1, -1, 12, 18]
+        )
+        self.assertEqual(accept_indices.tolist(), [[0, 3, 4, 5], [6, 10, 11, -1]])
+        self.assertEqual(num_correct_drafts.tolist(), [3, 2])
+
+    def test_verify_tree_greedy_tree_topk4(self):
+        # EAGLE config: topk=4, steps=3, draft_token_num=16
+        topk, num_steps, draft_token_num = 4, 3, 16
+        vocab = 3  # small vocab so accepts, rejects and sibling hops all happen
+        for bs in [1, 3]:
+            with self.subTest(bs=bs):
+                parent_list, selected_index = _gen_draft_tree(
+                    bs, topk, num_steps, draft_token_num
+                )
+                seq_lens = torch.randint(4, 32, (bs,), dtype=torch.int64)
+                _, _, retrieve_index, retrieve_next_token, retrieve_next_sibling = (
+                    _run_build_tree_kernel(
+                        parent_list,
+                        selected_index,
+                        seq_lens,
+                        topk,
+                        num_steps,
+                        draft_token_num,
+                        TreeMaskMode.QLEN_ONLY,
+                    )
+                )
+                # the random trees must branch, otherwise sibling traversal
+                # would go untested
+                self.assertGreater(int((retrieve_next_sibling != -1).sum()), 0)
+                candidates = torch.randint(
+                    0, vocab, (bs, draft_token_num), dtype=torch.int64
+                )
+                target_predict = torch.randint(
+                    0, vocab, (bs, draft_token_num), dtype=torch.int64
+                )
+                _, _, num_correct_drafts = self._run_and_check(
+                    candidates,
+                    retrieve_index,
+                    retrieve_next_token,
+                    retrieve_next_sibling,
+                    target_predict,
+                    num_steps + 1,
+                )
+                # deterministic with the fixed seed: the accept path is taken
+                self.assertGreater(int(num_correct_drafts.sum()), 0)
+
+
+class TestBuildTreeKernelEfficient(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def _check_against_reference(
+        self, parent_list, selected_index, seq_lens, topk, num_steps, draft_token_num
+    ):
+        bs = seq_lens.numel()
+        (
+            retrieve_index_ref,
+            retrieve_next_token_ref,
+            retrieve_next_sibling_ref,
+            positions_ref,
+            tree_mask_ref,
+        ) = _ref_build_tree(
+            parent_list, selected_index, seq_lens, topk, draft_token_num
+        )
+
+        # QLEN_ONLY
+        (
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+        ) = _run_build_tree_kernel(
+            parent_list,
+            selected_index,
+            seq_lens,
+            topk,
+            num_steps,
+            draft_token_num,
+            TreeMaskMode.QLEN_ONLY,
+        )
+        torch.testing.assert_close(retrieve_index, retrieve_index_ref, atol=0, rtol=0)
+        torch.testing.assert_close(
+            retrieve_next_token, retrieve_next_token_ref, atol=0, rtol=0
+        )
+        torch.testing.assert_close(
+            retrieve_next_sibling, retrieve_next_sibling_ref, atol=0, rtol=0
+        )
+        torch.testing.assert_close(positions, positions_ref, atol=0, rtol=0)
+        torch.testing.assert_close(
+            tree_mask.view(bs, draft_token_num, draft_token_num),
+            tree_mask_ref,
+            atol=0,
+            rtol=0,
+        )
+
+        # FULL_MASK: same retrieve_*/positions; the mask additionally carries the
+        # (all-true) committed-prefix columns per row.
+        (
+            full_mask,
+            positions_f,
+            retrieve_index_f,
+            retrieve_next_token_f,
+            retrieve_next_sibling_f,
+        ) = _run_build_tree_kernel(
+            parent_list,
+            selected_index,
+            seq_lens,
+            topk,
+            num_steps,
+            draft_token_num,
+            TreeMaskMode.FULL_MASK,
+        )
+        torch.testing.assert_close(retrieve_index_f, retrieve_index_ref, atol=0, rtol=0)
+        torch.testing.assert_close(
+            retrieve_next_token_f, retrieve_next_token_ref, atol=0, rtol=0
+        )
+        torch.testing.assert_close(
+            retrieve_next_sibling_f, retrieve_next_sibling_ref, atol=0, rtol=0
+        )
+        torch.testing.assert_close(positions_f, positions_ref, atol=0, rtol=0)
+        offset = 0
+        for bid in range(bs):
+            seq_len = int(seq_lens[bid])
+            chunk = full_mask[
+                offset : offset + draft_token_num * (seq_len + draft_token_num)
+            ].view(draft_token_num, seq_len + draft_token_num)
+            self.assertTrue(chunk[:, :seq_len].all().item())
+            torch.testing.assert_close(
+                chunk[:, seq_len:], tree_mask_ref[bid], atol=0, rtol=0
+            )
+            offset += draft_token_num * (seq_len + draft_token_num)
+
+        return tree_mask_ref, positions_ref
+
+    def test_build_tree_chain_topk1(self):
+        # MTP config: topk=1, steps=3, draft_token_num=4
+        num_steps, draft_token_num = 3, 4
+        bs = 2
+        seq_lens = torch.tensor([7, 12], dtype=torch.int64)
+        parent_list, selected_index = _topk1_chain_inputs(bs, num_steps)
+        tree_mask_ref, positions_ref = self._check_against_reference(
+            parent_list, selected_index, seq_lens, 1, num_steps, draft_token_num
+        )
+        # A chain must yield a causal (lower-triangular) mask and consecutive positions.
+        tril = torch.tril(
+            torch.ones(draft_token_num, draft_token_num, dtype=torch.bool)
+        )
+        for bid in range(bs):
+            torch.testing.assert_close(tree_mask_ref[bid], tril, atol=0, rtol=0)
+            self.assertEqual(
+                positions_ref[
+                    bid * draft_token_num : (bid + 1) * draft_token_num
+                ].tolist(),
+                [int(seq_lens[bid]) + i for i in range(draft_token_num)],
+            )
+
+    def test_build_tree_chain_mtp_step1(self):
+        # MTP steps=1: topk=1, draft_token_num=2, depth=1. parent_list is the
+        # empty (bs, 0) tensor organize_draft_results emits when there are no
+        # non-root parents; the kernel must accept it (regression guard).
+        num_steps, draft_token_num = 1, 2
+        bs = 2
+        seq_lens = torch.tensor([7, 12], dtype=torch.int64)
+        parent_list, selected_index = _topk1_chain_inputs(bs, num_steps)
+        self.assertEqual(parent_list.shape, (bs, 0))
+        tree_mask_ref, positions_ref = self._check_against_reference(
+            parent_list, selected_index, seq_lens, 1, num_steps, draft_token_num
+        )
+        tril = torch.tril(
+            torch.ones(draft_token_num, draft_token_num, dtype=torch.bool)
+        )
+        for bid in range(bs):
+            torch.testing.assert_close(tree_mask_ref[bid], tril, atol=0, rtol=0)
+
+    def test_build_tree_topk2_hand_case(self):
+        # topk=2, steps=2, draft_token_num=4. Nodes: 1 and 2 are children of the
+        # root; node 3 is a child of node 1.
+        topk, num_steps, draft_token_num = 2, 2, 4
+        parent_list = torch.tensor([[-1, 0, 1]], dtype=torch.int64)
+        selected_index = torch.tensor([[0, 1, 2]], dtype=torch.int64)
+        seq_lens = torch.tensor([5], dtype=torch.int64)
+        (
+            tree_mask,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+        ) = _run_build_tree_kernel(
+            parent_list,
+            selected_index,
+            seq_lens,
+            topk,
+            num_steps,
+            draft_token_num,
+            TreeMaskMode.QLEN_ONLY,
+        )
+        self.assertEqual(retrieve_index.tolist(), [[0, 1, 2, 3]])
+        self.assertEqual(retrieve_next_token.tolist(), [[1, 3, -1, -1]])
+        self.assertEqual(retrieve_next_sibling.tolist(), [[-1, 2, -1, -1]])
+        self.assertEqual(positions.tolist(), [5, 6, 6, 7])
+        self.assertEqual(
+            tree_mask.view(draft_token_num, draft_token_num).int().tolist(),
+            [
+                [1, 0, 0, 0],
+                [1, 1, 0, 0],
+                [1, 0, 1, 0],
+                [1, 1, 0, 1],
+            ],
+        )
+        # Cross-check the python reference on the same hand-built case.
+        self._check_against_reference(
+            parent_list, selected_index, seq_lens, topk, num_steps, draft_token_num
+        )
+
+    def test_build_tree_upstream_golden(self):
+        # Captured EAGLE trace ported from the CUDA UT
+        # test/registered/spec/utils/test_build_eagle_tree.py::TestBuildEagleTree::
+        # test_build_tree_kernel_efficient (device swapped to CPU). The inputs go
+        # through the same organize_draft_results; the expected positions and
+        # retrieve_* vectors are the CUDA kernel's.
+        score_list = [
+            torch.tensor(
+                [
+                    [[7.1127e-01, 2.8292e-01, 2.2995e-03, 1.7357e-03]],
+                    [[9.7476e-01, 2.2219e-02, 6.5031e-04, 1.3212e-04]],
+                ],
+                dtype=torch.float32,
+            ),
+            torch.tensor(
+                [
+                    [
+                        [6.9142e-01, 1.2863e-02, 1.6873e-03, 1.1871e-03],
+                        [2.4787e-01, 1.8818e-02, 1.4204e-02, 9.2235e-04],
+                        [2.2971e-03, 1.6700e-06, 1.8737e-07, 8.3146e-08],
+                        [1.2771e-03, 2.4374e-04, 1.7832e-04, 1.1947e-05],
+                    ],
+                    [
+                        [8.4832e-02, 6.6068e-02, 5.8304e-02, 5.7851e-02],
+                        [2.3616e-03, 1.1243e-03, 5.4368e-04, 2.7768e-04],
+                        [2.5286e-04, 1.5578e-04, 2.8817e-05, 1.2888e-05],
+                        [1.2834e-04, 2.5417e-06, 1.1279e-06, 1.6088e-08],
+                    ],
+                ],
+                dtype=torch.float32,
+            ),
+            torch.tensor(
+                [
+                    [
+                        [6.6438e-01, 2.6997e-02, 2.4236e-05, 4.0821e-06],
+                        [2.4402e-01, 2.8409e-03, 5.0935e-04, 2.9022e-04],
+                        [1.6178e-02, 2.0567e-03, 4.5892e-04, 3.0034e-05],
+                        [1.3023e-02, 5.0497e-04, 3.6371e-04, 8.7750e-05],
+                    ],
+                    [
+                        [2.3263e-02, 2.0054e-02, 9.3990e-03, 2.7783e-03],
+                        [6.4156e-02, 5.5506e-04, 1.0429e-04, 9.7211e-05],
+                        [4.9950e-02, 5.0630e-03, 9.0068e-04, 3.3656e-04],
+                        [7.5817e-03, 8.5731e-04, 6.9972e-04, 6.0793e-04],
+                    ],
+                ],
+                dtype=torch.float32,
+            ),
+            torch.tensor(
+                [
+                    [
+                        [6.6420e-01, 1.0525e-04, 6.5864e-05, 1.2253e-06],
+                        [1.3019e-01, 1.0461e-01, 5.2083e-03, 1.6777e-03],
+                        [2.0103e-02, 6.7335e-03, 1.2625e-04, 1.0364e-05],
+                        [1.5142e-02, 7.0819e-04, 9.6595e-05, 8.7951e-05],
+                    ],
+                    [
+                        [5.8608e-02, 1.8840e-03, 7.8535e-04, 4.4400e-04],
+                        [1.2185e-02, 2.0684e-03, 1.7418e-03, 1.4327e-03],
+                        [6.2455e-03, 6.1487e-03, 2.6862e-03, 1.8034e-03],
+                        [1.8590e-03, 1.6151e-03, 1.2481e-03, 3.6038e-04],
+                    ],
+                ],
+                dtype=torch.float32,
+            ),
+        ]
+        token_list = [
+            torch.tensor(
+                [[29896, 29906, 29900, 29945], [13, 2, 29871, 28956]],
+                dtype=torch.int64,
+            ),
+            # fmt: off
+            torch.tensor(
+                [
+                    [29889, 29974, 29945, 29900, 29974, 29922, 29930, 29958,
+                     29889, 29974, 29930, 29945, 29974, 29922, 29930, 29958],
+                    [22550, 4136, 16492, 8439, 29871, 2, 3001, 13,
+                     2, 13, 29906, 29946, 2, 13, 29871, 259],
+                ],
+            ),
+            torch.tensor(
+                [
+                    [29946, 29945, 29953, 29906, 29896, 29945, 29900, 29906,
+                     29896, 29945, 29906, 29953, 29896, 29945, 29906, 29946],
+                    [29871, 2, 29901, 29889, 29871, 2, 395, 259,
+                     29901, 29871, 2, 29889, 3001, 1234, 7146, 2186],
+                ],
+            ),
+            torch.tensor(
+                [
+                    [29946, 29974, 29945, 29930, 29889, 29922, 29974, 29930,
+                     29974, 29946, 29930, 29922, 29889, 29974, 29945, 29922],
+                    [29941, 29906, 2, 29946, 29871, 450, 319, 14990,
+                     29946, 29941, 2, 29906, 29871, 2, 3001, 13],
+                ],
+            ),
+            # fmt: on
+        ]
+        parents_list = [
+            torch.tensor([[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]], dtype=torch.int64),
+            torch.tensor([[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64),
+            torch.tensor([[20, 24, 21, 28], [24, 28, 20, 21]], dtype=torch.int64),
+            torch.tensor([[36, 40, 41, 44], [36, 40, 44, 45]], dtype=torch.int64),
+        ]
+        seq_lens = torch.tensor([5, 10], dtype=torch.int64)
+        topk, depth, draft_token_num = 4, 4, 8
+
+        parent_list, selected_index, draft_tokens = organize_draft_results(
+            score_list, token_list, parents_list, draft_token_num
+        )
+        # The CUDA golden interleaves the bonus tokens [29974, 13] at positions
+        # 0 and 8; on CPU that concat happens outside the kernel, so compare the
+        # drafts only.
+        self.assertEqual(
+            draft_tokens.tolist(),
+            [
+                [29896, 29906, 29889, 29974, 29946, 29896, 29946],
+                [13, 22550, 4136, 16492, 8439, 29871, 29941],
+            ],
+        )
+
+        (
+            _,
+            positions,
+            retrieve_index,
+            retrieve_next_token,
+            retrieve_next_sibling,
+        ) = _run_build_tree_kernel(
+            parent_list,
+            selected_index,
+            seq_lens,
+            topk,
+            depth,
+            draft_token_num,
+            TreeMaskMode.QLEN_ONLY,
+        )
+        self.assertEqual(
+            positions.tolist(),
+            [5, 6, 6, 7, 7, 8, 8, 9, 10, 11, 12, 12, 12, 12, 13, 14],
+        )
+        self.assertEqual(
+            retrieve_index.tolist(),
+            [
+                [0, 1, 2, 3, 4, 5, 6, 7],
+                [8, 9, 10, 11, 12, 13, 14, 15],
+            ],
+        )
+        self.assertEqual(
+            retrieve_next_token.tolist(),
+            [
+                [1, 3, 4, 5, 6, 7, -1, -1],
+                [1, 2, -1, 6, -1, -1, 7, -1],
+            ],
+        )
+        self.assertEqual(
+            retrieve_next_sibling.tolist(),
+            [
+                [-1, 2, -1, -1, -1, -1, -1, -1],
+                [-1, -1, 3, 4, 5, -1, -1, -1],
+            ],
+        )
+        # Cross-check the python reference (incl. tree masks, which the CUDA UT
+        # does not assert) on the same trace.
+        self._check_against_reference(
+            parent_list, selected_index, seq_lens, topk, depth, draft_token_num
+        )
+
+    def test_build_tree_topk4_random(self):
+        # EAGLE config: topk=4, steps=3, draft_token_num=16
+        topk, num_steps, draft_token_num = 4, 3, 16
+        bs = 3
+        seq_lens = torch.tensor([9, 3, 21], dtype=torch.int64)
+        parent_list, selected_index = _gen_draft_tree(
+            bs, topk, num_steps, draft_token_num
+        )
+        tree_mask_ref, positions_ref = self._check_against_reference(
+            parent_list, selected_index, seq_lens, topk, num_steps, draft_token_num
+        )
+        # Invariants: causality (mask support only on self + earlier nodes) and
+        # node depth bounded by the number of speculative steps.
+        for bid in range(bs):
+            mask = tree_mask_ref[bid]
+            self.assertTrue(mask.diagonal().all().item())
+            self.assertTrue(torch.triu(mask.int(), diagonal=1).sum().item() == 0)
+            depths = positions_ref[
+                bid * draft_token_num : (bid + 1) * draft_token_num
+            ] - int(seq_lens[bid])
+            self.assertEqual(int(depths[0]), 0)
+            self.assertTrue(bool((depths[1:] >= 1).all()))
+            self.assertTrue(bool((depths <= num_steps).all()))
+            # row depth == number of visible draft tokens (self + ancestors) - 1
+            torch.testing.assert_close(
+                depths, mask.sum(dim=1).to(torch.int64) - 1, atol=0, rtol=0
+            )
+
+
+class TestReconstructIndicesFromTreeMask(CustomTestCase):
+    """reconstruct_indices_from_tree_mask on CPU (NGRAM verify metadata),
+    called through the device-dispatching sgl_kernel wrapper."""
+
+    def test_tree_and_chain(self):
+        from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
+
+        bs, draft_token_num = 2, 4
+        seq_lens = torch.tensor([12, 5], dtype=torch.int64)
+        # Request 0: root(0) -> {1, 2}, 2 -> 3 (golden case from
+        # sgl-kernel/tests/speculative/test_ngram_utils.py).
+        # Request 1: plain chain 0 -> 1 -> 2 -> 3.
+        tree_mask = torch.tensor(
+            # fmt: off
+            [
+                1, 0, 0, 0,
+                1, 1, 0, 0,
+                1, 0, 1, 0,
+                1, 0, 1, 1,
+
+                1, 0, 0, 0,
+                1, 1, 0, 0,
+                1, 1, 1, 0,
+                1, 1, 1, 1,
+            ],
+            # fmt: on
+            dtype=torch.bool,
+        )
+        retrieve_index = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+        retrieve_next_token = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+        retrieve_next_sibling = torch.full((bs, draft_token_num), -1, dtype=torch.int64)
+        positions = torch.empty((bs * draft_token_num,), dtype=torch.int64)
+
+        reconstruct_indices_from_tree_mask(
+            tree_mask,
+            seq_lens,
+            positions,  # mutable
+            retrieve_index,  # mutable
+            retrieve_next_token,  # mutable
+            retrieve_next_sibling,  # mutable
+            bs,
+            draft_token_num,
+        )
+
+        self.assertEqual(retrieve_index.tolist(), [[0, 1, 2, 3], [4, 5, 6, 7]])
+        self.assertEqual(retrieve_next_token.tolist(), [[1, -1, 3, -1], [1, 2, 3, -1]])
+        self.assertEqual(
+            retrieve_next_sibling.tolist(), [[-1, 2, -1, -1], [-1, -1, -1, -1]]
+        )
+        self.assertEqual(positions.tolist(), [12, 13, 13, 14, 5, 6, 7, 8])
+
+
+class TestFillKernels(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def test_fill_bonus_tokens(self):
+        bs, accept_stride = 4, 4
+        accept_tokens = torch.randint(0, 1000, (bs, accept_stride), dtype=torch.int32)
+        accept_lens = torch.tensor([1, 4, 2, 3], dtype=torch.int32)
+        bonus_tokens = torch.empty((bs,), dtype=torch.int32)
+        torch.ops.sgl_kernel.fill_bonus_tokens_cpu(
+            accept_tokens, accept_lens, bonus_tokens, accept_stride
+        )
+        bonus_tokens_ref = accept_tokens[
+            torch.arange(bs), accept_lens.to(torch.int64) - 1
+        ]
+        torch.testing.assert_close(bonus_tokens, bonus_tokens_ref, atol=0, rtol=0)
+
+    def test_fill_accept_out_cache_loc(self):
+        bs, num_spec_step, draft_token_num = 3, 4, 16
+        accept_indices = torch.tensor(
+            [
+                [0, 1, -1, -1],
+                [16, 18, 21, -1],
+                [32, -1, -1, -1],
+            ],
+            dtype=torch.int32,
+        )
+        out_cache_loc = torch.randperm(4096, dtype=torch.int64)[: bs * draft_token_num]
+        size = bs * num_spec_step
+        accept_out_cache_loc = torch.zeros(size, dtype=torch.int64)
+        torch.ops.sgl_kernel.fill_accept_out_cache_loc_cpu(
+            accept_indices, out_cache_loc, accept_out_cache_loc
+        )
+        valid = accept_indices.flatten()[accept_indices.flatten() > -1].to(torch.int64)
+        ref = torch.zeros(size, dtype=torch.int64)
+        ref[: valid.numel()] = out_cache_loc[valid]
+        torch.testing.assert_close(accept_out_cache_loc, ref, atol=0, rtol=0)
+
+
+class TestAssignCacheLocKernels(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def test_assign_draft_cache_locs_contiguous(self):
+        bs, topk, num_steps, pool_len, num_reqs = 3, 2, 3, 32, 5
+        req_pool_indices = torch.tensor([4, 0, 2], dtype=torch.int64)
+        req_to_token = torch.randint(0, 10000, (num_reqs, pool_len), dtype=torch.int32)
+        seq_lens = torch.tensor([5, 0, 17], dtype=torch.int64)
+        out_cache_loc = torch.empty((bs * topk * num_steps,), dtype=torch.int64)
+        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,
+        )
+        copy_len = topk * num_steps
+        ref = torch.empty_like(out_cache_loc)
+        for i in range(bs):
+            start = int(seq_lens[i])
+            ref[i * copy_len : (i + 1) * copy_len] = req_to_token[
+                int(req_pool_indices[i]), start : start + copy_len
+            ].to(torch.int64)
+        torch.testing.assert_close(out_cache_loc, ref, atol=0, rtol=0)
+
+        # int32 index tensors are accepted natively (per-tensor index dispatch)
+        out_cache_loc_i32 = torch.empty_like(out_cache_loc)
+        torch.ops.sgl_kernel.assign_draft_cache_locs_contiguous_cpu(
+            req_pool_indices,
+            req_to_token,
+            seq_lens.to(torch.int32),
+            out_cache_loc_i32,
+            pool_len,
+            topk,
+            num_steps,
+        )
+        torch.testing.assert_close(out_cache_loc_i32, ref, atol=0, rtol=0)
+
+        # dtype contract: int32 req_to_token, int64 out_cache_loc (TORCH_CHECK)
+        with self.assertRaises(RuntimeError):
+            torch.ops.sgl_kernel.assign_draft_cache_locs_contiguous_cpu(
+                req_pool_indices,
+                req_to_token.to(torch.int64),
+                seq_lens,
+                out_cache_loc,
+                pool_len,
+                topk,
+                num_steps,
+            )
+        with self.assertRaises(RuntimeError):
+            torch.ops.sgl_kernel.assign_draft_cache_locs_contiguous_cpu(
+                req_pool_indices,
+                req_to_token,
+                seq_lens,
+                out_cache_loc[:-1],  # wrong numel
+                pool_len,
+                topk,
+                num_steps,
+            )
+
+    def test_assign_extend_cache_locs(self):
+        bs, pool_len, num_reqs = 4, 32, 5
+        req_pool_indices = torch.tensor([1, 3, 0, 2], dtype=torch.int64)
+        req_to_token = torch.randint(0, 10000, (num_reqs, pool_len), dtype=torch.int32)
+        # Last request has start == end (zero-length range, nothing copied).
+        start_offset = torch.tensor([4, 0, 11, 6], dtype=torch.int64)
+        end_offset = torch.tensor([9, 7, 12, 6], dtype=torch.int64)
+        total = int((end_offset - start_offset).sum())
+        out_cache_loc = torch.empty((total,), dtype=torch.int64)
+        torch.ops.sgl_kernel.assign_extend_cache_locs_cpu(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            pool_len,
+        )
+        ref = torch.cat(
+            [
+                req_to_token[
+                    int(req_pool_indices[i]),
+                    int(start_offset[i]) : int(end_offset[i]),
+                ].to(torch.int64)
+                for i in range(bs)
+            ]
+        )
+        torch.testing.assert_close(out_cache_loc, ref, atol=0, rtol=0)
+
+    def test_assign_req_to_token_pool(self):
+        bs, pool_len, num_reqs = 3, 32, 5
+        req_pool_indices = torch.tensor([2, 4, 1], dtype=torch.int32)
+        req_to_token = torch.zeros((num_reqs, pool_len), dtype=torch.int32)
+        req_to_token_ref = req_to_token.clone()
+        start_offset = torch.tensor([0, 3, 8], dtype=torch.int32)
+        end_offset = torch.tensor([6, 3, 15], dtype=torch.int32)
+        total = int((end_offset - start_offset).sum())
+        out_cache_loc = torch.randperm(4096, dtype=torch.int64)[:total]
+        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,
+        )
+        offset = 0
+        for i in range(bs):
+            length = int(end_offset[i]) - int(start_offset[i])
+            req_to_token_ref[
+                int(req_pool_indices[i]),
+                int(start_offset[i]) : int(end_offset[i]),
+            ] = out_cache_loc[offset : offset + length].to(torch.int32)
+            offset += length
+        torch.testing.assert_close(req_to_token, req_to_token_ref, atol=0, rtol=0)
+
+
+class TestCopyAllLayerKvCache(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def _make_buffers(self, num_slots=32, layer_num=3, head_num=2):
+        # K and V rows deliberately differ in width (v_head_dim < head_dim) so
+        # the per-buffer byte strides differ, like a real MHA pool.
+        k_buffer = [
+            torch.randn(num_slots, head_num, 64, dtype=torch.bfloat16)
+            for _ in range(layer_num)
+        ]
+        v_buffer = [
+            torch.randn(num_slots, head_num, 32, dtype=torch.bfloat16)
+            for _ in range(layer_num)
+        ]
+        return k_buffer + v_buffer
+
+    def _run_and_check(self, buffers, tgt_loc, src_loc):
+        # Reference: gather-then-scatter semantics (advanced indexing gathers
+        # buf[src] before assigning)
+        refs = []
+        for buf in buffers:
+            ref = buf.clone()
+            ref[tgt_loc.long()] = buf[src_loc.long()]
+            refs.append(ref)
+
+        data_ptrs = torch.tensor([b.data_ptr() for b in buffers], dtype=torch.uint64)
+        strides = torch.tensor(
+            [b[0].numel() * b.dtype.itemsize for b in buffers], dtype=torch.int64
+        )
+        torch.ops.sgl_kernel.copy_all_layer_kv_cache_cpu(
+            data_ptrs, strides, tgt_loc, src_loc
+        )
+        for buf, ref in zip(buffers, refs):
+            torch.testing.assert_close(buf, ref, atol=0, rtol=0)
+
+    def test_disjoint_locs(self):
+        buffers = self._make_buffers()
+        tgt_loc = torch.tensor([1, 5, 9, 13], dtype=torch.int64)
+        src_loc = torch.tensor([17, 21, 25, 29], dtype=torch.int64)
+        self._run_and_check(buffers, tgt_loc, src_loc)
+
+    def test_overlapping_locs_reverse_shift(self):
+        # tgt[i] == src[i + 1]: a sequential row-by-row copy would read
+        # already-overwritten rows; the kernel must stage all sources first.
+        buffers = self._make_buffers()
+        tgt_loc = torch.tensor([1, 2, 3, 4], dtype=torch.int64)
+        src_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
+        self._run_and_check(buffers, tgt_loc, src_loc)
+
+    def test_self_copy_noop(self):
+        # The pool warmup issues loc -> loc copies; buffers must be unchanged.
+        buffers = self._make_buffers()
+        loc = torch.zeros(1, dtype=torch.int64)
+        self._run_and_check(buffers, loc, loc)
+
+    def test_int32_locs(self):
+        buffers = self._make_buffers()
+        tgt_loc = torch.tensor([0, 3], dtype=torch.int32)
+        src_loc = torch.tensor([8, 11], dtype=torch.int32)
+        self._run_and_check(buffers, tgt_loc, src_loc)
+
+
+class TestExtendAttentionTreeMask(CustomTestCase):
+    H_Q, H_KV, D, DV = 4, 2, 64, 64
+
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def _build_target_verify_inputs(self, prefix_lens, draft_token_num):
+        # Mimic intel_amx_backend.forward_extend in TARGET_VERIFY mode: kv cache
+        # for the draft tokens is written first (set_kv_buffer), seq_lens are the
+        # committed lens + draft_token_num and extend lens are inferred uniform.
+        dtype = torch.bfloat16
+        bs = len(prefix_lens)
+        prefix_lens = torch.tensor(prefix_lens, dtype=torch.int64)
+        seq_lens = prefix_lens + draft_token_num  # seq_lens += draft_token_num
+        total_tokens = int(seq_lens.sum())
+
+        req_to_token = torch.zeros((bs, int(seq_lens.max())), dtype=torch.int32)
+        start = 0
+        for i in range(bs):
+            req_to_token[i, : int(seq_lens[i])] = torch.arange(
+                start, start + int(seq_lens[i]), dtype=torch.int32
+            )
+            start += int(seq_lens[i])
+
+        k_buffer = torch.randn((total_tokens, self.H_KV, self.D), dtype=dtype)
+        v_buffer = torch.randn((total_tokens, self.H_KV, self.DV), dtype=dtype)
+
+        # extend (draft) tokens were already written into the buffers
+        extend_token_num = bs * draft_token_num
+        k_extend = torch.empty((extend_token_num, self.H_KV, self.D), dtype=dtype)
+        v_extend = torch.empty((extend_token_num, self.H_KV, self.DV), dtype=dtype)
+        q_extend = torch.randn((extend_token_num, self.H_Q, self.D), dtype=dtype)
+        for i in range(bs):
+            buf_start = int(seq_lens[:i].sum()) + int(prefix_lens[i])
+            buf_end = buf_start + draft_token_num
+            k_extend[i * draft_token_num : (i + 1) * draft_token_num] = k_buffer[
+                buf_start:buf_end
+            ]
+            v_extend[i * draft_token_num : (i + 1) * draft_token_num] = v_buffer[
+                buf_start:buf_end
+            ]
+
+        req_pool_indices = torch.arange(bs, dtype=torch.int64)
+        extend_seq_lens = torch.full((bs,), draft_token_num, dtype=torch.int32)
+        extend_start_loc = torch.zeros((bs,), dtype=torch.int32)
+        if bs > 1:
+            extend_start_loc[1:] = torch.cumsum(extend_seq_lens[:-1], dim=0)
+        return (
+            q_extend,
+            k_extend,
+            v_extend,
+            k_buffer,
+            v_buffer,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+            extend_seq_lens,
+            extend_start_loc,
+        )
+
+    def _run_kernel(self, inputs, draft_token_num, tree_mask, with_trailing_arg=True):
+        (
+            q_extend,
+            k_extend,
+            v_extend,
+            k_buffer,
+            v_buffer,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+            extend_seq_lens,
+            extend_start_loc,
+        ) = inputs
+        o_extend = torch.empty(
+            (q_extend.shape[0], self.H_Q, self.DV), dtype=q_extend.dtype
+        )
+        sm_scale = 1.0 / (self.D**0.5)
+        args = [
+            q_extend,
+            k_extend,
+            v_extend,
+            o_extend,
+            k_buffer,
+            v_buffer,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+            extend_seq_lens,
+            extend_start_loc,
+            draft_token_num,  # max_len_extend
+            sm_scale,
+            0.0,  # logit_cap
+            False,  # is_cross_attn
+            0,  # sliding_window_size (layer.sliding_window_size + 1)
+            None,  # encoder_lens
+            None,  # sinks
+        ]
+        if with_trailing_arg:
+            args.append(tree_mask)
+        torch.ops.sgl_kernel.extend_attention_cpu(*args)
+        return o_extend
+
+    def _ref_sdpa(self, inputs, draft_token_num, qlen_masks):
+        (
+            q_extend,
+            _,
+            _,
+            k_buffer,
+            v_buffer,
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+            _,
+            _,
+        ) = inputs
+        bs = seq_lens.shape[0]
+        sm_scale = 1.0 / (self.D**0.5)
+        out = torch.empty(
+            (bs * draft_token_num, self.H_Q, self.DV), dtype=torch.float32
+        )
+        for i in range(bs):
+            kv_len = int(seq_lens[i])
+            tokens = req_to_token[int(req_pool_indices[i]), :kv_len].to(torch.int64)
+            keys = k_buffer[tokens].float().movedim(0, 1)  # [H_KV, kv_len, D]
+            values = v_buffer[tokens].float().movedim(0, 1)
+            queries = (
+                q_extend[i * draft_token_num : (i + 1) * draft_token_num]
+                .float()
+                .movedim(0, 1)
+            )  # [H_Q, qlen, D]
+            attn_mask = torch.cat(
+                [
+                    torch.ones(
+                        draft_token_num, kv_len - draft_token_num, dtype=torch.bool
+                    ),
+                    qlen_masks[i],
+                ],
+                dim=1,
+            )
+            o = F.scaled_dot_product_attention(
+                queries.unsqueeze(0),
+                keys.unsqueeze(0),
+                values.unsqueeze(0),
+                attn_mask=attn_mask,
+                scale=sm_scale,
+                enable_gqa=True,
+            ).squeeze(0)
+            out[i * draft_token_num : (i + 1) * draft_token_num] = o.movedim(0, 1)
+        return out
+
+    def test_no_tree_mask_matches_legacy_call(self):
+        # (a) with and without the trailing optional argument must be bitwise equal
+        draft_token_num = 16
+        inputs = self._build_target_verify_inputs([13, 7], draft_token_num)
+        o_legacy = self._run_kernel(
+            inputs, draft_token_num, None, with_trailing_arg=False
+        )
+        o_none = self._run_kernel(inputs, draft_token_num, None, with_trailing_arg=True)
+        self.assertTrue(torch.equal(o_legacy, o_none))
+
+    def test_chain_topk1_matches_causal_sdpa(self):
+        # (b) topk=1 chain (no tree mask passed): plain causal masking
+        draft_token_num = 4
+        inputs = self._build_target_verify_inputs([9, 5], draft_token_num)
+        o_extend = self._run_kernel(inputs, draft_token_num, None)
+        tril = torch.tril(
+            torch.ones(draft_token_num, draft_token_num, dtype=torch.bool)
+        )
+        o_ref = self._ref_sdpa(inputs, draft_token_num, [tril] * len(inputs[7]))
+        atol = rtol = precision[torch.bfloat16]
+        torch.testing.assert_close(o_extend.float(), o_ref, atol=atol, rtol=rtol)
+
+    def test_tree_topk4_matches_tree_sdpa(self):
+        # (c) topk=4 tree with a QLEN_ONLY mask from build_tree_kernel_efficient_cpu
+        topk, num_steps, draft_token_num = 4, 3, 16
+        bs = 2
+        prefix_lens = [13, 7]
+        inputs = self._build_target_verify_inputs(prefix_lens, draft_token_num)
+        parent_list, selected_index = _gen_draft_tree(
+            bs, topk, num_steps, draft_token_num
+        )
+        tree_mask, _, _, _, _ = _run_build_tree_kernel(
+            parent_list,
+            selected_index,
+            torch.tensor(prefix_lens, dtype=torch.int64),
+            topk,
+            num_steps,
+            draft_token_num,
+            TreeMaskMode.QLEN_ONLY,
+        )
+        o_extend = self._run_kernel(inputs, draft_token_num, tree_mask)
+        qlen_masks = tree_mask.view(bs, draft_token_num, draft_token_num)
+        o_ref = self._ref_sdpa(
+            inputs, draft_token_num, [qlen_masks[i] for i in range(bs)]
+        )
+        # 2x the shared bf16 tolerance: the sparse tree mask leaves few terms
+        # per row, so single-element bf16 rounding differences dominate.
+        torch.testing.assert_close(o_extend.float(), o_ref, atol=2e-2, rtol=2e-2)
+
+
+class TestMultiLayerEagleStateKernels(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def test_rotate_input_ids(self):
+        extend_seq_lens = torch.tensor([4, 1, 3], dtype=torch.int64)
+        extend_start_loc = torch.tensor([0, 4, 5], dtype=torch.int64)
+        topk_index = torch.tensor([100, 200, 300], dtype=torch.int64)
+        total = int(extend_seq_lens.sum())
+        input_ids_init = torch.arange(10, 10 + total, dtype=torch.int64)
+
+        # without select_index: shift each request left by 1, append the new token
+        input_ids = input_ids_init.clone()
+        torch.ops.sgl_kernel.rotate_input_ids_cpu(
+            input_ids, extend_start_loc, extend_seq_lens, topk_index, None
+        )
+        ref = input_ids_init.clone()
+        for i in range(extend_seq_lens.numel()):
+            start = int(extend_start_loc[i])
+            seq_len = int(extend_seq_lens[i])
+            ref[start : start + seq_len - 1] = input_ids_init[
+                start + 1 : start + seq_len
+            ]
+            ref[start + seq_len - 1] = topk_index[i]
+        torch.testing.assert_close(input_ids, ref, atol=0, rtol=0)
+
+        # int32 extend_* tensors are accepted natively (ForwardBatch supplies
+        # int32; the kernel dispatches on the index dtype, no cast needed)
+        input_ids = input_ids_init.clone()
+        torch.ops.sgl_kernel.rotate_input_ids_cpu(
+            input_ids,
+            extend_start_loc.to(torch.int32),
+            extend_seq_lens.to(torch.int32),
+            topk_index,
+            None,
+        )
+        torch.testing.assert_close(input_ids, ref, atol=0, rtol=0)
+
+        # with select_index: the new token lands at the given global slots
+        select_index = torch.tensor([2, 4, 7], dtype=torch.int64)
+        input_ids = input_ids_init.clone()
+        torch.ops.sgl_kernel.rotate_input_ids_cpu(
+            input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index
+        )
+        ref = input_ids_init.clone()
+        for i in range(extend_seq_lens.numel()):
+            start = int(extend_start_loc[i])
+            seq_len = int(extend_seq_lens[i])
+            ref[start : start + seq_len - 1] = input_ids_init[
+                start + 1 : start + seq_len
+            ]
+        for i in range(extend_seq_lens.numel()):
+            ref[int(select_index[i])] = topk_index[i]
+        torch.testing.assert_close(input_ids, ref, atol=0, rtol=0)
+
+        # int64 contract: int32 input_ids must be rejected
+        with self.assertRaises(RuntimeError):
+            torch.ops.sgl_kernel.rotate_input_ids_cpu(
+                input_ids_init.to(torch.int32),
+                extend_start_loc,
+                extend_seq_lens,
+                topk_index,
+                None,
+            )
+
+
+class TestBuildDraftDecodeMetadata(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(1234)
+
+    def _ref_build_metadata(
+        self, req_to_token, req_pool_indices, seq_lens, topk, num_steps
+    ):
+        """Pure-python reference. Draft row b*topk+tk holds the committed prefix
+        followed by candidate tk's draft slots, which live contiguously in the
+        source row at sl + tk*num_steps + s (assign_draft_cache_locs_contiguous
+        layout); only the first sl + num_steps entries per row are defined.
+        """
+        rows = []
+        for b in range(seq_lens.numel()):
+            src_row = req_to_token[int(req_pool_indices[b])]
+            sl = int(seq_lens[b])
+            for tk in range(topk):
+                draft_start = sl + tk * num_steps
+                rows.append(
+                    torch.cat(
+                        [src_row[:sl], src_row[draft_start : draft_start + num_steps]]
+                    )
+                )
+        return rows
+
+    def _run_and_check(self, seq_lens, topk, num_steps, pool_len, num_reqs):
+        num_seqs = seq_lens.numel()
+        # distinct slot values so any misplaced copy is caught exactly
+        req_to_token = (
+            torch.randperm(num_reqs * pool_len).to(torch.int32).view(num_reqs, pool_len)
+        )
+        req_pool_indices = torch.randperm(num_reqs, dtype=torch.int64)[:num_seqs]
+        req_to_token_draft = torch.ops.sgl_kernel.build_draft_decode_metadata_cpu(
+            req_to_token,
+            req_pool_indices,
+            seq_lens,
+            topk,
+            num_steps,
+            pool_len,
+        )
+        self.assertEqual(
+            req_to_token_draft.shape, torch.Size([num_seqs * topk, pool_len])
+        )
+        self.assertEqual(req_to_token_draft.dtype, req_to_token.dtype)
+        ref_rows = self._ref_build_metadata(
+            req_to_token, req_pool_indices, seq_lens, topk, num_steps
+        )
+        for b in range(num_seqs):
+            sl = int(seq_lens[b])
+            for tk in range(topk):
+                flat = b * topk + tk
+                # Only the first sl + num_steps entries are the kernel's
+                # contract: the output is allocated uninitialized and the
+                # draft-decode consumer reads at most seq_len + num_steps
+                # slots per candidate, so the row tail stays unasserted.
+                torch.testing.assert_close(
+                    req_to_token_draft[flat, : sl + num_steps],
+                    ref_rows[flat],
+                    atol=0,
+                    rtol=0,
+                )
+
+    def test_build_metadata_topk1_chain(self):
+        # MTP config: topk=1 -> the single candidate's drafts sit right after
+        # the prefix, so each draft row is the source row's first
+        # sl + num_steps slots verbatim.
+        seq_lens = torch.tensor([5, 0, 17], dtype=torch.int64)
+        for num_steps in [1, 2, 3]:
+            with self.subTest(num_steps=num_steps):
+                self._run_and_check(seq_lens, 1, num_steps, pool_len=32, num_reqs=5)
+
+    def test_build_metadata_topk4(self):
+        # EAGLE config: topk=4 -> candidate tk's drafts come from the strided
+        # source range [sl + tk*num_steps, sl + (tk+1)*num_steps) but always
+        # land right after the prefix in the expanded row.
+        seq_lens = torch.tensor([9, 0, 3, 21], dtype=torch.int64)
+        for num_steps in [1, 2, 3]:
+            with self.subTest(num_steps=num_steps):
+                self._run_and_check(seq_lens, 4, num_steps, pool_len=64, num_reqs=6)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py
index 25f3384d0..c7bae0a9a 100644
--- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py
+++ b/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):
diff --git a/test/registered/unit/spec/test_spec_cpu_overlap_constraint.py b/test/registered/unit/spec/test_spec_cpu_overlap_constraint.py
new file mode 100644
index 000000000..6e21a69be
--- /dev/null
+++ b/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()