[feat]Qwen4Exp PLE INT8 量化、CPU offload/UVA 零拷贝与 prefetch stream 合入

[guanyu1998]

变更概要

将 guanyu-hcu-upstream-54371-56208-v0.28.1-dev 的 8 个 commit 合入 v0.28.1-dev：

1. a9925a4 feat(qwen4_exp): 新增 Qwen4Exp PLE embedding compressed-tensors INT8 支持
2. 987b7d1 fix(qwen4_exp): PLE INT8 查表即反量化，修复 custom allreduce dtype 报错
3. af8b8f1 feat(qwen4_exp_ple): 新增 PLE INT8 CPU offload 与 UVA 零拷贝支持
4. 1e8c1bd feat(qwen4_exp_ple): 新增 PLE INT8 UVA 查表与模型计算重叠的 prefetch stream
5. ffc81ca feat(qwen4_exp_ple): PLE INT8 UVA 查表与模型计算重叠的 prefetch stream 优化及测试
6. 60cbd30 feat(qwen4_exp): engram 模块与 PLE INT8 offload 相关调整及测试
7. 8ec6851 feat(qwen4_exp): 新增 PLE ETP 与 HCU Engram 配置支持
8. 604b381 feat(qwen4_exp_ple): 兼容 support_torch_compile 包装的模型构造函数

改动内容

PLE INT8 量化支持与修复（commit 1 & 2）

• 上游 vLLM 快照中 Qwen4Exp PLE 表仍为未量化的 PLEVocabParallelEmbedding，compressed-tensors INT8（per-channel）checkpoint 无法加载与推理；新增 patch_qwen4_exp_ple_int8 补丁，worker 启动时替换模块内 PLE 存储类与 shard loader，INT8 权重 + 每 vocab 行一个 BF16 scale，TP 下 scale 走 all_reduce
• 修复：原实现基类 forward 后才反量化，TP>1 时 int8 查表结果先进入 custom allreduce（仅支持 fp32/fp16/bf16），cudagraph capture 阶段报错；改为查表后立即按 scale 反量化，直接返回浮点张量再走基类 all-reduce 路径

PLE INT8 CPU offload 与 UVA 零拷贝（commit 3 & 6）

• PLE INT8 ngram 表常驻 GPU，每 TP rank 约 12.36 GiB；新增 VLLM_HCU_PLE_CPU_OFFLOAD=1 将表 offload 到 CPU pinned 内存
• UVA 优先：注册 get_cuda_view_from_cpu_tensor 算子，embedding lookup 零拷贝直读 host 内存，无 D2H 同步，兼容 CUDA graph capture；算子不可用时自动降级为显式 host gather + H2D staging
• UVA view 惰性创建并缓存（绕开 process_weights_after_loading 阶段权重临时在 GPU 的问题）

prefetch stream 计算重叠（commit 4 & 5 & 8）

• UVA 零拷贝查表在当前 stream 内串行执行，无法被前序 decoder layer 计算掩盖；新增 VLLM_HCU_PLE_PREFETCH_STREAM=True 开关，在 decoder layer 前向之前于独立 stream 提前发起 UVA 查表
• custom op schema 扩展为带 prefetch workspace buffer 的固定签名，piecewise cudagraph capture 前同步 prefetch stream 保证捕获正确性
• 新增 _is_torch_compile_init 识别 @support_torch_compile 包装的模型构造函数，hcu_init 改为 (*args, vllm_config=None, prefix="", **kwargs) 透传形式，避免 patch 应用失败

ETP 与 HCU Engram 配置（commit 7 & 6）

• 上游 Engram 校验仅接受 CUDA，HCU 显式 cpu_offload 配置被拒绝；新增 HCU CUDA-like Engram 兼容门禁并保留上游模型约束
• PLE 接入 ETP group、跨 DP ID gather、padding 与本地结果切片，PLE all-reduce 使用 ETP group，prefetch buffer 扩展并适配 CUDA Graph 路径

影响的模块

• vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_int8.py：新增（INT8 存储/loader/量化方法，UVA/offload 工厂与降级逻辑）
• vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_prefetch.py：新增（prefetch stream 注册、layer 前置查表、torch_compile 兼容）
• vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_cudagraph.py：新增（capture 前 prefetch stream join）
• vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_conv.py：修改（custom op 扩展 prefetch workspace 参数并统一 schema）
• vllm_hcu/models/qwen4_exp/amd/ple_layer.py：新增（HCU PLE layer 替换实现，含跨 DP ETP 查表与 prefetch workspace）
• vllm_hcu/models/qwen4_exp/engram.py：新增（engram 模块）
• vllm_hcu/patch/platform/core_fix/patch_engram_config.py：新增（HCU Engram 配置兼容门禁）
• vllm_hcu/patch/module_exchange.py：修改（按 VLLM_HCU_PLE_PREFETCH_STREAM 条件注册 ple_layer 替换）
• vllm_hcu/platforms/envs.py：新增环境变量 VLLM_HCU_PLE_CPU_OFFLOAD、VLLM_HCU_PLE_PREFETCH_STREAM
• vllm_hcu/platforms/hcu.py：修改（注册 UVA 零拷贝算子）
• vllm_hcu/patch/worker/__init__.py、patch/worker/core_fix/__init__.py、patch/platform/core_fix/__init__.py：注册新 patch 回调
• 新增测试：tests/runtime_patch/test_qwen4_exp_ple_int8_offload.py、tests/runtime_patch/test_qwen4_exp_ple_prefetch.py；更新 test_qwen4_exp_ple_conv.py、test_module_exchange.py

验证方式

• 新增/更新单测：PLE INT8 offload、prefetch、conv、module_exchange 相关测试
• 真机 8×BW100 HCU、TP=4：eager 与 CUDA graph（FULL_AND_PIECEWISE）模式端到端通过，capture 无报错，推理数值正常，每卡省约 12.36 GiB
• TP>1 INT8 PLE 推理数值对齐回归（对比未量化路径）
• VLLM_HCU_PLE_PREFETCH_STREAM=True/False 两种模式下的 cudagraph capture 与解码吞吐对比
• 跨 DP + ETP 场景端到端推理回归

风险评估

• 中：prefetch/offload 均为 opt-in（默认关闭），默认路径不受影响；但 PLE INT8 存储类替换了模块内原始 PLE 类，需关注非 INT8 checkpoint 加载路径的兼容性
• 已知问题（本分支未修复）：M-RoPE RopeState.prefill_positions 按 (max_num_reqs × num_dims × max_model_len) 申请 pinned 内存（TP4 约 1.5 GiB），可能触发 hipErrorOutOfMemory，当前只能靠 --max-num-seqs / --max-model-len 缓解，待后续跟进

[hygon-ai-ai-reviewer]

AI Review

本次变更中未发现有明确证据的正确性或安全性问题。

[wics1224]

Review findings

1. [P1] embedding_across_dp 只有同时启用 PLE prefetch 才会生效
   module_exchange.py:106 仅在 VLLM_HCU_PLE_PREFETCH_STREAM 为真时注册 _QWEN4_EXP_PLE 替换模块：
   if _ple_prefetch_requested():
   return (*_ALL_GROUPS, _QWEN4_EXP_PLE)
   但这个替换模块不只是实现 prefetch，还包含本 PR 的 ETP/cross-DP PLE 逻辑，例如 get_etp_group()、parallel_group=etp_group、跨 DP 的 ID gather/reduce/slice，以及按 TP×DP 切分 PLE embedding。
   因此，用户只设置：
   --engram-config.embedding_across_dp=true
   而不打开 prefetch 时，参数会被新的 EngramConfig patch 接受，但真正的 PLE 层不会被替换，仍走上游 TP-only 实现。最终表现是配置被静默忽略，PLE 权重依然在不同 DP rank 之间复制，既没有获得预期的显存收益，也没有获得预期的跨 DP 拓扑。
   现有测试 test_qwen4_exp_ple_exchange_is_strictly_opt_in_and_lazy 反而固化了这个错误的绑定关系。
   建议将 PLE replacement 的注册条件改为：
   • prefetch 开启，或
   • embedding_across_dp 开启；
     更稳妥的方式是始终注册兼容版本的 PLE 层，并确保 TP-only、不开 prefetch 时行为与上游一致。也可以把 ETP 与 prefetch 拆成独立 patch，避免两个本应独立的功能相互作为开关。
2. [P1] CPU offload 在模型加载阶段仍会把完整 PLE 权重临时搬到 GPU，可能在启动时 OOM
   patch_qwen4_exp_ple_int8.py:229 中的 HcuQwen4ExpPLEInt8UVAEmbeddingMethod，以及下面的 H2D offload method，都没有覆盖 QuantizeMethodBase.requires_device_loading。
   目标 vLLM 版本里该属性默认是 True，所以在调用 process_weights_after_loading() 前，device_loading_context 会先把 CPU parameter 临时移动到 accelerator，处理结束后才移回 CPU。代码里的注释也承认参数在这里可能暂时位于 accelerator。
   对这个 PR 所描述的约 12.36 GiB PLE 表来说，结果是：
   • 稳态运行时显存确实能够下降；
   • 但模型加载阶段仍需要额外容纳完整 PLE 表；
   • 本来依赖 offload 才能装进显存的模型，可能在服务启动时直接 OOM。
     而这两个方法的 post-load 处理目前只是重建 view、检查布局和记录日志，并不需要把整个参数搬到 GPU。
     建议在 UVA 和 H2D offload quant method 上显式设置：
     requires_device_loading = False
     同时增加回归测试，确认调用 post-load hook 时 PLE 参数仍保持在 CPU，并且不会出现一次完整 CPU→GPU→CPU staging。这里应特别覆盖文档中的大 PLE shape，而不能只用很小的单元测试 tensor，因为小 tensor 无法暴露峰值显存问题。
3. [P1] 新增 patch 没有满足仓库的 patch-test coverage contract，质量门确定性失败
   新增的 patch_engram_config.py:18 声明了一个 patch adapter，但没有测试直接引用这个 patch。
   我在 PR merge tree 上执行仓库自己的检查，结果为：
   patch coverage: files=116 adapters=115 helpers=1 untested=1 invalid_contracts=0
   missing direct test reference: vllm_hcu.patch.platform.core_fix.patch_engram_config
   返回码是 1。同一检查在 base tree 上通过：
   patch coverage: files=112 adapters=111 helpers=1 untested=0 invalid_contracts=0
   所以这是本 PR 新引入的确定性回归，不是本地环境差异。
   建议补充一个直接导入或引用 vllm_hcu.patch.platform.core_fix.patch_engram_config 的 contract test，并至少验证：
   • HCU CUDA-like 平台允许目标 Qwen4Exp 模型使用 embedding_across_dp；
   • 非目标模型不会被意外放宽；
   • 不支持的平台仍保留上游拒绝行为；
   • patch 的注册和应用发生在预期生命周期内，而不只是通过 AST 检查上游类是否存在。
4. [P2] UVA op 缺失时自动选择的 H2D fallback 与默认 CUDA Graph 路径不兼容
   patch_qwen4_exp_ple_int8.py:490 在 UVA lookup op 不可用时，打印 warning 后自动退回 HcuQwen4ExpPLEInt8OffloadEmbeddingMethod。
   但该 fallback 在 lookup 内执行同步的 GPU→CPU ID 搬运、CPU index_select 和 CPU→GPU embedding 搬运；实现本身已经注明这条路径会阻碍 CUDA Graph capture。默认配置没有因此关闭 CUDA Graph，也没有在启动阶段提前拒绝该组合。
   这意味着“自动 fallback”对默认运行方式未必真正可用：服务可能一直初始化到 graph capture 阶段才失败，错误发生得晚且不容易和 UVA op 缺失关联。
   建议二选一：
   • 检测到只能走 H2D fallback 且启用了 CUDA Graph 时尽早报错，并明确提示使用 eager 模式；
   • 或由配置层可靠地强制该模型进入 eager 模式，并在日志中说明原因。
     仅打印 warning 不够，因为后续路径仍可能确定性失败。

[guanyu1998]

1. [P1] embedding_across_dp 只有同时启用 PLE prefetch 才会生效 ： 目前用不到DP，暂时对外开放使用

[guanyu1998]

PLE CPU offload 权重不再进入默认的 post-load CPU→设备→CPU 往返搬运流程。权重会始终保留在 pinned CPU memory 中，并在首次使用时延迟创建 UVA view。

• 为 patch_engram_config 补充了直接的 contract 和行为测试。
• 完全移除了 H2D fallback。启用 PLE CPU offload 但 UVA 不可用时，现在会在模型初始化阶段直接给出明确错误，不再延迟到 CUDA Graph capture 阶段失败。

验证结果：

• PLE INT8 offload 测试：13 passed
• PLE prefetch 测试：26 passed
• Platform core 测试：18 passed
• Patch coverage：untested=0，invalid_contracts=0

[hygon-ai-ai-reviewer]

AI Review

🟡 建议修改

发现有明确证据的问题，建议核对并处理。

已确认 1 个问题，其中 0 个已添加到对应代码行。

变更概览

本次变更主要包含：

• 新增 Qwen4Exp PLE 模块交换测试，验证注册仅在指定环境变量开启时生效且保持目标模块不被提前导入，规范模块已被导入时报错
• 新增 patch_engram_config 契约测试，覆盖 HCU 支持判定与包装合约、上游拒绝场景保留及签名/源码形状漂移检测；并更新 ngram 动态预处理 fake run 签名以校验新增的四个空 buffer 参数
• 新增针对 Qwen4Exp PLE INT8 CPU offload 路径的测试，覆盖 UVA/驻留方法选择、engram 配置与环境变量优先级、create_weights 的 CPU 分配以及加载后权重驻留 CPU 的行为验证。
• 新增 Qwen4Exp PLE INT8 offload 的测试，覆盖 UVA 方法与驻留方法的子类关系、UVA 零拷贝查表与 GPU 驻留反量化数值一致性，以及设备视图跨调用缓存复用。

文件审查摘要

文件	变更	审查结果
vllm_hcu/models/qwen4_exp/amd/ple_layer.py	新增 · +1461/-0	P2 × 1
tests/runtime_patch/test_qwen4_exp_ple_prefetch.py	新增 · +592/-0	—
vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_int8.py	新增 · +503/-0	—
vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_prefetch.py	新增 · +373/-0	—
tests/runtime_patch/test_qwen4_exp_ple_int8_offload.py	新增 · +329/-0	—
vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_cudagraph.py	新增 · +131/-0	—
vllm_hcu/patch/platform/core_fix/patch_engram_config.py	新增 · +128/-0	—
tests/runtime_patch/test_platform_core.py	修改 · +121/-1	—
vllm_hcu/patch/worker/core_fix/patch_qwen4_exp_ple_conv.py	修改 · +95/-8	—
vllm_hcu/patch/module_exchange.py	修改 · +35/-3	—
tests/patch/test_module_exchange.py	修改 · +31/-0	—
vllm_hcu/models/qwen4_exp/engram.py	新增 · +25/-0	—
tests/runtime_patch/test_qwen4_exp_ple_conv.py	修改 · +15/-1	—
vllm_hcu/platforms/envs.py	修改 · +15/-0	—
vllm_hcu/patch/worker/core_fix/init.py	修改 · +6/-0	—
vllm_hcu/models/qwen4_exp/init.py	新增 · +3/-0	—
vllm_hcu/models/qwen4_exp/amd/init.py	新增 · +3/-0	—
vllm_hcu/patch/worker/init.py	修改 · +3/-0	—
vllm_hcu/patch/platform/core_fix/init.py	修改 · +2/-0	—
vllm_hcu/platforms/hcu.py	修改 · +1/-0	—

无法定位到 diff 行的问题（1）

P2 · 填充行映射到槽 0 后 index_copy_ 重复索引可能覆盖真实请求状态

• 文件：vllm_hcu/models/qwen4_exp/amd/ple_layer.py（无法安全定位到当前 diff 行）
• 触发条件：批次前 num_prefills 行中存在 NULL_BLOCK_ID 填充行（代码注释自述该路径会处理填充行），且同一批次有真实请求占用缓存槽 0。
• 证据与影响：state_indices 先用 torch.where 把无效索引替换为 0，随后用 conv_state.index_copy_(0, state_indices, existing_state) 回写。填充行的 update_mask 为 False，只会把槽 0 的旧值再写回槽 0，与真实请求写入的新状态形成同索引竞争。 index_copy_ 对重复索引的结果在加速卡上不确定：真实请求位于槽 0 的新 conv 状态可能被旧值覆盖，后续解码沿用过期卷积状态，产生难以复现的输出错误。
• 修改建议：回写前先按 valid_state 压缩索引与数据（如 state_indices[valid_state] 与对应 next_state 子集）再执行唯一索引的 index_copy_，避免重复索引写入。

审查信息

• 变更统计：20 个文件，+3872/-13。
• 覆盖情况：共 20 个文件，已完整审查 20 个。
• 候选问题：1 项；证据复核过滤：0 项；发布前敏感信息保护：0 项。
• 本服务只审查 GitHub 提供的 PR diff，未执行代码或重跑测试；结论仍需维护者核验。

提交修复并推送新 commit 后，将自动审查 PR 的最新版本。

[wics1224]

[P1] embedding_across_dp 仍然只有开启 PLE prefetch 时才真正生效
问题仍位于 module_exchange.py:106：
def _enabled_groups() -> tuple[tuple[_Entry, ...], ...]:
if _ple_prefetch_requested():
return (*_ALL_GROUPS, _QWEN4_EXP_PLE)
return _ALL_GROUPS
register_qwen4_exp_ple_exchange() 也使用相同的条件：
if not _ple_prefetch_requested():
return ()
这意味着只有：
VLLM_HCU_PLE_PREFETCH_STREAM=1
时，HCU 的 vllm_hcu.models.qwen4_exp.amd.ple_layer 才会替换上游 AMD PLE 实现。
但这个 replacement 不只是 prefetch 实现，它还承载了本 PR 的全部 ETP/cross-DP 功能，包括：

• get_etp_group()；
• parallel_group=etp_group；
• etp_data_parallel_size；
• 跨 DP ID gather；
• 不同 token 数的 padding；
• ETP all-reduce；
• 每个 DP rank 的本地结果切片。
  因此下面这个完全合法、也符合配置语义的组合仍然不能工作：
  --engram-config.embedding_across_dp=true
  VLLM_HCU_PLE_PREFETCH_STREAM=0
  在这个组合下：

1. EngramConfig patch 会接受 HCU 上的配置；
2. vLLM 会根据 embedding_across_dp=true 创建更大的 ETP group；
3. 但 HCU PLE replacement 没有注册；
4. AMD 原始 PLE 层仍使用 TP-only embedding，没有跨 DP gather/slice；
5. 结果是配置被静默接受，却没有按照 TP×DP 对 PLE 表分片。
   实际影响不只是少了一项优化，而是：

• PLE 权重仍会在各 DP rank 间复制；
• 预期的跨 DP 显存节省不会出现；
• 配置表达的执行拓扑与实际执行拓扑不一致；
• 用户很难从启动日志判断 embedding_across_dp 已经被忽略；
• PR 自己声明支持的“ETP 与 cross-DP PLE”只有绑定另一个独立功能开关后才存在。
  现有测试仍明确固化这一错误条件：
  def test_qwen4_exp_ple_exchange_is_strictly_opt_in_and_lazy(...):
  monkeypatch.delenv("VLLM_HCU_PLE_PREFETCH_STREAM", raising=False)
  assert register_qwen4_exp_ple_exchange(coordinator) == ()
  建议将“加载 HCU PLE 实现”和“是否启用 prefetch”拆开。最直接的修复是始终注册 _QWEN4_EXP_PLE replacement，再由 replacement 内部现有的运行时 guard 决定是否启用 prefetch。当前 replacement 在 prefetch 关闭时本来就有 inline lookup 路径，因此 module exchange 没必要由 prefetch 环境变量控制。
  至少需要新增以下回归组合：
• embedding_across_dp=true、prefetch=false：确认使用 ETP group，并执行跨 DP gather/reduce/slice；
• embedding_across_dp=true、prefetch=true：确认 ETP prefetch 路径；
• embedding_across_dp=false、prefetch=false：确认普通 TP 行为不变；
• embedding_across_dp=false、prefetch=true：确认普通 TP+prefetch 行为不变。
  其中第一项是本次缺失的核心场景，也正是 PR 描述里尚未完成的 cross-DP+ETP E2E。