[PR #583] Capture PyTorch profiling trace in HybridEPBuffer #644

Open
opened 2026-10-27 21:04:49 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/583
Author: @shengf-nv
Created: 2/18/2026
Status: 🔄 Open

Base: hybrid-ep ← Head: shengf/hybrid-ep-chakra


📝 Commits (10+)

📊 Changes

5 files changed (+662 additions, -7 deletions)

View changed files

📝 csrc/hybrid_ep/config.cuh (+64 -0)
📝 csrc/hybrid_ep/executor/executor.cuh (+20 -0)
📝 csrc/hybrid_ep/pybind_hybrid_ep.cu (+26 -7)
➕ csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh (+381 -0)
➕ tests/test_hybrid_ep_collect_et.py (+171 -0)

📄 Description

Capture PyTorch profiling trace in HybridEPBuffer

To test it, run "python tests/test_hybrid_ep_collect_et.py", you can find et_[0-3].json are generated for the PyTorch chakra trace.

HybridEPBuffer::combine is recorded as the following:

{
"id": 55, "name": "combine", "ctrl_deps": 2,
"inputs": {"values": [[7168,4096,8,4,1,512,108,2112,"uint8_t",10,8,128,32,false,true,10,2,128,4,32,2,false,true],[13,16,0,109355008,2,"cuda:0"],"",[28,8,0,16384,4,"cuda:0"],[35,9,0,4096,1,"cuda:0"],[37,10,0,0,1,"cuda:0"],4096,false], "shapes": [],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[,[15256,7168],[],[4096,4],[4096,1],[4096,0],[],[]], "types": ["Tuple[Int,Int,Int,Int,Int,Int,Int,Int,String,Int,Int,Int,Int,Bool,Bool,Int,Int,Int,Int,Int,Int,Bool,Bool]","Tensor(c10::BFloat16)","None","Tensor(int)","Tensor(bool)","Tensor(bool)","Int","Bool"], "strides": [],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[,[7168,1],[],[4,1],[1,1],[1,1],[],[]]},
"outputs": {"values": 47,17,0,29360128,2,"cuda:0"],[56,0,0,0,0,"", "shapes": 4096,7168],[, "types": ["Tensor(c10::BFloat16)","Tensor(nullptr (uninitialized))"], "strides": 7168,1],[},
"attrs": [{"name": "rf_id", "type": "uint64", "value": 39},{"name": "fw_parent", "type": "uint64", "value": 0},{"name": "seq_id", "type": "int64", "value": -1},{"name": "scope", "type": "uint64", "value": 0},{"name": "tid", "type": "uint64", "value": 1},{"name": "fw_tid", "type": "uint64", "value": 0},{"name": "op_schema", "type": "string", "value": ""},{"name": "kernel_backend", "type": "string", "value": ""},{"name": "kernel_file", "type": "string", "value": ""},{"name": "tensor_range", "type": "string", "value": ""}]
},


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/583 **Author:** [@shengf-nv](https://github.com/shengf-nv) **Created:** 2/18/2026 **Status:** 🔄 Open **Base:** `hybrid-ep` ← **Head:** `shengf/hybrid-ep-chakra` --- ### 📝 Commits (10+) - [`c2746f5`](https://github.com/deepseek-ai/DeepEP/commit/c2746f5db8ef54c6d6fce6e3087fcc4963f59ac1) first implementation for collecting ET - [`9c13226`](https://github.com/deepseek-ai/DeepEP/commit/9c1322686ca19eb334a4d6387dec25792c7b6bd4) cleanup code - [`b0dc913`](https://github.com/deepseek-ai/DeepEP/commit/b0dc913e97b1a1b5eb0655dd99df651c35a4ffc1) use c10::impl::GenericList - [`bb5bc8f`](https://github.com/deepseek-ai/DeepEP/commit/bb5bc8ff4199d35b038a067a3701d134277b2b76) support all APIs - [`f8480e7`](https://github.com/deepseek-ai/DeepEP/commit/f8480e7a2cd87997a06aa0fcdf2d2045794313d2) put new APIs into a new file - [`43db483`](https://github.com/deepseek-ai/DeepEP/commit/43db4837847e62c8d7902d74f76ff3376ef00e09) put new APIs into a new file - [`8fc1510`](https://github.com/deepseek-ai/DeepEP/commit/8fc15108af11250be721821d9ecbd670a8764361) clean up head files - [`a09d4db`](https://github.com/deepseek-ai/DeepEP/commit/a09d4dba959ae6cf38e4a046954dd55e241eb312) get pg id - [`ba9a0ab`](https://github.com/deepseek-ai/DeepEP/commit/ba9a0ab62ba1e87017d02fddc38ae38c04e1910b) get pg id - [`cc98dd8`](https://github.com/deepseek-ai/DeepEP/commit/cc98dd8e7d47ba55bb9011c8718b3114f054ebe1) add schema ### 📊 Changes **5 files changed** (+662 additions, -7 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/config.cuh` (+64 -0) 📝 `csrc/hybrid_ep/executor/executor.cuh` (+20 -0) 📝 `csrc/hybrid_ep/pybind_hybrid_ep.cu` (+26 -7) ➕ `csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh` (+381 -0) ➕ `tests/test_hybrid_ep_collect_et.py` (+171 -0) </details> ### 📄 Description Capture PyTorch profiling trace in HybridEPBuffer To test it, run "python tests/test_hybrid_ep_collect_et.py", you can find et_[0-3].json are generated for the PyTorch chakra trace. HybridEPBuffer::combine is recorded as the following: { "id": 55, "name": "combine", "ctrl_deps": 2, "inputs": {"values": [[7168,4096,8,4,1,512,108,2112,"uint8_t",10,8,128,32,false,true,10,2,128,4,32,2,false,true],[13,16,0,109355008,2,"cuda:0"],"<None>",[28,8,0,16384,4,"cuda:0"],[35,9,0,4096,1,"cuda:0"],[37,10,0,0,1,"cuda:0"],4096,false], "shapes": [[[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[]],[15256,7168],[],[4096,4],[4096,1],[4096,0],[],[]], "types": ["Tuple[Int,Int,Int,Int,Int,Int,Int,Int,String,Int,Int,Int,Int,Bool,Bool,Int,Int,Int,Int,Int,Int,Bool,Bool]","Tensor(c10::BFloat16)","None","Tensor(int)","Tensor(bool)","Tensor(bool)","Int","Bool"], "strides": [[[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[]],[7168,1],[],[4,1],[1,1],[1,1],[],[]]}, "outputs": {"values": [[47,17,0,29360128,2,"cuda:0"],[56,0,0,0,0,""]], "shapes": [[4096,7168],[]], "types": ["Tensor(c10::BFloat16)","Tensor(nullptr (uninitialized))"], "strides": [[7168,1],[]]}, "attrs": [{"name": "rf_id", "type": "uint64", "value": 39},{"name": "fw_parent", "type": "uint64", "value": 0},{"name": "seq_id", "type": "int64", "value": -1},{"name": "scope", "type": "uint64", "value": 0},{"name": "tid", "type": "uint64", "value": 1},{"name": "fw_tid", "type": "uint64", "value": 0},{"name": "op_schema", "type": "string", "value": ""},{"name": "kernel_backend", "type": "string", "value": ""},{"name": "kernel_file", "type": "string", "value": ""},{"name": "tensor_range", "type": "string", "value": ""}] }, --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#644
No description provided.