BSA: fix the cute-dsl bulk-copy elect gate for 4.6.2 and 4.7

[Anerudhan]

Authored by XiaoDong (@xiaoD, xiaod@nvidia.com) — ported from an internal change; the commit preserves the original authorship.

Fixes BSA blk128 backward hangs on SM100 with cute-dsl 4.6.2 and 4.7.0a0 (passes on 4.6.1).

Affected tests (all the same kernel — supported_block_size(backward=True) is 128 on SM10x):

• test_bsa_attention_backward_fixed_blocks
• test_bsa_attention_backward_blk128_dk_zero_init_accumulate_transition[1] / [4]

Under pytest -n=4 --timeout=300 the hang surfaces as worker 'gwN' crashed … node down.

Cause

The elect_one() guard was gated on cutlass.__version__[:2] >= (4, 6), matching a cute-dsl change that made the raw cp.async.bulk.* lowering elect a lane internally. That change has since been reverted on every cute-dsl line (shipped in 4.6.2 and 4.7.0a0), restoring the pre-4.6 contract that the caller elects — see NVIDIA/cutlass#3391.

Because the gate only compares major.minor, 4.6.2 and 4.7 still take the "the DSL elects for me" branch. Nothing elects, so all 32 lanes issue the same cp.async.bulk…complete_tx::bytes for the LSE/dPsum stats copies. producer_acquire sized expect_tx for a single copy, so the barrier is over-completed 32x and the pipeline wedges.

Fix

Only 4.6.0 and 4.6.1 ever shipped the internal election. Match that release range instead of >= 4.6; pre-releases (e.g. 4.6.0.dev0) and everything else fall through to "the caller elects".

Why not just revert the original gate

On 4.6.0/4.6.1 a caller-side elect_one() nests with the internal one and deadlocks — that is what the original gate fixed. Both are inside the declared support range (nvidia-cutlass-dsl>=4.5.0), so the gate has to stay; it just has to name the right versions. Measured on B200/SM100a with a full-warp cute.copy(CopyBulkG2SOp, mbar_ptr=…):

cute-dsl	caller elects	elect.sync in PTX	launch
4.6.0 / 4.6.1	no	1	ok
4.6.0 / 4.6.1	yes	2	wedges
4.7.0a0	no	0	wedges
4.7.0a0	yes	1	ok

Verification

Fresh clone + real pip install per leg, nothing patched in place:

cute-dsl	before	with this change
4.7.0a0	hang, killed at timeout	3 passed · fe_api/bsa 17 passed
4.6.1	—	3 passed · fe_api/bsa 17 passed
4.6.0	—	3 passed · fe_api/bsa 17 passed

Full fe_api under nightly settings: 80 failed / 2179 passed / 1668 skipped, zero hangs; residual failures match the baseline family-for-family (gemm_swiglu 64, sdpa_bwd 12, dsa 4) with only the 3 BSA ones removed.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved compatibility for block-sparse attention across supported CUTLASS versions.
  • Prevented potential deadlocks or barrier synchronization issues during bulk-copy operations.
  • Refined version detection to handle patch releases and pre-release versions more accurately.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 483663b7-e30d-4883-a8be-d67f3cd7d348

📥 Commits

Reviewing files that changed from the base of the PR and between 57dc579 and 4775612.

📒 Files selected for processing (1)

• python/cudnn/block_sparse_attention/csrc/utils/copy_utils.py

---

📝 Walkthrough

Walkthrough

The bulk-copy utility now uses strict regex-based Cute DSL version detection, enabling self-election only for non-prerelease versions 4.6.0 and 4.6.1. Its election behavior documentation now describes the effects of selecting the wrong election mode.

Changes

Bulk-copy election detection

Layer / File(s)	Summary
Strict version detection and election documentation python/cudnn/block_sparse_attention/csrc/utils/copy_utils.py	Adds regex version parsing, limits self-election detection to versions 4.6.0–4.6.1 without prerelease suffixes, and documents the resulting deadlock or barrier-desynchronization failure modes.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description is detailed, but it does not follow the required template and omits sections like Affected area, Related issues, and API impact.	Reformat the PR body to match the template and add the missing sections, especially Affected area, Related issues, API and compatibility impact, and Testing.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: tightening the cute-dsl bulk-copy election gate for affected versions.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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