Skip to content

swa-decode-window

Source: fixes/swa-decode-window/

Make the sliding window exact in paged decode for gpt-oss, instead of 129 to 256 visible tokens for a 128-token window.

Affects gpt-oss (20B and 120B)
Releases v0.26.0, v0.30.0
Patches vllm_gaudi/v1/worker/hpu_model_runner.py
Upstream not filed yet
Test test_swa_decode_window.py (CPU, needs torch)

Problem

gpt-oss sliding layers attend to the last 128 tokens, the current one included (keys k with k > q - 128). For those layers the HPU decode path hands the sequence's last window // block_size + 1 KV blocks to attention (two blocks at block size 128) and leaves every filled slot of both visible. A sliding layer therefore sees between 129 and 256 tokens, depending on where the current token sits in its block.

Symptom

Decode drifts from the reference model as soon as a sequence is longer than the window: slightly worse long-generation quality and tool-call accuracy, and prefill (which masks the window exactly) disagrees with decode. No error is raised.

Who is affected

gpt-oss on HPU, for every decode step past the first window. Other sliding-window models take window // block_size blocks in decode and are not changed by this fix.

The fix

When the window is a multiple of the block size and a sequence's window block list has full length, the oldest block's usage is stored as -lbu, where lbu is the usage of the sequence's last block. The decode bias reads a negative usage -u as "slots [u, block_size) are valid". Those are exactly the window positions inside the oldest block, so the visible set is (q - window, q]. Positive usages build the same mask as before, so every other model and path is unchanged. When the last block is full, the oldest block is masked completely; gpt-oss always has sinks, so its softmax stays finite.

Validated on hardware with v0.26.0. On v0.30.0 the same edit applies (the window call also passes force_non_contiguous=True); the logic is identical.

Verify

python3 fixes/swa-decode-window/apply.py --check   # 0 = applied, 1 = not applied, 2 = source drifted
pytest fixes/swa-decode-window

The test runs the inserted usage and bias code on CPU and checks that the visible positions equal the window for every decode position, at block sizes 64 and 128.