Add graph runner support with torch compile on CPU

[CaoE]

Motivation

Inspired by mingfeima#73. We add CPU graph runner with torch compile to reduce python overhead to speed up decoding on CPU.

Profiling with disabling torch compile:

Profiling with enabling torch compile:

From the profiling results above, we can see that torch compile can reduce python overhead by reducing the module call stack.

Modifications

• Fix the mutation errors in kernel registration on CPU.
• Register fake functions for custom kernels to support torch.compile on CPU.
• Convert custom parameters to native parameters of torch to avoid graph break on CPU.
• Add graph runner support on CPU. Note: cpu graph needs to be enabled using --enable-torch-compile

Performance:

Enable compile vs. Disable compile:

Accuracy:

		enable torch compile	disable torch compile
model id	dtype	mmlu	mmlu
meituan--DeepSeek-R1-Channel-INT8	int8	0.87	0.871
Qwen--Qwen3-14B-FP8	fp8	0.788	0.788
meta-llama--Llama-3.2-3B-Instruct	bf16	0.607	0.607
microsoft--Phi-4-multimodal-instruct	bf16	0.689	0.688

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.

[gemini-code-assist]

Summary of Changes

Hello @CaoE, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the SGLang runtime by integrating torch.compile support for CPU inference. The primary goal is to minimize Python overhead during model execution on CPU, mirroring the performance benefits previously achieved with CUDA graphs. This involves extensive modifications to enable graph capture, ensure compatibility with custom kernels and parameter types, and streamline distributed operations within a compiled environment, ultimately leading to more efficient CPU model execution.

Highlights

• CPU Graph Runner Integration: Introduced a generalized GraphRunner (formerly CudaGraphRunner) to enable graph capture and replay for CPU inference, leveraging torch.compile to reduce Python overhead. This involves adapting the existing CUDA graph infrastructure to support CPU-specific optimizations.
• torch.compile Compatibility for Custom Kernels: Added a new module (cpu_register_fake.py) that registers 'fake' implementations for various custom sgl_kernel CPU operations. This crucial step allows torch.compile to correctly trace the computational graph involving these operations without needing their actual execution during compilation, enabling end-to-end optimization.
• Enhanced Parameter Handling for torch.compile: Extended ChannelQuantScaleParameter with necessary __tensor_flatten__, __torch_dispatch__, and copy methods. These additions ensure that custom torch.nn.Parameter subclasses are properly handled and optimized by torch.compile's graph tracing mechanism.
• Distributed Operations and torch.compile Alignment: Adjusted the shm_allreduce operation and registered a fake shm_allgather to ensure seamless compatibility with torch.compile when distributed operations are performed on CPU, preventing tracing errors and enabling optimization of communication primitives.
• General torch.compile Infrastructure Improvements: Included various minor fixes and enhancements across the codebase, such as adjusting torch._inductor configurations for CPU, improving memory capacity reporting for CPU, and handling context managers during compilation to prevent tracing issues and improve robustness.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces graph runner support for CPU to reduce Python overhead. I've identified a critical issue in the replay logic that needs to be addressed, along with a few medium-severity issues related to code clarity and potential bugs.

[Alcanderian]

Hi @Alcanderian @zhyncs Can this PR be landed since the NPU graph runner has been re-landed ? Recently, some AMD-GPU and NPU-related tests have failed, most of which seem unrelated to this PR. Thanks.

please fix xeon CI

[CaoE]

please fix xeon CI

Fixed https://github.com/sgl-project/sglang/actions/runs/17227399712/job/48876607152?pr=7843.

[CaoE]

Hi @FlamingoPg @Alcanderian Can this PR be merged?

[CaoE]

@FlamingoPg @Alcanderian https://github.com/sgl-project/sglang/actions/runs/17231906249/job/48964007907?pr=7843 doesn't seem to be related to this PR.

[CaoE]

Hi @zhyncs Could you please help merge this PR ? Thank you.

[FlamingoPg]

Hi @zhyncs Could you please help merge this PR ? Thank you.

@zhyncs LGTM, any other comment?

[zhyncs]

can we keep the old name

[CaoE]

Changing can_run_cuda_graph to can_run_graph is simply to indicate that graph is also supported on the CPU, thus reducing confusion. Depending on the device, it can be determined whether it is a CUDA graph or a CPU graph. Do you think we should keep all can_run_cuda_graph? Are there any unconsidered implications of changing can_run_cuda_graph to can_run_graph? This is fine for us, and I can change all can_run_graph back to can_run_cuda_graph.

[zhyncs]

I hope to minimize the impact on downstream forks as much as possible, usually new hardware changes are best made independently, with minimal changes to existing NVIDIA GPUs.

[CaoE]

Thanks for your comments. I'll change them back.

[CaoE]

Modified.

[CaoE]

Hi @zhyncs https://github.com/sgl-project/sglang/actions/runs/17372540433/job/49363628357?pr=7843 doesn't seem to be related to this PR. Can this PR be merged?

[CaoE]

Hi @zhyncs Could you please take another look? Thank you.