Handle CUDA Graph node zero annotations

Author: drisspg

test/test_perfetto.py

@@ -21,7 +21,7 @@ def _duration_events(trace):
 def test_cuda_graph_annotations_become_contiguous_gpu_boxes():
     graph_id = 2
     annotations = {
-        (graph_id << 32) | 1: [{"name": "attention"}],
+        (graph_id << 32) | 0: [{"name": "attention"}],
         (graph_id << 32) | 2: [{"name": "attention"}],
         (graph_id << 32) | 3: [{"name": "loss"}],
     }
@@ -35,7 +35,7 @@ def test_cuda_graph_annotations_become_contiguous_gpu_boxes():
                 "tid": 7,
                 "ts": 10,
                 "dur": 3,
-                "args": {"graph id": graph_id, "graph node id": 1},
+                "args": {"graph id": graph_id, "graph node id": 0},
             },
             {
                 "ph": "X",

transformer_nuggets/utils/perfetto.py

@@ -128,7 +128,7 @@ def add_cuda_graph_annotation_boxes(
         args = event.get("args", {})
         graph_id = args.get("graph id")
         graph_node_id = args.get("graph node id")
-        if graph_id is None or not graph_node_id:
+        if graph_id is None or graph_node_id is None:
             continue
 
         entries = None