BSA: remove dead SM100 block-128 fragment allocations

[jiayus-nvidia]

Summary

• Remove an unused TMEM scale-load setup chain from the SM100 block-128 correction warp.
• Remove an RMEM scale tensor allocation whose result was never consumed.
• Remove an outer tOrO_frg allocation that was unconditionally overwritten before use.
• Keep the per-iteration make_rmem_tensor allocation to preserve the original fragment layout, dtype, and lifetime.

Motivation

nvidia-cutlass-dsl 4.6 removed the deprecated top-level cute.make_fragment API and private cute.core.ThrMma type. The BSA implementation had
already migrated its live register tensors and type annotations to cute.make_rmem_tensor and cute.ThrMma.

That migration exposed several pre-existing dead allocations in the SM100 block-128 correction path. This change removes only those unused
constructs, following the cleanup in FlashInfer PR #3922.

No copies, barriers, reductions, indexing operations, TMEM loads/stores, or architecture instructions are changed.

Compatibility

Verified with:

• nvidia-cutlass-dsl==4.5.2
• nvidia-cutlass-dsl==4.6.1

The block_sparse_attention tree contains no remaining uses of:

• cute.make_fragment(...)
• cute.core.ThrMma
• cute.core.ThrCopy

MMA object methods such as make_fragment_A/B/C remain unchanged because they are still valid public APIs.

Validation

Tested on NVIDIA B200 / SM100:

• BSA forward and backward correctness tests: 17 passed with CUTLASS DSL 4.5.2.
• BSA forward and backward correctness tests: 17 passed with CUTLASS DSL 4.6.1.
• Cold-cache 4.6.1 JIT trace and compilation of the SM100 block-128 forward kernel succeeded.
• All 39 files under python/cudnn/block_sparse_attention parsed and imported successfully with both versions.
• Black 26.3.1 formatting check passed.
• git diff --check passed.

The tests cover SM100 block-128 and block-64 forward/backward kernels. SM90 and SM120 modules were statically imported under both DSL versions but
were not executed because the validation machine was SM100.

[coderabbitai]

📝 Walkthrough

Walkthrough

The correction path removes intermediate vectorized softmax-scale setup and allocation, while changing tOrO_frg temporary storage from one outer allocation to per-fragment allocation.

Changes

Correction rescale changes

Layer / File(s)	Summary
Adjust correction temporaries python/cudnn/block_sparse_attention/csrc/fwd/sm100_blk128/bsa_fwd_sm100.py	correction_loop no longer constructs or allocates vector softmax-scale tensors, and correction_rescale allocates tOrO_frg within each fragment iteration.

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

🚥 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	The title clearly matches the change: removing dead SM100 block-128 fragment allocations.
Description check	✅ Passed	The description covers summary, motivation, compatibility, and validation with concrete test results.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-392-f482017
Pipeline: 58189329