perf(musa): fuse DeepSeek-V4 shared-expert SwiGLU FP8 down-projection

[Joey-gvwal]

Summary

This PR adds a MUSA fast path for the DeepSeek-V4 dense/shared-expert MLP. It
fuses the clamp-aware SwiGLU activation and per-token group-128 FP8
quantization that feed the existing DeepGEMM down-projection.

The change removes the materialized BF16 activation and replaces the previous
sequence of clamp, sigmoid, multiply, copy, absmax, scale, and FP8 conversion
kernels with one native MUSA kernel. The generated FP8 activations and FP32
scales are then consumed by the existing fp8_gemm_nt down-projection, and the
existing tensor-parallel all-reduce semantics are preserved.

No environment-variable gate is added. Unsupported shapes or configurations
return to the original vLLM activation plus linear path.

Motivation

When the DeepSeek-V4 dense/shared-expert MLP path runs, it executes:

gate_up_proj
  -> clamp(gate), clamp(up)
  -> SiLU(gate) * up
  -> BF16 activation tensor
  -> per-token group-128 FP8 quantization
  -> FP8 DeepGEMM down-projection
  -> TP all-reduce

For BS1 decode, the activation and quantization stages are small, launch-bound
kernels repeated once per layer. A four-step profile attributed roughly 5 ms
to the related ClipOp, SigmoidOp, MulOp, quantization, and activation-copy
family. The intermediate BF16 activation was written to global memory and
immediately read back by the quantizer.

The new path executes:

gate_up_proj
  -> fused clamp + SiLU + multiply + BF16 rounding + group absmax + FP8 quant
  -> existing FP8 DeepGEMM down-projection
  -> TP all-reduce

This reduces kernel-launch and intermediate-memory traffic without changing
the DeepGEMM implementation or the model's tensor-parallel contract.

Implementation

• Add silu_and_mul_clamp_per_token_group_fp8_quant, a clamp-aware variant of
  the native MUSA activation/quantization kernel.
• Preserve DeepSeek-V4 numerical ordering, including BF16-domain rounding
  before group absmax calculation.
• Add musa_silu_clamp_deepgemm_fp8_op, which invokes the fused quantization
  kernel and then the existing fp8_gemm_nt provider.
• Add a narrow RowParallelLinear.forward_swiglu_clamp() OOT hook and preserve
  reduce_results through tensor_model_parallel_all_reduce().
• Patch DeepseekV4MLP.forward() to try the optional hook before materializing
  the ordinary activation.
• Add source-contract tests for kernel registration, dispatch constraints,
  inverse-scale selection, fallback behavior, model-hook ordering, and TP
  reduction.

The fast path is selected only when all of the following hold:

• MUSA DeepGEMM block-FP8 row-parallel linear;
• input is already tensor-parallel and the projection has no bias;
• contiguous BF16 gate_up input;
• weight block size is (128, 128) and activation group size is 128;
• swiglu_limit == 10.0;
• batch-invariant mode is disabled.

All other configurations use the unchanged fallback path.

Single-batch service command

cd /workspace/musa3508-tps60-c13c3a38-699a7b28
unset PYTHONPATH

export SAFETENSORS_FAST_GPU=1
export VLLM_DEEP_GEMM_WARMUP=skip
export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
export VLLM_USE_DEEP_GEMM=1
export VLLM_USE_DEEP_GEMM_E8M0=0
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export PYTHONUNBUFFERED=1
export VLLM_MUSA_DEEPSEEK_V4_TP8_PROFILE=aggressive_long_prefill
export VLLM_MUSA_DEEPSEEK_V4_MOE_DEEPGEMM_PREFILL=1
export VLLM_MUSA_FUSED_AR_RMSNORM=0

unset TORCHDYNAMO_DISABLE
unset VLLM_ATTENTION_BACKEND
unset VLLM_MUSA_DEEPSEEK_V4_DISABLE_AUX_OVERLAP
unset VLLM_MUSA_DEEPSEEK_V4_MHC_PRE_IMPL
unset VLLM_MUSA_DEEPSEEK_V4_MHC_PRE_DECODE_PRENORM_IMPL
unset VLLM_MUSA_GEMV_MOE_BLOCK
unset VLLM_MUSA_FUSED_AR_RMSNORM_GRAPH_REGISTERED_INPUT

vllm serve /home/dist/models/DeepSeek-V4-Flash-Base \
  --served-model-name deepseek \
  --gpu-memory-utilization 0.95 \
  --max-model-len 6144 \
  --max-num-seqs 1 \
  --max-num-batched-tokens 8195 \
  --tensor-parallel-size 8 \
  -ac.backend FLASHMLA \
  --port 25606 \
  --host 0.0.0.0 \
  --kv-cache-dtype fp8 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --compilation-config \
  '{"mode":"NONE","cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":1,"cudagraph_capture_sizes":[1],"cudagraph_copy_inputs":false}'

Benchmark command

RESULT_ROOT=/home/dist/zuoyu/dsv4-bench/musa3508_shared_swiglu_bs1

for run in 1 2 3; do
  out_dir="$RESULT_ROOT/run$run"
  mkdir -p "$out_dir"

  vllm bench serve \
    --backend openai-chat \
    --dataset-name random \
    --seed 0 \
    --random-input-len 4096 \
    --random-output-len 1024 \
    --ignore-eos \
    --num-prompts 1 \
    --num-warmups 5 \
    --request-rate 10 \
    --max-concurrency 1 \
    --temperature 0 \
    --model deepseek \
    --tokenizer /home/dist/models/DeepSeek-V4-Flash-Base \
    --base-url http://127.0.0.1:25606 \
    --endpoint /v1/chat/completions \
    --save-result \
    --save-detailed \
    --result-dir "$out_dir" \
    2>&1 | tee "$out_dir/terminal.log"
done

Performance

Controlled BS1 A/B

The control and candidate were built from the same
upstream/v0.24.0-dev@11825fc1 source and measured with the same TP8,
capture-[1], no-MTP service configuration.

Variant	Output TPS runs	Mean output TPS	Mean TTFT	Mean TPOT
Control	48.8846 / 48.8838	48.8842 tok/s	783.27 ms	19.6125 ms
Fused shared-expert path	51.1668 / 50.9628	51.0648 tok/s	779.92 ms	18.7415 ms
Delta	-	+2.1806 tok/s (+4.46%)	-0.43%	-4.44%

Latest-upstream integration confirmation

After rebasing the optimization as 699a7b28 onto
upstream/v0.24.0-dev@c13c3a38, the same BS1 direct-path command produced:

Run	Output TPS	TTFT	TPOT	Result
1	51.0441 tok/s	781.45 ms	18.7479 ms	1/1, 1024 tokens
2	51.0860 tok/s	780.33 ms	18.7329 ms	1/1, 1024 tokens
3	51.2709 tok/s	780.87 ms	18.6617 ms	1/1, 1024 tokens
Mean	51.1337 tok/s	780.88 ms	18.7141 ms	3/3, zero failures

The latest-upstream rows are integration/no-regression evidence. The claimed
optimization gain is taken from the controlled A/B table above.

[Joey-gvwal]

I will clean up the DeepSeek-V4 related environment variables in this PR and submit the corresponding commit later.

[Joey-gvwal]

Vllm service command after DeepSeek-V4 env cleanup:

export SAFETENSORS_FAST_GPU=1
export VLLM_DEEP_GEMM_WARMUP=skip
export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
export VLLM_USE_DEEP_GEMM=1
export VLLM_USE_DEEP_GEMM_E8M0=0
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export PYTHONUNBUFFERED=1
export VLLM_MUSA_FUSED_AR_RMSNORM=0

unset TORCHDYNAMO_DISABLE
unset VLLM_ATTENTION_BACKEND
unset VLLM_MUSA_GEMV_MOE_BLOCK
unset VLLM_MUSA_FUSED_AR_RMSNORM_GRAPH_REGISTERED_INPUT

exec vllm serve /home/dist/models/DeepSeek-V4-Flash-Base \
  --served-model-name deepseek \
  --gpu-memory-utilization 0.95 \
  --max-model-len 6144 \
  --max-num-seqs 1 \
  --max-num-batched-tokens 8195 \
  --tensor-parallel-size 8 \
  -ac.backend FLASHMLA \
  --port 25606 \
  --host 0.0.0.0 \
  --kv-cache-dtype fp8 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --compilation-config \
  '{"mode":"NONE","cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":1,"cudagraph_capture_sizes":[1],"cudagraph_copy_inputs":false}'

and bench command：

vllm bench serve \
    --backend openai-chat \
    --dataset-name random \
    --seed 0 \
    --random-input-len 4096 \
    --random-output-len 1024 \
    --ignore-eos \
    --num-prompts 1 \
    --num-warmups 5 \
    --request-rate 10 \
    --max-concurrency 1 \
    --temperature 0 \
    --model deepseek \
    --tokenizer /home/dist/models/DeepSeek-V4-Flash-Base \
    --base-url http://127.0.0.1:25606 \
    --endpoint /v1/chat/completions \
    --save-result \
    --save-detailed