Skip to content

perf(kernel): split triton decode attention by context length and sm count #729

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jason-fxz merged 2 commits into from
Oct 17, 2026
+172 −90
Merged

perf(kernel): split triton decode attention by context length and sm count #729

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
855aedd
perf(kernel): split triton decode attention by context length and sm …
jason-fxz Oct 16, 2026
2bd6c33
perf(kernel): merge triton decode splits a tile at a time
jason-fxz Oct 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 26 additions & 39 deletions python/freetoken/attention/triton.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import torch
from freetoken.core import Batch, get_global_ctx
from freetoken.kernel.triton.attention import decode_kv_split_cap

from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
from .utils import BaseCaptureData
Expand All @@ -16,9 +17,9 @@
@dataclass
class TritonCaptureData(BaseCaptureData):
q_to_req: torch.Tensor
# flat decode split scratch: each graph batch size views it with its own split count
attn_logits: torch.Tensor
attn_lse: torch.Tensor
num_kv_splits: torch.Tensor
swa_page_table: torch.Tensor | None = None

@classmethod
Expand All @@ -30,7 +31,7 @@ def create(
*,
num_q_heads: int,
max_head_dim: int,
max_kv_splits: int,
scratch_rows: int,
**kwargs,
):
return cls(
Expand All @@ -40,22 +41,8 @@ def create(
cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device),
attn_logits=torch.empty(
(max_bs, num_q_heads, max_kv_splits, max_head_dim),
dtype=torch.float32,
device=device,
),
attn_lse=torch.empty(
(max_bs, num_q_heads, max_kv_splits),
dtype=torch.float32,
device=device,
),
num_kv_splits=torch.full(
(max_bs,),
max_kv_splits,
dtype=torch.int32,
device=device,
),
attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device),
attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device),
**kwargs,
)

Expand All @@ -72,7 +59,6 @@ class TritonMetadata(BaseAttnMetadata):
max_q_len: int
attn_logits: torch.Tensor | None = None
attn_lse: torch.Tensor | None = None
num_kv_splits: torch.Tensor | None = None
swa_indices: torch.Tensor | None = None

def get_last_indices(self, bs: int) -> torch.Tensor:
Expand All @@ -87,14 +73,18 @@ def __init__(self, config: ModelConfig):
self.capture: TritonCaptureData | None = None
self.capture_bs: List[int] = []
self.max_graph_bs = 0
self.max_kv_splits = 8
self.prefill_tile_min_q = 128
self.num_q_heads = int(getattr(config, "num_qo_heads", 1))
kv_groups = getattr(config, "kv_cache_group_specs", lambda: ())()
self.max_head_dim = max(
(group.head_dim for group in kv_groups),
default=int(getattr(config, "head_dim", 1)),
)
self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))]

def _split_cap(self, bs: int) -> int:
"""The most decode KV splits any layer group takes at this batch size."""
return max(decode_kv_split_cap(bs, max(1, self.num_q_heads), heads, self.device.index) for heads in self.kv_heads)

def _ensure_decode_scratch(
self,
Expand All @@ -103,32 +93,27 @@ def _ensure_decode_scratch(
num_q_heads: int,
head_dim: int,
) -> None:
splits = self._split_cap(bs)
if (
metadata.attn_logits is not None
and metadata.attn_lse is not None
and metadata.num_kv_splits is not None
and metadata.attn_logits.shape[0] >= bs
and metadata.attn_logits.shape[1] >= num_q_heads
and metadata.attn_logits.shape[2] >= splits
and metadata.attn_logits.shape[3] >= head_dim
):
return
scratch_head_dim = max(self.max_head_dim, head_dim)
metadata.attn_logits = torch.empty(
(bs, num_q_heads, self.max_kv_splits, scratch_head_dim),
(bs, num_q_heads, splits, scratch_head_dim),
dtype=torch.float32,
device=self.device,
)
metadata.attn_lse = torch.empty(
(bs, num_q_heads, self.max_kv_splits),
(bs, num_q_heads, splits),
dtype=torch.float32,
device=self.device,
)
metadata.num_kv_splits = torch.full(
(bs,),
self.max_kv_splits,
dtype=torch.int32,
device=self.device,
)

def forward(
self,
Expand Down Expand Up @@ -167,7 +152,6 @@ def forward(
self._ensure_decode_scratch(metadata, bs, q.shape[1], q.shape[-1])
assert metadata.attn_logits is not None
assert metadata.attn_lse is not None
assert metadata.num_kv_splits is not None
return decode_paged_attention(
q=q,
k_cache=k_cache,
Expand All @@ -177,8 +161,6 @@ def forward(
q_positions=metadata.q_positions,
attn_logits=metadata.attn_logits[:bs],
attn_lse=metadata.attn_lse[:bs],
num_kv_splits=metadata.num_kv_splits[:bs],
max_kv_splits=self.max_kv_splits,
sm_scale=scale,
sliding_window=spec.sliding_window,
sinks=spec.sinks,
Expand Down Expand Up @@ -280,7 +262,7 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
self.device,
num_q_heads=max(1, self.num_q_heads),
max_head_dim=max(1, self.max_head_dim),
max_kv_splits=self.max_kv_splits,
scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list),
)
if self._swa_capture_enabled():
self.capture.swa_page_table = torch.zeros(
Expand Down Expand Up @@ -324,14 +306,20 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None:
metadata.indices = indices
metadata.q_to_req = self.capture.q_to_req[: metadata.q_to_req.numel()]
metadata.q_positions = self.capture.positions[:q_tokens]
metadata.attn_logits = self.capture.attn_logits[:bs]
metadata.attn_lse = self.capture.attn_lse[:bs]
metadata.num_kv_splits = self.capture.num_kv_splits[:bs]
metadata.attn_logits, metadata.attn_lse = self._capture_scratch(bs)

def _capture_scratch(self, bs: int) -> tuple[torch.Tensor, torch.Tensor]:
assert self.capture is not None
splits, heads, dim = self._split_cap(bs), max(1, self.num_q_heads), max(1, self.max_head_dim)
logits = self.capture.attn_logits[: bs * heads * splits * dim].view(bs, heads, splits, dim)
lse = self.capture.attn_lse[: bs * heads * splits].view(bs, heads, splits)
return logits, lse

def prepare_for_capture(self, batch: Batch) -> None:
bs = batch.size
assert bs in self.capture_bs and self.capture is not None
capture = self.capture
attn_logits, attn_lse = self._capture_scratch(bs)
batch.attn_metadata = TritonMetadata(
cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1],
indptr=capture.cu_seqlens_k[: bs + 1],
Expand All @@ -341,9 +329,8 @@ def prepare_for_capture(self, batch: Batch) -> None:
is_decode=True,
prefix_lens=capture.positions[:bs],
max_q_len=1,
attn_logits=capture.attn_logits[:bs],
attn_lse=capture.attn_lse[:bs],
num_kv_splits=capture.num_kv_splits[:bs],
attn_logits=attn_logits,
attn_lse=attn_lse,
swa_indices=self._capture_swa_indices(),
)

Expand Down
Loading
Loading