[310P] Size attention masks by actual context to avoid long-ctx OOM (v2, rebased on c460a15af) - #24
Open
adeepn wants to merge 1 commit into
Open
[310P] Size attention masks by actual context to avoid long-ctx OOM (v2, rebased on c460a15af)#24adeepn wants to merge 1 commit into
adeepn wants to merge 1 commit into
Conversation
The base AscendAttentionMetadataBuilder builds a dense attention mask sized to max_model_len for every attention state. On 310P only PrefillNoCache consumes a dense mask; at long context the O(max_model_len^2) allocation OOMs the NPU (e.g. 131K -> tens of GiB fp16, observed as a worker crash on the first request). Extract mask construction into a _build_attn_mask() hook on the base builder (default behaviour unchanged) and override it in the 310P builder to return the dense mask only for PrefillNoCache, sized to the batch's actual max context (rounded up to the NZ tile) via a new actual_max_seqlen argument. Build the SplitFuse (chunked-prefill) mask directly as [num_query_tokens, key_len] instead of materializing and index-selecting from a [max_seqlen, max_seqlen] tensor. Grow the causal / non-causal mask caches monotonically. Add unit tests for the tile rounding, the direct splitfuse mask, and the 310P mask-dispatch hook. Signed-off-by: Viacheslav Bocharov <v@baodeep.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
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
Re-derive the 310P long-context attention-mask OOM fix onto current upstream (base
upstream-main@c460a15af).The base
AscendAttentionMetadataBuilder.build()builds a dense attention mask sized tomax_model_lenfor every attention state. On 310P onlyPrefillNoCacheconsumes a dense mask; atlong context the
O(max_model_len^2)allocation OOMs the NPU (observed: worker crash, "NPU out ofmemory, tried to allocate 16 GiB", on a 131K-context model — the first inference dies in
get_attention_mask)._build_attn_mask()hook on the base builder (default behaviourunchanged for all other devices).
PrefillNoCache, sized to thebatch's actual max context via a new
actual_max_seqlenarg. The size accounts for the alignmentpadding
forward_prefill_310folds into the last sequence (seq_len[-1] += num_actual_tokens - sum(seq_lens)), so the mask never under-covers the padded length.[num_query_tokens, key_len]instead ofmaterializing and index-selecting from a
[max_seqlen, max_seqlen]tensor; grow the mask cachesmonotonically.
Why a new PR
Supersedes #22. Upstream rewrote the 310P attention/mask path (new compressed-mask / splitfuse-v2
code); #22 no longer applies. This re-derives the same fix onto
c460a15af. Note: the upstreamcompressed-mask path (
_npu_flash_attention_v3) is unavailable on this CANN build, so the dense /splitfuse paths are still taken and the fix remains required.
Test
tests/ut/_310p/attention/test_attention_mask_310.py(tile rounding, direct splitfuse mask == index_select equivalence) and
tests/ut/_310p/attention/test_attention_v1_310.py(mask only for PrefillNoCache; sizing coversthe alignment-padded last sequence). 23 ut pass.
nightly-main-310p(c460a15af): withoutthis fix the first request OOMs in mask build; with it, loads, serves, KV budget ~1M tokens @131k,
needle retrieval at short context, decode 4.4 t/s.