Skip to content

Reuse tile prefix in recurrence kernel - #29

Open
0z5a wants to merge 1 commit into
MoonshotAI:masterfrom
0z5a:codex/kernel2-tile-prefix
Open

0z5a wants to merge 1 commit into
MoonshotAI:masterfrom
0z5a:codex/kernel2-tile-prefix

Conversation

@0z5a

@0z5a 0z5a commented Aug 14, 2026 •

Copy link
Copy Markdown

Incremental review: Files changed against master.

Summary

  • Pass the varlen tile prefix built before kernel1 to kernel2.
  • Replace kernel2's per-CTA linear scan of preceding sequence lengths with a direct tile_prefix[seq_idx] lookup.
  • Keep the fixed-length path, workspace size, and launch count unchanged.

Validation

  • Built on RTX 5080 (SM120).
  • Exact Torch-reference tests: fixed length and ragged H=96.
  • 8/8 N=8 varlen state/dtype combinations passed.
  • N=256, T=8192, H=16: kernel2 563.84 us -> 502.66 us (10.9%); full forward median 0.782 ms -> 0.731 ms (6.5%).
  • N=1 full-forward median: 1.721 ms -> 1.720 ms (no regression).

Notes

Zero-length sequences are an existing unsupported input: they hang on the baseline too. This focused optimization does not change that behavior.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
@0z5a
0z5a force-pushed the codex/kernel2-tile-prefix branch from 4a13001 to d51fa1e Compare September 24, 2026 05:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant