初始化项目，由ModelHub XC社区提供模型

Model: jiamingshan/AHA-L2A-Qwen3-1.7B-repro Source: Original Platform
2026-07-21 11:06:13 +08:00
commit c13d63b439
48 changed files with 165513 additions and 0 deletions
--- a/recipe/scripts/eval_cell.sh
+++ b/recipe/scripts/eval_cell.sh
@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -ne 4 ]]; then
+  echo "Usage: $0 METHOD MODEL_PATH GPU BENCHMARK" >&2
+  exit 2
+fi
+
+METHOD="$1"
+MODEL_PATH="$2"
+GPU="$3"
+BENCHMARK="$4"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+EVAL_ROOT="${EVAL_ROOT:-$REPO/outputs/eval}"
+OUT="$EVAL_ROOT/$METHOD"
+LOG="$EVAL_ROOT/logs/$METHOD"
+TOKENIZER="${VANILLA_DIR:-$REPO}"
+HELMET_ROWS="$RECIPE/data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl"
+MRCR_ROWS="$RECIPE/data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl"
+RULER_LIMIT="${RULER_LIMIT:-20}"
+BABILONG_LIMIT="${BABILONG_LIMIT:-50}"
+AHA_GATE_HARD_THRESHOLD="${AHA_GATE_HARD_THRESHOLD:-0.5}"
+mkdir -p "$OUT" "$LOG"
+
+if [[ -f "$OUT/${BENCHMARK}.DONE" ]]; then
+  exit 0
+fi
+
+export CUDA_VISIBLE_DEVICES="$GPU"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export AHA_SPARSITY_STATS_PATH="$OUT/sparsity_${BENCHMARK}.json"
+export AHA_GATE_HARD_THRESHOLD
+export AHA_FORCE_FULL_DECODE=0
+export AHA_FORCE_FULL_PREFILL_TAIL=0
+export AHA_FORCE_FULL_HEADS=
+export AHA_FORCE_LOW_ALPHA_FULL_HEADS=0
+export AHA_DUO_PREFILL_FULL=0
+
+model_args="pretrained=${MODEL_PATH},trust_remote_code=True,dtype=bfloat16,max_length=16384,attn_implementation=sdpa"
+split_a="niah_single_1,niah_single_3,niah_multikey_2,niah_multiquery,ruler_vt,ruler_fwe,ruler_qa_hotpot"
+split_b="niah_single_2,niah_multikey_1,niah_multikey_3,niah_multivalue,ruler_cwe,ruler_qa_squad"
+
+case "$BENCHMARK" in
+  ruler_a|ruler_b)
+    tasks="$split_a"
+    [[ "$BENCHMARK" == ruler_b ]] && tasks="$split_b"
+    python -m lm_eval \
+      --model hf --model_args "$model_args" \
+      --tasks "$tasks" --metadata '{"max_seq_lengths":[8192]}' \
+      --batch_size 1 --limit "$RULER_LIMIT" --log_samples \
+      --output_path "$OUT/ruler8k_splits/${BENCHMARK#ruler_}/lm_eval" \
+      >"$LOG/${BENCHMARK}.log" 2>&1
+    ;;
+  babilong)
+    python -m lm_eval \
+      --model hf --model_args "$model_args" \
+      --tasks babilong_longctx --metadata '{"max_seq_lengths":"8k"}' \
+      --num_fewshot 2 --batch_size 1 --limit "$BABILONG_LIMIT" --log_samples \
+      --output_path "$OUT/babilong8k_qa1_qa5_n50" \
+      >"$LOG/babilong.log" 2>&1
+    ;;
+  helmet)
+    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
+      --benchmark helmet_icl --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
+      --attn-implementation sdpa --rows "$HELMET_ROWS" \
+      --method "$METHOD" --output "$OUT/helmet_icl8k_n50.jsonl" \
+      >"$LOG/helmet.log" 2>&1
+    ;;
+  mrcr)
+    python "$SCRIPT_DIR/eval_qwen_external_longctx_pilot.py" \
+      --benchmark mrcr --model-path "$MODEL_PATH" --tokenizer-path "$TOKENIZER" \
+      --attn-implementation sdpa --rows "$MRCR_ROWS" \
+      --method "$METHOD" --output "$OUT/mrcr_8k_2_4_8needle_n10.jsonl" \
+      >"$LOG/mrcr.log" 2>&1
+    ;;
+  *) echo "Unknown benchmark: $BENCHMARK" >&2; exit 2 ;;
+esac
+
+touch "$OUT/${BENCHMARK}.DONE"
--- a/recipe/scripts/eval_qwen_external_longctx_pilot.py
+++ b/recipe/scripts/eval_qwen_external_longctx_pilot.py
@@ -0,0 +1,275 @@
+#!/usr/bin/env python3
+"""Aligned 8K HELMET-ICL and OpenAI-MRCR pilot for Qwen checkpoints.
+
+This is deliberately a small diagnostic runner.  HELMET prompt construction
+matches the ICL schedule used by the local OmniServe reproduction; MRCR uses
+the official OpenAI rows selected by that reproduction, but renders the
+messages with the Qwen chat template.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import json
+import math
+import random
+import re
+import string
+from collections import defaultdict
+from difflib import SequenceMatcher
+from pathlib import Path
+
+import pandas as pd
+import torch
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+HELMET_SPECS = {
+    "trec_coarse": ("icl_trec_coarse_400shot_balance", 6),
+    "trec_fine": ("icl_trec_fine_400shot_balance", 50),
+    "banking77": ("icl_banking77_360shot_balance", 77),
+    "clinic150": ("icl_clinic150_440shot_balance", 151),
+    "nlu": ("icl_nlu_510shot_balance", 68),
+}
+
+
+def balanced(data, shots: int, label_field: str, seed: int):
+    rng = random.Random(seed)
+    by_label = defaultdict(list)
+    for item in data:
+        by_label[item[label_field]].append(item)
+    rounds = math.ceil(shots / len(by_label))
+    selected_rounds = [[] for _ in range(rounds)]
+    for examples in by_label.values():
+        indices = rng.sample(range(len(examples)), rounds % len(examples))
+        while len(indices) < rounds:
+            indices += rng.sample(range(len(examples)), min(rounds - len(indices), len(examples)))
+        for index, example_index in enumerate(indices):
+            selected_rounds[index].append(examples[example_index])
+    for examples in selected_rounds:
+        rng.shuffle(examples)
+    return [item for group in selected_rounds for item in group][:shots]
+
+
+def helmet_source(family: str, seed: int):
+    if family == "trec_coarse":
+        source = load_dataset("CogComp/trec", trust_remote_code=True)
+        return source["train"], source["test"], "text", "coarse_label"
+    if family == "trec_fine":
+        source = load_dataset("CogComp/trec", trust_remote_code=True)
+        return source["train"], source["test"], "text", "fine_label"
+    if family == "banking77":
+        source = load_dataset("PolyAI/banking77", trust_remote_code=True)
+        return source["train"], source["test"], "text", "label"
+    if family == "clinic150":
+        source = load_dataset("clinc/clinc_oos", "plus")
+        return source["train"], source["validation"], "text", "intent"
+    if family == "nlu":
+        source = load_dataset("xingkunliuxtracta/nlu_evaluation_data", trust_remote_code=True)["train"]
+        split = source.train_test_split(test_size=0.1, seed=seed)
+        return split["train"], split["test"], "text", "label"
+    raise ValueError(family)
+
+
+def build_helmet(seed: int, limit: int):
+    rows = []
+    user_template = (
+        'Use the provided mapping from the text to label to assign a label to the text. '
+        'Only output "label: {{label}}" and nothing else. \n\n{context}\n\n{question}'
+    )
+    for family, (dataset_name, num_labels) in HELMET_SPECS.items():
+        shots = int(dataset_name.split("shot")[0].split("_")[-1])
+        train, test, text_field, label_field = helmet_source(family, seed)
+        samples = balanced(test, limit, label_field, seed)
+        for sample_index, sample in enumerate(samples):
+            local_seed = (int(hashlib.sha256(sample[text_field].encode()).hexdigest(), 16) + seed) % 2**31
+            demos = balanced(train, shots, label_field, local_seed)
+            mapping = list(range(num_labels))
+            random.Random(local_seed).shuffle(mapping)
+            context = "\n\n".join(
+                f"{demo[text_field]}\nlabel: {mapping[int(demo[label_field])]}" for demo in demos
+            )
+            rows.append({
+                "row_id": f"helmet_icl:8k:{family}:{sample_index}",
+                "benchmark": "helmet_icl",
+                "config": family,
+                "prompt": user_template.format(context=context, question=sample[text_field]) + "\nlabel:",
+                "answer": str(mapping[int(sample[label_field])]),
+                "max_new_tokens": 20,
+                "metadata": {"dataset": dataset_name, "shots": shots},
+            })
+    return rows
+
+
+def build_mrcr(tokenizer, selection_path: Path, limit: int):
+    selection = [json.loads(line) for line in selection_path.read_text().splitlines() if line.strip()]
+    chosen = sorted(
+        (row for row in selection if int(row["context_k"]) == 8),
+        key=lambda row: int(row["sample_index"]),
+    )[:limit]
+    if len(chosen) < limit:
+        raise ValueError(f"MRCR selection contains {len(chosen)} 8K rows, requested {limit}")
+    files = [
+        hf_hub_download("openai/mrcr", filename=f"2needle/2needle_{shard}.parquet", repo_type="dataset")
+        for shard in (0, 1)
+    ]
+    frame = pd.concat([pd.read_parquet(path) for path in files], ignore_index=True)
+    rows = []
+    for selected in chosen:
+        source = frame.iloc[int(selected["source_index"])]
+        messages = json.loads(source["prompt"])
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+        answer = str(source["answer"])
+        sample_index = int(selected["sample_index"])
+        rows.append({
+            "row_id": f"mrcr:8k:2needle:{sample_index}",
+            "benchmark": "mrcr",
+            "config": "2needle",
+            "prompt": prompt,
+            "answer": answer,
+            "prefix": str(source["random_string_to_prepend"]),
+            "max_new_tokens": min(768, len(tokenizer(answer, add_special_tokens=False).input_ids) + 64),
+            "metadata": {"source_index": int(selected["source_index"]), "official_tokens": int(selected["official_tokens"])},
+        })
+    return rows
+
+
+def normalize(text: str) -> str:
+    text = text.lower()
+    text = re.sub(r"\b(a|an|the)\b", " ", text)
+    text = "".join(ch for ch in text if ch not in string.punctuation)
+    return " ".join(text.split())
+
+
+def score(row, prediction: str) -> float:
+    if row["benchmark"] == "helmet_icl":
+        first = prediction.strip().splitlines()[0] if prediction.strip() else ""
+        parsed = re.sub(r"^label:", "", first, flags=re.I).strip()
+        return float(normalize(parsed) == normalize(row["answer"]))
+    prefix = row["prefix"]
+    if not prediction.startswith(prefix):
+        return 0.0
+    response = prediction.removeprefix(prefix).strip()
+    reference = row["answer"].removeprefix(prefix).strip()
+    return float(SequenceMatcher(None, response, reference).ratio())
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--benchmark", choices=("helmet_icl", "mrcr"), required=True)
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--tokenizer-path", required=True)
+    parser.add_argument(
+        "--attn-implementation",
+        default="flash_attention_2",
+        choices=("eager", "sdpa", "flash_attention_2"),
+        help="Transformers attention backend; AHA 4-D local masks require SDPA or eager.",
+    )
+    parser.add_argument("--selection", type=Path)
+    parser.add_argument("--rows", type=Path, help="Frozen aligned prompt rows (JSONL)")
+    parser.add_argument("--method", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--seed", type=int, default=20260710)
+    parser.add_argument("--limit", type=int, default=1, help="Samples per benchmark config")
+    parser.add_argument(
+        "--stop-new-line",
+        action="store_true",
+        help="Use HELMET's newline stop-token policy during generation.",
+    )
+    parser.add_argument(
+        "--duo-sink-size",
+        type=int,
+        help="Optional inference-only override for the AHA sink token count.",
+    )
+    parser.add_argument(
+        "--duo-recent-size",
+        type=int,
+        help="Optional inference-only override for the AHA recent-window token count.",
+    )
+    args = parser.parse_args()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, trust_remote_code=True, use_fast=False)
+    if args.rows:
+        rows = [json.loads(line) for line in args.rows.read_text().splitlines() if line.strip()]
+    else:
+        rows = (
+            build_helmet(args.seed, args.limit)
+            if args.benchmark == "helmet_icl"
+            else build_mrcr(tokenizer, args.selection, args.limit)
+        )
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        attn_implementation=args.attn_implementation,
+    )
+    for field, value in (
+        ("duo_sink_size", args.duo_sink_size),
+        ("duo_recent_size", args.duo_recent_size),
+    ):
+        if value is not None:
+            if value < 0:
+                parser.error(f"--{field.replace('_', '-')} must be non-negative")
+            previous = getattr(model.config, field, None)
+            setattr(model.config, field, value)
+            print(f"[eval] {field} override: {previous!r} -> {value!r}", flush=True)
+    model = model.to("cuda").eval()
+    stop_token_ids = model.generation_config.eos_token_id
+    stop_token_ids = list(stop_token_ids) if isinstance(stop_token_ids, list) else [stop_token_ids]
+    if args.stop_new_line:
+        newline_tokens = ["\n", "Ċ", "ĊĊ", "<0x0A>"]
+        stop_token_ids += [tokenizer.convert_tokens_to_ids(token) for token in newline_tokens]
+    stop_token_ids = sorted(
+        {
+            token_id
+            for token_id in stop_token_ids
+            if token_id is not None and token_id != tokenizer.unk_token_id
+        }
+    )
+    print(f"[eval] stop_token_ids={stop_token_ids}", flush=True)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    outputs = []
+    for ordinal, row in enumerate(rows, 1):
+        encoded = tokenizer(row["prompt"], return_tensors="pt", add_special_tokens=False)
+        prompt_tokens = int(encoded.input_ids.shape[1])
+        print(f"[{args.benchmark}] {ordinal}/{len(rows)} {row['config']} tokens={prompt_tokens}", flush=True)
+        encoded = {key: value.to("cuda") for key, value in encoded.items()}
+        with torch.inference_mode():
+            generated = model.generate(
+                **encoded,
+                do_sample=False,
+                max_new_tokens=row["max_new_tokens"],
+                use_cache=True,
+                eos_token_id=stop_token_ids,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+        prediction = tokenizer.decode(generated[0, prompt_tokens:], skip_special_tokens=True).strip()
+        outputs.append({
+            **row,
+            "method": args.method,
+            "prompt_tokens": prompt_tokens,
+            "prompt_sha256": hashlib.sha256(row["prompt"].encode()).hexdigest(),
+            "prediction": prediction,
+            "score": score(row, prediction),
+        })
+    args.output.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in outputs))
+    summary = args.output.with_suffix(".summary.csv")
+    with summary.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=["method", "benchmark", "config", "n", "score", "prompt_tokens"])
+        writer.writeheader()
+        for row in outputs:
+            writer.writerow({"method": args.method, "benchmark": row["benchmark"], "config": row["config"], "n": 1, "score": row["score"], "prompt_tokens": row["prompt_tokens"]})
+    print(summary)
+
+
+if __name__ == "__main__":
+    main()
--- a/recipe/scripts/eval_sweep.sh
+++ b/recipe/scripts/eval_sweep.sh
@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
+EVAL_ROOT="${EVAL_ROOT:-$OUTPUT_ROOT/eval}"
+VANILLA_MODEL="${VANILLA_DIR:-$REPO}"
+AHA_MODEL="${AHA_MODEL:-$OUTPUT_ROOT/aha/stage2/checkpoint-25}"
+L2A_MODEL="${L2A_MODEL:-$OUTPUT_ROOT/l2a_style/stage2/checkpoint-25}"
+GPU_LIST="${GPU_LIST:-0}"
+THRESHOLDS=(0.45 0.50 0.525 0.55 0.575 0.60 0.625 0.65)
+BENCHMARKS=(ruler_a ruler_b babilong helmet mrcr)
+IFS=',' read -r -a GPUS <<< "$GPU_LIST"
+declare -A SLOT_PIDS=()
+mkdir -p "$EVAL_ROOT"
+
+for model in "$VANILLA_MODEL" "$AHA_MODEL" "$L2A_MODEL"; do
+  [[ -f "$model/config.json" ]] || { echo "Missing model: $model" >&2; exit 2; }
+done
+
+jobs=()
+for benchmark in "${BENCHMARKS[@]}"; do
+  jobs+=("vanilla|$VANILLA_MODEL|0.5|$benchmark")
+done
+for threshold in "${THRESHOLDS[@]}"; do
+  slug="${threshold/./}"
+  for benchmark in "${BENCHMARKS[@]}"; do
+    jobs+=("token_kv_head_t${slug}|$AHA_MODEL|$threshold|$benchmark")
+    jobs+=("token_t${slug}|$L2A_MODEL|$threshold|$benchmark")
+  done
+done
+
+launch() {
+  local gpu="$1" method="$2" model="$3" threshold="$4" benchmark="$5"
+  EVAL_ROOT="$EVAL_ROOT" RULER_LIMIT=20 BABILONG_LIMIT=50 \
+    AHA_GATE_HARD_THRESHOLD="$threshold" \
+    "$SCRIPT_DIR/eval_cell.sh" "$method" "$model" "$gpu" "$benchmark" &
+  LAST_PID="$!"
+}
+
+for i in "${!jobs[@]}"; do
+  slot=$((i % ${#GPUS[@]}))
+  if [[ -n "${SLOT_PIDS[$slot]:-}" ]]; then
+    wait "${SLOT_PIDS[$slot]}"
+  fi
+  IFS='|' read -r method model threshold benchmark <<< "${jobs[$i]}"
+  launch "${GPUS[$slot]}" "$method" "$model" "$threshold" "$benchmark"
+  SLOT_PIDS[$slot]="$LAST_PID"
+done
+
+status=0
+for pid in "${SLOT_PIDS[@]}"; do
+  wait "$pid" || status=1
+done
+[[ "$status" == 0 ]] || exit "$status"
+
+python "$SCRIPT_DIR/summarize_qwen1p7b_router_granularity_20260714.py" \
+  --repo "$REPO" \
+  --eval-root "$EVAL_ROOT" \
+  --baseline-root "$EVAL_ROOT/vanilla" \
+  --input-dir "$RECIPE/data/eval_inputs" \
+  --output-dir "$OUTPUT_ROOT/summary"
--- a/recipe/scripts/evaluate.sh
+++ b/recipe/scripts/evaluate.sh
@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
+IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
+GPU_COUNT="$(nvidia-smi -L | wc -l)"
+if [[ "$GPU_COUNT" -lt 1 ]]; then
+  echo "At least one NVIDIA GPU is required." >&2
+  exit 2
+fi
+GPU_LIST="$(seq -s, 0 $((GPU_COUNT - 1)))"
+
+docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
+docker run --rm --gpus all --ipc=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -e GPU_LIST="$GPU_LIST" \
+  -v "$REPO:/workspace" \
+  -w /workspace/recipe \
+  "$IMAGE" bash scripts/eval_sweep.sh
--- a/recipe/scripts/export_tuned_vanilla_aha_hotstart.py
+++ b/recipe/scripts/export_tuned_vanilla_aha_hotstart.py
@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""Convert a tuned vanilla Qwen3 checkpoint into a quality-safe AHA hot start.
+
+All vanilla weights are preserved. Dynamic gate weights start at zero and a
+trainable bias initializes every gate to the requested full-attention
+probability, so hard routing is initially exactly vanilla full attention.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import torch
+from torch import nn
+from transformers import AutoTokenizer
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from modeling_aha_qwen3 import AHAQwen3Config, AHAQwen3ForCausalLM
+
+
+def add_zero_bias(linear: nn.Linear) -> nn.Linear:
+    new = nn.Linear(
+        linear.in_features,
+        linear.out_features,
+        bias=True,
+        device=linear.weight.device,
+        dtype=linear.weight.dtype,
+    )
+    with torch.no_grad():
+        new.weight.copy_(linear.weight)
+        new.bias.zero_()
+    return new
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--vanilla-path", required=True)
+    parser.add_argument("--output-path", required=True)
+    parser.add_argument("--window-size", type=int, default=128)
+    parser.add_argument("--gate-init-full-prob", type=float, default=0.90)
+    parser.add_argument(
+        "--local-kind", choices=("sliding_window", "sink_recent"), default="sliding_window"
+    )
+    parser.add_argument(
+        "--router-granularity",
+        choices=("token", "token_kv_head"),
+        default="token_kv_head",
+        help="Native dynamic gate shape. token is one shared gate per layer/token.",
+    )
+    args = parser.parse_args()
+
+    if not 0.5 < args.gate_init_full_prob < 1.0:
+        raise ValueError("--gate-init-full-prob must be strictly between 0.5 and 1")
+    output = Path(args.output_path)
+    output.mkdir(parents=True, exist_ok=True)
+
+    AHAQwen3Config.register_for_auto_class()
+    AHAQwen3ForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    model = AHAQwen3ForCausalLM.from_pretrained_qwen3(
+        args.vanilla_path,
+        aha_window_size=args.window_size,
+        aha_local_kind=args.local_kind,
+        aha_router_granularity=args.router_granularity,
+        aha_mode="dynamic",
+        aha_gate_target=0.70,
+        aha_reg_weight=-1.0,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="sdpa",
+    )
+
+    # Qwen3 normally has bias-free projections. A gate bias gives a stable,
+    # token-independent full-attention hot start while zero biases on q/k/v/o
+    # keep the original attention computation bit-for-bit unchanged.
+    if not model.config.attention_bias:
+        for layer in model.model.layers:
+            attn = layer.self_attn
+            attn.q_proj = add_zero_bias(attn.q_proj)
+            attn.k_proj = add_zero_bias(attn.k_proj)
+            attn.v_proj = add_zero_bias(attn.v_proj)
+            attn.o_proj = add_zero_bias(attn.o_proj)
+        model.config.attention_bias = True
+
+    q_rows = model.config.num_attention_heads * model.config.head_dim
+    gate_logit = math.log(args.gate_init_full_prob / (1.0 - args.gate_init_full_prob))
+    with torch.no_grad():
+        for layer in model.model.layers:
+            q_proj = layer.self_attn.q_proj
+            q_proj.weight[q_rows:].zero_()
+            q_proj.bias[q_rows:].fill_(gate_logit)
+
+    # Keep a distinct LM head so the custom checkpoint has an explicit,
+    # self-contained state dict under safetensors.
+    model.config.tie_word_embeddings = False
+    if model.lm_head.weight.data_ptr() == model.model.embed_tokens.weight.data_ptr():
+        model.lm_head.weight = nn.Parameter(model.lm_head.weight.detach().clone())
+    model.config.aha_hotstart_source = str(Path(args.vanilla_path).resolve())
+    model.config.aha_gate_init_full_prob = args.gate_init_full_prob
+
+    tokenizer = AutoTokenizer.from_pretrained(args.vanilla_path, trust_remote_code=True)
+    model.save_pretrained(output, safe_serialization=True)
+    tokenizer.save_pretrained(output)
+    gate_rows = model.model.layers[0].self_attn.aha_router_outputs
+    gate_parameters = len(model.model.layers) * gate_rows * (
+        model.config.hidden_size + 1
+    )
+    (output / "aha_hotstart_manifest.json").write_text(
+        json.dumps(
+            {
+                "source_checkpoint": str(Path(args.vanilla_path).resolve()),
+                "router_granularity": args.router_granularity,
+                "native_gate_rows_per_layer": gate_rows,
+                "effective_gate_parameters": gate_parameters,
+                "gate_init_full_probability": args.gate_init_full_prob,
+                "gate_weight_init": "zeros",
+                "gate_bias_logit": gate_logit,
+                "local_attention": {
+                    "kind": args.local_kind,
+                    "sink_size": model.config.duo_sink_size,
+                    "recent_size": model.config.duo_recent_size,
+                },
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+    print(
+        f"saved={output} init_full_prob={args.gate_init_full_prob:.4f} "
+        f"gate_logit={gate_logit:.6f} hard_sparsity=0.0 "
+        f"router_granularity={args.router_granularity} gate_params={gate_parameters:,}"
+    )
+
+
+if __name__ == "__main__":
+    main()
--- a/recipe/scripts/reproduce.sh
+++ b/recipe/scripts/reproduce.sh
@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO="$(cd "$SCRIPT_DIR/../.." && pwd)"
+IMAGE="${IMAGE:-aha-l2a-qwen3-repro:torch2.9.1-cu128}"
+GPU_COUNT="$(nvidia-smi -L | wc -l)"
+if [[ "$GPU_COUNT" -lt 1 ]]; then
+  echo "At least one NVIDIA GPU is required." >&2
+  exit 2
+fi
+GPU_AHA="${GPU_AHA:-0}"
+if [[ -z "${GPU_L2A+x}" ]]; then
+  GPU_L2A=0
+  if [[ "$GPU_COUNT" -ge 2 ]]; then
+    GPU_L2A=1
+  fi
+fi
+
+for gpu in "$GPU_AHA" "$GPU_L2A"; do
+  if [[ ! "$gpu" =~ ^[0-9]+$ ]] || (( gpu >= GPU_COUNT )); then
+    echo "Invalid GPU index $gpu; nvidia-smi reports $GPU_COUNT visible GPU(s)." >&2
+    exit 2
+  fi
+done
+
+echo "Training mode: one GPU per experimental variant (no DDP/FSDP/DeepSpeed)."
+echo "AHA variant GPU: $GPU_AHA"
+echo "L2A-style variant GPU: $GPU_L2A"
+if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
+  echo "The two variants will run sequentially on the same GPU."
+else
+  echo "The two variants will run concurrently on two independent GPUs."
+fi
+if [[ "$GPU_COUNT" -gt 2 ]]; then
+  echo "Visible GPUs: $GPU_COUNT; this recipe intentionally does not use all GPUs."
+fi
+
+docker build -t "$IMAGE" -f "$REPO/recipe/Dockerfile" "$REPO/recipe"
+docker run --rm --gpus all --ipc=host \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  -e GPU_AHA="$GPU_AHA" -e GPU_L2A="$GPU_L2A" \
+  -v "$REPO:/workspace" \
+  -w /workspace/recipe \
+  "$IMAGE" bash scripts/train_both.sh
--- a/recipe/scripts/summarize_qwen1p7b_router_granularity_20260714.py
+++ b/recipe/scripts/summarize_qwen1p7b_router_granularity_20260714.py
@@ -0,0 +1,403 @@
+#!/usr/bin/env python3
+"""Audit and summarize the matched shared-gate vs KV-head-gate sweep."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import glob
+import hashlib
+import json
+import subprocess
+from collections import defaultdict
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+
+
+SUITES = ("RULER-local-full13", "BabiLong", "HELMET-ICL", "MRCR")
+THRESHOLDS = (0.45, 0.50, 0.525, 0.55, 0.575, 0.60, 0.625, 0.65)
+ARMS = ("token", "token_kv_head")
+SLUGS = {
+    "RULER-local-full13": "ruler",
+    "BabiLong": "babilong",
+    "HELMET-ICL": "helmet",
+    "MRCR": "mrcr",
+}
+SCORE_DEFINITIONS = {
+    "RULER-local-full13": "unweighted macro of lm-eval 8192 string-match scores over the local 13-config set",
+    "BabiLong": "unweighted macro exact-match accuracy over qa1-qa5",
+    "HELMET-ICL": "unweighted macro label exact-match over five ICL configurations",
+    "MRCR": "unweighted macro of prefix-check plus SequenceMatcher scores over 2/4/8-needle configurations",
+}
+
+
+def threshold_slug(value: float) -> str:
+    return {
+        0.45: "045",
+        0.50: "050",
+        0.525: "0525",
+        0.55: "055",
+        0.575: "0575",
+        0.60: "060",
+        0.625: "0625",
+        0.65: "065",
+    }[value]
+
+
+def one(pattern: str) -> Path:
+    paths = [Path(path) for path in glob.glob(pattern, recursive=True)]
+    if len(paths) != 1:
+        raise RuntimeError(f"expected one path for {pattern}, found {paths}")
+    return paths[0]
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
+
+
+def collect_scores(root: Path) -> dict[tuple[str, str], dict]:
+    rows: dict[tuple[str, str], dict] = {}
+    for split in ("a", "b"):
+        result = one(str(root / f"ruler8k_splits/{split}/lm_eval/**/results_*.json"))
+        payload = json.loads(result.read_text())
+        for config, metrics in payload["results"].items():
+            if config in payload["n-samples"]:
+                rows[("RULER-local-full13", config)] = {
+                    "score": float(metrics["8192,none"]),
+                    "n": int(payload["n-samples"][config]["effective"]),
+                }
+    result = one(str(root / "babilong8k_qa1_qa5_n50/**/results_*.json"))
+    payload = json.loads(result.read_text())
+    for config, metrics in payload["results"].items():
+        if config in payload["n-samples"]:
+            rows[("BabiLong", config)] = {
+                "score": float(metrics["acc,none"]),
+                "n": int(payload["n-samples"][config]["effective"]),
+            }
+    for suite, filename in (
+        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
+        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
+    ):
+        grouped: dict[str, list[dict]] = defaultdict(list)
+        for row in read_jsonl(root / filename):
+            grouped[row["config"]].append(row)
+        for config, values in grouped.items():
+            rows[(suite, config)] = {
+                "score": sum(float(row["score"]) for row in values) / len(values),
+                "n": len(values),
+            }
+    return rows
+
+
+def sample_signatures(root: Path, score_rows: dict[tuple[str, str], dict]) -> dict[str, str]:
+    signatures: dict[str, str] = {}
+    for suite, rel_dirs in (
+        ("RULER-local-full13", ("ruler8k_splits/a/lm_eval", "ruler8k_splits/b/lm_eval")),
+        ("BabiLong", ("babilong8k_qa1_qa5_n50",)),
+    ):
+        configs = [config for row_suite, config in score_rows if row_suite == suite]
+        for config in configs:
+            matches = []
+            for rel_dir in rel_dirs:
+                matches.extend(root.glob(f"{rel_dir}/**/samples_{config}_*.jsonl"))
+            if len(matches) != 1:
+                raise RuntimeError(f"expected one sample log for {suite}/{config}: {matches}")
+            values = [
+                (
+                    row.get("doc_id"),
+                    row.get("doc_hash"),
+                    row.get("prompt_hash"),
+                    row.get("target_hash"),
+                )
+                for row in read_jsonl(matches[0])
+            ]
+            signatures[f"{suite}/{config}"] = hashlib.sha256(
+                json.dumps(values, sort_keys=True).encode()
+            ).hexdigest()
+    for suite, filename in (
+        ("HELMET-ICL", "helmet_icl8k_n50.jsonl"),
+        ("MRCR", "mrcr_8k_2_4_8needle_n10.jsonl"),
+    ):
+        grouped: dict[str, list[tuple]] = defaultdict(list)
+        for row in read_jsonl(root / filename):
+            grouped[row["config"]].append(
+                (row["row_id"], row.get("prompt_sha256"), row.get("target"))
+            )
+        for config, values in grouped.items():
+            signatures[f"{suite}/{config}"] = hashlib.sha256(
+                json.dumps(values, sort_keys=True).encode()
+            ).hexdigest()
+    return signatures
+
+
+def read_sparsity(root: Path, suite: str) -> dict:
+    if suite == "RULER-local-full13":
+        parts = [
+            json.loads((root / f"sparsity_ruler_{split}.json").read_text())
+            for split in ("a", "b")
+        ]
+        sparse = sum(int(part["sparse_decisions"]) for part in parts)
+        total = sum(int(part["total_decisions"]) for part in parts)
+        granularities = {part.get("router_granularity") for part in parts}
+        native = sum(int(part.get("native_router_decisions", 0)) for part in parts)
+        phases: dict[str, dict[str, int]] = defaultdict(lambda: {"sparse": 0, "total": 0})
+        for part in parts:
+            for phase, values in part.get("by_phase", {}).items():
+                phases[phase]["sparse"] += int(values["sparse_decisions"])
+                phases[phase]["total"] += int(values["total_decisions"])
+    else:
+        part = json.loads((root / f"sparsity_{SLUGS[suite]}.json").read_text())
+        sparse = int(part["sparse_decisions"])
+        total = int(part["total_decisions"])
+        granularities = {part.get("router_granularity")}
+        native = int(part.get("native_router_decisions", 0))
+        phases = {
+            phase: {
+                "sparse": int(values["sparse_decisions"]),
+                "total": int(values["total_decisions"]),
+            }
+            for phase, values in part.get("by_phase", {}).items()
+        }
+    return {
+        "sparse_decisions": sparse,
+        "total_decisions": total,
+        "sparsity": sparse / total,
+        "full_attention_usage": 1.0 - sparse / total,
+        "router_granularities": sorted(value for value in granularities if value),
+        "native_router_decisions": native or None,
+        "by_phase": phases,
+    }
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def command_output(command: list[str], cwd: Path) -> str:
+    try:
+        return subprocess.run(
+            command, cwd=cwd, check=True, text=True, capture_output=True
+        ).stdout.strip()
+    except (OSError, subprocess.CalledProcessError):
+        return "unavailable in downloaded HF snapshot"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--repo", type=Path, default=Path("/data/sjm/AHA/AHA-Qwen3")
+    )
+    parser.add_argument(
+        "--eval-root",
+        type=Path,
+        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_router_granularity_20260714/eval"),
+    )
+    parser.add_argument(
+        "--baseline-root",
+        type=Path,
+        default=Path("/data/sjm/AHA/AHA-Qwen3/experiments/qwen3_1p7b_fourbench_8k_broad_20260713/vanilla"),
+    )
+    parser.add_argument("--output-dir", type=Path)
+    parser.add_argument("--input-dir", type=Path)
+    args = parser.parse_args()
+    output = args.output_dir or args.eval_root.parent / "summary"
+    output.mkdir(parents=True, exist_ok=True)
+
+    baseline = collect_scores(args.baseline_root)
+    baseline_signatures = sample_signatures(args.baseline_root, baseline)
+    curve_rows = []
+    per_config = []
+    points = []
+    alignment = {}
+    done_files = ("ruler_a.DONE", "ruler_b.DONE", "babilong.DONE", "helmet.DONE", "mrcr.DONE")
+
+    for arm in ARMS:
+        for threshold in THRESHOLDS:
+            method = f"{arm}_t{threshold_slug(threshold)}"
+            root = args.eval_root / method
+            missing = [name for name in done_files if not (root / name).exists()]
+            if missing:
+                raise RuntimeError(f"incomplete {method}: missing {missing}")
+            scores = collect_scores(root)
+            if set(scores) != set(baseline):
+                raise RuntimeError(f"configuration mismatch for {method}")
+            if any(scores[key]["n"] != baseline[key]["n"] for key in baseline):
+                raise RuntimeError(f"sample-count mismatch for {method}")
+            signatures = sample_signatures(root, scores)
+            mismatches = sorted(
+                key for key in baseline_signatures if signatures.get(key) != baseline_signatures[key]
+            )
+            alignment[method] = {"aligned": not mismatches, "mismatches": mismatches}
+            if mismatches:
+                raise RuntimeError(f"prompt/target hash mismatch for {method}: {mismatches}")
+
+            suite_rows = []
+            sparse_sum = total_sum = 0
+            for suite in SUITES:
+                keys = sorted(key for key in baseline if key[0] == suite)
+                baseline_score = sum(baseline[key]["score"] for key in keys) / len(keys)
+                score = sum(scores[key]["score"] for key in keys) / len(keys)
+                sparsity = read_sparsity(root, suite)
+                expected_granularity = arm
+                if sparsity["router_granularities"] != [expected_granularity]:
+                    raise RuntimeError(
+                        f"{method}/{suite} recorded {sparsity['router_granularities']}, "
+                        f"expected {[expected_granularity]}"
+                    )
+                sparse_sum += sparsity["sparse_decisions"]
+                total_sum += sparsity["total_decisions"]
+                row = {
+                    "arm": arm,
+                    "router_granularity": arm,
+                    "threshold": threshold,
+                    "benchmark": suite,
+                    "config_count": len(keys),
+                    "samples_per_config": ",".join(
+                        str(value) for value in sorted({baseline[key]["n"] for key in keys})
+                    ),
+                    "total_samples": sum(baseline[key]["n"] for key in keys),
+                    "tuned_vanilla_score": baseline_score,
+                    "quality_floor_95pct": 0.95 * baseline_score,
+                    "score": score,
+                    "retention": score / baseline_score if baseline_score else None,
+                    "quality_pass": score >= 0.95 * baseline_score,
+                    **{key: sparsity[key] for key in (
+                        "sparse_decisions", "total_decisions", "sparsity",
+                        "full_attention_usage", "native_router_decisions", "by_phase"
+                    )},
+                }
+                curve_rows.append(row)
+                suite_rows.append(row)
+                for key in keys:
+                    per_config.append({
+                        "arm": arm,
+                        "threshold": threshold,
+                        "benchmark": suite,
+                        "config": key[1],
+                        "n": baseline[key]["n"],
+                        "tuned_vanilla_score": baseline[key]["score"],
+                        "score": scores[key]["score"],
+                    })
+            points.append({
+                "arm": arm,
+                "threshold": threshold,
+                "all_suite_quality_pass": all(row["quality_pass"] for row in suite_rows),
+                "decision_weighted_sparsity": sparse_sum / total_sum,
+                "full_attention_usage": 1.0 - sparse_sum / total_sum,
+                "sparse_decisions": sparse_sum,
+                "total_decisions": total_sum,
+                "suites": suite_rows,
+            })
+
+    headline = {}
+    for arm in ARMS:
+        valid = [
+            point for point in points
+            if point["arm"] == arm and point["all_suite_quality_pass"]
+        ]
+        if valid:
+            best = max(valid, key=lambda point: point["decision_weighted_sparsity"])
+            headline[arm] = {"found": True, **best}
+        else:
+            headline[arm] = {
+                "found": False,
+                "statement": "No measured threshold preserved at least 95% of tuned vanilla on every suite.",
+            }
+
+    input_dir = args.input_dir or args.baseline_root.parent / "inputs"
+    frozen_inputs = [
+        input_dir / "helmet_icl_8k_n50_per_config.jsonl",
+        input_dir / "mrcr_8k_2_4_8needle_n10_per_config.jsonl",
+    ]
+    payload = {
+        "protocol": {
+            "model": "Qwen3-1.7B tuned vanilla",
+            "context_length": 8192,
+            "arms": {
+                "token": "L2A-style shared-gate: one native gate per token/layer",
+                "token_kv_head": "AHA: one native gate per token/KV-head/layer",
+            },
+            "local_attention": {"sink_tokens": 64, "recent_tokens": 256},
+            "inference": "strict AHA routing in both prefill and decode; no force-full heads or full-decode fallback",
+            "thresholds": list(THRESHOLDS),
+            "quality_rule": "every suite macro score >= 95% of aligned tuned-vanilla macro score",
+            "headline_rule": "highest decision-weighted measured sparsity among thresholds passing every suite",
+            "sparsity_definition": "hard local routes / token x KV-head x layer effective decisions",
+            "suite_scope": "RULER local full13 means this repository's fixed 13-config set, not complete upstream RULER",
+        },
+        "score_definitions": SCORE_DEFINITIONS,
+        "headline": headline,
+        "points": points,
+        "curves": curve_rows,
+        "per_config": per_config,
+        "alignment": {"all_aligned": all(value["aligned"] for value in alignment.values()), "methods": alignment},
+        "appendix_note": (
+            "Upstream L2A results are not merged into this table because architecture, "
+            "training data, objective, and evaluation protocols differ."
+        ),
+        "reproducibility": {
+            "git_commit": command_output(["git", "rev-parse", "HEAD"], args.repo),
+            "git_status_short": command_output(["git", "status", "--short"], args.repo),
+            "container": "dmtd-repro",
+            "container_image": command_output(
+                ["docker", "inspect", "-f", "{{.Config.Image}}", "dmtd-repro"], args.repo
+            ),
+            "baseline_root": str(args.baseline_root),
+            "eval_root": str(args.eval_root),
+            "frozen_input_sha256": {
+                str(path): sha256(path) for path in frozen_inputs if path.exists()
+            },
+        },
+    }
+    (output / "router_granularity_results.json").write_text(
+        json.dumps(payload, indent=2) + "\n"
+    )
+
+    flat_curve_rows = [
+        {key: value for key, value in row.items() if key != "by_phase"}
+        for row in curve_rows
+    ]
+    with (output / "quality_sparsity_curves.csv").open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(flat_curve_rows[0]))
+        writer.writeheader()
+        writer.writerows(flat_curve_rows)
+    with (output / "per_config_scores.csv").open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(per_config[0]))
+        writer.writeheader()
+        writer.writerows(per_config)
+
+    figure, axes = plt.subplots(2, 2, figsize=(12, 9), constrained_layout=True)
+    colors = {"token": "#d62728", "token_kv_head": "#1f77b4"}
+    labels = {"token": "L2A-style shared-gate", "token_kv_head": "AHA KV-head gate"}
+    for axis, suite in zip(axes.flat, SUITES):
+        for arm in ARMS:
+            values = sorted(
+                (row for row in curve_rows if row["arm"] == arm and row["benchmark"] == suite),
+                key=lambda row: row["sparsity"],
+            )
+            axis.plot(
+                [100 * row["sparsity"] for row in values],
+                [100 * row["retention"] for row in values],
+                marker="o", color=colors[arm], label=labels[arm],
+            )
+            for row in values:
+                axis.annotate(f"{row['threshold']:.3g}", (100 * row["sparsity"], 100 * row["retention"]), fontsize=7)
+        axis.axhline(95, color="black", linestyle="--", linewidth=1)
+        axis.set_title(suite)
+        axis.set_xlabel("Measured effective sparsity (%)")
+        axis.set_ylabel("Retention vs tuned vanilla (%)")
+        axis.grid(alpha=0.25)
+    axes.flat[0].legend()
+    figure.suptitle("Qwen3-1.7B matched router-granularity quality–sparsity curves")
+    figure.savefig(output / "quality_sparsity_curves.png", dpi=180)
+    plt.close(figure)
+    print(json.dumps(headline, indent=2))
+
+
+if __name__ == "__main__":
+    main()
--- a/recipe/scripts/train_aha.sh
+++ b/recipe/scripts/train_aha.sh
@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "$SCRIPT_DIR/train_one.sh" aha "${1:-0}"
--- a/recipe/scripts/train_both.sh
+++ b/recipe/scripts/train_both.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+GPU_AHA="${GPU_AHA:-0}"
+GPU_L2A="${GPU_L2A:-1}"
+cd "$RECIPE"
+
+echo "AHA = per-token, per-KV-head gate; assigned visible GPU $GPU_AHA."
+echo "L2A-style = per-token, head-shared gate; assigned visible GPU $GPU_L2A."
+echo "Each variant is a separate world-size-1 run with global batch size 1."
+
+python "$SCRIPT_DIR/validate_release.py"
+python -m unittest -v tests.test_router_granularity
+
+if [[ "$GPU_AHA" == "$GPU_L2A" ]]; then
+  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA"
+  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A"
+else
+  "$SCRIPT_DIR/train_one.sh" aha "$GPU_AHA" & aha_pid=$!
+  "$SCRIPT_DIR/train_one.sh" l2a_style "$GPU_L2A" & l2a_pid=$!
+  status=0
+  wait "$aha_pid" || status=1
+  wait "$l2a_pid" || status=1
+  [[ "$status" == 0 ]] || exit "$status"
+fi
+
+echo "AHA selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/aha/stage2/checkpoint-25"
+echo "L2A-style selected checkpoint: ${OUTPUT_ROOT:-$RECIPE/../outputs}/l2a_style/stage2/checkpoint-25"
--- a/recipe/scripts/train_l2a_style.sh
+++ b/recipe/scripts/train_l2a_style.sh
@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "$SCRIPT_DIR/train_one.sh" l2a_style "${1:-0}"
--- a/recipe/scripts/train_one.sh
+++ b/recipe/scripts/train_one.sh
@@ -0,0 +1,127 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -lt 1 || $# -gt 2 ]]; then
+  echo "Usage: $0 aha|l2a_style [GPU_ID]" >&2
+  exit 2
+fi
+
+ARM="$1"
+GPU="${2:-0}"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RECIPE="$(cd "$SCRIPT_DIR/.." && pwd)"
+REPO="$(cd "$RECIPE/.." && pwd)"
+VANILLA_DIR="${VANILLA_DIR:-$REPO}"
+DATA_DIR="${DATA_DIR:-$RECIPE/data/am_distilled_long_mix}"
+OUTPUT_ROOT="${OUTPUT_ROOT:-$REPO/outputs}"
+
+case "$ARM" in
+  aha) GRANULARITY=token_kv_head ;;
+  l2a_style) GRANULARITY=token ;;
+  *) echo "Unknown arm: $ARM (expected aha or l2a_style)" >&2; exit 2 ;;
+esac
+
+ARM_OUT="$OUTPUT_ROOT/$ARM"
+HOTSTART="$ARM_OUT/hotstart"
+STAGE1="$ARM_OUT/stage1"
+STAGE2="$ARM_OUT/stage2"
+LOG_DIR="$ARM_OUT/logs"
+mkdir -p "$LOG_DIR"
+
+export CUDA_VISIBLE_DEVICES="$GPU"
+export PYTHONUNBUFFERED=1
+export TOKENIZERS_PARALLELISM=false
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+
+if [[ ! -f "$VANILLA_DIR/model.safetensors" ]]; then
+  echo "Missing tuned-vanilla checkpoint: $VANILLA_DIR/model.safetensors" >&2
+  exit 2
+fi
+
+if [[ ! -f "$HOTSTART/config.json" ]]; then
+  python "$SCRIPT_DIR/export_tuned_vanilla_aha_hotstart.py" \
+    --vanilla-path "$VANILLA_DIR" \
+    --output-path "$HOTSTART" \
+    --window-size 128 \
+    --local-kind sink_recent \
+    --gate-init-full-prob 0.90 \
+    --router-granularity "$GRANULARITY" \
+    2>&1 | tee "$LOG_DIR/hotstart.log"
+fi
+
+if [[ ! -f "$STAGE1/checkpoint-300/config.json" ]]; then
+  AHA_TRAIN_GATE_HARD_THRESHOLD=0.50 \
+  python "$RECIPE/dynamic_duo_train.py" \
+    --aha_checkpoint "$HOTSTART" \
+    --model_path "$VANILLA_DIR" \
+    --output_dir "$STAGE1" \
+    --data_source am_distilled \
+    --am_dataset_path "$DATA_DIR" \
+    --am_dataset_split train \
+    --am_label_mode full \
+    --max_length 8192 \
+    --num_steps 300 \
+    --warmup_ratio 0.10 \
+    --lr 3e-5 \
+    --reg_weight 0.1 \
+    --ce_weight 0.0 \
+    --batch_size 1 \
+    --grad_accum 1 \
+    --save_steps 100 \
+    --log_steps 10 \
+    --seed 42 \
+    --dtype bfloat16 \
+    --attn_impl sdpa \
+    --aha_local_kind sink_recent \
+    --router_granularity "$GRANULARITY" \
+    2>&1 | tee "$LOG_DIR/stage1.log"
+fi
+
+if [[ ! -f "$STAGE2/checkpoint-75/config.json" ]]; then
+  MODEL_PATH="$VANILLA_DIR" \
+  AHA_CHECKPOINT_PATH="$STAGE1/checkpoint-300" \
+  DATASET_PATH="$DATA_DIR" \
+  DATASET_SPLIT=train \
+  OUTPUT_DIR="$STAGE2" \
+  AHA_MODE=dynamic \
+  AHA_ROUTER_GRANULARITY="$GRANULARITY" \
+  AHA_LOCAL_KIND=sink_recent \
+  AHA_CE_WEIGHT=1.0 \
+  AHA_DISTILL_WEIGHT=0.5 \
+  AHA_REG_WEIGHT=0.01 \
+  AHA_TRAIN_GATE_HARD_THRESHOLD=0.58 \
+  GROUPED_LR=1 \
+  GATE_ONLY=0 \
+  LEARNING_RATE=3e-6 \
+  GATE_LEARNING_RATE=3e-6 \
+  BACKBONE_LEARNING_RATE=3e-7 \
+  LR_SCHEDULER_TYPE=constant_with_warmup \
+  WARMUP_RATIO=0.10 \
+  WEIGHT_DECAY=0.0 \
+  MAX_SEQ_LENGTH=8192 \
+  PER_DEVICE_TRAIN_BATCH_SIZE=1 \
+  GRADIENT_ACCUMULATION_STEPS=1 \
+  MAX_STEPS=75 \
+  LOGGING_STEPS=5 \
+  SAVE_STEPS=25 \
+  SAVE_TOTAL_LIMIT=3 \
+  FREEZE_EMBEDDINGS_LM_HEAD=1 \
+  REPORT_TO=none \
+  SEED=47 \
+  python "$RECIPE/sft.py" \
+    2>&1 | tee "$LOG_DIR/stage2.log"
+fi
+
+python - "$STAGE2/checkpoint-25" "$GRANULARITY" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+checkpoint = Path(sys.argv[1])
+expected = sys.argv[2]
+config = json.loads((checkpoint / "config.json").read_text())
+actual = config.get("aha_router_granularity", "token_kv_head")
+if actual != expected:
+    raise RuntimeError(f"checkpoint granularity {actual!r} != {expected!r}")
+print(f"selected_checkpoint={checkpoint} router_granularity={actual}")
+PY
--- a/recipe/scripts/validate_release.py
+++ b/recipe/scripts/validate_release.py
@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Fail fast if a downloaded release is incomplete or has drifted."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from collections import Counter
+from pathlib import Path
+
+from datasets import load_from_disk
+
+
+RECIPE = Path(__file__).resolve().parents[1]
+REPO = RECIPE.parent
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
+
+
+def main() -> None:
+    manifest = json.loads((RECIPE / "manifest.json").read_text())
+    checked = []
+    for relative, expected in manifest["files"].items():
+        path = (RECIPE / relative).resolve()
+        if not path.exists():
+            raise FileNotFoundError(path)
+        if path.name == "model.safetensors" and os.environ.get("SKIP_LARGE_HASH") == "1":
+            continue
+        actual = sha256(path)
+        if actual != expected:
+            raise RuntimeError(f"SHA256 mismatch for {path}: {actual} != {expected}")
+        checked.append(str(path.relative_to(REPO)))
+
+    config = json.loads((REPO / "config.json").read_text())
+    if config.get("model_type") != "qwen3":
+        raise RuntimeError(f"unexpected tuned-vanilla model_type: {config.get('model_type')}")
+    dataset = load_from_disk(str(RECIPE / "data/am_distilled_long_mix"))
+    if len(dataset["train"]) != 1024:
+        raise RuntimeError(f"expected 1024 training rows, found {len(dataset['train'])}")
+
+    helmet = read_jsonl(RECIPE / "data/eval_inputs/helmet_icl_8k_n50_per_config.jsonl")
+    mrcr = read_jsonl(RECIPE / "data/eval_inputs/mrcr_8k_2_4_8needle_n10_per_config.jsonl")
+    helmet_counts = Counter(row["config"] for row in helmet)
+    mrcr_counts = Counter(row["config"] for row in mrcr)
+    if sorted(helmet_counts.values()) != [50] * 5:
+        raise RuntimeError(f"unexpected HELMET counts: {helmet_counts}")
+    if sorted(mrcr_counts.values()) != [10] * 3:
+        raise RuntimeError(f"unexpected MRCR counts: {mrcr_counts}")
+
+    print(json.dumps({
+        "status": "ok",
+        "checked_sha256": checked,
+        "model_type": config["model_type"],
+        "training_rows": len(dataset["train"]),
+        "helmet_rows": len(helmet),
+        "mrcr_rows": len(mrcr),
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()