Support SWA and sink attention in dynamic inference - #5249
Conversation
…hing path The dynamic batching engine calls flash_decode_and_prefill, which dispatched to FA2/FA3/FA4 with window_size hardcoded to full attention. The static path already honors config.window_size via the TE wrapper; this change brings the dynamic path to parity. For each call, resolve the per-layer window via is_layer_window_attention (the same helper the TE static path uses), then plumb the (left, right) tuple to every kernel: flash_attn4_varlen_func, the FA3 _flash_attn_forward wrapper, flash_attn_varlen_func, flash_attn3_with_kvcache, and flash_attn_with_kvcache (decode). FlashMLA does not support SWA, so the MLA branch asserts window_size == (-1, -1). Tested on H100/FA3 via cog with three SWA configs, including the gpt-oss configuration (window_size=(127, 0), window_attn_skip_freq=2). The regression test_simple (no-SWA path) still passes. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
…200. Register the test in moe-dynamic-inference recipes and bootstrap platform goldens; apply YaRN config in gpt_builders when running inference without post-training GPT-OSS flags.
Gptoss 20b swa functional test
|
/ok to test f81aa44 |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if need_lse: | ||
| # FA2/FA3 *_with_kvcache return (out, softmax_lse) when | ||
| # return_softmax_lse=True. | ||
| output_total, softmax_lse = kvcache_ret | ||
| else: | ||
| output_total = kvcache_ret | ||
| softmax_lse = None | ||
| if need_lse: | ||
| # output_total: (B, S, H, D); softmax_lse: (B, H, S) | ||
| output_total = self._apply_sink_softmax_correction_bshd( | ||
| output_total, softmax_lse, softmax_offset | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Simplification] Two consecutive if need_lse: blocks — the first unpacks the return tuple and the second applies the correction. These can be merged into one block:
if need_lse:
# FA2/FA3 *_with_kvcache return (out, softmax_lse) when
# return_softmax_lse=True.
output_total, softmax_lse = kvcache_ret
# output_total: (B, S, H, D); softmax_lse: (B, H, S)
output_total = self._apply_sink_softmax_correction_bshd(
output_total, softmax_lse, softmax_offset
)
else:
output_total = kvcache_ret
There was a problem hiding this comment.
Review Summary
CRITICAL: 0 | IMPORTANT: 1 | SUGGESTION: 3
Overview
This PR adds sliding-window attention (SWA) and sink/off-by-one/learnable softmax correction to the dynamic-batching inference path, along with YaRN RoPE mscale plumbing and a softmax_scale fix for the FA2/FA3 kvcache decode path. The implementation is solid — the core sink-softmax math (out *= sigmoid(lse - offset)) is correct, the NaN/inf guard is well-reasoned, and all flash-attention backends (FA2, FA3, FA4, FlashMLA) are handled consistently.
Most impactful finding
Hardcoded YaRN hyperparameters in gpt_builders.py — yarn_original_max_position_embeddings=4096, yarn_beta_fast=32.0, yarn_beta_slow=1.0 are model-specific (GPT-OSS-20B) but applied unconditionally whenever position_embedding_type == 'yarn', including when the config was built from YAML. Other YaRN params (mscale, mscale_all_dim) are correctly read from args, making this inconsistent.
What looks good
- The sink-softmax post-correction derivation is mathematically clean and matches the canonical
SoftmaxOneformulation used by the static path. - The
isfiniteguard correctly distinguishes NaN LSE (padding artifacts → preserve row) from-infLSE (no attended keys → zero the row). - SWA window resolution via
is_layer_window_attention()correctly handles all three modes: globally disabled, every layer, and per-layer skip frequency. - The FA3
_flash_attn_forwardwrapper's version-adaptive signature introspection handles bothwindow_sizetuple andwindow_size_left/window_size_rightscalar styles. - The
softmax_scaleaddition to the FA2/FA3 kvcache decode path is a legitimate bugfix (was missing before this PR). - The FlashMLA assertion guarding against SWA is the right approach — fail fast with a clear message.
- Test coverage is thorough: math-only unit tests, monkeypatched wiring verification, and functional tests with golden values on both H100 and GB200.
Risk assessment: Low
No critical issues. The changes are well-scoped to inference-only paths (assert not self.training), so training behavior is completely unaffected. The new softmax_offset parameter defaults to None, preserving existing behavior when softmax_type == "vanilla". The only behavioral change for existing users is the softmax_scale bugfix in the FA2/FA3 kvcache decode path — but this was already broken for custom-scale models, so the fix is strictly an improvement.
…pr5249 refactor(attention): cache YaRN concentration factor in Attention.__i…
|
/ok to test a994af9 |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28064465483 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28065394213 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28068207953 |
Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Signed-off-by: Vitaly Kurin <vitalyk@nvidia.com>
Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com> Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
8 participants
Summary
This is a clean combined branch for PR #5138 and PR #5140, rebased onto current
origin/main.It includes:
mscalethrough dynamic key and query rotary-embedding application.softmax_scaleinto the non-MLA KV-cache kernels.Validation
git diff --check origin/main..HEADunsloth/gpt-oss-20b-BF16weights:374prompt tokens[127, 0],window_attn_skip_freq=2[623]/" The"[623]/" The"unsloth/gpt-oss-20b-BF16weights:374prompt tokensmax_new_tokens=8[623, 1825, 10648, 1606, 290, 2461, 50005, 4580]" The user wants only the run‑time"