GPT-OSS: Hoist Layer-Invariant Attention Mask Preparation Outside Decoder Subfunctions

[abhishek-singh591]

GPT-OSS: Hoist Layer-Invariant Attention Mask Preparation Outside Decoder Loop

Summary

This change hoists GPT-OSS attention-mask preparation out of QEffGptOssDecoderLayer and out of the
decoder-layer loop. The model now computes the reusable full-attention and sliding-window additive
masks once before iterating over decoder layers, then passes the already-prepared mask matching each
layer’s attention type.

Problem

For GPT-OSS 120B decode, enabling ONNX subfunctions showed an approximately 13% inference performance
drop compared with non-subfunction compilation. Inspection of the exported decoder function showed
redundant, layer-invariant mask work inside QEffGptOssDecoderLayer; for example, nodes such as
Gather_604 appeared inside the repeated decoder subfunction.

GPT-OSS uses both full_attention and sliding_attention, but the actual full/sliding masks are
layer-invariant for a given forward call. Only the choice of which mask to pass depends on
config.layer_types[layer_idx] / self_attn.sliding_window.

Approach

The model forward now builds the base causal mask and sliding-window causal mask once, before the
decoder-layer loop.

It then prepares reusable additive masks outside the loop:

• full_attention_mask for full-attention layers.
• sliding_attention_mask for sliding-attention layers.
• For chunked prefill, the same sliding-window index slice is applied once before the loop.

Inside the decoder-layer loop, the code only selects the already-prepared mask based on the layer’s
attention type.

Old in-layer behavior:

torch.where(mask, MIN_MASKED_ATTENTION_VALUE, scores)

New prepared-mask behavior:

scores + prepared_mask

These are equivalent because prepared_mask contains 0 for allowed positions and
MIN_MASKED_ATTENTION_VALUE for masked positions.

This follows the same idea as the Qwen2.5 vision optimization: compute reusable masks outside the
repeated layer/block loop, then pass the selected prepared mask into each layer.