From d10de26bb01e0daa06a063e422c155c4c112c3d2 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Wed, 11 Nov 2026 17:38:19 +0800 Subject: [PATCH] fix(flex): avoid decoding kernel selection for short GQA Default the shared FlexAttention wrapper to BACKEND=TRITON to avoid PyTorch 2.13 AUTO lowering short GQA rows to an invalid decoding kernel. Preserve explicit backend and legacy force selectors without mutating caller options. Cover option handling in CPU CI and compare short-row outputs and Q/K/V gradients with SDPA under wrapped and nested compilation. --- .github/workflows/ci.yml | 1 + aspec/models/ops/flex_attention.py | 9 ++++++ docs/user-guide/configuration.md | 8 +++++ tests/test_flex_attention.py | 36 +++++++++++++++++++++ tests/test_flex_attention_backend.py | 48 ++++++++++++++++++++++++++++ 5 files changed, 102 insertions(+) create mode 100644 tests/test_flex_attention_backend.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e1a5480..dedef4a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,6 +70,7 @@ jobs: tests/test_eagle4_moe.py \ tests/test_eagle4_validation.py \ tests/test_fa4_import_compat.py \ + tests/test_flex_attention_backend.py \ tests/test_open_source_entrypoints.py \ tests/test_release_audit.py \ tests/test_shipped_config_defaults.py \ diff --git a/aspec/models/ops/flex_attention.py b/aspec/models/ops/flex_attention.py index 0893c9e..94d3a06 100644 --- a/aspec/models/ops/flex_attention.py +++ b/aspec/models/ops/flex_attention.py @@ -18,6 +18,10 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. +# Adapted from TorchSpec 6c042a87140a84d13839e341ece2c5c3ada918bc, +# torchspec/models/ops/flex_attention.py. ASpec defaults to the standard Triton +# kernel while preserving explicit backend/legacy selectors: PyTorch 2.13 AUTO +# can select decoding with no valid kernel choice for short GQA training rows. from functools import lru_cache import torch @@ -82,6 +86,11 @@ def compile_friendly_flex_attention( value: torch.Tensor, **kwargs, ) -> torch.Tensor: + kernel_options = dict(kwargs.get("kernel_options") or {}) + # PyTorch rejects BACKEND combined with the legacy force-attention selector. + if "FORCE_USE_FLEX_ATTENTION" not in kernel_options: + kernel_options.setdefault("BACKEND", "TRITON") + kwargs["kernel_options"] = kernel_options # First call initialise singleton wrapper object, second call invokes the object method to return compiled flex attention # Do not use compiled version if already compiling forward (it raises issues) flex_attention_compiled = ( diff --git a/docs/user-guide/configuration.md b/docs/user-guide/configuration.md index 7a0f39b..13b0af8 100644 --- a/docs/user-guide/configuration.md +++ b/docs/user-guide/configuration.md @@ -104,3 +104,11 @@ or optimizer initialization. ### Draft activation and RoPE Dense DSpark/DFlash FFNs execute the configured `hidden_act` through Transformers' activation registry; the default remains SiLU. Their standard GQA path supports unscaled, linear and YaRN RoPE through `rope_scaling` or `rope_parameters`, using the pinned Transformers scaling functions. YaRN uses `original_max_position_embeddings` when provided and otherwise the configured maximum context. Unsupported GQA scaling types and partial rotary dimensions fail during configuration loading. MLA and compressed/indexer rotary contracts remain structure-specific. + +### FlexAttention kernel selection + +The shared `flex_attention` wrapper defaults to the standard Triton kernel +(`BACKEND=TRITON`). This avoids PyTorch 2.13's automatic decoding selection failing +with `NoValidChoicesError` for short GQA training rows. Explicit `BACKEND` choices +and the legacy `FORCE_USE_FLEX_ATTENTION` selector are preserved. GPU regressions +compare outputs and Q/K/V gradients with SDPA, including nested `torch.compile`. diff --git a/tests/test_flex_attention.py b/tests/test_flex_attention.py index f35caa6..905d1e8 100644 --- a/tests/test_flex_attention.py +++ b/tests/test_flex_attention.py @@ -513,5 +513,41 @@ def test_single_block_shape_still_builds(self): torch.testing.assert_close(built.to_dense(), eager.to_dense()) +@pytest.mark.parametrize("kernel_options", [None, {"FORCE_USE_FLEX_ATTENTION": True}], + ids=["default", "legacy-force"]) +@pytest.mark.parametrize("outer_compile", [False, True], ids=["wrapper", "nested-compile"]) +def test_short_gqa_training_matches_sdpa_in_output_and_gradients(kernel_options, outer_compile): + """AUTO selected decoding with no valid choice for this PyTorch 2.13 shape.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + from aspec.models.dflash import _create_dflash_mask_mod + + anchors = torch.arange(20, 36, device="cuda")[None] + keep = torch.ones_like(anchors, dtype=torch.bool) + mask_mod = _create_dflash_mask_mod(anchors, keep, 150, 7) + block_mask = create_block_mask(mask_mod, 1, None, 112, 262, device="cuda") + dense_mask = mask_mod(0, 0, torch.arange(112, device="cuda")[:, None], + torch.arange(262, device="cuda")[None, :])[None, None] + generator = torch.Generator(device="cuda").manual_seed(123) + actual = [torch.randn(*shape, device="cuda", dtype=torch.bfloat16, requires_grad=True, + generator=generator) + for shape in ((1, 32, 112, 128), (1, 8, 262, 128), (1, 8, 262, 128))] + reference = [value.detach().clone().requires_grad_() for value in actual] + attention = (torch.compile(compile_friendly_flex_attention, fullgraph=True) + if outer_compile else compile_friendly_flex_attention) + output = attention(*actual, block_mask=block_mask, enable_gqa=True, kernel_options=kernel_options) + with sdpa_kernel(SDPBackend.MATH): + expected = torch.nn.functional.scaled_dot_product_attention( + *reference, attn_mask=dense_mask, enable_gqa=True) + output.float().square().mean().backward() + expected.float().square().mean().backward() + for value, target in [(output, expected), *zip( + [value.grad for value in actual], [value.grad for value in reference], strict=True + )]: + assert torch.isfinite(value).all() + relative = (value.float() - target.float()).norm() / target.float().norm().clamp_min(1e-12) + assert relative < 0.01, float(relative) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_flex_attention_backend.py b/tests/test_flex_attention_backend.py new file mode 100644 index 0000000..3dc03d6 --- /dev/null +++ b/tests/test_flex_attention_backend.py @@ -0,0 +1,48 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Backend policy must preserve explicit choices in both compilation paths.""" + +import copy + +import pytest +import torch + +import aspec.models.ops.flex_attention as flex_mod + + +@pytest.mark.parametrize("compiling", [False, True], ids=["wrapper", "nested-compile"]) +@pytest.mark.parametrize("options,expected", [ + ({}, {"BACKEND": "TRITON"}), + ({"kernel_options": None}, {"BACKEND": "TRITON"}), + ({"kernel_options": {}}, {"BACKEND": "TRITON"}), + ({"kernel_options": {"BLOCK_M": 64}}, {"BLOCK_M": 64, "BACKEND": "TRITON"}), + ({"kernel_options": {"BACKEND": "AUTO"}}, {"BACKEND": "AUTO"}), + ({"kernel_options": {"BACKEND": "TRITON"}}, {"BACKEND": "TRITON"}), + ({"kernel_options": {"BACKEND": "TRITON_DECODE"}}, {"BACKEND": "TRITON_DECODE"}), + ({"kernel_options": {"BACKEND": "FLASH"}}, {"BACKEND": "FLASH"}), + ({"kernel_options": {"FORCE_USE_FLEX_ATTENTION": True}}, {"FORCE_USE_FLEX_ATTENTION": True}), + ({"kernel_options": {"FORCE_USE_FLEX_ATTENTION": False}}, {"FORCE_USE_FLEX_ATTENTION": False}), +], ids=["omitted", "none", "empty", "tuning", "auto", "triton", "decode", "flash", + "legacy-force", "legacy-auto"]) +def test_kernel_backend_policy_preserves_caller_choices(monkeypatch, compiling, options, expected): + original = copy.deepcopy(options) + inputs = [torch.empty(1) for _ in range(3)] + result = object() + calls = [] + + def attention(*args, **kwargs): + calls.append((args, kwargs)) + return result + + monkeypatch.setattr(flex_mod, "is_torchdynamo_compiling", lambda: compiling) + monkeypatch.setattr(flex_mod, "WrappedFlexAttention", lambda: lambda: attention) + monkeypatch.setattr(flex_mod, "flex_attention", attention) + assert flex_mod.compile_friendly_flex_attention(*inputs, enable_gqa=True, **options) is result + assert len(calls) == 1 + args, kwargs = calls[0] + assert all(actual is supplied for actual, supplied in zip(args, inputs, strict=True)) + assert kwargs == {"enable_gqa": True, "kernel_options": expected} + assert options == original + if options.get("kernel_options") is not None: + assert kwargs["kernel_options"] is not options["kernel_options"]