Skip to content

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

Open
endenis wants to merge 1 commit into
FlashML-org:mainfrom
endenis:fix/sm-89-hang-in-torch-scale-mm
Open

fix(deps): use PyTorch 2.13 to fix SM89 hangs#184
endenis wants to merge 1 commit 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

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.

Bug in PyTorch 2.11 causes SM89 hangs

1 participant