[2/4] offline: prepare target features through the capture plugin

docs/sections/basic_usage/AMD/amd_rocm.md

@@ -87,20 +87,18 @@ python -m pip install wandb
 python -m pip install --no-deps tensorboard tensorboard-data-server absl-py markdown werkzeug
 ```
 
-### Step 4: Apply the capture patch (online runs only)
+### Step 4: Install the capture plugin (online and offline preparation)
 
-These images pin SGLang to exactly `0.5.18` (editable at `/sgl-workspace/sglang`),
-so the online capture patch applies unchanged. Skip this step for offline
-training, which reads features from disk and needs no capture service:
+The images' stock SGLang 0.5.18 lacks the generic capture extension APIs.
+Use a separate ROCm-compatible SGLang environment with the APIs documented in
+[the plugin guide](../../../../plugins/sglang-spec-capture/README.md), then install:
 
 ```bash
-cd /workspace/SpecForge
-scripts/apply_sglang_spec_capture_patch.sh --target v0.5.18
+python -m pip install --no-deps plugins/sglang-spec-capture
 ```
 
-The patch adds the `--enable-spec-capture`, `--spec-capture-method`, and
-`--spec-capture-aux-layer-ids` server flags plus the `sglang.srt.spec_capture_sink`
-module used by online capture.
+Set `SPECFORGE_SPEC_CAPTURE=1` when launching the capture server. Offline
+training from existing feature files needs neither SGLang nor this plugin.
 
 ### Step 5: Attention backends on ROCm
 
@@ -140,21 +138,26 @@ see the [Data Preparation](../data_preparation.md) guide.
 
 Offline training reads target features from disk, so the trainer only has to fit
 the draft model. It uses more storage but keeps target inference out of the
-training loop, and needs no capture patch or Mooncake. This section trains a
+training loop, and needs no capture service or Mooncake. This section trains a
 **Qwen3.5-4B DFlash** draft (`configs/qwen3.5-4b-dflash.json`).
 
 ### Step 1: Capture hidden states
 
 Feature preparation is a data-processing step, not a second training entry point.
 Qwen3.5-4B is a hybrid linear-attention/Mamba target, so run the capture under
-**AITER** and pass `--sglang-disable-radix-cache`. Without it, SGLang's Mamba
+**AITER** and pass `--disable-radix-cache` to the capture server. Without it, SGLang's Mamba
 radix cache selects the `extra_buffer` strategy, which asserts CUDA/MUSA/NPU
 (FLA) at server init and fails on ROCm:
 
+Start the compatible plugin server and Mooncake following the
+[offline guide](../../../../examples/offline/README.md), with
+`SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1`
+and server flags `--attention-backend aiter --disable-radix-cache
+--mem-fraction-static 0.8 --context-length 2560`. Then run the CPU client:
+
 ```bash
-SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1 \
-torchrun --standalone --nproc_per_node 1 \
-  scripts/prepare_hidden_states.py \
+python scripts/prepare_hidden_states.py \
+  --server-urls http://127.0.0.1:30000 \
   --target-model-path Qwen/Qwen3.5-4B \
   --strategy dflash \
   --draft-model-config configs/qwen3.5-4b-dflash.json \
@@ -163,12 +166,7 @@ torchrun --standalone --nproc_per_node 1 \
   --output-path ./cache/hidden_states/qwen3.5-4b-dflash-sharegpt \
   --chat-template qwen3.5 \
   --max-length 2048 \
-  --tp-size 1 \
-  --batch-size 8 \
-  --sglang-attention-backend aiter \
-  --sglang-disable-radix-cache \
-  --sglang-mem-fraction-static 0.8 \
-  --sglang-context-length 2560
+  --batch-size 8
 ```
 
 The output path matches `data.hidden_states_path` in the checked-in offline
@@ -219,7 +217,7 @@ for the complete contract.
 
 ## 4. Online disaggregated training
 
-Online training captures target features live from a patched SGLang server and
+Online training captures target features live from a SGLang server running the capture plugin and
 streams them through Mooncake to the trainer. Every online run is
 **disaggregated**: a producer drives prompts through the capture server and a
 consumer trains the draft model. The choices below change service ownership and
@@ -263,7 +261,7 @@ mooncake_master --enable_http_metadata_server=true \
   --default_kv_lease_ttl=500 &
 ```
 
-Start the patched capture server on GPU 0. **The `--spec-capture-aux-layer-ids`
+Start the plugin-enabled capture server on GPU 0. **The `--return-aux-hidden-state-layer-ids`
 must match the draft's `target_layer_ids`** — for DFlash these are read straight
 from `configs/qwen3.5-4b-dflash.json`: `1 8 15 22 29` (this is not the EAGLE3
 `[1, num_layers//2 - 1, num_layers - 4]` formula). A mismatch produces zero
@@ -279,6 +277,7 @@ MOONCAKE_METADATA_SERVER=http://127.0.0.1:35880/metadata \
 MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:35551 \
 MOONCAKE_PROTOCOL=tcp \
 MOONCAKE_GLOBAL_SEGMENT_SIZE=$((32<<30)) \
+SPECFORGE_SPEC_CAPTURE=1 \
 python -m sglang.launch_server \
   --model-path Qwen/Qwen3.5-4B \
   --trust-remote-code \
@@ -289,8 +288,8 @@ python -m sglang.launch_server \
   --attention-backend aiter \
   --chunked-prefill-size -1 \
   --disable-radix-cache \
-  --enable-spec-capture --spec-capture-method dflash \
-  --spec-capture-aux-layer-ids 1 8 15 22 29 \
+  --enable-return-aux-hidden-states --return-hidden-states-mode full --speculative-algorithm DFLASH \
+  --return-aux-hidden-state-layer-ids 1 8 15 22 29 \
   --host 127.0.0.1 --port 30000 &
 ```
 
@@ -378,7 +377,7 @@ specforge train -c run.yaml --role consumer --node-rank 1   # trainer-1
 
 A fresh attempt requires fresh control and consumer-state directories, and every
 capture server must use the same target model, revision, capture method, and
-auxiliary layer ids — for this recipe `--spec-capture-method dflash` with aux ids
+auxiliary layer ids — for this recipe `--speculative-algorithm DFLASH` with aux ids
 `1 8 15 22 29`, and on ROCm each must run under AITER with `--disable-radix-cache`
 (see External services, Step 2). For external-service prerequisites, freshness
 rules, multi-server capture, inbox relays, and resume, see the
@@ -405,8 +404,8 @@ The SGLang side (offline capture and online capture server) runs the hybrid
 Mamba target under **AITER** on ROCm — export
 `SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1 AITER_FLYDSL_FORCE=1` and use
 `--attention-backend aiter`. A ROCm-specific requirement: the target needs
-the radix cache disabled (offline capture: `--sglang-disable-radix-cache`;
-external online server: `--disable-radix-cache`; managed-local config:
+the radix cache disabled (offline/online capture server:
+`--disable-radix-cache`; managed-local config:
 `model.sglang_disable_radix_cache: true`). SGLang's Mamba radix cache
 auto-selects the `extra_buffer` strategy, which asserts CUDA/MUSA/NPU (FLA) at
 server init and fails on ROCm; disabling the radix cache bypasses that path.

docs/sections/basic_usage/data_preparation.md

@@ -257,6 +257,13 @@ specforge train --config ./my-eagle3-disaggregated.yaml
 
 ## Prepare offline target features
 
+First start a compatible SGLang plugin capture server and Mooncake using the
+[offline capture guide](../../../examples/offline/README.md). The commands below
+are CPU clients; `--server-urls` must point to servers configured for the same
+target, capture method and layer IDs. TP and attention-backend flags belong to
+the server. Optional `torchrun --standalone --nproc_per_node=N` launches N CPU
+clients with non-overlapping dataset shards.
+
 Offline EAGLE3, DFlash, DFlash2, Domino, and DSpark runs consume
 target features created by `scripts/prepare_hidden_states.py`. Select the same
 strategy and draft
@@ -265,16 +272,15 @@ the target layers captured and the checkpoint schema. For example, prepare the
 checked-in Qwen3-8B DFlash recipe with:
 
 ```bash
-torchrun --nproc_per_node=8 \
-    scripts/prepare_hidden_states.py \
+python scripts/prepare_hidden_states.py \
+    --server-urls http://127.0.0.1:30000 \
     --strategy dflash \
     --target-model-path Qwen/Qwen3-8B \
     --draft-model-config configs/qwen3-8b-dflash.json \
     --data-path ./cache/dataset/sharegpt_train.jsonl \
     --output-path ./cache/hidden_states/qwen3-8b-dflash-sharegpt \
     --chat-template qwen \
     --max-length 3072 \
-    --tp-size 1 \
     --batch-size 32
 ```
 
@@ -322,8 +328,8 @@ For preformatted input, add `--is-preformatted` to the same command and keep
 `--chat-template` aligned with the template already applied to the text:
 
 ```bash
-torchrun --nproc_per_node=8 \
-    scripts/prepare_hidden_states.py \
+python scripts/prepare_hidden_states.py \
+    --server-urls http://127.0.0.1:30000 \
     --strategy eagle3 \
     --target-model-path meta-llama/Llama-3.1-8B-Instruct \
     --draft-model-config configs/llama3-8B-eagle3.json \

docs/sections/examples/llama3-eagle3-offline.md

@@ -15,17 +15,20 @@ python ./scripts/prepare_data.py --dataset sharegpt
 Feature preparation is a data-processing step, not a second training entry
 point:
 
+Start the plugin server and Mooncake as described in the
+[offline capture guide](../../../examples/offline/README.md), using the same
+Llama target and EAGLE3 layer IDs. Then run the CPU preparation client:
+
 ```bash
-torchrun --standalone --nproc_per_node 8 \
-  scripts/prepare_hidden_states.py \
+python scripts/prepare_hidden_states.py \
+  --server-urls http://127.0.0.1:30000 \
   --strategy eagle3 \
   --target-model-path meta-llama/Llama-3.1-8B-Instruct \
   --draft-model-config configs/llama3-8B-eagle3.json \
   --data-path ./cache/dataset/sharegpt_train.jsonl \
   --output-path ./cache/hidden_states/sharegpt_train_Llama-3.1-8B-Instruct \
   --chat-template llama3 \
   --max-length 4096 \
-  --tp-size 1 \
   --batch-size 32
 ```
 

examples/configs/offline/colocated/qwen3.5-4b-dflash-amd.yaml

@@ -1,7 +1,7 @@
 # Qwen3.5-4B DFlash draft, offline (features captured to disk) — AMD ROCm.
 # Trainer runs on ROCm unchanged (flex_attention). Capture the hidden states with
-# scripts/prepare_hidden_states.py using the AITER + --sglang-disable-radix-cache
-# flags (see docs/sections/basic_usage/AMD/amd_rocm.md, Section 3).
+# scripts/prepare_hidden_states.py --server-urls pointing at an AITER plugin
+# server with --disable-radix-cache (see docs/sections/basic_usage/AMD/amd_rocm.md, Section 3).
 model:
   target_model_path: "Qwen/Qwen3.5-4B"
   draft_model_config: "configs/qwen3.5-4b-dflash.json"

examples/offline/README.md

@@ -0,0 +1,98 @@
+# Prepare offline features with the capture plugin
+
+Offline preparation uses a CPU client, SGLang capture servers, and Mooncake:
+
+`dataset -> plugin prefill -> Mooncake -> CPU client -> .ckpt/.ckpt.gz`
+
+Offline colocated and disaggregated **training** continue to read these files;
+they do not need a capture server or a local ModelRunner. The preparation client
+does not import SGLang or load target model weights. It still needs the target's
+config/tokenizer and the same draft config used for training.
+
+## Start the services
+
+Use an SGLang environment with the native APIs described in the
+[plugin installation guide](../../plugins/sglang-spec-capture/README.md), and
+install the current plugin there. Use SGLang commit `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0`; released 0.5.23
+still lacks the capture APIs. The server and SpecForge client may use the same
+environment with a compatible Mooncake wheel, or separate environments.
+
+For a single-host TCP example, start a Mooncake master in a separate terminal:
+
+```bash
+mooncake_master --enable-http-metadata-server=true
+```
+
+Export the following in both the server and client terminals:
+
+```bash
+export MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:50051
+export MOONCAKE_METADATA_SERVER=http://127.0.0.1:8080/metadata
+export MOONCAKE_LOCAL_HOSTNAME=127.0.0.1
+export MOONCAKE_PROTOCOL=tcp
+```
+
+Across hosts, replace loopback addresses with reachable addresses and set each
+process's `MOONCAKE_LOCAL_HOSTNAME` to its own reachable address. The offline CPU
+client supports TCP/RDMA; NVLink GPU receive buffers are not used here.
+
+For `configs/qwen3-8b-dflash.json`, launch this server in the SGLang environment:
+
+```bash
+SPECFORGE_SPEC_CAPTURE=1 SGLANG_SPEC_CAPTURE_GPU_PUT=0 \
+python -m sglang.launch_server \
+  --model-path Qwen/Qwen3-8B --skip-tokenizer-init \
+  --port 30000 --tp-size 1 \
+  --speculative-algorithm DFLASH \
+  --enable-return-aux-hidden-states \
+  --return-aux-hidden-state-layer-ids 1 9 17 25 33 \
+  --return-hidden-states-mode full --chunked-prefill-size -1
+```
+
+Choose the method and layers from the resolved draft capture plan. Domino uses
+the DFLASH convention; EAGLE3 and DSpark use EAGLE3 and DSPARK respectively.
+The client checks every server's `/server_info` for the exact `model_path`,
+capture method, and ordered layer IDs. Use the same model identifier/path on
+both sides, and reinstall the plugin on existing servers to advertise these
+fields. TP size, device selection, attention backend, context length and radix
+cache options belong to the server command.
+
+## Capture and resume
+
+From the SpecForge environment:
+
+```bash
+python scripts/prepare_hidden_states.py \
+  --server-urls http://127.0.0.1:30000 \
+  --target-model-path Qwen/Qwen3-8B \
+  --strategy dflash --draft-model-config configs/qwen3-8b-dflash.json \
+  --data-path ./cache/dataset/sharegpt_train.jsonl \
+  --output-path ./cache/hidden_states/qwen3-8b-dflash-sharegpt \
+  --chat-template qwen --max-length 3072 --batch-size 32
+```
+
+Add `--compress` for `.ckpt.gz`. The output retains the existing
+`rows_START-END/data_INDEX.ckpt[.gz]` layout and vocabulary mapping. Algorithms
+own the persisted feature names. Input IDs and loss masks remain client-side;
+the service only returns hidden-state artifacts.
+
+Multiple `--server-urls` are used in round-robin order. To run concurrent CPU
+clients, replace `python` with `torchrun --standalone --nproc_per_node=4`.
+Workers partition samples without sampler padding or duplicate submissions;
+this worker count is independent of server TP size.
+
+Rerun the same command to skip completed files. Changing the batch size or CPU
+worker count preserves sample IDs. Keep dataset contents/order, preprocessing,
+strategy, target, draft config, and file-group size unchanged when resuming;
+use a new output directory for a different dataset or capture configuration.
+
+Writes use temporary files and atomic rename. The client waits for writes to
+finish before releasing temporary Mooncake objects, and attempts cleanup on
+request, read, or write errors. A response timeout fails the batch instead of
+silently retrying it. This is not a durable job/lease protocol: abrupt process
+termination or a server publication arriving after a timed-out request can
+leave objects until the store is reclaimed. Size the store for concurrent
+batches and use `--capture-timeout` appropriate to the target.
+
+The service still runs SGLang's normal prefill/logits path with
+`max_new_tokens=0`; this migration does not bypass the target LM head.

plugins/sglang-spec-capture/README.md

@@ -15,7 +15,10 @@ generic extension points it uses:
 - `Scheduler.defer_response(req, completion)`: SGLang holds a finished
   request until its objects are written, then reads the completion on the
   scheduler thread before returning capture metadata. Cancelling one request
-  does not cancel a shared write needed by other requests.
+  does not cancel a shared write needed by other requests. It queues reclamation
+  of that request's objects on the writer after publication, so a client cleanup
+  that ran before the objects existed cannot leave them orphaned. Cleanup RPC
+  failures are retried and logged if they persist.
 
 ## Install
 

plugins/sglang-spec-capture/specforge_sglang_capture/capture.py

@@ -304,7 +304,13 @@ def consume(self, batch, commits) -> None:
             self.runtime.defer_response(
                 req,
                 CaptureCompletion(
-                    future, index, len(samples), req.rid, spec, self.runtime.results
+                    future,
+                    index,
+                    len(samples),
+                    req.rid,
+                    spec,
+                    self.runtime.results,
+                    self.runtime.sink.discard_sample,
                 ),
             )
 
@@ -315,6 +321,8 @@ class CaptureCompletion:
 
     SGLang calls these methods on the scheduler thread. Cancelling one request
     leaves the shared write running for its siblings and suppresses its metadata.
+    Cleanup is queued on the sink's writer, after publication, even if the
+    client has already tried to remove keys that did not exist yet.
     """
 
     future: Future
@@ -323,6 +331,7 @@ class CaptureCompletion:
     rid: str
     spec: Dict[str, Any]
     results: Dict[str, Dict[str, Any]]
+    discard_sample: Callable[[Dict[str, Any]], Future]
     _cancelled: bool = False
     _resolved: bool = False
 
@@ -349,6 +358,9 @@ def result(self) -> None:
     def cancel(self) -> bool:
         if self._resolved:
             return False
+        if self._cancelled:
+            return True
+        self.discard_sample(self.spec)
         self._cancelled = True
         return True