fix duplicate entries in the tp > 1 case

neuralmagic · Sep 17, 2024 · a1271fa · a1271fa
1 parent 344595b
commit a1271fa
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py
@@ -161,6 +161,11 @@ def _build_module_tree(self):
 
         def _df_traversal(event: _ProfilerEvent,
                           curr_node: Optional[_ModuleTreeNode] = None):
+
+            # For the tensor parallel case for now only look at task 1
+            if event.start_tid != 1:
+                return
+
             if event_has_module(event):
                 node = _ModuleTreeNode(event=event, parent=curr_node)
                 if curr_node: