fix: Share the target lm_head for DSpark drafts without their own head

[rain-sicoreai]

What does this PR do?

Type of change: Bug fix (direct builder, DSpark speculative decoding)

Overview:

RadixArk/Qwen3.8-27B-DSpark is listed under DSpark Draft Models as the pair for RadixArk/Qwen3.8-27B-NVFP4. The draft checkpoint ships five proposal layers and no lm_head; it shares the target model's head. The ONNX export path already handles this (tensorrt_edgellm/models/dspark/modeling_dspark_draft.py documents lm_head as "shared with base"), but the direct builder rejected the checkpoint:

ValueError: DSpark draft checkpoint must provide lm_head

so tensorrt-edgellm-serve could not build the pair.

DSparkDraftModel now mirrors Eagle3DraftModel: when the draft checkpoint has no head, from_config opens the paired target checkpoint (--target-model-dir, which the CLI already passes for paired drafts) as a second Weights source with checkpoint_source="target", builds lm_head from the target's head in the target's precision (NVFP4 here), and releases the target weights in close(). Drafts that carry their own head are unchanged.

Verified on Jetson AGX Thor with the Qwen3.8-27B NVFP4 (ModelOpt) base and Qwen3.8-27B-DSpark, block size 7, batch 4, max input 32768, KV capacity 65536:

• bundle builds (base 9.5 min, draft 11 min, visual 42 s) and serves through tensorrt-edgellm-serve with --speculative-config '{"method":"dspark","model":<draft>,"num_speculative_tokens":7}';
• coding prompts decode at 44–57 tok/s single-stream and 128 tok/s aggregate at 4 concurrent requests; num_speculative_tokens from 1 to 7 works at runtime;
• image requests keep working alongside DSpark.

Usage

tensorrt-edgellm-serve RadixArk/Qwen3.8-27B-NVFP4 \
  --cache-dir /data/edgellm-cache \
  --speculative-config '{"method":"dspark","model":"RadixArk/Qwen3.8-27B-DSpark","num_speculative_tokens":7}'

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit.
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed (no unit test: the change mirrors the existing EAGLE3 target-head path, which has none either; verified end to end as described above).
• All tests are passing.

📄 Documentation

• No documentation change needed; the pair is already listed in supported-models.md.

⚙️ Compatibility

• The change is backward compatible: drafts with their own lm_head take the unchanged path.

Additional Information

Independent of #199; it only touches experimental/builder/models/dspark/modeling_dspark_draft.py.

🤖 Generated with Claude Code