perf(cuda): use PDL for paged attention and V4 prefill

[grimoire]

Summary

Use CUDA Programmatic Dependent Launch (PDL) to reduce launch gaps in two adjacent Triton kernel pipelines:

• split-K paged-attention producer → reducer during decode;
• DeepSeek-V4 prefill score_kv → fill_compress_state.

This also raises the minimum supported Triton version from 3.0 to 3.4 and removes obsolete compatibility branches.

Changes

Split-K paged attention

• Signal dependents before the producer’s final scratch stores.
• Launch the reducer with launch_pdl=True.
• Wait immediately before the reducer reads producer scratch.
• Cover standard, quantized, and TurboQuant reducer paths.
• Enable on NVIDIA SM90+.

DeepSeek-V4 prefill compressor

• Add a wrapper that owns the adjacent score and state-fill launches.
• Signal after the score kernel’s final read of the old ring-buffer state.
• Let the fill kernel preload independent inputs before waiting.
• Wait immediately before overwriting compressor state.
• Automatically enable PDL only for:
  • NVIDIA SM90+;
  • ratio-4 overlap prefill;
  • at least 4096 total prefill tokens.
• Keep ratio-128 on the sequential path.
• Leave decode unchanged because score and state fill are already fused there.

Triton cleanup

• Require Triton >=3.4.0.
• Remove Triton 2.x/early-3.x compatibility branches.
• Use device-specific compute capability queries.
• Simplify TMA descriptor and Triton math helper selection.

Performance

H200 split-K paged-attention CUDA graph latency:

Shape	Before	PDL	Change
Batch 1, context 1K	12.096 µs	11.520 µs	-4.8%
Batch 1, context 8K	22.080 µs	21.248 µs	-3.8%
Batch 8, context 1K	21.376 µs	20.704 µs	-3.1%
Batch 32, context 8K	260.384 µs	259.296 µs	-0.4%

All nine measured CUDA graph cases improved. The eager batch-1/context-1K case regressed from 14.464 µs to 16.352 µs; paged attention is used for decode, which uses CUDA graphs by default. Eager mode remains opt-in.

For V4 ratio-4 prefill at 4K–8K tokens, eager p50 improved by approximately 1.3%–3.2%. CUDA graph results were neutral to slightly better. Smaller prefills and ratio-128 do not automatically enable PDL.

[Copilot]

Pull request overview

This PR introduces CUDA Programmatic Dependent Launch (PDL) wiring to reduce launch gaps between adjacent Triton kernel pipelines in decode (split‑K paged attention) and DeepSeek‑V4 prefill (score → state fill), while also raising the minimum supported Triton version to simplify compatibility code paths.

Changes:

• Add PDL producer/consumer signaling + waits to split‑K paged attention (producer signals, reducer launches with launch_pdl=True and waits before scratch reads) on SM90+.
• Add a new prefill wrapper (score_and_fill_state_prefill) to optionally overlap score_kv with fill_compress_state, plus a CUDA SM90+/token-count gating heuristic and a correctness test.
• Bump minimum Triton to >=3.4.0 and remove legacy version branches / imports across multiple CUDA kernels and env checks.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/pytorch/kernel/test_v4_compressor.py	Adds a prefill test validating PDL vs sequential behavior for the V4 compressor.
requirements/runtime_cuda.txt	Raises Triton minimum version to >=3.4.0 for CUDA runtime installs.
lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py	Removes Triton version branching and uses a single rounding implementation.
lmdeploy/pytorch/kernels/cuda/v4_compressor.py	Implements PDL signaling/waiting between score_kv and fill_compress_state, and adds the new prefill wrapper + heuristic.
lmdeploy/pytorch/kernels/cuda/utils.py	Simplifies capability/TMA support detection and exposes compute_capability via get_device_props.
lmdeploy/pytorch/kernels/cuda/pagedattention.py	Adds PDL producer/consumer synchronization to split‑K paged attention kernels (SM90+).
lmdeploy/pytorch/kernels/cuda/flashattention.py	Removes version/capability globals and switches to device-props-based capability selection.
lmdeploy/pytorch/kernels/cuda/blocked_gemm_fp8.py	Updates TMA gating to use the new supports_tma(device) API and imports TensorDescriptor directly.
lmdeploy/pytorch/kernels/cuda/bitonic_topk.py	Cleans up Triton constexpr helper handling and renames _concate→_concat.
lmdeploy/pytorch/kernels/cuda/activation.py	Removes Triton version branching and standardizes fast_expf selection.
lmdeploy/pytorch/check_env/triton.py	Updates minimum Triton version check to 3.4.0.
lmdeploy/pytorch/backends/cuda/v4_compressor.py	Switches prefill path to use the new combined score_and_fill_state_prefill wrapper.

---

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

[grimoire]

This is intentional. The compressor dispatch defines decode as kv.size(0) == state_ids.size(0) and routes that case to the fused score_and_fill_state_decode kernel.
score_and_fill_state_prefill is called only from the opposite branch, so this assertion documents the wrapper’s internal call contract. A batch containing one token per request does not reach this wrapper.