FlexGptOssInference mirrors FlexMoEInference with three arch-specific
pieces:
1. Attention sinks via flex_attention(return_lse=True) + post-softmax
sigmoid(lse - sinks[h]) scaling. Same math as
unsloth_zoo.flex_attention.attention_sink.
2. Per-layer sliding window: the walker passes both a full and a
sliding-128 BlockMask; each attention forward picks one based on
self.sliding_window.
3. gpt-oss rotary: first/second-half split with head_dim/2-sized
cos/sin, not the Llama-style rotate_half on full-dim.
Reuses:
- refresh_moe_lora_merge_from_pristine for stacked-expert LoRA merge
(transposed orientation branch covers gpt-oss's (E, H, 2I) layout)
- forward_native_grouped_mm — the GptOssExperts branch already handles
the interleaved gate/up split + gate * sigmoid(gate * 1.702) activation
bnb-4bit: GptOssExpertsBnb4bit uses an nn.ModuleList per-expert loop
that can't be CUDA-graph-captured. __init__ detects it and disables
capture; decode still benefits from paged KV + flex_attention.
Also fixes a pre-existing deepcopy recursion in _LazyFlexEngineSentinel
that surfaced when the flex inference deepcopy ran on a model still
holding the sentinel (vision.py path doesn't pre-seed the inference
copy).
Bench on B200 (unsloth/gpt-oss-20b-BF16, max_new_tokens=64, chat):
| bs | HF naive | flex walker |
|---:|---------:|------------:|
| 8 | 258.9 | 832.1 |
| 32 | — | 1595.3 |
| 48 | — | 2146.4 |
3.2x over HF naive at bs=8. Parity: 24/24 greedy-token match vs HF on
3 chat prompts; cudagraph capture replay bitwise-matches eager decode.
1b5ff90e49