NSA compression: suppress aliased O1 store for static short queries

[mgoldfarb-nvidia]

Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.
• I added GitHub labels: one cat-*, one or more mod-*, and one orig-*. Maintainer action requested: cat-bug, mod-cutedsl, and orig-nv-eng; the contributor token cannot add upstream labels.

Affected area

FE OSS kernels or CuTeDSL

Summary

• Compile out the second 128-row NSA compression output store only when the query extent is statically known to contain one MMA tile or less.
• Keep both correction/epilogue pipeline stages consumed and released, with the TMA wait count adjusted to the number of committed stores.
• Add a JAX static-specialization regression at query lengths 127, 128, 129, and 300.

Why

Each compression CTA processes two 128-row query tiles. When a statically specialized query extent contains only one tile, flat_divide canonicalizes the singleton residual tile mode to stride zero. The unconditional O1 TMA store therefore aliases O0 and overwrites valid output instead of remaining out of bounds.

The guard checks the layout extent with cute.is_static and evaluates the tile-count condition with cutlass.const_expr. Dynamic extents and static multi-tile extents retain the original O1 store/wait path; only the static singleton specialization prunes the aliased store. Both correction pipeline handles are still consumed and released unconditionally.

Related issues

None.

API and compatibility impact

None. This is a correctness fix for the SM100 NSA compression kernel and does not change its public API. The guard folds during tracing, so the dynamic path adds no runtime comparison or branch. The short static specialization removes one invalid TMA store and commit.

Testing

• NVIDIA B200, CUDA 13.0, CuTeDSL 4.7.0, JAX 0.10.2:
  • Before the fix, seqlen_q=128 failed against the FP32 reference with max absolute error 0.27612305; seqlen_q=129 passed.
  • cd test/python && pytest -q fe_api/nsa/test_NSA_compression_attention_jax.py -m L0: 4 passed (127, 128, 129, 300).
  • The supplied bridge/direct repro produced zero difference at 128, 200, 256, and 300.
  • Direct dynamic outputs at 128, 129, and 300 were bit-identical to the unmodified cuDNN Frontend 1.25 kernel.
• uvx pre-commit run --files python/cudnn/native_sparse_attention/compression/fmha.py test/python/fe_api/nsa/test_NSA_compression_attention_jax.py: passed.
• python3 -m py_compile python/cudnn/native_sparse_attention/compression/fmha.py test/python/fe_api/nsa/test_NSA_compression_attention_jax.py: passed.
• git diff --check: passed.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b7e9d382-6887-4495-890c-9e0696282d95

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}