[liuanji]

Summary

The small-batch CUDA sum forward (layer/kernels/c/smallbatch_forward_sum.cu, used below batch 64 when the autotuner picks it over Triton) returned NaN whenever a child value was −inf. This PR fixes it with no speed cost.

The bug

The kernel computes each node's logsumexp online, with the running max starting at −inf:

float nmx = fmaxf(m, emar);
l = l * __expf(m - nmx) + p * __expf(emar - nmx);   // m = emar = -inf  ->  exp(-inf - (-inf)) = NaN

If the first children a warp reads are −inf, the running sum becomes NaN and stays NaN, even after finite children follow. The final merge across warps has the same problem when a whole split is −inf.

Measured on main (a85d063):

input	batch ≤ 32 (CUDA path)	batch 64 (Triton)
10% of a layer's children at −inf	100% of outputs NaN	correct
HMM (1024 latents) whose emission table has exact zeros, pc(x)	every row NaN	correct

Any PC whose product values can be −inf is affected at small batch: zero-probability evidence, zero parameters, Indicator / Literal leaves. The other CUDA sum kernels and the BlockScale small-batch forward already handle −inf; this one never did.

The fix

The running max starts at -FLT_MAX instead of −inf. Every update is then well defined:

• a −inf child adds p * exp(-inf) = 0;
• the first finite child scales the empty sum by exp(-FLT_MAX - x) = 0;
• a node whose children are all −inf gets log(0) = -inf.

Verification

Old vs fixed kernel, each built as its own extension and called with the exact arguments the sum layer passes (hidden size 256 / 1024 / 2048, batch 1–48, all four SPLIT configs):

• Finite inputs: bit-identical outputs.
• Inputs with −inf (leading −inf runs, 10% scattered, an all −inf column): within 7e-7 of an fp64 reference, never NaN.
• Speed: unchanged, within ±1% in every case. A guard on the new max inside the loop was also correct but measured 10–110% slower, because it stops the edge loop from pipelining.
• End to end, the HMM with zero emissions: no NaN at batch 1–32. The results differ from the batch-64 rows by at most 4.6e-4; a dense HMM shows 2.2e-4, so this is the usual difference between the two kernel paths.

New test: test_small_batch_forward_cuda_neg_inf_children calls the kernel directly with every config, so it doesn't depend on which kernel the autotuner picks, and compares against fp64. It fails on the old kernel. The existing test_small_batch_forward_cuda_matches_triton passes whenever the autotuner picks Triton, so it could not catch this.

Full suite: 1249 passed, 30 skipped on one GPU. The 4 extra skips compared with a 4-GPU run are the tests that need at least 2 GPUs.

🤖 Generated with Claude Code