Add offline replay and staged DSA workflows

MANIFEST.in

@@ -15,7 +15,7 @@ include pyproject.toml
 recursive-include .github *.yml *.md
 recursive-include configs *.yaml *.md
 recursive-include docs *.md *.svg *.png *.json
-recursive-include examples *.jsonl *.sh
+recursive-include examples *.jsonl *.sh *.md
 recursive-include aspec/config *.json
 recursive-include aspec/compat/patches *.patch
 recursive-include licenses *

README.md

@@ -59,6 +59,8 @@ Check [supported paths and limitations](docs/support.md), then browse the [recip
 3. [Choose a training recipe](docs/user-guide/recipes.md) and [prepare training data](docs/user-guide/data.md).
 
 Recipes cover Qwen3-8B, Qwen3.5-9B Native MTP, Qwen3.8-27B long context, and Ling tiny/flash 3.0/flash 3.1.
+For complete feature workflows, see [offline feature generation and replay](examples/qwen3-8b-offline/README.md)
+and [DSA dense → indexer → sparse training](examples/qwen3-8b-dsa/README.md).
 
 From the repository root:
 

aspec/offline/generate.py

@@ -143,7 +143,8 @@ def _cleanup_materialization(
         # Only release the reservation after source teardown has acknowledged.
         if inference_pg is not None and cleanup.get("state") == "completed":
             try:
-                ray.util.remove_placement_group(inference_pg)
+                pg, _, _ = inference_pg
+                ray.util.remove_placement_group(pg)
             except Exception as exc:
                 errors.append(f"placement_cleanup: {type(exc).__name__}: {exc}")
     if errors:

configs/examples/README.md

@@ -21,6 +21,11 @@ or disable evaluation/checkpointing; inspect each file before use.
 | [Qwen role-separated TCP](qwen3_8b_4k_multinode_tcp.yaml) | Separate Target and Draft nodes with shared storage |
 | [Qwen Target-generated quality protocol](qwen3_8b_dspark_target_generated.yaml) | Ten-pass training with held-out Target-generated data |
 | [TorchSpec-style input](sglang_qwen3_8b_dspark.yaml) | Fixed compatibility input, not rolling upstream support |
+| [Offline materialization](qwen3_8b_offline_materialize.yaml) / [replay](qwen3_8b_offline_replay.yaml) | Generate DSpark features once, then train with zero Target GPUs |
+
+For complete source-checkout workflows, follow the
+[offline generation, training and recovery example](../../examples/qwen3-8b-offline/README.md)
+or the [DSA three-stage script](../../examples/qwen3-8b-dsa/README.md).
 
 The 9B and 27B DSpark examples each start with one eight-GPU node: four TP1 Target engines and
 four Draft GPUs, cuDNN attention, MBS1, CP1 and 4K context. Set `ASPEC_TARGET_MODEL` to the matching local checkpoint

configs/examples/qwen3_8b_offline_materialize.yaml

@@ -0,0 +1,57 @@
+# Qwen3-8B DSpark feature materialization; see examples/qwen3-8b-offline/README.md.
+# Uses the retained TorchSpec-derived offline generator and ASpec lifecycle validation.
+# One Target GPU is used during generation; no Draft trainer is started.
+model:
+  target_model_path: ${oc.env:ASPEC_TARGET_MODEL}
+  native_context_length: 32768
+  trust_remote_code: false
+  draft_model_config: aspec/config/dspark_draft_config.json
+dataset:
+  train_data_path: ${oc.env:ASPEC_TRAIN_DATA}
+  eval_data_path: ${oc.env:ASPEC_EVAL_DATA}
+  chat_template: qwen3-instruct
+  prompt_key: conversations
+  min_loss_tokens: 14
+  shuffle_dataset: false
+training:
+  attention_backend: cudnn_flex_attention
+  draft_context_parallel_size: 1
+  dflash_packing: false
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  training_num_gpus_per_node: 1
+  training_num_nodes: 1
+  fsdp_strategy: REPLICATE
+  max_seq_length: 4096
+  num_epochs: 1
+  dspark_num_target_layers: 5
+  dspark_num_anchors: 64
+  dflash_block_size: 7
+inference:
+  inference_engine_type: sgl
+  inference_num_gpus: 1
+  inference_num_gpus_per_engine: 1
+  inference_num_gpus_per_node: 1
+  target_prefill_memory_optimized: false
+  store_last_hidden_states: true
+  inference_batch_size: 1
+  inference_fetch_batch: 2
+  inference_buffer_threshold: 2
+  max_sample_pool_size: 4
+  sglang:
+    tp_size: 1
+    mem_fraction_static: 0.5
+mooncake:
+  data_plane: store
+  protocol: tcp
+  enable_gpu_direct: false
+  require_nvlink: false
+  hidden_dim: 4096
+  max_seq_len: 4096
+  global_segment_size: 8GB
+  local_buffer_size: 2GB
+  host_buffer_size: 2GB
+logging:
+  report_to: none
+output_dir: outputs/qwen3-8b-materialize
+cache_dir: cache/qwen3-8b-materialize

configs/examples/qwen3_8b_offline_replay.yaml

@@ -0,0 +1,67 @@
+# Single-GPU DSpark replay; see examples/qwen3-8b-offline/README.md.
+# Keep the Target checkpoint and capture settings identical to materialization.
+model:
+  target_model_path: ${oc.env:ASPEC_TARGET_MODEL}
+  native_context_length: 32768
+  trust_remote_code: false
+  draft_model_config: aspec/config/dspark_draft_config.json
+dataset:
+  train_data_path: null
+  eval_data_path: null
+  renderer: null
+  min_loss_tokens: 14
+  eval_interval: 4
+  eval_micro_batch_size: 1
+  shuffle_dataset: false
+training:
+  attention_backend: cudnn_flex_attention
+  draft_context_parallel_size: 1
+  dflash_packing: false
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  training_num_gpus_per_node: 1
+  training_num_nodes: 1
+  fsdp_strategy: REPLICATE
+  max_seq_length: 4096
+  num_epochs: 100
+  num_train_steps: 8
+  lr_total_steps: 8
+  learning_rate: 0.0006
+  min_lr: 0.00006
+  warmup_ratio: 0.0
+  seed: 42
+  dspark_num_target_layers: 5
+  dspark_num_anchors: 64
+  dflash_block_size: 7
+  prefetch_depth: 1
+  streaming_dispatch_with_training: true
+  pipeline_dispatch_with_training: true
+  save_interval: 4
+  save_final_checkpoint: true
+  save_per_epoch: false
+inference:
+  inference_engine_type: offline
+  inference_num_gpus: 0
+  target_prefill_memory_optimized: false
+  store_last_hidden_states: true
+  inference_batch_size: 1
+  inference_fetch_batch: 2
+  inference_buffer_threshold: 2
+  max_sample_pool_size: 4
+  offline:
+    data_path: ${oc.env:ASPEC_OFFLINE_DATA}
+    num_engines: 1
+mooncake:
+  data_plane: store
+  protocol: tcp
+  enable_gpu_direct: false
+  require_nvlink: false
+  hidden_dim: 4096
+  max_seq_len: 4096
+  global_segment_size: 8GB
+  local_buffer_size: 2GB
+  host_buffer_size: 2GB
+logging:
+  report_to: none
+output_dir: outputs/qwen3-8b-offline
+cache_dir: cache/qwen3-8b-offline

docs/user-guide/structures.md

@@ -172,6 +172,10 @@ compressor/indexer, head, mapping and mixed dtype. See [export](evaluation.md).
 
 ## DSA: token-level sparse attention
 
+The [three-stage runnable example](../../examples/qwen3-8b-dsa/README.md) connects
+the existing dense, indexer and sparse configurations, with stage-local recovery
+instructions and native export.
+
 DSA uses a weighted-ReLU Lightning Indexer to select original context tokens **for each Draft query**.
 It neither compresses tokens like CSA/QSA nor shares one selection across a whole Draft block.
 The selected original K/V and the complete current bidirectional Draft block feed attention;

docs/user-guide/training.md

@@ -34,6 +34,9 @@ For a preplanned split, keep `training.lr_total_steps` equal to the **full** hor
 even when the first `num_train_steps` is smaller. Do not silently change data/objective/topology/model
 semantics across exact resume. Iteration names are one ahead of completed steps. Hidden `.incomplete`
 staging is not published recovery state. Uncommitted features may be recomputed.
+With asynchronous packing, those samples may form different rows after restart;
+checkpoint-state restoration does not guarantee the same document count or
+numerical trajectory at a later optimizer step as an uninterrupted run.
 
 ## New phases and initial weights
 
@@ -131,6 +134,10 @@ The GQA path executes `default`, `linear` and `yarn` RoPE from `rope_scaling` or
 
 ## Offline Target features
 
+The [complete offline example](../../examples/qwen3-8b-offline/README.md) provides
+paired materialization/replay configs and commands for evaluation, exact resume and
+native export on one GPU used sequentially.
+
 Materialize Target features once, then train Drafts from disk with zero Target GPUs. Draft architectures can share the dataset when they request the same Target feature contract. The current path uses single-node, single-Draft, CP1 training and CPU Store/TCP transport.
 
 Use a single-Draft SGLang config with local Target weights and train/eval data. Set `inference.target_prefill_memory_optimized=false`, `training.draft_context_parallel_size=1`, `training.dflash_packing=false`, and the `mooncake` settings shown below. Generate into a new or empty directory:

examples/qwen3-8b-dsa/README.md

@@ -0,0 +1,82 @@
+# Qwen3-8B DSA: three dependent training stages
+
+Run the existing full DSA Draft configuration through dense training, indexer
+warmup and sparse training. Each stage uses a new optimizer/scheduler and output
+directory, initializes from the preceding complete checkpoint, and retains source
+progress.
+
+Use a source checkout with the [isolated Linux environment](../../docs/user-guide/installation.md)
+activated. The default allocation is one Target GPU and one Draft GPU on a node
+with NVLink, using FA4 packing at CP1. Supply disjoint train/eval inputs prepared
+for Qwen3-8B, with sequences up to 4K and at least 32 supervised tokens per usable
+sample. See [data preparation](../../docs/user-guide/data.md).
+
+```bash
+export ASPEC_TARGET_MODEL=/path/to/Qwen3-8B
+export ASPEC_TRAIN_DATA=/path/to/train.parquet
+export ASPEC_EVAL_DATA=/path/to/eval.parquet
+export CUDA_VISIBLE_DEVICES=0,1
+export RAY_ADDRESS=local
+
+bash examples/qwen3-8b-dsa/run.sh outputs/dsa-workflow
+```
+
+The script uses [the existing DSA YAML](../../configs/examples/qwen3_8b_dsa_dspark.yaml)
+and three model JSONs. It checks each Plan, trains 16 updates per stage, evaluates
+every four updates, saves checkpoints and exports the final sparse Draft as a native
+asset. These short defaults demonstrate the workflow, not convergence or sparse
+attention speed. Pass ordinary overrides after the output root, for example
+`training.num_train_steps=8 training.lr_total_steps=8` for eight updates per stage.
+The same overrides apply to all stages; the script fixes stage model, checkpoint,
+source-cursor policy, capture requirements and output paths afterward.
+
+The output root must be new. Failed attempts remain available for diagnosis; the
+script stops on the first failed command and never skips an existing stage.
+The generated tree contains `dense/`, `indexer/`, `sparse/` and `native-asset/`.
+
+| Stage | Trainable state | Initialization |
+| --- | --- | --- |
+| Dense | Main Draft; indexer frozen | Fresh model |
+| Indexer warmup | Indexer; main Draft frozen | Dense checkpoint, weights-only |
+| Sparse | Main Draft and indexer | Indexer checkpoint, weights-only |
+
+Warmup does not request last hidden states needed by the normal DSpark objective;
+the sparse stage requests them again. Warmup evaluation reports indexer KL rather than
+a directly comparable language-model CE. Stage changes are intentional new phases,
+not exact checkpoint recovery.
+
+## Resume within a stage
+
+If training was interrupted after a complete checkpoint, use that stage's saved
+`config.yaml`. Clear **both** phase-transition flags for exact resume and select a
+fresh output directory. Keep the original total step limit and LR horizon.
+
+Recovery restores checkpoint state and the committed source frontier. Asynchronous
+packing can regroup uncommitted samples after restart, so a later optimizer-step
+count can correspond to a different document count and numerical trajectory than
+an uninterrupted run. Check restored state and source continuity separately from
+later trajectory equivalence.
+
+```bash
+aspec --config outputs/dsa-workflow/indexer/config.yaml \
+  training.load_path=outputs/dsa-workflow/indexer/checkpoints \
+  training.continual_training=false training.continue_source_cursor=false \
+  output_dir=outputs/dsa-indexer-resumed
+aspec --report outputs/dsa-indexer-resumed
+```
+
+After that stage finishes, start the sparse phase from its recovered checkpoint:
+
+```bash
+aspec --config outputs/dsa-indexer-resumed/config.yaml \
+  model.draft_model_config=aspec/config/dspark_qwen3_8b_dsa_sparse_draft_config.json \
+  training.load_path=outputs/dsa-indexer-resumed/checkpoints \
+  training.continual_training=true training.continue_source_cursor=true \
+  inference.store_last_hidden_states=true output_dir=outputs/dsa-sparse-after-resume
+```
+
+This uses the same configured horizon for the new stage. Review each Report's
+actual progress, required capabilities and lifecycle. DSA assets use
+`aspec export --format aspec`; strict native loading is separate from downstream
+serving integration. See [DSA semantics](../../docs/user-guide/structures.md#dsa-token-level-sparse-attention)
+and [recovery](../../docs/user-guide/training.md).

examples/qwen3-8b-dsa/run.sh

@@ -0,0 +1,59 @@
+#!/usr/bin/env bash
+# Copyright 2026 Ant Group Co., Ltd. Licensed under Apache-2.0.
+# Run dense, indexer warmup and sparse training with native checkpoint export.
+set -euo pipefail
+
+if [[ $# -lt 1 || ${1:-} == --help || ${1:-} == -h ]]; then
+  echo 'Usage: bash examples/qwen3-8b-dsa/run.sh NEW_OUTPUT_ROOT [KEY=VALUE ...]'
+  echo 'Set ASPEC_TARGET_MODEL, ASPEC_TRAIN_DATA, ASPEC_EVAL_DATA and CUDA_VISIBLE_DEVICES.'
+  echo 'Runs dense -> indexer warmup -> sparse, then exports a native asset.'
+  exit 0
+fi
+: "${ASPEC_TARGET_MODEL:?Set an absolute local Qwen3-8B checkpoint path}"
+: "${ASPEC_TRAIN_DATA:?Set an absolute training data path}"
+: "${ASPEC_EVAL_DATA:?Set an absolute disjoint evaluation data path}"
+: "${CUDA_VISIBLE_DEVICES:?Select two allocated GPUs explicitly}"
+
+# Resolve the output before changing directory so paths are relative to the caller.
+output_root=$(python -c 'import pathlib,sys; print(pathlib.Path(sys.argv[1]).resolve())' "$1")
+shift
+repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
+cd "$repo_root"
+mkdir -p -- "$(dirname -- "$output_root")"
+mkdir -- "$output_root"
+export RAY_ADDRESS="${RAY_ADDRESS:-local}"
+config=configs/examples/qwen3_8b_dsa_dspark.yaml
+previous=
+
+for stage in dense indexer sparse; do
+  continuation=false
+  last_hidden=true
+  if [[ $stage != dense ]]; then
+    continuation=true
+  fi
+  if [[ $stage == indexer ]]; then
+    last_hidden=false
+  fi
+  args=(--config "$config"
+    "model.target_model_path=$ASPEC_TARGET_MODEL"
+    "dataset.train_data_path=$ASPEC_TRAIN_DATA"
+    "dataset.eval_data_path=$ASPEC_EVAL_DATA"
+    dataset.eval_interval=4 "$@"
+    "model.draft_model_config=aspec/config/dspark_qwen3_8b_dsa_${stage}_draft_config.json"
+    "training.continual_training=$continuation"
+    "training.continue_source_cursor=$continuation"
+    "training.load_path=${previous:-null}"
+    training.save_final_checkpoint=true
+    "inference.store_last_hidden_states=$last_hidden"
+    "output_dir=$output_root/$stage"
+    "cache_dir=$output_root/cache")
+  aspec "${args[@]}" --dry-run
+  aspec "${args[@]}"
+  aspec --report "$output_root/$stage"
+  previous="$output_root/$stage/checkpoints"
+done
+
+# Export the final published checkpoint selected by the existing tracker.
+iteration=$(cat "$previous/latest_checkpointed_iteration.txt")
+printf -v checkpoint '%s/iter_%07d' "$previous" "$iteration"
+aspec export --format aspec --input-dir "$checkpoint" --output-dir "$output_root/native-asset"