fix(flex): avoid decoding kernel selection for short GQA

.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 \

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 = (

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`.

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)

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"]