Skip to content

[Perf][DSv4.1] Trim the last KV-source layer in decoder replay - #59894

Merged
ivanium merged 8 commits into
vllm-project:mainfrom
gcanlin:perf/dsv41-cut-layer-trim
Oct 9, 2026
Merged

ivanium merged 8 commits into
vllm-project:mainfrom
gcanlin:perf/dsv41-cut-layer-trim

Conversation

@gcanlin

@gcanlin gcanlin commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

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.

  • The cut layer now writes its KV for every row (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.
  • PIECEWISE graphs now decide the replay break (and write_kv) by graph size (DecoderReplayLayers.replays()) instead of replay_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 serve random. Baseline: #59532 with PIECEWISE trimming working, layer 20 not cut.

baseline this PR
8k in / 1 out, C16, total tok/s 236,952 245,556 +3.6%
32k in / 1 out, C16 235,791 245,996 +4.3%
128k in / 1 out, C4 194,012 204,764 +5.5%
32k in / 256 out, C32, output tok/s 1,281 1,303 +1.7%

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):

CI this PR
output tok/s 5,446 5,575 +2.4%
TTFT mean / p50 / p99 (ms) 293 / 163 / 2,558 255 / 155 / 2,004
ITL mean / p50 / p99 (ms) 8.66 / 7.60 / 29.3 8.21 / 7.36 / 24.0
per-user output tok/s 132.0 136.6 +3.5%
requests completed 18,804 19,572 +4.1%

Prefix cache hit 95.2% in both.

🤖 Generated with Claude Code

@mergify mergify Bot added deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models labels Oct 3, 2026
@gcanlin
gcanlin marked this pull request as ready for review October 3, 2026 18:36

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@gcanlin

gcanlin commented Oct 3, 2026

Copy link
Copy Markdown
Contributor Author

@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.

@Juntian777

Juntian777 commented Oct 6, 2026 •

Copy link
Copy Markdown
Contributor

@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!

@mergify

mergify Bot commented Oct 6, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gcanlin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Oct 6, 2026
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>
@gcanlin
gcanlin force-pushed the perf/dsv41-cut-layer-trim branch from 526f37f to 60aaace Compare October 6, 2026 18:36
@mergify mergify Bot removed the needs-rebase label Oct 6, 2026
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>
@gcanlin gcanlin changed the title [Perf][DSV4.1] Decoder replay: trim the cut layer's query side and FFN [Perf][DSV4.1] Decoder replay: replay the last KV source; fix PIECEWISE capture mode Oct 7, 2026
@mergify mergify Bot added nvidia mrv2 Model Runner V2 specific labels Oct 7, 2026
@gcanlin gcanlin changed the title [Perf][DSV4.1] Decoder replay: replay the last KV source; fix PIECEWISE capture mode [Perf][DSV4.1] Decoder replay: replay the last KV source with pig Oct 7, 2026
@gcanlin gcanlin changed the title [Perf][DSV4.1] Decoder replay: replay the last KV source with pig [Perf][DSV4.1] Decoder replay: replay the last KV source with PCG Oct 7, 2026
Comment thread vllm/v1/worker/gpu/cudagraph_utils.py Outdated
attn_metadata = model_state.prepare_attn(
input_batch,
CUDAGraphMode.NONE,
CUDAGraphMode.PIECEWISE if piecewise_cudagraph else CUDAGraphMode.NONE,

@gcanlin gcanlin Oct 7, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
    )

@Juntian777 Juntian777 Oct 8, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

@Juntian777 Juntian777 Oct 8, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@mergify

mergify Bot commented Oct 8, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gcanlin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Canlin <canlinguosdu@gmail.com>
@gcanlin gcanlin changed the title [Perf][DSV4.1] Decoder replay: replay the last KV source with PCG [Perf][DSV4.1] Decoder replay: replay the last KV source, break out by graph size Oct 8, 2026
@mergify mergify Bot removed the needs-rebase label Oct 8, 2026
@gcanlin gcanlin changed the title [Perf][DSV4.1] Decoder replay: replay the last KV source, break out by graph size [Perf][DSv4.1] Decoder replay: trim the last KV-source layer, fix PIECEWISE trimming Oct 8, 2026
@gcanlin gcanlin changed the title [Perf][DSv4.1] Decoder replay: trim the last KV-source layer, fix PIECEWISE trimming [Perf][DSv4.1] Trim the last KV-source layer in decoder replay Oct 8, 2026
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extract them to def graph_replays(self, ....) now.

gcanlin and others added 2 commits October 8, 2026 14:33
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>
@gcanlin
gcanlin force-pushed the perf/dsv41-cut-layer-trim branch from 3844ac9 to ffd52d5 Compare October 9, 2026 03:04
@ivanium
ivanium enabled auto-merge (squash) October 9, 2026 04:07
@github-actions github-actions Bot added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 9, 2026
@WoosukKwon

Copy link
Copy Markdown
Collaborator

/ci run

@github-actions

github-actions Bot commented Oct 9, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #93763 for commit afea5427a490.

@ivanium
ivanium merged commit ad73a47 into vllm-project:main Oct 9, 2026
94 checks passed
frida-andersson added a commit to frida-andersson/vllm that referenced this pull request Oct 9, 2026
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>
frida-andersson added a commit to frida-andersson/vllm that referenced this pull request Oct 12, 2026
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>
wangyicong52 pushed a commit to bytedance-iaas/vllm that referenced this pull request Oct 15, 2026
…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)
michael-gpu pushed a commit to practical-inference/vllm-ampere-extended that referenced this pull request Oct 20, 2026
…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>
doctorpangloss added a commit to AppMana/forks-vllm-consumer-nvidia-platforms that referenced this pull request Oct 29, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models mrv2 Model Runner V2 specific nvidia ready ONLY add when PR is ready to merge/full CI is needed

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

4 participants