Conversation
Convert the flash-attention prefill kernel from block_ptr/pointer form to the tl.make_tensor_descriptor API (td-convert skill): - block_ptr.py -> tensor_descriptor.py (_fwd_kernel on descriptors; masks handled by descriptor block_shape / shape bounds). - wrapper.py updated to launch the TD kernel. - conversion-notes.md records the conversion decisions. - tests/triton/test_prefill_attention_td.py: numerical-equivalence test. - Remove the superseded block_ptr kernel, its kernel.ktir, and the old block_ptr-era tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Ohad Eytan <ohad.eytan1@ibm.com>
fabianlim
reviewed
Jul 9, 2026
| acc = acc / l_i[:, None] | ||
| acc = acc.to(Out.dtype.element_ty) | ||
|
|
||
| # The store clamps at shape=[cur_batch_seq_len, ...], so the partial tail |
Collaborator
There was a problem hiding this comment.
would be good to try to lower it with the triton fork to see if these index comparisons and tl.where will be lowered properly.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a tensor-descriptor (TD) version of the
prefill_attentionflash-attention kernel, converting it from theblock_ptr/pointer-arithmetic form to thetl.make_tensor_descriptorAPI via thetd-convertskill.Changes
block_ptr.py→tensor_descriptor.py:_fwd_kernelrewritten on tensor descriptors. Pointer arithmetic and masked loads/stores are replaced with descriptors whoseshape/block_shapehandle out-of-bounds rows and columns (zero-fill on load, clamp on store).wrapper.py: updated to launch the TD kernel.conversion-notes.md: records the conversion decisions.tests/triton/test_prefill_attention_td.py: numerical-equivalence test against the original.block_ptrkernel, itskernel.ktir, and the old block_ptr-era tests.Notes
This TD kernel is intended as a clean baseline: it will serve as the shared starting point for a follow-up spyre-aware conversion (the two operate on the same descriptor form).
🤖 Generated with Claude Code