From 2396d91e931295df877cdad5e2ac0de4e35dab9f Mon Sep 17 00:00:00 2001 From: guybd Date: Thu, 25 Jun 2026 10:32:48 +0300 Subject: [PATCH] [CPU][Spec Decode] Enable DFlash SD for CPU (#44029) Signed-off-by: guybd Signed-off-by: Guy Boudoukh Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- csrc/cpu/spec_decode_utils.cpp | 83 +++++++++++++ csrc/cpu/torch_bindings.cpp | 23 ++++ docs/design/attention_backends.md | 2 +- vllm/model_executor/models/qwen3_dflash.py | 2 +- vllm/utils/cpu_triton_utils.py | 134 +++++++++++++++++++++ vllm/v1/attention/backends/cpu_attn.py | 4 + vllm/v1/spec_decode/dflash.py | 10 +- vllm/v1/worker/cpu_model_runner.py | 16 ++- 8 files changed, 266 insertions(+), 8 deletions(-) diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp index a76b8bc6937..30192196b95 100644 --- a/csrc/cpu/spec_decode_utils.cpp +++ b/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& 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(); + const int64_t* target_pos_ptr = target_positions.data_ptr(); + const int32_t* block_table_ptr = block_table.data_ptr(); + const int32_t* query_start_ptr = query_start_loc.data_ptr(); + const int64_t* rejected_ptr = + has_num_rejected && num_rejected_tokens.has_value() + ? num_rejected_tokens.value().data_ptr() + : nullptr; + + int64_t* out_ids_ptr = out_input_ids.data_ptr(); + int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr(); + int64_t* out_query_pos_ptr = out_query_positions.data_ptr(); + int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr(); + int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr(); + int32_t* out_token_idx_ptr = out_token_indices.data_ptr(); + + 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(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, diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 9cef2d0d535..bc02511eb80 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/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& 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, " diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index f965127cbfb..d268d5b4db2 100644 --- a/docs/design/attention_backends.md +++ b/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 | diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 820260f795c..36c0a357878 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/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 diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index c3cedf9a7ba..3b5012d0175 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/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 ) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b2e186ac3b7..056107c364d 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/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, diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index f76305d0857..bae6935cef8 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/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 diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 45ebe8a4da5..87f8cb154dc 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/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 = (