dungnv/qwen36-27b-length-traces

Qwen3.6-27B generation-length prediction: heads, calibrations and workloads

Artifacts for conformal length-aware LLM scheduling on Qwen/Qwen3.6-27B — predicting a request's remaining generation length from a hidden layer during decoding, wrapping it in a split-conformal interval, and scheduling with SRPT inside vLLM. Extends TRAIL (*Don't Stop Me Now*, ICLR'25) to a hybrid-attention reasoning model.

This repo contains the derived artifacts, not the raw activations. The 3250 decode-step traces are ~322 GB and are not uploaded. Everything needed to use the predictors is here; everything needed to retrain them is reproducible with the commands below (~20 GPU-hours).

Setup

Qwen3.6-27B is hybrid-attention: 64 layers with full_attention_interval=4, so only 16 carry a token-proportional KV cache and 48 hold fixed-size recurrent state. Consequences: KV is 64 KiB/token (vs 192 for a comparable dense 14B), but each request commits ~144 MiB of recurrent state at admission — equal to ~2300 tokens of KV. Per-request memory therefore barely grows with age until well past the median generation.

Files

Heads — all trained on the same 1000 base traces, so they differ only in target/readout. Held-out metrics on 300 base traces (mean T = 4625):

trail_mlp.pt is the faithful TRAIL baseline: the paper's Linear(h→512)→ReLU→Linear(512→bins) with 80 bins over 32768. trail.pt keeps the original 20-bins-over-8000 ceiling, which saturates on 7.5% of the short workload (34% of MATH) and understates the baseline — prefer trail_mlp.pt.

Calibrations (split-conformal, α=0.5, q_t pooled per decode step):

⚠️ `q_by_step[0]` is NaN in every file. The collector records only true decode steps (nst == 1), so the first token — produced by the prefill forward — is never captured and step 0 has no residuals to take a quantile over. Every waiting request has g = 0, so a scheduler that indexes q_by_step[g] without a guard gets NaN, and NaN compares False against everything: sorted() silently leaves the admission queue in arbitrary order. Note beta * q does not save you — 0.0 * nan = nan. Fall back to the nearest finite step.

⚠️ The augmented calibrations are not exchangeable with base traffic. Olympiad's median true_len is 13722 vs base's 1585, so its residuals inflate q_t by +20% (abs) / +38% (prog), and the extra support between step 17185 and 28750 comes almost entirely from Olympiad traces. Use *_base_* when serving base-distribution traffic; the augmented files are provided for the M0-vs-exchangeability comparison, not as a drop-in improvement.

Splits — split_base.json (train 1000 / calib 495 / eval 300, base only), split_aug.json (same train, calib 731 = 495 base + 236 Olympiad, eval 300 Olympiad). manifest_L32.json maps every trace path to (T, dataset).

Workloads — served request files with text, true_len, pred_len, dataset:

pred_len comes from ψ (prompt-only prefill predictor) and is never a copy of true_len — 0.0% identical in both. It is packed into the vLLM request id as the static admission priority, so setting it to true_len would hand the oracle key to every predictor policy.

Prompts — collect_prompts.json (1800: gsm8k 500, mmlupro 500, bbh 400, math 400), `benchmarkprompts.json (850, **verified zero overlap** with the collection pool), olympiad600prompts.jsonl` (OlympiadBench `OETO*enCOMP`, text-only English), `longcot600prompts.jsonl` (LongHorizonReasoning/longcot, 5 domains × 3 difficulties).

Prefill features — prefill_{collect,bench,olymp}.pt, mean-pooled layer-32 prompt states, one row per prompt, the input to ψ.

Layer sweep — layer_sweep_qwen36.json (all 100 traces) and ..._uncapped.json (59 uncapped).

Choosing the probe layer

8 layers swept, 4 full_attention / 4 linear_attention, depth 0.17–0.75, ranked by held-out cross-request Spearman over 5 resampled splits:

On clean data L32 sits on a plateau with {32, 35, 43} — the choice within it is not load-bearing. The depth-ratio heuristic (0.5625 → L36) would have been meaningfully worse; its nearest swept neighbour L40 ranks 6th of 8. Layer type mattered little (linear 0.663 vs full 0.644).

Trace schema (for reproduction)

Each .pt is {"L32": float16 (n_steps, 5120), "T": int, "g": int32 (n_steps,), "dataset": str, "prompt": str}.

• —`g` is the 0-based index of the token an activation produced. Remaining is T − (g+1). Recording starts at g=1 because the first token comes from the prefill forward, which is skipped.
• —A trace is not a contiguous step sequence — always index through g, never arange(len(H)).
• —Capped traces are mislabelled: a generation that never emits </think> stores T = cap. Base capped 1% (5/500); Olympiad ~22%. Drop them from calibration — the residual |R_true − R̂| uses R_true = T − t, so a capped trace injects a too-small residual and makes q_t overconfident exactly in the tail.

Reproduction

export SCHED_MODEL=Qwen/Qwen3.6-27B SCHED_LAYER=32
python collect_activations_vllm.py --layers 32 --prompt_file collected/collect_prompts.json \
    --max_tokens 32768 --max_model_len 40960 --gpu_mem 0.85 --chunk 100 --out <traces>
python build_split.py --init --base-glob '<traces>/*.pt' --sizes 1000,500,300 --out split_base.json
python train_head.py --head absolute --arch linear --split split_base.json --out abs.pt
python calibrate.py --ckpt abs.pt --split split_base.json --alphas 0.5 --out conf_abs_base_a5.npz

Collection is ~20 GPU-hours on one 80–96 GB card. enforce_eager=True is mandatory — CUDA graphs bypass forward hooks, so the collector silently records nothing without it.

Limitations

• —Raw activations are not included (~322 GB).
• —Metrics are single-seed; head training variance was not measured.
• —The layer sweep used max_tokens 2500, where 41% of traces capped; the full collection used 32768, where only 1% of base traces cap.
• —Scheduling results are not included in this repo and remain under evaluation. One measured property worth stating: with max_num_seqs=64 this workload never saturates the concurrency cap (~41 of 64 concurrent), so no queue forms and scheduling cannot help; a queue requires max_num_seqs=24. Preemption was exactly zero across ~100k schedule calls, so TRAIL's limited-preemption parameter c is inert here.