feat: enable NPU aclgraph backend for Python model executor. - #2053
feat: enable NPU aclgraph backend for Python model executor.#2053maojunx99 wants to merge 3 commits into
Conversation
b3cb2e0 to
102afaa
Compare
maojunx99
changed the title
ea1ff1a to
23cb195
Compare
8e9e010 to
1216ae1
Compare
| max_seqs_per_batch, | ||
| int(config["max_position_embeddings"]), | ||
| ) | ||
| elif graph_backend == "aclgraph": |
There was a problem hiding this comment.
这里需要手动配置吗？不是在npu device上配置了graph模式开启，就导向aclgraph？
There was a problem hiding this comment.
不需要了，现在npu+python+graph自动推导backend为aclgraph
python ttft比c++快？TPOT慢的原因能不能看一下，直接解决了？方便贴一下decode的profiling图吗，看看哪里还有bubble？ |
|
有精度测试报告吗？没有影响精度吧 |
| metadata: AttentionMetadata, | ||
| *, | ||
| graph_mode: bool = False, | ||
| static_block_table: torch.Tensor | None = None, |
There was a problem hiding this comment.
改成接口不变，用 ForwardContext 传递 会不会更好
There was a problem hiding this comment.
已修改。删除了 static_block_table、padded_batch_size、static_slot_mapping 等 NPU 专用参数，恢复通用 prepare(metadata, graph_mode) 接口。ACL runner 直接传入完整的静态 metadata；capture 状态则通过 ForwardContext 传递。
| self._in_capture = False | ||
| return list(self._task_handles) | ||
|
|
||
| def update_replay_args( |
There was a problem hiding this comment.
update_replay_args 这种东西是不是 runner自己持有，自己update比较好？而不是一个 attention 类方法？
There was a problem hiding this comment.
已修改。task handle、update stream、event 以及 graph_task_update_begin/end 现在都由 ACL graph runner 管理。attention backend 只注册 FIA 参数更新 callback，因为具体 FIA 调用参数仍属于 attention 实现。
| return output.reshape(num_tokens, self.num_heads * self.head_dim) | ||
|
|
||
| def begin_capture(self, stream: torch.npu.Stream) -> None: | ||
| self._in_capture = True |
There was a problem hiding this comment.
这种字段 _in_capture / _capture_stream / _task_handles 能删除吗？
改为通过 ForwardContext 传递信号是否可行？
attention 拥有图的生命周期管理权是否合理
There was a problem hiding this comment.
已修改。上述字段以及 begin_capture/end_capture/update_replay_args 已从 attention backend 删除。capture stream 和 task records 通过 ForwardContext 传递，图的 capture、update 和 replay 生命周期统一由 runner 管理。
| ForwardContext(self.attention_backend, self.device) | ||
| ): | ||
| if not first_capture: | ||
| self.attention_backend.update_replay_args( |
There was a problem hiding this comment.
这里建议避免由 ACL graph runner 直接依赖 attention_backend.update_replay_args()，而是抽象一层通用的 graph-task update callback。具体来说，算子在 capture 时通过 ForwardContext/recorder 注册 task record（handle + operator-specific update callback），runner 持有这些 records，并在每次 replay 前统一执行 graph_task_update_begin -> callback -> graph_task_update_end；callback 内仍由 attention 实现重新发射 FIA，因为具体调用参数属于 attention 的算子知识。
这个抽象不应只服务于 FIA。后续 Qwen3.5 的 Conv1D Ascend-C 算子也包含 host tensor 参数，在 ACL graph replay 前同样需要走类似的 task 参数更新流程。如果这里把更新入口固定为 attention backend，后续接入 Conv1D 等算子时，runner 中会不断增加算子专用分支。建议现在就将 graph orchestration 与 operator-specific update 解耦，形成可复用的 callback/record 机制。
3 participants
Description
Add a new "aclgraph" option for --python_graph_backend that captures and replays decode-step graphs on Ascend NPU using native torch.npu.* APIs (NPUGraph, graph_task_group_begin/end, graph_task_update_begin/end).
Also optimize Qwen3 row-parallel weight layout for pytorch path, transform O proj, Down proj from ND to NZ.
Change Type
Performance Comparision (from server log)
Qwen3 1.7B, tp=1, 2k input + 1k output, batch=1
Qwen3 8B, tp=1, 2k input + 1k output, batch=1
Qwen3 8B, tp=1, 2k input + 1k output, batch=8
Accuracy
Unit Tests