DSA backward SM100: support zero-length top-k rows in the kernel

[jiayus-nvidia]

Summary

Add native SM100 kernel support for sparse-attention backward rows whose
topk_length is zero.

The kernel now detects an empty row before initializing any asynchronous
pipeline or allocating TMEM. All threads in the CTA cooperatively write zero
to the corresponding dQ row and then exit. The row contributes nothing to
dKV or d_sink, matching the mathematical gradient of a query attending
to no KV entries.

Malformed negative lengths take the same defensive early-exit path so they
cannot enter the zero-tile pipeline path.

Why

The SM100 backward kernel previously computed:

tile_count = cute.ceil_div(topk, self.block_tile)

For topk == 0, this produces zero tiles, but the warp-specialized pipeline
assumes that at least one tile is handed between its producer and consumer
stages. On B200 this can deadlock the launch at 100% SM utilization. If
execution reached the epilogue, dQ could also be written from TMEM
accumulators that no MMA instruction had initialized.

Handling the row inside the kernel:

• avoids a per-call device reduction and device-to-host synchronization;
• behaves consistently during CUDA graph capture;
• gives empty rows a defined and mathematically natural result;
• leaves the existing positive-topk_length path unchanged.

Implementation

For each query-token/head-block CTA, the kernel reads the CTA-uniform
topk_length before setting up pipelines or TMEM.

When the value is nonpositive, it:

1. writes zero to every valid dQ element owned by the CTA;
2. skips all KV, MMA, reduction, and pipeline work;
3. exits the CTA before any asynchronous synchronization state is created.

The normal path reuses the same topk value, so positive rows do not incur
an additional load.

API and compatibility impact

• topk_length == 0 is now supported on SM100.
• The corresponding dQ row is exactly zero.
• The row contributes zero to dKV and d_sink.
• Calls with positive lengths are unchanged.
• Negative lengths are handled defensively as empty rows rather than being
  allowed to enter the deadlocking pipeline path.

Related work

This is an alternative kernel-side resolution for #433.

Thanks to @zkyue for identifying the zero-tile failure mode, providing the
B200 hang reproducer, and analyzing the likely pipeline deadlock. That
investigation directly motivated this implementation.

Unlike the interface guard proposed in #433, this change supports empty rows
without adding a host synchronization or capture-specific behavior.

Testing

Tested on an NVIDIA B200 (SM100):

PYTHONPATH=/code/github/cudnn-frontend/python
pytest -q -rs
fe_api/dsa/test_DSA_sparse_attention_backward.py
-k sm100_zero_topk_length

Result:

2 passed, 6 deselected

Coverage includes:

• mixed lengths around the 64-row tile boundary;
• zero and defensive negative lengths;
• all-empty inputs;
• head dimensions 512 and 576;
• full and partial head CTAs;
• caller-provided dQ and dKV buffers;
• numerical comparison with the PyTorch reference;
• exact-zero checks for all-empty gradients.

Summary by CodeRabbit

• Bug Fixes

  • Fixed sparse attention backward processing for empty or invalid attention rows.
  • Empty rows now produce zero gradients without affecting valid rows.
  • Improved handling of invalid indices in empty rows and prevented potential GPU synchronization issues.

• Tests

  • Added regression coverage for zero and negative attention lengths across multiple configurations.

[coderabbitai]

📝 Walkthrough

Walkthrough

The SM100 sparse-attention backward kernel now handles zero or negative topk_length rows by zeroing mdQ and exiting early. A parametrized GPU regression test validates outputs, gradients, buffer overwrites, invalid indices, all-empty cases, and synchronization.

Changes

SM100 empty sparse-row handling

Layer / File(s)	Summary
Kernel empty-row fast path python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py	The kernel computes topk before pipeline setup, zeros mdQ for non-positive rows, exits the CTA, and removes the later redundant initialization.
Empty-row regression coverage test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py	A parametrized SM100 test validates zero and negative lengths, empty-row outputs and gradients, caller-owned buffer overwrites, ignored invalid indices, all-empty gradients, and synchronization.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested labels: cat-enhancements

Suggested reviewers: zkyue, anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	Title clearly summarizes the main change: SM100 sparse-attention backward now supports zero-length top-k rows.
Description check	✅ Passed	The description covers summary, why, API impact, related issue, and testing; only template-specific headings like Affected area are omitted.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[zkyue]

Thanks @jiayus-nvidia — this is a better resolution than the interface guard in #433. Handling the empty row inside the kernel avoids the per-call D2H sync, keeps CUDA-graph capture behavior consistent, and the early exit before any pipeline/TMEM setup addresses the deadlock at its root.

I also looked at the sm100 score_recompute path while reviewing this: its block-skip loop participates in the mbarrier handshake for every skipped slot, so it is already safe for zero-length rows — with this PR the whole family is closed for empty rows.

Closing #433 in favor of this.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-439-ed665d8
Pipeline: 59802840