fix: prevent NaN outputs for short/non-multiple sequence lengths via CUDA-level zero-fill

[Onkitova]

Summary

This PR fixes a bug where SageAttention produces NaN outputs when the sequence length is not a multiple of the block size (e.g., short sequences like seq_len = 9 in Krea 2 Turbo models, or non-multiple shapes like Q_len = 1001, KV_len = 503).

Root Cause

In csrc/qattn/attn_utils.cuh, the predicated asynchronous global-to-shared memory copy helper load_global_to_share was configured with cp_async::SharedMemFillMode::kNoFill:

smem.load_128b_async<cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len);

When base_idx >= max_len (i.e. out-of-bounds elements in the boundary/tail block), kNoFill skips loading, leaving the shared memory tile with uninitialized garbage from previous GPU operations.

When the Tensor Cores perform the QK matrix multiplication, they read this garbage (which can contain NaNs or subnormal values). Although the out-of-bounds attention scores are later correctly masked to -5000000.0f in apply_out_of_bound_mask, NaN values produced during the intermediate matrix multiplication contaminate the calculations and propagate through arithmetic operations (since 0 * NaN = NaN).

Fix

We changed the fill mode to cp_async::SharedMemFillMode::kFillZero:

smem.load_128b_async<cp_async::SharedMemFillMode::kFillZero>(smem_offset, *lane_ptr, base_idx < max_len);

This instructs the hardware (cp.async instruction) to automatically zero-fill the out-of-bound bytes. Since the out-of-bounds scores are masked later anyway, loading clean zeros instead of uninitialized garbage prevents NaN generation and propagation entirely.

Note

Note on FP16 vs FP8 paths:
This fix resolves the NaN issue for Q, K, and V loading in the FP16 paths, and Q and K loading in the FP8 paths. The FP8 V loading path (load_fp8_V_global_to_share) remains unpredicated because the FP8 V tensor is already pre-padded to a multiple of 64 during Python-level quantization preprocessing, making it inherently memory-safe.

Performance Impact

The performance overhead is negligible (~0.0%):

• For in-bound elements (which represent the vast majority of operations), both modes compile to the exact same 16-byte async copy instructions from global memory.
• For out-of-bound elements (which only occur in the boundary/tail block of a sequence), kFillZero performs a cp.async with 0 source bytes to clear the shared memory tile. This is handled entirely within the fast SM shared memory pipeline and does not consume any global memory bandwidth.

Verification & Tests

1. Added a comprehensive test suite in tests/test_sageattn.py that:
   • Sweeps sequence lengths from 1 to 65.
   • Tests non-multiple layouts (Q=1001, KV=503) across both HND and NHD tensor layouts.
2. Verified that all tests now pass successfully with 0 NaNs on SM89 (RTX 40-series) and SM80 GPUs.
3. Confirmed that the Krea 2 Turbo model (which previously produced color noise / NaNs in ComfyUI due to its short text/time pooling layers of length 9 and 12) now runs flawlessly.