# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """Breakable CUDA graph capture for variable-shape (prefill / extend) forwards. A *breakable* CUDA graph captures a forward as an ordered list of zero-arg callables -- each is either a captured ``CUDAGraph.replay`` (a "graph segment") or an eager Python function (a "break"). At designated break points (attention / KV-cache ops, whose metadata is data-dependent and cannot be captured) the current stream capture is ended, the op runs eagerly, and a fresh segment begins capturing the remainder. Replay simply calls each segment in order. This is the ``torch.compile``-free alternative to piecewise CUDA graphs. The design is gratefully adapted from vLLM and SGLang, who pioneered the breakable prefill graph: vLLM's ``BreakableCUDAGraphWrapper`` (the homogeneous segment-list structure + the ``set_forward_context``/``get_forward_context`` ambient pattern we mirror in :func:`active_forward`/:func:`current_forward_ctx`) and SGLang's breakable prefill graph (the eager-copy output handoff at each break). Unlike a full prefill graph, attention -- the only batch/length-aware op and the source of the host-side ``max_seq_len_q`` scalar -- stays eager, so it never enters a graph. Keeping all KV-cache reads/writes in the eager breaks also makes them honor the per-layer transfer consumer index naturally. Address-stability contract (the load-bearing invariant): * All segments share one CUDA mempool, so graph-allocated intermediates keep stable device addresses across replays. * The runner must copy live inputs into the *same* static input buffers used at capture before calling :meth:`BreakableCapture.replay`. * Break-point outputs must land at the *same* address each replay. We achieve this by allocating a destination buffer in the captured segment (pool-pinned) and copying the eager op's result into it; the next segment reads that address. """ from __future__ import annotations import functools import gc from collections import deque from collections.abc import Callable, Iterator from contextlib import contextmanager from typing import Any import torch from tokenspeed_kernel.ops.transform.weak_ref import weak_ref_tensor as _kernel_weak_ref __all__ = [ "BreakableCapture", "active_forward", "break_here", "break_point", "current_forward_ctx", "is_breakable_capture_active", "scrub_padding_tail", "slice_to_real_tokens", "weak_ref_tensor", ] # Ambient per-forward ctx; plain module state (one launch thread per rank). _ambient_ctx: Any = None @contextmanager def active_forward(ctx: Any) -> Iterator[None]: """Publish ``ctx`` as the ambient forward context for the enclosed block. An eager break runs once at capture and again on every replay, and the args it closed over at capture are the *dummy* batch's, hence stale. Rather than thread the live context through ``replay()`` (which would conflate graph mechanics with forward semantics), the runner wraps capture and each replay in this, and breaks rebind their captured context arg to the ambient one by identity (see :func:`break_here`) -- so break bodies read live ``ctx`` fields exactly like the eager path. Re-entrant (saves/restores the previous value). """ global _ambient_ctx prev = _ambient_ctx _ambient_ctx = ctx try: yield finally: _ambient_ctx = prev def current_forward_ctx() -> Any: """The ambient forward context, or ``None`` outside an :func:`active_forward`.""" return _ambient_ctx def weak_ref_tensor(t: Any) -> Any: """Reference a break-point tensor without pinning its cudagraph mempool slot. CUDA tensors are wrapped in a non-owning view (``tokenspeed_kernel`` ``ops.transform.weak_ref``, an ``at::from_blob`` alias -- the vLLM/sglang approach), so break closures do not pin pool blocks and graph capture memory stays ~peak-live instead of scaling with the bucket sum. Safe because replay is stream-ordered: the captured segment rewrites the aliased address before the break reads it. Non-tensors and CPU tensors pass through; if the kernel extension is unavailable this degrades to the identity (strong ref -- correct, more memory). """ if isinstance(t, torch.Tensor) and t.is_cuda: return _kernel_weak_ref(t) return t class BreakableCapture: """Context manager that captures a breakable graph. Usage:: cap = BreakableCapture(pool=shared_pool) with cap: model_forward(...) # attention calls hit break_here() # later, after copying live inputs into the static buffers: cap.replay() Args: pool: An optional CUDA mempool id (as returned by ``torch.cuda.graph_pool_handle()`` or ``CUDAGraph.pool()``) shared by all segments. If ``None``, the first segment allocates a fresh pool and the rest reuse it. stream: An optional dedicated capture stream. CUDA forbids stream capture on the default stream; if ``None``, a single class-level side stream is (lazily) created and SHARED by all captures. Sharing one capture stream is load-bearing for memory: the caching allocator's pool blocks are stream-keyed, so captures on different streams can never reuse each other's freed blocks -- a fresh stream per capture makes graph-pool memory grow with the SUM of bucket sizes instead of the max (measured: 2478MB -> 564MB for buckets [8192,4096,2048,1024] on the repro). This mirrors ``torch.cuda.graph``'s shared ``default_capture_stream`` and its documented "pass the same stream for effective memory sharing". """ _active: BreakableCapture | None = None _default_capture_stream: torch.cuda.Stream | None = None def __init__( self, pool: Any | None = None, stream: torch.cuda.Stream | None = None ) -> None: self.pool = pool self.segments: list[Callable[[], Any]] = [] self._current_graph: torch.cuda.CUDAGraph | None = None self._capturing = False if stream is None: if BreakableCapture._default_capture_stream is None: BreakableCapture._default_capture_stream = torch.cuda.Stream() stream = BreakableCapture._default_capture_stream self._stream = stream self._stream_ctx: Any | None = None # Break-output handoff buffers keyed by (shape, dtype, device); see break_point. self._handoff: dict[Any, torch.Tensor] = {} @classmethod def current(cls) -> BreakableCapture | None: return cls._active # -- capture lifecycle ------------------------------------------------- def __enter__(self) -> BreakableCapture: if BreakableCapture.current() is not None: raise RuntimeError("Nested BreakableCapture is not supported.") # A GC run during capture invalidates it: destructors of collected # CUDA graphs call reset, which is illegal while a stream is capturing. # Clear pending garbage, then keep automatic GC off for the whole # capture window (restored in __exit__). gc.collect() self._gc_was_enabled = gc.isenabled() gc.disable() # The capture stream must observe prior entry-stream work (warmup, buffers). self._stream.wait_stream(torch.cuda.current_stream()) self._stream_ctx = torch.cuda.stream(self._stream) self._stream_ctx.__enter__() BreakableCapture._active = self self._begin_segment() return self def __exit__(self, *exc: object) -> bool: try: self._end_segment() finally: BreakableCapture._active = None if self._stream_ctx is not None: self._stream_ctx.__exit__(*exc) self._stream_ctx = None # Eager breaks ran on the side stream; entry stream must observe them. torch.cuda.current_stream().wait_stream(self._stream) if self._gc_was_enabled: gc.enable() return False def _begin_segment(self) -> None: assert not self._capturing graph = torch.cuda.CUDAGraph() if self.pool is not None: graph.capture_begin(pool=self.pool) else: graph.capture_begin() self._current_graph = graph self._capturing = True def _end_segment(self) -> None: if not self._capturing: return assert self._current_graph is not None self._current_graph.capture_end() self.segments.append(self._current_graph.replay) # All segments share one pool so intermediate addresses stay stable. if self.pool is None: self.pool = self._current_graph.pool() self._current_graph = None self._capturing = False def add_eager(self, fn: Callable[[], Any]) -> Any: """End the current segment, run ``fn`` eagerly, record it, start a new one. ``fn`` is a zero-arg callable that performs the break-point op and writes its result into a stable (pool-pinned) address. It is stored verbatim and re-invoked on every :meth:`replay`. """ assert self._capturing, "add_eager called outside an active capture" self._end_segment() result = fn() self.segments.append(fn) self._begin_segment() return result # -- replay ------------------------------------------------------------ def replay(self) -> None: """Replay all segments in order. Breaks read the live forward context from the ambient :func:`active_forward` scope (the runner wraps replay in it), so this stays a pure graph primitive. """ deque((run() for run in self.segments), maxlen=0) @property def num_segments(self) -> int: return len(self.segments) def is_breakable_capture_active() -> bool: """True while a :class:`BreakableCapture` is open AND currently capturing.""" cap = BreakableCapture.current() return cap is not None and cap._capturing def _record_break( cap: BreakableCapture, fn: Callable[..., torch.Tensor], resolve_dst: Callable[[torch.Tensor], torch.Tensor], args: tuple, kwargs: dict, ) -> torch.Tensor: """Record ``fn(*args, **kwargs)`` as an eager break on ``cap`` (the one closure builder shared by :func:`break_here` and :func:`break_point`). Args/kwargs are bound once at capture time, with two live exceptions: (1) tensor args alias persistent storage (the static input buffers / pool-pinned segment intermediates), so they carry live values at replay -- ``weak_ref_tensor`` is the (currently identity) hook to avoid pinning their pool slots; (2) the per-forward ``ForwardContext`` is rebound by identity to the live ambient context each replay (see :func:`active_forward`), so ``fn`` may read live ``ctx`` fields exactly like the eager path. **Other (loose) non-tensor scalars are frozen** to their capture-time value -- route per-request quantities through ``ctx`` / ``forward_*_metadata`` rather than a bare scalar arg. ``resolve_dst(result)`` maps the break's first output to its stable handoff buffer; it is called once (on the capture-time invocation) and the buffer is reused verbatim on every replay, where the (possibly shorter, see :func:`_land_in`) live result is copied into it. """ weak_args = tuple(weak_ref_tensor(a) for a in args) weak_kwargs = {k: weak_ref_tensor(v) for k, v in kwargs.items()} # Capture-time ambient ctx (the dummy batch's), rebound live at replay. captured_ctx = current_forward_ctx() state: dict[str, torch.Tensor] = {} def replay_fn() -> torch.Tensor: live_ctx = current_forward_ctx() def sub(a: Any) -> Any: return live_ctx if a is captured_ctx else a result = fn( *(sub(a) for a in weak_args), **{k: sub(v) for k, v in weak_kwargs.items()}, ) dst = state.get("dst") if dst is None: dst = state["dst"] = resolve_dst(result) _land_in(dst, result) return dst return cap.add_eager(replay_fn) def break_here( fn: Callable[..., torch.Tensor], dst: torch.Tensor, *args: Any, **kwargs: Any, ) -> torch.Tensor: """Run ``fn(*args, **kwargs)`` as an eager break, landing its result in ``dst``. The low-level explicit-destination primitive underneath :func:`break_point` (which is the decorator every model actually uses -- prefer it; this exists for callers that must control the handoff buffer's placement themselves, e.g. a pool-pinned ``dst`` allocated in the current captured segment, and for exercising the break mechanics directly in unit tests). ``dst`` must have a replay-stable address (pool-pinned or persistently owned). At capture and on every replay, ``fn`` runs eagerly and its result is copied into ``dst`` (unless ``fn`` already wrote ``dst`` in place and returned it); the following graph segment reads ``dst``. Outside an active capture this is a transparent pass-through. Argument binding/freezing semantics are those of :func:`_record_break`. Returns: ``dst`` (the stable handoff buffer). """ cap = BreakableCapture.current() if cap is None or not cap._capturing: _land_in(dst, fn(*args, **kwargs)) return dst weak_dst = weak_ref_tensor(dst) return _record_break(cap, fn, lambda _result: weak_dst, args, kwargs) def break_point(method: Callable | None = None) -> Callable: """Mark a sequence-mixing method as an eager breakable-graph break point. Decorate a sequence-mixing method (attention / MLA / linear-mixer / sparse indexer ``forward``) and it runs as an eager break under a breakable capture -- the surrounding token-shaped compute (norms, MoE, projections, collectives) is captured around it automatically, while everything inside the method stays eager -- or a zero-overhead direct call when not capturing. This is the one decorator every model uses to mark a break. Use it bare: ``@break_point``. The handoff buffer's shape/dtype/device are **inferred from the method's actual output** at capture time (the break runs during capture regardless), so no output spec is needed -- it works uniformly for breaks whose output matches no input (MLA: ``[tokens, heads*v_head_dim]`` vs ``q``'s ``[tokens, heads*qk_head_dim]``) and for one wrapper that returns different shapes per call (e.g. hybrid full-attn q-shaped vs GDN z-shaped). Buffers live in a per-capture shape-keyed cache (:attr:`BreakableCapture._handoff`), shared across same-shape breaks. That sharing relies on break outputs having strictly sequential lifetimes -- break K's output is consumed by K's following segment before break K+1 runs (true for transformer topology, where attention output feeds the adjacent o-proj/residual). A model whose break output is read by a LATER segment must not share its shape with an intervening break, or replay silently corrupts. Inside the method ``ctx`` is live (rebound by identity at replay), so write the body exactly like the eager path. Loose non-tensor scalar args are frozen to their capture-time value -- route per-request quantities through ``ctx`` / metadata. The decorator never skips the method: 0-row / idle batches remain each decorated model method's own explicit guard, on the eager path and under capture alike. """ def decorator(method: Callable) -> Callable: @functools.wraps(method) def wrapper(*args: Any, **kwargs: Any) -> Any: # Zero-overhead passthrough off the capture path (no 0-row skipping here). if not is_breakable_capture_active(): return method(*args, **kwargs) cap = BreakableCapture.current() def resolve_dst(result: torch.Tensor) -> torch.Tensor: # Handoff buffer inferred from the capture-time output, shape-keyed. key = (tuple(result.shape), result.dtype, result.device) dst = cap._handoff.get(key) if dst is None: dst = cap._handoff[key] = torch.empty( result.shape, dtype=result.dtype, device=result.device ) return dst return _record_break(cap, method, resolve_dst, args, kwargs) return wrapper return decorator(method) if method is not None else decorator # -- padded-input helpers (the two strategies of the prefill padding contract) -- def scrub_padding_tail(num_real_tokens: int, *tensors: torch.Tensor | None) -> None: """Zero the padded tail rows ``[num_real_tokens:]`` of token-shaped tensors in place. Under a padded-bucket replay, an eager break receives ``bucket`` rows whose tail holds garbage (it grows across layers and can overflow to NaN through projections / FP8 quantize). Zeroing suits breaks whose kernel honors the live cu-seqlens but whose surrounding ops (varlen attention read, recurrent scan writeback, FP8 quantize) would otherwise touch the garbage rows. Pass the real token count from the live metadata's CPU mirror (sync-free on the launch thread); a no-op on unpadded forwards, and ``None`` tensors are skipped. """ for t in tensors: if t is not None and num_real_tokens < t.shape[0]: t[num_real_tokens:].zero_() def slice_to_real_tokens(num_real_tokens: int, *tensors: torch.Tensor | None): """Return ``tensors`` (in order) each sliced to the real leading rows ``[:num_real_tokens]``. The slice-strategy counterpart to :func:`scrub_padding_tail`, for coarse breaks whose kernel asserts the input row count equals the live metadata token count (e.g. DSA sparse attention). A tensor already the right length (or ``None``) is returned unchanged. """ return tuple( t[:num_real_tokens] if (t is not None and num_real_tokens < t.shape[0]) else t for t in tensors ) def _land_in(dst: torch.Tensor, result: torch.Tensor) -> None: """Copy ``result`` into ``dst`` at a stable address. ``dst`` is the (possibly token-padded) handoff buffer the next graph segment reads. ``result`` may cover only the real (unpadded) leading rows -- e.g. a varlen attention kernel writes only ``sum(cu_seqlens_q)`` rows -- so we copy into the matching leading slice. Padded rows are left as-is (discarded by the final output slice). No-op when the op already wrote ``dst`` in place. """ if result is dst: return if result.shape == dst.shape: dst.copy_(result) else: dst.narrow(0, 0, result.shape[0]).copy_(result)