perf(pytorch): reduce speculative decoding pre/post-processing overhead

[grimoire]

Motivation

GLM-5.2 MTP profiling showed that work outside model forward was a material
part of each speculative decoding step. In particular, build_step_context,
fill_buffers_cudagraph, rejection sampling, and stopping criteria launched
many small CUDA kernels.

This PR reduces that overhead without changing model-forward implementations,
adding CUDA graphs, or introducing additional streams.

Changes

• Fuse common CUDA graph input-buffer updates into one Triton kernel.
  • Preserve randomized padded input_ids for MoE load balancing.
  • Fill token IDs, position IDs, block offsets, sequence lengths, and
    cumulative lengths together.
• Move DSA indexer metadata preparation under the CUDA DSA backend.
  • Alias kv_seqlens for single-token graph replay.
  • Fuse causal KV-length and expanded block-table preparation for multi-token
    graph and eager decode.
  • Keep prefill and non-CUDA paths on the portable Torch implementation.
• Reduce regular-layout StepContext metadata kernels.
• Dispatch speculative rejection sampling through the selected backend.
  • Remove the CUDA-to-host .item() synchronization.
  • Add an optimized all-greedy path while retaining mixed and random sampling.
• Reuse rejection results in speculative stopping criteria instead of scanning
  placeholder tokens repeatedly.
• Skip stop-word work when no effective stop words are configured.

The change preserves the existing CUDA graph count, DP ordering,
sleep/wakeup behavior, and weight-loading lifecycle.

Performance

Measured with a four-H200, TP=4, batch-size-128, MTP-5 reduced GLM-5.2 model
using FP8 KV cache:

• Combined target-plus-five-draft graph-buffer fill p50:
  2.019 ms -> 1.102 ms.
• GPU activities in those fills:
  118 -> 8.
• Stable unprofiled request time improved by approximately 0.49%.
• All-greedy rejection-sampling p50:
  1.286 ms -> 0.705 ms.
• EOS-enabled stopping-criteria p50:
  0.505 ms -> 0.371 ms.

The direct range and launch reductions are repeatable. End-to-end improvements
are smaller because model forward remains dominant.

[Copilot]

Pull request overview

This PR targets PyTorch speculative decoding performance by reducing non-forward overhead (CUDA graph buffer fills, DSA/NSA metadata prep, rejection sampling, and stopping criteria), primarily via kernel fusion and backend routing while keeping model-forward behavior unchanged.

Changes:

• Introduces fused Triton kernels to batch-update CUDA graph input buffers and DSA indexer metadata, reducing kernel launch counts.
• Refactors speculative rejection sampling into a backend-dispatched implementation (CUDA + portable Torch), including an all-greedy fast path and host-side greedy detection (has_greedy) to avoid syncs.
• Updates speculative stopping criteria to reuse rejection-sampling results (num_rejected_tokens) instead of rescanning placeholder tokens; expands/updates unit tests accordingly.

Reviewed changes

Copilot reviewed 25 out of 25 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/pytorch/spec_decode/test_strategies.py	Extends sampling strategy tests to cover host-side greedy summary (has_greedy).
tests/pytorch/spec_decode/test_spec_agent.py	Adds rejection-sampling delegation test; updates sampling-input selection tests.
tests/pytorch/spec_decode/test_reject_sample.py	Reworks rejection-sampling tests to use backend RejectionSampler instead of direct kernels/helpers.
tests/pytorch/spec_decode/test_guided_spec_decode.py	Removes slicing-related tests now that slicing helper is dropped; updates docstring/comments.
tests/pytorch/engine/test_model_agent.py	Updates spec-agent build expectations around backend config and proposer/follower behavior.
lmdeploy/pytorch/strategies/ar/sampling.py	Computes and stores has_greedy summary in SamplingInputs.
lmdeploy/pytorch/strategies/ar_spec/model_agent.py	Updates stopping criteria to use num_rejected_tokens and reduces metadata work.
lmdeploy/pytorch/spec_decode/spec_agent.py	Removes _slice_sampling_inputs; routes decode rejection sampling through RejectionSampler.
lmdeploy/pytorch/spec_decode/reject_sampler.py	New backend-dispatched rejection-sampling wrapper + bonus-policy selection helper.
lmdeploy/pytorch/spec_decode/base.py	Removes base-class initialization of rejection sampler (now owned by concrete agent).
lmdeploy/pytorch/models/utils/cudagraph.py	Replaces multiple buffer updates with a single fused kernel call for common inputs.
lmdeploy/pytorch/model_inputs.py	Adds _get_q_start_loc fast path for uniform layouts; gates window-size adjustment.
lmdeploy/pytorch/kernels/cuda/step_metadata/fill_graph_common_inputs.py	Adds fused Triton kernel to fill common CUDA graph buffers in one launch.
lmdeploy/pytorch/kernels/cuda/step_metadata/fill_dsa_indexer_metadata.py	Adds Triton kernel to build DSA indexer KV-lengths and optional expanded block tables.
lmdeploy/pytorch/kernels/cuda/step_metadata/init.py	Introduces package init for new CUDA step-metadata kernels.
lmdeploy/pytorch/kernels/cuda/rejection_sampling.py	Adds Triton kernels for greedy/random rejection sampling and recovered-token sampling.
lmdeploy/pytorch/engine/logits_process.py	Extends SamplingInputs with has_greedy host-side routing hint.
lmdeploy/pytorch/backends/rejection_sampling.py	Adds backend interface + shared placeholder constant for rejection sampling.
lmdeploy/pytorch/backends/nsa.py	Allows caller-supplied indexer_kv_seqlens to avoid redundant computation.
lmdeploy/pytorch/backends/default/rejection_sampling.py	Implements portable Torch rejection sampling backend.
lmdeploy/pytorch/backends/default/op_backend.py	Registers OpType.RejectionSampling for the default backend.
lmdeploy/pytorch/backends/cuda/rejection_sampling.py	Implements CUDA rejection sampling backend using the new Triton kernels.
lmdeploy/pytorch/backends/cuda/op_backend.py	Registers OpType.RejectionSampling for the CUDA backend.
lmdeploy/pytorch/backends/cuda/nsa.py	Moves DSA indexer metadata prep into CUDA backend and enables graph-buffer aliasing.
lmdeploy/pytorch/backends/base.py	Adds OpType.RejectionSampling enum entry.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.