NCP_ArchPreview_dolmo3_8.9B_Stage2_DFlash2_NCPFlash

This repository contains NCPFlash, an NCP-conditioned block-parallel Drafter developed for NCP_ArchPreview_dolmo3_8.9B_Stage2. It proposes a block of 16 draft tokens per forward pass and is intended for lossless speculative decoding with the paired Target model.

The Drafter is not a standalone language model. It reuses the Target tokenizer, token embedding, output head, intermediate hidden representations, and causal NCP concept representation at inference time.

Model architecture

The model uses a five-layer, 1.105B-parameter BF16 Drafter with hidden size 4096, 32 attention heads, and SwiGLU intermediate size 11008. Its design combines:

• DFlash2-style two-tap dynamic convolution and top-16 path selection;
• layer-wise fusion of five Target hidden representations from decoder layers [1, 4, 7, 10, 13];
• a causal residual derived from the NCP representation of the last fully completed Target chunk;
• a proposal horizon of 16 draft tokens.

The causal NCP residual adds 40,960 parameters (approximately 0.04M) to the Drafter and does not introduce additional Target-model layers. The five Drafter layers are initialized from Target decoder layers [0, 4, 8, 12, 15].

The checkpoint intentionally does not duplicate the Target token embedding or output head. It must therefore be paired with the exact Target model and tokenizer named above.

Training

The Drafter was trained with online Target distillation on a fixed 5B-token subset, repeated for 10 epochs. Training used:

• sequence length: 8192;
• training anchors per sequence: 512;
• global batch size: 512;
• proposal horizon: 16 draft tokens;
• CE, hidden-representation distillation, and path-selector objectives;
• FP32 optimizer master parameters with BF16 model parameters;
• checkpoint: step 11,920.

The training data are fixed corpus sequences; online means that Target representations are computed during training rather than read from a precomputed cache. This is not on-policy data generation.

Evaluation

We evaluate with exact lossless speculative verification using the DeepSpec / DSpark round-weighted definition of mean accepted length (MAL):

[ \operatorname{MAL} = \frac{\sum_r c_r}{R}, ]

where (c_r) is the number of tokens committed in verification round (r), including the Target correction or bonus token when present. The initial Target token generated during prefill is excluded.

Results below use the same Target model, task prompts, sampling configuration, BF16 Target route, proposal horizon, and matched non-degenerate evaluation set for both Drafters.

Benchmark	Baseline	NCPFlash	Relative gain
GSM8K	6.351	6.537	+2.93%
MATH	6.105	6.240	+2.22%
HumanEval	5.432	5.845	+7.59%
MBPP	5.844	6.099	+4.37%
Macro average	5.933	6.180	+4.17%

The matched set excludes generations for which either compared model enters a repeated-output loop or reaches the generation-length limit without a valid stopping condition. These numbers are intended for matched Drafter comparison; they should not be compared directly with results obtained using a different Target model, prompt template, sampling protocol, proposal-horizon convention, or MAL aggregation rule.

Speculative decoding preserves the Target distribution under exact verification. Task accuracy is therefore a correctness check rather than a measure of Drafter quality.

Loading the checkpoint

The architecture is packaged as Transformers remote code:

import torch
from transformers import AutoModel

model_id = (
    "ArchSpace-Collection/"
    "NCP_ArchPreview_dolmo3_8.9B_Stage2_DFlash2_NCPFlash"
)

drafter = AutoModel.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).eval().cuda()

Loading the module alone does not provide text generation. A complete speculative-decoding runtime must also load the paired Target model, capture the five Target hidden representations and the causal NCP representation, and apply exact Target verification to every proposed block.

The Drafter forward interface consumes:

• aux_hidden_states: [batch, sequence, 5, 4096];
• anchor_embeddings: [batch, anchors, 4096];
• mask_embedding: [4096];
• anchor_positions: [batch, anchors];
• sequence_lengths: [batch];
• hlm_hidden_states: [batch, anchors, 4096].

Use trust_remote_code=True only after reviewing the included configuration_dflash.py and modeling_dflash.py files.