Skip to content

fix(deps): use PyTorch 2.13 to fix SM89 hangs - #184

Closed
endenis wants to merge 2 commits into
FlashML-org:mainfrom
endenis:fix/sm-89-hang-in-torch-scale-mm
Closed

fix(deps): use PyTorch 2.13 to fix SM89 hangs#184
endenis wants to merge 2 commits into
FlashML-org:mainfrom
endenis:fix/sm-89-hang-in-torch-scale-mm

Conversation

@endenis

@endenis endenis commented Aug 25, 2026

Copy link
Copy Markdown

Fixes #182

Possibly related to #72 and #123

Summary

PyTorch 2.11 can hang indefinitely during prefill on SM89 GPUs.

I reproduced this with Qwen3.6-35B-A3B-NVFP4 on an RTX 4070: the request stops making progress while GPU utilization remains near 100%.

This PR updates FreeToken to PyTorch 2.13 and along with these dependencies:

torch>=2.13,<2.14
triton==3.7.1
sglang-kernel==0.4.6.post1

Root cause

I isolated the non-completing GPU submission to the fused QKV projection’s row-wise FP8 torch._scaled_mm call:

(2010, 2048) @ (2048, 9216) -> (2010, 9216)
FP8 E4M3 inputs, BF16 output
scale_a=(2010, 1)
scale_b=(1, 9216)
  • CUDA events showed that the GPU reached this call but never completed it.
  • CUDA_LAUNCH_BLOCKING=1 avoided the hang but only by changing the timing (not a real fix).
  • Splitting the projection into separate Q, K, and V calls also avoided the hang. This suggests the problem is in the row-wise FP8 path, not FP8 matrix multiplication in general.

The behavior matches an upstream PyTorch SM89 bug fixed by pytorch/pytorch@252bb4a
The fix is present in PyTorch 2.12 and 2.13.

I made two reproduction scripts:

  1. Standalone PyTorch reproducer
  2. FreeToken reproducer

You can find them in the description of the #182 issue.

Changes

  • Require PyTorch 2.13 for FreeToken at runtime and during extension builds.
  • Update Triton to 3.7.1 (required by PyTorch 2.13).
  • Update sglang-kernel to 0.4.6.post1 which is built against the PyTorch 2.13.
  • Apply the same PyTorch constraint to freetoken-kernel-cache.

There are no kernel or synchronization workarounds in this change.

Testing

Tested on:

GPU: RTX 4070 (Ada, SM89)
Driver: 580.142
CUDA: 13.0
Model: Qwen3.6-35B-A3B-NVFP4

Standalone _scaled_mm reproducer (attached in the #182 issue):

PyTorch 2.11: 5/5 runs hung
PyTorch 2.13: 0/5 runs hung

Notes

I only have access to an RTX 4070. It would be useful to test on other cards as well.
Don't hesitate to reach out if you need more details

@jason-fxz

Copy link
Copy Markdown
Collaborator

Closing: #243 fixes the sm_89 hang without bumping torch, so this is no longer needed for #182. Thanks for isolating the root cause.

@jason-fxz jason-fxz closed this Aug 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

2 participants