Repository navigation
Conversation
…ail-row bucket Eager prefill steps under decoder SWA bounded replay run layers after the last kv_source layer on each request's last 128 rows. Those rows now replay a breakable CUDA graph captured per 128-row bucket on first use; the SWA KV store, attention and low-ratio sources stay eager breaks reading the live tail metadata. Off by default: SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS=0.
Kernels that size cached workspaces by row count reallocate them on a larger step. A graph captured before the first full prefill chunk kept reading the freed workspace and faulted once flush_cache emptied the CUDA cache.
The TP-padded query buffer is cached on the tail metadata of the step that first needs it, here an eager warmup of the capture. The graph writes it by address on every replay, so the metadata must outlive that step.
…d replay The bank captures tail-row buckets up to 2048 rows (16 buckets, 0.11 GiB per GPU on GB300 TP4) whenever --enable-decoder-swa-bounded-replay is set; SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS=0 keeps the late layers eager.
…full-layer graphs The decoder-replay bank becomes one instance of EagerReplayGraphs; a second instance replays the full-width layers before the tail on eager prefill steps, per 256-token bucket (SGLANG_DSV4_FULL_LAYER_GRAPH_MAX_TOKENS). Every bucket is captured at startup from dummy eager forwards, after a full-chunk forward has grown every row-sized workspace; a missing bucket runs eagerly. Each replay asserts that the buffers its kernels read by address have not moved.
…ager get_is_capture_mode() is true under any breakable-graph flag. Inside the eager replay graphs it sent ratio-0 layers down the multi-stream prepare, whose fused K store and request-window gather were captured with the capture step's slots, and it skipped the window history check, leaving the next eager layer to check after this step's commits.
Skip zeroing the pad rows of the attention break's output (no real row reads them) and compute the step's token-to-request map once for all low-ratio layers instead of once per layer. SGLANG_DSV4_EAGER_GRAPH_LEAN_BREAKS=0 restores the previous behavior.
…ss layers The request, position, K-slot, request-start and compressed-length tensors of the prefill indexer data depend on the step and the compress ratio only; build them once per step and ratio instead of once per index layer. Gated by SGLANG_DSV4_EAGER_GRAPH_LEAN_BREAKS.
…p blocks The full-layer graphs and the decoder-replay bank never run concurrently, so they capture into one private pool. After startup capture the allocator still cached the warmups' activation blocks of every bucket size; empty_cache returns them, leaving headroom for long-context indexer scoring.
…s' pool With the prefill CUDA graph also on, its pool and these graphs' own pool do not both fit beside a 0.8 static memory fraction (waves of 16 hits on a 64K prefix ran out of memory). The two kinds of graph never run concurrently, so SGLANG_DSV4_EAGER_GRAPH_GLOBAL_POOL=1 captures into the global pool. Off by default in this commit.
…th the global pool The caching allocator hands a pool's free blocks only to allocations on the stream that freed them. On a private stream the eager replay graphs reused none of the prefill graphs' pool: startup free memory stayed at 7.16 GiB.
…efault Validated at --mem-fraction-static 0.8 with the prefill CUDA graph also on: free memory after startup capture 11.58 GiB (7.16 with a private pool), a wave of 16 hits on a 64K prefix runs, and GSM8K-200 and needle first-token KL match the bank-only reference.
…ore and the attention The decoder replay graphs broke twice per late layer: the SWA K store (live slots) and the attention. Nothing between them reads the SWA cache, so the store now runs at the start of the attention break, which halves the tail's breaks and segment launches (SGLANG_DSV4_EAGER_GRAPH_MERGED_KV_STORE).
The eager attention break wraps the plain attention body instead of the registered op wrapper, so a break pays one Python call.
…use the K store's rope view A decoder replay bucket the tail fills exactly carries no pad rows, so the break no longer slices query, KV and out_cache_loc or slices the output twice to compare element counts. The fused K norm-rope store reuses the real view of each layer's (static) freqs_cis instead of rebuilding it.
Under --enable-decoder-swa-bounded-replay the full-layer graphs now capture 256-token buckets up to the chunked prefill size, at most 8192 tokens (32 buckets, 12.6 GiB per GPU on GB300 TP4). SGLANG_DSV4_FULL_LAYER_GRAPH_MAX_TOKENS sets another cap; 0 keeps layers 0-20 eager.
YAMY1234
force-pushed
the
dsv41-full-layer-graphs
branch
from
October 12, 2026 18:07
e1d433f to
af22d99
Compare
Conflicts in models/deepseek_v4.py: - DeepseekV4Model.__init__: keep main's hash-layer assert and late_layer_vl_routing, then build the eager replay graphs; skip them under DP attention, whose late layers gather rows across ranks per step. - forward: keep the eager-step early return from the decoder replay graphs before main's late-layer DP switch (never both: no graphs under DP). - sgl-project#43833 removed seam_open from forward_hc_pre_from_prev; the full-layer and late-layer graph bodies drop it too.
…play graphs Under the breakable prefill graph, decoder SWA bounded replay ran layers past the last kv_source layer over a fixed block of min(max_running_requests * 128, bucket) rows, i.e. the whole bucket at mrr 64. They now run as one eager break: switch onto the live tail rows' metadata, replay the decoder replay graph of that row bucket (eager past the last bucket or with multimodal inputs), and hand the rows back first in the block, which the logits already slice by real_rows_count. The fixed block stays the path under DP attention, DSpark capture and with SGLANG_DSV4_DECODER_REPLAY_GRAPH_MAX_ROWS=0. The per-step caches of the eager breaks (token-to-request rows, indexer step rows) stay off on prefill-graph metadata, which outlives the step.
…e the window The late layers' eager break runs on fresh tail metadata, while the captured candidate-source layer publishes no candidate table. Inside a prefill-graph step every block is a candidate (the captured indexer already applies the full top-k there), so the late consumers do the same.
The decoder replay graphs entered enable_breakable_cuda_graph() around every run. Inside a prefill-graph step the flag is already set, and the nested scope cleared it on exit, so the rest of that step's warmup ran its low-ratio sources eagerly against the graph's metadata. Enter the scope only when it is not already open.
A replayed batch is padded to the captured token count, so the late layers' gather of the earlier layers' top-k rows sliced a single request's contiguous tail past the extend into the padding. The live tail now carries real_rows_count, as the graph's own tail block does.
YAMY1234
changed the title
This was referenced Oct 15, 2026
Closed
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Stacked on #43534; review that one first. This PR adds one commit range on top of it.
With the late layers replayed from graphs (#43534), layers 0–20 of an eager prefill step still run at full width with one launch per kernel, so cold prefill is still paced by the host. Repro: the same launch and
prefill_repro.pyas in #43534.Modifications
EagerReplayGraphsgeneralizes the bank to a layer range. Two ranges use it:SGLANG_DSV4_EAGER_GRAPH_GLOBAL_POOL=1, default), because they never run at the same time as the prefill graphs.SGLANG_DSV4_EAGER_GRAPH_LEAN_BREAKS=0disables them.SGLANG_DSV4_FULL_LAYER_GRAPH_MAX_TOKENS=-1= the chunked prefill size, at most 8,192 (32 graphs);0turns the full-layer graphs off.Memory (per GPU, defaults,
--mem-fraction-static 0.8): free memory after startup capture drops from 49.8 to 23.3 GiB, of which 15.0 GiB is logged as graph memory; the largest bucket is the biggest share. KV capacity is unchanged, because the KV pool is sized before capture. Startup capture takes about 19 s.SGLANG_DSV4_FULL_LAYER_GRAPH_MAX_TOKENS=4096captures half the buckets and saves about half of that memory; hit steps up to 4K new tokens still replay, and 8K cold chunks run layers 0–20 eagerly.Accuracy Tests
Same setup as #43534:
Speed Tests and Profiling
Same setup and node as #43534, two rounds in reversed order, averages of both rounds.
With every bucket captured at startup, free memory after startup goes from 49.79 to 23.33 GiB per GPU (14.9 GiB logged as full-layer graphs, 0.1 GiB as the bank). KV capacity is unchanged (62,784,256 tokens). Startup capture takes 17–27 s.