[CPU][Spec Decode] Enable DFlash SD for CPU (#44029)

Signed-off-by: guybd <[email protected]>
Signed-off-by: Guy Boudoukh <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

csrc/cpu/spec_decode_utils.cpp

@@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl(
   }
 }
 
+void copy_and_expand_dflash_inputs_kernel_impl(
+    const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
+    torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
+    torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
+    torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
+    const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
+    const std::optional<torch::Tensor>& num_rejected_tokens,
+    const int64_t parallel_drafting_token_id, const int64_t block_size,
+    const int64_t num_query_per_req, const int64_t num_speculative_tokens,
+    const int64_t total_input_tokens, const bool has_num_rejected) {
+  const int64_t num_reqs = query_start_loc.size(0) - 1;
+
+  const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
+  const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
+  const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
+  const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
+  const int64_t* rejected_ptr =
+      has_num_rejected && num_rejected_tokens.has_value()
+          ? num_rejected_tokens.value().data_ptr<int64_t>()
+          : nullptr;
+
+  int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
+  int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>();
+  int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>();
+  int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>();
+  int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>();
+  int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>();
+
+  const int64_t block_table_stride = block_table.stride(0);
+
+#pragma omp parallel for
+  for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
+    int32_t ctx_start = query_start_ptr[req_idx];
+    int32_t ctx_end = query_start_ptr[req_idx + 1];
+    int64_t num_ctx = ctx_end - ctx_start;
+    int64_t valid_ctx_end = ctx_end;
+    if (rejected_ptr != nullptr) {
+      valid_ctx_end -= rejected_ptr[req_idx];
+    }
+    // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that
+    // valid_ctx_end - 1 never reads before the request's context range.
+    valid_ctx_end =
+        std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1));
+
+    int64_t last_pos = target_pos_ptr[valid_ctx_end - 1];
+
+    for (int64_t j = 0; j < num_ctx; ++j) {
+      int64_t ctx_idx = ctx_start + j;
+      int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1);
+      int64_t position = target_pos_ptr[ctx_pos_idx];
+      int64_t block_num = position / block_size;
+      block_num = std::min(block_num, block_table_stride - 1);
+      int32_t block_id =
+          block_table_ptr[req_idx * block_table_stride + block_num];
+      int64_t slot = block_id * block_size + (position % block_size);
+
+      out_ctx_pos_ptr[ctx_idx] = position;
+      out_ctx_slot_ptr[ctx_idx] = slot;
+    }
+
+    for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) {
+      int64_t query_out = req_idx * num_query_per_req + query_off;
+      int64_t position = last_pos + 1 + query_off;
+      int64_t block_num = position / block_size;
+      block_num = std::min(block_num, block_table_stride - 1);
+      int32_t block_id =
+          block_table_ptr[req_idx * block_table_stride + block_num];
+      int64_t slot = block_id * block_size + (position % block_size);
+
+      out_query_pos_ptr[query_out] = position;
+      out_query_slot_ptr[query_out] = slot;
+      out_ids_ptr[query_out] =
+          query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id;
+
+      if (query_off > 0) {
+        int64_t sample_out_idx =
+            req_idx * num_speculative_tokens + (query_off - 1);
+        out_token_idx_ptr[sample_out_idx] = query_out;
+      }
+    }
+  }
+}
+
 void rejection_greedy_sample_kernel_impl(
     torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,

csrc/cpu/torch_bindings.cpp

@@ -237,6 +237,16 @@ void copy_and_expand_eagle_inputs_kernel_impl(
     const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
     const int64_t total_input_tokens,
     const int64_t num_padding_slots_per_request, const bool shift_input_ids);
+void copy_and_expand_dflash_inputs_kernel_impl(
+    const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
+    torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
+    torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
+    torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
+    const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
+    const std::optional<torch::Tensor>& num_rejected_tokens,
+    const int64_t parallel_drafting_token_id, const int64_t block_size,
+    const int64_t num_query_per_req, const int64_t num_speculative_tokens,
+    const int64_t total_input_tokens, const bool has_num_rejected);
 void rejection_greedy_sample_kernel_impl(
     torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
     const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
@@ -599,6 +609,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
       "bool shift_input_ids) -> ()",
       &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
+  ops.def(
+      "copy_and_expand_dflash_inputs_kernel_impl("
+      "Tensor next_token_ids, Tensor target_positions, "
+      "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, "
+      "Tensor(a4!) out_query_positions, "
+      "Tensor(a5!) out_context_slot_mapping, "
+      "Tensor(a6!) out_query_slot_mapping, "
+      "Tensor(a7!) out_token_indices, Tensor block_table, "
+      "Tensor query_start_loc, Tensor? num_rejected_tokens, "
+      "SymInt parallel_drafting_token_id, SymInt block_size, "
+      "SymInt num_query_per_req, SymInt num_speculative_tokens, "
+      "SymInt total_input_tokens, bool has_num_rejected) -> ()",
+      &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl);
   ops.def(
       "rejection_greedy_sample_kernel_impl("
       "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "

docs/design/attention_backends.md

@@ -159,7 +159,7 @@ Priority is **1 = highest** (tried first).
 
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
-| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A |
 | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |

vllm/model_executor/models/qwen3_dflash.py

@@ -131,7 +131,7 @@ class DFlashQwen3Attention(nn.Module):
         with the context K/V from the target model's hidden states. This forward op
         computes attention for the query tokens only.
         See also: precompute_and_store_context_kv"""
-        qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias)
+        qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
 
         # Per-head RMSNorm

vllm/utils/cpu_triton_utils.py

@@ -197,6 +197,133 @@ def _copy_and_expand_eagle_inputs_kernel_impl(
         out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype))
 
 
+def _copy_and_expand_dflash_inputs_kernel_impl(
+    next_token_ids_ptr,
+    target_positions_ptr,
+    out_input_ids_ptr,
+    out_context_positions_ptr,
+    out_query_positions_ptr,
+    out_context_slot_mapping_ptr,
+    out_query_slot_mapping_ptr,
+    out_token_indices_ptr,
+    block_table_ptr,
+    block_table_stride,
+    query_start_loc_ptr,
+    num_rejected_tokens_ptr,
+    parallel_drafting_token_id,
+    block_size,
+    num_query_per_req,
+    num_speculative_tokens,
+    total_input_tokens,
+    BLOCK_SIZE=None,
+    HAS_NUM_REJECTED=False,
+):
+    """Adapter between the DFlash Triton launch and the C++ CPU op."""
+    assert block_table_stride == block_table_ptr.stride(0), (
+        "block_table_stride mismatch: "
+        f"{block_table_stride} vs {block_table_ptr.stride(0)}"
+    )
+
+    orig_ids_dtype = out_input_ids_ptr.dtype
+    orig_context_positions_dtype = out_context_positions_ptr.dtype
+    orig_query_positions_dtype = out_query_positions_ptr.dtype
+    orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype
+    orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype
+    out_ids_i64 = _ensure_int64(out_input_ids_ptr)
+    out_context_positions_i64 = _ensure_int64(out_context_positions_ptr)
+    out_query_positions_i64 = _ensure_int64(out_query_positions_ptr)
+    out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr)
+    out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr)
+    rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None
+
+    if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"):
+        torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl(
+            _ensure_int64(next_token_ids_ptr),
+            _ensure_int64(target_positions_ptr),
+            out_ids_i64,
+            out_context_positions_i64,
+            out_query_positions_i64,
+            out_context_slot_mapping_i64,
+            out_query_slot_mapping_i64,
+            out_token_indices_ptr,
+            block_table_ptr,
+            query_start_loc_ptr,
+            rejected_i64,
+            parallel_drafting_token_id,
+            block_size,
+            num_query_per_req,
+            num_speculative_tokens,
+            total_input_tokens,
+            HAS_NUM_REJECTED,
+        )
+    else:
+        next_ids_i64 = _ensure_int64(next_token_ids_ptr)
+        target_positions_i64 = _ensure_int64(target_positions_ptr)
+        block_table_stride = block_table_ptr.stride(0)
+        num_reqs = query_start_loc_ptr.shape[0] - 1
+
+        for req_idx in range(num_reqs):
+            ctx_start = int(query_start_loc_ptr[req_idx].item())
+            ctx_end = int(query_start_loc_ptr[req_idx + 1].item())
+            num_ctx = ctx_end - ctx_start
+            valid_ctx_end = ctx_end
+            if rejected_i64 is not None:
+                valid_ctx_end -= int(rejected_i64[req_idx].item())
+            # Guard against out-of-bounds: ensure valid_ctx_end > ctx_start.
+            valid_ctx_end = max(valid_ctx_end, ctx_start + 1)
+
+            last_pos = int(target_positions_i64[valid_ctx_end - 1].item())
+
+            for j in range(num_ctx):
+                ctx_idx = ctx_start + j
+                ctx_pos_idx = min(ctx_idx, total_input_tokens - 1)
+                position = int(target_positions_i64[ctx_pos_idx].item())
+                block_num = min(position // block_size, block_table_stride - 1)
+                block_id = int(block_table_ptr[req_idx, block_num].item())
+                slot = block_id * block_size + (position % block_size)
+
+                out_context_positions_i64[ctx_idx] = position
+                out_context_slot_mapping_i64[ctx_idx] = slot
+
+            for query_off in range(num_query_per_req):
+                query_out = req_idx * num_query_per_req + query_off
+                position = last_pos + 1 + query_off
+                block_num = min(position // block_size, block_table_stride - 1)
+                block_id = int(block_table_ptr[req_idx, block_num].item())
+                slot = block_id * block_size + (position % block_size)
+
+                out_query_positions_i64[query_out] = position
+                out_query_slot_mapping_i64[query_out] = slot
+                out_ids_i64[query_out] = (
+                    int(next_ids_i64[req_idx].item())
+                    if query_off == 0
+                    else parallel_drafting_token_id
+                )
+
+                if query_off > 0:
+                    sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1)
+                    out_token_indices_ptr[sample_out_idx] = query_out
+
+    if orig_ids_dtype != torch.int64:
+        out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
+    if orig_context_positions_dtype != torch.int64:
+        out_context_positions_ptr.copy_(
+            out_context_positions_i64.to(orig_context_positions_dtype)
+        )
+    if orig_query_positions_dtype != torch.int64:
+        out_query_positions_ptr.copy_(
+            out_query_positions_i64.to(orig_query_positions_dtype)
+        )
+    if orig_context_slot_mapping_dtype != torch.int64:
+        out_context_slot_mapping_ptr.copy_(
+            out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype)
+        )
+    if orig_query_slot_mapping_dtype != torch.int64:
+        out_query_slot_mapping_ptr.copy_(
+            out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype)
+        )
+
+
 def _rejection_greedy_sample_kernel_impl(
     output_token_ids,
     cu_num_draft_tokens,
@@ -303,6 +430,10 @@ def _sample_recovered_tokens_kernel_impl(
     NO_DRAFT_PROBS=False,
     USE_FP64_GUMBEL=False,
 ):
+    # USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller
+    # has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes
+    # `inv_q` directly, so the flag is accepted for interface parity and the
+    # value is read at its existing dtype.
     # C++ reads integer tensors as int64_t*; ensure correct dtype.
     orig_dtype = output_token_ids.dtype
     output_i64 = _ensure_int64(output_token_ids)
@@ -330,6 +461,9 @@ eagle_prepare_next_token_padded_kernel = _FuncWrapper(
 copy_and_expand_eagle_inputs_kernel = _FuncWrapper(
     _copy_and_expand_eagle_inputs_kernel_impl
 )
+copy_and_expand_dflash_inputs_kernel = _FuncWrapper(
+    _copy_and_expand_dflash_inputs_kernel_impl
+)
 eagle_step_slot_mapping_metadata_kernel = _FuncWrapper(
     _eagle_step_slot_mapping_metadata_kernel_impl
 )

vllm/v1/attention/backends/cpu_attn.py

@@ -63,6 +63,10 @@ class CPUAttentionBackend(AttentionBackend):
     def get_name() -> str:
         return "CPU_ATTN"
 
+    @classmethod
+    def supports_non_causal(cls) -> bool:
+        return True
+
     @classmethod
     def supports_attn_type(cls, attn_type: str) -> bool:
         """CPU attention supports decoder,

vllm/v1/spec_decode/dflash.py

@@ -10,10 +10,12 @@ from typing_extensions import override
 from vllm.config import VllmConfig
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
-from vllm.triton_utils import triton
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
-from vllm.v1.spec_decode.utils import copy_and_expand_dflash_inputs_kernel
+from vllm.v1.spec_decode.utils import (
+    copy_and_expand_dflash_inputs_kernel,
+    next_power_of_2,
+)
 
 logger = init_logger(__name__)
 
@@ -126,8 +128,8 @@ class DFlashProposer(SpecDecodeBaseProposer):
         # and token_indices_to_sample
         max_ctx_per_req = cad.max_query_len
         max_tokens_per_req = max_ctx_per_req + num_query_per_req
-        BLOCK_SIZE = min(256, triton.next_power_of_2(max_tokens_per_req))
-        num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
+        BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))
+        num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE
         grid = (batch_size, num_blocks)
 
         has_num_rejected = num_rejected_tokens_gpu is not None

vllm/v1/worker/cpu_model_runner.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import sys
 from contextlib import contextmanager
 from typing import Any
 
@@ -78,7 +79,7 @@ class CPUModelRunner(GPUModelRunner):
         # Speculative decoding fallbacks
         import vllm.v1.sample.rejection_sampler
         import vllm.v1.spec_decode.llm_base_proposer
-        import vllm.v1.spec_decode.utils
+        import vllm.v1.spec_decode.utils as spec_decode_utils
 
         vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (
             cpu_tl.eagle_prepare_inputs_padded_kernel
@@ -89,7 +90,18 @@ class CPUModelRunner(GPUModelRunner):
         vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (
             cpu_tl.copy_and_expand_eagle_inputs_kernel
         )
-        vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel = (
+        spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (
+            cpu_tl.copy_and_expand_dflash_inputs_kernel
+        )
+        dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash")
+        if dflash_module is not None:
+            dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel"
+            setattr(
+                dflash_module,
+                dflash_kernel_name,
+                cpu_tl.copy_and_expand_dflash_inputs_kernel,
+            )
+        spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = (
             cpu_tl.eagle_step_slot_mapping_metadata_kernel
         )
         vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = (