[ROCm][Bugfix] Flatten 3D attn_out before o_proj in MLA forward pass

[psakhamo]

[ROCm] Bug fix: ROCMAiterMLASparseBackend returns attn_out as 3D
[num_tokens, num_heads, v_head_dim] instead of expected 2D
[num_tokens, num_heads * v_head_dim]. Flatten before o_proj.
Validated: GLM-5.1-FP8 on 8x MI300X (gfx942), TP=8 - PASS

Purpose

Fixes a shape mismatch crash introduced by the ROCm AITER sparse MLA
backend (ROCMAiterMLASparseBackend). Unlike the CUDA path, this backend
returns attention output as a 3D tensor [num_tokens, num_heads, v_head_dim]
instead of the expected 2D [num_tokens, num_heads * v_head_dim].

Without this fix, the unflattened tensor is passed directly into o_proj,
which expects a 2D input, causing a shape mismatch that surfaces as a
ValueError several layers later in DeepseekV2MoE.forward():

num_tokens, hidden_dim = hidden_states.shape

ValueError: too many values to unpack (expected 2)

This was root-caused via runtime shape probes on 8× MI300X (gfx942) at TP=8,
where hidden_states entering every MoE layer was confirmed to be
(num_tokens, 8, 6144) — 3D with num_heads=8 as the middle dimension.

The fix adds a defensive reshape before o_proj. It is a strict no-op for
2D inputs (the normal CUDA/non-ROCm path) and has no effect on model weights,
KV cache, or any other tensor. Affected models: all models using MLA with the
ROCm AITER sparse backend (DeepSeek-V2/V3/R1, GLM-5).

Closes vllm-project#39168 (related — ROCm MLA head-repeat padding)

Test Plan

# Start vLLM server with GLM-5.1-FP8 on 8x MI300X, TP=8
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 \
    --tool-call-parser glm47 \
    --reasoning-parser glm45 \
    --enable-auto-tool-choice \
    --enforce-eager --block-size 1

# Layer 1: basic inference
curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"glm-5-1-fp8",
       "messages":[{"role":"user","content":"What is 2+2?"}],
       "max_tokens":50}'

# Layer 2: streaming tool-call routing
curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"glm-5-1-fp8","stream":true,
       "messages":[{"role":"system","content":"You are a coding assistant."},
                   {"role":"user","content":"List files in current directory."}],
       "tools":[{"type":"function","function":{"name":"bash",
         "parameters":{"type":"object",
           "properties":{"command":{"type":"string"}},"required":["command"]}}}],
       "tool_choice":"auto","max_tokens":200}'

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:
ValueError: too many values to unpack (expected 2)
File ".../vllm/model_executor/models/deepseek_v2.py", line 412
num_tokens, hidden_dim = hidden_states.shape

Server crashes immediately after attention forward pass.

After fix:

• Layer 1 (basic inference): ✅ PASS — "content":"2+2 equals 4" returned correctly
• Layer 2 (streaming tool-call): ✅ PASS — finish_reason="tool_calls" confirmed,
  tool_calls array non-empty, bash(ls) function call routed correctly

No regression observed on standard inference path (2D attn_out case unchanged).

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR — fixes ROCm AITER sparse MLA backend 3D attn_out shape mismatch
• The test plan — server launch command + Layer 1 and Layer 2 curl tests provided
• The test results — before/after error + pass confirmation on 8x MI300X
• (Optional) Documentation update — not required for this bugfix

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