Keep FlashAttention inside captured CUDA graphs

docs/DEPLOYMENT.md

@@ -357,6 +357,17 @@ because their PEFT embedding adapter allocates tensors during each forward.
 eagerly under `acceleration.cuda_graphs`. Direct graphs cannot be combined with
 `JEVANY_COMPILE`.
 
+Transformers counts CUDA stream capture as tracing. Left alone, it materialises
+the causal mask inside every captured graph, so SDPA runs its memory-efficient
+kernel there instead of the FlashAttention kernel that eager inference uses.
+During capture JevAny therefore makes the eager mask decisions for an unpadded
+single row; rows whose positions do not rise by one run eagerly. On an
+A100-40GB with BF16 merged JevAny-Qwen3.5-4B, this made captured public-JevBench
+requests of 513 to 2,048 tokens 3–5% faster and left shorter ones unchanged.
+With a 4,096-token limit, requests of 2,049 to 4,096 tokens became 6–10% faster
+in graphs, but they were still no faster than eager inference, so the default
+limit stays at 2,048 tokens.
+
 Padded shapes select different kernels, so probabilities are close to, not
 identical with, the eager path. On an A100-40GB with FLA and causal-conv1d
 installed and SDPA, over every request of the development suites, median

jevany/cudagraphs.py

@@ -5,15 +5,20 @@
 Python between them. ``RowGraphs`` captures one graph per padded length bucket, all sharing one memory pool, and
 replays the smallest bucket that fits each row.
 
-Single-question requests (one row) replay a graph; other requests, rows longer than the largest bucket, and
-gradient-enabled calls run the eager path. Rows are right-padded with the pad token and run without an attention
-mask. Every layer used by row-mode backbones is causal (full or sliding-window attention, gated delta rule, short
-convolution), so tokens placed after the last real token cannot change a real token's hidden state. Pad positions
-continue the row's positions, keeping position ids increasing.
+Single-question requests (one row) replay a graph; other requests, rows longer than the largest bucket, rows whose
+positions do not rise by one, and gradient-enabled calls run the eager path. Rows are right-padded with the pad token
+and run without an attention mask. Every layer used by row-mode backbones is causal (full or sliding-window
+attention, gated delta rule, short convolution), so tokens placed after the last real token cannot change a real
+token's hidden state. Pad positions continue the row's positions, keeping position ids increasing.
+
+Capture makes the attention-mask decisions of eager inference (``_eager_mask_decisions``), so SDPA keeps its
+``is_causal`` path (FlashAttention) inside the graphs instead of reading a materialised causal mask.
 
 Results are not bit-identical to the eager path: padded shapes select different kernels. Replays reuse static
 buffers, so calls must not overlap; ``DecisionRuntime`` already serializes inference under its lock.
 """
+import contextlib
+import threading
 import time
 
 import torch
@@ -22,6 +27,68 @@
 LENGTHS = (128, 192, 256, 320, 384, 448, 512, 640, 768, 896, 1024, 1280, 1536, 1792, 2048, 2560, 3072, 4096,
            5120, 6144, 8192, 10240, 12288, 16384)
 
+_MASK_CAPTURE_LOCK = threading.RLock()
+_MASK_CAPTURE_STATE = threading.local()
+
+
+@contextlib.contextmanager
+def _eager_mask_decisions():
+    """Let transformers decide attention masks during capture as it does for an eager, unpadded single row.
+
+    transformers counts CUDA stream capture as tracing. It then cannot run its data-dependent check that the position
+    ids describe one sequence, so it keeps a packed-sequence mask; before transformers 5.18 it also declines SDPA's
+    ``is_causal`` shortcut for a missing padding mask. Either way the causal mask is materialised, which rules out
+    FlashAttention. A captured row has no padding mask and consecutive positions (``RowGraphs._bucket`` sends other
+    rows to the eager path), which is exactly when eager inference skips the mask. Capture contexts are serialized.
+    Only the owning thread's capture calls change; other forwards retain upstream mask and tracing decisions.
+    Nested contexts share the wrappers until the last context exits, including after a capture failure.
+    """
+    from transformers import masking_utils
+
+    with _MASK_CAPTURE_LOCK:
+        if not getattr(_MASK_CAPTURE_STATE, "depth", 0):
+            find_packed = getattr(masking_utils, "find_packed_sequence_indices", None)
+            ignore_causal = getattr(masking_utils, "_ignore_causal_mask_sdpa", None)
+            is_tracing = getattr(masking_utils, "is_tracing", None)
+
+            def capturing():
+                return getattr(_MASK_CAPTURE_STATE, "depth", 0) > 0 and torch.cuda.is_current_stream_capturing()
+
+            def single_sequence(*args, **kwargs):
+                return None if capturing() else find_packed(*args, **kwargs)
+
+            def tracing(*args, **kwargs):
+                return False if getattr(_MASK_CAPTURE_STATE, "ignore_tracing", False) else is_tracing(*args, **kwargs)
+
+            def causal_shortcut(padding_mask, *args, **kwargs):
+                if padding_mask is not None or not capturing():
+                    return ignore_causal(padding_mask, *args, **kwargs)
+                previous = getattr(_MASK_CAPTURE_STATE, "ignore_tracing", False)
+                _MASK_CAPTURE_STATE.ignore_tracing = True
+                try:
+                    return ignore_causal(padding_mask, *args, **kwargs)
+                finally:
+                    _MASK_CAPTURE_STATE.ignore_tracing = previous
+
+            replacements = {}
+            if find_packed is not None:
+                replacements["find_packed_sequence_indices"] = single_sequence
+            if ignore_causal is not None and is_tracing is not None:
+                replacements["_ignore_causal_mask_sdpa"] = causal_shortcut
+                replacements["is_tracing"] = tracing
+            _MASK_CAPTURE_STATE.originals = {name: getattr(masking_utils, name) for name in replacements}
+            for name, replacement in replacements.items():
+                setattr(masking_utils, name, replacement)
+        _MASK_CAPTURE_STATE.depth = getattr(_MASK_CAPTURE_STATE, "depth", 0) + 1
+        try:
+            yield
+        finally:
+            _MASK_CAPTURE_STATE.depth -= 1
+            if _MASK_CAPTURE_STATE.depth == 0:
+                for name, original in _MASK_CAPTURE_STATE.originals.items():
+                    setattr(masking_utils, name, original)
+                del _MASK_CAPTURE_STATE.originals
+
 
 class RowGraphs:
     """Captured row forwards for one DecisionModel. Construct, then ``capture()``; the model consults
@@ -64,7 +131,7 @@ def capture(self):
                     self._forward(ids, pos)
             current.wait_stream(side)
             graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph, pool=pool):
+            with _eager_mask_decisions(), torch.cuda.graph(graph, pool=pool):
                 out = self._forward(ids, pos)
             self.graphs[n] = (graph, ids, pos, out)
         torch.cuda.synchronize(self.device)
@@ -73,6 +140,8 @@ def capture(self):
 
     def _bucket(self, ids, pos):
         length = len(ids)
+        if pos != list(range(pos[0], pos[0] + length)):
+            return None   # eager inference would mask a position jump as a packed-sequence boundary
         for n in self.lengths:
             if n >= length:
                 # Absolute position tables (GPT-2) must also cover the padded tail.

tests/test_cudagraph_mask_scope.py

@@ -0,0 +1,122 @@
+"""Capture-only mask shortcuts must not escape into other model calls."""
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+import torch
+
+from jevany.cudagraphs import _eager_mask_decisions
+
+
+@pytest.fixture
+def masks(monkeypatch):
+    from transformers import LlamaConfig, masking_utils
+
+    config = LlamaConfig(hidden_size=8, intermediate_size=16, num_hidden_layers=1,
+                         num_attention_heads=2, num_key_value_heads=2, vocab_size=16)
+    config._attn_implementation = "sdpa"
+    embeddings = torch.zeros(1, 6, 8)
+    originals = {name: getattr(masking_utils, name) for name in
+                 ("find_packed_sequence_indices", "_ignore_causal_mask_sdpa", "is_tracing")
+                 if hasattr(masking_utils, name)}
+    # Restore upstream globals even if a regression leaves a wrapper installed.
+    for name, original in originals.items():
+        monkeypatch.setattr(masking_utils, name, original)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+
+    def mask(packed=False):
+        positions = torch.tensor([[0, 1, 2, 4, 5, 6]]) if packed else torch.arange(6)[None]
+        return masking_utils.create_causal_mask(
+            config=config, inputs_embeds=embeddings, attention_mask=None,
+            past_key_values=None, position_ids=positions,
+        )
+
+    return mask, masking_utils, originals
+
+
+def test_overlapping_capture_contexts_restore_globals(masks):
+    mask, masking_utils, originals = masks
+    first, second = _eager_mask_decisions(), _eager_mask_decisions()
+    first.__enter__()
+    second.__enter__()
+    try:
+        assert mask() is None
+    finally:
+        first.__exit__(None, None, None)
+        second.__exit__(None, None, None)
+    for name, original in originals.items():
+        assert getattr(masking_utils, name) is original
+    assert mask(packed=True) is not None
+
+
+def test_unrelated_capture_keeps_its_mask(masks):
+    mask, _, _ = masks
+    with ThreadPoolExecutor(max_workers=1) as pool, _eager_mask_decisions():
+        assert mask() is None
+        assert pool.submit(mask, True).result(timeout=5) is not None
+        assert pool.submit(mask).result(timeout=5) is not None
+
+
+def test_capture_failure_restores_nested_contexts(masks):
+    mask, masking_utils, originals = masks
+    with _eager_mask_decisions():
+        with pytest.raises(RuntimeError, match="capture failed"):
+            with _eager_mask_decisions():
+                assert mask() is None
+                raise RuntimeError("capture failed")
+        assert mask() is None
+    for name, original in originals.items():
+        assert getattr(masking_utils, name) is original
+    assert mask() is not None
+
+
+def test_concurrent_capture_contexts_are_serialized(masks):
+    mask, masking_utils, originals = masks
+    attempting, entered = Event(), Event()
+
+    def capture():
+        attempting.set()
+        with _eager_mask_decisions():
+            entered.set()
+            assert mask() is None
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        try:
+            with _eager_mask_decisions():
+                future = pool.submit(capture)
+                assert attempting.wait(5)
+                assert not entered.wait(0.05)
+                assert mask() is None
+        finally:
+            future.result(timeout=5)
+    assert entered.is_set()
+    for name, original in originals.items():
+        assert getattr(masking_utils, name) is original
+
+
+def test_causal_shortcut_does_not_change_other_threads_tracing(masks, monkeypatch):
+    _, masking_utils, _ = masks
+    if not hasattr(masking_utils, "_ignore_causal_mask_sdpa"):
+        pytest.skip("this Transformers version does not use the tracing-dependent shortcut")
+    original = masking_utils._ignore_causal_mask_sdpa
+    entered, release = Event(), Event()
+
+    def paused_shortcut(*args, **kwargs):
+        entered.set()
+        assert release.wait(5)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(masking_utils, "_ignore_causal_mask_sdpa", paused_shortcut)
+
+    def capture():
+        with _eager_mask_decisions():
+            return masking_utils._ignore_causal_mask_sdpa(None, 6, 6, 0, 0)
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        future = pool.submit(capture)
+        try:
+            assert entered.wait(5)
+            assert masking_utils.is_tracing()
+        finally:
+            release.set()
+            assert future.result(timeout=5)

tests/test_cudagraphs.py

@@ -1,9 +1,11 @@
 """CUDA-graph replay for row-mode inference: option plumbing everywhere, parity with the eager path on CUDA."""
+from types import SimpleNamespace
+
 import pytest
 import torch
 
 from jevany.checkpoint import LoadOptions
-from jevany.cudagraphs import RowGraphs
+from jevany.cudagraphs import RowGraphs, _eager_mask_decisions
 from jevany.model import DecisionModel, load_tokenizer
 from test_backbones import RECORD, make_base
 
@@ -61,6 +63,34 @@ def test_remote_decide_rejects_cuda_graphs(tmp_path, capsys):
     assert "require --checkpoint" in capsys.readouterr().err
 
 
+def test_capture_mask_decisions_follow_eager_inference(monkeypatch):
+    """transformers counts stream capture as tracing and would materialise the causal mask, ruling out FlashAttention."""
+    from transformers import LlamaConfig
+    from transformers.masking_utils import create_causal_mask
+    config = LlamaConfig(hidden_size=8, intermediate_size=16, num_hidden_layers=1, num_attention_heads=2,
+                         num_key_value_heads=2, vocab_size=16)
+    config._attn_implementation = "sdpa"
+    embeddings = torch.zeros(1, 6, 8)
+
+    def mask(positions):
+        return create_causal_mask(config=config, inputs_embeds=embeddings, attention_mask=None,
+                                  past_key_values=None, position_ids=positions)
+
+    consecutive, jump = torch.arange(6)[None], torch.tensor([[0, 1, 2, 4, 5, 6]])
+    assert mask(consecutive) is None and mask(jump) is not None   # eager: SDPA is_causal unless packed
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert mask(consecutive) is not None                           # capture counts as tracing
+    with _eager_mask_decisions():
+        assert mask(consecutive) is None
+    assert mask(consecutive) is not None                           # transformers is restored
+
+
+def test_rows_with_position_jumps_run_eagerly():
+    graphs = SimpleNamespace(lengths=(8, 16), window=None)
+    assert RowGraphs._bucket(graphs, [5] * 6, list(range(3, 9))) == 8
+    assert RowGraphs._bucket(graphs, [5] * 6, [0, 1, 2, 4, 5, 6]) is None
+
+
 ONE_QUESTION = {**RECORD, "questions": RECORD["questions"][:1]}
 
 
@@ -87,6 +117,23 @@ def test_graph_replay_matches_eager(tmp_path, family, dtype, decision_mode):
     assert (model.cuda_graphs.stats["graph_calls"], model.cuda_graphs.stats["eager_calls"]) == (2, 0)
 
 
+@cuda
+def test_capture_keeps_sdpa_causal_path(tmp_path, monkeypatch):
+    base = tmp_path / "base"
+    make_base(base, "qwen35", legacy=True)
+    model = DecisionModel(base, load_tokenizer(base), "cuda", lora=2, head_dim=8, dtype=torch.bfloat16).eval()
+    captured, sdpa = [], torch.nn.functional.scaled_dot_product_attention
+
+    def spy(*args, **kwargs):
+        if torch.cuda.is_current_stream_capturing():
+            captured.append((kwargs.get("attn_mask") is None, kwargs.get("is_causal", False)))
+        return sdpa(*args, **kwargs)
+
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", spy)
+    RowGraphs(model, lengths=(16, 32)).capture()
+    assert captured and all(mask_free and causal for mask_free, causal in captured)
+
+
 @cuda
 def test_multi_question_and_long_rows_run_eagerly(tmp_path):
     base = tmp_path / "base"