Add paged attention kernel for Spyre by kiszk · Pull Request #23 · torch-spyre/spyre-kernels

kiszk · 2026-06-26T06:20:40Z

This PR adds a paged attention kernel based on @ani300 's implementation.
This version uses a logical (PyTorch) tensor shape.

A test using ktir-cpu passed.

$ .venv/bin/python -m pytest tests/ktir/test_paged_attn.py -v
========================================================================================================================================================================================================================= test session starts =========================================================================================================================================================================================================================
platform linux -- Python 3.12.13, pytest-9.1.1, pluggy-1.6.0 -- /home/ishizaki/aiu-src/dt-inductor/spyre-kernels/.venv/bin/python
cachedir: .pytest_cache
rootdir: /home/ishizaki/aiu-src/dt-inductor/spyre-kernels
configfile: pyproject.toml
collected 2 items                                                                                                                                                                                                                                                                                                                                                                                                                                                     

tests/ktir/test_paged_attn.py::test_paged_attn_ktir[tensor_descriptor] SKIPPED (tensor_descriptor.ktir not generated yet (run scripts/gen_ktir.py paged_attn))                                                                                                                                                                                                                                                                                                  [ 50%]
tests/ktir/test_paged_attn.py::test_paged_attn_ktir[spyre_aware] PASSED                                                                                                                                                                                                                                                                                                                                                                                         [100%]

==================================================================================================================================================================================================================== 1 passed, 1 skipped in 0.16s =====================================================================================================================================================================================================================

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

fabianlim · 2026-06-26T13:06:06Z

+
+
+@triton.jit
+def _paged_attn_kernel_NHD_td(


I feel that having the _td is not really needed, if you need to import them into the same file we could do import as?

@fabianlim I follow this convention. Is it ok with removing _td? I am neutral on this.

ic sorry then lets leave it

fabianlim · 2026-06-26T13:18:15Z

+):
+    HD: tl.constexpr = H * BLOCK_D
+
+    cur_b = tl.program_id(0)


im abit confused with the original kernel, it is already in tensor descriptor, I thought this repo's scope is to rewrite kernels that were previously in pointer arithmetic

I see. I started working with tensor descriptor at first. I will drop original.py

fabianlim · 2026-06-26T13:20:36Z

+                    vv = tl.permute(v_g, (0, 2, 1, 3))             # (BLK_B, BLK_H, KV_BLOCK, D)
+
+                    # 4-D batched matmul over (BLK_B, BLK_H)
+                    scores = tl.dot(q, kT)                         # (BLK_B, BLK_H, BLOCK_Q, KV_BLOCK)


this is not standard triton. tl.dot cannot accept a 4D tensor. The tensor_descriptor.py needs to be runnable on other archs like GPU

sorry as clarified by @kiszk we have tl.core insert reshapes as per this code

fabianlim · 2026-06-26T17:05:54Z

+  straight in as the index, and `h_start` is the second index — yielding
+  `(BLK_B, KV_BLOCK, BLK_H, D)` directly, with **no reshape** of the index or of
+  the gathered result. This maps the gather onto the Spyre indirect-access tile
+  more directly than the flatten/reshape of the base form.


will be good to have a comment that need to revisit with the physical layouts

Thanks, I will add the comment.

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

kiszk · 2026-06-29T14:28:01Z

@fabianlim I addressed your comments except the first one.

fabianlim

comments are addressed

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

initial commit

c7e54a8

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

kiszk requested review from fabianlim and ohadeytan June 26, 2026 06:21

fabianlim reviewed Jun 26, 2026

View reviewed changes

address review comments

c59603e

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

fabianlim approved these changes Jun 30, 2026

View reviewed changes

retest

ba9a673

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>

fabianlim mentioned this pull request Jul 1, 2026

Add hand-written paged attention fixtures torch-spyre/triton#36

Closed

10 tasks

retest

3b224c7

Signed-off-by: Kazuaki Ishizaki <ishizaki@jp.ibm.com>



		@triton.jit
		def _paged_attn_kernel_NHD_td(

Uh oh!

Conversation

kiszk commented Jun 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiszk commented Jun 29, 2026

Uh oh!

fabianlim left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kiszk commented Jun 26, 2026 •

edited

Loading