Repository navigation
[Perf][DSv4.1] Trim the last KV-source layer in decoder replay - #59894
Conversation
|
@ivanium @Juntian777 Could you please take a look? When enabling decoder replay, we can cut some computation in 20th layer, like Sparse MLA and MoE. |
|
@gcanlin Hi Canlin, thanks for the contribution! We just merged #59532, which should further improve decoder replay performance. Would you mind rebasing this PR on top of the latest main, including #59532, and rerunning the benchmarks with a larger CUDA graph size (e.g. 8192) so we can evaluate the incremental gains on top of it? Also, while rebasing, could you please do a quick cleanup of the PR where applicable, removing any changes that are no longer needed or are now covered by #59532, so that the remaining diff is easier to review? Thanks! |
|
This pull request has merge conflicts that must be resolved before it can be |
The decoder replay layers run on each prefill's last window rows; the last KV-source layer before them (the cut layer) still ran on every row, yet for the trimmed rows only its KV is read: the compressed (ratio-1) KV and index K by the layers after it and later steps, the SWA KV by this step's replay rows. Its indexer scoring and top-k, sparse attention, wo and MoE feed only the replay layers. The cut layer now first writes its KV for every row (write_kv: its attention input mHC and Attention.forward_kv, in the batch's graph), and then runs as the first replay layer, in the replay batch and its CUDA graphs, rewriting the replay rows' KV with the same values. Its sliding-window metadata keeps the batch's window starts, since every row's window KV is there, so it stays exact. Its indexer now writes the replay rows' top-k and candidate blocks in place, so the row-buffer compaction goes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
526f37f to
60aaace
Compare
prepare_inputs_to_capture handed CUDAGraphMode.NONE to ModelState.prepare_attn for every capture, while PIECEWISE steps at runtime hand PIECEWISE. A model state that prepares what the graph runs by its mode then captured something else: DeepSeek-V4.1 builds the decoder replay batch for PIECEWISE graphs at the trim threshold or above (vllm-project#59532), so its graphs were captured without the replay break and ran the replay layers on the whole batch, though every step still prepared a replay batch. PIECEWISE captures now pass PIECEWISE. The other model states branch on FULL only, so they build the same metadata as before. DSV4.1-Flash DEP4 on B200, CUDA graphs up to 8192 tokens: 8k-token prefills 142.7k -> 236.9k total tok/s, 32k in / 256 out 772 -> 1,281 output tok/s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
| attn_metadata = model_state.prepare_attn( | ||
| input_batch, | ||
| CUDAGraphMode.NONE, | ||
| CUDAGraphMode.PIECEWISE if piecewise_cudagraph else CUDAGraphMode.NONE, |
There was a problem hiding this comment.
IIUC, we need to add this PIECEWISE parameter to enable building replay batch, otherwise the path can't go into graph for replay decoder. Can you double-check it? I just let agent to find the inferenceX CI config and code. It's different from the merged version(#59532). I failed enabling graph for replay decoder without this fix. @Juntian777
def _prepare_replay_batch(self, input_batch, cudagraph_mode, ..., capture_desc=None):
....
force_replay_batch = capture_desc is not None or (
cudagraph_mode == CUDAGraphMode.PIECEWISE
and self._trim_threshold is not None
and input_batch.num_tokens_after_padding >= self._trim_threshold
)
There was a problem hiding this comment.
Could we go back to deciding the break by graph size, to avoid changing vllm/v1/worker/gpu/cudagraph_utils.py? In DecoderReplayLayers.__call__ (vllm/models/deepseek_v41/decoder_replay_layers.py), replace the self.replay_batch is not None check with something like:
if in_piecewise_cudagraph():
desc = get_forward_context().batch_descriptor
assert desc is not None
if (
self.trim_threshold is not None
and desc.num_tokens >= self.trim_threshold
):
return self._run_in_graph_break(hidden_states, *states)
return self._run(hidden_states, *states)| ) | ||
| late_aux: list[torch.Tensor] = [] | ||
| if self.decoder_replay_layers is not None: | ||
| if self.decoder_replay_layers.replay_batch is not None: |
There was a problem hiding this comment.
Same as above, we should also avoid using self.decoder_replay_layers.replay_batch as the condition here. forward gates write_kv on replay_batch is not None, which is None at capture but set at runtime, same as the break.
There was a problem hiding this comment.
Fixed! Please take a look.
A PIECEWISE model graph fixes at capture whether it breaks out to the decoder replay batch, and with it whether the first replay layer writes every row's KV first. Both were decided by whether a replay batch was set, which the model state does only at runtime: captures prepare attention as eager, so the graphs were captured without the break and ran the replay layers on the whole batch. Decide both by DecoderReplayLayers.replays(): in a PIECEWISE graph, whether its token count reaches the trim threshold, as the model state's runtime replay batch does (the threshold now lives on the layers for both); elsewhere, whether a replay batch is set. This replaces the runner change that handed captures their graph mode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
| and self._trim_threshold is not None | ||
| and input_batch.num_tokens_after_padding >= self._trim_threshold | ||
| and layers.trim_threshold is not None | ||
| and input_batch.num_tokens_after_padding >= layers.trim_threshold |
There was a problem hiding this comment.
Could we share the PIECEWISE graph-size predicate between replays() and _prepare_replay_batch()? Capture and runtime should use the same policy, but currently they implement it separately.
There was a problem hiding this comment.
Extract them to def graph_replays(self, ....) now.
forward_kv projected Q through wq_b for every row only because the fused SWA insert takes a query along; nothing reads it there. Hand it a placeholder, as DSpark's context KV insert does. An 8k prefill step drops a 254 us GEMM on its whole batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
DecoderReplayLayers.graph_replays(num_tokens) now decides whether a PIECEWISE model graph breaks out to the replay batch, for both replays() (fixed at capture) and the model state's runtime replay batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
[DSv4.1] Simplify the cut-layer replay and pin its PIECEWISE policy Signed-off-by: Canlin <canlinguosdu@gmail.com>
3844ac9 to
ffd52d5
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #93763 for commit |
vllm-project#59894 writes that layer's KV for every token and replays from there. ROCm forward_kv uses the fused uint8 insert. Signed-off-by: Frida Andersson <fanderss@amd.com>
vllm-project#59894 writes that layer's KV for every token and replays from there. ROCm forward_kv uses the fused uint8 insert. Signed-off-by: Frida Andersson <fanderss@amd.com>
…project#59894) Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Juntian Liu <Juntianl777@gmail.com> (cherry picked from commit ad73a47)
…project#59894) Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Juntian Liu <Juntianl777@gmail.com>
Merged with `git merge --no-commit --no-ff` per MERGING.md from the tag pre-upstream-merge-2026-10-28 (d186259dc9); merge base bfa7e5b. 57 files conflicted, 152 files changed on both sides since the base. Resolutions by area: Build system - Dockerfile: upstream's package sites re-gated on BUILD_OS; VLLM_SKIP_FLASH_ATTN_BUILD, BUILD_QUTLASS, BUILD_FLASHKDA and BUILD_DEEPEP flags kept; patched NCCL stage and soname assertions kept; KV-connector install keeps LMCache and restores OpenTelemetry >=1.40 pins; runtime verification stage imports the DSV4 kernel registry, consumer_int8_kernels and GenAI trace handler. - Dockerfile.python-overlay and test_docker_build_metadata.py: removed vllm/v1/executor/ray_executor.py (deleted upstream in vllm-project#60760). - CMakeLists.txt: kept 12.1 in CUDA_SUPPORTED_ARCHS, the arch-gated add_custom_target stubs (_deepgemm_C, _fmha_sm100_C, _flashmla_C, _flashmla_extension_C, _flashmla_dense_C, _flashkda_C, _qutlass_C, _tml_fa4_C, _deepselect_C), FA3 arch gating, and updated AllSpark sources from csrc/libtorch_stable/quantization/gptq_allspark/* to csrc/quantization/gptq_allspark/* following vllm-project#60721. - pyproject.toml, requirements/{common,cuda}.txt: kept setuptools-scm v[0-9]* tag regex, optional model_hosting_container_standards dependency, and OpenTelemetry / vllm-flash-attn pins. - docker/versions.json regenerated with tools/generate_versions_json.py. Native kernels - Upstream vllm-project#60721 (d79b52b) consolidated csrc/libtorch_stable/* back into csrc/*. Re-homed the fork's four declarations (top_k_per_row_prefill_candidates, top_k_per_row_merge_candidates, allspark_repack_weight, allspark_w8a16_gemm) into csrc/ops.h, five operator bindings (including indexer_k_quant_and_cache_int8) into csrc/torch_bindings.cpp, and moved csrc/libtorch_stable/quantization/gptq_allspark/* to csrc/quantization/gptq_allspark/* (updating build_cutlass_arch_probe.py). - csrc/cache_kernels.cu, csrc/attention/mla/sm100_cutlass_mla_kernel.cu, vllm/_custom_ops.py: unioned upstream's head/tile updates and new custom ops with the fork's int8 indexer cache writer, fp8_ds_mla gather, streaming prefill top-k, SM120 gates, and AllSpark wrappers. KV cache layout - Upstream vllm-project#61183 (df6c7ad) split KV cache planning out of vllm/v1/core/kv_cache_utils.py into vllm/v1/core/kv_cache_planning.py. Re-ported the empty projected-group guards (if not group.layer_names: continue) in get_kv_cache_config_from_groups, _max_memory_usage_bytes_from_groups and _get_hisparse_kv_cache_bytes_per_block, is_eagle_group=group.is_eagle_group and bool(worker_layer_names), worker KV sizing error logging, and drop=("extra_retained_tokens", "bounded_replay") in _promote_local_kv_cache_specs into kv_cache_planning.py. - Upstream vllm-project#52670 moved is_quantized_kv_cache / get_kv_quant_mode from vllm/utils/torch_utils.py into vllm/v1/kv_cache_interface.py. Added KVQuantMode.INT8_DS_MLA = 15 and mapped "int8_ds_mla" in get_kv_quant_mode so is_quantized_kv_cache("int8_ds_mla") remains True; kept SlidingWindowMLASpec.merge forwarding kv_quant_mode and alignment; kept "int8_ds_mla" in CacheDType. DeepSeek V4 model and attention - Upstream vllm-project#59894 deleted vllm/transformers_utils/configs/deepseek_v4.py. Removed it and updated resolve_layer_compress_ratio in vllm/models/deepseek_v4/attention.py and _dsv4_compress_ratios in vllm/models/deepseek_v4/common/ops/cache_utils.py to use get_compress_ratios from vllm/transformers_utils/deepseek_v4_utils.py (with getattr guard on num_nextn_predict_layers). - attention.py: kept the opaque torch.ops.vllm.deepseek_v4_attention custom op (with optional output_scale and QuantAct dtype preservation), _run_indexer_op, checkpoint kernel-registry attention class resolution (nvidia_imma / nvidia_sm12x), state_content_bytes=528 and alignment=528 for int8_ds_mla, and _SPARSE_MLA_WARMUP_BACKENDS including FLASHINFER_SPARSE_MLA and SPARKINFER_SPARSE_MLA. - compressor.py, common/ops/{cache_utils,fused_compress_quant_cache, fused_indexer_q_rope_quant,o_proj}.py, nvidia/model.py, amd/mono/stages/moe.py: kept the 528-byte int8_ds_mla compressor and cache writers, C128 CUDA-graph capture skip guard, kv_proj_with_mhc / fused_q_kv_rmsnorm / wo_a_Linear kernel-registry hooks, 4-tensor mHC PP stream boundaries, aux_hidden_{j} PP relay (reserve_aux_intermediate_tensor_slots pinned to 0), and hoisted stage_words out of Smem's field annotation in amd/mono/stages/moe.py for pyflakes static name resolution. - v1/attention/backends/mla/{indexer,sparse_mla,sparse_swa}.py, sparse_attn_indexer.py: kept INT8 indexer K cache (int8_scale_bytes = 4), TP row-sharded streaming/exact prefill top-k, SM120 FlashInfer/CuteDSL gates, int8_ds_mla 528-byte SWA spec, and DSpark/MTP decode threshold handling. Pipeline parallel and executor - Upstream vllm-project#60760 removed v1 vllm/v1/executor/ray_executor.py; deleted it and kept VLLM_RAY_WORKER_IP_ORDER placement-group sorting and bounded Nsight capture in ray_utils.py and vllm/v1/worker/gpu/executor/ray_executor.py. - vllm/v1/worker/gpu/{model_runner,cudagraph_utils}.py: kept the single PP payload broadcast, replayed_pp_anchor_req_ids physical-anchor replay, model_Step profiler ranges, model_uses_input_ids preservation across intermediate PP stages and CUDA-graph captures, and on-demand KVBlockZeroer. - vllm/platforms/cuda.py: kept CudaPlatformBase.is_arch_support_pdl uncached at the classmethod wrapper so Dynamo fullgraph tracing of Triton kernel launchers does not trip on functools._lru_cache_wrapper (31d143c); kept GB10 memory fallback and SPARKINFER_SPARSE_MLA backend priority. - vllm/config/{vllm,load}.py, default_loader.py: kept DSV4 exclusion from breakable-cudagraph defaults, checkpoint kernel-block application, flashpack load format, and rank-local PP weight filtering; upstream's allow_patterns_overrides guard in default_loader.py supersedes our allow_partial argument at the call site. Speculative decoding and scheduler - dflash/speculator.py, dspark/{speculator,utils}.py, eagle/utils.py, rejection_sampler.py: kept _store_hidden_states before _precompute_context_kv on chunked-prefill steps, DSpark PP support, PP embedding-sharing guards (spec_decode_needs_target_embed), and _expand_for_deferred in RejectionSampler. - vllm/v1/core/sched/scheduler.py: kept replayed_pp_anchor_req_ids, defer_block_free, was_deferred_verify, and speculative draft-token budget accounting alongside upstream's update_from_output refactoring. Warmup, entrypoints, quantization, and tests - kernel_warmup.py, deep_gemm_warmup.py, v1/worker/gpu/warmup.py: kept DSV4 mHC, vision, Marlin, gather and top-k warmups, PPHandler disable/enable around graph warmup, and omitted deepseek_v4_sparse_mla_attention_warmup. - entry.py, envs.py: kept startup-probe stages and fork env compile factors. - activation.py, quantization/__init__.py, fused_moe/oracle/nvfp4.py, models/{qwen3_next,utils}.py: kept dsv4_int, dsv41_gguf, dsv4_mxfp4_int8, qwen_int8, SM86 native clamped SiLU, NVFP4 per-expert scale check, and spec_decode_needs_target_embed. - Tests: unioned all conflicted test files and updated fixtures for tokens_per_state, kv_cache_planning imports, get_fused_shared_expert_layers, and CompressorMetadata.tail_slot_mapping. Signed-off-by: doctorpangloss <2229300+doctorpangloss@users.noreply.github.com>
Purpose
In decoder replay steps, the last KV-source layer (layer 20 on DSV4.1-Flash) still ran on every row, though for the trimmed rows only its KV (SWA, compressed, index K) is read. Its indexer scoring and top-k, sparse MLA, wo and MoE feed only the replay layers.
DeepseekV4DecoderLayer.write_kv→DeepseekV4Attention.forward_kv), then runs as the first replay layer (decoder_replay_start = cut) on the replay rows, in the replay batch and its CUDA graphs. Its sliding-window metadata keeps the batch's window starts, so it stays exact. The indexer row-buffer compaction goes.write_kv) by graph size (DecoderReplayLayers.replays()) instead ofreplay_batch is not None, which is unset at capture: before, the break was never captured and PIECEWISE steps did not trim.Results
DSV4.1-Flash DEP4 on B200, CUDA graphs up to 8192 tokens,
vllm bench serverandom. Baseline: #59532 with PIECEWISE trimming working, layer 20 not cut.GSM8K 5-shot: 0.9606 both.
AgentX C128 (InferenceX DEP4 recipe, synthetic AL 3.51, 3,600 s), nightly 43b4aae + this PR vs. the InferenceX CI run (37070984585, dev image ac9126):
Prefix cache hit 95.2% in both.
🤖 Generated with Claude Code