[ROCm][INTERNAL][WIP] Cap profiling token count to fix SP all-gather OOM for MLA+DSA at TP>1

[psakhamo]

[ROCm][WIP] When fused_qkv_a_proj uses sequence-parallel mode (disable_tp=True),
the profiling run all-gathers hidden_states across TP ranks, producing tensors
with tp_size * local_tokens. Cap profiling token count to avoid shape mismatch
OOM in MoE layers during startup.
Validated: GLM-5.1-FP8 on 8x MI300X (gfx942), TP=8 — PASS (internal review only)

⚠️ INTERNAL REVIEW — Do not merge to vllm-project/vllm without AMD team sign-off

This PR is filed against psakhamo/vllm for internal review. The fix works but
the token cap strategy (hardcoded 64) needs discussion before upstream submission.
See "Open Questions" below.

Purpose

Fixes a cascade of OOM / shape mismatch errors during the vLLM profiling pass
(determine_num_available_blocks) for models using MLA + DSA attention at TP > 1.

Root cause: DeepSeekV2FusedQkvAProj is initialized with disable_tp=True
(sequence-parallel mode). During the profiling pass, fused_qkv_a_proj performs
an all-gather of hidden_states across all TP ranks, expanding the token
dimension from local_tokens (e.g. 8,192) to tp_size * local_tokens
(e.g. 65,536 at TP=8). All surrounding tensors — positions, k_pe,
residual — remain at local token count.

This mismatch triggers a cascade of errors:

Rotary embedding shape mismatch

RuntimeError: query, key and positions must have the same number of tokens
(got 65536, 8192, 8192)

MoE GEMM OOM (65536 * 64 = 4,194,304 token-expert pairs)

torch.OutOfMemoryError: HIP out of memory. Tried to allocate 6.00 GiB

MoE hidden_states unpack (3D tensor from attn_out flatten issue)

ValueError: too many values to unpack (expected 2)

Fix: Cap the token count passed to _dummy_run during the profiling pass
to 64. The profiling pass only needs to measure GPU memory and warm up code
paths — it does not need max_num_tokens. At 64 tokens, the all-gather
produces 512 tokens (64 × 8 GPUs), well within memory and matching no shape
constraint.

Affected models: GLM-5 (GlmMoeDsaForCausalLM), DeepSeek-V3, DeepSeek-R1
— any model using DeepSeekV2FusedQkvAProj with disable_tp=True at TP > 1.

Open Questions for Reviewers

1. Is 64 the right cap? Should this be derived from TP size or made configurable?
2. Scope: Does capping tokens affect cudagraph capture or any other warmup
   that expects max_num_tokens specifically?
3. Detection: Should we detect the SP all-gather path explicitly
   (e.g. check fused_qkv_a_proj.disable_tp) rather than capping globally?
4. Upstream readiness: Is there a cleaner architectural fix in
   DeepSeekV2FusedQkvAProj itself (skip all-gather when is_profile=True)?

Test Plan

# Launch server — previously crashed in profiling pass before KV cache init
docker run --device /dev/kfd --device /dev/dri \
  --group-add video --group-add render \
  --ipc host --network host \
  -e VLLM_ROCM_USE_AITER=1 \
  -e HIP_FORCE_DEV_KERNARG=1 \
  vllm-glm5-full-fix:local \
  zai-org/GLM-5.1-FP8 \
    --tensor-parallel-size 8 \
    --gpu-memory-utilization 0.85 \
    --port 8080 --enforce-eager --block-size 1

curl -s http://localhost:8080/health && echo "SERVER UP"

Test Result

Tested on: GLM-5.1-FP8 (GlmMoeDsaForCausalLM, 744B MoE) on 8× AMD Instinct
MI300X (gfx942), TP=8, vLLM v0.17.1 ROCm build, AITER enabled.

Before fix:
RuntimeError: query, key and positions must have the same number of tokens
(got 65536, 8192, 8192)
File ".../vllm/model_executor/layers/mla.py", line 164, in forward

Profiling pass crashes, KV cache never allocated, server exits.

After fix:

• Profiling pass: ✅ PASS — completes in ~177s (dominated by model load)
• AITER kernel JIT compilation: ✅ PASS — all kernels compiled and cached
• Server startup: ✅ PASS — Uvicorn listening on port 8080
• Layer 1 inference: ✅ PASS
• Layer 2 streaming tool-call: ✅ PASS — finish_reason="tool_calls" confirmed

JIRA: GPUAI-6471 (P1 Cohere/OCI engagement, MI300X GLM-5 enablement)

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR — fixes SP all-gather OOM/shape crash in MLA+DSA profiling pass
• The test plan — server launch + health check + inference validation
• The test results — before crash traceback + after full pass confirmation
• (Optional) Documentation update — not required; open questions logged above

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing