[simondanielsson]

(the text below is written by myself)

Purpose

1. increase standardization without imposing restrictions on kernel implementations
2. limit LoC changes in model files as a result from adding numerous specialized decode kernels

The kernels themselves will naturally remain difficult to review thoroughly, but ideally model code should remain easy to review. The abstraction layer should deliberately be made as thin as possible,

Note: you can run with either --monokernel or --specialized-decode monokernel

Claims

1. Add DecodeSpec and SpecializedDecode abstractions in vllm/models/common/specialized_decode/
2. Refactor DSv41 on gfx950 and gfx942 to use these abstractions
3. No regression

For details on point 1, see design doc in the end of the PR description.

Validation

Check no regression on dsv41

1. DeepSeek-v41 on MI350

On MI350, force max 8 seqs

VLLM_ROCM_USE_AITER=1 \
vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 4 \
  --tokenizer-mode deepseek_v41 \
  --trust-remote-code \
  --language-model-only \
  --max-model-len 8192 \
  --max-num-seqs 8 \
  --block-size 128 \
  --attention_config.indexer_kv_dtype=mxfp4 \
  --no-enable-prefix-caching \
  --gpu-memory-utilization 0.9 \
  --moe-backend aiter \
  --speculative-config "{\"method\":\"dspark\",\"num_speculative_tokens\":5}" \
  [--specialized-decode]

1. Performance non-regression check

vllm bench serve \
  --base-url http://localhost:8000 \
  --model deepseek-ai/DeepSeek-V4.1-Flash \
  --tokenizer-mode deepseek_v41 --trust-remote-code \
  --dataset-name random --random-input-len 65536 --random-output-len 1024 \
  --max-concurrency "$CONC" --num-prompts "$((CONC*10))" --num-warmups "$((CONC*2))" \
  --ignore-eos --seed 0 \
  --percentile-metrics ttft,tpot,itl,e2el

1k/512

build	flag	conc	median ITL (ms)	output tok/s	acceptance len
upstream nightly	--monokernel	1	5.91	249.50	1.55
branch	--specialized-decode	1	5.93	247.89	1.55
upstream nightly	(none)	1	9.37	176.78	1.72
branch	(none)	1	9.40	168.85	1.64
upstream nightly	--monokernel	4	8.86	482.20	1.78
branch	--specialized-decode	4	8.84	457.21	1.68
upstream nightly	(none)	4	11.73	427.99	2.03
branch	(none)	4	11.77	398.59	1.81

64k/1k

build	flag	conc	median ITL (ms)	output tok/s	acceptance len
upstream nightly	--monokernel	1	6.12	158.07	1.15
branch	--specialized-decode	1	6.16	149.34	1.09
upstream nightly	(none)	1	9.63	98.23	1.05
branch	(none)	1	9.60	104.48	1.12
upstream nightly	--monokernel	64	11.57	493.53	1.24
branch	--specialized-decode	64	11.55	489.41	1.23
upstream nightly	(none)	64	13.57	447.70	1.23
branch	(none)	64	13.58	451.17	1.25

Accuracy

# gsm8k
lm-eval --model local-completions \
  --model_args "model=$M,base_url=http://localhost:$PORT/v1/completions,tokenized_requests=False,trust_remote_code=True,num_concurrent=16,max_retries=5,timeout=3600" \
  --tasks gsm8k --num_fewshot 5 \
  --output_path "/acc/$ARM/gsm8k"

# gpqa
lm-eval --model local-chat-completions \
  --model_args "model=$M,base_url=http://localhost:$PORT/v1/chat/completions,tokenized_requests=False,num_concurrent=16,think_end_token=</think>,max_retries=5,timeout=3600" \
  --tasks gpqa_diamond_cot_zeroshot --batch_size auto --apply_chat_template \
  --gen_kwargs '{"max_gen_toks":64000}' --seed 42 \
  --output_path "/acc/$ARM/gpqa"

# niah @ 128k
lm-eval --model local-completions \
  --model_args "model=$M,base_url=http://localhost:$PORT/v1/completions,tokenized_requests=False,trust_remote_code=True,num_concurrent=8,max_length=140000,max_retries=5,timeout=3600" \
  --tasks niah_single_2 --metadata '{"max_seq_lengths":[131072]}' --seed 1234 \
  --output_path "/acc/$ARM/niah"

1, gsm8k

Main, specialized-decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9249	±	0.0073
		strict-match	5	exact_match	↑	0.9249	±	0.0073

Main no mono

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9242	±	0.0073
		strict-match	5	exact_match	↑	0.9249	±	0.0073

This branch. specialized decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9363	±	0.0067
		strict-match	5	exact_match	↑	0.9371	±	0.0067

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9257	±	0.0072
		strict-match	5	exact_match	↑	0.9257	±	0.0072

2. gpqa_diamond_cot_zeroshot

Main, specialized-decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8333	±	0.0266
		strict-match	0	exact_match	↑	0.0000	±	0.0000

Main no mono

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8485	±	0.0255
		strict-match	0	exact_match	↑	0.0000	±	0.0000

This branch, specialized decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8485	±	0.0255
		strict-match	0	exact_match	↑	0.0051	±	0.0051

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8535	±	0.0252
		strict-match	0	exact_match	↑	0.0000	±	0.0000

3. niah_single_2 @ 128k

Main, specialized-decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

Main no mono

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

This branch. specialized decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

2. DSv41 on MI300

export VLLM_ROCM_USE_AITER=1
vllm serve deepseek-ai/DeepSeek-V4.1-Flash \
  --tensor-parallel-size 4 \
  --tokenizer-mode deepseek_v41 \
  --trust-remote-code \
  --language-model-only \
  --max-num-seqs 7 \
  --block-size 64 \
  --no-enable-prefix-caching \
  --gpu-memory-utilization 0.97 \
  --moe-backend aiter \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5}' \
  --specialized-decode

vllm bench serve:

build	flag	workload	conc	median ITL (ms)	output tok/s	acceptance len
main	--monokernel	1k/512	1	7.32	218.62	1.67
branch	--specialized-decode*	1k/512	1	7.54	215.63	1.69
main	off	1k/512	1	14.90	98.90	1.50
branch	off	1k/512	1	14.90	104.68	1.60
main	--monokernel	1k/512	4	13.54	446.72	1.71
branch	--specialized-decode*	1k/512	4	13.95	419.52	1.80
main	off	1k/512	4	22.48	288.84	1.74
branch	off	1k/512	4	22.12	300.26	1.77
main	--monokernel	64k/1k	1	7.53	106.13	1.10
branch	--specialized-decode*	64k/1k	1	7.73	102.89	1.08
main	off	64k/1k	1	15.51	61.57	1.14
branch	off	64k/1k	1	15.51	58.53	1.07

gsm8k

Main, specialized dedode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9227	±	0.0074
		strict-match	5	exact_match	↑	0.9234	±	0.0073

This branch, specialized dedode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9234	±	0.0073
		strict-match	5	exact_match	↑	0.9242	±	0.0073

Main

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9295	±	0.0071
		strict-match	5	exact_match	↑	0.9287	±	0.0071

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9227	±	0.0074
		strict-match	5	exact_match	↑	0.9227	±	0.0074

gpqa diamond zero shot

Main, specialized dedode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8434	±	0.0259
		strict-match	0	exact_match	↑	0.0000	±	0.0000

This branch, specialized dedode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8586	±	0.0248
		strict-match	0	exact_match	↑	0.0000	±	0.0000

Main

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8384	±	0.0262
		strict-match	0	exact_match	↑	0.0000	±	0.0000

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gpqa_diamond_cot_zeroshot	2.2	flexible-extract	0	exact_match	↑	0.8636	±	0.0245
		strict-match	0	exact_match	↑	0.0000	±	0.0000

niah 128k

Main

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

This branch

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

Main specialized decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

This branch specialized decode

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
niah_single_2	1	none	0	131072	↑	1	±	N/A
		none	0	4096	↑	-1	±	N/A

1.1 Points from the RFC not covered

1. Replacing only a layer, i.e. the folllowing. This can easily be achieved through the "whole layer" mode anyways

Replace the Decode Layer within the forward function with the specific kernel to run, and let the forward function decide.

[mergify]

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

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

[Fangzhou-Ai]

Hi @simondanielsson thanks for the work! For now I have one concern over the specialized_decode everywhere stlye abstraction. mono kernel is only a subset of "specialized decode" which is dedicated to small batch size inference. using --specialized-decode to enable mono kernel feature would be a blocker and leads to confusions in the future if we add more specialized decoding method. I would recomment 1. still keep --monokernel as a valid option as it's straightforward, and keep our recipes remain coorect, 2. make the flag as an config choice, e.g. --specialized-decode monokernel for future compatability.

[simondanielsson]

Hi @simondanielsson thanks for the work! For now I have one concern over the specialized_decode everywhere stlye abstraction. mono kernel is only a subset of "specialized decode" which is dedicated to small batch size inference. using --specialized-decode to enable mono kernel feature would be a blocker and leads to confusions in the future if we add more specialized decoding method. I would recomment 1. still keep --monokernel as a valid option as it's straightforward, and keep our recipes remain coorect, 2. make the flag as an config choice, e.g. --specialized-decode monokernel for future compatability.

@Fangzhou-Ai Great point, you're right.

I'll do:

1. keep --monokernel
2. revert changes to recipes
3. --specialized-decode [auto|monokernel], both current map to same as --monokernel

[ChuanLi1101]

One check needs to come back. _record_view treats the KV cache as [blocks, block, RECORD] and steps to the next block with kv.stride(0). That stride has to be at least block * RECORD, or the next block starts inside the current one. refuse_caches only checks stride(0) < RECORD. Please validate kv.stride(0) >= block * RECORD before creating the view.

On MI350 the monokernel ITL matches main. The tok/s gaps follow acceptance on a short random bench. GSM8K and GPQA stay within stderr.

[shen-shanshan]

It's better to add the introduction/tutorial in the PR body into https://github.com/vllm-project/vllm/tree/main/docs/design.

[mergify]

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

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

[mergify]

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

[simondanielsson]

It's better to add the introduction/tutorial in the PR body into https://github.com/vllm-project/vllm/tree/main/docs/design.

@shen-shanshan Good idea, added and refined a bit. PTAL 🙏

[mergify]

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

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