spec_decode: consolidate DSpark correctness and load-aware serving - #109
spec_decode: consolidate DSpark correctness and load-aware serving#109voipmonitor wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughThis PR adds DSpark capacity-based speculative verification with varlen CUDA-graph dispatch, confidence and FP8 draft heads, DFlash cache-aware drafting, SPS profiling, expanded configuration, scheduler updates, and comprehensive GPU and unit-test coverage. ChangesDSpark configuration and model integration
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Canonical merge order, cross-repository dependencies, post-merge validation, and release provenance rules are tracked in local-inference-lab/rtx6kpro#29. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (5)
vllm/v1/core/sched/scheduler.py (1)
1587-1589: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer direct attribute access.
Since
self.acceptance_length_controlleris explicitly initialized toNonein__init__, you can access it directly instead of usinggetattr.♻️ Proposed refactor
- acceptance_length_controller = getattr( - self, "acceptance_length_controller", None - ) + acceptance_length_controller = self.acceptance_length_controller🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/core/sched/scheduler.py` around lines 1587 - 1589, In the scheduler logic around acceptance_length_controller, replace the getattr(self, "acceptance_length_controller", None) lookup with direct access to self.acceptance_length_controller, preserving the existing None behavior established by initialization.tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py (2)
610-627: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the in-place compaction named by this test.
This currently verifies only that
warmup()does not raise. Assert the expected buffer or metadata mutations, or rename it as an explicit smoke test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py` around lines 610 - 627, Update test_varlen_capacity_manager_warmup_compacts_inputs_in_place to verify the expected in-place buffer or metadata changes after handler.warmup, rather than only synchronizing and checking for exceptions; if no observable mutation can be asserted, rename the test to clearly identify it as a warmup smoke test.
12-13: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep host-only tests runnable without CUDA.
The module-level skip also disables controller, indexer, and dispatch tests that do not use a GPU. Split those tests into a host-compatible module or apply CUDA skips only to GPU-dependent cases.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py` around lines 12 - 13, Limit the CUDA availability skip in test_gpu_model_runner_v2_draft_capacity.py to GPU-dependent draft capacity tests instead of applying it at module scope. Move controller, indexer, and dispatch tests into a host-compatible module or mark only the relevant GPU test cases with the CUDA skip, keeping non-GPU tests runnable without CUDA.vllm/v1/worker/gpu/spec_decode/capacity.py (2)
825-827: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEnforce the capacity-array length invariant.
Use
strict=Trueso a mismatched request/capacity array cannot silently leave rows unprocessed.for req_idx, (num_draft_tokens, capacity) in enumerate( - zip(input_batch.num_draft_tokens_per_req, capacities) + zip(input_batch.num_draft_tokens_per_req, capacities, strict=True) ):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/worker/gpu/spec_decode/capacity.py` around lines 825 - 827, Update the zip call in the loop over input_batch.num_draft_tokens_per_req and capacities to use strict=True, ensuring mismatched arrays raise an error instead of silently skipping entries.Source: Linters/SAST tools
901-907: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse a non-flagged checksum for TP diagnostics.
Ruff treats both MD5 calls as errors. SHA-256 preserves the debug comparison semantics and avoids security-linter exceptions.
- return hashlib.md5(x.cpu().numpy().tobytes()).hexdigest() + return hashlib.sha256(x.cpu().numpy().tobytes()).hexdigest() ... - hashlib.md5(capacities[slots_np].tobytes()).hexdigest(), + hashlib.sha256(capacities[slots_np].tobytes()).hexdigest(),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vllm/v1/worker/gpu/spec_decode/capacity.py` around lines 901 - 907, Replace both MD5 checksum calls in gpu_md5 and the payload construction with SHA-256 hashing, preserving the existing byte inputs and hexadecimal digest comparison behavior for TP diagnostics.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/v1/spec_decode/test_dspark_fp8_draft_head.py`:
- Around line 141-144: Replace the exact argmax equality assertion in the CUDA
draft-choice check with the top-2 margin versus observed-error validation used
by test_draft_argmax_agreement_bf16_vs_fp8_emulated. Preserve the existing
real/emulated tensors and ensure near-tied rows are accepted when their margin
is within the measured error bound.
In `@vllm/config/speculative.py`:
- Around line 1379-1398: Remove the process-wide VLLM_MOE_SKIP_PADDING mutation
from SpeculativeConfig validation. Propagate the mask-mode requirement as
per-engine configuration, and enable padding skipping only after validating that
the selected MoE backend supports the -1 sentinel, preserving the existing
environment override behavior where applicable.
In `@vllm/model_executor/layers/fp8_draft_head.py`:
- Around line 78-89: Update the activation scaling in the logits computation to
cast hidden_states to float32 before abs().amax() and before calculating the FP8
scale, ensuring zero rows do not produce inf or NaN. Preserve the existing FP8
conversion and output dtype behavior in the surrounding logits path.
In `@vllm/models/deepseek_v4/common/ops/cache_utils.py`:
- Around line 760-770: Update the sparse-index construction around the
documented sparse_indices layout so the complete row width, swa_index_width plus
padded top-k capacity, is rounded to a multiple of four int32 elements. Apply
this consistently wherever padded_topk, sparse_indices shapes, offsets, or
related slices are computed, including decode and prefill paths, while
preserving the SWA-prefix and compressed/top-k regions.
In `@vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py`:
- Around line 312-315: Fix the SWA-only branch by updating the
prefill_topk_indices initialization to pass a single shape tuple to
decode_swa_indices.new_empty, preserving the intended shape of
(num_prefill_tokens, 0).
In `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py`:
- Around line 312-315: Update the parameterized docstrings in
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py at lines 312-315 and
1006-1009 to use Google-style sections: add an Args: section documenting
num_cached_tokens and num_cached_tokens_np at the first site, and all five
parameters at the second site.
- Around line 319-325: Update _has_unaligned_cached_prefix to check alignment
only for KV groups identified by self.draft_kv_cache_group_ids, matching the
group selection used by set_attn; exclude unrelated target KV groups so they
cannot disable drafting.
- Around line 506-513: Resolve and validate active_num_speculative_steps before
the unaligned cached-prefix early return in the drafting flow. When drafting is
disabled, return only the active draft-width slice filled with -1 values, rather
than the configured maximum-width slice, while preserving the existing warning
and behavior for normal drafting.
---
Nitpick comments:
In `@tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py`:
- Around line 610-627: Update
test_varlen_capacity_manager_warmup_compacts_inputs_in_place to verify the
expected in-place buffer or metadata changes after handler.warmup, rather than
only synchronizing and checking for exceptions; if no observable mutation can be
asserted, rename the test to clearly identify it as a warmup smoke test.
- Around line 12-13: Limit the CUDA availability skip in
test_gpu_model_runner_v2_draft_capacity.py to GPU-dependent draft capacity tests
instead of applying it at module scope. Move controller, indexer, and dispatch
tests into a host-compatible module or mark only the relevant GPU test cases
with the CUDA skip, keeping non-GPU tests runnable without CUDA.
In `@vllm/v1/core/sched/scheduler.py`:
- Around line 1587-1589: In the scheduler logic around
acceptance_length_controller, replace the getattr(self,
"acceptance_length_controller", None) lookup with direct access to
self.acceptance_length_controller, preserving the existing None behavior
established by initialization.
In `@vllm/v1/worker/gpu/spec_decode/capacity.py`:
- Around line 825-827: Update the zip call in the loop over
input_batch.num_draft_tokens_per_req and capacities to use strict=True, ensuring
mismatched arrays raise an error instead of silently skipping entries.
- Around line 901-907: Replace both MD5 checksum calls in gpu_md5 and the
payload construction with SHA-256 hashing, preserving the existing byte inputs
and hexadecimal digest comparison behavior for TP diagnostics.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 565db8b1-3c58-4871-9fa4-017e1e9ca7c7
📒 Files selected for processing (61)
benchmarks/profile_dspark_sps_curve.pytests/engine/test_arg_utils.pytests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yamltests/test_config.pytests/v1/attention/test_deepseek_v4_dspark_metadata.pytests/v1/spec_decode/test_acceptance_length_controller.pytests/v1/spec_decode/test_dflash_cudagraph_lifetime.pytests/v1/spec_decode/test_dflash_prefix_cache_masking.pytests/v1/spec_decode/test_dspark_fp8_draft_head.pytests/v1/spec_decode/test_dynamic_sd_cug.pytests/v1/worker/test_gpu_block_table.pytests/v1/worker/test_gpu_model_runner_v2_draft_capacity.pytests/v1/worker/test_gpu_sampling_states_seed.pytests/v1/worker/test_mixed_warmup_gate.pytools/pre_commit/generate_attention_backend_docs.pyvllm/config/speculative.pyvllm/engine/arg_utils.pyvllm/envs.pyvllm/model_executor/layers/fp8_draft_head.pyvllm/model_executor/layers/sparse_attn_indexer.pyvllm/model_executor/models/qwen3_dspark.pyvllm/models/deepseek_v4/common/ops/cache_utils.pyvllm/models/deepseek_v4/nvidia/dspark.pyvllm/models/deepseek_v4/nvidia/flashinfer_sparse.pyvllm/models/deepseek_v4/sparse_mla.pyvllm/platforms/cuda.pyvllm/utils/deep_gemm.pyvllm/v1/attention/backend.pyvllm/v1/attention/backends/flashinfer.pyvllm/v1/attention/backends/mla/indexer.pyvllm/v1/attention/backends/mla/sparse_swa.pyvllm/v1/attention/backends/utils.pyvllm/v1/core/sched/async_scheduler.pyvllm/v1/core/sched/output.pyvllm/v1/core/sched/scheduler.pyvllm/v1/worker/gpu/attn_utils.pyvllm/v1/worker/gpu/block_table.pyvllm/v1/worker/gpu/cudagraph_utils.pyvllm/v1/worker/gpu/dp_utils.pyvllm/v1/worker/gpu/input_batch.pyvllm/v1/worker/gpu/model_runner.pyvllm/v1/worker/gpu/model_states/default.pyvllm/v1/worker/gpu/sample/bad_words.pyvllm/v1/worker/gpu/sample/gumbel.pyvllm/v1/worker/gpu/sample/sampler.pyvllm/v1/worker/gpu/sample/states.pyvllm/v1/worker/gpu/spec_decode/autoregressive/speculator.pyvllm/v1/worker/gpu/spec_decode/capacity.pyvllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.pyvllm/v1/worker/gpu/spec_decode/dflash/cudagraph.pyvllm/v1/worker/gpu/spec_decode/dflash/speculator.pyvllm/v1/worker/gpu/spec_decode/dspark/capacity.pyvllm/v1/worker/gpu/spec_decode/dspark/online_sts.pyvllm/v1/worker/gpu/spec_decode/dspark/speculator.pyvllm/v1/worker/gpu/spec_decode/dspark/utils.pyvllm/v1/worker/gpu/spec_decode/rejection_sampler.pyvllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.pyvllm/v1/worker/gpu/spec_decode/speculator.pyvllm/v1/worker/gpu/states.pyvllm/v1/worker/gpu/warmup.pyvllm/v1/worker/ubatch_utils.py
| # The CUDA path must make exactly the same draft choice as the emulation. | ||
| # Quantization-induced flips against BF16 are covered separately by the | ||
| # top-2 margin bound in test_draft_argmax_agreement_bf16_vs_fp8_emulated. | ||
| torch.testing.assert_close(real.argmax(dim=-1), emulated.argmax(dim=-1)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 'test_fp8_draft_head_logits_cuda_matches_emulation|_scaled_mm' \
tests/v1/spec_decode/test_dspark_fp8_draft_head.py \
vllm/model_executor/layers/fp8_draft_head.py
rg -n -C2 'test_dspark_fp8_draft_head|SM89|H100|L4' .github tests || trueRepository: local-inference-lab/vllm
Length of output: 21145
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' tests/v1/spec_decode/test_dspark_fp8_draft_head.py
printf '\n---\n'
sed -n '1,180p' vllm/model_executor/layers/fp8_draft_head.pyRepository: local-inference-lab/vllm
Length of output: 9907
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tests/v1/spec_decode/test_dspark_fp8_draft_head.py")
text = path.read_text()
for i, line in enumerate(text.splitlines(), 1):
if 120 <= i <= 150:
print(f"{i:4}: {line}")
PYRepository: local-inference-lab/vllm
Length of output: 1385
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tests/v1/spec_decode/test_dspark_fp8_draft_head.py")
tree = ast.parse(path.read_text())
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == "test_fp8_draft_head_logits_cuda_matches_emulation":
print(f"Function: {node.name} @ lines {node.lineno}-{node.end_lineno}")
for stmt in node.body:
if hasattr(stmt, 'lineno'):
print(f"{stmt.lineno:4}-{getattr(stmt, 'end_lineno', stmt.lineno):4}: {type(stmt).__name__}")
PYRepository: local-inference-lab/vllm
Length of output: 435
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re
text = Path("tests/v1/spec_decode/test_dspark_fp8_draft_head.py").read_text().splitlines()
for i, line in enumerate(text, 1):
if "assert_close(real.argmax" in line or "top2" in line or "max_error" in line:
print(f"{i}: {line}")
PYRepository: local-inference-lab/vllm
Length of output: 342
🏁 Script executed:
#!/bin/bash
set -euo pipefail
nl -ba tests/v1/spec_decode/test_dspark_fp8_draft_head.py | sed -n '1,220p'Repository: local-inference-lab/vllm
Length of output: 202
🏁 Script executed:
#!/bin/bash
set -euo pipefail
nl -ba tests/v1/spec_decode/test_dspark_fp8_draft_head.py | sed -n '120,150p'
printf '\n---\n'
nl -ba vllm/model_executor/layers/fp8_draft_head.py | sed -n '68,96p'Repository: local-inference-lab/vllm
Length of output: 202
Relax the final argmax check
real.argmax(dim=-1) can differ from emulated.argmax(dim=-1) on near-tied rows even when the max-error bound passes. Use the same top-2 margin vs. observed-error check as the BF16 test instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/v1/spec_decode/test_dspark_fp8_draft_head.py` around lines 141 - 144,
Replace the exact argmax equality assertion in the CUDA draft-choice check with
the top-2 margin versus observed-error validation used by
test_draft_argmax_agreement_bf16_vs_fp8_emulated. Preserve the existing
real/emulated tensors and ensure near-tied rows are accepted when their margin
is within the measured error bound.
| if ( | ||
| self.method == "dspark" | ||
| and self.dspark_capacity_verification_mode == "mask" | ||
| and ( | ||
| self.dspark_confidence_threshold > 0.0 | ||
| or self.dspark_budget_frac < 1.0 | ||
| or self.dspark_sps_curve is not None | ||
| ) | ||
| and "VLLM_MOE_SKIP_PADDING" not in os.environ | ||
| ): | ||
| # Mask mode keeps pruned verify rows in the batch as padding; the | ||
| # pruning only saves work if MoE kernels skip those rows. Set | ||
| # here (frontend) so spawned workers inherit it before their env | ||
| # caches freeze. Set VLLM_MOE_SKIP_PADDING=0 to override. | ||
| logger.info( | ||
| "DSpark mask capacity mode: defaulting VLLM_MOE_SKIP_PADDING=1 " | ||
| "so MoE kernels skip pruned verify rows." | ||
| ) | ||
| os.environ["VLLM_MOE_SKIP_PADDING"] = "1" | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Do not globally force MoE padding skipping from SpeculativeConfig.
Line 1397 changes process-wide behavior for every engine, while vllm/envs.py Lines 1647-1651 explicitly state that not all MoE kernels support the -1 sentinel. An unrelated or unsupported MoE backend may consequently misroute tokens or fail.
Propagate this as per-engine configuration and enable it only after validating the selected MoE backend’s capability.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/config/speculative.py` around lines 1379 - 1398, Remove the process-wide
VLLM_MOE_SKIP_PADDING mutation from SpeculativeConfig validation. Propagate the
mask-mode requirement as per-engine configuration, and enable padding skipping
only after validating that the selected MoE backend supports the -1 sentinel,
preserving the existing environment override behavior where applicable.
| act_max = hidden_states.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) | ||
| act_fp8 = (hidden_states * (_FP8_MAX / act_max)).to(torch.float8_e4m3fn) | ||
| logits = torch._scaled_mm( | ||
| act_fp8, | ||
| head.weight_fp8.t(), | ||
| scale_a=head.unit_scale, | ||
| scale_b=head.unit_scale, | ||
| out_dtype=hidden_states.dtype, | ||
| ) | ||
| logits = logits * head.row_scale | ||
| logits = logits * (act_max / _FP8_MAX).to(hidden_states.dtype) | ||
| return logits |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
For PyTorch 2.11 CUDA tensors, is 448.0 / float16_tensor evaluated in float16, allowing overflow before conversion to float8?
💡 Result:
In PyTorch 2.11, the expression 448.0 / float16_tensor is evaluated in float16, which can lead to overflow if the intermediate values or the final result exceed the representable range of float16 (approximately ±65504). PyTorch uses specific type promotion rules for arithmetic operations. When a Python float scalar (like 448.0) is used in an operation with a tensor, the scalar is typically treated as having the same dtype as the tensor, or it is promoted to the tensor's dtype if it does not introduce a higher-precision category [1]. Because float16 is a floating-point type, and the scalar 448.0 is a floating-point number, the operation is performed using float16 precision [1][2]. This behavior is well-documented in issues where operations involving float16 tensors and scalars (or tensors of other types that are cast to float16) result in silent overflow to inf rather than promoting the computation to a wider type like float32 [3][4][5]. Consequently, the division is not automatically promoted to float32, and if the result exceeds the limits of float16, it will overflow before any potential subsequent conversion to another format [3][6]. To avoid this, you should explicitly cast the tensor to float32 before performing the division (e.g., 448.0 / float16_tensor.float()), which will ensure the calculation is performed in float32.
Citations:
- 1: https://docs.pytorch.org/docs/2.3/tensor_attributes.html
- 2: https://github.com/pytorch/pytorch/blob/master/test/test_type_promotion.py
- 3: torch.tensor([0.01], dtype=torch.float16) * torch.tensor(65536, dtype=torch.float32) returns INF pytorch/pytorch#98691
- 4: [fix] torch.where silently overflows fp16 scalars (issue #187429) pytorch/pytorch#187908
- 5:
torch.cliphas checks for float16 scalar overflow on CPU but not on GPU pytorch/pytorch#171356 - 6: Inconsistent float16 overflow behavior between CPU and CUDA devices pytorch/pytorch#152731
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the target file and inspect the relevant section.
ast-grep outline vllm/model_executor/layers/fp8_draft_head.py --view expanded || true
printf '\n--- file excerpt ---\n'
sed -n '1,180p' vllm/model_executor/layers/fp8_draft_head.py
printf '\n--- search for callers / dtype assumptions ---\n'
rg -n "fp8_draft_head|draft head|float8_e4m3fn|unit_scale|row_scale|hidden_states" vllm/model_executor -g '*.py'Repository: local-inference-lab/vllm
Length of output: 50381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("vllm/model_executor/layers/fp8_draft_head.py")
print(path.exists())
if path.exists():
for i, line in enumerate(path.read_text().splitlines(), 1):
if 60 <= i <= 110:
print(f"{i:4d}: {line}")
PY
printf '\n--- references ---\n'
rg -n "fp8_draft_head|float8_e4m3fn|row_scale|unit_scale|scaled_mm" vllm -g '*.py'Repository: local-inference-lab/vllm
Length of output: 50380
Compute the activation scale in float32. hidden_states may be float16 here, so _FP8_MAX / act_max can overflow to inf before the FP8 cast; an all-zero row then becomes 0 * inf and produces NaN. Cast to float32 before the reduction and scaling.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/model_executor/layers/fp8_draft_head.py` around lines 78 - 89, Update
the activation scaling in the logits computation to cast hidden_states to
float32 before abs().amax() and before calculating the FP8 scale, ensuring zero
rows do not produce inf or NaN. Preserve the existing FP8 conversion and output
dtype behavior in the surrounding logits path.
| Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width + | ||
| padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the | ||
| rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length | ||
| per token). Decode tokens read precomputed SWA/compressed indices; prefill | ||
| tokens derive their SWA window from the position and translate local | ||
| compressed indices to global slots via the block tables. | ||
| """ | ||
| assert decode_swa_indices.dtype == torch.int32 | ||
| assert decode_swa_indices.dim() == 2 | ||
| assert decode_swa_indices.shape[-1] == window_size | ||
| swa_index_width = decode_swa_indices.shape[-1] | ||
| assert swa_index_width >= window_size |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Align the complete sparse-index row, not only the top-k suffix.
A dynamic swa_index_width need not be divisible by four. Rounding only
padded_topk therefore leaves subsequent rows misaligned, violating the
documented 16-byte FlashInfer load requirement.
Proposed alignment fix
- padded_topk = max(topk, decode_compressed_topk)
- padded_topk = (padded_topk + 3) // 4 * 4
+ min_topk_width = max(topk, decode_compressed_topk)
+ row_width = (swa_index_width + min_topk_width + 3) // 4 * 4
+ padded_topk = row_width - swa_index_widthAlso applies to: 816-820, 829-832, 863-869, 931-935, 945-950, 993-1010, 1026-1028, 1072-1078
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/models/deepseek_v4/common/ops/cache_utils.py` around lines 760 - 770,
Update the sparse-index construction around the documented sparse_indices layout
so the complete row width, swa_index_width plus padded top-k capacity, is
rounded to a multiple of four int32 elements. Apply this consistently wherever
padded_topk, sparse_indices shapes, offsets, or related slices are computed,
including decode and prefill paths, while preserving the SWA-prefix and
compressed/top-k regions.
| if swa_only: | ||
| assert self.topk_indices_buffer is not None | ||
| compressed_kv_cache = swa_k_cache | ||
| decode_compressed_indices = None | ||
| prefill_topk_indices = self.topk_indices_buffer[ | ||
| num_decode_tokens:num_tokens, :0 | ||
| ] | ||
| prefill_topk_indices = decode_swa_indices.new_empty(num_prefill_tokens, 0) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
set -e
python3 - <<'PY'
from pathlib import Path
path = Path('vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py')
lines = path.read_text().splitlines()
for i in range(300, 321):
print(f"{i+1}: {lines[i]}")
PYRepository: local-inference-lab/vllm
Length of output: 1131
🏁 Script executed:
python3 - <<'PY'
import inspect
try:
import torch
except Exception as e:
print(f"IMPORT_ERROR: {e}")
raise SystemExit(0)
print("torch_version:", torch.__version__)
sig = inspect.signature(torch.Tensor.new_empty)
print("signature:", sig)
print("doc_head:", (torch.Tensor.new_empty.__doc__ or "").splitlines()[:8])
x = torch.zeros(1)
tests = [
("tuple", lambda: x.new_empty((2, 0))),
("variadic", lambda: x.new_empty(2, 0)),
]
for name, fn in tests:
try:
t = fn()
print(name, "OK", tuple(t.shape), t.dtype, t.device)
except Exception as e:
print(name, type(e).__name__, e)
PYRepository: local-inference-lab/vllm
Length of output: 200
🏁 Script executed:
python3 - <<'PY'
import inspect
try:
import torch
except Exception as e:
print(f"IMPORT_ERROR: {e}")
raise SystemExit(0)
print("torch_version:", torch.__version__)
print("signature:", inspect.signature(torch.Tensor.new_empty))
print("doc_head:")
for line in (torch.Tensor.new_empty.__doc__ or "").splitlines()[:12]:
print(line)
x = torch.zeros(1)
for name, fn in [
("tuple", lambda: x.new_empty((2, 0))),
("variadic", lambda: x.new_empty(2, 0)),
]:
try:
t = fn()
print(name, "OK", tuple(t.shape))
except Exception as e:
print(name, type(e).__name__, e)
PYRepository: local-inference-lab/vllm
Length of output: 200
🌐 Web query:
PyTorch Tensor.new_empty signature size positional arguments docs
💡 Result:
The PyTorch method torch.Tensor.new_empty has the following function signature [1][2]: Tensor.new_empty(size, , dtype=None, device=None, requires_grad=False, layout=torch.strided, pin_memory=False) The size argument is a positional argument that defines the shape of the output tensor [1][2]. It can be provided as a list, tuple, or torch.Size of integers [1][2]. All other arguments (dtype, device, requires_grad, layout, and pin_memory) are keyword-only arguments, indicated by the asterisk () in the signature [1][2].
Citations:
- 1: https://docs.pytorch.org/docs/stable/generated/torch.Tensor.new_empty.md
- 2: https://docs.pytorch.org/docs/main/generated/torch.Tensor.new_empty.html
Pass a shape tuple to new_empty. Tensor.new_empty takes a single size argument, so decode_swa_indices.new_empty(num_prefill_tokens, 0) will raise TypeError in the SWA-only path. Use decode_swa_indices.new_empty((num_prefill_tokens, 0)).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py` around lines 312 - 315,
Fix the SWA-only branch by updating the prefill_topk_indices initialization to
pass a single shape tuple to decode_swa_indices.new_empty, preserving the
intended shape of (num_prefill_tokens, 0).
| """Register the runner's per-request-slot cache-restored token counts. | ||
|
|
||
| Indexed by req_state_idx; see the buffer comment in __init__. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use Google-style sections in the newly added parameterized docstrings.
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py#L312-L315: documentnum_cached_tokensandnum_cached_tokens_npunderArgs:.vllm/v1/worker/gpu/spec_decode/dflash/speculator.py#L1006-L1009: document all five parameters underArgs:.
As per coding guidelines, “Use Google-style docstrings in Python code, with Args:/Returns:/Raises: sections.”
📍 Affects 1 file
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py#L312-L315(this comment)vllm/v1/worker/gpu/spec_decode/dflash/speculator.py#L1006-L1009
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py` around lines 312 - 315,
Update the parameterized docstrings in
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py at lines 312-315 and
1006-1009 to use Google-style sections: add an Args: section documenting
num_cached_tokens and num_cached_tokens_np at the first site, and all five
parameters at the second site.
Source: Coding guidelines
| def _has_unaligned_cached_prefix(self, input_batch: InputBatch) -> bool: | ||
| req_state_indices = input_batch.idx_mapping_np[: input_batch.num_reqs] | ||
| cached = self.num_cached_tokens_np[req_state_indices] | ||
| return any( | ||
| np.any(cached % block_size != 0) | ||
| for block_size in self.block_tables.kernel_block_sizes | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Check alignment against draft KV groups only.
This iterates every target KV group, so an unrelated group can falsely disable drafting even when all draft groups align. Filter by self.draft_kv_cache_group_ids, matching the invariant enforced by set_attn.
return any(
np.any(cached % block_size != 0)
- for block_size in self.block_tables.kernel_block_sizes
+ for gid in self.draft_kv_cache_group_ids
+ for block_size in (self.block_tables.kernel_block_sizes[gid],)
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _has_unaligned_cached_prefix(self, input_batch: InputBatch) -> bool: | |
| req_state_indices = input_batch.idx_mapping_np[: input_batch.num_reqs] | |
| cached = self.num_cached_tokens_np[req_state_indices] | |
| return any( | |
| np.any(cached % block_size != 0) | |
| for block_size in self.block_tables.kernel_block_sizes | |
| ) | |
| def _has_unaligned_cached_prefix(self, input_batch: InputBatch) -> bool: | |
| req_state_indices = input_batch.idx_mapping_np[: input_batch.num_reqs] | |
| cached = self.num_cached_tokens_np[req_state_indices] | |
| return any( | |
| np.any(cached % block_size != 0) | |
| for gid in self.draft_kv_cache_group_ids | |
| for block_size in (self.block_tables.kernel_block_sizes[gid],) | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py` around lines 319 - 325,
Update _has_unaligned_cached_prefix to check alignment only for KV groups
identified by self.draft_kv_cache_group_ids, matching the group selection used
by set_attn; exclude unrelated target KV groups so they cannot disable drafting.
| if not dummy_run and self._has_unaligned_cached_prefix(input_batch): | ||
| logger.warning_once( | ||
| "DFlash/DSpark drafting is disabled for a batch containing a " | ||
| "block-unaligned cache-restored prefix because draft KV is " | ||
| "not available for the restored partial block." | ||
| ) | ||
| self.draft_tokens[:num_reqs].fill_(-1) | ||
| return self.draft_tokens[:num_reqs] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Preserve the active draft width when drafting is disabled.
This early return occurs before dynamic depth is resolved and returns the configured maximum width. For DSpark, a smaller num_speculative_tokens therefore becomes a full-width row of -1 placeholders, which downstream scheduling may still count by width. Compute and validate active_num_speculative_steps first, then return:
- return self.draft_tokens[:num_reqs]
+ return self.draft_tokens[
+ :num_reqs, :active_num_speculative_steps
+ ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vllm/v1/worker/gpu/spec_decode/dflash/speculator.py` around lines 506 - 513,
Resolve and validate active_num_speculative_steps before the unaligned
cached-prefix early return in the drafting flow. When drafting is disabled,
return only the active draft-width slice filled with -1 values, rather than the
configured maximum-width slice, while preserving the existing warning and
behavior for normal drafting.
1 participant
Summary
Consolidate the DSpark/DFlash work previously split across #105, #106, and
#107 into one reviewable change against the current
dev/fathomless-firmamentbranch.The three commits deliberately separate:
At default settings, the capacity controller and FP8 draft head remain
disabled and the configured fixed draft width is preserved.
Correctness and replay safety
deterministic across TP ranks;
graph capture combinations.
The block-table barrier fixes a reproduced cross-lane read/write race. The
long-row stress test completed 100 consecutive iterations without corruption.
Optional rowwise-FP8 draft head
VLLM_DSPARK_FP8_DRAFT_HEAD=1stores a rowwise FP8-E4M3 draft-only copy ofthe shared vocabulary head. Dynamic per-token activation quantization and
torch._scaled_mmreduce draft-head GEMM time. Target verification continuesto use the original target weights, so a draft mismatch can only reduce
acceptance; it cannot directly change an accepted target token.
Prior SM120 measurement: draft-head GEMM
2.73 ms -> 1.45 ms, producing anapproximately 3-5% single-stream DSpark decode gain.
Optional load-aware compact verification
The controller keeps the configured logical K and target acceptance semantics,
but can reduce physical draft and verification work once load reaches a
profiled activation knee.
The main controls are:
VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH=1VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW=<N>VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE=<N>--dspark-capacity-verification-mode varlen|maskdspark_*fields in--speculative-configMeasured TP2 Lucifer throughput, fixed K5 versus load-aware capacity:
The three-run C64 result was
2725.1 +/- 11.2 tok/s. A 120-second C64 runmeasured
2751.4 tok/s.Consolidation map
This PR supersedes:
The following changes intentionally remain separate because they are not
DSpark-specific:
#108 should merge before this PR because the capacity path touches the same
indexer module and current FF already contains the duplicate definition. The
cleanup itself is not included here. An integration merge of all PRs above was
performed in the listed dependency order and completed without conflicts.
Rejected or excluded work
CUTLASS target compute.
the hot path or were unsafe across asynchronous graph replays.
InstantTensor defaults, and release-specific KV-memory tuning are independent
work and are not hidden in this source PR.
Validation
ruff checkandruff format --checkpass for the complete integratedchanged-file set;
git diff --checkpasses;366 passedand714 skipped; it exposed one missing-attribute guard in the capacityscheduler path, which is fixed here and passed its 15-test focused rerun;
GPU and passed when the platform count was set to two;
test_scheduler_kv_connector_statsfixture failure reproducesunchanged on current FF and is not introduced by this PR;
Summary by CodeRabbit
New Features
Bug Fixes
Tests