feat: add Ascend NPU support for VP-Drafter

configs/qwen3-8b-dta.json

@@ -0,0 +1,47 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [1, 9, 17, 25, 33],
+    "training_mode": "vp_drafter",
+    "prefix_weight_base": 0.9
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 36,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

configs/qwen3.5-4b-dta.json

@@ -0,0 +1,49 @@
+{
+    "architectures": [
+      "DFlashDraftModel"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+      "AutoModel": "dflash.DFlashDraftModel"
+    },
+    "block_size": 16,
+    "bos_token_id": 248043,
+    "dflash_config": {
+      "mask_token_id": 248070,
+      "target_layer_ids": [1, 8, 15, 22, 29],
+      "training_mode": "vp_drafter",
+      "prefix_weight_base": 0.9
+    },
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "initializer_range": 0.02,
+    "intermediate_size": 9728,
+    "layer_types": [
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "max_window_layers": 5,
+    "model_type": "qwen3",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 5,
+    "num_key_value_heads": 8,
+    "num_target_layers": 32,
+    "pad_token_id": 248044,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 10000000,
+    "sliding_window": null,
+    "tie_word_embeddings": true,
+    "transformers_version": "4.57.1",
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 248320
+  }

examples/run_qwen3.5_4b_dta_online_npu.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+# VP-Drafter (D2SD) training for Qwen3.5-4B on Ascend NPU
+# Backend: HF + SDPA attention + HCCL distributed
+#
+# Required environment variables (override before invoking the script):
+#   TARGET_MODEL_PATH  Path to Qwen3.5-4B weights
+#   TRAIN_DATA_PATH    Path to the training jsonl
+
+set -eu
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname "$SCRIPT_DIR")
+
+: "${TARGET_MODEL_PATH:?Set TARGET_MODEL_PATH to the Qwen3.5-4B weights directory}"
+: "${TRAIN_DATA_PATH:?Set TRAIN_DATA_PATH to the training jsonl file}"
+
+NPU_DEVICES=${1:-0,1,2,3,4,5,6,7}
+NUM_DEVICES=$(echo "$NPU_DEVICES" | tr ',' '\n' | wc -l)
+
+export ASCEND_RT_VISIBLE_DEVICES=$NPU_DEVICES
+export PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:32
+
+torchrun \
+    --standalone \
+    --nproc_per_node "$NUM_DEVICES" \
+    "$ROOT_DIR/scripts/train_dflash.py" \
+    --target-model-path "$TARGET_MODEL_PATH" \
+    --target-model-backend hf \
+    --draft-config-path "$ROOT_DIR/configs/qwen3.5-4b-dta.json" \
+    --train-data-path "$TRAIN_DATA_PATH" \
+    --output-dir "$ROOT_DIR/outputs/qwen3.5-4b-dta-npu" \
+    --num-epochs 6 \
+    --batch-size 1 \
+    --accumulation-steps 4 \
+    --learning-rate 6e-4 \
+    --warmup-ratio 0.04 \
+    --max-grad-norm 1.0 \
+    --max-length 1024 \
+    --chat-template qwen3.5 \
+    --attention-backend sdpa \
+    --num-anchors 16 \
+    --loss-decay-gamma 7.0 \
+    --log-interval 50 \
+    --save-interval 3000 \
+    --report-to tensorboard \
+    --block-size 16 \
+    --embedding-key model.language_model.embed_tokens.weight \
+    --trust-remote-code

examples/run_qwen3_8b_dta_online.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname $SCRIPT_DIR)
+export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
+export SPECFORGE_DATA_NUM_PROC=32
+NUM_GPUS=${1:-8}
+
+ATTENTION_BACKEND=${2:-flex_attention}
+
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
+    $ROOT_DIR/scripts/train_dflash.py \
+    --target-model-path Qwen/Qwen3-8B \
+    --target-model-backend sglang \
+    --draft-config-path $ROOT_DIR/configs/qwen3-8b-dta.json \
+    --train-data-path $ROOT_DIR/cache/dataset/perfectblend_qwen3-8b_regen.jsonl \
+    --output-dir $ROOT_DIR/outputs/qwen3-8b-dta-perfectblend \
+    --num-epochs 6 \
+    --batch-size 4 \
+    --learning-rate 6e-4 \
+    --warmup-ratio 0.04 \
+    --max-grad-norm 1.0 \
+    --max-length 3072 \
+    --chat-template qwen \
+    --attention-backend $ATTENTION_BACKEND \
+    --loss-decay-gamma 7.0 \
+    --log-interval 50 \
+    --save-interval 1000 \
+    --report-to wandb \
+    --wandb-project specforge-qwen3-8b-dta \
+    --block-size 16 \
+    --num-anchors 512 \
+    --wandb-name qwen3-8b-dta-perfectblend

scripts/train_dflash.py

@@ -92,21 +92,34 @@ def parse_args():
     model_group.add_argument(
         "--loss-type",
         type=str,
-        default="dflash",
+        default=None,
         choices=[
             "dflash",
+            "vp_drafter",
             "dpace",
             "dpace-cumulative-confidence-only",
             "dpace-continuation-value-only",
         ],
-        help=("Loss variant. Use dpace for Dynamic Position-Aware Cross-Entropy."),
+        help=(
+            "Training objective. If omitted, reads dflash_config.training_mode or "
+            "dflash_config.loss_type from the draft config, defaulting to dflash."
+        ),
     )
     model_group.add_argument(
         "--dpace-alpha",
         type=float,
         default=0.5,
         help="Smoothing alpha for D-PACE position weights.",
     )
+    model_group.add_argument(
+        "--prefix-weight-base",
+        type=float,
+        default=None,
+        help=(
+            "VP-Drafter prefix length sampling base. Values below 1 prefer shorter "
+            "visible prefixes; defaults to dflash_config.prefix_weight_base or 0.9."
+        ),
+    )
     model_group.add_argument(
         "--embedding-key",
         type=str,
@@ -218,8 +231,20 @@ def build_models(args) -> Tuple[DFlashTargetModel, DFlashDraftModel]:
     if not hasattr(draft_config, "dflash_config") or draft_config.dflash_config is None:
         draft_config.dflash_config = {}
 
+    args.loss_type = (
+        args.loss_type
+        or draft_config.dflash_config.get("training_mode")
+        or draft_config.dflash_config.get("loss_type")
+        or "dflash"
+    )
+    if args.prefix_weight_base is None:
+        args.prefix_weight_base = draft_config.dflash_config.get(
+            "prefix_weight_base", 0.9
+        )
+
     draft_config._attn_implementation = args.attention_backend
     print_on_rank0(f"Using attention backend: {args.attention_backend}")
+    print_on_rank0(f"Using DFlash training loss_type: {args.loss_type}")
 
     draft_model = DFlashDraftModel(draft_config).to(device=device, dtype=torch.bfloat16)
 
@@ -453,6 +478,8 @@ def main():
     print_on_rank0(f"Total training steps: {total_steps}")
 
     print_on_rank0("Loading target embeddings and head...")
+    device = get_local_device()
+    device_type = device.type
     target_components = TargetEmbeddingsAndHead.from_pretrained(
         args.target_model_path,
         embed_key=args.embedding_key,
@@ -472,6 +499,7 @@ def main():
         loss_decay_gamma=args.loss_decay_gamma,
         loss_type=args.loss_type,
         dpace_alpha=args.dpace_alpha,
+        prefix_weight_base=args.prefix_weight_base,
     )
 
     # Wrap each transformer block as its own FSDP unit so that all-gather /