Skip to content

Add page_table support to cutlass blackwell#116

Open
cthi wants to merge 1 commit intometa-pytorch:mainfrom
cthi:export-D91059471
Open

Add page_table support to cutlass blackwell#116
cthi wants to merge 1 commit intometa-pytorch:mainfrom
cthi:export-D91059471

Conversation

@cthi
Copy link
Contributor

@cthi cthi commented Jan 21, 2026

Summary: When using the fmha API we don't support page_table and the paged biases for cutlass blackwell. This diff adds support for PagedBlockDiagonalCausalLocalPaddedKeysMask & PagedBlockDiagonalCausalWithOffsetPaddedKeysMask. Note that we are notably missing support for the decode kernel, but that's okay.

Differential Revision: D91059471

@meta-cla meta-cla bot added the cla signed label Jan 21, 2026
@meta-codesync
Copy link

meta-codesync bot commented Jan 21, 2026

@cthi has exported this pull request. If you are a Meta employee, you can view the originating Diff in D91059471.

Summary:

When using the fmha API we don't support `page_table` and the paged biases for cutlass blackwell. This diff adds support for `PagedBlockDiagonalCausalLocalPaddedKeysMask` & `PagedBlockDiagonalCausalWithOffsetPaddedKeysMask`.

Limitations:
- Missing support for the specialized decode kernel
- Only support `page_size == 128` due to paged being implemented with TMA.

Differential Revision: D91059471
@cthi cthi force-pushed the export-D91059471 branch from ed413a2 to 58aa6fb Compare January 21, 2026 16:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant