[Bugfix][MUSA] Support Gemma4 heterogeneous attention · Pull Request !1 · Moore Threads Corporation/vLLM-MT

## 目的

为 MUSA 平台补齐 Gemma4 异构注意力支持，解决以下运行时阻断：

- MUSA 伪 SM100 capability 导致 FA4 支持被误判，Gemma4 被错误路由到不支持 head_dim=512 的 MUSA FlashAttention v3。
- D0930 镜像内 Triton 3.2.0 缺少 make_tensor_descriptor，现有 unified attention kernel 在 JIT cache-key 阶段即失败。
- MUSA 伪 SM100 错误触发 CUDA B200 的大 tile 调优，prefill 所需共享内存超过 S4000 硬件上限。

本 PR 采用最小范围修复：

1. MUSA 平台明确禁用 FA4 capability 判断，保留 CUDA/ROCm/XPU 原有行为。
2. 将 Triton unified attention 拆分为 TD 与非 TD 两个独立 JIT 入口，确保非 TD 路径的词法依赖闭包不包含 make_tensor_descriptor。
3. 将 B200 大 tile 调优显式限制为 CUDA 平台。
4. 增加 MUSA FA4 误判、TD 依赖隔离、双入口选择及 head_dim=512 覆盖测试。

当前官方仓库 Pull Requests 为 0，未发现同目标重复 PR。

## 测试计划

- Python 语法编译与 git diff 检查。
- FA4 capability、TD 依赖隔离、双入口选择离线回归。
- D0930 镜像 + S4000 上运行 Triton 256/512 shape 探针。
- Gemma4 E4B-it BF16、TP1、eager 模式启动和 Completion/Chat 功能验证。

## 测试结果

- 离线测试：24/24 PASS；py_compile PASS；git diff --check PASS。
- Triton 运行时：256/512 decode、prefill、GQA、block-size 16/32、sliding window、同进程异构 head dim 全部通过；输出 finite 且数值对照通过。
- 模型级验证：Gemma4 E4B-it BF16 TP1 eager 服务正常 ready，实际选择 TRITON_ATTN；英文 Chat、中文 Chat、官方 chat template Completion 均通过；无 HTTP 500、EngineDead、NaN、OOM、MUSA 设备错误或 Triton 编译错误。

## 范围与限制

- 已验证范围：D0930 镜像、MTT S4000、Gemma4 E4B-it、BF16、TP1、eager。
- 尚未覆盖：26B-A4B、TP>1、graph、量化和性能结论。
- Gemma4-it 裸 Completion 若不使用官方 chat template 可能出现重复输出；推荐使用 Chat Completion 或先应用官方模板。这属于提示词使用约束，不是 attention kernel 故障。

## AI 辅助说明

本 PR 的问题分析、测试设计和补丁审查使用了 Codex/Trae 辅助；所有代码变更与测试结果均由提交者人工复核，并在 MUSA 实机环境验证。

类型	指派人员	状态
审查	Flecky-MT 已审查通过	已完成 (1/0人)
测试	Hobbey-MT 已测试通过 Flecky-MT 已测试通过	已完成 (2/1人)

Moore Threads Corporation / vLLM-MT .gitee-modal { width: 500px !important; }

搜索帮助

Moore Threads Corporation / vLLM-MT