Port DeepSeek-V4 optimization PR #174 to v0.26 and fix MTP multi-batch Graph staging

[Joey-gvwal]

Summary

The PR is therefore a v0.26 port plus compatibility/robustness fixes; the
Graph-staging changes are not a separate performance optimization series.

1. a0f9a1e3 DSV4 optimization (#174)
   Whole-commit cherry-pick of becb11a3 onto v0.26.0-dev.
   Any v0.26 source-layout/model-hook adaptation belongs in this port commit.

2. db2c2bc5 fix(musa): handle v0.26 Graph profiling for DSV4 CAR staging
   Do not seal the persistent CAR arena during throwaway profiling capture.

3. 1e80f150 fix(musa): cover DSV4 MTP BS32 and BS64 CAR Graph staging
   Add the (160,32) and (320,64) descriptors so BS64 does not fall back to
   the hanging standard-collective Graph path.

The five kernel optimizations and the required v0.26 compatibility changes are
part of commit 1. Commits 2 and 3 are correctness/robustness fixes for the
v0.26 port, not additional independent performance claims.

Optimizations carried from PR #174

Area	Ported change	Effect
Attention cache write	Parallelize inverse-RoPE FP8 cache-write chunks	Replaces serial chunk processing with parallel work and reduces the decode cache-write critical path.
FP8 quantization	Keep FP32 scale groups contiguous	Improves scale access locality and avoids unnecessary regrouping.
Attention O-projection	Select calibrated GEMV tiles for DeepSeek-V4 shapes	Uses shape-specific tiles for the small-M decode O-projection.
Attention O-projection	Write the GEMV result directly to the caller-owned output	Removes an intermediate output and the following copy.
Shared expert MoE	Route TP8 shared gate-up M=1 to the GEMV path	Uses the low-overhead kernel for single-token decode rather than a larger GEMM path.
Shared expert MoE	Keep the fused shared SwiGLU/FP8 path available in unified multi-batch deployments	Avoids disabling the BS1 optimization merely because the server captures BS1/4/16/64.
MTP sparse attention and CAR	Preserve sparse-prefill headroom, mixed-prefill fencing, sparse direct output, and fixed-address custom-all-reduce staging	Keeps the MTP target and draft Graph paths memory-safe while avoiding materialized sparse output where the direct path applies.

The optimized paths remain narrowly guarded to the validated
DeepSeek-V4-Flash-Base TP8, FP8 KV-cache, FLASHMLA, FULL_DECODE_ONLY contract.
Generic models and non-matching configurations fail closed to their existing
paths.

v0.26 port compatibility

The v0.26 branch does not contain the v0.24 PR #166 model-wide DeepSeek-V4
optimization contract on which part of the source PR was based. These are
mechanical adaptations required to make the one-commit cherry-pick usable on
v0.26; they are not additional independent optimizations:

1. Add a local, shape-checked DeepSeek-V4 runtime policy for the PR DSV4 optimization #174 MTP
   paths instead of importing the unavailable v0.24 contract.
2. Restore the DeepSeek-V4 shared-SwiGLU model hook through a v0.26 patch so
   the fused shared-expert path is actually installed into vLLM 0.26.
3. Adapt the MTP sparse-prefill and mixed-prefill patches to the v0.26 vLLM
   source layout.
4. Preserve the v0.26 third-party vLLM and FlashInfer pins while carrying the
   vllm-musa kernels and dispatch policies from PR DSV4 optimization #174.

Root cause of the v0.26 MTP BS64 failure

The failure was caused by an incomplete v0.26 adaptation, not by an inherent
BS64 or DeepSeek-V4 model limit.

1. Temporary Graph profiling sealed the persistent CAR arena

vLLM 0.26 first captures throwaway graphs in
profile_cudagraph_memory() and clears them before the persistent
capture_model() pass. The original PR #174 staging logic assumed a single
persistent capture and sealed the custom-all-reduce arena after the profiling
capture.

The later persistent capture then failed with:

RuntimeError: MUSA custom AR graph staging arena cannot be reused while
previously captured graphs may still reference its slots

The port now recognizes the temporary profiling phase through
compilation_counter.num_gpu_runner_capture_triggers == 0. A successful
profiling capture may release and reuse its staging slots; only the persistent
model capture seals the arena against unsafe recapture.

2. MTP staging descriptors stopped at BS16

MTP4 schedules five model tokens per request: one target token plus four draft
tokens. A production capture ladder for request batches
[1, 2, 4, 8, 16, 32, 64] therefore uses Graph token sizes
[5, 10, 20, 40, 80, 160, 320].

The original staging policy only covered:

(5,1), (10,2), (20,4), (40,8), (80,16)

BS32 and BS64 fell back to the standard collective path inside the captured
Graph. In the BS64 test that path stopped making useful progress and produced
only 226 of the required 65,536 output tokens before the engine was
terminated.

The final policy covers all seven descriptors:

(5,1), (10,2), (20,4), (40,8), (80,16), (160,32), (320,64)

The complete communicator arena remains exactly 512 MiB, which is the existing
DeepSeek-V4 MTP upper bound. The fix does not raise that limit.

Files changed by the v0.26 Graph follow-up

vllm_musa/deepseek_v4_policy.py
vllm_musa/distributed/device_communicators/musa_jit_custom_all_reduce.py
tests/test_deepseek_v4_mtp_policy.py
tests/test_musa_jit_custom_all_reduce_graph.py

The vendored vLLM checkout was reset to its pinned commit before installation;
the expected TorchAda/vllm-musa source mappings and patches were then applied
by the source install. The local vllm-musa sources and the four modified files
inside the validation container were hash-checked before serving.

Common environment

export MUSA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
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

PORT=25606
MODEL=/home/dist/models/DeepSeek-V4-Flash-Base
SERVED_MODEL=deepseek

No PR #174 optimization requires an additional DeepSeek-V4 A/B-test
environment gate in these commands.

No-MTP serve command

The no-MTP service captures the complete request-batch ladder in one process;
BS1 is not run through a dedicated capture-[1]-only service.

COMPILATION_CONFIG='{"mode":"NONE","cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":64,"cudagraph_capture_sizes":[1,2,4,8,16,32,64],"cudagraph_copy_inputs":false}'

vllm serve "$MODEL" \
  --served-model-name "$SERVED_MODEL" \
  --gpu-memory-utilization 0.94 \
  --max-model-len 6144 \
  --max-num-seqs 64 \
  --max-num-batched-tokens 8195 \
  --tensor-parallel-size 8 \
  -ac.backend FLASHMLA \
  --port "$PORT" \
  --host 0.0.0.0 \
  --kv-cache-dtype fp8 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --compilation-config "$COMPILATION_CONFIG"

Startup evidence:

GPU KV cache size: 69,408 tokens
Capturing CUDA graphs (decode, FULL): 7/7
Graph capturing finished in 15 secs, took 0.44 GiB

MTP4 serve command

COMPILATION_CONFIG='{"mode":"NONE","cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":320,"cudagraph_capture_sizes":[5,10,20,40,80,160,320],"cudagraph_copy_inputs":false}'
SPECULATIVE_CONFIG='{"method":"mtp","num_speculative_tokens":4}'

vllm serve "$MODEL" \
  --served-model-name "$SERVED_MODEL" \
  --gpu-memory-utilization 0.95 \
  --max-model-len 6144 \
  --max-num-seqs 64 \
  --max-num-batched-tokens 8195 \
  --tensor-parallel-size 8 \
  -ac.backend FLASHMLA \
  --port "$PORT" \
  --host 0.0.0.0 \
  --kv-cache-dtype fp8 \
  --no-enable-prefix-caching \
  --async-scheduling \
  --compilation-config "$COMPILATION_CONFIG" \
  --speculative-config "$SPECULATIVE_CONFIG"

Final MTP startup evidence:

Graph descriptors:        (5,1), (10,2), (20,4), (40,8), (80,16), (160,32), (320,64)
Profiling capture slots:   174 temporary slots
Persistent capture slots: 609 disjoint slots
Communicator arena:       512 MiB
Available KV cache:       35.08 GiB
GPU KV cache size:        66,150 tokens
Graph capture:            7/7
Graph memory:             0.59 GiB

Benchmark command

Each row used random inputs of exactly 4,096 tokens and generated exactly
1,024 tokens per successful request.

BS=16
RATE=inf
RESULT_DIR=/home/dist/zuoyu/dsv4-bench/musa3509-v026-pr174-774e0572-20260811/mtp4-full-car-v2/bs${BS}

vllm bench serve \
  --backend openai-chat \
  --base-url "http://127.0.0.1:${PORT}" \
  --endpoint /v1/chat/completions \
  --model "$SERVED_MODEL" \
  --tokenizer "$MODEL" \
  --dataset-name random \
  --seed 0 \
  --random-input-len 4096 \
  --random-output-len 1024 \
  --random-range-ratio 0.0 \
  --ignore-eos \
  --num-prompts "$BS" \
  --num-warmups 5 \
  --request-rate "$RATE" \
  --max-concurrency "$BS" \
  --temperature 0 \
  --save-result \
  --save-detailed \
  --result-dir "$RESULT_DIR"

No-MTP results

All requests completed successfully with zero failures and the exact requested
output length.

Batch	Runs	Output throughput (tok/s)	Mean TTFT (ms)	Mean TPOT (ms)	Completed
1	3	57.10 mean (56.87 / 57.09 / 57.34)	1,208.97 mean	16.250 mean	1/1 each
4	1	149.71	3,465.90	23.087	4/4
16	1	304.70	11,550.29	41.135	16/16
64	1	387.94	40,098.37	124.906	64/64

[Joey-gvwal]

While validating MTP, I found that the acceptance rate with multiple batches is significantly lower than in v0.24.0-dev. This requires further investigation.

[Joey-gvwal]

After fixing the MTP issue, the acceptance rate is now stable：

Batch	Runs	Output throughput (tok/s)	Mean TTFT (ms)	Mean TPOT (ms)	Acceptance rate	Acceptance length	Completed
1	1	91.36	1,148.70	10.778	54.43%	3.18	1/1
4	1	130.50	3,649.32	26.127	55.94%	3.24	4/4
16	1	288.49	11,825.13	43.565	58.93%	3.36	16/16
64	1	372.52	54,795.50	167.917	57.06%	3.28	64/64