[DSv4 Perf] Adaptive topk width for dsv4, making #50004 back

[yewentao256]

Purpose

#50004 was reverted by #51318

This PR get the similar performance gain but doesn't hurt the accuracy

Test

Perf

E2E perf see #50004

Kernel level perf can be seen in this AI generated script

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Benchmark adaptive-width DeepSeek-V4 C128A prefill metadata."""

import torch

from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata
from vllm.triton_utils import triton

NUM_TOKENS = [512, 2048, 8192, 32768]
CONTEXT_LEN = 32768
COMPRESS_RATIO = 128
CAPACITY_WIDTH = 8192
ACTIVE_WIDTH = max(triton.next_power_of_2(CONTEXT_LEN // COMPRESS_RATIO), 128)


def benchmark(num_tokens: int, width: int) -> float:
    device = torch.device("cuda")
    positions = torch.arange(
        CONTEXT_LEN - num_tokens,
        CONTEXT_LEN,
        dtype=torch.int64,
        device=device,
    )
    global_decode_buffer = torch.empty(
        (1, CAPACITY_WIDTH), dtype=torch.int32, device=device
    )
    prefill_buffer = torch.empty(
        (num_tokens, CAPACITY_WIDTH), dtype=torch.int32, device=device
    )
    decode_lens_buffer = torch.empty(1, dtype=torch.int32, device=device)
    token_to_req_indices = torch.zeros(num_tokens, dtype=torch.int32, device=device)
    slot_mapping = torch.zeros(num_tokens, dtype=torch.int64, device=device)
    block_table = torch.zeros((1, 1), dtype=torch.int32, device=device)

    def run() -> None:
        build_c128a_topk_metadata(
            positions,
            COMPRESS_RATIO,
            0,
            token_to_req_indices,
            block_table,
            2,
            slot_mapping,
            global_decode_buffer,
            decode_lens_buffer,
            prefill_buffer,
            max_compressed_tokens=width,
        )

    run()
    return triton.testing.do_bench(run, warmup=100, rep=500, return_mode="median")


print(f"context={CONTEXT_LEN}, capacity={CAPACITY_WIDTH}, active={ACTIVE_WIDTH}")
print(f"{'tokens':>8} {'fixed us':>10} {'adaptive us':>12} {'speedup':>8}")
for num_tokens in NUM_TOKENS:
    fixed_ms = benchmark(num_tokens, CAPACITY_WIDTH)
    adaptive_ms = benchmark(num_tokens, ACTIVE_WIDTH)
    print(
        f"{num_tokens:>8} {fixed_ms * 1000:>10.2f} "
        f"{adaptive_ms * 1000:>12.2f} {fixed_ms / adaptive_ms:>7.2f}x"
    )

And we can get

context=32768, capacity=8192, active=256
  tokens   fixed us  adaptive us  speedup
     512       9.22         7.20    1.28x
    2048      15.39         7.20    2.14x
    8192      44.00        11.30    3.90x
   32768     150.56        25.79    5.84x

Acc

GPT 5.6 sol helps doing the acc test as #52448 mentioned, results below:

• Hardware: 2 x NVIDIA B200 allocated by chg (physical GPUs 5 and 6).
• Model: deepseek-ai/DeepSeek-V4-Flash-0731, exact snapshot
  7872f01b1d1fe23eabc4c98b48bffcef5a386062.
• Trigger configuration: TP2, expert parallel, DSpark with 7 speculative
  tokens, greedy draft, prefix caching, breakable CUDA graphs, and
  FULL_AND_PIECEWISE graph mode.
• Load: concurrency 32, 8 warm-up waves followed by 8 measured waves, 512
  requests total. Each request used the issue payload, temperature=0, and
  max_tokens=16384.
• Issue hit definition: finish_reason == "length", empty content, and no
  tool calls.
• Warm-up: 0/256 hits, 0 errors, 256/256 tool_calls finishes.
• Measured after the first 256 requests: 0/256 hits, 0 errors, 256/256
  tool_calls finishes.
• Total: 0/512 hits and 0/512 transport/server errors.
• Server error scan: zero CUDA, illegal-memory, assertion, graph-error, fatal,
  or traceback matches.

All 512 responses had non-empty content and tool calls using declared tool
names with JSON-decodable arguments. Tool-call count was 1 for 75 responses,
4 for 27 responses, and 6 for 410 responses. This validates absence of the
specific length-cap/empty-turn failure; it is not a full semantic task-accuracy
evaluation.

The prompt was 6,543 tokens and the largest prompt-plus-completion was 9,037
tokens. Therefore this run exercised logical C128A width 128 while retaining
the capacity row stride 640, rather than falling back to the full-width path.

repro-52448-adaptive-stride-20260818.zip

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84463 for commit c2450d668857.

[yewentao256]

E2E lm eval added

export PYTHONPATH=/home/yewentao256/vllm-source
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1

MODEL=/data/LucasWilkinson/hub_cache/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01b1d1fe23eabc4c98b48bffcef5a386062
OUT=/home/yewentao256/dsv4-c128a-lm-eval-results
mkdir -p "$OUT"

exec /home/yewentao256/.venv/bin/python \
  -m vllm.entrypoints.cli.main serve "$MODEL" \
  --served-model-name dsv4-flash \
  --tensor-parallel-size 2 \
  --enable-expert-parallel \
  --max-model-len 1048576 \
  --max-num-batched-tokens 8192 \
  --max-num-seqs 512 \
  --gpu-memory-utilization 0.95 \
  --kv-cache-dtype fp8_ds_mla \
  --no-enable-prefix-caching \
  --tokenizer-mode deepseek_v4 \
  --trust-remote-code \
  --host 127.0.0.1 \
  --port 18004

lm_eval   --model local-completions   --model_args "model=dsv4-flash,tokenizer=/data/LucasWilkinson/hub_cache/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01b1d1fe23eabc4c98b48bffcef5a386062,base_url=http://127.0.0.1:18004/v1/completions,num_concurrent=32,tokenized_requests=False,trust_remote_code=True,timeout=7200"   --tasks gpqa_diamond_cot_zeroshot   --gen_kwargs "max_gen_toks=32768,temperature=0.0,do_sample=False"   --output_path /home/yewentao256/dsv4-c128a-lm-eval-results/reasoning   --log_samples   --seed 0

And we get

# now
|          Tasks          |Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-------------------------|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gpqa_diamond_cot_zeroshot|    2.2|flexible-extract|     0|exact_match|↑  |0.5455|±  |0.0355|
|                         |       |strict-match    |     0|exact_match|↑  |0.0758|±  |0.0189|
|gsm8k                    |    3.0|flexible-extract|     5|exact_match|↑  |0.9538|±  |0.0058|
|                         |       |strict-match    |     5|exact_match|↑  |0.9538|±  |0.0058|

# main
|gpqa_diamond_cot_zeroshot|    2.2|flexible-extract|     0|exact_match|↑  |0.5051|±  |0.0356|
|                         |       |strict-match    |     0|exact_match|↑  |0.1010|±  |0.0215|
|gsm8k                    |    3.0|flexible-extract|     5|exact_match|↑  |0.9462|±  |0.0062|
|                         |       |strict-match    |     5|exact_match|↑  |0.9469|±  |0.0062|

Also CC @zyongye