[CPU][Spec] Enable DFLASH speculative decoding on Xeon (intel_amx)

python/sglang/srt/arg_groups/speculative_hook.py

@@ -186,11 +186,13 @@ def _handle_dflash(server_args: ServerArgs) -> None:
     cfg = resolving_view(server_args)
     from sglang.srt.arg_groups.overrides import resolved_view
 
-    if not (cfg.device.startswith("cuda") or cfg.device == "npu"):
+    if not (cfg.device.startswith("cuda") or cfg.device in ("npu", "cpu")):
         raise ValueError(
-            "DFLASH speculative decoding only supports CUDA and NPU devices."
+            "DFLASH speculative decoding only supports CUDA, NPU and CPU devices."
         )
 
+    _disable_overlap_schedule_for_cpu(server_args)
+
     if resolved_view(server_args).enable_dp_attention:
         raise ValueError(
             "Currently DFLASH speculative decoding does not support dp attention."
@@ -573,9 +575,14 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
         "triton",
         "trtllm_mha",
         "ascend",
+        "intel_amx",
     )
-    # Use triton on ROCm (no FlashInfer), flashinfer on CUDA.
-    fallback_backend = "triton" if is_hip() else "flashinfer"
+    # Use triton on ROCm (no FlashInfer), flashinfer on CUDA, and the AMX
+    # kernels on CPU.
+    if cfg.device == "cpu":
+        fallback_backend = "intel_amx"
+    else:
+        fallback_backend = "triton" if is_hip() else "flashinfer"
 
     draft_backend = cfg.speculative_draft_attention_backend
     if draft_backend is None:

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

@@ -5,6 +5,7 @@
 import torch
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -61,6 +62,10 @@ def __init__(self, model_runner: ModelRunner):
 
         self._attn_logits_buffers: dict[tuple[int, int], torch.Tensor] = {}
 
+        # All-visible qlen masks for non-causal (ENCODER_ONLY) spec blocks,
+        # keyed by (batch_size, draft_token_num).
+        self._non_causal_masks: dict[tuple[int, int], torch.Tensor] = {}
+
         # speculative decoding params
         self.num_draft_tokens = get_spec().speculative_num_draft_tokens
 
@@ -108,7 +113,7 @@ def _build_extend_metadata(self, forward_batch: ForwardBatch):
             # 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.
+            # (irregular tree); both need the mask. DFLASH is a chain (== 1).
             if spec_info.tree_topk != 1:
                 custom_mask = spec_info.custom_mask
                 if custom_mask is not None and custom_mask.numel() > 0:
@@ -214,6 +219,19 @@ def forward_extend(
         # verify batches carry no extend_* fields; see _build_extend_metadata).
         seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata
 
+        if (
+            tree_mask is None
+            and not layer.is_cross_attention
+            and layer.attn_type == AttentionType.ENCODER_ONLY
+            and forward_batch.forward_mode.is_target_verify()
+        ):
+            # The kernel's implicit mask is causal. A non-causal layer (DFLASH
+            # draft blocks) needs every query in the block to see every key in
+            # it, which an all-visible qlen mask expresses.
+            tree_mask = self._get_non_causal_qlen_mask(
+                forward_batch.batch_size, forward_batch.spec_info.draft_token_num
+            )
+
         _, max_extend_len = self.forward_metadata
         if seq_lens.dtype != torch.int64:
             seq_lens = seq_lens.to(torch.int64)
@@ -303,6 +321,16 @@ def forward_decode(
         )
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
+    def _get_non_causal_qlen_mask(self, bs: int, draft_token_num: int) -> torch.Tensor:
+        key = (int(bs), int(draft_token_num))
+        mask = self._non_causal_masks.get(key)
+        if mask is None:
+            mask = torch.ones(
+                key[0] * key[1] * key[1], dtype=torch.bool, device=self.device
+            )
+            self._non_causal_masks[key] = mask
+        return mask
+
     def _get_attn_logits_buffer(
         self, num_seqs: int, num_heads: int, v_head_dim: int
     ) -> torch.Tensor:

python/sglang/srt/models/dflash.py

@@ -39,17 +39,19 @@
 from sglang.srt.runtime_context import get_parallel, get_spec
 from sglang.srt.speculative.dflash_utils import (
     can_dflash_slice_qkv_weight,
+    dflash_head_logits,
     get_dflash_attention_sliding_window_size,
     get_dflash_layer_types,
     is_dense_head_weight,
     is_nemotron_35_draft_config,
     parse_dflash_draft_config,
 )
-from sglang.srt.utils import is_npu, set_weight_attrs
+from sglang.srt.utils import is_cpu, is_npu, set_weight_attrs
 from sglang.srt.utils.common import get_compiler_backend
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 _is_npu = is_npu()
+_is_cpu = is_cpu()
 if _is_npu:
     from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope
 logger = logging.getLogger(__name__)
@@ -90,8 +92,14 @@ def _project_candidate_logits(
 ) -> torch.Tensor:
     """Project draft hiddens through the target head, restricted to the org vocab."""
     if not use_quant_head:
-        weight = lm_head.weight
-        return torch.matmul(hidden.to(weight.dtype), weight[:num_org].T)
+        logits = dflash_head_logits(lm_head, hidden, 0, num_org)
+        if logits.shape[-1] <= num_org:
+            return logits
+        # An AMX-prepacked head can only produce the full local shard, so the
+        # padded tail is masked out of the top-k like the quantized path below.
+        logits = logits.contiguous()
+        logits[:, num_org:] = float("-inf")
+        return logits
     # A packed weight can't be row-sliced to the org vocab like the dense path,
     # and flashinfer's radix top-k rejects the crop view (non-contiguous), so
     # mask the padded tail out of the top-k instead.
@@ -219,6 +227,8 @@ def __init__(
         rotary = self.rotary_emb
         self.use_table_qk_norm_rope = (
             not _is_npu
+            # table_qk_norm_rope_ is a Triton kernel with no CPU equivalent.
+            and not _is_cpu
             and hasattr(rotary, "cos_sin_cache")
             and getattr(rotary, "rotary_dim", None) == head_dim
             and getattr(rotary, "is_neox_style", False)

python/sglang/srt/server_args.py

@@ -254,6 +254,7 @@
     "triton",
     "ascend",
     "trtllm_mha",
+    "intel_amx",
 ]
 add_draft_attention_backend_choices = DRAFT_ATTENTION_BACKEND_CHOICES.extend
 

python/sglang/srt/speculative/dflash_info.py

@@ -56,6 +56,12 @@ def __post_init__(self):
             self.num_tokens_per_req = int(self.draft_token_num)
         self.num_tokens_for_logprob_per_req = int(self.draft_token_num)
 
+    @property
+    def tree_topk(self) -> int:
+        # DFLASH proposals are a linear chain, so backends that branch on tree
+        # width (e.g. intel_amx) can take their mask-free causal path.
+        return int(self.topk)
+
     def prepare_for_verify(
         self,
         batch: ScheduleBatch,

python/sglang/srt/speculative/dflash_utils.py

@@ -19,6 +19,7 @@
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.speculative.spec_utils import sample_simulated_acc_len
 from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu
+from sglang.srt.utils.common import use_intel_amx_backend
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
 
@@ -34,6 +35,9 @@
         "TritonAttnBackend",
         "TRTLLMHAAttnBackend",
         "TRTLLMMLABackend",
+        # The CPU kernel applies causal masking to the verify block on its own
+        # when the proposal is a linear chain (tree_topk == 1).
+        "IntelAMXAttnBackend",
     }
 )
 
@@ -735,9 +739,44 @@ def can_dflash_slice_qkv_weight(qkv_proj: Any) -> Tuple[bool, str]:
         )
     if not hasattr(qkv_proj, "weight"):
         return False, "qkv weight tensor is missing"
+    if use_intel_amx_backend(qkv_proj):
+        # The AMX path repacks the weight into a blocked VNNI layout, so its rows
+        # no longer map to q/k/v output ranges and a plain F.linear against it
+        # would read the packed bytes as if they were a dense matrix.
+        return False, "AMX-prepacked qkv_proj weight cannot be sliced"
     return True, ""
 
 
+def dflash_head_logits(
+    lm_head: Any,
+    hidden_states: torch.Tensor,
+    start: int = 0,
+    end: Optional[int] = None,
+) -> torch.Tensor:
+    """Logits for rows ``[start:end)`` of a dense target lm_head weight.
+
+    On CPU with AMX the head weight is prepacked, so the rows cannot be sliced
+    before the matmul; the packed kernel produces the full local shard and the
+    requested range is taken from the output instead.
+    """
+    weight = lm_head.weight
+    if hidden_states.dtype != weight.dtype:
+        hidden_states = hidden_states.to(weight.dtype)
+    if use_intel_amx_backend(lm_head):
+        logits = torch.ops.sgl_kernel.weight_packed_linear(
+            hidden_states.contiguous(),
+            weight,
+            None,  # bias
+            True,  # is_vnni
+        )
+        if start != 0 or end is not None:
+            logits = logits[:, start : logits.shape[-1] if end is None else end]
+        return logits
+    if start != 0 or end is not None:
+        weight = weight[start:end]
+    return torch.matmul(hidden_states, weight.T)
+
+
 def can_dflash_use_fused_qkv_proj(qkv_proj: Any) -> Tuple[bool, str]:
     """Validate whether a QKV layer is eligible for DFlash fused KV materialization."""
     eligible, reason = can_dflash_slice_qkv_weight(qkv_proj)

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -49,6 +49,7 @@
     can_dflash_use_fused_qkv_proj,
     compute_dflash_correct_drafts_and_bonus,
     compute_dflash_sampling_correct_drafts_and_bonus,
+    dflash_head_logits,
     is_dense_head_weight,
     is_dflash_sampling_verify_available,
     parse_dflash_draft_config,
@@ -70,9 +71,10 @@
     assign_req_to_token_pool_func,
     build_grammar_vocab_mask,
 )
-from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu
+from sglang.srt.utils import get_available_gpu_memory, is_cpu, is_cuda, is_hip, is_npu
 
 _is_npu = is_npu()
+_is_cpu = is_cpu()
 
 
 logger = logging.getLogger(__name__)
@@ -462,7 +464,8 @@ def init_attention_backends(self):
 
     def init_cuda_graphs(self):
         capture_decode_cuda_graph = (
-            get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
+            not _is_cpu
+            and get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
         )
         if is_cuda() and capture_decode_cuda_graph:
             available_mem = get_available_gpu_memory(self.device, self.gpu_id)
@@ -1109,7 +1112,7 @@ def _cast_hs(x: torch.Tensor) -> torch.Tensor:
             for start in range(0, num_tokens, int(chunk_size)):
                 end = min(num_tokens, start + int(chunk_size))
                 hs = _cast_hs(hidden_states[start:end])
-                logits = torch.matmul(hs, weight.T)
+                logits = dflash_head_logits(lm_head, hs)
                 out_tokens[start:end] = torch.argmax(logits, dim=-1).to(torch.long)
             return out_tokens
 
@@ -1164,7 +1167,7 @@ def _ensure_local_reduce_buffers(
                 end = min(num_tokens, start + fast_chunk_size)
                 hs = _cast_hs(hidden_states[start:end])
                 if num_org > 0:
-                    base_logits = torch.matmul(hs, weight[:num_org].T)
+                    base_logits = dflash_head_logits(lm_head, hs, 0, num_org)
                     local_max, local_arg = _ensure_local_reduce_buffers(
                         end - start, base_logits.dtype, hs.device
                     )
@@ -1182,7 +1185,7 @@ def _ensure_local_reduce_buffers(
 
             # Base vocab logits.
             if num_org > 0:
-                base_logits = torch.matmul(hs, weight[:num_org].T)
+                base_logits = dflash_head_logits(lm_head, hs, 0, num_org)
                 local_max, local_arg = _ensure_local_reduce_buffers(
                     chunk_len, base_logits.dtype, hs.device
                 )
@@ -1202,8 +1205,8 @@ def _ensure_local_reduce_buffers(
             if num_added > 0:
                 added_slice_start = num_org_padded
                 added_slice_end = num_org_padded + num_added
-                added_logits = torch.matmul(
-                    hs, weight[added_slice_start:added_slice_end].T
+                added_logits = dflash_head_logits(
+                    lm_head, hs, added_slice_start, added_slice_end
                 )
                 added_max, added_arg = torch.max(added_logits, dim=-1)
                 use_added = added_max > local_max
@@ -1754,10 +1757,11 @@ def forward_batch_generation(
             )
 
         # `seq_lens` is carried over from the previous overlap iteration and may have been
-        # produced on another stream.
-        batch.seq_lens.record_stream(
-            torch.get_device_module(self.device).current_stream()
-        )
+        # produced on another stream. CPU tensors have no stream to record against.
+        if not _is_cpu:
+            batch.seq_lens.record_stream(
+                torch.get_device_module(self.device).current_stream()
+            )
 
         bs = len(batch.seq_lens)
         device = self.device

python/sglang/srt/speculative/draft_worker_common.py

@@ -9,7 +9,7 @@
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
-from sglang.srt.runtime_context import attention_backends, get_spec
+from sglang.srt.runtime_context import attention_backends, get_device, get_spec
 from sglang.srt.server_args import DRAFT_ATTENTION_BACKEND_CHOICES, ServerArgs
 from sglang.srt.speculative.dflash_info import DFlashVerifyInput
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -29,6 +29,13 @@ class DraftWorkerBundle(msgspec.Struct, frozen=True):
     resolved_attention_backend: str
 
 
+def _default_draft_attention_backend() -> str:
+    # A GPU default on a CPU run would fail at the first draft forward.
+    if get_device().device == "cpu":
+        return "intel_amx"
+    return "triton" if torch.version.hip else "flashinfer"
+
+
 def _resolve_draft_attention_backend_fallback(*, algo_label: str) -> str:
     """The draft's attention backend, from the published leaves.
 
@@ -40,9 +47,9 @@ def _resolve_draft_attention_backend_fallback(*, algo_label: str) -> str:
     if draft_backend is None:
         draft_backend, _ = attention_backends()
     if draft_backend is None:
-        return "triton" if torch.version.hip else "flashinfer"
+        return _default_draft_attention_backend()
     if draft_backend not in DRAFT_ATTENTION_BACKEND_CHOICES:
-        fallback = "triton" if torch.version.hip else "flashinfer"
+        fallback = _default_draft_attention_backend()
         logger.warning(
             "%s draft worker only supports attention_backend in %s for now, "
             "but got %r. Falling back to '%s'.",