#55272 [Qwen3.8-Flash-Next] Remove torch.compile for NVIDIA implementation
原始 PR · 作者 gau-nernst · 合并时间 2026-09-07 12:57
移除 Qwen3.8-Flash-Next NVIDIA 路径的 torch.compile 依赖,改用可中断 CUDA 图。
此 PR 值得精读,因为它不仅修复了一个重要的生产问题(FP8 OOM),还展示了如何将模型执行路径从 `torch.compile` 的隐式优化迁移至更显式、可控制的 CUDA 图捕获模式。关键的设计决策包括: 1. **选择性 eager break**:仅在需要运行时元数据(如 `attn_metadata`)的函数上使用 `@eager_break_during_capture`,而不是在整个模型前向传播中。 2. **缓冲区地址稳定性**:通过确保 `ngram_context` 和 `query_start_loc` 缓冲区在捕获和重放期间使用相同的内存地址,避免了 CUDA 图的重新捕获。 3. **参数传递重构**:将 Q/K 投影提前到 `Qwen4ExpQSAAttention.forward` 中,使 `QSAIndexer.forward` 接收预投影的 `projected_qk`,简化了数据流。 推荐关注以下文件:`vllm/models/qwen4_exp/nvidia/qsa.py`(QSA 注意力路径的重构)、`vllm/models/qwen4_exp/nvidia/model_state.py`(缓冲区稳定性逻辑)和 `vllm/models/qwen4_exp/nvidia/ple_layer.py`(PLE 层的 eager break 添加)。
参与讨论