From 96cb307bed956fc2f3f905a917c6c1f814133094 Mon Sep 17 00:00:00 2001 From: Huanxin Sheng Date: Mon, 5 Oct 2026 02:30:17 +0400 Subject: [PATCH 1/2] Keep FlashAttention inside captured CUDA graphs transformers counts CUDA stream capture as tracing. While RowGraphs captured a row it kept a packed-sequence mask (it cannot run its data-dependent single-sequence check) and, before transformers 5.18, also declined SDPA's is_causal shortcut. Captured attention therefore read a materialised causal mask and used the memory-efficient kernel instead of the FlashAttention kernel eager inference picks for the same row. Capture now makes the eager mask decisions for an unpadded single row: the two masking_utils functions are wrapped only for calls made while the current stream is capturing, and are restored afterwards. Rows whose positions do not rise by one run eagerly, because eager inference would mask such a jump as a packed-sequence boundary. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EpJAgfodeE1m4Fk2tDiLfB --- jevany/cudagraphs.py | 62 ++++++++++++++++++++++++++++++++++++---- tests/test_cudagraphs.py | 49 ++++++++++++++++++++++++++++++- 2 files changed, 104 insertions(+), 7 deletions(-) diff --git a/jevany/cudagraphs.py b/jevany/cudagraphs.py index 24be570..0bc189f 100644 --- a/jevany/cudagraphs.py +++ b/jevany/cudagraphs.py @@ -5,15 +5,19 @@ 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 time import torch @@ -23,6 +27,50 @@ 5120, 6144, 8192, 10240, 12288, 16384) +@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. Only calls made while the current + stream is capturing change; the original functions are restored on exit. + """ + from transformers import masking_utils + + find_packed = getattr(masking_utils, "find_packed_sequence_indices", None) + ignore_causal = getattr(masking_utils, "_ignore_causal_mask_sdpa", None) + + def single_sequence(*args, **kwargs): + return None if torch.cuda.is_current_stream_capturing() else find_packed(*args, **kwargs) + + def causal_shortcut(padding_mask, *args, **kwargs): + if padding_mask is not None or not torch.cuda.is_current_stream_capturing(): + return ignore_causal(padding_mask, *args, **kwargs) + is_tracing = masking_utils.is_tracing + masking_utils.is_tracing = lambda *_, **__: False # without a padding mask nothing data-dependent is read + try: + return ignore_causal(padding_mask, *args, **kwargs) + finally: + masking_utils.is_tracing = is_tracing + + replacements = {} + if find_packed is not None: + replacements["find_packed_sequence_indices"] = single_sequence + if ignore_causal is not None and hasattr(masking_utils, "is_tracing"): + replacements["_ignore_causal_mask_sdpa"] = causal_shortcut + originals = {name: getattr(masking_utils, name) for name in replacements} + for name, replacement in replacements.items(): + setattr(masking_utils, name, replacement) + try: + yield + finally: + for name, original in originals.items(): + setattr(masking_utils, name, original) + + class RowGraphs: """Captured row forwards for one DecisionModel. Construct, then ``capture()``; the model consults ``forward_rows_batch`` before its eager path and ``stats`` reports lengths, capture time and call counts.""" @@ -64,7 +112,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 +121,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. diff --git a/tests/test_cudagraphs.py b/tests/test_cudagraphs.py index 6979f0d..5ff5323 100644 --- a/tests/test_cudagraphs.py +++ b/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" From 3b8fca2008c2342ad14cb725541cfc084d190a32 Mon Sep 17 00:00:00 2001 From: Huanxin Sheng Date: Mon, 5 Oct 2026 09:07:03 +0400 Subject: [PATCH 2/2] docs: describe CUDA-graph capture mask decisions Explain why capture now keeps FlashAttention and record the A100 JevAny-Qwen3.5-4B measurements, including why the 2,048-token default limit stays. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EpJAgfodeE1m4Fk2tDiLfB --- docs/DEPLOYMENT.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index 42a6c00..75ce05f 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -356,6 +356,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