feat: shard hy4 linear_gate across PCP ranks

[dyldu]

HYV4 PCP Linear Gate 分片支持

1. 背景

在 HYV4 模型启用 PCP 时，linear_gate 默认仍使用完整输入维度权重。对于长序列和 PCP 多卡场景，这会导致每张卡保留较大的权重和中间激活，增加显存占用，并限制可用 KV Cache 空间。
本次修改将 linear_gate 按 PCP 维度进行切分，使其与 PCP 的输入 token 分片方式匹配。

2. 主要修改

1. linear_gate 按 PCP 进行 K 维切分
   新增 PCPShardedGateLinear：

• 按 PCP size 切分 linear_gate 的输入维度；
• 每个 PCP rank 只加载对应的局部权重；
• 使用 all-to-all 分发输入；
• 各 rank 执行局部矩阵乘；
• 使用 reduce-scatter 汇总并重新分发结果。
  当 PCP=8 时，每个 rank 的 linear_gate 输入权重维度为原来的 1/8。

2. 通过单一环境变量控制
   仅保留以下环境变量：
   export VLLM_HCU_ENABLE_LINEAR_GATE_PCP_SHARD=1
   行为说明：

• 设置为 1、true、yes 或 on：启用 PCP linear_gate 分片；
• 未设置或设置为 0：使用原始 linear_gate 实现。
  已移除旧的 mode 和 overlap 相关环境变量。

3. 放宽 HYV4 PCP 的 PP 限制
   更新 patch_vllm_config.py：

• 不再强制要求 PP=2；
• PP=1 配置可以继续使用 HYV4 PCP；
• 保留必要的 MTP、PCP、DeepEP、DeepGEMM 和 FP8 KV Cache 配置校验。

3. 精度结果

[hygon-ai-ai-reviewer]

AI Review

⚪ 审查未完成

部分内容未完成审查，当前结论不代表全部改动。

变更概览

本次变更主要包含：

• 新增环境开关 VLLM_HCU_ENABLE_LINEAR_GATE_PCP_SHARD 与 PCPShardedGateLinear 类，在 PCP 维度按输入维 K 分片 gated MLA 的 linear_gate 权重（要求 TP=1 且 PCP>1），并加入分片可整除、量化支持与权重元素数的校验。
• 为 HY V4 attention 的 linear_gate 增加按 PCP rank 的输入列分片：权重加载时从完整检查点权重切出本 rank 的输入列，前向计算仅取本地 K 切片并输出完整 gate 维度结果，同时按 PCP 配置记录一次性分片日志。
• 在 attention.py 中新增 PCPShardedGateLinear：启用开关时将 hy_v4 注意力的 linear_gate 替换为跨 PCP 分片实现，经 all_to_all 分发输入切片、本地部分 GEMM 后 reduce_scatter 聚合；未启用时沿用原 ColumnParallelLinear。
• 本批将 _require_hyv4_pcp_mtp_contract 中对 PP2 加精确 native MTP3 的契约校验整段注释，该函数不再对不匹配的 speculative 配置抛出 ValueError，函数主体变为无实际检查。

文件审查摘要

文件	变更	审查结果
vllm_hcu/models/hy_v4/attention.py	修改 · +180/-8	—
vllm_hcu/patch/platform/core_fix/patch_vllm_config.py	修改 · +8/-8	—

审查信息

• 变更统计：2 个文件，+188/-16。
• 覆盖情况：共 2 个文件，已完整审查 1 个，部分审查 1 个。
• 候选问题：1 项；证据复核过滤：1 项；发布前敏感信息保护：0 项。
• 未完成原因：候选复核：模型调用失败（1 批）
• 本服务只审查 GitHub 提供的 PR diff，未执行代码或重跑测试；结论仍需维护者核验。

[hygon-ai-ai-reviewer]

AI CI 失败分析

工作流：HCU PR CI
状态：失败

总结

• Static gate and test selection：日志直接证明 8 个用例失败：4 个因校验未抛出 ValueError（DID NOT RAISE），4 个抛出 PatchCompatibilityError（L634，位于 patch_vllm_config.py:42，其前因为 L403 的 draft_model_config AttributeError）。与 diff 中注释掉 _require_hyv4_pcp_mtp_contract 精确 MTP3 校验直接对应。

• ci-gate：尚未定位可核对的直接异常。

与本次改动的关系

相关：Run SHA 与 PR SHA 相同；报错文件正是被修改文件，且修改点恰为被注释的 ValueError raise；8 个失败用例均针对 MTP3/PCP 契约，日志中未抛 ValueError 与被注释的校验构成可核对触发路径。

建议处理

1. 最小方向：确认注释该校验是否有意变更。若非有意，恢复 _require_hyv4_pcp_mtp_contract 的契约校验；若有意放行，需同步修订这些契约测试，并核对后续 draft_model_config 校验的异常类型与调用顺序，再由人工核验。

[wics1224]

1. P1：PP1 放行方式扩大了配置支持面

PR 描述中的意图是“不再强制 PP=2”，但当前实现直接注释了整个 guard：

• 删除了 pipeline_parallel_size == 2 限制，这是预期变化。
• 同时删除了 num_speculative_tokens == 3 限制，这不是单纯的 PP 放宽。
• method == "mtp" 虽然后面还有一层检查，但 HYV4 的 speculative token 数没有后续兜底；后面的深度检查只覆盖 not is_hyv4。
• 函数注释、错误信息以及后续多个检查仍然明确写着 “MTP3”，与实际可接受配置已经不一致。
• 原有测试也明确将 HYV4 的 speculative token 数 1、2、4 视为非法邻近配置。

这个问题已经被当前 GitHub CI 实际触发：Static gate and test selection 失败，自动分析显示共 8 个相关失败，其中包括：

• 4 个预期 ValueError 但没有抛出的用例；
• 4 个因继续访问缺失的 draft_model_config，最终抛出 PatchCompatibilityError 的用例。

建议将契约改成两个独立维度：

if (
    speculative.method != "mtp"
    or speculative.num_speculative_tokens != 3
):
    raise ValueError(
        "HYV4 PCP speculative decoding requires checkpoint-native MTP3."
    )

PP1/PP2 的支持范围则由拓扑校验单独表达。测试侧应该：

• 新增一个完整、合法的 PP1 + PCP + native MTP3 正向用例；
• 从旧的 “neighbors stay closed” 中只移除 PP1；
• 保留 speculative token 数 1、2、4 的反向用例；
• 继续拒绝没有 native draft 配置的输入，并确保异常类型是面向用户的 ValueError。

2. P2：新实现降低权重常驻显存，但显著放大运行时峰值激活

分片权重本身是有效的：每个 PCP rank 只保留 [output_size, input_size / PCP]。但前向路径为了完成 K-shard 归约，在每张卡上构造了完整 global-token 的 partial 输出：

received: [PCP, local_tokens, hidden_size / PCP]
partial:  [PCP × local_tokens, gate_output_size]
output:   [local_tokens, gate_output_size]

因此相对于旧实现，临时 gate 输出从：

local_tokens × gate_output_size

变成：

PCP × local_tokens × gate_output_size

此外还有以下同时存活的张量：

• 原始 input_；
• transpose 后生成的连续 send；
• received；
• partial；
• reduce-scatter 新分配的最终输出。

以仓库中的 HYV4 配置计算：

• hidden_size = 2816
• num_attention_heads = 32
• v_head_dim = 256
• elementwise gate 输出维度为 32 × 256 = 8192
• PCP=8、BF16 时：

每 rank 的 local tokens | 旧 gate 输出 | 新 partial -- | -- | -- 4,096 | 64 MiB | 512 MiB 16,384 | 256 MiB | 2 GiB 32,768 | 512 MiB | 4 GiB

这不一定会在所有批次中抵消跨全部层累计节省的权重显存，但当前实现没有限制单步 token 数；在 PR 最关注的长上下文场景里，峰值显存并非单调改善。尤其 vLLM 的 KV Cache 容量取决于模型运行时峰值，额外的几 GiB 临时张量可能直接减少可分配 block 数。

可选修复方向：

• 按 token block 分块计算 partial 并逐块 reduce-scatter，最后拼接本 rank 的输出；
• 使用支持 HCU/RCCL 的融合 GEMM + reduce-scatter；
• 如果短期只能保留当前实现，至少依据可证明安全的显存预算限制 local_tokens，并在最大支持的 max_num_batched_tokens 下提供峰值显存 A/B，而不只是精度截图。

1. P1：PP1 放行方式扩大了配置支持面 PR 描述中的意图是“不再强制 PP=2”，但当前实现直接注释了整个 guard： - 删除了 pipeline_parallel_size == 2 限制，这是预期变化。 - 同时删除了 num_speculative_tokens == 3 限制，这不是单纯的 PP 放宽。 - method == "mtp" 虽然后面还有一层检查，但 HYV4 的 speculative token 数没有后续兜底；后面的深度检查只覆盖 not is_hyv4。 - 函数注释、错误信息以及后续多个检查仍然明确写着 “MTP3”，与实际可接受配置已经不一致。 - 原有测试也明确将 HYV4 的 speculative token 数 1、2、4 视为非法邻近配置。 这个问题已经被当前 GitHub CI 实际触发：[Static gate and test selection](https://github.com/HYGON-AI/vllm-plugin-das/actions/runs/34971077619/job/104387255233) 失败，自动分析显示共 8 个相关失败，其中包括： - 4 个预期 ValueError 但没有抛出的用例； - 4 个因继续访问缺失的 draft_model_config，最终抛出 PatchCompatibilityError 的用例。 建议将契约改成两个独立维度： if ( speculative.method != "mtp" or speculative.num_speculative_tokens != 3 ): raise ValueError( "HYV4 PCP speculative decoding requires checkpoint-native MTP3." ) PP1/PP2 的支持范围则由拓扑校验单独表达。测试侧应该： - 新增一个完整、合法的 PP1 + PCP + native MTP3 正向用例； - 从旧的 “neighbors stay closed” 中只移除 PP1； - 保留 speculative token 数 1、2、4 的反向用例； - 继续拒绝没有 native draft 配置的输入，并确保异常类型是面向用户的 ValueError。 2. P2：新实现降低权重常驻显存，但显著放大运行时峰值激活 分片权重本身是有效的：每个 PCP rank 只保留 [output_size, input_size / PCP]。但前向路径为了完成 K-shard 归约，在每张卡上构造了完整 global-token 的 partial 输出： received: [PCP, local_tokens, hidden_size / PCP] partial: [PCP × local_tokens, gate_output_size] output: [local_tokens, gate_output_size] 因此相对于旧实现，临时 gate 输出从： local_tokens × gate_output_size 变成： PCP × local_tokens × gate_output_size 此外还有以下同时存活的张量： - 原始 input_； - transpose 后生成的连续 send； - received； - partial； - reduce-scatter 新分配的最终输出。 以仓库中的 HYV4 配置计算： - hidden_size = 2816 - num_attention_heads = 32 - v_head_dim = 256 - elementwise gate 输出维度为 32 × 256 = 8192 - PCP=8、BF16 时： 每 rank 的 local tokens 旧 gate 输出 新 partial 4,096 64 MiB 512 MiB 16,384 256 MiB 2 GiB 32,768 512 MiB 4 GiB

这不一定会在所有批次中抵消跨全部层累计节省的权重显存，但当前实现没有限制单步 token 数；在 PR 最关注的长上下文场景里，峰值显存并非单调改善。尤其 vLLM 的 KV Cache 容量取决于模型运行时峰值，额外的几 GiB 临时张量可能直接减少可分配 block 数。
可选修复方向：

• 按 token block 分块计算 partial 并逐块 reduce-scatter，最后拼接本 rank 的输出；
• 使用支持 HCU/RCCL 的融合 GEMM + reduce-scatter；
• 如果短期只能保留当前实现，至少依据可证明安全的显存预算限制 local_tokens，并在最大支持的 max_num_batched_tokens 下提供峰值显存 A/B，而不只是精度截图。