fix(dflash): infer draft dimensions from safetensors

[easel]

Summary

Fixes Qwen3.6 DFlash draft loading by deriving the draft architecture dimensions from the draft .safetensors tensor shapes instead of inheriting the verifier/target model's attention metadata.

The failure reproduced as:

draft load: safetensors: 'layers.0.self_attn.k_norm.weight' shape[0]=128 expected 256

A later mismatch could also appear on projection tensors, for example o_proj.weight shape[1]=4096 expected 3072.

Current understanding

This is not a bug in upstream safetensors. The safetensors archive already records and returns the correct tensor shape, e.g. layers.0.self_attn.q_norm.weight has shape [128]. The bug was in lucebox: before validating the draft tensor header, load_draft_safetensors copied target model attention dimensions into DraftWeights.

That is wrong for Qwen3.6 because the target and draft are different transformers:

Dimension	Qwen3.6 target	DFlash draft
head_dim	256	128
n_head	24	32
n_head_kv	4	8
q_dim	6144	4096
kv_dim	1024	1024

The old behavior happened to work for earlier drafts where these dimensions matched closely enough, which made the target-metadata inheritance bug latent.

What changed

load_draft_safetensors now infers draft dimensions directly from the safetensors header after parsing it:

• hidden_norm.weight / norm.weight -> n_embd
• fc.weight -> captured target-layer count
• layers.0.mlp.gate_proj.weight -> n_ff
• layers.0.self_attn.q_norm.weight / k_norm.weight -> head_dim
• layers.0.self_attn.q_proj.weight -> n_head
• layers.0.self_attn.k_proj.weight / v_proj.weight -> n_head_kv

The loader still cross-checks target-dependent invariants when target metadata is available:

• draft hidden size must match target hidden size
• fc.weight's inferred captured-layer count must match target->n_capture_layers
• mask_token_id is still copied from the target

Why this replaces the earlier config.json approach

The previous version of this PR parsed the draft model's adjacent config.json and used that as the authoritative source for draft dimensions. After reviewing the failure more carefully, that was more complicated than necessary for this bug.

For the #195 failure, the safetensors header itself is the authoritative source. Parsing config.json may still be useful later for non-tensor architecture metadata such as SWA layer types or sliding-window policy, but it is not required to fix the tensor-shape mismatch and should be handled as a separate feature if needed.

Compatibility

• Drafts without config.json are supported; the fix no longer depends on that file.
• No-target callers still work because dimensions are inferred from the archive header itself.
• Existing target/draft compatibility checks are stricter where the draft must agree with the target, especially hidden size and captured-layer count.
• The patch does not rename the historical DFLASH27B_TARGET_* constants, even though some of them describe draft defaults. That rename would be mechanical but noisy and is intentionally out of scope.

Test plan

• Proved upstream safetensors preserves the relevant tensor shape with a direct local test: a BF16 tensor named like layers.0.self_attn.q_norm.weight with shape [128] deserializes as [128].
• g++ -std=c++17 -fsyntax-only -I dflash/include -I dflash/src -I dflash/src/common -I dflash/src/draft -I dflash/deps/llama.cpp/ggml/include dflash/src/draft/draft_safetensors_loader.cpp
• CUDACXX=/usr/local/cuda/bin/nvcc cmake -S dflash -B /tmp/lbh-pr195-shape/build -DDFLASH27B_GPU_BACKEND=cuda -DDFLASH27B_FA_ALL_QUANTS=OFF -DCMAKE_CUDA_ARCHITECTURES=86
• cmake --build /tmp/lbh-pr195-shape/build --target smoke_load_draft -j 8
• LD_LIBRARY_PATH=/tmp/lbh-pr195-shape/build/deps/llama.cpp/ggml/src:/tmp/lbh-pr195-shape/build/deps/llama.cpp/ggml/src/ggml-cuda /tmp/lbh-pr195-shape/build/smoke_load_draft /home/erik/Projects/lucebox-hub/dflash/models/draft/model.safetensors

The real local draft smoke test loaded 58 tensors / 3.22 GiB and completed with OK.

[cubic-dev-ai]

No issues found across 5 files

_{Re-trigger cubic}