Cap unified-attention num_stages to fit gfx11 LDS in 3D decode path

[zohaibm-amd]

Problem

On gfx11 (RDNA3/3.5, 64KB LDS), the Triton unified-attention 3D decode path hard-codes num_stages=3 for GQA models. For wide attention heads (e.g. gemma-4 full-attention layers, head_size=512), the per-stage K/V software-pipeline buffers exceed the 64KB shared-memory limit and Triton aborts at engine startup with OutOfResources: shared memory.

The 2D path already caps num_stages to fit the LDS budget; the 3D decode path had no equivalent check.

Fix

Extract the existing 2D Navi LDS cap into a shared helper and apply it to the 3D decode path. Shapes that already fit keep their current num_stages; previously-overflowing wide-head configs drop to fewer pipeline stages and compile.

Validation

• Affected model now starts and serves instead of crashing at engine init.
• Narrow-head and 2D num_stages selections are unchanged.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[mgehre-amd]

nit: I wonder whether we can query triton for this number instead of hardcoding it here; that would also make it applicable for more devices.
We can also solve this question when upstreaming, no need on this PR.