diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index bfd0fcb..acff268 100644 --- a/docs/DEPLOYMENT.md +++ b/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 diff --git a/jevany/cudagraphs.py b/jevany/cudagraphs.py index fe34bd6..5c10750 100644 --- a/jevany/cudagraphs.py +++ b/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. diff --git a/tests/test_cudagraph_mask_scope.py b/tests/test_cudagraph_mask_scope.py new file mode 100644 index 0000000..44ce86b --- /dev/null +++ b/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) 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"