Ragged prefill wrapper - #203
Open
NSagan271 wants to merge 7 commits into
Open
Conversation
Covers both stateless runners end-to-end: a graph captured once per bucket, re-planned before each replay against a different varlen layout with the segment count padded out by zero-length segments. Checked against a per-segment SDPA reference at head_dim=72 so the head-dim padding is in the loop too. Also pins the contract edges: capture-fns that predate ragged attention must not receive static_attn, a bucket that captured ragged attention must reject a replay with no cu_seqlens rather than replay a stale plan, and a config without max_segments_per_request skips the capture instead of failing it. Docs get the two capture bounds, the segment-vs-request distinction, and the plan-outside/run-inside rule with the seam each runner provides. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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.
What does this PR do?
Addresses an issue found during the review of #150: neither engine has support for a ragged attention flashinfer wrapper (needed to cuda-graph ragged attention), meaning that the encoder submodule had to do all of the plumbing for flashinfer ragged attention planing itself.
This PR adds an engine-provided
AttentionState, which is analogous to theBatchedCacheManager, except it only has functionsplanandrun, and is not tied to a static KV cache (e.g., for one-time encoding attention). Submodules declare whether they need ragged attention support via a newBaseSubmodulemethodget_ragged_attention_config. Right now, only the stateless engine, its corresponding cuda graph runner, and the piecewise cudagraph runner can provide andAttentionState, but it would be simple to migrate the KV cache engine as well.How was it tested?
A lot of unit tests from Claude.
Migrated BAGEL ViT encoder, benchmarked T2I on coriander and checked outputs were correct + that throughput increased slightly with cuda graphs.
Checklist
ruff check .passes