fix(cute_dsl): migrate off APIs removed in nvidia-cutlass-dsl 4.6 (keeps 4.5.x compat)

[YangXu1990uiuc]

📌 Description

nvidia-cutlass-dsl 4.6.0 (released 2026-07-02) executed its long-deprecated API clean-up:

• cute.make_fragment removed → cute.make_rmem_tensor
• cute.core.ThrMma / cute.core.ThrCopy removed → cute.ThrMma / cute.ThrCopy

Since requirements.txt floats on nvidia-cutlass-dsl>=4.5.0, a fresh install now resolves to 4.6.0, and the blk128 sparse-attention path and the SM120 b12x dense block-scaled GEMM path raise AttributeError at trace time.

This PR is a pure rename (±31 lines, 4 files):

• cute.make_fragment( → cute.make_rmem_tensor( — 18 sites
• cute.core.ThrMma → cute.ThrMma — 13 sites

No version guard and no requirements change needed: the replacement names already exist in 4.5.0/4.5.1/4.5.2 (verified against the v4.5.x tags), so the same code works across 4.5.x and 4.6.x. cute.make_fragment_like (79 sites) is not touched — it is still shipped in 4.6.0 and in the current internal nightly.

✅ Verification

• Live probes on both versions: make_rmem_tensor/ThrMma/ThrCopy present in 4.5.2 and 4.6.0; make_fragment confirmed absent in 4.6.0.
• A make_rmem_tensor smoke kernel traces, compiles, and runs with correct output on SM100 under both cutlass-dsl 4.5.2 and 4.6.0.
• Repo-wide grep for the removed spellings is clean (flashinfer/, tests/, benchmarks/).
• pre-commit (ruff, mypy, formatters) passes.

Not covered here (needs CI/QA hardware): full blk128/VSA suite on SM100 (requires quack; note quack must be version-paired — quack 0.5.x for DSL 4.5.x, quack ≥0.6 pins ==4.6.0), and SM120 b12x GEMM/MoE suites.

🔍 Related Issues

Follow-up (separate PR): flashinfer/cute_dsl/sparse/blk128/mma_sm100_desc.py references cutlass.FloatE4M3FN/cutlass.FloatE5M2, which do not exist in any of 4.5.x/4.6.0 — latent AttributeError on the fp8 branch, independent of this migration.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved compatibility and reliability across sparse attention and dense GEMM kernels.
  • Updated internal temporary tensor handling used by softmax, correction, reduction, and epilogue paths to improve consistent execution on supported GPU workloads.
  • Adjusted predicate and intermediate tensor staging to maintain existing computation while improving stability under varied shapes and configurations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f1c8d640-1607-45c8-9c10-b467ee057c0c

📥 Commits

Reviewing files that changed from the base of the PR and between 871a177 and cb7ec1b.

📒 Files selected for processing (1)

• flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py

💤 Files with no reviewable changes (1)

• flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py

---

📝 Walkthrough

Walkthrough

CuTe DSL temporaries are migrated from fragments to RMEM tensors across sparse attention, GQA packing, utility reductions, and dense GEMM predicates. MMA annotations use cute.ThrMma, and an unused correction-loop temporary is removed.

Changes

CuTe RMEM migration

Layer / File(s)	Summary
Sparse helper and packing allocations flashinfer/cute_dsl/sparse/blk128/utils.py, flashinfer/cute_dsl/sparse/blk128/pack_gqa.py	Reduction, predicate, conversion, scalar, MMA helper, and GQA pointer temporaries use RMEM tensors or public MMA annotations.
FlashAttention pipeline temporaries flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py	FlashAttention SM100 methods update MMA annotations, use RMEM tensors for softmax and correction intermediates, and remove an intermediate correction-loop temporary.
Dense GEMM predicate storage flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py	Row predicate tensors use RMEM allocation while preserving predicate initialization and conditional copy operations.

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

Possibly related PRs

• flashinfer-ai/flashinfer#3430: Updates related CuTe tensor allocation APIs and MMA annotations.
• flashinfer-ai/flashinfer#3473: Applies similar RMEM temporary construction changes in kernel paths.

Suggested labels: op: attention

Suggested reviewers: saltyminty

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly names the DSL 4.6 API migration and 4.5.x compatibility goal.
Description check	✅ Passed	The description covers the change, related issue, and verification, though it doesn't use the full template checklist headings.
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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[gemini-code-assist]

Code Review

This pull request refactors the codebase by replacing cute.core.ThrMma with cute.ThrMma and updating calls from cute.make_fragment to cute.make_rmem_tensor across several files. The review feedback highlights two important issues in flash_fwd_sm100.py: an unused tensor allocation that represents dead code and should be removed, and a redundant outer tensor allocation that is immediately overwritten inside a loop, which can be optimized to avoid unnecessary register pressure.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The allocation cute.make_rmem_tensor(tSrScale_t2r_shape, Float32) is completely unused and represents dead code. Removing it cleans up the implementation and avoids any potential confusion.

Suggested change

	cute.make_rmem_tensor(tSrScale_t2r_shape, Float32)
	# q_stage=1 correction loop
	# q_stage=1 correction loop

[YangXu1990uiuc]

Fixed in cb7ec1b. Pulling this thread revealed the flagged line was the last reference into a whole vestigial tmem scale-load scaffold (tStScale_layout / tStScales / tScScale / thr_tmem_load_vec and a discarded list-comprehension of partition_S calls), so the entire reference-less chain is removed (−26 lines), not just this line — all trace-time-pure layout math, no copies/barriers touched. Note this dead code predates this PR (it was equally dead as cute.make_fragment). Re-verified on SM100: VSA suite 30 passed / 9 skipped under both cutlass-dsl 4.5.2 and 4.6.1, identical to pre-cleanup.

[gemini-code-assist]

The outer allocation of tOrO_frg with shape (tOrO_t2r_shape, frg_count) is immediately overwritten inside the loop with a smaller shape tOrO_t2r_shape. This wastes register allocation and increases register pressure unnecessarily. Allocating tOrO_frg outside the loop with the correct shape tOrO_t2r_shape and reusing it inside the loop avoids redundant allocations.

        tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype)\n        for i in cutlass.range_constexpr(frg_count):

[YangXu1990uiuc]

Fixed in cb7ec1b by deleting the dead outer allocation (it was rebound on the first loop iteration before any use, so removal is a no-op on the generated kernel). Deliberately NOT hoisting the in-loop allocation out as suggested: that changes fragment lifetime/reuse semantics in a hot correction loop for no verified gain, and this PR aims to stay behavior-preserving. Also pre-existing before this PR's rename. Re-verified on SM100: 30 passed / 9 skipped under both 4.5.2 and 4.6.1.

[waynehacking8]

Ran this branch on SM120 (RTX PRO 6000, CUDA 13): tests/moe/test_b12x_fused_moe.py is 142/142 under nvidia-cutlass-dsl 4.5.2 and 4.6.0, and the mm_fp4 b12x n=k=128 slice (96 cases) passes identically on both, so the b12x half holds up on real hardware -- the blk128 sparse files are SM100-only so I couldn't exercise those. One catch for the 4.6 story: the 4.6.0 DSL's kwargs-wrapper path needs apache-tvm-ffi>=0.1.11, and on 0.1.6 (still allowed by flashinfer's >=0.1.6,<0.2 floor; the 4.6.0 wheels don't declare a tvm-ffi dep themselves) the whole moe suite dies at the first kernel build with TypeError: make_kwargs_wrapper() got an unexpected keyword argument 'map_dataclass_to_tuple' from cutlass/cutlass_dsl/tvm_ffi_provider.py -- its ImportError guard names >=0.1.11 but only catches a missing module, not the old signature. Probably worth bumping the requirements floor alongside this migration or the eventual 4.6 pin bump.

[YangXu1990uiuc]

Thanks @waynehacking8 for the SM120 verification — that covered exactly the half we couldn't reach. We've now run the other half plus an old-arch sweep on our dev box (A100 / L40S / H100 / SM100-Blackwell), this time against 4.6.1:

SM100 — VSA blk128 suite (tests/attention/test_vsa_block_sparse.py, the SM100-only files this PR touches):

• 4.5.2 + quack 0.5.0: 30 passed / 9 skipped
• 4.6.1 + quack 0.6.1: 30 passed / 9 skipped — identical result set

SM80 / SM89 / SM90 — norm suite under 4.6.1 (tests/utils/test_norm.py; relevant because the CuTe-DSL norm kernels are the default backend on every arch, not just Blackwell):

• A100 (SM80): 1374 passed / 1344 skipped (arch-gated skips)
• L40S (SM89): 1374 passed / 1344 skipped
• H100 (SM90): 2718 passed / 0 skipped

Two observations from the 4.6.1 runs:

1. +1 on the apache-tvm-ffi>=0.1.11 floor bump — our env has 0.1.11 so we didn't hit the make_kwargs_wrapper TypeError, but the failure mode @waynehacking8 describes is real and the current >=0.1.6 floor permits it. Happy to fold that into this PR or a follow-up.
2. 4.6.1 starts shipping the next deprecation wave early: tcgen05.OperandMajorMode now warns on import (tests/conftest.py:16 triggers it), and the VSA suite's warning count goes 274 → 394 vs 4.5.2. Alias migration (OperandMajorMode, and later SmemAllocator & friends per the CuTe-DSL deprecation page) is planned as a separate follow-up.

Verification on our side is complete: 4.5.x and 4.6.x are both green with this branch on SM80/89/90/100 (+ SM120 from the report above).

🤖 Generated with Claude Code

[coderabbitai]

🧹 Nitpick comments (2)

flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py (2)

2373-2375: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove redundant tensor allocation.

The outer allocation for tOrO_frg is immediately and unconditionally overwritten in the first iteration of the loop, making it dead code.

♻️ Proposed fix

-        tOrO_frg = cute.make_rmem_tensor((tOrO_t2r_shape, frg_count), self.pv_acc_dtype)
         for i in cutlass.range_constexpr(frg_count):
             tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py` around lines 2373 -
2375, Remove the outer `cute.make_rmem_tensor((tOrO_t2r_shape, frg_count),
self.pv_acc_dtype)` allocation for `tOrO_frg`; initialize it only within the
`for i in cutlass.range_constexpr(frg_count)` loop using the existing
per-iteration shape and dtype.

---

2120-2120: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove unassigned expression.

This cute.make_rmem_tensor allocation is not assigned to a variable and acts as dead code.

♻️ Proposed fix

-                cute.make_rmem_tensor(tSrScale_t2r_shape, Float32)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py` at line 2120, Remove
the unassigned cute.make_rmem_tensor call in the surrounding function, as it
allocates a tensor without using or storing the result. Preserve any assigned
tensor allocations and the surrounding control flow.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py`:
- Around line 2373-2375: Remove the outer
`cute.make_rmem_tensor((tOrO_t2r_shape, frg_count), self.pv_acc_dtype)`
allocation for `tOrO_frg`; initialize it only within the `for i in
cutlass.range_constexpr(frg_count)` loop using the existing per-iteration shape
and dtype.
- Line 2120: Remove the unassigned cute.make_rmem_tensor call in the surrounding
function, as it allocates a tensor without using or storing the result. Preserve
any assigned tensor allocations and the surrounding control flow.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 464617e5-f6c5-4965-9f26-3aaf5df40208

📥 Commits

Reviewing files that changed from the base of the PR and between d527621 and 871a177.

📒 Files selected for processing (4)

• flashinfer/cute_dsl/sparse/blk128/flash_fwd_sm100.py
• flashinfer/cute_dsl/sparse/blk128/pack_gqa.py
• flashinfer/cute_dsl/sparse/blk128/utils.py
• flashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py

[bkryu]

Approving b12x changes.

Sparse attention changes should still be reviewed

[YangXu1990uiuc]

/bot run

[flashinfer-bot]

GitLab MR !972 has been created, and the CI pipeline #58066677 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[FAILED] Pipeline #58066677: 4/20 passed

[YangXu1990uiuc]

CI triage for the record — the internal bot's "4/20 passed" verdict decomposes entirely into infra + pre-existing issues; no failure traces to this PR:

• 8 jobs = runner_system_failure (h100 ×2, rtx-pro-6000 ×2, multi_gpu_b300 ×2 died in <95s; 5090 ×2 reclaimed at ~4h) — same legs failed identically on pipeline #58064562 (different code, 30 min earlier).
• b300 ×2: 282/284 files passed. Fails: tests/moe/test_unified_moe.py (also failing on #58064562 → pre-existing) + test_trtllm_gen_fused_moe 7200s per-file timeout (known chronic).
• gb200 ×2: 283/284 passed; only test_unified_moe again.
• gb300 ×2: all 279 parallel-phase files passed; job then hit the SLURM wall-clock limit mid-way through the memory-heavy sequential MoE tail (CANCELLED DUE TO TIME LIMIT while running test_trtllm_cutlass_fused_moe) — no test failures observed.
• 4 passed = multi-node gb200/gb300 ×2.

Everything this PR touches passed on every leg that reached it: test_vsa_block_sparse ✅ on B300/GB200/GB300 (both cu129/cu130), b12x suites ✅ (plus the SM120 report above). External GitHub CI is 23/23 green including the H100 full-suite leg.

🤖 Generated with Claude Code