Skip to content

feat(inference): support sliding-window attention in the dynamic batc… #5138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shanmugamr1992 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(inference): support sliding-window attention in the dynamic batc… #5138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions megatron/core/transformer/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from megatron.core.transformer.identity_op import IdentityOp
from megatron.core.transformer.module import MegatronModule
from megatron.core.transformer.torch_norm import L2Norm, LayerNormBuilder
from megatron.core.transformer.utils import is_layer_window_attention
from megatron.core.typed_torch import apply_module, not_none
from megatron.core.utils import (
deprecate_inference_params,
Expand Down Expand Up @@ -775,6 +776,7 @@ def _flash_attention_3_forward_wrapper(
seqlens_k,
block_table,
softmax_scale,
window_size: Tuple[int, int] = (-1, -1),
):
"""
Wrapper for calling the FA3 _flash_attn_forward function.
Expand Down Expand Up @@ -809,9 +811,9 @@ def _flash_attention_3_forward_wrapper(
"causal": True,
"attention_chunk": 0,
"softcap": 0.0,
"window_size": (-1, -1),
"window_size_left": -1,
"window_size_right": -1,
"window_size": window_size,
"window_size_left": window_size[0],
"window_size_right": window_size[1],
"rotary_interleaved": True,
"scheduler_metadata": None,
"num_splits": 0 if not self.batch_invariant_mode else 1,
Expand Down Expand Up @@ -865,6 +867,18 @@ def flash_decode_and_prefill(
assert not self.training
assert block_table is not None

# Resolve sliding-window-attention size for this layer.
# `config.window_size` is a (left, right) tuple, where -1 means infinite
# window in that direction (i.e. full attention). When SWA is not active
# for this layer (either globally disabled, or the layer is a "full
# attention" layer per `window_attn_skip_freq`), fall back to (-1, -1).
if is_layer_window_attention(
self.config.window_size, self.config.window_attn_skip_freq, self.layer_number
):
window_size = self.config.window_size

@cuichenx cuichenx Jun 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

While validating GPT-OSS/YARN on this dynamic path, I found another required plumbing fix adjacent to the SWA window plumbing: dynamic RoPE currently needs the YaRN concentration factor passed through. Static and dynamic matched through hidden input and QKV projection, then diverged immediately after RoPE unless both dynamic key and query RoPE calls pass mscale=_yarn_get_concentration_factor_from_config(self.config) into inference_context.apply_rotary_emb_key(...) / apply_rotary_emb_query(...). Can you include that here or land it as a prerequisite? Without it, SWA+YARN models can diverge before the attention kernel sees matching Q/K.

else:
window_size = (-1, -1)

# Flash attn kernel.
if not is_decode_only:
q = q.squeeze(1)
Expand All @@ -884,6 +898,7 @@ def flash_decode_and_prefill(
page_table=block_table,
softmax_scale=softmax_scale,
causal=True,
window_size=window_size,
num_splits=1,
)
elif HAVE_FA3:
Expand All @@ -899,6 +914,7 @@ def flash_decode_and_prefill(
seqlens_k,
block_table,
softmax_scale,
window_size=window_size,
)
else:
assert (
Expand All @@ -914,6 +930,7 @@ def flash_decode_and_prefill(
max_seqlen_k,
softmax_scale=softmax_scale,
causal=True,
window_size=window_size,
block_table=block_table,
)
output_total = output_total.unsqueeze(1)
Expand All @@ -929,6 +946,11 @@ def flash_decode_and_prefill(
# The `softmax_scale` attribute check is to find out whether this is an MLA layer or
# standard Attention.
if isinstance(self.config, MLATransformerConfig) and hasattr(self, "softmax_scale"):
# FlashMLA does not currently support sliding window attention.
assert window_size == (-1, -1), (
"FlashMLA decode kernel does not support sliding window attention. "
"Set config.window_size = None or use a non-MLA attention layer."
)
softmax_scale = self.softmax_scale

num_heads_k = 1 # Only a single head for MLA Flash
Expand Down Expand Up @@ -974,6 +996,7 @@ def flash_decode_and_prefill(
page_table=block_table,
softmax_scale=softmax_scale,
causal=True,
window_size=window_size,
num_splits=1,
)
# Reshape back to (B, S, H, D)
Expand All @@ -987,6 +1010,7 @@ def flash_decode_and_prefill(
"v_cache": v,
"cache_seqlens": seqlens_k,
"causal": True,
"window_size": window_size,
"page_table" if HAVE_FA3 else "block_table": block_table,
"num_splits": 0 if not self.batch_invariant_mode else 1,
}
Expand Down
42 changes: 42 additions & 0 deletions tests/unit_tests/inference/engines/test_dynamic_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,12 @@ class DynamicEngineTestConfig:
num_speculative_tokens: int = 0
position_embedding_type: str = "learned_absolute"
sampling_backend: str = 'torch'
# Sliding-window attention config. When `window_size` is None, SWA is
# disabled and all layers do full causal attention. When set to a
# `(left, right)` tuple, layers selected by `window_attn_skip_freq` use a
# local window of `left` past tokens and `right` future tokens.
window_size: Optional[Tuple[int, int]] = None
window_attn_skip_freq: Optional[int] = None

def __post_init__(self):

Expand Down Expand Up @@ -371,6 +377,8 @@ def _build_test_env(cls, test_config):
else "LayerNorm"
),
# inference optimized currently only supports RMS Norm
window_size=test_config.window_size,
window_attn_skip_freq=test_config.window_attn_skip_freq,
)
if test_config.fp8 or test_config.transformer_impl == "transformer_engine":
layer_spec = get_gpt_layer_with_transformer_engine_spec()
Expand Down Expand Up @@ -860,6 +868,40 @@ def test_multi_add(self, model_provider: str) -> None:
skip_if_mamba_sequence_packing_not_available(model_provider)
self._run_test(num_gap_steps=0, model_provider=model_provider)

@pytest.mark.internal
@pytest.mark.skipif(
not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
)
@pytest.mark.parametrize(
# Cover three regimes:
# - SWA active on every layer (window_attn_skip_freq=None)
# - SWA active on a subset of layers (gpt-oss style: every other layer)
# - window smaller than the longest sequence we generate, so the
# kernel actually applies the local-attention mask.
"window_size,window_attn_skip_freq",
[((4, 0), None), ((4, 0), 2), ((127, 0), 2)],
)
def test_sliding_window_attention(
self, window_size: Tuple[int, int], window_attn_skip_freq: Optional[int]
) -> None:
"""Exercise SWA on the dynamic batching (FA2/FA3/FA4) attention path.

This mirrors the gpt-oss configuration (window 127 to the left, no
future tokens, applied every other layer) at a much smaller scale.
The test only checks that decoding runs end-to-end and produces the
expected number of tokens; numerical correctness of the SWA kernels
themselves is owned by the upstream flash-attention test suites.
"""
self._run_test(
model_provider="gpt",
num_gap_steps=0,
window_size=window_size,
window_attn_skip_freq=window_attn_skip_freq,
# Disable CUDA graphs: this test only validates the SWA plumbing
# through the attention kernel, not the CG capture path.
num_cuda_graphs=None,
)

@pytest.mark.internal
@pytest.mark.skipif(
not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching"
Expand Down