Sliding-window path in FusedScaleMaskSoftmax silently discards the caller-provided attention mask
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
In `megatron/core/fusions/fused_softmax.py::forward_torch_softmax` the sliding-window branch **unconditionally overwrites** the caller-provided mask:
```python
# Generate causal mask if not given
sq, sk = input.size(2), input.size(3)
if self.window_size is not None:
mask = get_sliding_window_causal_mask(sq, sk, self.window_size) # <- discards the incoming mask
elif self.attn_mask_type == AttnMaskType.causal and mask is None and sq > 1:
...
mask = get_default_causal_mask(sq)
```
Note the asymmetry: the plain-causal branch is correctly guarded by `mask is None`, the sliding-window branch is not.
## Impact
For any model with sliding-window layers, a padding mask (micro_batch_size > 1) or a packed-sequence block-diagonal mask passed by the caller is silently ignored on the SWA layers: pad tokens / cross-document tokens receive attention, the loss is quietly polluted, and nothing raises.
We hit this while running full SFT of MiMo-V2.5 (39 of 48 layers are SWA with this torch-fallback path): the only safe configurations today are micro_batch_size=1 and no sequence packing, which is a significant efficiency loss for variable-length multimodal data.
## Repro sketch
Build `FusedScaleMaskSoftmax(..., attn_mask_type=AttnMaskType.causal, scaled_masked_softmax_fusion=False, window_size=(4, 0))`, pass zero scores of shape [2, 1, 8, 8] and a bool padding mask masking the last 3 keys of sample 1 — the returned probabilities at the padded keys are non-zero.
## Proposed fix
Compose instead of overwrite (both masks are bool with True = masked; `[b,1,sq,sk] | [sq,sk]` broadcasts):
```python
swa_mask = get_sliding_window_causal_mask(sq, sk, self.window_size)
mask = swa_mask if mask is None else (mask | swa_mask)
```
Behavior is unchanged for every caller that passes `mask=None`. PR with this fix + a unit test incoming (scope: the torch fallback path only; fused/TE kernel paths are untouched).
Contributor guide
Research direction
Start in megatron/core/fusions/fused_softmax.py at forward_torch_softmax and reproduce the sliding-window case with a caller-provided padding mask. Check the existing causal-mask handling and add coverage for combining the caller mask with the sliding-window mask; done means padded or cross-document keys remain masked while mask=None behavior is unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100