NFC: replace deprecated API: cute.make_fragment

[brandon-yujie-sun]

📌 Description

More mechanical change to replace deprecated API:

• cute.make_fragment --> cute.make_rmem_tensor

🔍 Related Issues

N/A

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

N/A

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

N/A

Reviewer Notes

N/A

Summary by CodeRabbit

• Refactor
  • Updated internal tensor representation in GPU kernel epilogue operations for block-scaled matrix computations to improve implementation consistency.

[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: 956e022b-c347-4eef-bad6-0053ad726ff9

📥 Commits

Reviewing files that changed from the base of the PR and between fc12ef2 and b0213f9.

📒 Files selected for processing (2)

• flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py
• flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py

---

📝 Walkthrough

Walkthrough

This PR migrates epilogue register tensor constructions from cute.make_fragment(...) to cute.make_rmem_tensor(...) in SM100 GEMM kernels. Changes affect C-output, accumulator, and top-k intermediate tensor representations across dense_blockscaled_gemm_sm100.py and grouped_gemm_masked_blackwell.py, updating downstream copy/partition/store operations without altering control flow.

Changes

Register Tensor API Migration in SM100 Kernels

Layer / File(s)	Summary
C-output register tensor conversion flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py, flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py	tTR_rC is migrated from cute.make_fragment(...) to cute.make_rmem_tensor(...) in epilogue warps, affecting register/shared/global copy and partition operations.
Top-k intermediate register tensor conversion flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py	topk_rFrag is converted from fragment-based to rmem-backed tensor via cute.make_rmem_tensor(...), changing the tensor type used in combine-fusion accumulation and retiling.
Accumulator register tensor conversion in epilogue helper flashinfer/gemm/kernels/dense_blockscaled_gemm_sm100.py, flashinfer/gemm/kernels/grouped_gemm_masked_blackwell.py	tTR_rAcc returned from epilog_tmem_copy_and_partition is migrated from cute.make_fragment(...) to cute.make_rmem_tensor(...), updating tensor representation in tmem-to-register epilogue path.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#3430: Contains parallel cute.make_fragment → cute.make_rmem_tensor API replacements in GEMM epilogue tensor constructions.

Suggested labels

run-ci

Suggested reviewers

• bkryu
• dhiraj113
• yzh119
• aleozlx

Poem

🐰 A rabbit hops through GPU code,
Where tensors find a new abode—
From fragments old to rmem's embrace,
The SM100 picks up pace!
Two kernels bloom with cleaner ways, ✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: replacing the deprecated cute.make_fragment API with cute.make_rmem_tensor.
Description check	✅ Passed	The description includes the required template structure with a clear description section explaining the mechanical API replacement, though pre-commit and test checklist items are marked N/A.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 and usage tips.}

[gemini-code-assist]

Code Review

This pull request replaces several calls to cute.make_fragment with cute.make_rmem_tensor in the GEMM kernels dense_blockscaled_gemm_sm100.py and grouped_gemm_masked_blackwell.py. No review comments were provided, and there is no additional feedback.

[aleozlx]

/bot run /tests/gemm

[flashinfer-bot]

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

[flashinfer-bot]

[FAILED] Pipeline #53265395: 9/20 passed