Skip to content

perf: Fix python API overhead when CUDAGraph is not enabled#969

Merged
yzh119 merged 9 commits intoflashinfer-ai:mainfrom
yzh119:cpu-overhead
Mar 24, 2025
Merged

perf: Fix python API overhead when CUDAGraph is not enabled#969
yzh119 merged 9 commits intoflashinfer-ai:mainfrom
yzh119:cpu-overhead

Conversation

@yzh119
Copy link
Copy Markdown
Collaborator

@yzh119 yzh119 commented Mar 23, 2025

This PR fixes issue #960 , we identifies several performance bottlenecks for our python APIs when kernels are not captured by CUDAGraph:

  1. The device guard in Python is slow (with input.device as device:)
  2. Get current cuda stream in Python is time-consuming.

These issues were introduced in JIT refactor after v0.1.6 (mainly for accelerating JIT compilation speed). In this PR, we changed back to get stream and device guard in C++).

@MichoChan @xiaoqi35

@yzh119
Copy link
Copy Markdown
Collaborator Author

yzh119 commented Mar 23, 2025

After the changes, here is the norm python API performance (cpu and python overhead included, https://gist.github.com/yzh119/d9bf2abbb667abcbb806979f4bbea633):

A100 Performance

This PR

w/o CUDAGraph 0.006144285202026367
w/  CUDAGraph 0.003793001174926758

v0.1.6

w/o CUDAGraph 0.004994630813598633
w/  CUDAGraph 0.0035665035247802734

H100 Performance

This PR (no PDL)

w/o CUDAGraph 0.004047870635986328
w/  CUDAGraph 0.0024733543395996094

This PR (w/ PDL)

w/o CUDAGraph 0.003877878189086914
w/  CUDAGraph 0.0021266937255859375

v0.1.6 (no PDL)

w/o CUDAGraph 0.0033130645751953125
w/  CUDAGraph 0.0023775100708007812

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant