[ROCm][Perf][Bugfix] DSv4 indexer: use platform FP8 dtype (fnuz) for Q-quant on gfx942

[akii96]

Motivation

On gfx942 the DeepSeek-V4 Flash indexer quantizes Q and K to different FP8 types. K already uses the platform type (e4m3fnuz on gfx942, via current_platform.fp8_dtype()), but the fused RoPE+quant kernel in fused_indexer_q.py hardcodes Q to e4m3fn. The FP8 logits kernel then gets fnuz K with fn Q and falls back to a mixed-dtype path on every call. This change derives Q's type from current_platform.fp8_dtype() as well, so on gfx942 both are fnuz and the logits kernel runs its native fnuz/fnuz path.

This is gfx942-specific by design. is_fp8_fnuz() is true only for gfx94x, so on gfx950 the platform type is OCP e4m3fn and Q/K are already fn/fn. Nothing changes there, and the NVIDIA cutedsl and MXFP4 paths are untouched (the two new kernel constexprs are defaulted).

The fnuz quant max is set to 224.0 to match get_fp8_min_max() in quant_utils.py, the value the K cache already uses.

Results

End to end serving of DeepSeek-V4 Flash (TP4, gfx942 / MI300), mean TTFT, prefill-heavy (OSL=27, concurrency 4):

ISL	OSL	main (mixed fp8 q=fn / k=fnuz)	This PR (fnuz q/k)	Speedup
8,192	27	3283 ms	1260 ms	2.6x
32,768	27	39744 ms	5350 ms	7.4x

Bonus correctness check: the existing kernel test tests/kernels/test_fused_indexer_q_rope_quant.py matches the unfused reference bit for bit on 9 of 10 gfx942 shapes. The one miss is 3 of 8,380,416 values at float32 / 1023 tokens, from fused vs unfused RoPE rounding at FP8 boundaries, not a dtype or scale error.

Note

The same indexer dtype handling was included in the larger ROCm enablement PRs #41601 and #42033, both stalled on rebase since May. This is a minimal, standalone version of just that fix for the gfx942 path; main as of today still hardcodes Q to e4m3fn.

Repro: serve + bench commands (DeepSeek-V4 Flash, TP4, gfx942)

Serve (4x MI325):

HIP_VISIBLE_DEVICES=0,1,2,3 VLLM_ROCM_USE_AITER=1 \
vllm serve deepseek-ai/DeepSeek-V4-Flash \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.85 \
  --kv-cache-dtype fp8_e4m3 \
  --block-size 256 \
  --max-model-len 132096 \
  --max-num-batched-tokens 16384 \
  --max-num-seqs 156 \
  --async-scheduling \
  --no-enable-prefix-caching \
  --tokenizer-mode deepseek_v4 \
  --reasoning-parser deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --enable-auto-tool-choice \
  --disable-log-stats \
  --host 0.0.0.0 --port 8000
  
  vllm bench serve --backend vllm \
  --model deepseek-ai/DeepSeek-V4-Flash \
  --host localhost --port 8000 \
  --dataset-name random --ignore-eos --trust-remote-code \
  --seed 5678 \
  --random-input-len 8192 \
  --random-output-len 27 \
  --max-concurrency 4 --num-prompts 12 --num-warmups 4

# 32K ISL: same command with --random-input-len 32768

[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.

[akii96]

Ping here also @tjtanaa since you were reviewing the predecessor PRs! This is a much smaller scoped PR which has already big impact on the prefill slowness

[akii96]

I rebased earlier today and now the CI looks green 🎉 @tjtanaa

[tjtanaa]

@akii96 Please also provide an end to end model lmeval score as well?

after you launch your vllm server.

Use this command to test the vllm server

#!/bin/bash

# SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
# LM_EVAL_PYTHON=${LM_EVAL_PYTHON:-"${SCRIPT_DIR}/hfvenv/bin/python"}
MODEL=deepseek-ai/DeepSeek-V4-Flash
lm_eval --model local-completions --model_args model=$MODEL,base_url=http://0.0.0.0:8001/v1/completions,num_concurrent=256,max_retries=10,max_gen_toks=2048,max_length=1048576,timeout=60000,trust_remote_code=True --batch_size auto --tasks gsm8k --num_fewshot 20 \
  --output_path ./results_dsv4_numshot20_cc256 \
  --log_samples \
| tee lmeval_dsv4_numshot20_cc256.log

We don't have a e2e model test on CI at the moment.

[akii96]

@tjtanaa apologies for not having this when I opened the PR!

Here is the lm-eval GSM8K results (20-shot) before and after the PR, using your command:

Baseline (main, mixed fp8 q=fn / k=fnuz):

Tasks	Version	Filter	n-shot	Metric	Value	Stderr
gsm8k	3	flexible-extract	20	exact_match	0.9227	± 0.0074
		strict-match	20	exact_match	0.9151	± 0.0077

With PR (fnuz q/k matched):

Tasks	Version	Filter	n-shot	Metric	Value	Stderr
gsm8k	3	flexible-extract	20	exact_match	0.9280	± 0.0071
		strict-match	20	exact_match	0.9189	± 0.0075

I used the same serve command as in the main PR body + the latest nightly for my tests vllm/vllm-openai-rocm:nightly-a16dbd5b8572d4128be9f10b9dcff4999b594b25

[tjtanaa]

LGTM