[Perf][DSV4] Optimize global top-k index kernel with compile-time constants

[chaunceyjiang]

Purpose

Optimize global top-k index kernel with compile-time constants

Test Plan

vllm serve deepseek-ai/DeepSeek-V4-Flash-0731 --trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 8 --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --no-enable-prefix-caching --max-num-batched-tokens 16384

vllm bench serve --backend vllm --base-url http://localhost:8000 --model deepseek-ai/DeepSeek-V4-Flash-0731 --dataset-name random --random-input-len 1024 --random-output-len 64 --num-prompts 128 --num-warmups 8 --request-rate inf --ignore-eos --temperature 0 --seed 701 --save-result

Test Result

Serving Benchmark Results

Version	Mean Output Throughput	Mean TPOT	Completed Requests	Relative Change
main	634.97 tokens/s	103.33 ms	128/128	Baseline
this pr	638.12 tokens/s	102.32 ms	128/128	Throughput +0.50%, TPOT -0.98%

Kernel Microbenchmark

import torch
import triton

from vllm.models.deepseek_v4.common.ops.cache_utils import (
    _compute_global_topk_indices_and_lens_kernel,
)

num_tokens = 512
topk = 512
block_size = 256
triton_block_size = 1024

global_topk_indices = torch.empty(
    (num_tokens, topk), device="cuda", dtype=torch.int32
)
global_topk_lens = torch.empty(num_tokens, device="cuda", dtype=torch.int32)
topk_indices = torch.randint(
    0, 8192, (num_tokens, topk), device="cuda", dtype=torch.int32
)
block_table = torch.arange(
    8192, device="cuda", dtype=torch.int32
).reshape(num_tokens, -1)


def run():
    _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
        global_topk_indices,
        global_topk_indices.stride(0),
        global_topk_lens,
        topk_indices,
        topk_indices.stride(0),
        topk,
        block_table,
        block_table.stride(0),
        block_size,
        TRITON_BLOCK_SIZE=triton_block_size,
    )


for _ in range(20):
    run()
torch.cuda.synchronize()

latency_us = triton.testing.do_bench(run, warmup=100, rep=500)
print(f"Latency: {latency_us:.3f} us")

Kernel Version	Latency	Relative Change
Original implementation	~7.50 us	Baseline
Compile-time-constant optimization	~6.37 us	15.1% faster

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

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.