[fix] add graph capture patch(like vLLM) for sglang+atom path

[zhuyuhua-v]

Motivation

• add a shared graph capture patch that wraps the framework's GroupCoordinator.graph_capture with aiter's ca_comm.capture()
• enable the patch for the SGLang + ATOM plugin path during model preparation
• refactor the existing ATOM+vLLM-specific implementation to reuse the same shared helper

Technical Details

When ATOM runs as a plugin backend, the model uses aiter collectives such as tensor_model_parallel_fused_allreduce_rmsnorm, but the host framework only enters its own graph capture context. As a result, aiter does not enter capture mode, falls back to the unregistered path, and triggers extra hipMemcpyAsync calls.

This change makes the framework graph capture also enter aiter's ca_comm.capture() so the aiter collectives can run in the expected capture mode and avoid the extra memory copy overhead.

Impact

• fixes the graph capture mismatch for the SGLang + ATOM path
• keeps the vLLM path aligned by sharing the same implementation
• reduces unnecessary hipMemcpyAsync overhead in SGLang plugin mode

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Adds a shared “graph capture patch” helper so plugin backends (vLLM and SGLang) enter both the host framework’s GroupCoordinator.graph_capture context and aiter’s ca_comm.capture() context, preventing aiter collectives from falling back to non-capture paths (and avoiding extra hipMemcpyAsync overhead) in plugin mode.

Changes:

• Introduce a shared patching utility (atom.plugin.graph_capture_patch) to wrap a framework’s GroupCoordinator.graph_capture with aiter capture.
• Refactor the vLLM-specific patch module to delegate to the shared helper.
• Add a SGLang patch module and invoke it during SGLang plugin prepare_model.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
atom/plugin/graph_capture_patch.py	New shared helper to patch a framework’s GroupCoordinator.graph_capture to nest aiter capture.
atom/plugin/vllm/graph_capture_patch.py	Refactored to delegate to the shared helper for vLLM.
atom/plugin/sglang/graph_capture_patch.py	New delegating patch module for SGLang.
atom/plugin/prepare.py	Applies the SGLang patch during model preparation in SGLang plugin mode.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

In the wrapper, graph_capture_context is forwarded positionally (original_graph_capture(self, graph_capture_context, **kwargs)). To keep this patch resilient to frameworks that define graph_capture_context as keyword-only, pass it by name (e.g., graph_capture_context=graph_capture_context) when calling the original method.

Suggested change

	with original_graph_capture(self, graph_capture_context, **kwargs) as ctx:
	with original_graph_capture(
	self, graph_capture_context=graph_capture_context, **kwargs
	) as ctx:

[Copilot]

_GRAPH_CAPTURE_PATCH_APPLIED is set to the boolean returned by the shared helper. Since the shared helper returns False both when the patch cannot be applied and when it is already patched, this wrapper will keep retrying on every call (re-importing and re-checking). Consider marking _GRAPH_CAPTURE_PATCH_APPLIED = True after the first attempt, or adjusting the shared helper to return True when it detects an existing patch.

Suggested change

	_GRAPH_CAPTURE_PATCH_APPLIED = _apply("vllm.distributed.parallel_state")
	_apply("vllm.distributed.parallel_state")
	_GRAPH_CAPTURE_PATCH_APPLIED = True

[Copilot]

_GRAPH_CAPTURE_PATCH_APPLIED is set to the boolean returned by the shared helper. Since the shared helper returns False both when the patch cannot be applied and when it is already patched, this wrapper will keep retrying on every call (re-importing and re-checking). Consider marking _GRAPH_CAPTURE_PATCH_APPLIED = True after the first attempt, or adjusting the shared helper to return True when it detects an existing patch.

Suggested change

	_GRAPH_CAPTURE_PATCH_APPLIED = _apply("sglang.srt.distributed.parallel_state")
	_apply("sglang.srt.distributed.parallel_state")
	_GRAPH_CAPTURE_PATCH_APPLIED = True

[Copilot]

prepare_model now unconditionally attempts to apply the SGLang graph-capture patch in the sglang path. There are existing unit tests for prepare_model in tests/plugin/test_sglang_prepare_model.py, but none assert that this patch hook is invoked (or that failures are handled). Add a test that injects a fake atom.plugin.sglang.graph_capture_patch module and asserts apply_graph_capture_patch() is called once on the happy path.