diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..24e115f --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +*.patch -whitespace diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1d93174..b58f0fd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,6 +20,19 @@ jobs: run: python -m compileall -q dashboard benchmarks overlay - name: Shell syntax run: find scripts dashboard -type f -name '*.sh' -print0 | xargs -0 -n1 bash -n + - name: vLLM patch applicability + run: | + source upstream.lock + git clone --filter=blob:none --sparse --no-checkout \ + "$VLLM_REPOSITORY" .build/vllm + git -C .build/vllm sparse-checkout set \ + vllm/v1/worker/gpu/spec_decode/dflash/speculator.py + git -C .build/vllm fetch --depth=1 origin "$VLLM_COMMIT" + git -C .build/vllm checkout --detach FETCH_HEAD + ./scripts/apply-vllm-patches.sh .build/vllm + git -C .build/vllm diff --check + python -m py_compile \ + .build/vllm/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py - name: Overlay license markers run: | missing="$(find overlay/vllm -type f -name '*.py' \ diff --git a/patches/vllm/0001-dspark-graph-replay-safety.patch b/patches/vllm/0001-dspark-graph-replay-safety.patch new file mode 100644 index 0000000..81f932e --- /dev/null +++ b/patches/vllm/0001-dspark-graph-replay-safety.patch @@ -0,0 +1,114 @@ +From dfecbb52ce1801d12e99c1d8bfb6e37a5530cd20 Mon Sep 17 00:00:00 2001 +From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> +Date: Thu, 6 Aug 2026 21:42:50 +0800 +Subject: [PATCH] [Bugfix][Spec Decode] Make DSpark graph replay and draft KV + writes safe + +Two defects on the DSpark draft path: + +- `sample_idx_mapping` was zero-filled, so CUDA graph capture executed the full + buffer before a real batch had populated it and every padding row scattered + into request slot 0. Use -1 to mark an inert row. +- Draft KV could be written into physical block 0, the null block, when a + sliding-window block table carried evicted or padding entries and for rejected + suffix rows. Route both to PAD_SLOT_ID. + +Co-authored-by: OpenAI Codex +Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> +--- + .../gpu/spec_decode/dflash/speculator.py | 42 +++++++++++++++---- + 1 file changed, 33 insertions(+), 9 deletions(-) + +diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +index 97c284c03..eb5dc470b 100644 +--- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py ++++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +@@ -84,8 +84,11 @@ class DFlashSpeculator(DraftModelSpeculator): + self.sample_pos = torch.zeros( + max_num_sampled_tokens, dtype=torch.int64, device=device + ) +- self.sample_idx_mapping = torch.zeros( +- max_num_sampled_tokens, dtype=torch.int32, device=device ++ # -1 marks an inert sampling row. CUDA graph capture can execute the ++ # full buffer before a real batch has populated it, so zero would make ++ # every padding row scatter into request slot 0. ++ self.sample_idx_mapping = torch.full( ++ (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device + ) + # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into + # draft_logits[req, step, :]. +@@ -255,7 +258,6 @@ class DFlashSpeculator(DraftModelSpeculator): + num_tokens_across_dp, + cudagraph_runtime_mode, + ) +- + num_sample = num_reqs * self.num_speculative_steps + sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] + # sample_pos is the predicted token's position Q; verification keys +@@ -520,6 +522,7 @@ def _prepare_dflash_inputs_kernel( + + num_rejected = tl.load(num_rejected_ptr + req_idx) + valid_ctx_end = ctx_end - num_rejected ++ num_valid_ctx = valid_ctx_end - ctx_start + + num_sampled = tl.load(num_sampled_ptr + req_idx) + if num_sampled > 0: +@@ -533,20 +536,34 @@ def _prepare_dflash_inputs_kernel( + + j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + is_ctx = j < num_ctx +- is_query = (j >= num_ctx) & (j < num_ctx + num_query_per_req) +- query_off = j - num_ctx ++ is_valid_ctx = j < num_valid_ctx ++ is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req) ++ query_off = j - num_valid_ctx + + # --- Context positions / slots --- + ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) +- ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0) ++ ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) + ctx_block_num = ctx_pos // block_size + ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) + ctx_block_id = tl.load( + block_table_ptr + req_idx * block_table_stride + ctx_block_num, +- mask=is_ctx, ++ mask=is_valid_ctx, + other=0, + ).to(tl.int64) +- ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size) ++ # Block 0 is the null block. Old sliding-window context positions can map ++ # to it after eviction; rejected suffix rows are invalid context as well. ++ # Neither kind of row may write draft KV into physical block 0. ++ ctx_resident = is_valid_ctx & (ctx_block_id != 0) ++ ctx_slot = tl.where( ++ ctx_resident, ++ ctx_block_id * block_size + (ctx_pos % block_size), ++ PAD_SLOT_ID, ++ ) ++ # Stored over the full [0, num_ctx) span while the loads above are masked to ++ # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and ++ # PAD_SLOT_ID. That is intentional — those rows write no KV and their ++ # positions are never consumed, but the span must stay fully initialized so ++ # a replayed graph cannot observe a stale value from an earlier batch. + tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) + tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) + +@@ -563,7 +580,14 @@ def _prepare_dflash_inputs_kernel( + mask=is_query, + other=0, + ).to(tl.int64) +- q_slot = q_block_id * block_size + (query_pos % block_size) ++ # A null block is never a writable cache slot. This can occur when a ++ # sliding-window block table contains evicted/global padding entries. ++ q_resident = is_query & (q_block_id != 0) ++ q_slot = tl.where( ++ q_resident, ++ q_block_id * block_size + (query_pos % block_size), ++ PAD_SLOT_ID, ++ ) + + tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) + clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) +-- +2.43.0 + diff --git a/patches/vllm/README.md b/patches/vllm/README.md new file mode 100644 index 0000000..d1a6c99 --- /dev/null +++ b/patches/vllm/README.md @@ -0,0 +1,11 @@ +# vLLM patches + +These patches are applied to the exact `VLLM_COMMIT` from `upstream.lock` +before the repository overlay is copied into the vLLM source tree. + +`0001-dspark-graph-replay-safety.patch` is an unmodified, path-limited +backport of vLLM commit +[`dfecbb52ce1801d12e99c1d8bfb6e37a5530cd20`](https://github.com/vllm-project/vllm/commit/dfecbb52ce1801d12e99c1d8bfb6e37a5530cd20), +merged through [vLLM PR #51538](https://github.com/vllm-project/vllm/pull/51538). +It applies cleanly to the pinned vLLM commit and retains the upstream authorship +and DCO trailers in the patch header. diff --git a/scripts/apply-vllm-patches.sh b/scripts/apply-vllm-patches.sh new file mode 100755 index 0000000..498770e --- /dev/null +++ b/scripts/apply-vllm-patches.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT +set -euo pipefail + +source_dir="${1:?usage: apply-vllm-patches.sh VLLM_SOURCE_DIR}" +root="$(cd "$(dirname "$0")/.." && pwd)" + +if ! git -C "$source_dir" rev-parse --git-dir >/dev/null 2>&1; then + echo "vLLM source is not a git checkout: $source_dir" >&2 + exit 1 +fi + +shopt -s nullglob +patches=("$root"/patches/vllm/*.patch) +for patch in "${patches[@]}"; do + echo "Applying vLLM patch: $(basename "$patch")" + git -C "$source_dir" apply --check --whitespace=error-all "$patch" + git -C "$source_dir" apply --whitespace=error-all "$patch" +done diff --git a/scripts/build-image.sh b/scripts/build-image.sh index cc6980b..260bacb 100755 --- a/scripts/build-image.sh +++ b/scripts/build-image.sh @@ -18,6 +18,7 @@ git -C "$source_dir" fetch --tags origin git -C "$source_dir" checkout --detach "$VLLM_COMMIT" git -C "$source_dir" reset --hard "$VLLM_COMMIT" git -C "$source_dir" clean -fdx +"$root/scripts/apply-vllm-patches.sh" "$source_dir" rsync -a "$root/overlay/vllm/" "$source_dir/vllm/" # vLLM 0.25 keeps its primary Dockerfile under docker/, not at repository root.