[Bugfix] Pin FlashInfer bmm_fp8 to cuBLAS on sm_12x to avoid cuDNN hot-path stalls - #48210
Majid-Taheri wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # build_cudnn_gemm_fp8_graph call inside the serving hot path, | ||
| # stalling the engine 12-17s per new prompt length. Fall back to | ||
| # CUTLASS on sm_12x until FlashInfer builds cuDNN plans off the | ||
| # hot path. |
There was a problem hiding this comment.
Turning off Flashinfer capability all backends for bmm_fp8 on SM120+ because of cuDNN backend alone doesn't seem to be the right approach. Can we wait for newer flashinfer version to fix this like the issue suggests?
There was a problem hiding this comment.
Thanks for the review. I agree, turning off the whole kernel is too much. I test another way today: keep the FlashInfer kernel but use backend="cublas" instead of "auto". Then there is no stall on sm_120. New shapes run in under 1 ms. cuDNN is only in the path because vLLM asks for "auto"; flashinfer's default is "cublas". I can update the PR this way. Is that ok? Waiting for a flashinfer fix leaves sm_12x at 48 tok/s instead of 210.
There was a problem hiding this comment.
Datapoint on the wait-for-newer-flashinfer option: the crash half is already fixed in the 0.6.13 that vLLM pins (flashinfer#3437's tactic clamp; A/B'd on GB10 in flashinfer-ai/flashinfer#3566), but the stall half is structural and unchanged on flashinfer main -- build_cudnn_gemm_fp8_graph is an lru_cache keyed on the exact gemm shape (gemm_base.py:3098 on main, :2915 in v0.6.13) and cudnn still joins the auto list on sm12x, so new prompt lengths keep paying the graph build inside the engine loop. FWIW backend="cublas" is also flashinfer's own default for bmm_fp8 (v0.6.13 gemm_base.py:6415); the "auto" comes from the vLLM side in flashinfer_scaled_fp8_mm.
There was a problem hiding this comment.
@pavanimajety I redo the fix the smaller way you asked (force-pushed). The kernel stays on for sm_12x. vLLM now uses backend="cublas" there, not "auto", so cuDNN is not picked. "auto" stays for all other GPUs. I also update the AsyncTP fusion patterns to use the same backend helper, so they still match. On sm_120 the new-shape delay drops from 12-17s to under 1.3s. Thanks @waynehacking8 for confirming the cause.
…d=auto flashinfer_scaled_fp8_mm passes backend="auto" to FlashInfer's bmm_fp8. With the cudnn python module importable (a transitive dependency via nvidia-cudnn-frontend in requirements/cuda.txt), auto selection includes the cuDNN runner, and FlashInfer's fp8_gemm_sm100 builds a cuDNN GEMM graph lazily per exact (batch, seqlen) shape on the host, inside the serving hot path. On the sm_12x (consumer/workstation Blackwell) family this causes: 1. 12-17s engine stalls (GPU at 0%) on every novel prompt length. On a ShareGPT serving benchmark (Nemotron-3-Super-120B NVFP4, RTX PRO 6000 Blackwell) this cut output throughput from ~210 tok/s to ~48 tok/s. 2. Execute-time plan rejection: "Plan index N is invalid" (flashinfer-ai/flashinfer#3566). Instead of gating FlashInferFP8ScaledMMLinearKernel off sm_12x, keep the kernel enabled and pin the backend to cuBLAS (FlashInfer's own default for bmm_fp8) on sm_12x via a new bmm_fp8_backend() helper. "auto" is preserved on other archs (e.g. sm_10x datacenter Blackwell), where the cuDNN runner is autotuned at warmup and beneficial. The AsyncTP fusion patterns (FlashInferBMMFP8ReduceScatterPattern, FlashInferAllGatherBMMFP8Pattern) match torch.ops.vllm.bmm_fp8 on the exact backend string literal, so they now derive it from the same helper — otherwise the fusion would silently stop matching on sm_12x. Tests: parametrized is_supported over cc in {90,100,103,120,121} (sm_12x stays supported) and bmm_fp8_backend over the platform family (cublas on sm_12x, auto elsewhere). Signed-off-by: Majid Taheri Andani <tahemaji@amazon.com>
7783012 to
971b2aa
Compare
Majid-Taheri
changed the title
waynehacking8
left a comment
waynehacking8
left a comment
There was a problem hiding this comment.
Reread the full diff at 971b2aa and ran the two new tests on this branch locally - 7/7 pass, and on my RTX PRO 6000 bmm_fp8_backend() resolves to cublas as intended. The pin matches the stall mechanism I traced earlier: the per-exact-shape cuDNN graph build is still on the serving hot path in current flashinfer main, so keeping the kernel on and pinning cublas on sm_12x only is the right shape for this fix. Deriving the AsyncTP pattern literal from the same helper instead of hardcoding "auto" also avoids the silent pattern-match break. LGTM.
|
Hi @pavanimajety, thanks for the earlier look. A reviewer approved this and the two new tests pass. Can you approve and add the |
3 participants
Purpose
FlashInferFP8ScaledMMLinearKernelruns FP8 linears throughbmm_fp8(backend="auto"). On sm_12x GPUs (RTX PRO 6000, RTX 50, GB10),"auto"picks cuDNN, which builds a GEMM graph for every new shape inside the serving loop. This stalls the engine 12-17 s per new prompt length, GPU at 0%."cublas"is FlashInfer's own default forbmm_fp8. The"auto"comes from vLLM. This PR keeps the kernel on and pins"cublas"on sm_12x only. Other GPUs keep"auto".(Reworked from the first version that gated the kernel off sm_12x. Per review, the backend pin keeps the kernel in use.)
Change (3 files, +99/-4)
flashinfer.py: newbmm_fp8_backend()helper returns"cublas"on sm_12x, else"auto". Uses@torch.compiler.assume_constant_result(functools.cachecrashestorch.compilehere).collective_fusion.py: the two AsyncTP patterns use the same helper, so they still matchbmm_fp8.is_supportedover cc {90,100,103,120,121}, plus the backend pin.Test Result (RTX PRO 6000, sm_120, Nemotron-3-Super-120B-NVFP4, ShareGPT, c=8)
Kernel stays FlashInfer (no CUTLASS fallback). Tests 12/12 pass.
Relation to flashinfer#3566
The crash half is already fixed in flashinfer 0.6.13. The stall is structural on flashinfer main. This pin fixes the stall on the vLLM side and can be relaxed once flashinfer builds cuDNN plans off the hot path.