[Plugin][SGLang] Carry Flash PLE history through the mamba radix cache; guard DFLASH on the resolved backend

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.
@@ -504,7 +512,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_(

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 ---------------------------------------------------------

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):

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

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
+            )