Skip to content

benchmark_inference: Add CLI option to enable thunder CUDAGraph Transform #2697

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kshitij12345 merged 9 commits into from
Nov 4, 2025
Merged

benchmark_inference: Add CLI option to enable thunder CUDAGraph Transform #2697

Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
62585f5
benchmark_inference: Add CLI option to enable thunder CUDAGraph Trans…
kshitij12345 Oct 27, 2025
686d594
update
kshitij12345 Oct 27, 2025
438768d
move import at module level
kshitij12345 Oct 27, 2025
60800ab
Merge branch 'main' into ksh/bench-inf-cudagraph
kshitij12345 Oct 27, 2025
0e89d2f
fix usage with nv_enable_linear
kshitij12345 Oct 28, 2025
6647fa2
fix merge conflict
kshitij12345 Oct 30, 2025
1c9e8f6
Merge branch 'ksh/bench-inf-cudagraph' of https://github.com/Lightnin…
kshitij12345 Oct 30, 2025
b5b62bd
Merge branch 'main' into ksh/bench-inf-cudagraph
kshitij12345 Nov 3, 2025
87eb65c
retrigger CI
kshitij12345 Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion thunder/benchmarks/benchmark_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,7 @@ class InferenceBenchmarkConfig:
mode: str
disable_moe_replacement: bool
profile: bool
enable_thunder_cudagraph: bool


@dataclass
Expand Down Expand Up @@ -281,7 +282,7 @@ def __init__(self, config: InferenceBenchmarkConfig):
def _thunder_jit_options(self) -> dict[str, Any]:
# `nv_enable_linear=True` might fail with distributed run
# ref: https://github.com/NVIDIA/Fuser/issues/4507
res = {}
res = {"transforms": []}
if self.config.enable_nv_linear:
res = {"nv_enable_linear": True, "nv_enable_matmul": True}
if self.config.mode == "thunderjit":
Expand All @@ -291,6 +292,10 @@ def _thunder_jit_options(self) -> dict[str, Any]:
self._mask_transform = SDPAMaskTransform()
res["transforms"] = [self._mask_transform]
res["executors"] = [self._mask_transform.get_executor(), *thunder.get_default_executors()]
if self.config.enable_thunder_cudagraph:
from thunder.transforms.cudagraph import CUDAGraphTransform

res["transforms"].append(CUDAGraphTransform())
return res

def _compile_model(self, model):
Expand Down Expand Up @@ -676,6 +681,7 @@ def parse_args() -> argparse.Namespace:

parser.add_argument("--save-results", action="store_true", help="Save results to JSON file")
parser.add_argument("--output-dir", type=str, default="./results", help="Directory to save results")
parser.add_argument("--enable-thunder-cudagraph", action="store_true", help="Pass CUDAGraphTransform to Thunder")

args = parser.parse_args()
return args
Expand Down Expand Up @@ -708,6 +714,7 @@ def main():
enable_nv_linear=args.enable_nv_linear,
disable_moe_replacement=args.disable_moe_replacement,
profile=args.profile,
enable_thunder_cudagraph=args.enable_thunder_cudagraph,
)
benchmark = InferenceBenchmark(config)

Expand Down
Loading