[CPU][Spec Decode] Enable DFlash SD for CPU (#44029)
Signed-off-by: guybd <guy.boudoukh@intel.com> Signed-off-by: Guy Boudoukh <guy.boudoukh@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
guybd committed
Jun 25, 2026 at 10:32 UTC
2396d91e931295df877cdad5e2ac0de4e35dab9f
8 files changed
+266
-8
csrc/cpu/spec_decode_utils.cpp
+83
@@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl(
208
}
209
}
210
211
+void copy_and_expand_dflash_inputs_kernel_impl(
212
+ const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
213
+ torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
214
+ torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
215
+ torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
216
+ const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
217
+ const std::optional<torch::Tensor>& num_rejected_tokens,
218
+ const int64_t parallel_drafting_token_id, const int64_t block_size,
219
+ const int64_t num_query_per_req, const int64_t num_speculative_tokens,
220
+ const int64_t total_input_tokens, const bool has_num_rejected) {
221
+ const int64_t num_reqs = query_start_loc.size(0) - 1;
222
+
223
+ const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
224
+ const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
225
+ const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
226
+ const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
227
+ const int64_t* rejected_ptr =
228
+ has_num_rejected && num_rejected_tokens.has_value()
229
+ ? num_rejected_tokens.value().data_ptr<int64_t>()
230
+ : nullptr;
231
+
232
+ int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
233
+ int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>();
234
+ int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>();
235
+ int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>();
236
+ int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>();
237
+ int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>();
238
+
239
+ const int64_t block_table_stride = block_table.stride(0);
240
+
241
+#pragma omp parallel for
242
+ for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
243
+ int32_t ctx_start = query_start_ptr[req_idx];
244
+ int32_t ctx_end = query_start_ptr[req_idx + 1];
245
+ int64_t num_ctx = ctx_end - ctx_start;
246
+ int64_t valid_ctx_end = ctx_end;
247
+ if (rejected_ptr != nullptr) {
248
+ valid_ctx_end -= rejected_ptr[req_idx];
249
+ }
250
+ // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that
251
+ // valid_ctx_end - 1 never reads before the request's context range.
252
+ valid_ctx_end =
253
+ std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1));
254
+
255
+ int64_t last_pos = target_pos_ptr[valid_ctx_end - 1];
256
+
257
+ for (int64_t j = 0; j < num_ctx; ++j) {
258
+ int64_t ctx_idx = ctx_start + j;
259
+ int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1);
260
+ int64_t position = target_pos_ptr[ctx_pos_idx];
261
+ int64_t block_num = position / block_size;
262
+ block_num = std::min(block_num, block_table_stride - 1);
263
+ int32_t block_id =
264
+ block_table_ptr[req_idx * block_table_stride + block_num];
265
+ int64_t slot = block_id * block_size + (position % block_size);
266
+
267
+ out_ctx_pos_ptr[ctx_idx] = position;
268
+ out_ctx_slot_ptr[ctx_idx] = slot;
269
+ }
270
+
271
+ for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) {
272
+ int64_t query_out = req_idx * num_query_per_req + query_off;
273
+ int64_t position = last_pos + 1 + query_off;
274
+ int64_t block_num = position / block_size;
275
+ block_num = std::min(block_num, block_table_stride - 1);
276
+ int32_t block_id =
277
+ block_table_ptr[req_idx * block_table_stride + block_num];
278
+ int64_t slot = block_id * block_size + (position % block_size);
279
+
280
+ out_query_pos_ptr[query_out] = position;
281
+ out_query_slot_ptr[query_out] = slot;
282
+ out_ids_ptr[query_out] =
283
+ query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id;
284
+
285
+ if (query_off > 0) {
286
+ int64_t sample_out_idx =
287
+ req_idx * num_speculative_tokens + (query_off - 1);
288
+ out_token_idx_ptr[sample_out_idx] = query_out;
289
+ }
290
+ }
291
+ }
292
+}
293
+
294
void rejection_greedy_sample_kernel_impl(
295
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
296
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
csrc/cpu/torch_bindings.cpp
+23
@@ -237,6 +237,16 @@ void copy_and_expand_eagle_inputs_kernel_impl(
237
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
238
const int64_t total_input_tokens,
239
const int64_t num_padding_slots_per_request, const bool shift_input_ids);
240
+void copy_and_expand_dflash_inputs_kernel_impl(
241
+ const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
242
+ torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
243
+ torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
244
+ torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
245
+ const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
246
+ const std::optional<torch::Tensor>& num_rejected_tokens,
247
+ const int64_t parallel_drafting_token_id, const int64_t block_size,
248
+ const int64_t num_query_per_req, const int64_t num_speculative_tokens,
249
+ const int64_t total_input_tokens, const bool has_num_rejected);
250
void rejection_greedy_sample_kernel_impl(
251
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
252
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
@@ -599,6 +609,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
609
"SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
610
"bool shift_input_ids) -> ()",
611
&cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
612
+ ops.def(
613
+ "copy_and_expand_dflash_inputs_kernel_impl("
614
+ "Tensor next_token_ids, Tensor target_positions, "
615
+ "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, "
616
+ "Tensor(a4!) out_query_positions, "
617
+ "Tensor(a5!) out_context_slot_mapping, "
618
+ "Tensor(a6!) out_query_slot_mapping, "
619
+ "Tensor(a7!) out_token_indices, Tensor block_table, "
620
+ "Tensor query_start_loc, Tensor? num_rejected_tokens, "
621
+ "SymInt parallel_drafting_token_id, SymInt block_size, "
622
+ "SymInt num_query_per_req, SymInt num_speculative_tokens, "
623
+ "SymInt total_input_tokens, bool has_num_rejected) -> ()",
624
+ &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl);
625
ops.def(
626
"rejection_greedy_sample_kernel_impl("
627
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
docs/design/attention_backends.md
+1
-1
@@ -159,7 +159,7 @@ Priority is **1 = highest** (tried first).
159
160
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
161
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
162
-| `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 |
162
+| `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 |
163
| `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 |
164
| `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 |
165
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
vllm/model_executor/models/qwen3_dflash.py
+1
-1
@@ -131,7 +131,7 @@ class DFlashQwen3Attention(nn.Module):
131
with the context K/V from the target model's hidden states. This forward op
132
computes attention for the query tokens only.
133
See also: precompute_and_store_context_kv"""
134
- qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias)
134
+ qkv, _ = self.qkv_proj(hidden_states)
135
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
136
137
# Per-head RMSNorm
vllm/utils/cpu_triton_utils.py
+134
@@ -197,6 +197,133 @@ def _copy_and_expand_eagle_inputs_kernel_impl(
197
out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype))
198
199
200
+def _copy_and_expand_dflash_inputs_kernel_impl(
201
+ next_token_ids_ptr,
202
+ target_positions_ptr,
203
+ out_input_ids_ptr,
204
+ out_context_positions_ptr,
205
+ out_query_positions_ptr,
206
+ out_context_slot_mapping_ptr,
207
+ out_query_slot_mapping_ptr,
208
+ out_token_indices_ptr,
209
+ block_table_ptr,
210
+ block_table_stride,
211
+ query_start_loc_ptr,
212
+ num_rejected_tokens_ptr,
213
+ parallel_drafting_token_id,
214
+ block_size,
215
+ num_query_per_req,
216
+ num_speculative_tokens,
217
+ total_input_tokens,
218
+ BLOCK_SIZE=None,
219
+ HAS_NUM_REJECTED=False,
220
+):
221
+ """Adapter between the DFlash Triton launch and the C++ CPU op."""
222
+ assert block_table_stride == block_table_ptr.stride(0), (
223
+ "block_table_stride mismatch: "
224
+ f"{block_table_stride} vs {block_table_ptr.stride(0)}"
225
+ )
226
+
227
+ orig_ids_dtype = out_input_ids_ptr.dtype
228
+ orig_context_positions_dtype = out_context_positions_ptr.dtype
229
+ orig_query_positions_dtype = out_query_positions_ptr.dtype
230
+ orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype
231
+ orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype
232
+ out_ids_i64 = _ensure_int64(out_input_ids_ptr)
233
+ out_context_positions_i64 = _ensure_int64(out_context_positions_ptr)
234
+ out_query_positions_i64 = _ensure_int64(out_query_positions_ptr)
235
+ out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr)
236
+ out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr)
237
+ rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None
238
+
239
+ if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"):
240
+ torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl(
241
+ _ensure_int64(next_token_ids_ptr),
242
+ _ensure_int64(target_positions_ptr),
243
+ out_ids_i64,
244
+ out_context_positions_i64,
245
+ out_query_positions_i64,
246
+ out_context_slot_mapping_i64,
247
+ out_query_slot_mapping_i64,
248
+ out_token_indices_ptr,
249
+ block_table_ptr,
250
+ query_start_loc_ptr,
251
+ rejected_i64,
252
+ parallel_drafting_token_id,
253
+ block_size,
254
+ num_query_per_req,
255
+ num_speculative_tokens,
256
+ total_input_tokens,
257
+ HAS_NUM_REJECTED,
258
+ )
259
+ else:
260
+ next_ids_i64 = _ensure_int64(next_token_ids_ptr)
261
+ target_positions_i64 = _ensure_int64(target_positions_ptr)
262
+ block_table_stride = block_table_ptr.stride(0)
263
+ num_reqs = query_start_loc_ptr.shape[0] - 1
264
+
265
+ for req_idx in range(num_reqs):
266
+ ctx_start = int(query_start_loc_ptr[req_idx].item())
267
+ ctx_end = int(query_start_loc_ptr[req_idx + 1].item())
268
+ num_ctx = ctx_end - ctx_start
269
+ valid_ctx_end = ctx_end
270
+ if rejected_i64 is not None:
271
+ valid_ctx_end -= int(rejected_i64[req_idx].item())
272
+ # Guard against out-of-bounds: ensure valid_ctx_end > ctx_start.
273
+ valid_ctx_end = max(valid_ctx_end, ctx_start + 1)
274
+
275
+ last_pos = int(target_positions_i64[valid_ctx_end - 1].item())
276
+
277
+ for j in range(num_ctx):
278
+ ctx_idx = ctx_start + j
279
+ ctx_pos_idx = min(ctx_idx, total_input_tokens - 1)
280
+ position = int(target_positions_i64[ctx_pos_idx].item())
281
+ block_num = min(position // block_size, block_table_stride - 1)
282
+ block_id = int(block_table_ptr[req_idx, block_num].item())
283
+ slot = block_id * block_size + (position % block_size)
284
+
285
+ out_context_positions_i64[ctx_idx] = position
286
+ out_context_slot_mapping_i64[ctx_idx] = slot
287
+
288
+ for query_off in range(num_query_per_req):
289
+ query_out = req_idx * num_query_per_req + query_off
290
+ position = last_pos + 1 + query_off
291
+ block_num = min(position // block_size, block_table_stride - 1)
292
+ block_id = int(block_table_ptr[req_idx, block_num].item())
293
+ slot = block_id * block_size + (position % block_size)
294
+
295
+ out_query_positions_i64[query_out] = position
296
+ out_query_slot_mapping_i64[query_out] = slot
297
+ out_ids_i64[query_out] = (
298
+ int(next_ids_i64[req_idx].item())
299
+ if query_off == 0
300
+ else parallel_drafting_token_id
301
+ )
302
+
303
+ if query_off > 0:
304
+ sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1)
305
+ out_token_indices_ptr[sample_out_idx] = query_out
306
+
307
+ if orig_ids_dtype != torch.int64:
308
+ out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
309
+ if orig_context_positions_dtype != torch.int64:
310
+ out_context_positions_ptr.copy_(
311
+ out_context_positions_i64.to(orig_context_positions_dtype)
312
+ )
313
+ if orig_query_positions_dtype != torch.int64:
314
+ out_query_positions_ptr.copy_(
315
+ out_query_positions_i64.to(orig_query_positions_dtype)
316
+ )
317
+ if orig_context_slot_mapping_dtype != torch.int64:
318
+ out_context_slot_mapping_ptr.copy_(
319
+ out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype)
320
+ )
321
+ if orig_query_slot_mapping_dtype != torch.int64:
322
+ out_query_slot_mapping_ptr.copy_(
323
+ out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype)
324
+ )
325
+
326
+
327
def _rejection_greedy_sample_kernel_impl(
328
output_token_ids,
329
cu_num_draft_tokens,
@@ -303,6 +430,10 @@ def _sample_recovered_tokens_kernel_impl(
430
NO_DRAFT_PROBS=False,
431
USE_FP64_GUMBEL=False,
432
):
433
+ # USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller
434
+ # has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes
435
+ # `inv_q` directly, so the flag is accepted for interface parity and the
436
+ # value is read at its existing dtype.
437
# C++ reads integer tensors as int64_t*; ensure correct dtype.
438
orig_dtype = output_token_ids.dtype
439
output_i64 = _ensure_int64(output_token_ids)
@@ -330,6 +461,9 @@ eagle_prepare_next_token_padded_kernel = _FuncWrapper(
461
copy_and_expand_eagle_inputs_kernel = _FuncWrapper(
462
_copy_and_expand_eagle_inputs_kernel_impl
463
)
464
+copy_and_expand_dflash_inputs_kernel = _FuncWrapper(
465
+ _copy_and_expand_dflash_inputs_kernel_impl
466
+)
467
eagle_step_slot_mapping_metadata_kernel = _FuncWrapper(
468
_eagle_step_slot_mapping_metadata_kernel_impl
469
)
vllm/v1/attention/backends/cpu_attn.py
+4
@@ -63,6 +63,10 @@ class CPUAttentionBackend(AttentionBackend):
63
def get_name() -> str:
64
return "CPU_ATTN"
65
66
+ @classmethod
67
+ def supports_non_causal(cls) -> bool:
68
+ return True
69
+
70
@classmethod
71
def supports_attn_type(cls, attn_type: str) -> bool:
72
"""CPU attention supports decoder,
vllm/v1/spec_decode/dflash.py
+6
-4
@@ -10,10 +10,12 @@ from typing_extensions import override
10
from vllm.config import VllmConfig
11
from vllm.forward_context import set_forward_context
12
from vllm.logger import init_logger
13
-from vllm.triton_utils import triton
13
from vllm.v1.attention.backend import CommonAttentionMetadata
14
from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
16
-from vllm.v1.spec_decode.utils import copy_and_expand_dflash_inputs_kernel
15
+from vllm.v1.spec_decode.utils import (
16
+ copy_and_expand_dflash_inputs_kernel,
17
+ next_power_of_2,
18
+)
19
20
logger = init_logger(__name__)
21
@@ -126,8 +128,8 @@ class DFlashProposer(SpecDecodeBaseProposer):
128
# and token_indices_to_sample
129
max_ctx_per_req = cad.max_query_len
130
max_tokens_per_req = max_ctx_per_req + num_query_per_req
129
- BLOCK_SIZE = min(256, triton.next_power_of_2(max_tokens_per_req))
130
- num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
131
+ BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req))
132
+ num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE
133
grid = (batch_size, num_blocks)
134
135
has_num_rejected = num_rejected_tokens_gpu is not None
vllm/v1/worker/cpu_model_runner.py
+14
-2
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import sys
4
from contextlib import contextmanager
5
from typing import Any
6
@@ -78,7 +79,7 @@ class CPUModelRunner(GPUModelRunner):
79
# Speculative decoding fallbacks
80
import vllm.v1.sample.rejection_sampler
81
import vllm.v1.spec_decode.llm_base_proposer
81
- import vllm.v1.spec_decode.utils
82
+ import vllm.v1.spec_decode.utils as spec_decode_utils
83
84
vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = (
85
cpu_tl.eagle_prepare_inputs_padded_kernel
@@ -89,7 +90,18 @@ class CPUModelRunner(GPUModelRunner):
90
vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = (
91
cpu_tl.copy_and_expand_eagle_inputs_kernel
92
)
92
- vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel = (
93
+ spec_decode_utils.copy_and_expand_dflash_inputs_kernel = (
94
+ cpu_tl.copy_and_expand_dflash_inputs_kernel
95
+ )
96
+ dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash")
97
+ if dflash_module is not None:
98
+ dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel"
99
+ setattr(
100
+ dflash_module,
101
+ dflash_kernel_name,
102
+ cpu_tl.copy_and_expand_dflash_inputs_kernel,
103
+ )
104
+ spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = (
105
cpu_tl.eagle_step_slot_mapping_metadata_kernel
106
)
107
vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = (