[Qwen3.8-Flash-Next] Tune FP8 TP2/TP4 Triton MoE on B200

[gcanlin]

Purpose

Add missing B200 Triton MoE configurations for Qwen3.8-Flash-Next-FP8 at TP2 (local intermediate 320, FP8 blocks 64×64) and TP4 (160, 32×32), with E=512 and top-k=10. Tune 14 token counts from 1 to 8192 for each shape. These shapes differ from #43506 and are independent of the FlashInfer loader change in #55867. TP1 keeps native 128×128 blocks and defaults to TRTLLM on B200, so it is outside this Triton tuning change.

Test Plan

Deploy

MODEL=Qwen/Qwen3.8-Flash-Next-FP8
TP=4  # use TP=2 and CUDA_VISIBLE_DEVICES=0,1 for TP2
INPUT_LEN=8192 OUTPUT_LEN=1024 MAX_MODEL_LEN=10240
# Also test INPUT_LEN=128 or 2048 with OUTPUT_LEN=128.
# TP4 short-input runs use MAX_MODEL_LEN=4096; TP2 uses 10240 throughout.
# Run the base and PR checkouts serially on the same GPUs.
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve "$MODEL" \
  --served-model-name qwen-flash --host 127.0.0.1 --port 8000 \
  --tensor-parallel-size "$TP" --moe-backend triton \
  --max-model-len "$MAX_MODEL_LEN" --max-num-seqs 32 --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.85 --no-enable-prefix-caching --enable-flashinfer-autotune \
  --limit-mm-per-prompt '{"image":0,"video":0}' \
  --compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[1,2,4,8,16,32]}'

Benchmark

Online reproduction, in a second terminal with matching workload variables; the tables below are from offline fixed-batch measurements.

MODEL=Qwen/Qwen3.8-Flash-Next-FP8
INPUT_LEN=8192 OUTPUT_LEN=1024
TP=4  # match the server
RESULT_DIR=results-tp${TP}-triton-tuned  # use a separate directory for the base checkout
for concurrency in 1 8 32; do
  vllm bench serve --backend vllm --base-url http://127.0.0.1:8000 \
    --model "$MODEL" --served-model-name qwen-flash \
    --dataset-name random --random-input-len "$INPUT_LEN" \
    --random-output-len "$OUTPUT_LEN" --random-range-ratio 0 --random-prefix-len 0 \
    --num-prompts "$((5 * concurrency))" --num-warmups "$((3 * concurrency))" \
    --request-rate inf --max-concurrency "$concurrency" \
    --ignore-eos --temperature 0 --seed 42 \
    --save-result --result-dir "$RESULT_DIR" \
    --result-filename "input${INPUT_LEN}-output${OUTPUT_LEN}-concurrency${concurrency}.json"
done

Recorded offline commands using the local fixed-batch harness:

EXPERIMENT_DIR=/gpfs/mszn/workspace/qwen-fp8-tp-fi-experiments
bash "$EXPERIMENT_DIR/run_tuned_triton_tp4.sh"
/opt/canlin-qwen-fp8-tp/.venv/bin/python "$EXPERIMENT_DIR/compare_tuned_triton_tp4.py"

PY=/opt/canlin-qwen-fp8-tp/.venv/bin/python
export CUDA_VISIBLE_DEVICES=0,1
export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 VLLM_WORKER_MULTIPROC_METHOD=spawn
unset VLLM_TUNED_CONFIG_FOLDER
"$PY" "$EXPERIMENT_DIR/benchmark_tp2.py" --backend triton --all-profiles --output "$EXPERIMENT_DIR/perf-tp2-triton.json"
export VLLM_TUNED_CONFIG_FOLDER="$EXPERIMENT_DIR/triton-tuning-tp2/configs"
"$PY" "$EXPERIMENT_DIR/benchmark_tp2.py" --backend triton --all-profiles --output "$EXPERIMENT_DIR/perf-tp2-tuned-triton.json"
"$PY" "$EXPERIMENT_DIR/compare_tuned_triton_tp2.py"

Test Result

TP4

4×B200, TP4, synthetic inputs, fixed output lengths, 3 warmups + 5 measured batches, median output throughput. Decode CUDA Graphs enabled; prefix caching, detokenization, and torch.compile disabled. Default measurements were recorded earlier on the same GPUs and runtime; tuned measurements load this exact JSON through VLLM_TUNED_CONFIG_FOLDER.

Input / output tokens	Batch	Default Triton tok/s	Tuned Triton tok/s	Throughput improvement
128 / 128	1	133.3	141.6	+6.2%
128 / 128	8	729.5	815.3	+11.8%
128 / 128	32	2198.2	2619.3	+19.2%
2048 / 128	1	133.5	141.3	+5.9%
2048 / 128	8	611.5	681.7	+11.5%
2048 / 128	32	1251.2	1465.0	+17.1%
8192 / 1024	1	147.2	157.6	+7.1%
8192 / 1024	8	730.6	859.3	+17.6%
8192 / 1024	32	1607.4	1984.0	+23.4%

Check	Result
Kernel output vs. default config, 14 token counts	Maximum relative L2 < 2.2e-6
Full-model smoke	42 / Paris / 77
Config lookup	All 14 entries loaded from the standard config directory
Preemptions	0 in all measured cases
8192/1024 batch-time CV	0.01% / 0.17% / 0.86% at batch 1 / 8 / 32
128/128, batch 8 TTFT	189.9 → 248.5 ms regression; tuned batch-time CV 4.91%

TP2

2×B200, same fixed-batch measurement method, with context 10240 for all workloads. Default and tuned runs were measured serially on the same GPUs; logs confirm the exact TP2 JSON was loaded.

Input / output tokens	Batch	Default Triton tok/s	Tuned Triton tok/s	Throughput improvement
128 / 128	1	136.8	139.3	+1.8%
128 / 128	8	748.7	722.9	-3.4%
128 / 128	32	2248.2	2349.0	+4.5%
2048 / 128	1	136.8	139.1	+1.7%
2048 / 128	8	663.3	685.9	+3.4%
2048 / 128	32	1321.9	1364.1	+3.2%
8192 / 1024	1	150.7	153.5	+1.8%
8192 / 1024	8	760.6	791.3	+4.0%
8192 / 1024	32	1684.4	1735.8	+3.0%

Check	Result
Kernel output vs. default config, 14 token counts	Maximum relative L2 < 1.6e-6
Full-model smoke, default and tuned	42 / Paris / 77
Preemptions	0 in all measured cases
8192/1024 tuned batch-time CV	0.01% / 0.21% / 0.54% at batch 1 / 8 / 32
128/128, batch 8 throughput	-3.4%; default/tuned batch-time CV 4.47%/4.87%, so a stable regression is not established

Measured on source 1217ddc5b00a with native image build 8a728663c1c3, Torch 2.13.0+cu130, FlashInfer 0.6.18, and Triton 3.7.1. The measured checkout contains #55867, whose explicit FlashInfer path is inactive for these Triton runs. The Triton kernel and config selection code are unchanged at this PR's base 13cf9e05c1ed. Validation includes synthetic kernel comparisons and smoke prompts, not a model-quality benchmark.

AI-assisted implementation and testing with OpenAI Codex.

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

[gau-nernst]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87693 for commit 7b3dcb9b9947.