Fix Sage2 causal attention for unequal qkv lengths

[qqq-tao]

Summary

Fix Sage2 causal attention when qo_len < kv_len, which is the common KV-cache/decode shape.

The previous Sage2 causal mask used top-left alignment:

kv_idx <= q_idx

For decode-style unequal q/kv lengths, the expected behavior is right-aligned causal masking:

kv_idx <= q_idx + (kv_len - qo_len)

This matches the Sage3 behavior and avoids incorrect attention windows when the KV sequence is longer than the query sequence.

This PR also hardens the varlen Triton path against NaN propagation from finite but extremely large logits.

Root Cause

• Sage2 CUDA kernels computed causal bounds as if qo_len == kv_len.
• The varlen Triton causal kernel used the same top-left mask alignment.
• The varlen softmax path could propagate NaNs when finite inputs produced non-finite intermediate scores.

Changes

• Add shared CUDA helpers for right-aligned causal KV offset and causal KV limit.
• Apply right-aligned causal masking in SM80, SM89/SM120, and SM90 Sage2 CUDA paths.
• Update varlen Triton causal attention to use right-aligned masking.
• Add causal input validation for qo_len <= kv_len.
• Add numerical guards in varlen Triton attention and quantization paths.
• Add regression tests for fixed-length CUDA, varlen masking, invalid lengths, finite-large logits, and reported large-shape smoke cases.

Before / After

With old installed SageAttention, using q/k = ones, v = normalized key index, qo_len=5, kv_len=9:

4D CUDA observed:          [0.0, 0.0625, 0.125, 0.1914, 0.2520]
varlen observed:           [0.0, 0.0625, 0.125, 0.1875, 0.25]
right-aligned expected:    [0.25, 0.3125, 0.375, 0.4375, 0.5]

With this patch:

4D CUDA observed:          [0.2520, 0.3125, 0.3789, 0.4434, 0.5039]
varlen observed:           [0.25, 0.3125, 0.375, 0.4375, 0.5]
varlen large-input NaNs:   4224 -> 0

Validation

Validated on an RTX 5090 / SM120 environment with CUDA 12.8.

Build:

TORCH_CUDA_ARCH_LIST=12.0 MAX_JOBS=8 \
  python -m pip install -v --no-build-isolation -e .

Regression tests:

CUDA_VISIBLE_DEVICES=0 \
  python -m pytest -q tests/test_cuda_causal_unequal_qkv.py -q
# 10 passed, 2 skipped

Large-shape smoke tests:

CUDA_VISIBLE_DEVICES=0 SAGEATTN_RUN_SLOW_TESTS=1 \
  python -m pytest -q \
  tests/test_cuda_causal_unequal_qkv.py::test_fp8_cuda_causal_user_reported_shape_is_finite \
  tests/test_cuda_causal_unequal_qkv.py::test_varlen_causal_user_reported_shape_is_finite -q
# 2 passed

Additional checks:

python -m compileall -q sageattention tests
git show --check

Notes

SM80 and SM89/SM120 were compiled and validated on the available SM120 machine. SM90 was updated with the same right-aligned causal logic, but was not runtime-tested because the validation machine cannot execute the SM90 path.