dspark: draft-pass cleanup (-10% draft time) + upstream rejection-utils consolidation with opt-in block verification - #71
Closed
voipmonitor wants to merge 5 commits into
Conversation
… FP8 kernel The b12x MXFP8 dense GEMM wins below M~2048 on DSV4-Flash TP2 shapes (0.56-0.96x vs DeepGEMM) but loses at prefill M (1.05-1.21x at M=4096-8192). Keep the b12x pack for decode/small batches and additionally run the DeepGEMM weight processing on the original checkpoint params (replaced in place, no extra memory); token batches >= VLLM_B12X_FP8_LINEAR_DG_PREFILL_MIN_TOKENS (default 2049) route to DeepGemmFp8BlockScaledMMKernel.apply_weights inside the opaque custom op. Kill switch: VLLM_B12X_FP8_LINEAR_DG_PREFILL=0. Full-B12X DS4 TP2 A8 prefill with this + the b12x tiled-quant fix: 13,556-13,684 / 13,034-13,175 / 12,037-12,151 tok/s @ 8k/64k/128k (vs 12,468/11,971/11,133 before; Lucifer 13,442/12,622/11,716; pure DeepGEMM-linear hybrid 13,811/13,201/12,182). Decode cc1 unchanged (ITL 7.715ms). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two Triton kernels (FC1+SiLU staged in fp32 gmem, FC2 with router-weighted fp32-atomic scatter) that read the b12x N256/K128 in-place-repacked FP4 expert weights and e8m0 sfb grids directly via their verified inverse bit mappings (b12x tests/test_w4a8_rp_inverse_mapping.py), with BF16 activations (no input quantization) and fp32 accumulation. Numerics: cos vs fp32 oracle 0.999999 (better than the w4a8 dynamic kernel's 0.9990, which quantizes activations); cos vs the dynamic kernel 0.999. Isolated graph-replay at DS4-Flash TP2 shapes (E=256 K=4096 N=1024 topk=6): 28.7 us/layer at M=1 vs 34.8 us for the dynamic grouped kernel (plus ~3 us of wrapper fills/copies this path also bypasses). Enabled via VLLM_B12X_W4A8_MX_TINY_DECODE=1; engages only for quant_mode=w4a8_mx, silu, w31 layout, M==1, shape multiples of 256; everything else falls through to the b12x dynamic path. First E2E serve measured 138.7 tok/s DS4 TP2 A8 decode cc1 (from 133.7 baseline); a restart-to-restart variance question is still open (see rtx6kpro wiki). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sem='relaxed' on the fp32 atomic_add scatters read neutral in the isolated graph-replay microbench (28.7 vs 28.8 us/layer, L2-warm identical routing) but cost ~7% E2E decode in real serving (131.0-131.3 vs 138.7-138.8 tok/s, DS4 TP2 A8 cc1). Default semantics restored; E2E reproducible at 138.7+. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…utils (opt-in), keep fp64 acceptance hardening Wholesale adoption of upstream's rejection_sampler_utils.py (post-vllm-project#46781/ vllm-project#47093/vllm-project#47383: block-verification kernels, refactored logsumexp helpers, int64-safe offsets) with our extras re-applied on top: fp64 acceptance uniforms under USE_FP64 (upstream is fp32-only) and the active-row-guarded gumbel_block_argmax signature. Adds tl_rand32 to gumbel.py, 'block' to RejectionSampleMethod, and wires use_block_verification through RejectionSampler with a host-side all-greedy skip (the block prep kernels cost ~5-7% of a step and are inert at temp=0). Validation (synthetic kernel harness, no server): - temp=0: block == standard exactly (valid-prefix comparison) - distribution preservation: TVD vs target within sampling bounds - accepted length on Gaussian-noise drafts: +7% to +47% over standard across temp in {0.6,0.8,1.0} x draft quality E2E on DS4-Flash-DSpark TP2 (12x800-token probes x3, temp 0.7): - block 54.0/56.4/54.2%, standard 54.1/56.0/59.3% -> the real-draft gain is within run-to-run noise (DSpark's errors are bimodal, not the smooth ratio spectrum the synthetic model rewards). Left OPT-IN, not default: rejection_sample_method='block'. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Three graph-captured reductions in the DSpark draft forward, validated by in-kernel-graph CUDA-event timing (VLLM_DSPARK_TIMING=1, prints rolling prepare/graph-forward ms every 200 steps): - skip the confidence head unless VLLM_DSPARK_CONFIDENCE=1 (the serving speculator discards it; it is an lm-head-sized projection) - build the layer-invariant dspark topk_idxs once per draft step (shared across the three layers via a per-step module dict, cleared at forward_spec entry; capture-safe) - replace the per-layer window gather + full-width torch.cat with a persistent [b, window+block] staging buffer (index_select into the window slice + small block copy; kills one window-wide copy and two allocations per layer per step) Measured on DS4-Flash-DSpark TP2 A8 (600+ steps): draft graph-forward 2.01 -> 1.81 ms/step (-10%), prepare 0.33 -> 0.31 ms; 30k coherence clean. Also from this campaign, measured-and-falsified for the record: k-sweep (k=5 beats k=3/k=2: 251 vs 218 vs 192 tok/s greedy probes - verify time scales weakly with M, so shorter blocks just lose accepted tokens) and A16-vs-A8 MoE for DSpark (244.5 vs 241.3, tie within trajectory noise). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR adds a "block" rejection sampling method with new Triton kernels for residual-mass and cumulative log-probability computation across the rejection sampler stack. It also independently adds DeepGEMM prefill routing for b12x linear layers, a tiny-decode MoE Triton path, DSpark attention buffer reuse/caching with timing instrumentation, and a Gumbel random-value helper. ChangesBlock Verification Rejection Sampling
B12x DeepGEMM Prefill Routing
B12x Tiny-Decode MoE Path
DSpark Attention Caching and Timing
Gumbel Random Value Helper
Estimated code review effort: 5 (Critical) | ~100 minutes Sequence Diagram(s)sequenceDiagram
participant RejectionSampler
participant rejection_sample
participant CumulativeLogPKernel as _compute_cumulative_log_p_kernel
participant ResidualMassKernel as _compute_local_residual_mass_kernel
participant RejectionKernel as _rejection_kernel
participant ResampleKernel as _resample_kernel
RejectionSampler->>RejectionSampler: derive use_block_verification (batch temps)
RejectionSampler->>rejection_sample: call with use_block_verification
alt use_block_verification enabled
rejection_sample->>CumulativeLogPKernel: launch (compute cumulative_log_p)
rejection_sample->>ResidualMassKernel: launch (compute local_residual_mass)
end
rejection_sample->>RejectionKernel: launch (cumulative_log_p, residual_mass, flags)
RejectionKernel-->>rejection_sample: accepted_length, rejected_idx
rejection_sample->>ResampleKernel: launch (cumulative_log_p, USE_BLOCK_VERIFICATION)
ResampleKernel-->>rejection_sample: resampled tokens
rejection_sample-->>RejectionSampler: sampled output tokens
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
vllm/model_executor/kernels/linear/scaled_mm/b12x.py (1)
120-124: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the parsed threshold instead of re-parsing every forward call.
_b12x_dg_prefill_min_tokens()is invoked from_apply_b12x_fp8_block_scaled_linearon every forward call (per layer, per token batch). Doing anos.getenv+int()parse each time is unnecessary overhead in a hot path, and a malformed env value (int(raw)with no try/except) will raise on the very first forward call rather than at start-up.♻️ Proposed fix using functools.lru_cache
+from functools import lru_cache + + +@lru_cache(maxsize=1) def _b12x_dg_prefill_min_tokens() -> int: raw = os.getenv("VLLM_B12X_FP8_LINEAR_DG_PREFILL_MIN_TOKENS") if raw is None: return _B12X_DG_PREFILL_DEFAULT_MIN_TOKENS return int(raw)🤖 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/kernels/linear/scaled_mm/b12x.py` around lines 120 - 124, The _b12x_dg_prefill_min_tokens() helper is parsing the environment variable on every _apply_b12x_fp8_block_scaled_linear forward call, which is unnecessary in this hot path and can fail late on a bad value. Cache the parsed result in _b12x_dg_prefill_min_tokens() (for example with a memoized helper) so the env lookup and int conversion happen only once, and keep the default fallback behavior for missing values.vllm/v1/worker/gpu/spec_decode/dspark/speculator.py (1)
525-545: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: de-duplicate the timing report blocks.
The graph-forward (Line 534, every 200 steps) and eager-forward (Line 579, every 50 steps) paths carry nearly identical accumulate-and-print logic with different intervals and labels. Extracting a small
_report_timing(phase_label, interval)helper would remove the duplication and make the interval discrepancy intentional rather than incidental. Behavior is otherwise correct and fully env-gated.Also applies to: 573-587
🤖 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/dspark/speculator.py` around lines 525 - 545, The timing accumulation and stderr print logic is duplicated in the speculative decode paths, making the graph-forward and eager-forward reports harder to maintain. Extract the shared report/update code from the speculator forward flow into a small helper such as _report_timing in Speculator, and have both the run_fullgraph path and the eager-forward path call it with their respective labels/intervals so the different cadence stays explicit.
🤖 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 `@vllm/model_executor/kernels/linear/scaled_mm/b12x.py`:
- Around line 133-141: The DeepGEMM prefill branch in b12x routes directly to
dg_kernel.apply_weights even though that path uses x.view(-1, x.shape[-1]) and
can fail on non-contiguous tensors. Update the dispatch logic in the b12x
scaled_mm path to add a contiguity check before calling dg_kernel.apply_weights,
and keep non-contiguous inputs on the existing b12x fallback path that already
normalizes them. Use the existing dg_kernel and apply_weights symbols to place
the guard in the prefill-regime branch.
In `@vllm/models/deepseek_v4/nvidia/dspark.py`:
- Line 1082: The condition in dspark.py exceeds the 88-character line limit and
should be wrapped. Reformat the long if statement in the shared/capturing check
so it spans multiple lines while preserving the same logic, using the nearby
conditional in the DeepSeek V4 NVidia code path to locate it.
- Around line 1058-1092: The tensor-debug layer-parts dump path in the dspark KV
staging logic still references cache_window even though that symbol is out of
scope after switching to the persistent all_kv buffer. Update the debug payload
in the code around the all_kv / topk_idxs construction to use the actual window
slice from all_kv (for example the first window_size columns) or remove that
field entirely so _dspark_dump_layer_parts_enabled() cannot raise NameError when
VLLM_DSPARK_TENSOR_DEBUG_LAYER_PARTS is enabled.
---
Nitpick comments:
In `@vllm/model_executor/kernels/linear/scaled_mm/b12x.py`:
- Around line 120-124: The _b12x_dg_prefill_min_tokens() helper is parsing the
environment variable on every _apply_b12x_fp8_block_scaled_linear forward call,
which is unnecessary in this hot path and can fail late on a bad value. Cache
the parsed result in _b12x_dg_prefill_min_tokens() (for example with a memoized
helper) so the env lookup and int conversion happen only once, and keep the
default fallback behavior for missing values.
In `@vllm/v1/worker/gpu/spec_decode/dspark/speculator.py`:
- Around line 525-545: The timing accumulation and stderr print logic is
duplicated in the speculative decode paths, making the graph-forward and
eager-forward reports harder to maintain. Extract the shared report/update code
from the speculator forward flow into a small helper such as _report_timing in
Speculator, and have both the run_fullgraph path and the eager-forward path call
it with their respective labels/intervals so the different cadence stays
explicit.
🪄 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: 0eeb706b-f148-4a1a-a71a-713c30837303
📒 Files selected for processing (9)
vllm/config/speculative.pyvllm/model_executor/kernels/linear/scaled_mm/b12x.pyvllm/model_executor/layers/fused_moe/b12x_moe.pyvllm/model_executor/layers/fused_moe/b12x_tiny_decode.pyvllm/models/deepseek_v4/nvidia/dspark.pyvllm/v1/worker/gpu/sample/gumbel.pyvllm/v1/worker/gpu/spec_decode/dspark/speculator.pyvllm/v1/worker/gpu/spec_decode/rejection_sampler.pyvllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
Comment on lines
+133
to
+141
| dg_kernel = getattr(layer, "b12x_dg_prefill_kernel", None) | ||
| if ( | ||
| dg_kernel is not None | ||
| and x.numel() // x.shape[-1] >= _b12x_dg_prefill_min_tokens() | ||
| ): | ||
| # Prefill-regime token counts run the DeepGEMM path on the | ||
| # dg-processed copy of the weights; decode/small batches stay on the | ||
| # b12x MXFP8 dense GEMM, which wins below the crossover. | ||
| return dg_kernel.apply_weights(layer, x, bias) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and nearby symbols first.
ast-grep outline vllm/model_executor/kernels/linear/scaled_mm/b12x.py --view expanded || true
printf '\n--- file excerpt ---\n'
sed -n '1,380p' vllm/model_executor/kernels/linear/scaled_mm/b12x.py
printf '\n--- search related kernels / apply_weights ---\n'
rg -n "class .*BlockScaledMMLinearKernel|def apply_weights|view\\(-1, x\\.shape\\[-1\\]\\)|torch\\.compiler\\.is_compiling|b12x_dg_prefill_kernel|dg_prefill" vllm/model_executor/kernels/linear -SRepository: local-inference-lab/vllm
Length of output: 20708
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '--- BlockScaledMMLinearKernel.py ---\n'
sed -n '1,220p' vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
printf '\n--- deep_gemm.py ---\n'
sed -n '1,240p' vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
printf '\n--- mxfp8/b12x.py compile-time path ---\n'
sed -n '280,380p' vllm/model_executor/kernels/linear/mxfp8/b12x.py
printf '\n--- any DG-specific x normalization ---\n'
rg -n "contiguous\\(|reshape\\(|view\\(-1, x\\.shape\\[-1\\]\\)|apply_weights\\(layer, x, bias\\)|b12x_dg_prefill_kernel" vllm/model_executor/kernels/linear/scaled_mm -SRepository: local-inference-lab/vllm
Length of output: 253
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- BlockScaledMMLinearKernel.py ---'
sed -n '1,220p' vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
echo
echo '--- deep_gemm.py ---'
sed -n '1,240p' vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
echo
echo '--- mxfp8/b12x.py compile-time path ---'
sed -n '280,380p' vllm/model_executor/kernels/linear/mxfp8/b12x.py
echo
echo '--- DG-specific x normalization search ---'
rg -n "contiguous\\(|reshape\\(|view\\(-1, x\\.shape\\[-1\\]\\)|apply_weights\\(layer, x, bias\\)|b12x_dg_prefill_kernel" vllm/model_executor/kernels/linear/scaled_mm -SRepository: local-inference-lab/vllm
Length of output: 18800
Add a contiguity guard before routing to DeepGEMM prefill.
dg_kernel.apply_weights() ultimately does x.view(-1, x.shape[-1]), so this branch can crash on non-contiguous inputs while the b12x fallback already normalizes them.
Proposed fix
- return dg_kernel.apply_weights(layer, x, bias)
+ return dg_kernel.apply_weights(layer, x.contiguous(), bias)📝 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.
Suggested change
| dg_kernel = getattr(layer, "b12x_dg_prefill_kernel", None) | |
| if ( | |
| dg_kernel is not None | |
| and x.numel() // x.shape[-1] >= _b12x_dg_prefill_min_tokens() | |
| ): | |
| # Prefill-regime token counts run the DeepGEMM path on the | |
| # dg-processed copy of the weights; decode/small batches stay on the | |
| # b12x MXFP8 dense GEMM, which wins below the crossover. | |
| return dg_kernel.apply_weights(layer, x, bias) | |
| dg_kernel = getattr(layer, "b12x_dg_prefill_kernel", None) | |
| if ( | |
| dg_kernel is not None | |
| and x.numel() // x.shape[-1] >= _b12x_dg_prefill_min_tokens() | |
| ): | |
| # Prefill-regime token counts run the DeepGEMM path on the | |
| # dg-processed copy of the weights; decode/small batches stay on the | |
| # b12x MXFP8 dense GEMM, which wins below the crossover. | |
| return dg_kernel.apply_weights(layer, x.contiguous(), bias) |
🤖 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/kernels/linear/scaled_mm/b12x.py` around lines 133 - 141,
The DeepGEMM prefill branch in b12x routes directly to dg_kernel.apply_weights
even though that path uses x.view(-1, x.shape[-1]) and can fail on
non-contiguous tensors. Update the dispatch logic in the b12x scaled_mm path to
add a contiguity check before calling dg_kernel.apply_weights, and keep
non-contiguous inputs on the existing b12x fallback path that already normalizes
them. Use the existing dg_kernel and apply_weights symbols to place the guard in
the prefill-regime branch.
Comment on lines
+1058
to
+1092
| # Persistent [b, window+block] staging buffer: one gather for the | ||
| # window + a small copy for the block KV, replacing the per-layer | ||
| # gather + full-width torch.cat (saves one window-wide copy and two | ||
| # allocations per layer per step). | ||
| if ( | ||
| getattr(self, "_dspark_allkv_buf", None) is None | ||
| or self._dspark_allkv_buf.shape[0] < batch_size | ||
| ): | ||
| self._dspark_allkv_buf = torch.empty( | ||
| batch_size, | ||
| window_size + block_size, | ||
| kv.shape[-1], | ||
| dtype=kv.dtype, | ||
| device=kv.device, | ||
| ) | ||
| all_kv = self._dspark_allkv_buf[:batch_size] | ||
| torch.index_select( | ||
| self.dspark_kv_cache, 0, cache_rows, out=all_kv[:, :window_size] | ||
| ) | ||
| all_kv[:, window_size:].copy_(kv) | ||
| # topk_idxs is layer-invariant: computed once per draft step by the | ||
| # first dspark layer and shared via the shared-state dict. | ||
| shared = _dspark_layer_shared_state | ||
| step_key = (positions.data_ptr(), batch_size, block_size) | ||
| if shared.get("key") != step_key or shared.get("capturing") is not torch.cuda.is_current_stream_capturing(): | ||
| shared["key"] = step_key | ||
| shared["capturing"] = torch.cuda.is_current_stream_capturing() | ||
| shared["topk"] = _build_dspark_topk_idxs( | ||
| window_size=window_size, | ||
| batch_size=batch_size, | ||
| block_size=block_size, | ||
| positions=positions.view(batch_size), | ||
| device=x.device, | ||
| ) | ||
| topk_idxs = shared["topk"] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="vllm/models/deepseek_v4/nvidia/dspark.py"
echo "== context around lines 1035-1125 =="
sed -n '1035,1125p' "$file" | cat -n
echo
echo "== all references to cache_window in file =="
rg -n "\bcache_window\b" "$file"
echo
echo "== debug dump helper references =="
rg -n "_dspark_dump_layer_parts_enabled|tensor_debug|layer_parts" "$file"Repository: local-inference-lab/vllm
Length of output: 7167
cache_window still breaks the tensor-debug dump path.
_dspark_dump_layer_parts_enabled() can still reach "cache_window": cache_window, but that name is no longer in scope. Enabling VLLM_DSPARK_TENSOR_DEBUG_LAYER_PARTS=1 will raise NameError; point it at all_kv[:, :window_size] or remove the field.
🤖 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/dspark.py` around lines 1058 - 1092, The
tensor-debug layer-parts dump path in the dspark KV staging logic still
references cache_window even though that symbol is out of scope after switching
to the persistent all_kv buffer. Update the debug payload in the code around the
all_kv / topk_idxs construction to use the actual window slice from all_kv (for
example the first window_size columns) or remove that field entirely so
_dspark_dump_layer_parts_enabled() cannot raise NameError when
VLLM_DSPARK_TENSOR_DEBUG_LAYER_PARTS is enabled.
| # first dspark layer and shared via the shared-state dict. | ||
| shared = _dspark_layer_shared_state | ||
| step_key = (positions.data_ptr(), batch_size, block_size) | ||
| if shared.get("key") != step_key or shared.get("capturing") is not torch.cuda.is_current_stream_capturing(): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Line exceeds the 88-character limit.
This condition is ~116 chars; wrap it. As per coding guidelines, "Python code must follow an 88-character line length limit."
✂️ Proposed wrap
- if shared.get("key") != step_key or shared.get("capturing") is not torch.cuda.is_current_stream_capturing():
+ capturing = torch.cuda.is_current_stream_capturing()
+ if shared.get("key") != step_key or shared.get("capturing") is not capturing:
shared["key"] = step_key
- shared["capturing"] = torch.cuda.is_current_stream_capturing()
+ shared["capturing"] = capturing📝 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.
Suggested change
| if shared.get("key") != step_key or shared.get("capturing") is not torch.cuda.is_current_stream_capturing(): | |
| capturing = torch.cuda.is_current_stream_capturing() | |
| if shared.get("key") != step_key or shared.get("capturing") is not capturing: | |
| shared["key"] = step_key | |
| shared["capturing"] = capturing |
🤖 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/dspark.py` at line 1082, The condition in
dspark.py exceeds the 88-character line limit and should be wrapped. Reformat
the long if statement in the shared/capturing check so it spans multiple lines
while preserving the same logic, using the nearby conditional in the DeepSeek V4
NVidia code path to locate it.
Source: Coding guidelines
Author
|
Closing this old The intended changes were audited before closing:
No production change from this PR is missing from the new FF submission. Any optimization for the current paged sparse-SWA implementation should be developed and measured separately against |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two commits from the DSpark speedup/consolidation campaign (full context: dspark-upstream-consolidation.md §6–8).
1.
16140bcb— draft-pass cleanup + timing instrumentation (the measured win)Three graph-captured reductions in the DSpark draft forward:
VLLM_DSPARK_CONFIDENCE=1— the serving speculator discards it; it's an lm-head-sized projection computed every draft steptopk_idxssharing across the three draft layers (layer-invariant; built once, shared via a module dict cleared atforward_specentry; capture-safe)[b, window+block]staging buffer replacing the per-layer window gather + full-widthtorch.cat(one window-wide copy and two allocations per layer per step gone)Measured with the new CUDA-event instrumentation (
VLLM_DSPARK_TIMING=1, rolling prepare/graph-forward ms — the only way to time inside FULL cudagraphs; measurement-only: its per-step host sync collapses cc64): draft graph-forward 2.01 → 1.81 ms/step, prepare 0.33 → 0.31 ms, 30k coherence clean (CJK 0).2.
5a72086— upstream rejection-utils consolidation + opt-in block verificationWholesale adoption of upstream's refactored
rejection_sampler_utils.py(post-vllm-project#46781/vllm-project#47093/vllm-project#47383: block-verification kernels, int64-safe offsets) with our extras re-applied on top: fp64 acceptance uniforms underUSE_FP64(upstream is fp32-only) and the active-row-guardedgumbel_block_argmaxsignature. Addstl_rand32to gumbel.py,"block"toRejectionSampleMethod, and a host-side all-greedy skip (the block prep kernels cost ~5–7 % of a step and are inert at temp 0).Honest verdict on block verification itself: synthetic kernel harness (drives
rejection_sample()directly with controlled draft divergence — in rtx6kprocode/tiny-decode/synth_rejection_test.py) shows temp-0 equivalence exact, distribution preservation, and +7–47 % accepted length on Gaussian-noise drafts; but on the real DSpark draft, E2E acceptance is unchanged (12×800-token probes ×3 per config @temp 0.7: block 54.8 % mean vs standard 56.5 %, spread ±3 — DSpark's errors are bimodal, no ratio slack to pool). Left opt-in, not default. The consolidation value stands regardless: one utils file aligned with upstream + our hardening, and the synth harness as a permanent regression tool.Measurements around this branch (DS4-Flash-DSpark TP2 A8, k=5)
cc1 197.3 / cc64 1,925 aggregate / prefill 12,972/12,875/12,067 (full B12X @f416b75). For the record, the paired Lucifer-CUTLASS comparison (cc64 +21 % — the dynamic w4a8 batch-M MoE gap) is in the wiki §8.
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes