Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an in-graph WAR-barrier read-done event mechanism for CUDA graphs. This allows backends whose captured metadata kernels index live shared buffers on every replay (such as trtllm_mha) to record the event inside the graph right after the metadata hook, ensuring correct synchronization. It adds the in_graph_metadata_reads_shared_buffers property across attention backends and updates DecodeCudaGraphRunner and EagleDraftExtendCudaGraphRunner to plant and manage this in-graph event. The review feedback suggests catching AttributeError in addition to TypeError when initializing the WAR-barrier read-done event to prevent potential crashes in environments where device_module does not have an Event attribute.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
/tag-and-rerun-ci |
| war_policy = self._war_read_done_policy( | ||
| self.draft_extend_attn_backend, forward_batch.forward_mode | ||
| ) | ||
| if war_policy == "pre_replay": |
There was a problem hiding this comment.
nit: maybe we could introduce a WarReadDonePolicy enum to avoid literal string
2 participants
trtllm_mha's fused in-graph metadata kernel re-reads the live req_to_token / full->SWA mapping on every graph replay -- after the WAR fastpath's pre-replay read-done event was recorded -- so scheduler-side writes could overwrite them mid-replay. Record the read-done event as an external-event node inside the captured graph at the snapshot-completion point (for backends with a no-op in-graph hook the node sits at the graph head, preserving today's semantics), with a new backend capability falling back to post-replay recording where in-graph capture is unavailable.
CI States
Latest PR Test (Base): 🚫 Run #29707770973
Latest PR Test (Extra): ❌ Run #29707770880