Repository navigation
Let continuous batching page the call, not the model - #48735
qgallouedec wants to merge 7 commits into
Conversation
`init_continuous_batching` put `paged|<impl>` on the model and only `stop()` took it off, so any ordinary forward while the engine ran landed in the paged kernel with no paged cache and raised. The engine already passes its cache and the packed-input metadata through the forward kwargs, so `flash_attention_forward` and `sdpa_attention_forward` can route to their paged counterpart when a call carries a `PagedAttentionCache` and stay ordinary otherwise. The model is then left on the implementation it was loaded with and serves both. `eager` keeps the `paged|` prefix: every model defines its own eager forward, so there is no single one to route. `attn_mask_is_needed` now keys off the implementation instead of the `paged|` prefix, which only eager still carries.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| # Dispatch to paged attention instead if there is a paged cache | ||
| if isinstance(kwargs.get("cache"), PagedAttentionCache): | ||
| return sdpa_attention_paged_forward( | ||
| module, query, key, value, attention_mask, dropout=dropout, scaling=scaling, **kwargs | ||
| ) |
There was a problem hiding this comment.
Do we really still need those redirects then? it just goes to the same issue but on a different condition as previously no?
There was a problem hiding this comment.
I don't think this can lead to any issues in the current form!
The root of the issue is that the attn_implementation flavour (paged vs. unpaged) is tied to the model. While I agree that the implementation itself (sdpa, etc) needs to be tied, the fact that it is paged or not has no reason to be. If there is a paged cache -> use paged fllavour. Otherwise -> don't.
Previously the fallback was awkward: fall back from paged to normal, without a set attention mask or cu_seqlen. Now, it's more logical, and I don't see how it could fail. Open to ideas tho!
There was a problem hiding this comment.
But this is still weird, why would the ordinary forward call carry a cache at all? What would happen if we did use_cache=False in the forward call
My concern is that there is again no clean separation, attention variations should not call between each other
There was a problem hiding this comment.
The ordinary forward never carries a paged cache. It carries a DynamicCache, or nothing with use_cache=False, and both take the plain path: the branch is isinstance(cache, PagedAttentionCache), and only the engine builds one of those. The test's train step runs with the default cache; use_cache=False means no cache kwarg at all, so it can't reach the branch.
On separation: the paged kernel cannot run without the paged cache's block table, so keying on that cache is keying on the one thing that decides the kernel. The alternative, an implementation switch per call, is a write to config._attn_implementation, which I don't think is a nice design.
If two dispatch sites is the objection, I can move it to one: AttentionInterface.__getitem__ returning a wrapper that routes sdpa/flash on the cache type, so the attention functions themselves stay untouched?
There was a problem hiding this comment.
I discussed with @remi-or on slack but essentially what we want to do here is to fuse paged| with their counterparts (e.g. eager is not working this current way), e.g. because switching between the internal property is not thread safe to switch just in between
This is a bandaid solution and it's essentially better to do this properly than to rush it now.
CI recapDashboard: View test results in Grafana |
What
init_continuous_batchingputspaged|<impl>on the model and onlystop()takes it off. So while the engine is alive, any ordinary forward on that model lands in the paged kernel with no paged cache:That blocks training a model while generating from it, which is the whole point of the pause mechanism in #48462 (finding 1 of my review there).
The change
Paged-ness is a property of the call, not of the model. The engine already passes its cache and all the packed-input metadata (
cu_seq_lens_q/k,read_index,write_index,block_table) through the forward kwargs. Soflash_attention_forwardandsdpa_attention_forwardroute to their paged counterpart when the call carries aPagedAttentionCache, and stay ordinary otherwise:The model then stays on the implementation it was loaded with and serves both paths.
switch_to_cb_friendly_attnstill upgrades to flash when available, it just no longer rewrites the implementation string.eagerkeeps thepaged|prefix: 334 modeling files define their owneager_attention_forwardand"eager"is never registered inALL_ATTENTION_FUNCTIONS, so there is no single forward to route.The mask needed no change. Flash builds its own from the cumulative seqlens, and for eager/sdpa the engine's 4D mask hits the early return at the top of
_preprocess_mask_arguments, socreate_causal_masknever runs.This is the direction that #48297 tried in a19f0d7 and dropped on review, flipped: the fallback there was paged -> plain when no cache, here it is plain -> paged when there is one.
Measured
Qwen2.5-0.5B, H100, generate 3 prompts / training forward+backward / generate again, all with the manager running:
Completions before the training step are byte-identical to main.