wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
657 lines
30 KiB
Python
657 lines
30 KiB
Python
# 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 graphs for prefill (extend) forwards.
|
|
|
|
:class:`PrefillGraph` holds one breakable graph per padded token bucket
|
|
(captured from a dummy bs=1 extend batch). The embedding lookup stays OUTSIDE
|
|
the captured region: graphs start from a static input-embeds buffer, filled at
|
|
replay by an eager ``embed_tokens`` gather (text) or by precomputed merged
|
|
embeddings (multimodal, via the model's ``multimodal_input_embeds`` seam).
|
|
Constructed after the decode
|
|
:class:`~tokenspeed.runtime.execution.cuda_graph_wrapper.CudaGraphWrapper`,
|
|
borrowing its capture stream; buckets share one private mempool, deliberately
|
|
not the decode graphs' pool (see :meth:`capture`). At serving time
|
|
the executor's target-forward dispatch is a flat
|
|
three-way -- decode & captured replays the decode graph (one level up, since
|
|
it captures the whole step), prefill & captured replays here (:meth:`can_run`
|
|
/ :meth:`replay`), everything else runs the eager model forward.
|
|
|
|
Unlike decode (whole forward captured, keyed by batch size), the captured
|
|
region here is purely token-shaped compute keyed by total token count:
|
|
attention runs as an eager break (see
|
|
:mod:`tokenspeed.runtime.execution.breakable_cuda_graph`), so one graph per
|
|
bucket serves any batch size at that token count, and a replayed forward is
|
|
finished with the model's eager logits tail.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import bisect
|
|
from contextlib import contextmanager
|
|
from typing import TYPE_CHECKING, NamedTuple
|
|
|
|
import torch
|
|
|
|
from tokenspeed.runtime.execution.breakable_cuda_graph import (
|
|
BreakableCapture,
|
|
active_forward,
|
|
)
|
|
from tokenspeed.runtime.execution.context import ForwardContext
|
|
from tokenspeed.runtime.execution.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardMode,
|
|
)
|
|
from tokenspeed.runtime.layers.logits_processor import LogitsMetadata
|
|
from tokenspeed.runtime.utils import get_colorful_logger
|
|
from tokenspeed.runtime.utils.common import maybe_inference_mode
|
|
|
|
logger = get_colorful_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from tokenspeed.runtime.execution.cuda_graph_wrapper import CudaGraphWrapper
|
|
from tokenspeed.runtime.execution.input_buffer import InputBuffers
|
|
from tokenspeed.runtime.execution.model_executor import ModelExecutorConfig
|
|
from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend
|
|
|
|
|
|
# Smallest prefill bucket; below this, denser rungs would only add capture time.
|
|
PREFILL_BUCKET_FLOOR: int = 16
|
|
|
|
# Relative rung spacing (largest pow2 <= size/8), bounding the padded tail at ~12.5%.
|
|
PREFILL_BUCKET_STEP_DIVISOR: int = 8
|
|
|
|
# Absolute rung-spacing cap, bounding the worst case at the top of the ladder.
|
|
PREFILL_BUCKET_MAX_STEP: int = 512
|
|
|
|
|
|
def get_prefill_token_buckets(config: ModelExecutorConfig) -> list[int]:
|
|
"""Padded token-count buckets to capture for the breakable prefill graph.
|
|
|
|
Unlike decode (keyed by batch size), the breakable prefill graph captures
|
|
pure token-shaped compute, so it is keyed by total token count. A live extend
|
|
forward is padded up to the smallest bucket >= its token count; forwards above
|
|
the largest bucket run eager.
|
|
|
|
Returns an empty list (graph disabled) when ``disable_prefill_graph`` is set or
|
|
``prefill_graph_max_tokens <= 0``. The largest bucket is clamped to the
|
|
chunked-prefill size: the scheduler's per-forward token budget
|
|
(``max_scheduled_tokens`` = chunked-prefill size) covers extends AND any fused
|
|
decode rows -- with mixed batching, decodes are scheduled first and each
|
|
decrements the budget, and the prefill chunk is sized to what remains
|
|
(scheduler ``newForwardOperation``/``push_op``) -- so no forward, mixed or
|
|
pure, ever exceeds the chunk. No headroom above it is needed.
|
|
|
|
The default ladder bounds RELATIVE padding waste: a forward pads its graphed
|
|
compute to the next bucket, so what matters is the gap as a fraction of the
|
|
size -- a flat stride is needlessly coarse for short prompts and needlessly
|
|
dense at the top. Each bucket's step is the largest power of two <= size/8
|
|
(padded tail at most ~12.5% anywhere on the ladder), floored at 16 tokens and
|
|
capped at 512 so the absolute worst case stays bounded at the top end. Dense
|
|
ladders are cheap: all captures share one stream + mempool, so graph memory
|
|
is ~the largest bucket's peak regardless of bucket count (see
|
|
``BreakableCapture``); the remaining cost is ~0.5s of startup capture per
|
|
bucket.
|
|
|
|
``prefill_graph_capture_sizes`` overrides the ladder with an explicit list
|
|
(mirroring decode's ``cudagraph_capture_sizes``) -- e.g. a short list for
|
|
faster startup on dev boots; sizes are clamped to the largest bucket.
|
|
|
|
Args:
|
|
config: The model-executor config carrying ``disable_prefill_graph``,
|
|
``prefill_graph_max_tokens``, ``prefill_graph_capture_sizes`` and
|
|
``chunked_prefill_size``.
|
|
|
|
Returns:
|
|
Sorted ascending list of token-bucket sizes (possibly empty).
|
|
"""
|
|
max_tokens = int(config.prefill_graph_max_tokens or 0)
|
|
if config.disable_prefill_graph or max_tokens <= 0:
|
|
return []
|
|
chunk = int(config.chunked_prefill_size or 0)
|
|
if chunk > 0:
|
|
max_tokens = min(max_tokens, chunk)
|
|
explicit = config.prefill_graph_capture_sizes
|
|
if explicit:
|
|
buckets = {int(b) for b in explicit if 0 < int(b) <= max_tokens}
|
|
buckets.add(max_tokens)
|
|
return sorted(buckets)
|
|
buckets = []
|
|
size = min(PREFILL_BUCKET_FLOOR, max_tokens)
|
|
while size < max_tokens:
|
|
buckets.append(size)
|
|
size += _prefill_bucket_step(size)
|
|
buckets.append(max_tokens)
|
|
return sorted(set(buckets))
|
|
|
|
|
|
def _prefill_bucket_step(size: int) -> int:
|
|
"""Distance from bucket ``size`` to the next rung.
|
|
|
|
The largest power of two <= ``size / PREFILL_BUCKET_STEP_DIVISOR`` (so the
|
|
padded tail stays within ~1/8 of the real token count), clamped between
|
|
``PREFILL_BUCKET_FLOOR`` and ``PREFILL_BUCKET_MAX_STEP``.
|
|
"""
|
|
relative = size // PREFILL_BUCKET_STEP_DIVISOR
|
|
if relative <= PREFILL_BUCKET_FLOOR:
|
|
return PREFILL_BUCKET_FLOOR
|
|
largest_pow2 = 1 << (relative.bit_length() - 1)
|
|
return min(largest_pow2, PREFILL_BUCKET_MAX_STEP)
|
|
|
|
|
|
class CapturedForward(NamedTuple):
|
|
"""A bucket's captured inner-forward outputs (stable pool addresses)."""
|
|
|
|
# Final hidden states with shape [bucket, hidden]; padded tail is garbage.
|
|
hidden_states: torch.Tensor
|
|
|
|
# Aux hidden states for drafting, each [bucket, hidden]; None when mode is NULL.
|
|
aux_hidden_states: list[torch.Tensor] | None
|
|
|
|
def sliced(self, num_tokens: int) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
|
|
"""The leading real-token rows, in the (hidden, aux) shape callers expect."""
|
|
hidden = self.hidden_states[:num_tokens]
|
|
if self.aux_hidden_states is None:
|
|
return hidden, None
|
|
return hidden, [a[:num_tokens] for a in self.aux_hidden_states]
|
|
|
|
|
|
class PrefillGraph:
|
|
"""The breakable prefill (extend) CUDA graphs.
|
|
|
|
A pure graph object -- :meth:`can_run` / :meth:`replay` -- holding no
|
|
reference to any other component. Constructed AFTER the decode
|
|
``CudaGraphWrapper`` and captures in ``__init__`` like it: the decode
|
|
wrapper is used transiently for its capture stream and dummy paged-cache
|
|
tables, not kept. (The executor's target-forward dispatch therefore mode-
|
|
checks before touching this object -- decode capture runs that dispatch
|
|
while this object does not exist yet.) The dispatch checks :meth:`can_run`
|
|
and calls :meth:`replay`; the eager path stays a direct
|
|
``model_runner.forward`` call at that call site. Capture failure degrades
|
|
to eager -- world-agreed, so DP/TP ranks stay in lockstep.
|
|
|
|
Args:
|
|
model_runner: The target ModelRunner. Supplies the loaded model
|
|
(multimodal wrappers are unwrapped internally: the graph wraps the
|
|
nested ``language_model``'s text transformer, image prefills run
|
|
eager) and ``is_generation`` (embedding models run eager).
|
|
attn_backend: Backend whose extend metadata the dummy capture batch sets.
|
|
token_to_kv_pool: KV pool the dummy batch points at (reserved dummy slot).
|
|
input_buffers: The shared static input buffers the graphs read from.
|
|
config: Model-executor config (buckets, DP/world topology, device).
|
|
req_to_page: Request page table; row 0 backs the dummy capture request.
|
|
drafter: If present, aux-hidden capture (EAGLE3/MTP) is baked into the
|
|
captured graphs.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_runner,
|
|
attn_backend: AttentionBackend,
|
|
token_to_kv_pool,
|
|
input_buffers: InputBuffers,
|
|
config: ModelExecutorConfig,
|
|
req_to_page: torch.Tensor | None,
|
|
drafter=None,
|
|
decode_wrapper: CudaGraphWrapper | None = None,
|
|
num_warmup: int = 3,
|
|
) -> None:
|
|
model = model_runner.model if model_runner is not None else None
|
|
# Multimodal seam: models whose multimodal path is embeds-only expose
|
|
# multimodal_input_embeds; others (e.g. deepstack) replay text only.
|
|
self._multimodal_input_embeds = getattr(model, "multimodal_input_embeds", None)
|
|
self.text_model = (
|
|
model.language_model if hasattr(model, "language_model") else model
|
|
)
|
|
self.inner_model = getattr(self.text_model, "model", None)
|
|
# Embedding runs eagerly OUTSIDE the graphs (see capture); the graphs
|
|
# read a static input-embeds buffer instead of gathering from input_ids.
|
|
self._embed_tokens = getattr(self.inner_model, "embed_tokens", None)
|
|
self._input_embeds_buf: torch.Tensor | None = None
|
|
self.attn_backend = attn_backend
|
|
self.token_to_kv_pool = token_to_kv_pool
|
|
self.input_buffers = input_buffers
|
|
self.config = config
|
|
self.req_to_page = req_to_page
|
|
self.drafter = drafter
|
|
self.num_warmup = num_warmup
|
|
self.dp_size = config.data_parallel_size
|
|
|
|
self.capture_buckets = get_prefill_token_buckets(config)
|
|
self.disable = (
|
|
config.enforce_eager
|
|
or config.disable_prefill_graph
|
|
or not self.capture_buckets
|
|
or self.inner_model is None
|
|
or self._embed_tokens is None
|
|
or model_runner is None
|
|
or not model_runner.is_generation
|
|
# DP replay decisions must come from replicated state, and a
|
|
# forward's multimodal-ness is rank-local: one rank running its mm
|
|
# prefill eager while text-only peers replay desyncs the EP
|
|
# collectives. Until the DP metadata gather carries a multimodal
|
|
# flag, keep the graph off for multimodal models under DP.
|
|
or (config.data_parallel_size > 1 and model_runner.is_multimodal)
|
|
)
|
|
|
|
self._ctx: ForwardContext | None = None
|
|
self._pool = None
|
|
self._engaged_logged: set[str] = set()
|
|
# Aux-capture mode baked into the graphs; mismatched live forwards run eager.
|
|
self._captured_hidden_mode = None
|
|
# One captured graph + bucket-sized output per padded token bucket.
|
|
self._captures: dict[int, BreakableCapture] = {}
|
|
self._outputs: dict[int, CapturedForward] = {}
|
|
|
|
if not self.disable:
|
|
self.capture(decode_wrapper)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Graph capture
|
|
# ------------------------------------------------------------------
|
|
|
|
def capture(self, decode_wrapper: CudaGraphWrapper | None = None) -> None:
|
|
"""Capture one breakable graph per token bucket (no-op when disabled).
|
|
|
|
Called from ``__init__``; ``decode_wrapper`` supplies the shared
|
|
capture stream and dummy paged-cache block tables (used here only,
|
|
not stored). Buckets share one PRIVATE mempool (first capture
|
|
allocates it), so graph memory stays ~the largest bucket's peak --
|
|
but never the decode graphs' pool: eager ops cache raw pointers to
|
|
buffers they lazily allocated inside a decode capture (flashinfer's
|
|
trtllm-gen MoE runner), and a prefill capture reusing those freed
|
|
blocks means every replay rewrites them, corrupting the next eager
|
|
call (IMA; A/B-proven on qwen3.5 MTP).
|
|
|
|
Runs under inference mode like serving forwards (in-place updates on
|
|
inference-mode model state buffers are only legal there). OOM fails
|
|
the boot LOUDLY (the graph pool did not fit next to weights + KV
|
|
cache; the operator decides: free headroom, lower
|
|
``--prefill-graph-max-tokens``, or 0 to disable). Any other failure
|
|
means the dummy-batch machinery doesn't cover this model family yet:
|
|
degrade to eager prefill instead of crashing the server, and agree on
|
|
that across the world (a MIN all-reduce over the success flag) --
|
|
replay force-sets ``global_num_tokens`` on every rank, so one eager
|
|
rank among replaying peers diverges the token counts and deadlocks
|
|
the next collective.
|
|
"""
|
|
if self.disable:
|
|
return
|
|
weight = self._embed_tokens.weight
|
|
self._input_embeds_buf = torch.zeros(
|
|
max(self.capture_buckets),
|
|
weight.shape[1],
|
|
dtype=weight.dtype,
|
|
device=weight.device,
|
|
)
|
|
captured_ok = True
|
|
try:
|
|
with maybe_inference_mode():
|
|
self._capture_all_buckets(decode_wrapper)
|
|
except torch.cuda.OutOfMemoryError:
|
|
logger.error(
|
|
"Prefill graph capture ran out of GPU memory. Free up "
|
|
"--gpu-memory-utilization headroom, lower "
|
|
"--prefill-graph-max-tokens (default %d), or set it to 0 to "
|
|
"disable the prefill graph.",
|
|
2048,
|
|
)
|
|
raise
|
|
except (NotImplementedError, AttributeError, KeyError, RuntimeError) as exc:
|
|
logger.warning(
|
|
"Prefill graph capture failed (%s: %s); falling back to eager "
|
|
"prefill. This model family may need dedicated dummy-batch support.",
|
|
type(exc).__name__,
|
|
exc,
|
|
)
|
|
captured_ok = False
|
|
if not self._capture_unanimous(captured_ok):
|
|
self.disable = True
|
|
|
|
def _capture_all_buckets(self, decode_wrapper: CudaGraphWrapper | None) -> None:
|
|
for bucket in sorted(self.capture_buckets, reverse=True):
|
|
self._ctx = self._make_dummy_batch(bucket, decode_wrapper)
|
|
self._land_input_embeds(
|
|
self._embed_tokens(self.input_buffers.input_ids_buf[:bucket]), bucket
|
|
)
|
|
self._captured_hidden_mode = self._ctx.capture_hidden_mode
|
|
# Breaks record the ambient dummy ctx; it is rebound live at replay.
|
|
try:
|
|
with active_forward(self._ctx):
|
|
self._capture_bucket(bucket, decode_wrapper)
|
|
finally:
|
|
self._ctx = None
|
|
if self.config.global_rank == 0:
|
|
sample = next(iter(self._captures.values()), None)
|
|
logger.info(
|
|
"prefill breakable graph: captured buckets %s (segments=%d, eager "
|
|
"attention breaks)",
|
|
sorted(self._captures),
|
|
sample.num_segments if sample is not None else 0,
|
|
)
|
|
|
|
def _capture_bucket(
|
|
self, bucket: int, decode_wrapper: CudaGraphWrapper | None
|
|
) -> None:
|
|
"""Warm up and capture the breakable graph for ``bucket`` from the buffers."""
|
|
for _ in range(self.num_warmup):
|
|
self._run_inner(bucket)
|
|
torch.cuda.synchronize()
|
|
stream = decode_wrapper.stream if decode_wrapper is not None else None
|
|
cap = BreakableCapture(pool=self._pool, stream=stream)
|
|
with cap:
|
|
self._outputs[bucket] = CapturedForward(*self._run_inner(bucket))
|
|
if self._pool is None:
|
|
self._pool = cap.pool # share the pool across all subsequent buckets
|
|
cap.replay() # capture records kernels without executing; smoke-test replay
|
|
self._captures[bucket] = cap
|
|
|
|
def _run_inner(self, num_tokens: int):
|
|
"""Run the inner model over the leading ``num_tokens`` of the static buffers.
|
|
|
|
``num_tokens`` is the padded bucket size; the padded tail [real:bucket] is
|
|
already scrubbed to safe values (embeds=0, positions=0,
|
|
out_cache_loc=dummy_kv_slot) by :meth:`_land_input_embeds` and
|
|
``InputBuffers.fill_input_buffers``. The embedding is NOT part of the
|
|
graph: the inner model starts from the static input-embeds buffer, so a
|
|
replay can take precomputed (e.g. merged multimodal) embeddings.
|
|
"""
|
|
ib = self.input_buffers
|
|
if self.config.model_is_mrope:
|
|
positions = ib.mrope_positions_buf[:, :num_tokens]
|
|
else:
|
|
positions = ib.positions_buf[:num_tokens]
|
|
return self.inner_model(
|
|
ib.input_ids_buf[:num_tokens],
|
|
positions,
|
|
self._ctx,
|
|
ib.out_cache_loc_buf[:num_tokens],
|
|
input_embeds=self._input_embeds_buf[:num_tokens],
|
|
)
|
|
|
|
def _land_input_embeds(self, embeds: torch.Tensor, bucket: int) -> None:
|
|
"""Copy ``embeds`` into the static buffer's leading rows, zero the tail.
|
|
|
|
The zeroed padded tail keeps the graphed compute over garbage-free rows
|
|
(RMSNorm of zeros is zeros; the tail is discarded by the output slice).
|
|
"""
|
|
num_tokens = embeds.shape[0]
|
|
self._input_embeds_buf[:num_tokens].copy_(embeds)
|
|
if num_tokens < bucket:
|
|
self._input_embeds_buf[num_tokens:bucket].zero_()
|
|
|
|
def _make_dummy_batch(
|
|
self, num_tokens: int, decode_wrapper: CudaGraphWrapper | None
|
|
) -> ForwardContext:
|
|
"""Populate the static buffers + attention metadata for a dummy bs=1 extend
|
|
forward of ``num_tokens`` tokens, and return its ForwardContext.
|
|
|
|
The prefill analogue of decode's ``_init_capture_metadata``. KV writes
|
|
go to the reserved dummy slot and the page table points at page 0, so
|
|
the forward runs (producing discarded garbage) without touching real
|
|
cache state. Backends with extra paged caches (DeepSeek-V4 DSA: SWA +
|
|
compressor + indexer state) also need per-cache block tables, or their
|
|
extend metadata comes up incomplete and the eager attention break
|
|
aborts the capture -- reuse the decode wrapper's dummy-table builder
|
|
(all zeros, the safe page 0) for those.
|
|
"""
|
|
ib = self.input_buffers
|
|
ib.input_ids_buf[:num_tokens].fill_(1)
|
|
ib.out_cache_loc_buf[:num_tokens].fill_(ib.dummy_kv_slot)
|
|
ib.positions_buf[:num_tokens].copy_(
|
|
torch.arange(num_tokens, device=self.config.device)
|
|
)
|
|
ib.req_pool_indices_buf[:1].zero_()
|
|
ib.seq_lens_buf[:1].fill_(num_tokens)
|
|
ib.extend_seq_lens_buf[:1].fill_(num_tokens)
|
|
ib.extend_seq_lens_cpu[:1].fill_(num_tokens)
|
|
ib.extend_prefix_lens_buf[:1].zero_()
|
|
ib.extend_prefix_lens_cpu[:1].zero_()
|
|
self.req_to_page[0].zero_() # dummy request's pages -> page 0 (valid memory)
|
|
|
|
ctx = ForwardContext(
|
|
attn_backend=self.attn_backend,
|
|
token_to_kv_pool=self.token_to_kv_pool,
|
|
req_to_page=self.req_to_page,
|
|
bs=1,
|
|
num_extends=1,
|
|
input_num_tokens=num_tokens,
|
|
forward_mode=ForwardMode.EXTEND,
|
|
capture_hidden_mode=(
|
|
CaptureHiddenMode.FULL
|
|
if self.drafter is not None
|
|
else CaptureHiddenMode.NULL
|
|
),
|
|
)
|
|
if self.dp_size > 1:
|
|
ctx.global_num_tokens = [num_tokens] * self.config.world_size
|
|
ctx.global_bs = [1] * self.config.world_size
|
|
extra_metadata_kwargs: dict = {}
|
|
if (
|
|
getattr(self.attn_backend, "uses_paged_cache_groups", False)
|
|
and decode_wrapper is not None
|
|
):
|
|
tables = decode_wrapper._capture_paged_cache_block_tables(
|
|
1, self.token_to_kv_pool
|
|
)
|
|
if tables is not None:
|
|
extra_metadata_kwargs["paged_cache_block_tables"] = tables
|
|
extra_metadata_kwargs["num_tokens"] = num_tokens
|
|
extra_metadata_kwargs["positions"] = ib.positions_buf[:num_tokens]
|
|
self.attn_backend.init_forward_metadata(
|
|
bs=1,
|
|
num_extends=1,
|
|
req_pool_indices=ib.req_pool_indices_buf[:1],
|
|
seq_lens=ib.seq_lens_buf[:1],
|
|
req_to_page=self.req_to_page,
|
|
forward_mode=ForwardMode.EXTEND,
|
|
extend_seq_lens=ib.extend_seq_lens_buf[:1],
|
|
extend_seq_lens_cpu=ib.extend_seq_lens_cpu[:1],
|
|
extend_prefix_lens=ib.extend_prefix_lens_buf[:1],
|
|
extend_prefix_lens_cpu=ib.extend_prefix_lens_cpu[:1],
|
|
**extra_metadata_kwargs,
|
|
)
|
|
return ctx
|
|
|
|
def _capture_unanimous(self, captured_ok: bool) -> bool:
|
|
"""MIN-reduce capture success across the world (see ``capture``)."""
|
|
if self.config.world_group is None or self.config.world_size <= 1:
|
|
return captured_ok
|
|
from tokenspeed.runtime.distributed.process_group_manager import (
|
|
process_group_manager as pg_manager,
|
|
)
|
|
|
|
cpu_group = pg_manager.get_process_group("gloo", self.config.world_group)
|
|
flag = torch.tensor([1 if captured_ok else 0], dtype=torch.int32)
|
|
torch.distributed.all_reduce(
|
|
flag, op=torch.distributed.ReduceOp.MIN, group=cpu_group
|
|
)
|
|
unanimous = bool(flag.item())
|
|
if not unanimous and captured_ok:
|
|
logger.warning(
|
|
"Prefill graph: a peer rank failed capture; falling back to "
|
|
"eager prefill on all ranks to keep DP/TP token counts in lockstep."
|
|
)
|
|
return unanimous
|
|
|
|
# ------------------------------------------------------------------
|
|
# Replay dispatch
|
|
# ------------------------------------------------------------------
|
|
|
|
def can_run(self, ctx: ForwardContext, multimodal_context=None) -> bool:
|
|
"""Whether this forward replays a captured graph (mirrors decode's can_run).
|
|
|
|
A forward carrying multimodal inputs replays only when the model
|
|
exposes the embeds-only ``multimodal_input_embeds`` seam; models with
|
|
extra per-layer inputs (deepstack) run eager.
|
|
"""
|
|
if multimodal_context is not None and self._multimodal_input_embeds is None:
|
|
return False
|
|
return self._replay_bucket(ctx) is not None
|
|
|
|
def replay(
|
|
self,
|
|
ctx: ForwardContext,
|
|
input_ids: torch.Tensor,
|
|
multimodal_context=None,
|
|
):
|
|
"""Replay the captured graph for ``ctx`` (caller checked :meth:`can_run`).
|
|
|
|
The embedding runs eagerly here, outside the graph: a plain text
|
|
prefill gathers ``embed_tokens(input_ids)`` into the static buffer; a
|
|
multimodal prefill builds the merged text+vision embeddings via the
|
|
model's ``multimodal_input_embeds`` seam (vision encoder included)
|
|
instead -- both replay the same graphs. Then the inner stack replays
|
|
over the padded bucket and the model's eager logits tail finishes on
|
|
the real-token rows.
|
|
"""
|
|
bucket = self._replay_bucket(ctx)
|
|
assert bucket is not None, "replay() called without can_run()"
|
|
self._log_engaged_once(bucket, ctx, multimodal_context is not None)
|
|
num_tokens = ctx.input_num_tokens
|
|
input_embeds = None
|
|
if multimodal_context is not None:
|
|
input_embeds = self._multimodal_input_embeds(
|
|
input_ids, ctx, multimodal_context
|
|
)
|
|
self._land_input_embeds(
|
|
input_embeds if input_embeds is not None else self._embed_tokens(input_ids),
|
|
bucket,
|
|
)
|
|
with self._padded_to(ctx, bucket):
|
|
self._captures[bucket].replay()
|
|
hidden_states, aux_hidden_states = self._outputs[bucket].sliced(num_tokens)
|
|
# The eager logits tail of BaseCausalLM.forward, on the replayed hidden states.
|
|
logits_metadata = LogitsMetadata.from_forward_context(ctx)
|
|
return self.text_model.logits_processor(
|
|
input_ids,
|
|
hidden_states,
|
|
self.text_model.lm_head,
|
|
logits_metadata,
|
|
aux_hidden_states,
|
|
)
|
|
|
|
def _replay_bucket(self, ctx: ForwardContext) -> int | None:
|
|
"""The captured bucket this forward replays, or ``None`` to run eager.
|
|
|
|
Pure-extend AND mixed extend+decode batches are eligible: the attention
|
|
break reads the LIVE ambient ctx and dispatches the prefill/decode
|
|
split itself, while the captured token-shaped compute is uniform over
|
|
all rows (pure decode is the decode graph's job). Two ctx fields are
|
|
baked into the captured segments rather than rebound at replay -- the
|
|
draft first-step row narrowing (keyed on ``accept_lengths``) and the
|
|
``capture_hidden_mode`` aux-hidden capture -- so a live forward carrying
|
|
different values falls back to eager rather than silently dropping the
|
|
reduce / mismatching aux. Prefix caching (cache hits and chunked-prefill
|
|
chunks 2+) IS eligible: the prefix affects only the ragged attention,
|
|
which runs entirely inside the eager break, and it adds zero new tokens,
|
|
so the padded bucket -- hence the baked EP all-to-all shape under DP --
|
|
is identical on prefix and non-prefix ranks.
|
|
"""
|
|
if self.disable or ctx.forward_mode is None:
|
|
return None
|
|
if ctx.num_extends <= 0:
|
|
return None
|
|
if not (ctx.forward_mode.is_extend() or ctx.forward_mode.is_mixed()):
|
|
return None
|
|
if ctx.accept_lengths is not None:
|
|
return None
|
|
if ctx.capture_hidden_mode != self._captured_hidden_mode:
|
|
return None
|
|
bucket = self._select_bucket(ctx)
|
|
if bucket is None or bucket not in self._captures:
|
|
return None
|
|
return bucket
|
|
|
|
def _select_bucket(self, ctx: ForwardContext) -> int | None:
|
|
"""The padded bucket for this forward, or ``None`` to run eager.
|
|
|
|
Under data parallelism the MoE expert-parallel all-to-all is a collective
|
|
across ALL ranks, sized from a replicated per-rank token list. The captured
|
|
graph bakes a uniform ``[bucket]*world_size`` layout, so every rank must
|
|
replay the SAME bucket or the collective desyncs (NCCL deadlock). Decide
|
|
purely from replicated global state -- the all-extend flag and the global
|
|
max token count -- so all ranks reach the identical decision/bucket with no
|
|
extra sync (mirrors the decode graph). Idle ranks run a DECODE forward, so
|
|
``all_extend`` is False whenever any rank is idle and the graph stays off
|
|
(e.g. warmup), correctly falling back to eager.
|
|
"""
|
|
if self.dp_size <= 1 or ctx.global_num_tokens is None:
|
|
return self._padded_bucket(ctx.input_num_tokens)
|
|
if not ctx.all_extend:
|
|
return None
|
|
return self._padded_bucket(max(ctx.global_num_tokens))
|
|
|
|
def _padded_bucket(self, num_tokens: int) -> int | None:
|
|
"""Smallest bucket >= ``num_tokens``, or ``None`` if over the largest.
|
|
|
|
With ``--disable-cuda-graph-padding``, only an exact bucket match
|
|
replays (mirroring the decode wrapper's no-padding semantics).
|
|
"""
|
|
idx = bisect.bisect_left(self.capture_buckets, num_tokens)
|
|
if idx == len(self.capture_buckets):
|
|
return None
|
|
bucket = self.capture_buckets[idx]
|
|
if self.config.disable_cuda_graph_padding and bucket != num_tokens:
|
|
return None
|
|
return bucket
|
|
|
|
@contextmanager
|
|
def _padded_to(self, ctx: ForwardContext, bucket: int):
|
|
"""Publish ``ctx`` as the ambient live context, pinned to the padded bucket.
|
|
|
|
The graph replays over ``bucket`` (padded) tokens; attention metadata stays
|
|
at the real count (set upstream), so the eager attention break only touches
|
|
real tokens and the padded rows produce discarded garbage. Pin
|
|
``input_num_tokens`` to the bucket and, under DP, ``global_num_tokens`` /
|
|
``global_bs`` to the captured uniform layout so any live read during the
|
|
break matches the baked EP shapes. The break reads ``forward_mode`` / ``bs``
|
|
/ ``num_extends`` LIVE off this same (ambient) ctx -- which we do NOT pin --
|
|
so models split prefill vs decode and dispatch the per-mode backend
|
|
correctly with no side channel.
|
|
"""
|
|
saved = (ctx.input_num_tokens, ctx.global_num_tokens, ctx.global_bs)
|
|
ctx.input_num_tokens = bucket
|
|
if self.dp_size > 1 and ctx.global_num_tokens is not None:
|
|
ctx.global_num_tokens = [bucket] * self.config.world_size
|
|
ctx.global_bs = [1] * self.config.world_size
|
|
try:
|
|
with active_forward(ctx):
|
|
yield
|
|
finally:
|
|
ctx.input_num_tokens, ctx.global_num_tokens, ctx.global_bs = saved
|
|
|
|
def _log_engaged_once(
|
|
self, bucket: int, ctx: ForwardContext, is_multimodal: bool
|
|
) -> None:
|
|
kind = "multimodal" if is_multimodal else "text"
|
|
if kind in self._engaged_logged:
|
|
return
|
|
self._engaged_logged.add(kind)
|
|
logger.info(
|
|
"prefill breakable graph ENGAGED (%s): bucket=%d dp=%s mode=%s "
|
|
"(mixed prefill+decode batches supported)",
|
|
kind,
|
|
bucket,
|
|
# The replay mode actually taken (mirrors _select_bucket), a DP-debug anchor.
|
|
self.dp_size > 1 and ctx.global_num_tokens is not None,
|
|
ctx.forward_mode,
|
|
)
|