feat: add Ascend NPU support for VP-Drafter

[curnane-lab]

Motivation

This PR adds Ascend NPU training support for the VP-Drafter (variable-prefix DFlash drafter used in D2SD) on Qwen3.5-4B.

The base VP-Drafter training feature is cherry-picked from #592. This PR layers the NPU enablement on top while keeping the base commits intact.

Modifications

• Added configs/qwen3.5-4b-dta.json: Qwen3.5-4B VP-Drafter draft config with training_mode=vp_drafter and prefix_weight_base=0.9.
• Added examples/run_qwen3.5_4b_dta_online_npu.sh: Ascend NPU launcher using HF backend, SDPA attention, and HCCL distributed training.
• Cherry-picked the base VP-Drafter training commits from [Feature] Add VP drafter training mode for DFlash #592

Related Issues

N/A

Accuracy Test

N/A (training-only change; inference still uses the existing DFlashDraftModel architecture).

Benchmark & Profiling

N/A

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Code Review

This pull request introduces VP-Drafter (variable-prefix) training support for DFlash, including new configuration files, training scripts, and core implementation changes in specforge/core/dflash.py to sample prefix lengths and calculate the corresponding loss. Feedback on these changes suggests minor optimizations in PyTorch operations, such as simplifying torch.pow with a scalar base and performing integer-precision operations before casting to float. Additionally, it is recommended to double-quote variable expansions in the newly added bash scripts to prevent word splitting.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In bash scripts, it is highly recommended to double-quote all variable expansions (especially paths like $ROOT_DIR and user inputs like $NUM_GPUS and $ATTENTION_BACKEND) to prevent word splitting and pathname expansion if they contain spaces or special characters. This is particularly important for paths that might be run in environments with non-standard directory names.

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

[gemini-code-assist]

Instead of creating a full tensor of the base value and then raising it to the power of the exponent tensor, you can directly use torch.pow with a scalar base and a tensor exponent. This is cleaner, more readable, and avoids allocating unnecessary intermediate tensors.

        weights = torch.pow(self.prefix_weight_base, prefix_ids.float())

[gemini-code-assist]

You can perform the subtraction and clamping in integer precision before casting to float. This is slightly more efficient and cleaner as it avoids casting both tensors to float beforehand.

Suggested change

	k = torch.arange(self.block_size, device=device).view(1, 1, -1)
	effective_pos = (
	k.float() - prefix_lengths.unsqueeze(-1).float()
	).clamp(min=0)
	k = torch.arange(self.block_size, device=device).view(1, 1, -1)
	effective_pos = (
	k - prefix_lengths.unsqueeze(-1)
	).clamp(min=0).float()