perf(qwen4): add opt-in layer-major long prefill

README.md

@@ -82,4 +82,7 @@ learned the design and reused code from the following projects:
 
 ## License
 
+For the optional Qwen3.8 pinned-PLE long-prefill execution path, see
+[layer-major prefill](docs/layer_major_prefill.md).
+
 [Apache License 2.0](https://github.com/FlashML-org/FreeToken/blob/main/LICENSE).

docs/layer_major_prefill.md

@@ -0,0 +1,56 @@
+# Opt-in layer-major prefill
+
+`--prefill-layer-major` changes the loop order of a long Qwen3.8 text prefill:
+each decoder layer processes every internal token chunk before the next layer
+starts. GPU offload can then reuse a streamed expert layer across those chunks.
+The flag is disabled by default.
+
+## Supported configuration
+
+The initial implementation requires Qwen3.8/Qwen4Exp, TP=1, NVFP4 experts,
+`--moe-strategy offload`, and **`--ple-backend pinned`**. A single pending text
+request is eligible when its new-token length exceeds `--max-prefill-length`
+and is at most 32,768. The existing prefill length remains the internal chunk
+size. This limit bounds the activations retained between layers; it is not a
+promise that every supported configuration fits every GPU.
+
+For an existing local checkpoint, add these options to the otherwise identical
+serve command used for a comparison:
+
+```sh
+--moe-strategy offload --ple-backend pinned \
+--quant-backend moe.nvfp4=triton --max-prefill-length 4096 \
+--prefill-layer-major
+```
+
+The default disk PLE backend is not supported by this execution path: its host
+staging is sized for a token chunk. Disk or other PLE modes, multimodal inputs,
+multiple pending requests, TP>1, other models or expert formats, short prefills,
+and prefills above the limit retain ordinary chunked execution.
+
+## State and numerical behavior
+
+The outer request owns checkpoint planning. Internal chunks keep absolute
+cached positions and project that plan without choosing new checkpoint slots.
+Conv, recurrent and PLE state advance in chunk order within their layer. The
+outer batch context is restored after normal completion and exceptions.
+
+Numerical comparisons must fix the internal chunk structure, model, quantized
+kernel and sampling parameters. This flag does not promise bitwise equivalence
+between different chunk structures, unsupported runtimes, or concurrent runs.
+An eligible long prefill is one scheduling turn; this implementation does not
+interleave other requests' decoding between its internal chunks.
+
+## Measuring the change
+
+Compare `main` and the candidate with the same pinned-PLE configuration,
+checkpoint, prompt bytes, 4K chunk size and output length. Record first-token
+latency with its exact streaming definition, expert H2D traffic, peak VRAM,
+generated token IDs and finish reason. Keep any NVFP4 kernel patch identical
+on both sides; kernel speedups are not layer-major speedups.
+
+This implementation is informed by [Layered Prefill](https://arxiv.org/abs/2510.08055).
+The repeated expert-streaming problem is also described in
+[#579](https://github.com/FlashML-org/FreeToken/issues/579).
+[#578](https://github.com/FlashML-org/FreeToken/pull/578) is a separate kernel
+optimization and is not included here.

python/freetoken/attention/linear.py

@@ -44,7 +44,9 @@ class FLAMetadata:
     track_boundary_row: torch.Tensor | None = None  # [nt] int64 forward-local row of the track boundary; states with their own left context (qwen4_exp PLE) derive their windows from it
 
 
-def build_fla_metadata(batch: "Batch", device: torch.device) -> FLAMetadata:
+def build_fla_metadata(
+    batch: "Batch", device: torch.device, *, include_track: bool = True
+) -> FLAMetadata:
     """Build the per-forward GDN metadata. Uses pinned host staging + non_blocking H2D
     (the input_ids/attn-metadata pattern), so the copies overlap the forward instead of
     stalling it.
@@ -78,7 +80,11 @@ def gdn_slot(r):
     fresh = [gdn_slot(r) for r in reqs if r.cached_len == 0]
     fresh_host = torch.tensor(fresh, dtype=torch.int64, **pin) if fresh else None
 
-    track = _build_track_metadata(reqs, cu_host, device, pin)
+    track = (
+        _build_track_metadata(reqs, cu_host, device, pin)
+        if include_track
+        else dict(track_dst=None, track_h_row=None, track_conv_src=None, track_boundary_row=None)
+    )
 
     return FLAMetadata(
         cu_seqlens=cu_host.to(device, non_blocking=True),

python/freetoken/engine/engine.py

@@ -1061,8 +1061,12 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         if batch.mm_gather_plan:
             self._run_mm_encoder(batch)
         use_graph = self.graph_runner.can_use_cuda_graph(batch)
-        with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph):
-            logits = self.graph_runner.replay(batch) if use_graph else self.model.forward()
+        if getattr(batch, "layer_major", False) and batch.is_prefill:
+            with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, False):
+                logits = self._forward_layer_major(batch)
+        else:
+            with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph):
+                logits = self.graph_runner.replay(batch) if use_graph else self.model.forward()
         if self.cpu_moe_executor is not None:
             # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator
             # -> stale expert outputs) as a loud error instead of silent corruption.
@@ -1078,6 +1082,100 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
         copy_done_event.record(self.stream)
         return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
 
+    def _build_layer_major_chunks(self, batch: Batch) -> list:
+        """Split one admitted long prefill into ordinary eager-forward shaped chunks."""
+        import types
+
+        from freetoken.attention.linear import build_fla_metadata
+        from freetoken.kernel.fla.chunk import CHUNK_SIZE
+
+        req = batch.reqs[0]
+        origin = req.cached_len
+        outer_fla = getattr(batch, "fla_metadata", None)
+        if outer_fla is None:
+            outer_fla = build_fla_metadata(batch, self.device)
+            batch.fla_metadata = outer_fla
+        track = None
+        if outer_fla.track_dst is not None:
+            assert outer_fla.track_dst.numel() == 1
+            boundary = origin + int(outer_fla.track_boundary_row[0])
+            assert origin < boundary < origin + req.extend_len
+            assert req.mamba_last_track_seqlen == boundary
+            assert req.linear_slot_idx == int(outer_fla.cache_indices[0])
+            track = {
+                "boundary": boundary,
+                "src": int(outer_fla.cache_indices[0]),
+                "dst": int(outer_fla.track_dst[0]),
+                "h_row": int(outer_fla.track_h_row[0]),
+                "conv_src": outer_fla.track_conv_src[0:1],
+                "boundary_row": int(outer_fla.track_boundary_row[0]),
+            }
+        carry = (
+            "table_idx", "uid", "linear_slot_idx", "max_device_len", "mamba_ping_pong",
+            "mamba_next_track_idx", "mamba_restore_src", "swa_evicted_seqlen",
+            "sampling_params", "mrope_delta",
+        )
+        chunks, lo, projections = [], 0, 0
+        while lo < req.extend_len:
+            hi = min(lo + self.config.max_extend_tokens, req.extend_len)
+            if track is not None:
+                boundary = track["boundary"]
+                if origin + lo < boundary < origin + hi and (boundary - origin - lo) % CHUNK_SIZE:
+                    hi = boundary - origin
+            child_req = types.SimpleNamespace(**{key: getattr(req, key, None) for key in carry})
+            child_req.cached_len = origin + lo
+            child_req.device_len = origin + hi
+            child_req.extend_len = hi - lo
+            child_req.mm_items = None
+            child_req.mrope_positions_full = None
+            child = types.SimpleNamespace(
+                phase="prefill",
+                input_ids=batch.input_ids[lo:hi],
+                positions=batch.positions[lo:hi],
+                out_loc=batch.out_loc[lo:hi],
+                mrope_positions=(batch.mrope_positions[:, lo:hi] if getattr(batch, "mrope_positions", None) is not None else None),
+                mm_embeds=None,
+                mm_rows=None,
+                mm_block_ends=None,
+                linear_table_idx=getattr(batch, "linear_table_idx", None),
+                reqs=[child_req],
+                padded_reqs=[child_req],
+            )
+            child.is_decode = False
+            child.is_prefill = True
+            child.size = 1
+            child.get_attn_positions = lambda child=child: (
+                child.mrope_positions if child.mrope_positions is not None else child.positions
+            )
+            child.fla_metadata = build_fla_metadata(child, self.device, include_track=False)
+            child.track_initial_src = child.track_initial_dst = None
+            if track is not None:
+                boundary = track["boundary"]
+                if boundary == origin + lo:
+                    child.track_initial_src, child.track_initial_dst = track["src"], track["dst"]
+                    projections += 1
+                elif origin + lo < boundary < origin + hi:
+                    local_boundary = boundary - origin - lo
+                    assert local_boundary % CHUNK_SIZE == 0
+                    child.fla_metadata.track_dst = torch.tensor([track["dst"]], dtype=torch.int64, device=self.device)
+                    child.fla_metadata.track_h_row = torch.tensor([track["h_row"] - lo // CHUNK_SIZE], dtype=torch.int64, device=self.device)
+                    child.fla_metadata.track_conv_src = (track["conv_src"] - lo).to(self.device)
+                    child.fla_metadata.track_boundary_row = torch.tensor([track["boundary_row"] - lo], dtype=torch.int64, device=self.device)
+                    projections += 1
+            self.attn_backend.prepare_metadata(child)
+            child.skip_begin_prefill = True
+            chunks.append(child)
+            lo = hi
+        if track is not None:
+            assert projections == 1, f"layer-major track projected {projections} times"
+        return chunks
+
+    def _forward_layer_major(self, batch: Batch) -> torch.Tensor:
+        chunks = self._build_layer_major_chunks(batch)
+        if self.moe_offload_cache is not None:
+            self.moe_offload_cache.begin_prefill()
+        return self.model.forward_layer_major(chunks)
+
     @torch.inference_mode()
     def _warmup_prefill(self) -> None:
         """Compile the Triton prefill path before the first real request.

python/freetoken/layers/moe.py

@@ -394,7 +394,7 @@ def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, .
         bank registration order; buffer position == expert id, so routing ids pass
         through unmapped). The caller runs ``release_prefill_layer`` after its GEMMs.
         """
-        if self.layer_id == 0:
+        if self.layer_id == 0 and not getattr(get_global_ctx().batch, "skip_begin_prefill", False):
             cache.begin_prefill()
         cache.prefetch_prefill_layer(self.layer_id)
         cache.prefetch_prefill_layer(self.layer_id + 1)

python/freetoken/models/qwen4_exp/model.py

@@ -35,6 +35,23 @@
     from freetoken.models.config import ModelConfig
 
 
+def _copy_layer_major_track_state(pool, layer, batch) -> None:
+    """Seed a child at the one outer prefix-checkpoint boundary."""
+    src = getattr(batch, "track_initial_src", None)
+    if src is None:
+        return
+    dst = batch.track_initial_dst
+    if layer._is_linear:
+        index = pool.local_index(layer.linear_attn.layer_id)
+        pool.recurrent_states[index, dst].copy_(pool.recurrent_states[index, src])
+        pool.conv_states[index, dst].copy_(pool.conv_states[index, src])
+    if layer.ple is not None:
+        from .config import PLE_CONV_STATE, PLE_NGRAM_STATE
+
+        pool.slot_state(PLE_CONV_STATE, layer.ple.layer_id)[dst].copy_(pool.slot_state(PLE_CONV_STATE, layer.ple.layer_id)[src])
+        pool.slot_state(PLE_NGRAM_STATE)[dst].copy_(batch.ple_meta.ngram_context[0])
+
+
 def build_linear_mixer(config: ModelConfig, layer_id: int, prefix: str) -> BaseOP:
     """GDN mixer of a linear_attention layer (Qwen3.5's GDN with a configurable output gate)."""
     from .gdn import Qwen4ExpGatedDeltaNet
@@ -126,6 +143,30 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor:
             commit_ngram_context(meta, getattr(batch, "fla_metadata", None))
         return self.hyper_connection_mixer.mix(hidden)[0]
 
+    def forward_layer_major(self, chunk_batches: List[Batch]) -> torch.Tensor:
+        """Run decoder layers outermost, retaining the normal chunk order within each layer."""
+        from .ple import build_ple_metadata, commit_ngram_context
+
+        ctx = get_global_ctx()
+        full_batch = ctx.batch
+        hidden = []
+        try:
+            for child in chunk_batches:
+                ctx._batch = child
+                hidden.append(embed_input_ids(self.embed_tokens, child.input_ids, child).repeat(1, self.hc_count))
+            for layer in self.layers.op_list:
+                for index, child in enumerate(chunk_batches):
+                    ctx._batch = child
+                    if layer.ple is not None:
+                        child.ple_meta = build_ple_metadata(child, layer.ple.args, child.input_ids.device)
+                    _copy_layer_major_track_state(ctx.linear_state_pool, layer, child)
+                    hidden[index] = layer.forward(hidden[index], child)
+                    if layer.ple is not None:
+                        commit_ngram_context(child.ple_meta, child.fla_metadata)
+            return self.hyper_connection_mixer.mix(hidden[-1])[0]
+        finally:
+            ctx._batch = full_batch
+
 
 class Qwen4ExpForCausalLM(BaseLLMModel):
     def __init__(self, config: ModelConfig) -> None:
@@ -212,6 +253,16 @@ def forward(self) -> torch.Tensor:
         batch = get_global_ctx().batch
         return self.lm_head.forward(self.model.forward(batch.input_ids, batch))
 
+    def forward_layer_major(self, chunk_batches: List[Batch]) -> torch.Tensor:
+        ctx = get_global_ctx()
+        previous_batch = ctx.batch
+        try:
+            hidden = self.model.forward_layer_major(chunk_batches)
+            ctx._batch = chunk_batches[-1]
+            return self.lm_head.forward(hidden)
+        finally:
+            ctx._batch = previous_batch
+
 
 class Qwen4ExpForConditionalGeneration(QwenVLVisionMixin, Qwen4ExpForCausalLM):
     def __init__(self, config: ModelConfig) -> None:

python/freetoken/models/qwen4_exp/ple.py

@@ -594,6 +594,8 @@ def forward(
     ) -> torch.Tensor:
         pending, self._pending = self._pending, None
         row_ids = None
+        if meta is None:
+            meta = getattr(batch, "ple_meta", None)
         if meta is None:
             if pending is not None:
                 meta, row_ids = pending

python/freetoken/scheduler/prefill.py

@@ -22,6 +22,8 @@
 
 logger = init_logger(__name__)
 
+_LAYER_MAJOR_MAX_EXTEND_TOKENS = 32_768
+
 
 def _maybe_pinned(t: torch.Tensor) -> torch.Tensor:
     """Pinning only buys the async H2D copy below; without a device it just raises."""
@@ -49,6 +51,7 @@ class PrefillAdder:
     encoder_cache: EncoderCache | None = None
     # end a chunk before an image it would cut; only models whose image spans attend in both directions need it
     keep_images_whole: bool = False
+    layer_major: bool = False
     # the whole budget of this pass; token_budget shrinks as requests are admitted
     pass_budget: int = 0
     # SWA-pool tokens charged to reqs admitted so far this pass. Mirrors reserved_size: swa is
@@ -146,6 +149,15 @@ def _add_one_req(
     ) -> Req | None:
         remain_len = pending_req.input_len - cached_len
         chunk_size = min(self.token_budget, remain_len)
+        if (
+            self.layer_major
+            and remain_len > self.token_budget
+            and remain_len <= _LAYER_MAJOR_MAX_EXTEND_TOKENS
+            and pending_req.mm_items is None
+            and not self.cache_manager.swa_paged
+        ):
+            # The engine splits this measured-size whole request internally.
+            chunk_size = remain_len
         if self.cache_manager.swa_paged:
             # Cap this chunk by the swa the pool can back this pass. swa is allocated per token in
             # allocate_paged, and token_budget (max_extend_tokens, default 8192) won't chunk a
@@ -277,6 +289,7 @@ class PrefillManager:
     decode_manager: DecodeManager
     encoder_cache: EncoderCache | None = None
     keep_images_whole: bool = False
+    layer_major: bool = False
     pending_list: List[PendingReq] = field(default_factory=list)
 
     def add_one_req(self, req: UserMsg) -> None:
@@ -303,6 +316,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
             table_manager=self.table_manager,
             encoder_cache=self.encoder_cache,
             keep_images_whole=self.keep_images_whole,
+            layer_major=self.layer_major and len(self.pending_list) == 1,
         )
         reqs: List[Req] = []
         chunked_list: List[PendingReq] = []