Skip to content

[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
upstream-mainfrom
jh/310p-longctx-mask-oom-c460
Open

[310P] Size attention masks by actual context to avoid long-ctx OOM (v2, rebased on c460a15af)#24
adeepn wants to merge 1 commit into
upstream-mainfrom
jh/310p-longctx-mask-oom-c460

Conversation

@adeepn

@adeepn adeepn commented Jun 16, 2026

Copy link
Copy Markdown
Member

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 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 (observed: worker crash, "NPU out of
memory, tried to allocate 16 GiB", on a 131K-context model — the first inference dies in
get_attention_mask).

  • Extract mask construction into a _build_attn_mask() hook on the base builder (default behaviour
    unchanged for all other devices).
  • Override it in the 310P builder to return the dense mask only for PrefillNoCache, sized to the
    batch's actual max context via a new actual_max_seqlen arg. The size accounts for the alignment
    padding forward_prefill_310 folds into the last sequence (seq_len[-1] += num_actual_tokens - sum(seq_lens)), so the mask never under-covers the padded length.
  • 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 mask caches
    monotonically.

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 upstream
compressed-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

  • Unit (CPU, run in the 310P image): 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 covers
    the alignment-padded last sequence). 23 ut pass.
  • Hardware: Qwen3.6-35B-A3B FP16, TP=4, 4×310P3, image nightly-main-310p (c460a15af): without
    this 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.

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>
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