fix(kernel): avoid nan in triton sampling softmax when a chunk starts with -inf

[ischencheng]

Fixes #647.

_sm_partial starts every column chunk with m = -inf. When the first BLOCK_SIZE block of a chunk is all -inf, new_m is still -inf, so exp(x - new_m) and exp(m - new_m) are exp(-inf - -inf) = NaN. _sm_finalize adds that partial to the row sum and the whole row comes out NaN.

This shifts by 0 while the running max is still -inf, so an all -inf prefix adds 0. Once the max is finite the shift is new_m as before, so a row without -inf goes through the same arithmetic. _sm_finalize needs no change: an all -inf chunk now stores pl = 0, and 0 * exp(-inf - gm) is 0 as long as the row has a finite logit.

flashinfer's softmax already returns finite probabilities for these inputs, so only the Triton path is affected. sample_impl uses it when flashinfer is not installed, which is always the case on ROCm. Nothing on main writes -inf into logits yet, so today this is latent; an allowed-token-ids, min_tokens or grammar mask would hit it.

The test goes into tests/kernels/test_triton_fallbacks.py: 3 rows of 151936 logits, two of them masked (the first 16384 columns, or everything except 3 token ids) and one unmasked, with per-row temperatures, against torch.softmax.

Tested on NVIDIA L4 (sm_89) on Modal, CUDA 13.0, driver 610.57.04 for the main test run and 580.95.05 for this PR's, torch 2.11.0+cu130, Triton 3.6.0, flashinfer 0.6.18.post1, on top of 2d8b9cc. No checkpoint is involved.

	main	this PR
pytest tests/kernels/test_triton_fallbacks.py	2 failed, 15 passed	17 passed
the repro in #647	NaN	finite
12 masked cases: first 16384 columns, 3 allowed ids, or a 20000-column block in the middle; B = 1, 3, 8; V = 151936, 201088, 262144	every masked row is all NaN	max rel err vs torch.softmax 3.2e-6 over probs > 1e-6 (flashinfer on the same inputs: 4.2e-6)
sample_impl with flashinfer turned off, 3 allowed ids, 16 rows x 64 draws	1024 of 1024 draws are not an allowed id (the ones I logged are all token 0)	all 1024 draws are allowed ids

Rows without -inf, same process, main's kernels against this PR's, both pinned to the same config:

• outputs are bit-identical for all 12 autotune configs at B=1, 4, 16, 64 (V = 151936, 151936, 201088, 262144), 48 cases;
• do_bench_cudagraph, median of 7 interleaved runs with the config main's autotune picks:

shape	main (µs)	this PR (µs)
B=1, V=151936	5.19	5.16
B=4, V=151936	10.60	10.57
B=16, V=201088	44.54	44.67
B=64, V=262144	811.89	811.58

Not tested: ROCm and sm_120.

A/B script, as run (from the repo root on this branch, with main's python/freetoken/kernel/triton/sampling.py copied to control/sampling_main.py)

"""Same-process A/B of main's softmax kernels (control/sampling_main.py) vs this branch's."""
import importlib.util
import json
import pathlib
import statistics

import torch
import triton
import triton.testing as tt

here = pathlib.Path(__file__).parent
art = here / "artifacts"
art.mkdir(exist_ok=True)
spec = importlib.util.spec_from_file_location("sampling_main", here / "sampling_main.py")
smain = importlib.util.module_from_spec(spec)
spec.loader.exec_module(smain)
from freetoken.kernel.triton import sampling as sfix

out = {"gpu": torch.cuda.get_device_name(), "triton": triton.__version__, "torch": torch.__version__}
shapes = [(1, 151936), (4, 151936), (16, 201088), (64, 262144)]


def pin(mod, cfg):
    for k in (mod._sm_partial, mod._sm_finalize):
        k.configs = [cfg]
        k.cache = {}


# 1. Bitwise: finite logits, both kernels pinned to the same config.
full = list(smain._SM_CFGS)
bitwise = {}
for cfg in full:
    pin(smain, cfg); pin(sfix, cfg)
    for B, V in shapes:
        g = torch.Generator(device="cuda").manual_seed(B * 7 + V)
        x = torch.randn((B, V), device="cuda", generator=g) * 4
        t = torch.rand(B, device="cuda", generator=g) + 0.5
        a, b = smain.softmax(x, t), sfix.softmax(x, t)
        bitwise[f"{cfg.kwargs['BLOCK_SIZE']}/w{cfg.num_warps}/s{cfg.num_stages} B{B} V{V}"] = bool(torch.equal(a, b))
out["bitwise_equal_finite"] = bitwise
out["bitwise_all_equal"] = all(bitwise.values())

# 2. Timing: same config per shape for both (the one main's autotune picks), interleaved.
timing = {}
for B, V in shapes:
    x = torch.randn((B, V), device="cuda") * 4
    t = torch.rand(B, device="cuda") + 0.5
    for k in (smain._sm_partial, smain._sm_finalize):
        k.configs = full; k.cache = {}
    smain.softmax(x, t)
    chosen = smain._sm_partial.best_config
    pin(sfix, chosen)
    fin = smain._sm_finalize.best_config
    sfix._sm_finalize.configs = [fin]; smain._sm_partial.configs = [chosen]; smain._sm_finalize.configs = [fin]
    smain._sm_partial.cache = {}; smain._sm_finalize.cache = {}
    ra, rb = [], []
    for _ in range(7):
        ra.append(tt.do_bench_cudagraph(lambda: smain.softmax(x, t)) * 1000)
        rb.append(tt.do_bench_cudagraph(lambda: sfix.softmax(x, t)) * 1000)
    timing[f"B{B} V{V}"] = {"partial_cfg": str(chosen), "main_us": round(statistics.median(ra), 2),
                           "fix_us": round(statistics.median(rb), 2),
                           "main_runs": [round(v, 2) for v in ra], "fix_runs": [round(v, 2) for v in rb]}
out["timing"] = timing
(art / "probe_ab.json").write_text(json.dumps(out, indent=1))
print(json.dumps({"bitwise_all_equal": out["bitwise_all_equal"], "n": len(bitwise),
                  "timing": {k: (v["main_us"], v["fix_us"]) for k, v in timing.items()}}, indent=1))