From c87f8c1e4990b992bae5f0e81fae9650f734807c Mon Sep 17 00:00:00 2001 From: xinyuazh Date: Wed, 21 Oct 2026 18:12:08 +0800 Subject: [PATCH 01/13] fix(sglang): keep Flash full decode graph correct with mamba radix. PLE conv state was captured at max_bs, so extra_buffer radix recycled a higher mamba index into an undersized buffer and faulted. Size it to the physical pool, keep Hybrid's static slot indices, and snapshot mamba state after decode-graph replay. Co-authored-by: Cursor --- .../sglang/attention_backend/attention_gdn.py | 198 ++++++++++++++++-- atom/plugin/sglang/qwen4_exp_bridge.py | 87 +++++++- atom/plugin/sglang/register.py | 30 +++ 3 files changed, 293 insertions(+), 22 deletions(-) diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py index 2f68a1f81b..b4ec99b3b0 100644 --- a/atom/plugin/sglang/attention_backend/attention_gdn.py +++ b/atom/plugin/sglang/attention_backend/attention_gdn.py @@ -49,11 +49,23 @@ def _align_flydsl_decode_slots(index: torch.Tensor, batch: int) -> torch.Tensor: - """Contiguous int32 prefix of length ``batch`` for FlyDSL decode.""" + """Contiguous int32 prefix of length ``batch`` for FlyDSL decode. + + Must stay a view into Hybrid's static cuda-graph buffer. A ``.clone()`` / + fresh ``.contiguous()`` allocation here bakes a throwaway address into the + decode graph (see ``_linear_static_ple_slots``). + """ if index.shape[0] != batch: index = index[:batch] - if index.dtype != torch.int32 or not index.is_contiguous(): - index = index.to(dtype=torch.int32).contiguous() + if index.dtype != torch.int32: + raise RuntimeError( + f"GDN decode slot index must be int32 static buffer, got {index.dtype}" + ) + if not index.is_contiguous(): + raise RuntimeError( + "GDN decode slot index must be a contiguous view of the Hybrid " + "cuda-graph mamba_cache_indices buffer" + ) return index @@ -171,6 +183,157 @@ def _flydsl_prefill_metadata( return metadata +def _track_mamba_radix_slots(forward_batch: Any, layer_id: int) -> None: + """Mirror community ``GDNAttnBackend`` mamba track after Native GDN compute. + + Community bakes ``_track_mamba_state_decode`` inside the decode CUDA graph. + Native GDN + FlyDSL has repeatedly faulted with that binding under + radix/extra_buffer, so decode-graph capture **skips** track here and + ``track_mamba_radix_after_decode_graph`` (post-replay epilogue) performs + the same snapshot with live Hybrid metadata. Eager prefill/extend still + tracks inline. Never host-sync on ``mask.any()``. + + ``check_freed_slots=True`` skips ``src/dst == -1`` (graph pad / tombstones); + the track kernel otherwise faults on ``ptr + (-1)*stride``. + """ + if getattr(forward_batch, "mamba_track_mask", None) is None: + return + # Do not bake track into the full decode graph; epilogue owns decode track. + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return + attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) + if attn_backend is None: + return + linear = SGLangGDNForwardContext._linear_attn_backend(attn_backend) + if linear is None or getattr(linear, "forward_metadata", None) is None: + return + fm = linear.forward_metadata + pool = getattr(linear, "req_to_token_pool", None) + if pool is None: + return + try: + layer_cache = pool.mamba2_layer_cache(layer_id) + except Exception: # noqa: BLE001 + return + cache_indices = getattr(fm, "mamba_cache_indices", None) + track_indices = getattr(fm, "mamba_track_indices", None) + if cache_indices is None or track_indices is None: + return + + bs = int(forward_batch.batch_size) + # Match PR #32555 / Hybrid replay: only the refreshed [:bs] prefix is live. + if cache_indices.shape[0] != bs: + cache_indices = cache_indices[:bs] + if track_indices.shape[0] != bs: + track_indices = track_indices[:bs] + track_mask = forward_batch.mamba_track_mask + if track_mask is not None and track_mask.shape[0] != bs: + track_mask = track_mask[:bs] + + mode = forward_batch.forward_mode + if mode.is_extend() and not mode.is_target_verify(): + # Community also writes conv at track boundaries inside causal_conv; + # Native only has the finished state. Apply final-src extend copies + # when no chunk ``h`` is required, then the same decode-style row copy. + track_extend = getattr(linear, "_track_mamba_state_extend", None) + h_src = getattr(fm, "track_ssm_h_src", None) + needs_h = h_src is not None and int(h_src.numel()) > 0 + if ( + track_extend is not None + and not needs_h + and getattr(fm, "has_mamba_track_mask", False) + ): + track_extend(forward_batch, None, layer_cache.temporal, fm) + + # Prefer community coalesced all-layers track at the last mamba layer + # (same as ``GDNAttnBackend.forward_decode(..., layer.layer_id)``). Always + # pass check_freed_slots=True: Flash full-graph pads use PAD_SLOT_ID=-1 and + # the track kernel faults on ``ptr + (-1)*stride`` when that flag is off. + from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + track_mamba_states_if_needed, + ) + + track_pools = getattr(linear, "_track_pools", None) + pools = track_pools() if callable(track_pools) else None + if pools is not None: + conv_pool, ssm_pool, last_pool_idx = pools + if pool.mamba_map.get(layer_id) == last_pool_idx: + from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( + track_mamba_states_all_layers, + ) + + track_mamba_states_all_layers( + conv_pool, + ssm_pool, + cache_indices, + track_mask, + track_indices, + bs, + check_freed_slots=True, + ) + return + + track_mamba_states_if_needed( + layer_cache.conv[0], + layer_cache.temporal, + cache_indices, + track_mask, + track_indices, + bs, + check_freed_slots=True, + ) + + +def track_mamba_radix_after_decode_graph( + forward_batch: Any, attn_backend: Any = None +) -> None: + """Post-replay mamba track (Native epilogue). + + Community bakes track inside the graph; Native also does that via + ``_track_mamba_radix_slots``. This epilogue re-runs a coalesced track with + the *live* Hybrid metadata after replay so radix/extra_buffer still gets a + correct snapshot if the captured track-save missed a buffer. Cheap relative + to the graphed decode body. + + ``attn_backend`` must be passed from ``DecodeCudaGraphRunner`` — after + ``execute`` returns, SGLang's forward context is gone so + ``resolve_attn_backend(forward_batch)`` is usually ``None``. + """ + if forward_batch is None: + return + if getattr(forward_batch, "mamba_track_mask", None) is None: + return + mode = getattr(forward_batch, "forward_mode", None) + if mode is None or not callable(getattr(mode, "is_decode_or_idle", None)): + return + if not mode.is_decode_or_idle(): + return + if attn_backend is None: + attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) + if attn_backend is None: + return + # Ensure resolve paths inside ``_track_mamba_radix_slots`` see the backend. + try: + if getattr(forward_batch, "attn_backend", None) is None: + forward_batch.attn_backend = attn_backend + except Exception: # noqa: BLE001, S110 + pass + SGLangGDNForwardContext._patch_forward_batch_pools(forward_batch, attn_backend) + linear = SGLangGDNForwardContext._linear_attn_backend(attn_backend) + pool = getattr(linear, "req_to_token_pool", None) if linear is not None else None + mamba_map = getattr(pool, "mamba_map", None) if pool is not None else None + if not mamba_map: + return + track_pools = getattr(linear, "_track_pools", None) + pools = track_pools() if callable(track_pools) else None + if pools is not None: + # Coalesced all-layers track fires only at the last mamba layer. + _track_mamba_radix_slots(forward_batch, int(max(mamba_map.keys()))) + return + for layer_id in mamba_map: + _track_mamba_radix_slots(forward_batch, int(layer_id)) + + class SGLangGatedDeltaNet(GatedDeltaNet): """Run batched ATOM GDN while filling SGLang's verify snapshots.""" @@ -214,6 +377,8 @@ def forward( and not draft_extend ): note_full_state_write(self._verify_cache_slots(forward_batch)) + if forward_batch is not None: + _track_mamba_radix_slots(forward_batch, int(self.layer_num)) return out attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) @@ -819,17 +984,26 @@ def _build_gdn_metadata( return None query_start_loc, idx = reconstructed device = query_start_loc.device - idx = idx.to(dtype=torch.int32, device=device) - # Native GDN pad clone. Hybrid already pads Flash like 2.4T; - # drop this after Native GDN consumes Hybrid buffers directly. - # See apply_gdn_pad_sentinels in qwen4_exp_gdn_pad.py. - idx, query_start_loc = apply_gdn_pad_sentinels( - forward_batch, idx, query_start_loc, mode, bs - ) + # Decode cuda-graph: keep Hybrid's static buffers (already pad=-1 on + # dead rows in ``_replay_metadata``). Cloning here bakes a throwaway + # address into the full decode graph and Native GDN keeps writing the + # capture-time slots → GPU fault under radix/extra_buffer recycle. + # Prefill/extend is eager; pad-sentinel clone is fine there. if mode.is_decode_or_idle(): - # FlyDSL checks the slot vector against q's batch axis. The graph - # buffer is the capture bucket; keep a contiguous int32 prefix. + # Never ``.to()`` / clone here — that bakes a throwaway address into + # the full decode graph. Hybrid ``state_indices_list`` is already + # int32 on device; ``_align_flydsl_decode_slots`` enforces the view. + if query_start_loc.device != device: + raise RuntimeError( + "GDN decode query_start_loc must already live on the Hybrid " + f"cuda-graph device ({device}), got {query_start_loc.device}" + ) idx = _align_flydsl_decode_slots(idx, bs) + else: + idx = idx.to(dtype=torch.int32, device=device) + idx, query_start_loc = apply_gdn_pad_sentinels( + forward_batch, idx, query_start_loc, mode, bs + ) common_kwargs = { "num_spec_decodes": 0, "num_spec_decode_tokens": 0, diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index 9d46ca42e9..81c64b3d2b 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1168,24 +1168,79 @@ def build_qsa_metadata( return qsa +def _mamba_physical_slots(pool: Any) -> int: + """Rows in the physical mamba conv/temporal tensors (0 if unresolved).""" + if pool is None: + return 0 + slots = 0 + mapping = getattr(pool, "req_index_to_mamba_index_mapping", None) + if torch.is_tensor(mapping) and mapping.numel(): + slots = max(slots, int(mapping.numel())) + slots = _max_int_attrs(pool, ("size", "max_num_reqs", "mamba_size"), slots) + mamba_pool = getattr(pool, "mamba_pool", None) or pool + mamba_cache = getattr(mamba_pool, "mamba_cache", None) + if mamba_cache is not None: + conv = getattr(mamba_cache, "conv", None) + if conv is not None and len(conv) > 0 and hasattr(conv[0], "shape"): + slots = max(slots, int(conv[0].shape[0])) + temporal = getattr(mamba_cache, "temporal", None) + if temporal is not None and hasattr(temporal, "shape"): + slots = max(slots, int(temporal.shape[0])) + return slots + + def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: """Native allocates PLE conv state for the whole per-req pool, not max_bs. Decode CUDA graphs bake ``conv_state``'s address. Growing after capture frees that storage; undersizing it makes a recycled mamba slot OOB on the next eager prefill (the conc=4 wave-2 HSA). + + Flash + ``extra_buffer`` radix keeps finished track slots in the same + physical pool; the capture-time floor must cover the full mamba pool, not + ``max_bs`` (otherwise unique-stress dies around the first wave of recycled + indices with a GPU fault under full decode graph). + + Capture often builds PLE before ``forward_batch.req_to_token_pool`` is + attached; fall back to the Hybrid linear child's pool and a high floor + (≥ allocated ``max_mamba_cache_size``, e.g. 323 here). """ - slots = max(int(_DECODE_GRAPH.max_bs), 16) + # High floor: capture may run before pool attrs resolve; never bake a + # max_bs-sized PLE buffer for radix/extra_buffer. + slots = max(int(_DECODE_GRAPH.max_bs), 512) + args = _server_args() + if args is not None: + for key in ("max_mamba_cache_size", "max_total_tokens"): + val = getattr(args, key, None) + if val: + slots = max(slots, int(val)) + running = int(getattr(args, "max_running_requests", 0) or 0) + if running > 0: + # extra_buffer ≈ 2x; radix retains finished track slots on top. + slots = max(slots, running * 64) + # SGLang resolves the real mamba capacity onto the schedule object after + # memory-pool init (launch-time server_args.max_mamba_cache_size stays None). + try: + from sglang.srt.runtime_context import get_schedule + + sched = get_schedule() + mamba_cap = getattr(sched, "max_mamba_cache_size", None) + if mamba_cap: + slots = max(slots, int(mamba_cap)) + except Exception: # noqa: BLE001 + pass backend = resolve_attn_backend(forward_batch) - pool = resolve_mamba_req_pool(forward_batch, backend) - if pool is not None: - mapping = getattr(pool, "req_index_to_mamba_index_mapping", None) - if torch.is_tensor(mapping) and mapping.numel(): - slots = max(slots, int(mapping.numel())) - slots = _max_int_attrs(pool, ("size", "max_num_reqs", "mamba_size"), slots) - slots = _max_int_attrs( - _req_to_token_pool(forward_batch), ("size", "max_num_reqs"), slots - ) + linear = getattr(backend, "linear_attn_backend", None) or backend + pool = resolve_mamba_req_pool(forward_batch, linear) + if pool is None: + pool = getattr(linear, "req_to_token_pool", None) + slots = max(slots, _mamba_physical_slots(pool)) + try: + slots = max( + slots, _mamba_physical_slots(_req_to_token_pool(forward_batch)) + ) + except Exception: # noqa: BLE001 - batch may lack the pool during early warmups + pass # Host indices only. A GPU ``idx.max().item()`` drains verify, which is # still queued when eager draft builds PLE state. Pool capacity above # already matches Native's whole-pool allocation. @@ -1249,6 +1304,7 @@ def _ensure_ple_states( hf = _hf_text_config(atom_config) ngram = max(int(getattr(hf, "ngram_size", 3)), 1) rows = max(num_slots, 1) + existing = getattr(model, "_atom_ple_conv_state", None) conv_state = _grow_once_tensor( model, "_atom_ple_conv_state", @@ -1267,6 +1323,17 @@ def _ensure_ple_states( dtype=torch.int64, fill=int(eos[0] if isinstance(eos, (list, tuple)) else eos), ) + if not getattr(_ensure_ple_states, "_logged", False) or ( + existing is not None and int(existing.shape[0]) < int(conv_state.shape[0]) + ): + logger.info( + "Flash PLE state pool rows=%s (requested=%s active=%s capturing=%s)", + int(conv_state.shape[0]), + rows, + bool(_DECODE_GRAPH.active), + bool(_is_capturing()), + ) + _ensure_ple_states._logged = True # type: ignore[attr-defined] return conv_state, ngram_state diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py index da241cb57f..ce1926bc15 100644 --- a/atom/plugin/sglang/register.py +++ b/atom/plugin/sglang/register.py @@ -145,6 +145,36 @@ def capture_one_shape_with_forward_context(self, *args, **kwargs): return original_capture_one_shape(self, *args, **kwargs) DecodeCudaGraphRunner.capture_one_shape = capture_one_shape_with_forward_context + + # Post-replay mamba track epilogue for Native GDN (community bakes track + # inside the graph; Native also bakes via ``_track_mamba_radix_slots``). + # Re-snapshot after replay so radix/extra_buffer never depends solely on + # the captured track-save binding. + if not getattr(DecodeCudaGraphRunner, "_atom_mamba_track_epilogue", False): + original_execute = DecodeCudaGraphRunner.execute + + def execute_with_mamba_track_epilogue(self, forward_batch, *args, **kwargs): + output = original_execute(self, forward_batch, *args, **kwargs) + try: + from atom.plugin.sglang.attention_backend.attention_gdn import ( + track_mamba_radix_after_decode_graph, + ) + + track_mamba_radix_after_decode_graph( + forward_batch, + attn_backend=self.model_runner.attn_backend, + ) + except Exception: # noqa: BLE001 - never fail the decode path on track + logger.exception("Native GDN post-replay mamba track failed") + return output + + DecodeCudaGraphRunner.execute = execute_with_mamba_track_epilogue + DecodeCudaGraphRunner._atom_mamba_track_epilogue = True + logger.info( + "Installed Native GDN post-replay mamba track epilogue on " + "DecodeCudaGraphRunner.execute" + ) + DecodeCudaGraphRunner._atom_forward_context_patched = True From a8ca307b1fb09e501f4bf530e309a127d4079e19 Mon Sep 17 00:00:00 2001 From: xinyuazh Date: Wed, 21 Oct 2026 18:28:47 +0800 Subject: [PATCH 02/13] fix(sglang): satisfy Black and Ruff on the Flash radix graph fix. Log the optional PLE pool lookups instead of swallowing them, and drop the unused noqa on the decode-graph track epilogue. Co-authored-by: Cursor --- atom/plugin/sglang/qwen4_exp_bridge.py | 12 +++++------- atom/plugin/sglang/register.py | 2 +- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index 81c64b3d2b..cad4d651bb 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1227,8 +1227,8 @@ def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: mamba_cap = getattr(sched, "max_mamba_cache_size", None) if mamba_cap: slots = max(slots, int(mamba_cap)) - except Exception: # noqa: BLE001 - pass + except Exception: + logger.debug("Flash PLE: schedule mamba capacity unavailable", exc_info=True) backend = resolve_attn_backend(forward_batch) linear = getattr(backend, "linear_attn_backend", None) or backend pool = resolve_mamba_req_pool(forward_batch, linear) @@ -1236,11 +1236,9 @@ def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: pool = getattr(linear, "req_to_token_pool", None) slots = max(slots, _mamba_physical_slots(pool)) try: - slots = max( - slots, _mamba_physical_slots(_req_to_token_pool(forward_batch)) - ) - except Exception: # noqa: BLE001 - batch may lack the pool during early warmups - pass + slots = max(slots, _mamba_physical_slots(_req_to_token_pool(forward_batch))) + except Exception: + logger.debug("Flash PLE: req_to_token_pool not attached yet", exc_info=True) # Host indices only. A GPU ``idx.max().item()`` drains verify, which is # still queued when eager draft builds PLE state. Pool capacity above # already matches Native's whole-pool allocation. diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py index ce1926bc15..4b1a301061 100644 --- a/atom/plugin/sglang/register.py +++ b/atom/plugin/sglang/register.py @@ -164,7 +164,7 @@ def execute_with_mamba_track_epilogue(self, forward_batch, *args, **kwargs): forward_batch, attn_backend=self.model_runner.attn_backend, ) - except Exception: # noqa: BLE001 - never fail the decode path on track + except Exception: logger.exception("Native GDN post-replay mamba track failed") return output From b2b554f51d2a09ffe768a7011233abf9690d230d Mon Sep 17 00:00:00 2001 From: zovonoir Date: Wed, 21 Oct 2026 18:31:17 +0800 Subject: [PATCH 03/13] [Plugin][SGLang] Snapshot GDN state for the mamba radix cache (extra_buffer) `--mamba-radix-cache-strategy extra_buffer --page-size 64` is what lets Qwen3.5 run prefix caching with paged KV and overlap scheduling. SGLang's HybridLinearAttnBackend owns that cache and computes the track metadata, but the snapshot calls live inside its own GDN kernels, which the plugin replaces wholesale. Nothing ever wrote the track slots, so a prefix hit resumed from whatever state a recycled slot last held: short retrieval answers usually survive that, through the full-attention layers, but logprobs drift and a 9k-token hit decodes to an empty string. `SGLangGatedDeltaNet` now makes those calls after the ATOM forward, through SGLang's own helpers so the snapshot semantics stay defined in one place: - extend: the conv window from the pre-conv input, and the SSM state from the chunk kernel's per-chunk `h` (unaligned) or the runtime slot (aligned), via `_track_mamba_state_extend`; - decode: `_track_mamba_state_decode` with the layer id, which collapses into one all-layers launch at the last GDN layer. `h` is kept only when SGLang asks for a snapshot: a new `GDNAttentionMetadata.keep_intermediate_states` lets the plugin request it without placing native `ssm_checkpoints`, and it is read back with the existing `pop_last_intermediate_states`. ReplaySSM is turned off whenever the radix cache is on. Under the ring a request's slot holds a checkpoint plus uncommitted records rather than its state, and the commit hook scatters only the conv window, so every state SGLang copied into a track slot would be stale. Speculative serving with the radix cache therefore verifies through the batched snapshot path; forcing `ATOM_ENABLE_REPLAYSSM=1` alongside the cache is rejected. Co-authored-by: Cursor --- atom/model_ops/attention_gdn.py | 4 +- atom/model_ops/attentions/gdn_attn.py | 4 + .../sglang/attention_backend/attention_gdn.py | 72 ++++++ .../sglang/attention_backend/gdn_replayssm.py | 27 ++- tests/plugin/test_sglang_gdn_radix_track.py | 227 ++++++++++++++++++ 5 files changed, 330 insertions(+), 4 deletions(-) create mode 100644 tests/plugin/test_sglang_gdn_radix_track.py diff --git a/atom/model_ops/attention_gdn.py b/atom/model_ops/attention_gdn.py index f6fb5bb439..8f8bdf0f38 100644 --- a/atom/model_ops/attention_gdn.py +++ b/atom/model_ops/attention_gdn.py @@ -504,7 +504,9 @@ def forward( ), # Only when there is somewhere to put them; `h` is large and is # dropped on return otherwise. - keep_intermediate_states=ckpt is not None, + keep_intermediate_states=( + ckpt is not None or gdn_metadata.keep_intermediate_states + ), ) # Init cache ssm_state.index_copy_( diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index 68b5e3665b..ae0c862a8c 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -140,6 +140,10 @@ class GDNAttentionMetadata: ssm_chunk_offsets: torch.Tensor | None = None # Conv tokens a checkpoint holds, `conv_kernel - 1`; see `_conv_state_len`. ssm_conv_state_len: int = 0 + # Keep the chunk kernel's per-chunk states for a caller that places its own + # checkpoints after the forward (the SGLang plugin's mamba radix cache), + # read back with `pop_last_intermediate_states`. `ssm_checkpoints` implies it. + keep_intermediate_states: bool = False # Qwen4Exp-only optional AITER schedule, built once outside the layer loop. flydsl_prefill_metadata: object | None = None # --- ReplaySSM --------------------------------------------------------- diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py index b4ec99b3b0..55a76ecd2c 100644 --- a/atom/plugin/sglang/attention_backend/attention_gdn.py +++ b/atom/plugin/sglang/attention_backend/attention_gdn.py @@ -361,6 +361,7 @@ def forward( self.dt_bias.detach().to(dtype=dtype), requires_grad=False ) out = super().forward(mixed_qkv, b, a, core_attn_out, layer_name) + self._track_mamba_states(forward_batch, mixed_qkv) # Chunked prefill rewrites the checkpoint in place. The ring # cursor has to drop the previous tenant's records or the next # verify folds them on top of that new state. @@ -464,6 +465,73 @@ def forward( return core_attn_out + def _track_mamba_states(self, forward_batch: Any, mixed_qkv: torch.Tensor) -> None: + """Fill SGLang's mamba radix cache (extra_buffer) track slots. + + SGLang's HybridLinearAttnBackend owns the prefix cache and computes the + track metadata, but the snapshot calls live inside its own GDN kernels, + which the plugin replaces. They are made here, after the ATOM forward, + through SGLang's helpers so the snapshot semantics stay defined there. + Every source is final by now: the conv window is read from the pre-conv + input, the SSM state from the chunk kernel's `h` or the runtime slot. + """ + if ( + forward_batch is None + or getattr(forward_batch, "mamba_track_mask", None) is None + ): + return + mode = forward_batch.forward_mode + attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) + if attn_backend is None: + return + linear_backend = SGLangGDNForwardContext._linear_attn_backend(attn_backend) + metadata = getattr(linear_backend, "forward_metadata", None) + if metadata is None: + return + layer_cache = linear_backend.req_to_token_pool.mamba2_layer_cache( + self.layer_num + ) + conv_states = layer_cache.conv[0] + # `h` is in the logical [K, V] layout; FlyDSL decode stores VK. + ssm_states = layer_cache.temporal + policy = getattr(self, "gdn_flydsl_policy", None) + if policy is not None and policy.decode: + ssm_states = ssm_states.transpose(-1, -2) + + if mode.is_decode_or_idle(): + linear_backend._track_mamba_state_decode( + forward_batch, + conv_states, + ssm_states, + metadata.mamba_cache_indices, + self.layer_num, + ) + return + if not mode.is_extend(): + return + + from atom.model_ops.fla_ops.chunk import pop_last_intermediate_states + + h = pop_last_intermediate_states() + if not metadata.has_mamba_track_mask: + return + track_conv_indices = metadata.track_conv_indices + if track_conv_indices is not None and track_conv_indices.numel() > 0: + # Only SGLang's GDN backend caches the destination slots; indexing + # the pool with a None would broadcast over every slot. + dst = getattr(metadata, "conv_states_mask_indices", None) + if dst is None: + rows = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] + dst = metadata.mamba_track_indices[rows] + # [rows, state_len, dim] -> the pool's [slot, dim, state_len]. + conv_states[dst] = mixed_qkv[track_conv_indices].transpose(1, 2) + if metadata.track_ssm_h_src.numel() > 0 and h is None: + raise RuntimeError( + "SGLang mamba radix cache needs the GDN chunk kernel's " + "intermediate states, but this prefill kept none." + ) + linear_backend._track_mamba_state_extend(forward_batch, h, ssm_states, metadata) + def _verify_cache_slots(self, forward_batch: Any) -> torch.Tensor | None: attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) if attn_backend is None: @@ -1068,6 +1136,10 @@ def _build_gdn_metadata( flydsl_prefill_metadata=_flydsl_prefill_metadata( query_start_loc, forward_batch ), + # SGLangGatedDeltaNet snapshots radix-cache states out of `h`. + keep_intermediate_states=bool( + getattr(fm, "has_mamba_track_mask", False) + ), **common_kwargs, ) diff --git a/atom/plugin/sglang/attention_backend/gdn_replayssm.py b/atom/plugin/sglang/attention_backend/gdn_replayssm.py index 95cc6a4119..012a86e431 100644 --- a/atom/plugin/sglang/attention_backend/gdn_replayssm.py +++ b/atom/plugin/sglang/attention_backend/gdn_replayssm.py @@ -65,20 +65,41 @@ def replayssm_enabled() -> bool: ``ATOM_ENABLE_REPLAYSSM=0`` keeps the per-step snapshot path. ``ATOM_ENABLE_REPLAYSSM=1`` forces the ring even without server args (tests). Unset follows whether this process is an MTP server. + + Off whenever SGLang's radix cache is on. The ring leaves a request's slot + holding a checkpoint plus uncommitted records rather than its state, and + the commit hook scatters only the conv window, so every state SGLang + copies into a radix track slot would be stale. """ + args = _server_args() + radix_cache = args is not None and not bool( + getattr(args, "disable_radix_cache", True) + ) override = envs.ATOM_ENABLE_REPLAYSSM if override is not None: + if override and radix_cache: + raise RuntimeError( + "ATOM_ENABLE_REPLAYSSM=1 cannot be combined with the SGLang " + "radix cache: ReplaySSM leaves no materialized GDN state for " + "the cache to snapshot. Unset it or pass --disable-radix-cache." + ) return bool(override) + if radix_cache: + return False return _server_is_speculative() -def _server_is_speculative() -> bool: +def _server_args() -> Any: try: from sglang.srt.runtime_context import get_server_args - args = get_server_args() + return get_server_args() except Exception: # noqa: BLE001 - unit tests have no server - return False + return None + + +def _server_is_speculative() -> bool: + args = _server_args() if args is None: return False if int(getattr(args, "speculative_num_steps", 0) or 0) > 0: diff --git a/tests/plugin/test_sglang_gdn_radix_track.py b/tests/plugin/test_sglang_gdn_radix_track.py new file mode 100644 index 0000000000..960c2984f9 --- /dev/null +++ b/tests/plugin/test_sglang_gdn_radix_track.py @@ -0,0 +1,227 @@ +"""SGLang mamba radix cache (extra_buffer) tracking on the plugin GDN path. + +The plugin replaces SGLang's GDN kernels, so it has to make the track calls +SGLang's own backend makes around them. These tests drive +`SGLangGatedDeltaNet._track_mamba_states` against a stand-in linear backend and +check which helper receives which buffers, plus the ReplaySSM radix guard. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from atom.model_ops.fla_ops import chunk as chunk_mod +from atom.plugin.sglang.attention_backend import attention_gdn, gdn_replayssm +from atom.plugin.sglang.attention_backend.attention_gdn import ( + SGLangGatedDeltaNet, + SGLangGDNForwardContext, +) + +CONV_DIM = 6 +STATE_LEN = 3 +NUM_SLOTS = 8 +LAYER_NUM = 5 + + +class _Mode: + def __init__(self, kind: str): + self.kind = kind + + def is_decode_or_idle(self): + return self.kind == "decode" + + def is_extend(self): + return self.kind == "extend" + + +class _LinearBackend: + def __init__(self, metadata): + self.forward_metadata = metadata + self.layer_cache = SimpleNamespace( + conv=[torch.zeros(NUM_SLOTS, CONV_DIM, STATE_LEN)], + temporal=torch.zeros(NUM_SLOTS, 2, 4, 4), + ) + self.req_to_token_pool = SimpleNamespace( + mamba2_layer_cache=lambda layer: self.layer_cache + ) + self.calls = [] + + def _track_mamba_state_decode(self, fb, conv, ssm, idx, layer_id): + self.calls.append(("decode", conv, ssm, idx, layer_id)) + + def _track_mamba_state_extend(self, fb, h, ssm, metadata): + self.calls.append(("extend", h, ssm, metadata)) + + +def _impl() -> SGLangGatedDeltaNet: + impl = SGLangGatedDeltaNet.__new__(SGLangGatedDeltaNet) + torch.nn.Module.__init__(impl) + impl.layer_num = LAYER_NUM + impl.gdn_flydsl_policy = None + return impl + + +@pytest.fixture +def backend(monkeypatch): + holder = {} + + def resolve(_fb): + return SimpleNamespace(linear_attn_backend=holder["backend"]) + + monkeypatch.setattr( + SGLangGDNForwardContext, "_resolve_attn_backend", staticmethod(resolve) + ) + + def install(metadata): + holder["backend"] = _LinearBackend(metadata) + return holder["backend"] + + return install + + +def test_no_track_mask_is_a_no_op(backend): + lb = backend(SimpleNamespace()) + fb = SimpleNamespace(mamba_track_mask=None, forward_mode=_Mode("decode")) + _impl()._track_mamba_states(fb, torch.zeros(4, CONV_DIM)) + assert lb.calls == [] + + +def test_decode_tracks_through_sglang_with_layer_id(backend): + idx = torch.tensor([2, 3]) + lb = backend(SimpleNamespace(mamba_cache_indices=idx)) + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True, False]), forward_mode=_Mode("decode") + ) + _impl()._track_mamba_states(fb, torch.zeros(2, CONV_DIM)) + (call,) = lb.calls + assert call[0] == "decode" + assert call[1] is lb.layer_cache.conv[0] + assert call[2] is lb.layer_cache.temporal + assert call[3] is idx + assert call[4] == LAYER_NUM + + +def _extend_metadata(track_conv_indices, dst, h_src_rows): + return SimpleNamespace( + has_mamba_track_mask=True, + track_conv_indices=track_conv_indices, + conv_states_mask_indices=dst, + mamba_track_indices=torch.tensor([6, 7]), + track_ssm_h_src=torch.arange(h_src_rows), + ) + + +def test_extend_copies_conv_window_from_pre_conv_input(backend, monkeypatch): + h = torch.randn(1, 3, 2, 4, 4) + monkeypatch.setattr(chunk_mod.ChunkGatedDeltaRuleFunction, "last_h", h) + # Two tracked rows whose windows end at tokens 4 and 9. + track = torch.tensor([[2, 3, 4], [7, 8, 9]]) + lb = backend(_extend_metadata(track, torch.tensor([6, 1]), 1)) + mixed_qkv = torch.arange(10 * CONV_DIM, dtype=torch.float32).view(10, CONV_DIM) + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True, True]), forward_mode=_Mode("extend") + ) + + _impl()._track_mamba_states(fb, mixed_qkv) + + conv = lb.layer_cache.conv[0] + torch.testing.assert_close(conv[6], mixed_qkv[2:5].T) + torch.testing.assert_close(conv[1], mixed_qkv[7:10].T) + assert conv[[0, 2, 3, 4, 5, 7]].abs().sum() == 0 + (call,) = lb.calls + assert call[0] == "extend" and call[1] is h + assert call[2] is lb.layer_cache.temporal + # Consumed, so a later forward cannot read a stale `h`. + assert chunk_mod.ChunkGatedDeltaRuleFunction.last_h is None + + +def test_extend_derives_conv_destination_when_backend_did_not(backend, monkeypatch): + monkeypatch.setattr( + chunk_mod.ChunkGatedDeltaRuleFunction, "last_h", torch.zeros(1, 1, 2, 4, 4) + ) + lb = backend(_extend_metadata(torch.tensor([[1, 2, 3]]), None, 0)) + mixed_qkv = torch.randn(4, CONV_DIM) + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([False, True]), forward_mode=_Mode("extend") + ) + + _impl()._track_mamba_states(fb, mixed_qkv) + + conv = lb.layer_cache.conv[0] + torch.testing.assert_close(conv[7], mixed_qkv[1:4].T) + assert conv[:7].abs().sum() == 0 + + +def test_extend_without_kept_states_fails_closed(backend, monkeypatch): + monkeypatch.setattr(chunk_mod.ChunkGatedDeltaRuleFunction, "last_h", None) + backend(_extend_metadata(torch.empty(0, 3, dtype=torch.long), None, 1)) + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True]), forward_mode=_Mode("extend") + ) + with pytest.raises(RuntimeError, match="intermediate states"): + _impl()._track_mamba_states(fb, torch.zeros(4, CONV_DIM)) + + +def test_extend_metadata_keeps_states_only_when_tracking(monkeypatch): + monkeypatch.setattr(attention_gdn, "_flydsl_prefill_metadata", lambda *a: None) + monkeypatch.setattr( + attention_gdn, + "apply_gdn_pad_sentinels", + lambda fb, idx, qsl, mode, bs: (idx, qsl), + ) + + def build(track: bool): + fm = SimpleNamespace( + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + mamba_cache_indices=torch.tensor([1], dtype=torch.int32), + has_mamba_track_mask=track, + ) + + class _ExtendMode: + def is_target_verify(self): + return False + + def is_decode_or_idle(self): + return False + + def is_extend(self): + return True + + fb = SimpleNamespace( + forward_mode=_ExtendMode(), + batch_size=1, + extend_prefix_lens=torch.tensor([0]), + ) + return SGLangGDNForwardContext._build_gdn_metadata( + fb, SimpleNamespace(forward_metadata=fm) + ) + + assert build(True).keep_intermediate_states is True + assert build(False).keep_intermediate_states is False + + +def _server(monkeypatch, **fields): + args = SimpleNamespace(**fields) + monkeypatch.setattr(gdn_replayssm, "_server_args", lambda: args) + + +def test_replayssm_follows_speculation_without_radix(monkeypatch): + monkeypatch.delenv("ATOM_ENABLE_REPLAYSSM", raising=False) + _server(monkeypatch, disable_radix_cache=True, speculative_algorithm="DFLASH") + assert gdn_replayssm.replayssm_enabled() is True + + +def test_replayssm_is_off_under_radix_cache(monkeypatch): + monkeypatch.delenv("ATOM_ENABLE_REPLAYSSM", raising=False) + _server(monkeypatch, disable_radix_cache=False, speculative_algorithm="DFLASH") + assert gdn_replayssm.replayssm_enabled() is False + + +def test_forcing_replayssm_under_radix_cache_is_rejected(monkeypatch): + monkeypatch.setenv("ATOM_ENABLE_REPLAYSSM", "1") + _server(monkeypatch, disable_radix_cache=False, speculative_algorithm="DFLASH") + with pytest.raises(RuntimeError, match="radix cache"): + gdn_replayssm.replayssm_enabled() From 148b038cdecb45ac8a85de1dcc0d47e9251d0c10 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Wed, 21 Oct 2026 18:31:17 +0800 Subject: [PATCH 04/13] [Plugin][SGLang] Reject Qwen3.5 configurations that corrupt silently ATOM's GDN chunk kernel exposes recurrent states every 64 tokens, while SGLang snapshots at `max(64, page_size)`. An extra_buffer page size other than 64 reads a valid but wrong state row and produces fluent corruption, so refuse it at model construction. DFLASH target verify on Aiter has no custom mask and needs Aiter's unified-attention path; refuse `SGLANG_USE_AITER_UNIFIED_ATTN` unset there rather than crashing on the first verify batch. Server args come from `sglang.srt.runtime_context.get_server_args`; `get_global_server_args` raises on SGLang 0.5.20, which would have turned both checks into silent no-ops. Co-authored-by: Cursor --- atom/plugin/sglang/models/qwen3_5.py | 67 +++++++++++++++++++ .../test_sglang_qwen35_runtime_contract.py | 49 ++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 tests/plugin/test_sglang_qwen35_runtime_contract.py diff --git a/atom/plugin/sglang/models/qwen3_5.py b/atom/plugin/sglang/models/qwen3_5.py index c611c7719e..b013065050 100644 --- a/atom/plugin/sglang/models/qwen3_5.py +++ b/atom/plugin/sglang/models/qwen3_5.py @@ -70,6 +70,72 @@ } +def _validate_qwen35_runtime_contract( + server_args: Any | None = None, + *, + unified_attn_enabled: bool | None = None, +) -> None: + """Reject Qwen3.5 plugin combinations that otherwise fail silently. + + ATOM's GDN chunk kernel exposes recurrent states every 64 tokens, while + SGLang snapshots at ``mamba_cache_chunk_size = max(64, page_size)``. An + ``extra_buffer`` page size other than 64 therefore reads a valid but wrong + recurrent-state row, producing fluent corruption instead of an exception. + + DFLASH target verification with Aiter has no ``custom_mask`` and must use + Aiter's unified-attention path. Detect that at model construction rather + than crashing on the first verify batch. + """ + if server_args is None: + try: + from sglang.srt.runtime_context import get_server_args + + server_args = get_server_args() + except Exception: # noqa: BLE001 - standalone imports have no server + return + if server_args is None: + return + + strategy = str(getattr(server_args, "mamba_radix_cache_strategy", "") or "") + radix_disabled = bool(getattr(server_args, "disable_radix_cache", False)) + if ( + not radix_disabled + and "extra_buffer" in strategy.lower() + and int(getattr(server_args, "page_size", 0) or 0) != 64 + ): + raise ValueError( + "ATOM Qwen3.5 mamba radix cache with extra_buffer requires " + "--page-size 64. Other page sizes select the wrong GDN snapshot " + "offset and can silently produce incorrect output." + ) + + algorithm = str(getattr(server_args, "speculative_algorithm", "") or "").upper() + backends = { + str(getattr(server_args, name, "") or "").lower() + for name in ( + "attention_backend", + "decode_attention_backend", + "prefill_attention_backend", + ) + } + if "DFLASH" not in algorithm or "aiter" not in backends: + return + + if unified_attn_enabled is None: + try: + from sglang.srt.environ import envs as sglang_envs + + unified_attn_enabled = sglang_envs.SGLANG_USE_AITER_UNIFIED_ATTN.get() + except Exception: # noqa: BLE001 - absence means SGLang's false default + unified_attn_enabled = False + if not unified_attn_enabled: + raise ValueError( + "ATOM Qwen3.5 DFLASH with Aiter requires " + "SGLANG_USE_AITER_UNIFIED_ATTN=1. Refusing to defer this " + "misconfiguration to a first-request verify crash." + ) + + def _apply_bf16_in_proj_mapping(mapping: dict, atom_config: Any) -> dict: if atom_config.quant_config.global_quant_config.quant_dtype != torch.bfloat16: return mapping @@ -390,6 +456,7 @@ def __init__( quant_config: SGLangQuantizationConfig | None = None, prefix: str = "", ) -> None: + _validate_qwen35_runtime_contract() self._prepare_sglang_root_config(config) stack_cls = _get_qwen35_language_model_stack_cls( type(self).atom_language_model_cls diff --git a/tests/plugin/test_sglang_qwen35_runtime_contract.py b/tests/plugin/test_sglang_qwen35_runtime_contract.py new file mode 100644 index 0000000000..8479d4106d --- /dev/null +++ b/tests/plugin/test_sglang_qwen35_runtime_contract.py @@ -0,0 +1,49 @@ +"""Qwen3.5 plugin configurations that would otherwise fail silently.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from atom.plugin.sglang.models.qwen3_5 import _validate_qwen35_runtime_contract + + +def _args(**overrides): + fields = { + "disable_radix_cache": False, + "mamba_radix_cache_strategy": "extra_buffer", + "page_size": 64, + "speculative_algorithm": None, + "attention_backend": "aiter", + "decode_attention_backend": None, + "prefill_attention_backend": None, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_extra_buffer_at_page_64_is_accepted(): + _validate_qwen35_runtime_contract(_args()) + + +@pytest.mark.parametrize("page_size", [1, 128, 256]) +def test_extra_buffer_rejects_other_page_sizes(page_size): + with pytest.raises(ValueError, match="--page-size 64"): + _validate_qwen35_runtime_contract(_args(page_size=page_size)) + + +def test_page_size_is_free_without_the_radix_cache(): + _validate_qwen35_runtime_contract(_args(disable_radix_cache=True, page_size=128)) + + +def test_dflash_on_aiter_requires_unified_attention(): + args = _args(speculative_algorithm="DFLASH") + with pytest.raises(ValueError, match="SGLANG_USE_AITER_UNIFIED_ATTN=1"): + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) + _validate_qwen35_runtime_contract(args, unified_attn_enabled=True) + + +def test_dflash_off_aiter_does_not_need_unified_attention(): + args = _args(speculative_algorithm="DFLASH", attention_backend="triton") + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) From ef310a416995ce04784d591c770adfd0e60ed6db Mon Sep 17 00:00:00 2001 From: zovonoir Date: Wed, 21 Oct 2026 18:31:17 +0800 Subject: [PATCH 05/13] [Test][SGLang] Add a GDN prefix-reuse acceptance check Drives a running server: a prefix hit at temperature 0 must reproduce a cold run. Byte equality alone passes a broken cache, because a quarter of Qwen3.5's layers attend over the intact KV cache and short answers survive a stale GDN state, so the warm run's output logprobs must also sit within the cold-to-cold spread. The cold arm runs several times since MoE reductions make it nondeterministic at near-ties. Every case reads back `cached_tokens` and is INCONCLUSIVE without a hit. The tree is populated with the bare prefix, because extra_buffer stores state only at the last 64-token boundary a request reached. The multi-turn case builds turn 2 from turn 1's token ids so the hit extends into generated tokens (the decode-side snapshot), and stops turn 1 just past a track boundary so the recomputed tail is too short for the forget gate to wash a stale state out. On Qwen3.5-397B-A17B-FP8 TP4 against origin/main without the snapshot fix, prompt-prefix cases fail (logprob gaps 0.12-0.43, and an empty completion at ~9k tokens); with it all cases pass. Co-authored-by: Cursor --- tests/plugin/check_gdn_prefix_reuse.py | 334 +++++++++++++++++++++++++ 1 file changed, 334 insertions(+) create mode 100755 tests/plugin/check_gdn_prefix_reuse.py diff --git a/tests/plugin/check_gdn_prefix_reuse.py b/tests/plugin/check_gdn_prefix_reuse.py new file mode 100755 index 0000000000..6582e72ac0 --- /dev/null +++ b/tests/plugin/check_gdn_prefix_reuse.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +"""Acceptance check for GDN prefix caching: a cache hit must not change output. + +Not a pytest unit test -- it drives a running SGLang server, hence the `check_` +prefix so collection skips it. Run it against a server started with +`--mamba-radix-cache-strategy extra_buffer --page-size 64`. + +The failure mode this guards against is silent. A hybrid GDN model's radix node +carries both KV pages and a recurrent state snapshot; if the snapshot is stale +or was never taken, a prefix hit resumes from the wrong state and the model +keeps generating fluent-but-wrong text. Accuracy benchmarks blur that into +noise, so compare byte-for-byte instead: + + warm: [flush] -> P (populates the tree) -> P+S (hits the P prefix) + cold: [flush] -> P+S (no prefix to hit), several times + +The tree is populated with P alone because extra_buffer only stores a GDN +state at the last 64-token boundary a request reached: populating with P+S1 +puts that boundary inside S1, past where P+S2 diverges, and nothing is reused. + +Two things make a byte-for-byte verdict trustworthy, and both are enforced here +rather than left to the reader: + +* **The warm arm must actually hit.** Without a hit the two arms are the same + computation and every case passes vacuously. The warm request's + `cached_tokens` is read back and a zero is reported INCONCLUSIVE. +* **The cold arm is not byte-deterministic at temperature 0**, with or without + DFLASH: MoE reductions and speculative verify batch shapes both move + near-ties, so the first tokens can flip between e.g. entering and skipping + the thinking block. The cold arm therefore runs several times and the warm + output must equal one of them byte for byte. It is a FAIL only when every + cold run agrees and the warm run differs; anything else is INCONCLUSIVE -- + the harness failing to measure, not the cache failing. + +Usage: check_gdn_prefix_reuse.py [--port 31000] [--prefix-tokens 2048 2080] + [--cold-runs 3] + [--multiturn] +""" + +import argparse +import json +import sys +import time +import urllib.error +import urllib.request + +PORT = 31000 +MODEL = "/model/Qwen3.5-397B-A17B-FP8" + +# Constrained enough that a near-tie on formatting cannot flip it; see the +# module docstring on why free-form answers are unusable under DFLASH. +BARE = " Reply with only the numeric value: no words, no punctuation, no markdown." + + +def post(path, payload=None, port=PORT, timeout=600, retries=30): + # /flush_cache answers 400 while warmup requests are still in flight. + if path == "/flush_cache": + for _ in range(retries): + try: + return _post(path, payload, port, timeout) + except urllib.error.HTTPError as e: + if e.code != 400: + raise + time.sleep(2) + return _post(path, payload, port, timeout) + + +def _post(path, payload, port, timeout): + data = json.dumps(payload).encode() if payload is not None else b"" + req = urllib.request.Request( + f"http://127.0.0.1:{port}{path}", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as r: + body = r.read().decode() + try: + return json.loads(body) + except json.JSONDecodeError: + return body + + +class Sample: + """One temperature-0 completion: text, per-token (token id, logprob), hit size.""" + + def __init__(self, text, tokens, cached): + self.text = text + self.tokens = tokens + self.cached = cached + + +def gen(prompt, port, max_new_tokens=32, input_ids=None): + payload = { + "sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens}, + "return_logprob": True, + } + if input_ids is not None: + payload["input_ids"] = input_ids + else: + payload["text"] = prompt + out = post("/generate", payload, port=port) + meta = out["meta_info"] + tokens = [(t[1], t[0]) for t in meta.get("output_token_logprobs") or []] + return Sample(out["text"], tokens, meta.get("cached_tokens")) + + +def logprob_gap(a, b): + """Max |logprob| difference over the tokens two samples agree on. + + Compared up to the first differing token: past a near-tie flip the two runs + condition on different text and their logprobs are no longer comparable. + Returns None when not even the first token agrees. + """ + gap = None + for (tok_a, lp_a), (tok_b, lp_b) in zip(a.tokens, b.tokens): + if tok_a != tok_b: + break + gap = max(gap or 0.0, abs(lp_a - lp_b)) + return gap + + +def build_prefix(approx_tokens): + """Deterministic, non-repeating filler so the radix key is a real prefix.""" + words = [] + n = 0 + i = 0 + while n < approx_tokens: + words.append(f"Item {i} records value {(i * 7919) % 10007}.") + n += 9 + i += 1 + return " ".join(words) + + +def verdict(label, warm, colds, logprob_tol): + """PASS / FAIL / INCONCLUSIVE, printed and returned. + + Text alone is too coarse: a quarter of Qwen3.5's layers are full attention + over the intact KV cache, so a short retrieval answer often survives a + corrupted GDN state. The warm run's logprobs must also sit within the + cold-to-cold spread (floored at ``logprob_tol``) of the nearest cold run. + + ``warm.cached`` is None when the endpoint does not report the hit size + (then the server log is the only evidence that the warm run hit). + """ + if warm.cached == 0: + print(f"[INCONCLUSIVE] {label} -- warm request hit no cached tokens") + return "inconclusive" + + noise = [ + g + for i, a in enumerate(colds) + for b in colds[i + 1 :] + if (g := logprob_gap(a, b)) is not None + ] + bound = max([logprob_tol] + [2 * g for g in noise]) + gaps = [g for c in colds if (g := logprob_gap(warm, c)) is not None] + warm_gap = min(gaps) if gaps else None + texts = [c.text for c in colds] + stats = ( + f"cached_tokens={warm.cached}, warm/cold logprob gap=" + f"{'n/a' if warm_gap is None else f'{warm_gap:.4f}'}, " + f"cold/cold max={max(noise) if noise else float('nan'):.4f}, bound={bound:.4f}" + ) + + if warm_gap is not None and warm_gap > bound: + print(f"[FAIL] {label} -- logprobs drift past the cold spread ({stats})") + print(f" warm: {warm.text[:160]!r}") + print(f" cold: {texts[0][:160]!r}") + return "fail" + if warm.text in texts: + print( + f"[PASS] {label} (matches {texts.count(warm.text)}/{len(texts)} " + f"cold runs; {stats})" + ) + return "pass" + if len(set(texts)) > 1: + print(f"[INCONCLUSIVE] {label} -- cold runs disagree and warm matches none") + for i, text in enumerate(texts): + print(f" cold #{i + 1}: {text[:160]!r}") + print(f" warm : {warm.text[:160]!r}") + return "inconclusive" + print(f"[FAIL] {label} ({stats})") + print(f" warm (cache hit): {warm.text[:160]!r}") + print(f" cold (no cache) : {texts[0][:160]!r}") + return "fail" + + +def check_prompt_prefix(port, approx, cold_runs, logprob_tol): + """Reuse a prompt prefix -- the prefill-side snapshot path.""" + prefix = build_prefix(approx) + suffix = "\n\nWhat is the single largest recorded value above?" + BARE + + post("/flush_cache", {}, port=port) + gen(prefix, port, max_new_tokens=1) # populate the tree with P + warm = gen(prefix + suffix, port) # should hit the P prefix + + colds = [] + for _ in range(cold_runs): + post("/flush_cache", {}, port=port) + colds.append(gen(prefix + suffix, port)) + + return verdict(f"prompt prefix ~{approx} tokens", warm, colds, logprob_tol) + + +def check_multiturn(port, cold_runs, logprob_tol, model, track_interval=256): + """Reuse *generated* tokens as a prefix -- the decode-side snapshot path. + + The prompt-prefix checks never exercise it: they only reuse the prompt, + whose state was snapshotted during prefill. A snapshot taken while decoding + only matters once a request's own output becomes someone else's prefix, + which is what the second turn of a conversation does. + + Turn 2 is built from token ids -- turn 1's prompt ids, then the ids turn 1 + generated, then the templated follow-up -- because re-rendering the history + through the chat template re-tokenizes (and may rewrite) the assistant turn, + leaving only the prompt as a shared prefix. The warm hit must then extend + past turn 1's prompt, or the decode-side snapshot was never used. + + Turn 1 generates just past a `track_interval` boundary (SGLang's + `--mamba-track-interval`, 256 by default) so the hit lands a few tokens + before turn 1's end. GDN's forget gate washes a stale resumed state out over + a long recomputed tail, which is what makes a long tail pass vacuously. + """ + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(model, trust_remote_code=True) + prompt = build_prefix(1024) + ( + "\n\nList every item above with its value, one per line, no commentary." + ) + ask = "Which of the items you just listed has the largest value?" + BARE + + def templated(messages): + return tok.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + # Thinking off, so turn 1 spends its budget on the answer. + enable_thinking=False, + ) + + turn1_ids = tok(templated([{"role": "user", "content": prompt}]))["input_ids"] + boundary = (len(turn1_ids) // track_interval + 2) * track_interval + post("/flush_cache", {}, port=port) + turn1 = gen( + None, port, max_new_tokens=boundary + 4 - len(turn1_ids), input_ids=turn1_ids + ) + generated = [t for t, _ in turn1.tokens] + # The follow-up exactly as the template renders it after an assistant turn. + sentinel = "QZXJSENTINEL" + full = templated( + [ + {"role": "user", "content": "x"}, + {"role": "assistant", "content": sentinel}, + {"role": "user", "content": ask}, + ] + ) + tail = full[full.index(sentinel) + len(sentinel) :] + # Turn 1 may already have emitted the end-of-turn token the tail opens with. + last = tok.convert_ids_to_tokens(generated[-1]) if generated else None + if last and tail.startswith(last): + tail = tail[len(last) :] + turn2_ids = turn1_ids + generated + tok(tail, add_special_tokens=False)["input_ids"] + + warm = gen(None, port, input_ids=turn2_ids) + colds = [] + for _ in range(cold_runs): + post("/flush_cache", {}, port=port) + colds.append(gen(None, port, input_ids=turn2_ids)) + + label = ( + f"multi-turn (turn 1: {len(turn1_ids)} prompt + {len(generated)} " + f"generated tokens)" + ) + if warm.cached is not None and warm.cached <= len(turn1_ids): + print( + f"[INCONCLUSIVE] {label} -- warm hit {warm.cached} tokens, none past " + "turn 1's prompt, so no decode-side snapshot was reused" + ) + return "inconclusive" + return verdict(label, warm, colds, logprob_tol) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, default=PORT) + ap.add_argument("--model", default=MODEL, help="tokenizer for the multi-turn check") + ap.add_argument( + "--prefix-tokens", + type=int, + nargs="+", + default=[2048, 2080], + help="approx prefix lengths; include one that is not a multiple of 64", + ) + ap.add_argument( + "--multiturn", + action="store_true", + help="also reuse generated tokens as a prefix (decode-side snapshot)", + ) + ap.add_argument( + "--cold-runs", + type=int, + default=3, + help="cold repetitions the warm output is compared against", + ) + ap.add_argument( + "--logprob-tol", + type=float, + # Cold-to-cold gaps up to ~0.057 are routine on Qwen3.5-397B-A17B-FP8. + default=0.06, + help="floor of the allowed warm/cold logprob gap", + ) + args = ap.parse_args() + + results = [ + check_prompt_prefix(args.port, a, args.cold_runs, args.logprob_tol) + for a in args.prefix_tokens + ] + if args.multiturn: + results.append( + check_multiturn(args.port, args.cold_runs, args.logprob_tol, args.model) + ) + + fails = results.count("fail") + unknown = results.count("inconclusive") + print( + f"\nRESULT: {results.count('pass')} pass, {fails} fail, " + f"{unknown} inconclusive" + ) + return 1 if fails or unknown else 0 + + +if __name__ == "__main__": + sys.exit(main()) From 0b2b4fc823a4e236f93e923b2407b73cabd16364 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Thu, 22 Oct 2026 15:57:29 +0800 Subject: [PATCH 06/13] [Plugin][SGLang] One owner per mode for GDN radix tracking (Qwen3.5 + Qwen3.8) The extend snapshot (chunk-boundary `h`) and the Flash decode-graph fix each tracked both modes, and together they broke each other: * Qwen3.8 prefills on FlyDSL, which kept no chunk states for SGLang, so the first tracked extend raised "this prefill kept none". * The decode-graph helper also row-copied the finished state on extend. Where the track point is unaligned, that overwrote the chunk-boundary snapshot with the state after the whole extend, and a later hit resumed from a state that had already consumed up to 63 of the tokens it then recomputed. Now extend is tracked only by `_track_mamba_states` from the prefill kernel's chunk states, Triton or FlyDSL (published through the same pop contract), and decode only by `_track_mamba_radix_slots`, inline when eager and from the post-replay epilogue for a replayed graph. That also drops the in-graph per-layer decode track that ran without `check_freed_slots` over padded slots. --- atom/model_ops/attention_gdn.py | 10 ++- atom/model_ops/fla_ops/chunk.py | 12 ++++ .../sglang/attention_backend/attention_gdn.py | 62 +++++++------------ tests/plugin/test_sglang_gdn_radix_track.py | 48 +++++++++++--- 4 files changed, 83 insertions(+), 49 deletions(-) diff --git a/atom/model_ops/attention_gdn.py b/atom/model_ops/attention_gdn.py index 8f8bdf0f38..5b56c3269f 100644 --- a/atom/model_ops/attention_gdn.py +++ b/atom/model_ops/attention_gdn.py @@ -468,10 +468,18 @@ def forward( ssm_state, non_spec_query_start_loc, flydsl_metadata, - keep_intermediate_states=ckpt is not None, + keep_intermediate_states=( + ckpt is not None or gdn_metadata.keep_intermediate_states + ), state_indices=non_spec_state_indices_in_tensor, has_initial_state=has_initial_state, ) + if gdn_metadata.keep_intermediate_states: + from atom.model_ops.fla_ops.chunk import ( + publish_intermediate_states, + ) + + publish_intermediate_states(flydsl_h) else: # Slots and mask come per step from the metadata builder, so # this is one gather and one masked fill per layer. diff --git a/atom/model_ops/fla_ops/chunk.py b/atom/model_ops/fla_ops/chunk.py index 6edbb0dab5..fbd1ec3a25 100644 --- a/atom/model_ops/fla_ops/chunk.py +++ b/atom/model_ops/fla_ops/chunk.py @@ -388,3 +388,15 @@ def pop_last_intermediate_states(): h = ChunkGatedDeltaRuleFunction.last_h ChunkGatedDeltaRuleFunction.last_h = None return h + + +def publish_intermediate_states(h): + """Hand per-chunk states from another prefill kernel to the popper. + + The FlyDSL prefill returns its chunk states instead of leaving them here. + Publishing them keeps one contract for consumers that only see the layer + from outside (the SGLang radix cache): same logical [K, V] layout and the + same row per chunk boundary over the flattened varlen batch, whichever + kernel ran the prefill. + """ + ChunkGatedDeltaRuleFunction.last_h = h diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py index 55a76ecd2c..d1fc039605 100644 --- a/atom/plugin/sglang/attention_backend/attention_gdn.py +++ b/atom/plugin/sglang/attention_backend/attention_gdn.py @@ -184,20 +184,27 @@ def _flydsl_prefill_metadata( def _track_mamba_radix_slots(forward_batch: Any, layer_id: int) -> None: - """Mirror community ``GDNAttnBackend`` mamba track after Native GDN compute. + """Decode-side mamba track after Native GDN compute. Community bakes ``_track_mamba_state_decode`` inside the decode CUDA graph. Native GDN + FlyDSL has repeatedly faulted with that binding under radix/extra_buffer, so decode-graph capture **skips** track here and ``track_mamba_radix_after_decode_graph`` (post-replay epilogue) performs - the same snapshot with live Hybrid metadata. Eager prefill/extend still - tracks inline. Never host-sync on ``mask.any()``. + the same snapshot with live Hybrid metadata. Eager decode tracks inline. + Never host-sync on ``mask.any()``. + + Decode only: a decode step's live state IS the state at its track point. + Extend is tracked by ``SGLangGatedDeltaNet._track_mamba_states``, which + reads the chunk-boundary state an unaligned track point needs; a row copy + of the finished state here would overwrite that snapshot. ``check_freed_slots=True`` skips ``src/dst == -1`` (graph pad / tombstones); the track kernel otherwise faults on ``ptr + (-1)*stride``. """ if getattr(forward_batch, "mamba_track_mask", None) is None: return + if not forward_batch.forward_mode.is_decode_or_idle(): + return # Do not bake track into the full decode graph; epilogue owns decode track. if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): return @@ -230,21 +237,6 @@ def _track_mamba_radix_slots(forward_batch: Any, layer_id: int) -> None: if track_mask is not None and track_mask.shape[0] != bs: track_mask = track_mask[:bs] - mode = forward_batch.forward_mode - if mode.is_extend() and not mode.is_target_verify(): - # Community also writes conv at track boundaries inside causal_conv; - # Native only has the finished state. Apply final-src extend copies - # when no chunk ``h`` is required, then the same decode-style row copy. - track_extend = getattr(linear, "_track_mamba_state_extend", None) - h_src = getattr(fm, "track_ssm_h_src", None) - needs_h = h_src is not None and int(h_src.numel()) > 0 - if ( - track_extend is not None - and not needs_h - and getattr(fm, "has_mamba_track_mask", False) - ): - track_extend(forward_batch, None, layer_cache.temporal, fm) - # Prefer community coalesced all-layers track at the last mamba layer # (same as ``GDNAttnBackend.forward_decode(..., layer.layer_id)``). Always # pass check_freed_slots=True: Flash full-graph pads use PAD_SLOT_ID=-1 and @@ -289,11 +281,10 @@ def track_mamba_radix_after_decode_graph( ) -> None: """Post-replay mamba track (Native epilogue). - Community bakes track inside the graph; Native also does that via - ``_track_mamba_radix_slots``. This epilogue re-runs a coalesced track with - the *live* Hybrid metadata after replay so radix/extra_buffer still gets a - correct snapshot if the captured track-save missed a buffer. Cheap relative - to the graphed decode body. + Community bakes track inside the graph; ``_track_mamba_radix_slots`` skips + capture, so for a replayed decode graph this epilogue is the only snapshot: + it runs the coalesced track with the *live* Hybrid metadata after replay. + Cheap relative to the graphed decode body. ``attn_backend`` must be passed from ``DecodeCudaGraphRunner`` — after ``execute`` returns, SGLang's forward context is gone so @@ -378,8 +369,6 @@ def forward( and not draft_extend ): note_full_state_write(self._verify_cache_slots(forward_batch)) - if forward_batch is not None: - _track_mamba_radix_slots(forward_batch, int(self.layer_num)) return out attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) @@ -473,7 +462,11 @@ def _track_mamba_states(self, forward_batch: Any, mixed_qkv: torch.Tensor) -> No which the plugin replaces. They are made here, after the ATOM forward, through SGLang's helpers so the snapshot semantics stay defined there. Every source is final by now: the conv window is read from the pre-conv - input, the SSM state from the chunk kernel's `h` or the runtime slot. + input, the SSM state from the prefill kernel's chunk states `h` (Triton + or FlyDSL) or the runtime slot. + + Decode goes through `_track_mamba_radix_slots`, which also covers the + replayed decode graph via the post-replay epilogue. """ if ( forward_batch is None @@ -481,6 +474,11 @@ def _track_mamba_states(self, forward_batch: Any, mixed_qkv: torch.Tensor) -> No ): return mode = forward_batch.forward_mode + if mode.is_decode_or_idle(): + _track_mamba_radix_slots(forward_batch, int(self.layer_num)) + return + if not mode.is_extend(): + return attn_backend = SGLangGDNForwardContext._resolve_attn_backend(forward_batch) if attn_backend is None: return @@ -498,18 +496,6 @@ def _track_mamba_states(self, forward_batch: Any, mixed_qkv: torch.Tensor) -> No if policy is not None and policy.decode: ssm_states = ssm_states.transpose(-1, -2) - if mode.is_decode_or_idle(): - linear_backend._track_mamba_state_decode( - forward_batch, - conv_states, - ssm_states, - metadata.mamba_cache_indices, - self.layer_num, - ) - return - if not mode.is_extend(): - return - from atom.model_ops.fla_ops.chunk import pop_last_intermediate_states h = pop_last_intermediate_states() diff --git a/tests/plugin/test_sglang_gdn_radix_track.py b/tests/plugin/test_sglang_gdn_radix_track.py index 960c2984f9..3c96ec32d5 100644 --- a/tests/plugin/test_sglang_gdn_radix_track.py +++ b/tests/plugin/test_sglang_gdn_radix_track.py @@ -3,7 +3,8 @@ The plugin replaces SGLang's GDN kernels, so it has to make the track calls SGLang's own backend makes around them. These tests drive `SGLangGatedDeltaNet._track_mamba_states` against a stand-in linear backend and -check which helper receives which buffers, plus the ReplaySSM radix guard. +check which helper receives which buffers, that extend and decode each have a +single owner, plus the ReplaySSM radix guard. """ from __future__ import annotations @@ -89,19 +90,46 @@ def test_no_track_mask_is_a_no_op(backend): assert lb.calls == [] -def test_decode_tracks_through_sglang_with_layer_id(backend): - idx = torch.tensor([2, 3]) - lb = backend(SimpleNamespace(mamba_cache_indices=idx)) +def test_decode_tracks_through_the_radix_slot_helper(backend, monkeypatch): + lb = backend(SimpleNamespace(mamba_cache_indices=torch.tensor([2, 3]))) + seen = [] + monkeypatch.setattr( + attention_gdn, + "_track_mamba_radix_slots", + lambda fb, layer_id: seen.append((fb, layer_id)), + ) fb = SimpleNamespace( mamba_track_mask=torch.tensor([True, False]), forward_mode=_Mode("decode") ) _impl()._track_mamba_states(fb, torch.zeros(2, CONV_DIM)) - (call,) = lb.calls - assert call[0] == "decode" - assert call[1] is lb.layer_cache.conv[0] - assert call[2] is lb.layer_cache.temporal - assert call[3] is idx - assert call[4] == LAYER_NUM + # One owner for decode: the helper that also backs the post-replay + # epilogue, never SGLang's in-graph per-layer call. + assert seen == [(fb, LAYER_NUM)] + assert lb.calls == [] + + +def test_radix_slot_helper_leaves_extend_to_chunk_states(backend): + """A finished-state row copy on extend would overwrite the chunk-boundary + snapshot an unaligned track point needs.""" + lb = backend( + SimpleNamespace( + has_mamba_track_mask=True, + mamba_cache_indices=torch.tensor([1, 2]), + mamba_track_indices=torch.tensor([6, 7]), + track_ssm_h_src=torch.empty(0, dtype=torch.long), + ) + ) + lb.layer_cache.conv[0][1:3] = 1.0 + lb.layer_cache.temporal[1:3] = 1.0 + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True, True]), + forward_mode=_Mode("extend"), + batch_size=2, + ) + attention_gdn._track_mamba_radix_slots(fb, LAYER_NUM) + assert lb.calls == [] + assert lb.layer_cache.conv[0][6:].abs().sum() == 0 + assert lb.layer_cache.temporal[6:].abs().sum() == 0 def _extend_metadata(track_conv_indices, dst, h_src_rows): From f4149c83fbff483ac621fee83475722f8aa1d56a Mon Sep 17 00:00:00 2001 From: zovonoir Date: Thu, 22 Oct 2026 17:41:21 +0800 Subject: [PATCH 07/13] [Plugin][SGLang] Address review: PLE rows from mamba slots, fail closed on decode snapshot * `_ple_state_pool_slots` floored the PLE row count at `max_total_tokens`, the KV-cache token budget. Each Flash PLE row is 184 KB, so an explicit `--max-total-tokens` allocated one row per KV token (65536 -> 12 GB) on top of a mamba pool of a few hundred slots. Size from mamba capacity only. * The post-replay decode track epilogue logged and swallowed failures. It is the only snapshot writer for a replayed decode graph, so a failure left a track slot unwritten that later prefix hits would resume from. Let it propagate. --- atom/plugin/sglang/qwen4_exp_bridge.py | 9 +-- atom/plugin/sglang/register.py | 26 ++++----- tests/plugin/test_sglang_gdn_radix_track.py | 57 ++++++++++++++++++- .../test_sglang_qwen4_exp_graph_padding.py | 23 ++++++++ 4 files changed, 96 insertions(+), 19 deletions(-) diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index cad4d651bb..7c5d9ddde5 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1210,10 +1210,11 @@ def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: slots = max(int(_DECODE_GRAPH.max_bs), 512) args = _server_args() if args is not None: - for key in ("max_mamba_cache_size", "max_total_tokens"): - val = getattr(args, key, None) - if val: - slots = max(slots, int(val)) + # Mamba slots only: `max_total_tokens` is the KV token budget, and as + # a row floor it allocated one PLE row (184 KB on Flash) per KV token. + mamba_cache_size = getattr(args, "max_mamba_cache_size", None) + if mamba_cache_size: + slots = max(slots, int(mamba_cache_size)) running = int(getattr(args, "max_running_requests", 0) or 0) if running > 0: # extra_buffer ≈ 2x; radix retains finished track slots on top. diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py index 4b1a301061..6c590ad2e5 100644 --- a/atom/plugin/sglang/register.py +++ b/atom/plugin/sglang/register.py @@ -146,26 +146,24 @@ def capture_one_shape_with_forward_context(self, *args, **kwargs): DecodeCudaGraphRunner.capture_one_shape = capture_one_shape_with_forward_context - # Post-replay mamba track epilogue for Native GDN (community bakes track - # inside the graph; Native also bakes via ``_track_mamba_radix_slots``). - # Re-snapshot after replay so radix/extra_buffer never depends solely on - # the captured track-save binding. + # Post-replay mamba track epilogue for Native GDN. Community bakes the + # decode track inside the graph; ``_track_mamba_radix_slots`` skips capture, + # so for a replayed decode this epilogue is the only snapshot writer. A + # failure propagates: a track slot left unwritten would be offered to later + # prefix hits as if it held this request's state. if not getattr(DecodeCudaGraphRunner, "_atom_mamba_track_epilogue", False): original_execute = DecodeCudaGraphRunner.execute def execute_with_mamba_track_epilogue(self, forward_batch, *args, **kwargs): output = original_execute(self, forward_batch, *args, **kwargs) - try: - from atom.plugin.sglang.attention_backend.attention_gdn import ( - track_mamba_radix_after_decode_graph, - ) + from atom.plugin.sglang.attention_backend.attention_gdn import ( + track_mamba_radix_after_decode_graph, + ) - track_mamba_radix_after_decode_graph( - forward_batch, - attn_backend=self.model_runner.attn_backend, - ) - except Exception: - logger.exception("Native GDN post-replay mamba track failed") + track_mamba_radix_after_decode_graph( + forward_batch, + attn_backend=self.model_runner.attn_backend, + ) return output DecodeCudaGraphRunner.execute = execute_with_mamba_track_epilogue diff --git a/tests/plugin/test_sglang_gdn_radix_track.py b/tests/plugin/test_sglang_gdn_radix_track.py index 3c96ec32d5..69650aded8 100644 --- a/tests/plugin/test_sglang_gdn_radix_track.py +++ b/tests/plugin/test_sglang_gdn_radix_track.py @@ -9,7 +9,8 @@ from __future__ import annotations -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch @@ -253,3 +254,57 @@ def test_forcing_replayssm_under_radix_cache_is_rejected(monkeypatch): _server(monkeypatch, disable_radix_cache=False, speculative_algorithm="DFLASH") with pytest.raises(RuntimeError, match="radix cache"): gdn_replayssm.replayssm_enabled() + + +def _patched_decode_runner(monkeypatch): + """Install the plugin's DecodeCudaGraphRunner patch on a stand-in class.""" + from atom.plugin.sglang import register + + class _Runner: + def __init__(self): + self.model_runner = SimpleNamespace(attn_backend="backend") + + def capture_one_shape(self, *args, **kwargs): + return None + + def execute(self, forward_batch, *args, **kwargs): + return ("replayed", forward_batch) + + forward_context = ModuleType("sglang.srt.model_executor.forward_context") + forward_context.ForwardContext = object + forward_context.forward_context = None + forward_context.has_forward_context = lambda: True + runner_module = ModuleType( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner" + ) + runner_module.DecodeCudaGraphRunner = _Runner + monkeypatch.setitem(sys.modules, forward_context.__name__, forward_context) + monkeypatch.setitem(sys.modules, runner_module.__name__, runner_module) + register._install_decode_graph_forward_context_patch() + assert _Runner._atom_mamba_track_epilogue + return _Runner() + + +def test_decode_graph_epilogue_snapshots_after_replay(monkeypatch): + seen = [] + monkeypatch.setattr( + attention_gdn, + "track_mamba_radix_after_decode_graph", + lambda fb, attn_backend: seen.append((fb, attn_backend)), + ) + runner = _patched_decode_runner(monkeypatch) + assert runner.execute("fb") == ("replayed", "fb") + assert seen == [("fb", "backend")] + + +def test_decode_graph_epilogue_failure_is_not_swallowed(monkeypatch): + """The epilogue is the only snapshot writer for a replayed decode; a + failure must not return the step as if the track slot had been written.""" + + def fail(fb, attn_backend): + raise RuntimeError("track launch failed") + + monkeypatch.setattr(attention_gdn, "track_mamba_radix_after_decode_graph", fail) + runner = _patched_decode_runner(monkeypatch) + with pytest.raises(RuntimeError, match="track launch failed"): + runner.execute("fb") diff --git a/tests/plugin/test_sglang_qwen4_exp_graph_padding.py b/tests/plugin/test_sglang_qwen4_exp_graph_padding.py index 312cd9aa86..edc6ae6cda 100644 --- a/tests/plugin/test_sglang_qwen4_exp_graph_padding.py +++ b/tests/plugin/test_sglang_qwen4_exp_graph_padding.py @@ -775,6 +775,29 @@ def test_ple_state_pool_slots_uses_mamba_pool_size(): assert _ple_state_pool_slots(fb, idx) >= 64 +def test_ple_state_pool_slots_ignores_kv_token_budget(monkeypatch): + """`--max-total-tokens` sizes the KV cache, not the mamba pool; as a PLE + row floor it allocated one 184 KB row per KV token.""" + from atom.plugin.sglang import qwen4_exp_bridge as bridge + + monkeypatch.setattr( + bridge, + "_server_args", + lambda: SimpleNamespace( + max_mamba_cache_size=330, max_total_tokens=65536, max_running_requests=4 + ), + ) + fb = SimpleNamespace( + batch_size=4, + attn_backend=None, + req_to_token_pool=SimpleNamespace(size=64, req_to_token=torch.zeros(64, 8)), + token_to_kv_pool=None, + ) + idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32) + slots = _ple_state_pool_slots(fb, idx) + assert 330 <= slots < 65536 + + def test_plugin_ple_metadata_is_native(): """Plugin PLE metadata is Native Qwen4ExpPLEMetadata, including num_accepted_tokens.""" from atom.model_ops.attentions.qwen4_exp_attn import Qwen4ExpPLEMetadata From ebcbe7f3ea7f30952f79b3b3ef1b512cbc0fc36b Mon Sep 17 00:00:00 2001 From: zovonoir Date: Thu, 22 Oct 2026 21:25:33 +0800 Subject: [PATCH 08/13] [Plugin][SGLang] Size Flash PLE state from the physical mamba slot axis `_mamba_physical_slots` read axis 0 of the full conv / temporal pools, which is the layer count; the slot axis is 1 and holds `size + 1` rows (SGLang's reserved row included). With `--max-mamba-cache-size` at or above the estimate floors, PLE got `size` rows while slot `size` is valid: eager PLE clamped it onto another request's row and the decode graph indexed past the buffer. A resolved pool now sizes PLE exactly. The estimate (512 floor, `max_running_requests * 64`, capacities + 1 for the reserved row) applies only while no pool is reachable. At 64 requests / 330 slots that is 331 rows instead of 4096. --- atom/plugin/sglang/qwen4_exp_bridge.py | 112 ++++++++++-------- .../test_sglang_qwen4_exp_graph_padding.py | 41 +++++++ 2 files changed, 102 insertions(+), 51 deletions(-) diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index 7c5d9ddde5..8c737747d5 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1169,44 +1169,71 @@ def build_qsa_metadata( def _mamba_physical_slots(pool: Any) -> int: - """Rows in the physical mamba conv/temporal tensors (0 if unresolved).""" + """Slot rows of the physical mamba pool, 0 if it is not attached yet. + + The full conv / temporal tensors are ``[num_mamba_layers, size + 1, ...]``: + axis 1 counts slots, SGLang's reserved row included, so every mamba index + (working, ping-pong track and cached slots alike) is below it. + """ if pool is None: return 0 - slots = 0 - mapping = getattr(pool, "req_index_to_mamba_index_mapping", None) - if torch.is_tensor(mapping) and mapping.numel(): - slots = max(slots, int(mapping.numel())) - slots = _max_int_attrs(pool, ("size", "max_num_reqs", "mamba_size"), slots) mamba_pool = getattr(pool, "mamba_pool", None) or pool mamba_cache = getattr(mamba_pool, "mamba_cache", None) - if mamba_cache is not None: - conv = getattr(mamba_cache, "conv", None) - if conv is not None and len(conv) > 0 and hasattr(conv[0], "shape"): - slots = max(slots, int(conv[0].shape[0])) - temporal = getattr(mamba_cache, "temporal", None) - if temporal is not None and hasattr(temporal, "shape"): - slots = max(slots, int(temporal.shape[0])) + if mamba_cache is None: + return 0 + slots = 0 + conv = getattr(mamba_cache, "conv", None) + if conv is not None and len(conv) > 0 and getattr(conv[0], "ndim", 0) >= 2: + slots = max(slots, int(conv[0].shape[1])) + temporal = getattr(mamba_cache, "temporal", None) + if temporal is not None and getattr(temporal, "ndim", 0) >= 2: + slots = max(slots, int(temporal.shape[1])) return slots def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: - """Native allocates PLE conv state for the whole per-req pool, not max_bs. + """PLE state rows: one per physical mamba slot, like Native's whole-pool state. - Decode CUDA graphs bake ``conv_state``'s address. Growing after capture - frees that storage; undersizing it makes a recycled mamba slot OOB on the - next eager prefill (the conc=4 wave-2 HSA). + Decode CUDA graphs bake ``conv_state``'s address, so the buffer cannot grow + after capture, and with ``extra_buffer`` radix any slot of the pool (track + slots included) can become a working slot. Sizing to the physical pool is + exact for both. - Flash + ``extra_buffer`` radix keeps finished track slots in the same - physical pool; the capture-time floor must cover the full mamba pool, not - ``max_bs`` (otherwise unique-stress dies around the first wave of recycled - indices with a GPU fault under full decode graph). + Capture can build PLE before the pool is reachable from ``forward_batch``; + the Hybrid linear child's pool is tried first. Only when no pool resolves + does a conservative estimate apply, which may overallocate but never + undersizes a configured pool. + """ + backend = resolve_attn_backend(forward_batch) + linear = getattr(backend, "linear_attn_backend", None) or backend + pool = resolve_mamba_req_pool(forward_batch, linear) + if pool is None: + pool = getattr(linear, "req_to_token_pool", None) + slots = _mamba_physical_slots(pool) + if not slots: + try: + slots = _mamba_physical_slots(_req_to_token_pool(forward_batch)) + except Exception: + logger.debug("Flash PLE: req_to_token_pool not attached yet", exc_info=True) + if not slots: + slots = _ple_state_pool_slots_estimate(forward_batch) + # Host indices only. A GPU ``idx.max().item()`` drains verify, which is + # still queued when eager draft builds PLE state. + if ( + torch.is_tensor(idx) + and idx.numel() + and not _is_capturing() + and idx.device.type == "cpu" + ): + slots = max(slots, int(idx.max().item()) + 1) + return slots + + +def _ple_state_pool_slots_estimate(forward_batch: Any) -> int: + """Row count when the physical mamba pool is not reachable yet. - Capture often builds PLE before ``forward_batch.req_to_token_pool`` is - attached; fall back to the Hybrid linear child's pool and a high floor - (≥ allocated ``max_mamba_cache_size``, e.g. 323 here). + Pool capacities count usable slots; the pool holds one more (reserved) row. """ - # High floor: capture may run before pool attrs resolve; never bake a - # max_bs-sized PLE buffer for radix/extra_buffer. slots = max(int(_DECODE_GRAPH.max_bs), 512) args = _server_args() if args is not None: @@ -1214,7 +1241,7 @@ def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: # a row floor it allocated one PLE row (184 KB on Flash) per KV token. mamba_cache_size = getattr(args, "max_mamba_cache_size", None) if mamba_cache_size: - slots = max(slots, int(mamba_cache_size)) + slots = max(slots, int(mamba_cache_size) + 1) running = int(getattr(args, "max_running_requests", 0) or 0) if running > 0: # extra_buffer ≈ 2x; radix retains finished track slots on top. @@ -1224,33 +1251,16 @@ def _ple_state_pool_slots(forward_batch: Any, idx: torch.Tensor | None) -> int: try: from sglang.srt.runtime_context import get_schedule - sched = get_schedule() - mamba_cap = getattr(sched, "max_mamba_cache_size", None) + mamba_cap = getattr(get_schedule(), "max_mamba_cache_size", None) if mamba_cap: - slots = max(slots, int(mamba_cap)) + slots = max(slots, int(mamba_cap) + 1) except Exception: logger.debug("Flash PLE: schedule mamba capacity unavailable", exc_info=True) - backend = resolve_attn_backend(forward_batch) - linear = getattr(backend, "linear_attn_backend", None) or backend - pool = resolve_mamba_req_pool(forward_batch, linear) - if pool is None: - pool = getattr(linear, "req_to_token_pool", None) - slots = max(slots, _mamba_physical_slots(pool)) - try: - slots = max(slots, _mamba_physical_slots(_req_to_token_pool(forward_batch))) - except Exception: - logger.debug("Flash PLE: req_to_token_pool not attached yet", exc_info=True) - # Host indices only. A GPU ``idx.max().item()`` drains verify, which is - # still queued when eager draft builds PLE state. Pool capacity above - # already matches Native's whole-pool allocation. - if ( - torch.is_tensor(idx) - and idx.numel() - and not _is_capturing() - and idx.device.type == "cpu" - ): - slots = max(slots, int(idx.max().item()) + 1) - return slots + pool = getattr(forward_batch, "req_to_token_pool", None) + mapping = getattr(pool, "req_index_to_mamba_index_mapping", None) + if torch.is_tensor(mapping) and mapping.numel(): + slots = max(slots, int(mapping.numel())) + return _max_int_attrs(pool, ("size", "max_num_reqs", "mamba_size"), slots) def _max_int_attrs(obj: Any, keys: tuple[str, ...], current: int) -> int: diff --git a/tests/plugin/test_sglang_qwen4_exp_graph_padding.py b/tests/plugin/test_sglang_qwen4_exp_graph_padding.py index edc6ae6cda..210d2a0971 100644 --- a/tests/plugin/test_sglang_qwen4_exp_graph_padding.py +++ b/tests/plugin/test_sglang_qwen4_exp_graph_padding.py @@ -798,6 +798,47 @@ def test_ple_state_pool_slots_ignores_kv_token_budget(monkeypatch): assert 330 <= slots < 65536 +def _batch_with_mamba_pool(slots_axis: int, num_layers: int = 3): + """A batch whose mamba pool tensors are [num_layers, size + 1, ...].""" + cache = SimpleNamespace( + conv=[torch.zeros(num_layers, slots_axis, 2, 3)], + temporal=torch.zeros(num_layers, slots_axis, 2, 2), + ) + pool = SimpleNamespace( + size=8, + req_to_token=torch.zeros(8, 8), + mamba_pool=SimpleNamespace(mamba_cache=cache), + ) + return SimpleNamespace( + batch_size=4, attn_backend=None, req_to_token_pool=pool, token_to_kv_pool=None + ) + + +@pytest.mark.parametrize( + "max_mamba_cache_size,max_running_requests", + [(1024, 4), (330, 64)], +) +def test_ple_state_pool_slots_follow_physical_slot_axis( + monkeypatch, max_mamba_cache_size, max_running_requests +): + """Axis 1 of the full pool counts slots (reserved row included); axis 0 + counts layers. A resolved pool sizes PLE exactly: covering slot `size`, + without the running-requests floor that overallocated 12x at 64/330.""" + from atom.plugin.sglang import qwen4_exp_bridge as bridge + + monkeypatch.setattr( + bridge, + "_server_args", + lambda: SimpleNamespace( + max_mamba_cache_size=max_mamba_cache_size, + max_running_requests=max_running_requests, + ), + ) + fb = _batch_with_mamba_pool(max_mamba_cache_size + 1) + idx = torch.tensor([0, 1, 2, 3], dtype=torch.int32) + assert _ple_state_pool_slots(fb, idx) == max_mamba_cache_size + 1 + + def test_plugin_ple_metadata_is_native(): """Plugin PLE metadata is Native Qwen4ExpPLEMetadata, including num_accepted_tokens.""" from atom.model_ops.attentions.qwen4_exp_attn import Qwen4ExpPLEMetadata From d15213dd7bc2c7796d96c0420de72105a0b11cc1 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Mon, 26 Oct 2026 20:05:58 +0800 Subject: [PATCH 09/13] [Test] Align the FlyDSL decode slot test with the static-buffer contract `_align_flydsl_decode_slots` stopped materializing a contiguous copy, so the decode graph cannot bake a throwaway address. The gfx942 test still asserted the old behaviour and failed on hardware; the CPU CI job skips it, so the break was invisible. It now expects the strided view to be rejected, covers the int32 check, and asserts that narrowing a padded buffer returns a view of the same storage. --- tests/test_aiter_flydsl_gdn.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/tests/test_aiter_flydsl_gdn.py b/tests/test_aiter_flydsl_gdn.py index e3ebde4c3b..c187b749fa 100644 --- a/tests/test_aiter_flydsl_gdn.py +++ b/tests/test_aiter_flydsl_gdn.py @@ -290,7 +290,7 @@ def test_qwen_backend_binds_zero_copy_vk_state( @requires_gfx942 -def test_align_flydsl_decode_accepts_strided_and_padded_indices(): +def test_align_flydsl_decode_rejects_strided_and_views_padded_indices(): q, k, v, a, b, log, bias = inputs(4) state = torch.randn(8, 24, 128, 128, device="cuda", dtype=torch.float32) state = state.transpose(-1, -2).contiguous().transpose(-1, -2) @@ -303,10 +303,21 @@ def test_align_flydsl_decode_accepts_strided_and_padded_indices(): _align_flydsl_decode_slots, ) - reads = _align_flydsl_decode_slots(strided, 4) + # Fail closed instead of materializing a contiguous copy: the fresh + # allocation would bake a throwaway address into the decode graph. + with pytest.raises(RuntimeError, match="contiguous view"): + _align_flydsl_decode_slots(strided, 4) + with pytest.raises(RuntimeError, match="int32"): + _align_flydsl_decode_slots(padded.to(torch.int64), 4) + + # The graph-bucket case: narrowing a padded contiguous buffer keeps the + # original storage, so the slot vector the graph recorded stays live. + reads = _align_flydsl_decode_slots(wide[:4], 4) writes = _align_flydsl_decode_slots(padded, 4) + assert reads.data_ptr() == wide.data_ptr() + assert writes.data_ptr() == padded.data_ptr() assert fly.decode_supported(q, k, v, a, b, state, log, bias, reads, writes) - assert reads.tolist() == [0, 2, 4, 6] + assert reads.tolist() == [0, 1, 2, 3] assert writes.tolist() == [0, 1, 2, 3] From 9091c3e04fbd9f13b9bece7cd5fe9185fdf58060 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Mon, 26 Oct 2026 21:32:08 +0800 Subject: [PATCH 10/13] [Plugin][SGLang] Guard DFLASH on the backend target verify actually uses The Qwen3.5 startup guard required SGLANG_USE_AITER_UNIFIED_ATTN whenever any of the three backend settings mentioned Aiter. Per-mode overrides made that a false positive: `--attention-backend aiter --prefill-attention-backend triton --decode-attention-backend triton` runs neither effective backend on Aiter and was still rejected. Resolve the one backend DFLASH verifies on -- the decode backend when `speculative_attention_mode` is "decode", the prefill backend otherwise, each falling back to the default -- and require unified attention only for that one. --- atom/plugin/sglang/models/qwen3_5.py | 22 ++++++---- .../test_sglang_qwen35_runtime_contract.py | 41 +++++++++++++++++++ 2 files changed, 54 insertions(+), 9 deletions(-) diff --git a/atom/plugin/sglang/models/qwen3_5.py b/atom/plugin/sglang/models/qwen3_5.py index b013065050..2d0db1d44b 100644 --- a/atom/plugin/sglang/models/qwen3_5.py +++ b/atom/plugin/sglang/models/qwen3_5.py @@ -110,15 +110,19 @@ def _validate_qwen35_runtime_contract( ) algorithm = str(getattr(server_args, "speculative_algorithm", "") or "").upper() - backends = { - str(getattr(server_args, name, "") or "").lower() - for name in ( - "attention_backend", - "decode_attention_backend", - "prefill_attention_backend", - ) - } - if "DFLASH" not in algorithm or "aiter" not in backends: + speculative_attention_mode = str( + getattr(server_args, "speculative_attention_mode", "") or "prefill" + ).lower() + mode_backend_name = ( + "decode_attention_backend" + if speculative_attention_mode == "decode" + else "prefill_attention_backend" + ) + default_backend = getattr(server_args, "attention_backend", "") or "" + verify_backend = str( + getattr(server_args, mode_backend_name, "") or default_backend + ).lower() + if "DFLASH" not in algorithm or verify_backend != "aiter": return if unified_attn_enabled is None: diff --git a/tests/plugin/test_sglang_qwen35_runtime_contract.py b/tests/plugin/test_sglang_qwen35_runtime_contract.py index 8479d4106d..ab017813aa 100644 --- a/tests/plugin/test_sglang_qwen35_runtime_contract.py +++ b/tests/plugin/test_sglang_qwen35_runtime_contract.py @@ -44,6 +44,47 @@ def test_dflash_on_aiter_requires_unified_attention(): _validate_qwen35_runtime_contract(args, unified_attn_enabled=True) +def test_dflash_uses_per_mode_overrides_instead_of_default_backend(): + args = _args( + speculative_algorithm="DFLASH", + attention_backend="aiter", + prefill_attention_backend="triton", + decode_attention_backend="triton", + ) + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) + + +def test_dflash_on_effective_prefill_aiter_requires_unified_attention(): + args = _args( + speculative_algorithm="DFLASH", + attention_backend="triton", + prefill_attention_backend="aiter", + ) + with pytest.raises(ValueError, match="SGLANG_USE_AITER_UNIFIED_ATTN=1"): + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) + + +@pytest.mark.parametrize( + ("prefill_backend", "decode_backend", "raises"), + [("triton", "aiter", True), ("aiter", "triton", False)], +) +def test_dflash_decode_mode_uses_decode_backend( + prefill_backend, decode_backend, raises +): + args = _args( + speculative_algorithm="DFLASH", + attention_backend="triton", + speculative_attention_mode="decode", + prefill_attention_backend=prefill_backend, + decode_attention_backend=decode_backend, + ) + if raises: + with pytest.raises(ValueError, match="SGLANG_USE_AITER_UNIFIED_ATTN=1"): + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) + else: + _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) + + def test_dflash_off_aiter_does_not_need_unified_attention(): args = _args(speculative_algorithm="DFLASH", attention_backend="triton") _validate_qwen35_runtime_contract(args, unified_attn_enabled=False) From 006ba237d34c849d5be26d4bc1c95b3e9c880b27 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Mon, 26 Oct 2026 21:32:08 +0800 Subject: [PATCH 11/13] [Plugin][SGLang] Carry Flash PLE history through the mamba radix cache Flash keeps its PLE history (n-gram ids, short-conv window) in plugin-owned buffers indexed by mamba slot. SGLang's equivalent pools are slot-lifecycle siblings and are snapshotted into the track slot by its own Qwen4-Exp model, but they stay disabled under the plugin, and the plugin's buffers joined neither contract. A prefix hit therefore resumed GDN from the cached state while PLE read whatever window the working slot's previous tenant left. * Register the buffers as a slot sibling, so SGLang copies them with the slot on a hit or fork and clears them when a slot is recycled. * Snapshot them at the boundary SGLang tracks GDN at. Extend gathers the window out of this chunk (the working slot holds the state after the whole chunk by then); decode copies the working slot, which already is the boundary state. Capture is skipped and the post-replay epilogue owns the replayed decode graph, as the GDN track already does. --- atom/model_ops/attentions/qwen4_exp_attn.py | 10 + atom/model_ops/qwen4_exp/ple_layer.py | 35 +++ atom/plugin/sglang/qwen4_exp_bridge.py | 178 +++++++++++++++- atom/plugin/sglang/register.py | 4 + .../plugin/test_sglang_qwen4_exp_ple_radix.py | 200 ++++++++++++++++++ 5 files changed, 426 insertions(+), 1 deletion(-) create mode 100644 tests/plugin/test_sglang_qwen4_exp_ple_radix.py diff --git a/atom/model_ops/attentions/qwen4_exp_attn.py b/atom/model_ops/attentions/qwen4_exp_attn.py index 7196ec5509..8c724dc755 100644 --- a/atom/model_ops/attentions/qwen4_exp_attn.py +++ b/atom/model_ops/attentions/qwen4_exp_attn.py @@ -91,6 +91,16 @@ class Qwen4ExpPLEMetadata: has_initial_state: torch.Tensor # [reqs] bool conv_state: torch.Tensor # [slots, channels, state_len] short-conv pool num_accepted_tokens: torch.Tensor | None = None + # SGLang mamba radix cache (extra_buffer). The tree caches the ping-pong + # track slot, not the working slot, so PLE's own history has to be written + # there at the same boundary the GDN snapshot uses. Extend gathers the + # boundary window out of this chunk's tokens (the working slot holds the + # state after the whole chunk by then); decode is snapshotted by the + # caller, whose slot state already is the boundary state. + track_dst: torch.Tensor | None = None # [tracked] int32 destination slots + track_src: torch.Tensor | None = None # [tracked] int32 source slots (decode) + track_conv_indices: torch.Tensor | None = None # [tracked, state_len] token pos + track_ngram_indices: torch.Tensor | None = None # [tracked, ngram - 1] token pos class Qwen4ExpBackend(GDNAttentionBackend): diff --git a/atom/model_ops/qwen4_exp/ple_layer.py b/atom/model_ops/qwen4_exp/ple_layer.py index b040dd0046..704e394349 100644 --- a/atom/model_ops/qwen4_exp/ple_layer.py +++ b/atom/model_ops/qwen4_exp/ple_layer.py @@ -521,6 +521,7 @@ def forward_with_state( hidden_states, input_ids, metadata.query_start_loc, ngram_context ) normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) + self._track_radix_state(metadata, input_ids, normalized) conv_output = self._short_conv( normalized, metadata.query_start_loc, @@ -531,3 +532,37 @@ def forward_with_state( metadata.num_accepted_tokens, ) return gated_value.flatten(-2) + conv_output + + @staticmethod + def _track_radix_state( + metadata, input_ids: torch.Tensor, conv_input: torch.Tensor + ) -> None: + """Write this chunk's boundary PLE history into the radix track slots. + + Extend reads the boundary rather than the working slot, which by now + holds the state after the whole chunk: the n-gram history is the input + ids ending there, and the convolution window is the same rows of the + convolution input that the kernel would have carried over. An aligned + boundary is no special case, so extend has one path. + + A decode step advances a single token, so there the working slot is + already the boundary state and the caller asks for a row copy instead. + The caller owns both choices; see ``Qwen4ExpPLEMetadata``. + """ + dst = metadata.track_dst + if dst is None or dst.numel() == 0: + return + src = metadata.track_src + if src is not None and src.numel(): + metadata.conv_state[dst] = metadata.conv_state[src] + metadata.ngram_state[dst] = metadata.ngram_state[src] + return + conv_indices = metadata.track_conv_indices + if conv_indices is not None and conv_indices.numel(): + # [tracked, state_len, channels] -> the pool's [slot, channels, state_len] + metadata.conv_state[dst] = conv_input[conv_indices].transpose(1, 2) + ngram_indices = metadata.track_ngram_indices + if ngram_indices is not None and ngram_indices.numel(): + metadata.ngram_state[dst] = input_ids[ngram_indices].to( + metadata.ngram_state.dtype + ) diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index 8c737747d5..c428e356f3 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1346,6 +1346,157 @@ def _ensure_ple_states( return conv_state, ngram_state +class _PLESlotState: + """Flash PLE history as an SGLang slot-lifecycle sibling. + + SGLang owns the mamba slot: it copies a slot when a prefix hit or a fork + hands one request's state to another, and clears it when a slot is reused. + Its own Qwen4-Exp PLE pools ride along through this protocol, but they stay + disabled under the plugin (the pool kwargs are keyed on SGLang's config + class), so Flash's history lives in these model buffers instead. Without + this registration a resumed request keeps whatever PLE window the slot's + previous tenant left behind. + + The buffers are read back from the model on every call because they are + allocated lazily and may be replaced before capture. + """ + + def __init__(self, model: Any, ngram_fill: int): + self._model = model + self._ngram_fill = ngram_fill + + def _states(self): + for attr in ("_atom_ple_conv_state", "_atom_ple_ngram_state"): + tensor = getattr(self._model, attr, None) + if torch.is_tensor(tensor) and tensor.numel(): + yield attr, tensor + + def reset_slots(self, indices: torch.Tensor) -> None: + if indices.numel() == 0: + return + for attr, tensor in self._states(): + tensor[indices] = 0 if attr.endswith("conv_state") else self._ngram_fill + + def copy_slots(self, src_index: torch.Tensor, dst_index: torch.Tensor) -> None: + if src_index.numel() == 0: + return + for _, tensor in self._states(): + tensor[dst_index] = tensor[src_index] + + def get_cpu_slots(self, indices: torch.Tensor) -> Any: + return [t[indices].to("cpu", non_blocking=True) for _, t in self._states()] + + def load_cpu_slots(self, data: Any, indices: torch.Tensor) -> None: + if not data: + return + for (_, tensor), rows in zip(self._states(), data): + tensor[indices] = rows.to(tensor.device, non_blocking=True) + + def iter_transfer_state_entries(self): + # P/D disaggregation registers per-layer buffers; PLE is a single + # layer's side state and is not part of that transfer contract yet. + return iter(()) + + +def _register_ple_slot_state(forward_batch: Any, model: Any, ngram_fill: int) -> None: + """Attach the PLE buffers to the mamba pool's slot lifecycle, once.""" + backend = resolve_attn_backend(forward_batch) + linear = getattr(backend, "linear_attn_backend", None) or backend + pool = resolve_mamba_req_pool(forward_batch, linear) + if pool is None: + pool = getattr(linear, "req_to_token_pool", None) + mamba_pool = getattr(pool, "mamba_pool", None) if pool is not None else None + if mamba_pool is None or getattr(mamba_pool, "_atom_ple_sibling", False): + return + register = getattr(mamba_pool, "register_slot_state", None) + if not callable(register): + return + register(_PLESlotState(model, ngram_fill)) + mamba_pool._atom_ple_sibling = True + logger.info("Registered Flash PLE history as a mamba slot-state sibling") + + +def _ple_track_indices( + forward_batch: Any, + query_start_loc: torch.Tensor, + conv_state_len: int, + ngram_width: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Destination slots and boundary token positions for the PLE snapshot. + + The window ends at the same aligned boundary SGLang snapshots GDN at, so + the two halves of a cached node can never describe different positions. + """ + mask = getattr(forward_batch, "mamba_track_mask", None) + track_indices = getattr(forward_batch, "mamba_track_indices", None) + if mask is None or track_indices is None: + return None + aligned = forward_batch.mamba_track_aligned_lens() + if aligned is None: + return None + rows = mask.nonzero(as_tuple=True)[0] + if rows.numel() == 0: + return None + device = query_start_loc.device + starts = query_start_loc[:-1].to(device=device, dtype=torch.long)[rows] + starts = starts + aligned.to(device=device, dtype=torch.long)[rows] + last = (query_start_loc[-1] - 1).to(device=device, dtype=torch.long) + + def window(width: int) -> torch.Tensor: + offsets = torch.arange(width, device=device, dtype=torch.long) - width + return (starts.unsqueeze(1) + offsets).clamp_(min=0).minimum(last) + + return track_indices[rows], window(conv_state_len), window(ngram_width) + + +def _ple_decode_track_slots( + forward_batch: Any, state_indices_out: torch.Tensor +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Destination and source slots for a decode-step PLE snapshot. + + A decode step advances one token, so the working slot already holds the + state at the track boundary and a row copy is exact. + """ + mask = getattr(forward_batch, "mamba_track_mask", None) + track_indices = getattr(forward_batch, "mamba_track_indices", None) + if mask is None or track_indices is None: + return None, None + rows = mask[: state_indices_out.shape[0]].nonzero(as_tuple=True)[0] + if rows.numel() == 0: + return None, None + src = state_indices_out[rows].to(torch.long) + dst = track_indices[rows].to(device=src.device, dtype=torch.long) + # Graph pads park dead rows on PAD_SLOT_ID; they own no history to track. + live = (src >= 0) & (dst >= 0) + if not bool(live.all()): + src, dst = src[live], dst[live] + if src.numel() == 0: + return None, None + return dst, src + + +def track_ple_radix_slots(forward_batch: Any) -> None: + """Post-replay PLE snapshot for a decode graph. + + Eager decode does this inside the PLE forward. A replayed graph cannot: + the slots would be baked at capture, the same reason the GDN track skips + capture, so the epilogue re-runs it against the live metadata. + """ + ple = _DECODE_GRAPH.last_ple + if ple is None or forward_batch is None: + return + mode = getattr(forward_batch, "forward_mode", None) + if mode is None or not mode.is_decode_or_idle(): + return + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return + dst, src = _ple_decode_track_slots(forward_batch, ple.state_indices_out) + if dst is None: + return + ple.conv_state[dst] = ple.conv_state[src] + ple.ngram_state[dst] = ple.ngram_state[src] + + def _build_decode_graph_ple_metadata( atom_config: Any, forward_batch: Any, @@ -1460,7 +1611,27 @@ def build_ple_metadata( has_initial = has_initial.clone() has_initial[live_bs:] = False - return Qwen4ExpPLEMetadata( + eos = getattr(hf, "eos_token_id", 0) + _register_ple_slot_state( + forward_batch, model, int(eos[0] if isinstance(eos, (list, tuple)) else eos) + ) + track_dst = track_src = track_conv = track_ngram = None + if is_prefill: + track = _ple_track_indices( + forward_batch, + query_start_loc, + int(conv_state.shape[-1]), + int(ngram_state.shape[-1]), + ) + if track is not None: + track_dst, track_conv, track_ngram = track + elif forward_batch.forward_mode.is_decode_or_idle() and not _is_capturing(): + # Eager decode snapshots itself here, once per forward. A replayed + # decode graph cannot: the slots would be baked, so the post-replay + # epilogue calls ``track_ple_radix_slots`` with the live metadata. + track_dst, track_src = _ple_decode_track_slots(forward_batch, idx) + + _DECODE_GRAPH.last_ple = Qwen4ExpPLEMetadata( query_start_loc=query_start_loc, ngram_state=ngram_state, state_indices_in=idx_in, @@ -1468,7 +1639,12 @@ def build_ple_metadata( has_initial_state=has_initial, conv_state=conv_state, num_accepted_tokens=accepted, + track_dst=track_dst, + track_src=track_src, + track_conv_indices=track_conv, + track_ngram_indices=track_ngram, ) + return _DECODE_GRAPH.last_ple def bind_qsa_caches(model: Any, forward_batch: Any, atom_config: Any) -> None: diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py index 6c590ad2e5..fdd74f67c6 100644 --- a/atom/plugin/sglang/register.py +++ b/atom/plugin/sglang/register.py @@ -159,11 +159,15 @@ def execute_with_mamba_track_epilogue(self, forward_batch, *args, **kwargs): from atom.plugin.sglang.attention_backend.attention_gdn import ( track_mamba_radix_after_decode_graph, ) + from atom.plugin.sglang.qwen4_exp_bridge import track_ple_radix_slots track_mamba_radix_after_decode_graph( forward_batch, attn_backend=self.model_runner.attn_backend, ) + # Flash keeps its PLE history beside the GDN state; a track slot + # with only half of it resumes from a mismatched window. + track_ple_radix_slots(forward_batch) return output DecodeCudaGraphRunner.execute = execute_with_mamba_track_epilogue diff --git a/tests/plugin/test_sglang_qwen4_exp_ple_radix.py b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py new file mode 100644 index 0000000000..a8f22bf8f7 --- /dev/null +++ b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py @@ -0,0 +1,200 @@ +"""Flash PLE history under SGLang's mamba radix cache (extra_buffer). + +The tree caches the ping-pong track slot, and SGLang moves a cached slot into +the working slot on a hit. Flash keeps its PLE history in plugin-owned buffers, +so it has to join both halves of that contract: ride the slot lifecycle, and be +snapshotted at the same boundary the GDN state is. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from atom.model_ops.attentions.qwen4_exp_attn import Qwen4ExpPLEMetadata +from atom.model_ops.qwen4_exp.ple_layer import Qwen4ExpPLELayer +from atom.plugin.sglang.qwen4_exp_bridge import ( + _ple_decode_track_slots, + _ple_track_indices, + _PLESlotState, +) + +CHANNELS = 6 +STATE_LEN = 4 +NGRAM_WIDTH = 2 +EOS = 7 + + +def _model(slots: int = 8): + return SimpleNamespace( + _atom_ple_conv_state=torch.zeros(slots, CHANNELS, STATE_LEN), + _atom_ple_ngram_state=torch.full((slots, NGRAM_WIDTH), EOS, dtype=torch.int64), + ) + + +def _metadata(model, **track): + return Qwen4ExpPLEMetadata( + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + ngram_state=model._atom_ple_ngram_state, + state_indices_in=torch.tensor([0], dtype=torch.int32), + state_indices_out=torch.tensor([0], dtype=torch.int32), + has_initial_state=torch.ones(1, dtype=torch.bool), + conv_state=model._atom_ple_conv_state, + **track, + ) + + +def test_slot_state_copies_both_buffers_on_a_hit(): + """A prefix hit hands a cached slot to the working slot; PLE must follow.""" + model = _model() + model._atom_ple_conv_state[5] = 3.0 + model._atom_ple_ngram_state[5] = 11 + sibling = _PLESlotState(model, EOS) + + sibling.copy_slots(torch.tensor([5]), torch.tensor([2])) + + assert torch.equal(model._atom_ple_conv_state[2], model._atom_ple_conv_state[5]) + assert torch.equal(model._atom_ple_ngram_state[2], model._atom_ple_ngram_state[5]) + assert model._atom_ple_conv_state[1].abs().sum() == 0 + + +def test_slot_state_reset_restores_the_cold_history(): + """A recycled slot must look unused, not like its previous tenant.""" + model = _model() + model._atom_ple_conv_state[3] = 2.0 + model._atom_ple_ngram_state[3] = 11 + + _PLESlotState(model, EOS).reset_slots(torch.tensor([3])) + + assert model._atom_ple_conv_state[3].abs().sum() == 0 + assert torch.equal( + model._atom_ple_ngram_state[3], torch.full((NGRAM_WIDTH,), EOS).long() + ) + + +def test_slot_state_follows_a_regrown_buffer(): + """The buffers are allocated lazily; the sibling must not pin the old one.""" + model = _model(slots=4) + sibling = _PLESlotState(model, EOS) + model._atom_ple_conv_state = torch.zeros(16, CHANNELS, STATE_LEN) + model._atom_ple_ngram_state = torch.full((16, NGRAM_WIDTH), EOS, dtype=torch.int64) + model._atom_ple_conv_state[12] = 5.0 + + sibling.copy_slots(torch.tensor([12]), torch.tensor([15])) + + assert model._atom_ple_conv_state[15].eq(5.0).all() + + +def test_slot_state_round_trips_through_host_memory(): + model = _model() + model._atom_ple_conv_state[6] = 4.0 + model._atom_ple_ngram_state[6] = 9 + sibling = _PLESlotState(model, EOS) + indices = torch.tensor([6]) + + saved = sibling.get_cpu_slots(indices) + sibling.reset_slots(indices) + sibling.load_cpu_slots(saved, indices) + + assert model._atom_ple_conv_state[6].eq(4.0).all() + assert model._atom_ple_ngram_state[6].eq(9).all() + + +def _extend_batch(aligned, track_indices, mask): + return SimpleNamespace( + mamba_track_mask=torch.tensor(mask), + mamba_track_indices=torch.tensor(track_indices, dtype=torch.int32), + mamba_track_aligned_lens=lambda: torch.tensor(aligned), + ) + + +def test_extend_track_window_ends_at_the_aligned_boundary(): + """The snapshot is the window ending at the boundary, not at the chunk end.""" + qsl = torch.tensor([0, 10, 20], dtype=torch.int32) + fb = _extend_batch(aligned=[8, 6], track_indices=[4, 5], mask=[True, True]) + + dst, conv_idx, ngram_idx = _ple_track_indices(fb, qsl, STATE_LEN, NGRAM_WIDTH) + + assert dst.tolist() == [4, 5] + # Request 0 starts at 0 and is tracked at 8; request 1 starts at 10. + assert conv_idx.tolist() == [[4, 5, 6, 7], [12, 13, 14, 15]] + assert ngram_idx.tolist() == [[6, 7], [14, 15]] + + +def test_extend_track_skips_untracked_rows(): + qsl = torch.tensor([0, 10, 20], dtype=torch.int32) + fb = _extend_batch(aligned=[8, 6], track_indices=[4, 5], mask=[False, True]) + + dst, conv_idx, _ = _ple_track_indices(fb, qsl, STATE_LEN, NGRAM_WIDTH) + + assert dst.tolist() == [5] + assert conv_idx.tolist() == [[12, 13, 14, 15]] + + +def test_extend_track_is_absent_without_a_mask(): + qsl = torch.tensor([0, 10], dtype=torch.int32) + fb = SimpleNamespace(mamba_track_mask=None, mamba_track_indices=None) + assert _ple_track_indices(fb, qsl, STATE_LEN, NGRAM_WIDTH) is None + + +def test_extend_snapshot_reads_the_boundary_window(): + model = _model() + conv_input = torch.arange(10 * CHANNELS, dtype=torch.float32).view(10, CHANNELS) + input_ids = torch.arange(100, 110) + md = _metadata( + model, + track_dst=torch.tensor([5]), + track_conv_indices=torch.tensor([[4, 5, 6, 7]]), + track_ngram_indices=torch.tensor([[6, 7]]), + ) + + Qwen4ExpPLELayer._track_radix_state(md, input_ids, conv_input) + + # [rows, state_len, channels] stored as the pool's [slot, channels, state_len]. + torch.testing.assert_close(model._atom_ple_conv_state[5], conv_input[4:8].T) + assert model._atom_ple_ngram_state[5].tolist() == [106, 107] + assert model._atom_ple_conv_state[4].abs().sum() == 0 + + +def test_decode_snapshot_copies_the_working_slot(): + """One token per step, so the working slot already is the boundary state.""" + model = _model() + model._atom_ple_conv_state[1] = 6.0 + model._atom_ple_ngram_state[1] = 13 + md = _metadata(model, track_dst=torch.tensor([7]), track_src=torch.tensor([1])) + + Qwen4ExpPLELayer._track_radix_state(md, torch.zeros(1, dtype=torch.long), None) + + assert model._atom_ple_conv_state[7].eq(6.0).all() + assert model._atom_ple_ngram_state[7].eq(13).all() + + +def test_snapshot_is_a_no_op_without_a_destination(): + model = _model() + md = _metadata(model) + Qwen4ExpPLELayer._track_radix_state(md, torch.zeros(1, dtype=torch.long), None) + assert model._atom_ple_conv_state.abs().sum() == 0 + + +@pytest.mark.parametrize("pad_slot", [-1]) +def test_decode_track_slots_drop_graph_pad_rows(pad_slot): + """Graph pads park dead rows on PAD_SLOT_ID; indexing with it would wrap.""" + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True, True]), + mamba_track_indices=torch.tensor([6, pad_slot], dtype=torch.int32), + ) + dst, src = _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32)) + assert dst.tolist() == [6] and src.tolist() == [2] + + +def test_decode_track_slots_absent_when_nothing_is_tracked(): + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([False, False]), + mamba_track_indices=torch.tensor([6, 7], dtype=torch.int32), + ) + assert _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32)) == ( + None, + None, + ) From b78fb249e39ddcc01493a8b331474aba297d597d Mon Sep 17 00:00:00 2001 From: zovonoir Date: Mon, 26 Oct 2026 22:07:48 +0800 Subject: [PATCH 12/13] [Plugin][SGLang] Fix PLE snapshot ordering and resolve replay slots live Two defects in the previous commit, both found in review: * The snapshot ran before `_short_conv`, so a decode row copy paired a convolution window from before this token with n-gram ids from after it. It now runs once both halves have advanced. * The post-replay epilogue read the last built PLE metadata. Replay does not rebuild it, so it still described whatever forward ran eagerly before -- typically a prefill, with its own slot vector and batch size. The epilogue now takes the buffers from the model and the slots from the live linear backend's `forward_metadata`, and nothing stashes the metadata any more. Both regressions fail against the previous commit. --- atom/model_ops/qwen4_exp/ple_layer.py | 6 +- atom/plugin/sglang/qwen4_exp_bridge.py | 39 +++++++--- atom/plugin/sglang/register.py | 6 +- tests/plugin/test_sglang_gdn_radix_track.py | 16 +++- .../plugin/test_sglang_qwen4_exp_ple_radix.py | 73 +++++++++++++++++++ 5 files changed, 127 insertions(+), 13 deletions(-) diff --git a/atom/model_ops/qwen4_exp/ple_layer.py b/atom/model_ops/qwen4_exp/ple_layer.py index 704e394349..ba2d117a36 100644 --- a/atom/model_ops/qwen4_exp/ple_layer.py +++ b/atom/model_ops/qwen4_exp/ple_layer.py @@ -521,7 +521,6 @@ def forward_with_state( hidden_states, input_ids, metadata.query_start_loc, ngram_context ) normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) - self._track_radix_state(metadata, input_ids, normalized) conv_output = self._short_conv( normalized, metadata.query_start_loc, @@ -531,6 +530,11 @@ def forward_with_state( metadata.has_initial_state, metadata.num_accepted_tokens, ) + # After both halves advance: the n-gram history moved in + # ``advance_ngram_state`` above and the convolution window in + # ``_short_conv``, so a decode row copy would otherwise mix a window + # from before this token with ids from after it. + self._track_radix_state(metadata, input_ids, normalized) return gated_value.flatten(-2) + conv_output @staticmethod diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index c428e356f3..80be9086fa 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1475,26 +1475,48 @@ def _ple_decode_track_slots( return dst, src -def track_ple_radix_slots(forward_batch: Any) -> None: +def track_ple_radix_slots( + forward_batch: Any, attn_backend: Any = None, model: Any = None +) -> None: """Post-replay PLE snapshot for a decode graph. Eager decode does this inside the PLE forward. A replayed graph cannot: the slots would be baked at capture, the same reason the GDN track skips - capture, so the epilogue re-runs it against the live metadata. + capture, so the epilogue re-runs it after replay. + + Everything is resolved live. The last built PLE metadata is not usable + here: replay does not rebuild it, so it still describes whatever forward + ran eagerly before this one -- typically a prefill, with its own slot + vector and batch size. """ - ple = _DECODE_GRAPH.last_ple - if ple is None or forward_batch is None: + if forward_batch is None or model is None: return mode = getattr(forward_batch, "forward_mode", None) if mode is None or not mode.is_decode_or_idle(): return if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): return - dst, src = _ple_decode_track_slots(forward_batch, ple.state_indices_out) + conv_state = getattr(model, "_atom_ple_conv_state", None) + ngram_state = getattr(model, "_atom_ple_ngram_state", None) + if not torch.is_tensor(conv_state) or not torch.is_tensor(ngram_state): + return + from atom.plugin.sglang.attention_backend.attention_gdn import ( + SGLangGDNForwardContext, + ) + + if attn_backend is None: + attn_backend = resolve_attn_backend(forward_batch) + linear = SGLangGDNForwardContext._linear_attn_backend(attn_backend) + metadata = getattr(linear, "forward_metadata", None) + cache_indices = getattr(metadata, "mamba_cache_indices", None) + if cache_indices is None: + return + bs = int(forward_batch.batch_size) + dst, src = _ple_decode_track_slots(forward_batch, cache_indices[:bs]) if dst is None: return - ple.conv_state[dst] = ple.conv_state[src] - ple.ngram_state[dst] = ple.ngram_state[src] + conv_state[dst] = conv_state[src] + ngram_state[dst] = ngram_state[src] def _build_decode_graph_ple_metadata( @@ -1631,7 +1653,7 @@ def build_ple_metadata( # epilogue calls ``track_ple_radix_slots`` with the live metadata. track_dst, track_src = _ple_decode_track_slots(forward_batch, idx) - _DECODE_GRAPH.last_ple = Qwen4ExpPLEMetadata( + return Qwen4ExpPLEMetadata( query_start_loc=query_start_loc, ngram_state=ngram_state, state_indices_in=idx_in, @@ -1644,7 +1666,6 @@ def build_ple_metadata( track_conv_indices=track_conv, track_ngram_indices=track_ngram, ) - return _DECODE_GRAPH.last_ple def bind_qsa_caches(model: Any, forward_batch: Any, atom_config: Any) -> None: diff --git a/atom/plugin/sglang/register.py b/atom/plugin/sglang/register.py index fdd74f67c6..45ebedfecb 100644 --- a/atom/plugin/sglang/register.py +++ b/atom/plugin/sglang/register.py @@ -167,7 +167,11 @@ def execute_with_mamba_track_epilogue(self, forward_batch, *args, **kwargs): ) # Flash keeps its PLE history beside the GDN state; a track slot # with only half of it resumes from a mismatched window. - track_ple_radix_slots(forward_batch) + track_ple_radix_slots( + forward_batch, + attn_backend=self.model_runner.attn_backend, + model=self.model_runner.model, + ) return output DecodeCudaGraphRunner.execute = execute_with_mamba_track_epilogue diff --git a/tests/plugin/test_sglang_gdn_radix_track.py b/tests/plugin/test_sglang_gdn_radix_track.py index 69650aded8..0f59718e41 100644 --- a/tests/plugin/test_sglang_gdn_radix_track.py +++ b/tests/plugin/test_sglang_gdn_radix_track.py @@ -262,7 +262,7 @@ def _patched_decode_runner(monkeypatch): class _Runner: def __init__(self): - self.model_runner = SimpleNamespace(attn_backend="backend") + self.model_runner = SimpleNamespace(attn_backend="backend", model="model") def capture_one_shape(self, *args, **kwargs): return None @@ -286,15 +286,27 @@ def execute(self, forward_batch, *args, **kwargs): def test_decode_graph_epilogue_snapshots_after_replay(monkeypatch): - seen = [] + """Both halves of a slot's state are snapshotted after replay: the GDN + state and Flash's PLE history, which lives in its own buffers.""" + from atom.plugin.sglang import qwen4_exp_bridge + + seen, ple_seen = [], [] monkeypatch.setattr( attention_gdn, "track_mamba_radix_after_decode_graph", lambda fb, attn_backend: seen.append((fb, attn_backend)), ) + monkeypatch.setattr( + qwen4_exp_bridge, + "track_ple_radix_slots", + lambda fb, attn_backend=None, model=None: ple_seen.append( + (fb, attn_backend, model) + ), + ) runner = _patched_decode_runner(monkeypatch) assert runner.execute("fb") == ("replayed", "fb") assert seen == [("fb", "backend")] + assert ple_seen == [("fb", "backend", "model")] def test_decode_graph_epilogue_failure_is_not_swallowed(monkeypatch): diff --git a/tests/plugin/test_sglang_qwen4_exp_ple_radix.py b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py index a8f22bf8f7..f6ac8d83af 100644 --- a/tests/plugin/test_sglang_qwen4_exp_ple_radix.py +++ b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py @@ -189,6 +189,79 @@ def test_decode_track_slots_drop_graph_pad_rows(pad_slot): assert dst.tolist() == [6] and src.tolist() == [2] +def _layer_stub(model, advance_conv_by): + """Minimal stand-in for the parts of the PLE layer `forward_with_state` uses. + + `_short_conv` advances the working slot the way the real kernel does, so a + snapshot taken before it would capture the previous token's window. + """ + return SimpleNamespace( + ple_embedding=SimpleNamespace(eos_token_id=EOS, ngram_size=NGRAM_WIDTH + 1), + norm_conv=None, + gated_memory=lambda *a, **k: torch.zeros(1, 1, CHANNELS), + _apply_norm=lambda norm, x: x, + _short_conv=lambda inputs, qsl, state, idx_in, idx_out, init, accepted: ( + state.__setitem__(idx_out.long(), advance_conv_by), + torch.zeros(1, CHANNELS), + )[1], + _track_radix_state=Qwen4ExpPLELayer._track_radix_state, + ) + + +def test_decode_snapshot_runs_after_both_halves_advance(monkeypatch): + """The window and the n-gram ids must describe the same token position.""" + from atom.model_ops.qwen4_exp import ple_layer as layer_mod + + model = _model() + advanced_ngram = 21 + + def fake_advance(input_ids, qsl, state, idx_in, idx_out, *a, **k): + state[idx_out.long()] = advanced_ngram + return torch.zeros(1, NGRAM_WIDTH) + + monkeypatch.setattr(layer_mod, "advance_ngram_state", fake_advance) + md = _metadata(model, track_dst=torch.tensor([7]), track_src=torch.tensor([0])) + + Qwen4ExpPLELayer.forward_with_state( + _layer_stub(model, advance_conv_by=5.0), + torch.zeros(1, 1, CHANNELS), + torch.zeros(1, dtype=torch.long), + md, + ) + + assert model._atom_ple_conv_state[7].eq(5.0).all() + assert model._atom_ple_ngram_state[7].eq(advanced_ngram).all() + + +def test_graph_epilogue_uses_live_slots_not_the_last_built_metadata(): + """Replay does not rebuild PLE metadata, so the stale one would name the + slots of whatever forward ran eagerly before -- usually a prefill.""" + from atom.plugin.sglang import qwen4_exp_bridge as bridge + + model = _model() + model._atom_ple_conv_state[3] = 8.0 + model._atom_ple_ngram_state[3] = 17 + live = SimpleNamespace( + linear_attn_backend=SimpleNamespace( + forward_metadata=SimpleNamespace( + mamba_cache_indices=torch.tensor([3, 4], dtype=torch.int32) + ) + ) + ) + fb = SimpleNamespace( + batch_size=2, + forward_mode=SimpleNamespace(is_decode_or_idle=lambda: True), + mamba_track_mask=torch.tensor([True, False]), + mamba_track_indices=torch.tensor([6, 7], dtype=torch.int32), + ) + + bridge.track_ple_radix_slots(fb, attn_backend=live, model=model) + + assert model._atom_ple_conv_state[6].eq(8.0).all() + assert model._atom_ple_ngram_state[6].eq(17).all() + assert model._atom_ple_conv_state[7].abs().sum() == 0 + + def test_decode_track_slots_absent_when_nothing_is_tracked(): fb = SimpleNamespace( mamba_track_mask=torch.tensor([False, False]), From 251726c81de9028a9a69bb1b75707a2fe00a0111 Mon Sep 17 00:00:00 2001 From: zovonoir Date: Mon, 26 Oct 2026 22:54:50 +0800 Subject: [PATCH 13/13] [Plugin][SGLang] Make the decode PLE snapshot reach the buffers, and drop its host syncs Three defects from review, all in the decode half: * The replay epilogue read the PLE buffers off the SGLang wrapper, but `build_ple_metadata` is called with the native model, which the wrapper keeps in `__dict__` rather than as a child module. The lookup returned None, so a replayed decode silently snapshotted nothing. * An uncaptured decode batch below the graph ceiling still builds its PLE metadata through the graph path, yet runs eagerly with no epilogue behind it. It carried no track destination, so GDN was tracked and PLE was not. That metadata now gets the decode track fields whenever it is not capture. * Selecting the tracked rows used `nonzero()` and a `.all()` check, two host synchronisations per decode step. The snapshot is now full-batch and branch-free: untracked rows and graph pads are routed to a scratch row, one past the pool, that nothing reads. --- atom/plugin/sglang/qwen4_exp_bridge.py | 82 ++++++++++++++----- .../plugin/test_sglang_qwen4_exp_ple_radix.py | 67 +++++++++++---- 2 files changed, 113 insertions(+), 36 deletions(-) diff --git a/atom/plugin/sglang/qwen4_exp_bridge.py b/atom/plugin/sglang/qwen4_exp_bridge.py index 80be9086fa..a7ec585ee5 100644 --- a/atom/plugin/sglang/qwen4_exp_bridge.py +++ b/atom/plugin/sglang/qwen4_exp_bridge.py @@ -1312,7 +1312,9 @@ def _ensure_ple_states( ) -> tuple[torch.Tensor, torch.Tensor]: hf = _hf_text_config(atom_config) ngram = max(int(getattr(hf, "ngram_size", 3)), 1) - rows = max(num_slots, 1) + # One row past the pool: the decode snapshot routes untracked rows and + # graph pads here so it can copy the whole batch without a host sync. + rows = max(num_slots, 1) + 1 existing = getattr(model, "_atom_ple_conv_state", None) conv_state = _grow_once_tensor( model, @@ -1449,30 +1451,53 @@ def window(width: int) -> torch.Tensor: return track_indices[rows], window(conv_state_len), window(ngram_width) +def _ple_history_buffers( + model: Any, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """The PLE history buffers, wherever they were allocated. + + ``build_ple_metadata`` is called with the native language model, so that + is what owns them. The SGLang wrapper holds it in ``__dict__`` rather than + as a child module, so attribute lookup does not reach through: a caller + holding the wrapper (the decode-graph epilogue gets it from + ``model_runner``) has to unwrap it explicitly. + """ + for candidate in (model, getattr(model, "_atom_lm", None)): + conv = getattr(candidate, "_atom_ple_conv_state", None) + ngram = getattr(candidate, "_atom_ple_ngram_state", None) + if torch.is_tensor(conv) and torch.is_tensor(ngram): + return conv, ngram + return None, None + + +def _ple_scratch_row(conv_state: torch.Tensor) -> int: + """Discard row of the PLE buffers; see ``_ensure_ple_states``.""" + return int(conv_state.shape[0]) - 1 + + def _ple_decode_track_slots( - forward_batch: Any, state_indices_out: torch.Tensor + forward_batch: Any, state_indices_out: torch.Tensor, scratch_row: int ) -> tuple[torch.Tensor | None, torch.Tensor | None]: """Destination and source slots for a decode-step PLE snapshot. A decode step advances one token, so the working slot already holds the state at the track boundary and a row copy is exact. + + Full-batch and free of host synchronisation, which a decode loop cannot + afford: selecting the tracked rows would need ``nonzero()``, and checking + for graph pads would need another. Untracked rows and pads (PAD_SLOT_ID) + are routed to the scratch row instead, so the copy runs unconditionally + and those writes land somewhere nothing reads. """ mask = getattr(forward_batch, "mamba_track_mask", None) track_indices = getattr(forward_batch, "mamba_track_indices", None) if mask is None or track_indices is None: return None, None - rows = mask[: state_indices_out.shape[0]].nonzero(as_tuple=True)[0] - if rows.numel() == 0: - return None, None - src = state_indices_out[rows].to(torch.long) - dst = track_indices[rows].to(device=src.device, dtype=torch.long) - # Graph pads park dead rows on PAD_SLOT_ID; they own no history to track. - live = (src >= 0) & (dst >= 0) - if not bool(live.all()): - src, dst = src[live], dst[live] - if src.numel() == 0: - return None, None - return dst, src + bs = int(state_indices_out.shape[0]) + src = state_indices_out.to(torch.long) + dst = track_indices[:bs].to(device=src.device, dtype=torch.long) + keep = mask[:bs].to(device=src.device) & (src >= 0) & (dst >= 0) + return torch.where(keep, dst, dst.new_full((), scratch_row)), src.clamp(min=0) def track_ple_radix_slots( @@ -1496,9 +1521,8 @@ def track_ple_radix_slots( return if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): return - conv_state = getattr(model, "_atom_ple_conv_state", None) - ngram_state = getattr(model, "_atom_ple_ngram_state", None) - if not torch.is_tensor(conv_state) or not torch.is_tensor(ngram_state): + conv_state, ngram_state = _ple_history_buffers(model) + if conv_state is None: return from atom.plugin.sglang.attention_backend.attention_gdn import ( SGLangGDNForwardContext, @@ -1512,11 +1536,19 @@ def track_ple_radix_slots( if cache_indices is None: return bs = int(forward_batch.batch_size) - dst, src = _ple_decode_track_slots(forward_batch, cache_indices[:bs]) + dst, src = _ple_decode_track_slots( + forward_batch, cache_indices[:bs], _ple_scratch_row(conv_state) + ) if dst is None: return conv_state[dst] = conv_state[src] ngram_state[dst] = ngram_state[src] + if not getattr(track_ple_radix_slots, "_logged", False): + # Proof the path is reachable at all: it silently did nothing for a + # whole round of testing when the buffer lookup landed on the SGLang + # wrapper instead of the native model it allocates them on. + logger.info("Flash PLE post-replay snapshot active (rows=%s)", int(bs)) + track_ple_radix_slots._logged = True # type: ignore[attr-defined] def _build_decode_graph_ple_metadata( @@ -1551,6 +1583,13 @@ def _build_decode_graph_ple_metadata( batch_size=bs, num_accepted_tokens=accepted, ) + if not _is_capturing() and forward_batch.forward_mode.is_decode_or_idle(): + # This metadata is also what an *uncaptured* decode batch gets, which + # runs eagerly and so has no replay epilogue behind it. Capture itself + # must stay untracked: the slots would be baked into the graph. + ple.track_dst, ple.track_src = _ple_decode_track_slots( + forward_batch, ple.state_indices_out, _ple_scratch_row(conv_state) + ) if not getattr(_build_decode_graph_ple_metadata, "_logged", False): logger.info( "Flash decode graph PLE aliases GDN mamba_cache_indices " @@ -1620,7 +1659,8 @@ def build_ple_metadata( num_slots = _ple_state_pool_slots(forward_batch, idx) conv_state, ngram_state = _ensure_ple_states(model, atom_config, num_slots) - last_slot = max(int(conv_state.shape[0]) - 1, 0) + # The last row is the snapshot's discard row, not a slot a request may use. + last_slot = max(int(conv_state.shape[0]) - 2, 0) if last_slot >= 0: idx = torch.where(idx < 0, idx, idx.clamp(max=last_slot)) idx_in = torch.where(idx_in < 0, idx_in, idx_in.clamp(max=last_slot)) @@ -1651,7 +1691,9 @@ def build_ple_metadata( # Eager decode snapshots itself here, once per forward. A replayed # decode graph cannot: the slots would be baked, so the post-replay # epilogue calls ``track_ple_radix_slots`` with the live metadata. - track_dst, track_src = _ple_decode_track_slots(forward_batch, idx) + track_dst, track_src = _ple_decode_track_slots( + forward_batch, idx, _ple_scratch_row(conv_state) + ) return Qwen4ExpPLEMetadata( query_start_loc=query_start_loc, diff --git a/tests/plugin/test_sglang_qwen4_exp_ple_radix.py b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py index f6ac8d83af..ba3c4c65ce 100644 --- a/tests/plugin/test_sglang_qwen4_exp_ple_radix.py +++ b/tests/plugin/test_sglang_qwen4_exp_ple_radix.py @@ -10,7 +10,6 @@ from types import SimpleNamespace -import pytest import torch from atom.model_ops.attentions.qwen4_exp_attn import Qwen4ExpPLEMetadata @@ -178,17 +177,6 @@ def test_snapshot_is_a_no_op_without_a_destination(): assert model._atom_ple_conv_state.abs().sum() == 0 -@pytest.mark.parametrize("pad_slot", [-1]) -def test_decode_track_slots_drop_graph_pad_rows(pad_slot): - """Graph pads park dead rows on PAD_SLOT_ID; indexing with it would wrap.""" - fb = SimpleNamespace( - mamba_track_mask=torch.tensor([True, True]), - mamba_track_indices=torch.tensor([6, pad_slot], dtype=torch.int32), - ) - dst, src = _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32)) - assert dst.tolist() == [6] and src.tolist() == [2] - - def _layer_stub(model, advance_conv_by): """Minimal stand-in for the parts of the PLE layer `forward_with_state` uses. @@ -262,12 +250,59 @@ def test_graph_epilogue_uses_live_slots_not_the_last_built_metadata(): assert model._atom_ple_conv_state[7].abs().sum() == 0 -def test_decode_track_slots_absent_when_nothing_is_tracked(): +def test_decode_track_slots_route_untracked_and_pad_rows_to_scratch(): + """No host sync in a decode loop: instead of selecting the tracked rows, + every row is copied and the rest are parked on the scratch row.""" fb = SimpleNamespace( - mamba_track_mask=torch.tensor([False, False]), - mamba_track_indices=torch.tensor([6, 7], dtype=torch.int32), + mamba_track_mask=torch.tensor([True, False, True]), + mamba_track_indices=torch.tensor([6, 7, -1], dtype=torch.int32), ) - assert _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32)) == ( + scratch = 9 + dst, src = _ple_decode_track_slots( + fb, torch.tensor([2, 3, 4], dtype=torch.int32), scratch + ) + # Row 0 is tracked; row 1 is not; row 2 is a graph pad (PAD_SLOT_ID). + assert dst.tolist() == [6, scratch, scratch] + assert src.tolist() == [2, 3, 4] + + +def test_decode_track_slots_absent_without_track_metadata(): + fb = SimpleNamespace(mamba_track_mask=None, mamba_track_indices=None) + assert _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32), 9) == ( None, None, ) + + +def test_decode_track_slots_do_not_synchronize(monkeypatch): + """`nonzero()` and `.all()` both drain the stream; a decode step cannot.""" + calls = [] + for name in ("nonzero", "item", "tolist"): + original = getattr(torch.Tensor, name) + monkeypatch.setattr( + torch.Tensor, + name, + lambda self, *a, _n=name, _o=original, **k: ( + calls.append(_n), + _o(self, *a, **k), + )[1], + ) + fb = SimpleNamespace( + mamba_track_mask=torch.tensor([True, False]), + mamba_track_indices=torch.tensor([6, 7], dtype=torch.int32), + ) + _ple_decode_track_slots(fb, torch.tensor([2, 3], dtype=torch.int32), 9) + assert calls == [] + + +def test_history_buffers_unwrap_the_native_model(): + """The metadata builder is called with the native model, so that is where + the buffers live; the SGLang wrapper keeps it out of the module tree.""" + from atom.plugin.sglang.qwen4_exp_bridge import _ple_history_buffers + + native = _model() + wrapper = SimpleNamespace(_atom_lm=native) + conv, ngram = _ple_history_buffers(wrapper) + assert conv is native._atom_ple_conv_state + assert ngram is native._atom_ple_ngram_state + assert _ple_history_buffers(SimpleNamespace()) == (None, None)