[Bugfix] Allocate MLA attention output buffer when None during profiling pass

[psakhamo]

[Bugfix] vLLM profiling run (_dummy_run) calls forward_impl without an
output buffer — it only needs shape info, not values. Allocate a zero
buffer when output is None instead of asserting.
Validated: GLM-5.1-FP8 on 8x MI300X (gfx942), TP=8 — server startup PASS

Purpose

Fixes a hard AssertionError crash during vLLM server startup on models
using MLAAttention.

The vLLM profiling warmup pass (_dummy_run, called from
determine_num_available_blocks) invokes MLAAttention.forward() without
providing an output tensor buffer — it only needs to measure GPU memory
consumption and warm up code paths; it does not need computed values.

forward_impl contains this assertion on the output argument:

assert output is not None, "Output tensor must be provided."

This fires unconditionally during the profiling pass because forward() at
line 573 allocates an output buffer, but the profiling pass path does not
call forward() — it calls forward_impl directly without one, causing:

AssertionError: Output tensor must be provided.

The fix replaces the hard assert with a conditional allocation: if output
is None, a correctly shaped zero tensor is allocated. During real inference,
callers always provide a pre-allocated buffer, so this branch is never taken
outside of startup profiling.

Affected: all models using MLAAttention under vLLM v0.17.1 on any hardware
where the profiling pass reaches forward_impl without a pre-allocated output
(observed on AMD MI300X TP=8, reproducible at startup before KV cache init).

Test Plan

# Launch vLLM server — previously crashed immediately at startup
docker run --device /dev/kfd --device /dev/dri \
  --group-add video --group-add render \
  --ipc host --network host \
  -e VLLM_ROCM_USE_AITER=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

# Confirm server came up (previously never reached this point)
curl -s http://localhost:8080/health && echo "SERVER UP"

# Basic inference to confirm forward pass works end-to-end
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}'

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.

Before fix:
AssertionError: Output tensor must be provided.
File ".../vllm/model_executor/layers/attention/mla_attention.py", line 618
assert output is not None, "Output tensor must be provided."

Server exits during determine_num_available_blocks before KV cache is
allocated. GPU memory profiling never completes.

After fix:

• Server startup: ✅ PASS — profiling pass completes, KV cache allocated
• Uvicorn listening on port 8080: ✅ PASS
• Basic inference ("What is 2+2?"): ✅ PASS — correct response returned
• Zero buffer is discarded immediately after profiling; no impact on inference

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR — fixes AssertionError crash during profiling warmup in MLAAttention
• The test plan — server launch + health check + inference curl provided
• The test results — before crash traceback + after pass confirmation
• (Optional) Documentation update — not required for this bugfix

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