feat(dspark): add Ascend NPU support for Qwen3.5-4B DSpark training

configs/qwen3-8b-dspark.json

@@ -0,0 +1,49 @@
+{
+  "architectures": [
+    "DSparkDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dspark.DSparkDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [1, 9, 17, 25, 33]
+  },
+  "dtype": "bfloat16",
+  "enable_confidence_head": true,
+  "confidence_head_with_markov": true,
+  "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"
+  ],
+  "markov_head_type": "vanilla",
+  "markov_rank": 256,
+  "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-dspark.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "DSparkDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dspark.DSparkDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 248043,
+  "dflash_config": {
+    "mask_token_id": 248070,
+    "target_layer_ids": [1, 8, 15, 22, 29]
+  },
+  "dtype": "bfloat16",
+  "enable_confidence_head": true,
+  "confidence_head_with_markov": true,
+  "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"
+  ],
+  "markov_head_type": "vanilla",
+  "markov_rank": 256,
+  "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,
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 248320
+}

examples/run_qwen3.5_4b_dspark_online_npu.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# DSpark online training for Qwen3.5-4B on Ascend NPU.
+#
+# DSpark = DFlash block-diffusion drafter + EAGLE-style Markov & confidence heads,
+# trained with cross-entropy + L1 distribution distillation + confidence BCE.
+# The L1 / confidence losses need the target model's FINAL hidden state; the HF
+# backend (used below) always surfaces it on 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
+export SPECFORGE_DATA_NUM_PROC=32
+
+torchrun \
+    --standalone \
+    --nproc_per_node "$NUM_DEVICES" \
+    "$ROOT_DIR/scripts/train_dspark.py" \
+    --target-model-path "$TARGET_MODEL_PATH" \
+    --draft-config-path "$ROOT_DIR/configs/qwen3.5-4b-dspark.json" \
+    --train-data-path "$TRAIN_DATA_PATH" \
+    --output-dir "$ROOT_DIR/outputs/qwen3.5-4b-dspark-npu" \
+    --num-epochs 10 \
+    --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 \
+    --target-model-backend hf \
+    --block-size 16 \
+    --embedding-key model.embed_tokens.weight \
+    --markov-rank 256 \
+    --enable-confidence-head \
+    --confidence-head-with-markov \
+    --ce-loss-alpha 0.1 \
+    --l1-loss-alpha 0.9 \
+    --confidence-head-alpha 1.0 \
+    --trust-remote-code

examples/run_qwen3_8b_dspark_online.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# DSpark online training for Qwen3-8B.
+#
+# DSpark = DFlash block-diffusion drafter + EAGLE-style Markov & confidence heads,
+# trained with cross-entropy + L1 distribution distillation + confidence BCE.
+# The L1 / confidence losses need the target model's FINAL hidden state, so the
+# target backend must surface it. The 'hf' backend (default below) always does;
+# the 'sglang' backend does when its runner returns both the captured aux stream
+# and the final hidden state. To train CE-only (no target final hidden state),
+# pass: --l1-loss-alpha 0 --no-confidence-head --ce-loss-alpha 1.0
+
+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}
+TARGET_BACKEND=${3:-hf}
+
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
+    $ROOT_DIR/scripts/train_dspark.py \
+    --target-model-path Qwen/Qwen3-8B \
+    --draft-config-path $ROOT_DIR/configs/qwen3-8b-dspark.json \
+    --train-data-path $ROOT_DIR/cache/dataset/perfectblend_qwen3-8b_regen.jsonl \
+    --output-dir $ROOT_DIR/outputs/qwen3-8b-dspark \
+    --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 4.0 \
+    --log-interval 50 \
+    --save-interval 1000 \
+    --report-to wandb \
+    --wandb-project specforge-qwen3-8b-dspark \
+    --target-model-backend $TARGET_BACKEND \
+    --block-size 16 \
+    --num-anchors 512 \
+    --markov-rank 256 \
+    --enable-confidence-head \
+    --confidence-head-with-markov \
+    --ce-loss-alpha 0.1 \
+    --l1-loss-alpha 0.9 \
+    --confidence-head-alpha 1.0 \
+    --wandb-name qwen3-8b-dspark-perfectblend