Bug in PyTorch 2.11 causes SM89 hangs

[endenis]

Summary

On an RTX 4070 (SM89), Qwen3.6-35B-A3B-NVFP4 can hang indefinitely during prefill with PyTorch 2.11 but works with PyTorch 2.13.

While hung, GPU utilization remains near 100% and the request never completes.

I suspect this to be the cause of the following existing issues:

• Prefill hangs on RTX 4060 Ti (SM89) for prompts beyond warmup length #72
• Request hang (silent, zero-log) with hybrid_radix cache + moe-backend offload + nvfp4 triton backend — reproduced on 2 independent models, not explained by #110 #123

An early clue was that the same workload completed when FreeToken was run with:

CUDA_LAUNCH_BLOCKING=1 uv run ft serve ...

I isolated the cause to a row-wise FP8 torch._scaled_mm call in the fused QKV projection:

(2010, 2048) @ (2048, 9216) -> (2010, 9216)
FP8 E4M3 inputs, BF16 output
scale_a=(2010, 1)
scale_b=(1, 9216)

This matches an upstream PyTorch SM89 bug fixed by pytorch/pytorch@252bb4a

Environment

• GPU: RTX 4070, SM89
• Driver: 580.142
• CUDA: 13.0
• Affected PyTorch: 2.11.0+cu130
• Model: Qwen3.6-35B-A3B-NVFP4

Reproduction scripts

I made two reproduction scripts

Standalone PyTorch reproducer

#!/usr/bin/env python3
"""Minimal reproducer for an SM89 row-wise FP8 torch._scaled_mm hang."""

import os
import time

import torch


# M is the number of rows, K the input size, and N the output size.
M, N, K = 2010, 9216, 2048

# Repetition makes the intermittent failure reliable:
# 5/5 locally at 128 calls on an RTX 4070.
OPS = 128
TIMEOUT_SECONDS = 10

torch.manual_seed(0)
a = (
    torch.randn((M, K))
    .mul_(10)
    .clamp_(-448, 448)
    .to(device="cuda", dtype=torch.float8_e4m3fn)
)

torch.cuda.manual_seed_all(1)
b = (
    torch.randn((N, K), device="cuda")
    .mul_(32)
    .clamp_(-448, 448)
    .to(torch.float8_e4m3fn)
)

# These scale shapes select the row-wise scaling path.
scale_a = torch.ones((M, 1), device="cuda")
scale_b = torch.ones((1, N), device="cuda")

ready = torch.cuda.Event()
ready.record()
while not ready.query():
    time.sleep(0.01)

stream = torch.cuda.Stream()
events = [torch.cuda.Event() for _ in range(OPS)]

with torch.cuda.stream(stream):
    for event in events:
        out = torch._scaled_mm(
            a,
            b.t(),
            scale_a=scale_a,
            scale_b=scale_b,
            out_dtype=torch.bfloat16,
        )
        event.record(stream)

deadline = time.monotonic() + TIMEOUT_SECONDS
completed = 0

while time.monotonic() < deadline:
    while completed < OPS and events[completed].query():
        completed += 1

    if completed == OPS:
        print(f"completed all {OPS} calls")
        raise SystemExit(0)

    time.sleep(0.1)

print(
    f"GPU stalled: {completed}/{OPS} calls completed; "
    f"call #{completed + 1} did not complete within "
    f"{TIMEOUT_SECONDS} seconds",
    flush=True,
)

# Normal interpreter shutdown would wait for the stuck CUDA operation.
os._exit(124)

FreeToken reproducer

Start a fresh server:

uv run ft serve \
    --model /path/to/Qwen3.6-35B-A3B-NVFP4 \
    --host 127.0.0.1 \
    --port 1919 \
    --attention-backend triton,fi \
    --moe-backend offload \
    --disable-moe-prefill-overlap \
    --cuda-graph-max-bs 0

Then run:

import json
import urllib.request


fragment = (
    "CUDA stream ordering and expert cache lifetime must remain correct "
    "across asynchronous prefill execution. "
)

body = json.dumps(
    {
        "model": "Qwen3.6-35B-A3B-NVFP4",
        "prompt": (
            "freetoken-prefill-hang-fp8-ab-fixed-0 unique-prefix 17\n"
            + fragment * 124
        ),
        "max_tokens": 1,
        "temperature": 0,
    }
).encode()

request = urllib.request.Request(
    "http://127.0.0.1:1919/v1/completions",
    data=body,
    headers={"Content-Type": "application/json"},
)

print(urllib.request.urlopen(request, timeout=60).read().decode())

With PyTorch 2.11, the request can time out while the GPU remains near 100% utilization.

Results

PyTorch 2.11: standalone reproducer hung in 5/5 runs
PyTorch 2.12: 0/5 hangs
PyTorch 2.13: 0/5 hangs

In FreeToken, the fused row-wise QKV operation hung in 5/5 fresh-server trials. Splitting it into tensor-wise Q, K, and V operations completed in 5/5 trials.

The split implementation was only a diagnostic control, not a proposed fix.

Proposed fix

Update FreeToken dependency stack to:

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

This includes the upstream fix and aligns with SGLang 0.5.18, which uses PyTorch 2.13 and sglang-kernel==0.4.6.post1