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

Reuse tile prefix in recurrence kernel#29
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

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.

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