From f76679d88d482aac9aa01a40b88c35e7e337e9fd Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Tue, 2 Jun 2026 17:18:07 -0700 Subject: [PATCH] feat(inference): support sliding-window attention in the dynamic batching 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 Signed-off-by: shanmugamr1992 --- megatron/core/transformer/attention.py | 30 +++++++++++-- .../inference/engines/test_dynamic_engine.py | 42 +++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index b27f90c53d0..480f21a4e75 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -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, @@ -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. @@ -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, @@ -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 + else: + window_size = (-1, -1) + # Flash attn kernel. if not is_decode_only: q = q.squeeze(1) @@ -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: @@ -899,6 +914,7 @@ def flash_decode_and_prefill( seqlens_k, block_table, softmax_scale, + window_size=window_size, ) else: assert ( @@ -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) @@ -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 @@ -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) @@ -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, } diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 503b0fa52ae..23e489d956a 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -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): @@ -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() @@ -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"