fix(kernel): avoid row-wise _scaled_mm stall on sm_89 with torch<2.12

[bernimccoy]

Fixes #182. Same root cause as #72 and the prefill hangs in #220; also covers #227.
Complements #184 (torch 2.13): this keeps the current pin and is a no-op on torch >= 2.12.

Problem

On sm_89, torch < 2.12 runs row-wise _scaled_mm through a CUTLASS kernel whose launch
ignores the current stream (pytorch/pytorch#177651, fixed in 2.12 by pytorch/pytorch@252bb4a;
isolated by @endenis in #182). fp8_pertensor_linear takes that path for fused projections
(q/k/v, GDN qkv|z: one scalar per part) from a side stream. At M >= 256 it stalls the GPU:
no Prefill batch line, 100 % util, kill -9 required; decode-sized M completes, so short
prompts work and long ones hang. Windows torch builds have no row-wise kernel at all (#227).

Fix

Tensor-wise scaling is unaffected. Where row-wise is unsafe, a fused
projection runs one tensor-wise _scaled_mm per part over its row slice (weight[s:e].t(),
still stride-only) and concatenates. Same W8A8 scheme (rel ~7e-4 to row-wise, accumulation order), one extra launch per part.

• rowwise_scaled_mm_ok(): False on sm_89 + torch < 2.12; otherwise a one-off row-wise probe
  on the default stream (safe there), so a build without the kernel fails at load. Warmed in
  load_state_dict, never under graph capture.
• weight_scale_segments(): the parts' row ranges from the load-time weight_scale
  run-lengths, next to the existing _uniform_scale check. A genuine per-row scale (> 8
  segments or not 16-row aligned) stays on the W8A16 triton path.
• FREETOKEN_FP8_ROWWISE_MM=0/1 forces either path.

Unchanged: uniform-scale weights, no input_scale, sm_90+/sm_120, sm_89 on torch >= 2.12.

Tested on

RTX 4070 SUPER 12 GB (sm_89), driver 591.86, WSL2, torch 2.11.0+cu130, triton 3.6.0,
flashinfer 0.6.17, sglang-kernel 0.4.5. nvidia/Qwen3.6-35B-A3B-NVFP4, HF safetensors:

ft serve --model ~/models/Qwen3.6-35B-A3B-NVFP4 --port 1919

	main (= FREETOKEN_FP8_ROWWISE_MM=1)	this branch
18-token prompt	ok	ok
626-token prompt, max_tokens=32	never returns (> 90 s, GPU 100 %, kill -9)	3 tokens, 4.4 s
streamed decode, 299 tokens, bs=1	n/a	26.4 tok/s

Sweep of the fused [8192, 512, 512] x 2048 projection on a side stream: row-wise completes at
M <= 128, stalls at M >= 256; the per-part path completes at every M from 1 to 2010.

tests/kernels/test_fp8_pertensor_linear.py: side-stream test at the #182 shape, in a
subprocess so a stall fails instead of hanging pytest; fails on main (rc=124), passes
here. Per-part path vs row-wise on the same fused input at M=1/4/64/300: rel 6.9e-4 to 8.5e-4.

Not tested on sm_90+/sm_120, torch >= 2.12, or Windows (#227 relies on the probe).

Caveats

Pre-existing on this GPU, unrelated: 3 test_w8a8_matches_w8a8_reference cases miss the 1e-2
tolerance (rel 0.0103-0.0107) on main and here alike, incl. the untouched uniform case.

Assisted by: Claude Code

[jason-fxz]

Verified on RTX 4060 Laptop (sm_89, torch 2.11.0+cu130). The fix itself LGTM. Three requests before merge:

Comments

Keep them short, and only where the code isn't self-explanatory — the header block can be 3 lines. Three statements need fixing:

• Header: "every other FP8-capable arch goes to cuBLASLt" is wrong. In torch 2.11 row-wise goes to CUTLASS on Ada and Blackwell; only Hopper (cuBLAS ≥ 12.9) uses cuBLASLt. The sm_89 gate is still right — only the sm89 kernel uses stream-K, whose initialize memsets a workspace on the wrong stream.
• _scaled_mm docstring: "exact, same W8A8 numerics" → same scheme, not bit-identical (accumulation order differs; ~7% of elements, rel ~7e-4).
• rowwise_scaled_mm_ok docstring: "Probe on the default stream" — the probe runs on the current stream. Drop the claim or actually switch to the default stream.

Tests

One test is enough: run the same fused input through the row-wise and per-part paths (monkeypatch rowwise_scaled_mm_ok) and compare them directly at a few M, rel < 2e-3 (measured ~7e-4). This replaces the three new tests that go through the W8A8 reference or check internal plumbing.

Commit message

Use Assisted-by: <assistant name> instead of Co-Authored-By:.

[bernimccoy]

Thanks for the review and the 4060 run. All three addressed in the amended commit:

• Comments: header cut to 3 lines and the cuBLASLt claim dropped (now names the stream-K launch); _scaled_mm docstring says same scheme, not bit-identical; the probe now actually runs on the default stream (torch.cuda.stream(torch.cuda.default_stream(dev))) and the docstring matches.
• Tests: the three reference/plumbing tests replaced by test_per_part_path_matches_rowwise, which monkeypatches rowwise_scaled_mm_ok both ways on the same fused input at M=1/4/64/300 and asserts rel < 2e-3 (measured 6.9e-4 to 8.5e-4 on the 4070 SUPER). The side-stream subprocess test stays.
• Commit trailer: Assisted-by:.

PR body updated to match.

[jason-fxz]

Nice work, thanks. Merging.