[Attention] Skip sparse indexer scoring for dense short prefills

[qianlihuang]

Purpose

Try extend #47327 to see some improvement.

This PR removes dead indexer scoring work before dense attention.

Test Plan

MODEL_NAME=zai-org/GLM-5.2-FP8
export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=1
export VLLM_MOE_SKIP_PADDING=1

vllm serve "$MODEL_NAME" \
  -tp 8 \
  --kv-cache-dtype fp8 \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}' 

vllm bench serve \
  --backend openai-chat \
  --model "$MODEL_NAME" \
  --tokenizer "$MODEL_NAME" \
  --endpoint /v1/chat/completions \
  --host 127.0.0.1 \
  --port 8000 \
  --dataset-name random \
  --num-prompts 128 \
  --random-input-len 1024 \
  --random-output-len 1 \
  --random-range-ratio '{"input":0.125,"output":0}' \
  --num-warmups 16 \
  --max-concurrency 1 \
  --save-result

--max-num-seqs 1
+
python3 - <<'PY'
import json
import runpy

m = runpy.run_path("tests/evals/gsm8k/gsm8k_eval.py")

result = m["evaluate_gsm8k"](
    num_shots=8, # 8-shot keeps all prompts within 1,156–1,319 tokens
    request_timeout_seconds=3600,
)

print(json.dumps(result, indent=2))
PY

Test Result

ISL=1024, OSL=1, concurrency=1, 128 requests:

Metric	Base (e6d1310b)	PR (46e170b7)	Delta
Request throughput	5.998 req/s	6.150 req/s	+2.54%
Median TTFT	167.43 ms	163.22 ms	−2.52%
P95 TTFT	168.76 ms	164.21 ms	−2.70%

gsm8k

main

{
  "accuracy": 0.9423805913570887,
  "invalid_rate": 0.0,
  "latency": 942.1709538325667,
  "questions_per_second": 1.3999582502884074,
  "total_output_tokens": 134486,
  "tokens_per_second": 142.7405498470711,
  "num_questions": 1319,
  "num_shots": 8,
  "max_tokens": 256,
  "timestamp": 1784696291.8677127
}

PR

{
  "accuracy": 0.9423805913570887,
  "invalid_rate": 0.0,
  "latency": 932.1877947188914,
  "questions_per_second": 1.4149509438682952,
  "total_output_tokens": 134486,
  "tokens_per_second": 144.26921352317783,
  "num_questions": 1319,
  "num_shots": 8,
  "max_tokens": 256,
  "timestamp": 1784694381.1623588
}

Profile

at::native::vectorized_elementwise_kernel, vllm::cp_gather_indexer_k_quant_cache_kernel, deep_gemm::sm90_fp8_mqa_logits, vllm::topKPerRowPrefill have been skipped.

Decode regression check

The table reports the median result across three runs.

Concurrency	Main mean TPOT	PR mean TPOT	Delta	Main median TPOT	PR median TPOT	Delta
1	5.41 ms	5.40 ms	−0.18%	5.42 ms	5.41 ms	−0.18%
8	6.44 ms	6.42 ms	−0.31%	6.44 ms	6.41 ms	−0.47%
32	9.37 ms	9.35 ms	−0.21%	9.34 ms	9.33 ms	−0.11%

No decode regression was observed. TPOT is unchanged or slightly lower with the PR at all tested concurrency levels.

Environment and commands

Base image:

vllm/vllm-openai:nightly-f25953cc59f9b4ba9b04b16228d2b86dcfbcbdb1

Server configuration:

MODEL_NAME=/glm52

export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=1
export VLLM_MOE_SKIP_PADDING=1

vllm serve "$MODEL_NAME" \
  --load-format dummy \
  --max-model-len 262144 \
  --max-num-seqs 64 \
  --max-cudagraph-capture-size 256 \
  -tp 8 \
  --kv-cache-dtype fp8 \
  --speculative-config '{
    "method": "mtp",
    "num_speculative_tokens": 3,
    "rejection_sample_method": "synthetic",
    "synthetic_acceptance_rates": [0.90, 0.75, 0.50]
  }' \
  --no-enable-prefix-caching

Benchmark command:

for concurrency in 1 8 32; do
  num_prompts=$((concurrency * 8))
  num_warmups=$((concurrency * 2))

  for run in 1 2 3; do
    vllm bench serve \
      --backend openai-chat \
      --endpoint /v1/chat/completions \
      --host 127.0.0.1 \
      --port 8000 \
      --model /glm52 \
      --tokenizer /glm52 \
      --dataset-name random \
      --random-input-len 1 \
      --random-output-len 1024 \
      --random-range-ratio 0 \
      --num-prompts "$num_prompts" \
      --num-warmups "$num_warmups" \
      --max-concurrency "$concurrency" \
      --save-result
  done
done

Detailed results

Concurrency 1

Variant	Run	Duration	Output tok/s	Mean TPOT	Median TPOT	P99 TPOT	Acceptance rate
Main	1	44.59 s	183.71	5.40 ms	5.42 ms	5.52 ms	72.40%
Main	2	44.87 s	182.59	5.44 ms	5.49 ms	5.64 ms	71.49%
Main	3	44.64 s	183.52	5.41 ms	5.42 ms	5.51 ms	72.11%
PR	1	44.42 s	184.41	5.39 ms	5.40 ms	5.50 ms	72.40%
PR	2	44.74 s	183.12	5.43 ms	5.48 ms	5.63 ms	71.49%
PR	3	44.47 s	184.20	5.40 ms	5.41 ms	5.50 ms	72.11%

Concurrency 8

Variant	Run	Duration	Output tok/s	Mean TPOT	Median TPOT	P99 TPOT	Acceptance rate
Main	1	54.05 s	1212.48	6.48 ms	6.48 ms	6.74 ms	71.58%
Main	2	53.92 s	1215.36	6.44 ms	6.44 ms	6.75 ms	71.63%
Main	3	53.63 s	1221.96	6.41 ms	6.41 ms	6.74 ms	71.64%
PR	1	53.45 s	1226.13	6.42 ms	6.41 ms	6.71 ms	71.58%
PR	2	53.83 s	1217.41	6.43 ms	6.44 ms	6.83 ms	71.63%
PR	3	53.46 s	1225.96	6.40 ms	6.40 ms	6.69 ms	71.64%

Concurrency 32

Variant	Run	Duration	Output tok/s	Mean TPOT	Median TPOT	P99 TPOT	Acceptance rate
Main	1	82.21 s	3188.56	9.35 ms	9.32 ms	9.88 ms	71.53%
Main	2	104.67 s	2504.37	9.46 ms	9.34 ms	10.06 ms	71.75%
Main	3	78.41 s	3343.34	9.37 ms	9.35 ms	9.89 ms	71.66%
PR	1	80.80 s	3244.37	9.35 ms	9.33 ms	9.97 ms	71.53%
PR	2	77.87 s	3366.26	9.29 ms	9.29 ms	9.71 ms	71.75%
PR	3	78.36 s	3345.30	9.36 ms	9.34 ms	10.01 ms	71.66%

The second main run at concurrency 32 had an unusually long benchmark duration and lower aggregate throughput, while its TPOT remained close to the other runs.

AI Assistance

AI assistance was used for design analysis, implementation, tests, review, and community research.

---

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.

[mergify]

Documentation preview: https://vllm--48407.org.readthedocs.build/en/48407/

[mergify]

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

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

[MatthewBonanni]

Hi @qianlihuang, thanks for the contribution! Is this ready for review? If so, please mark it as such

[qianlihuang]

@MatthewBonanni The implementation in this PR is mostly final. The only blocker is that I don’t have suitable GPUs to run some validation and profiling for now.

The expected production benefit is for pure short-prefill steps, such as the early turns of bs=1 multi-turn conversations, and some prefill dp ranks in PD-disaggregated deployments.

I also saw #48770. There may be some minor conflicts to resolve, but nothing fundamental.

Reviews of the current approach and implementation are very welcome.

[MatthewBonanni]

Thanks! Could you add a gsm8k eval to ensure no accuracy regression?

[qianlihuang]

@MatthewBonanni Added the gsm8k comparison to the PR description.

[MatthewBonanni]

LGTM other than a small nit, thanks!

[MatthewBonanni]

I disabled auto-merge because a similar PR for DSv4, #49486, actually showed a minor decode regression. Can you benchmark a decode-dominated workload to ensure no regression?

[qianlihuang]

@MatthewBonanni Added a decode-dominated benchmark with ISL=1, OSL=1024, and concurrency 1/8/32 in PR description. Across three runs, median TPOT is unchanged at every concurrency.

[MatthewBonanni]

Hi @qianlihuang, please stop updating the branch because the CI runs are costly. For each CI failure, determine whether it's present on main or not. If you can't find an exactly matching failure on main (some might be flaky), please fix the bug. Otherwise, if all failures are present on main, let me know or post in #pr-merge-requests in Slack

[qianlihuang]

Thanks, @MatthewBonanni. I did try to get a clean CI run, but each run failed in a different, unrelated test. For context, the only failure on cc54941 was caused by the issue fixed in #50060, so the checks relevant to this PR had effectively passed.