Reverse-Engineering the DFlash Drafter: Inspecting Raw Weights to Build a Model Config from Scratch

Introduction

In the ongoing effort to deploy speculative decoding for the Qwen3.6-27B model, the assistant reached a critical inflection point in message [msg 6917]. After successfully copying the gated z-lab/Qwen3.6-27B-DFlash drafter model's safetensors file to the target machine, the assistant faced a fundamental problem: the model was distributed as raw weight tensors with no accompanying configuration files. Without a config.json, neither vLLM nor any other serving framework could load the model. The message captures the assistant's pivot from deployment operations to investigative reverse-engineering — using Python to inspect the safetensors file's tensor shapes and names to deduce the model architecture, enabling the creation of a correct configuration from scratch.

This seemingly simple inspection command represents a crucial moment of problem-solving. The assistant could have attempted to download the config from HuggingFace (blocked by gated access), guessed parameters based on the target model, or tried to use the model without configuration. Instead, it chose the most reliable path: directly interrogating the weight tensors to determine the exact architectural dimensions.

The Message Content

The message begins with a reasoning preamble that frames the investigation:

Now I need to figure out the right config for the Qwen3.6-27B drafter. The target model is Qwen3.6-27B which uses qwen3_5 architecture with 64 layers, hidden_size=5120, 24 attention heads, 4 KV heads. Let me inspect the safetensors to understand the draft model's dimensions:

This is followed by a bash command that executes a Python script on the remote machine via SSH:

ssh root@10.1.230.172 '/root/ml-env/bin/python3 -c "
from safetensors import safe_open
f = safe_open(\"/root/models/Qwen3.6-27B-DFlash/dflash-q36-27b.safetensors\", framework=\"pt\")
for name in sorted(f.keys()):
    t = f.get_tensor(name)
    print(f\"{name:60s} {str(list(t.shape)):>30s} {t.dtype}\")
"'

The output reveals the model's tensor structure, starting with:

• fc.weight — shape [5120, 25600], dtype bfloat16
• hidden_norm.weight — shape [5120], dtype bfloat16
• Layer 0 components: input_layernorm.weight, mlp.down_proj.weight, mlp.gate_proj.weight, etc.

Why This Message Was Written: The Motivation and Context

The message was written to solve a concrete deployment blocker. The broader context reveals a multi-phase effort to deploy speculative decoding for Qwen3.6-27B. The team had already:

1. Migrated the Qwen3.6-27B deployment from one host (kpro6) to another (kpro5)
2. Achieved a baseline of 73.5 tok/s with MTP speculation on SGLang
3. Investigated DFlash and DDTree as more advanced speculative decoding methods
4. Acquired the gated DFlash drafter model as a 3.3GB safetensors file
5. Copied the file to the target machine at /root/models/Qwen3.6-27B-DFlash/dflash-q36-27b.safetensors The critical blocker was that the model card on HuggingFace labeled the drafter as "still under training" and the repository was gated — requiring HuggingFace authentication and license acceptance. The assistant had obtained the raw weight file (likely from a local copy at /tmp/dflash-q36-27b.safetensors on the host machine, as referenced in [msg 6911]), but this came without the configuration metadata that normally accompanies HuggingFace model repositories. Without a config.json, vLLM's model loading code would fail. The assistant needed to know: - The model's hidden dimension (hidden_size) - The number of layers - The intermediate size (for MLP projections) - The vocabulary size - The number of attention heads and KV heads - The block size for DFlash's block diffusion - The target layer IDs (which target model layers to extract hidden states from) - The mask token ID and other DFlash-specific parameters The reasoning preamble shows the assistant already had some knowledge about the target model (Qwen3.6-27B: 64 layers, hidden_size=5120, 24 attention heads, 4 KV heads), but the draft model is a separate 2B-parameter network with its own architecture. The draft model's dimensions could not be inferred from the target model alone.

How Decisions Were Made

The decision to inspect the safetensors file directly reflects a pragmatic engineering mindset. Several alternative approaches existed:

1. Download config from HuggingFace: The most obvious approach, but blocked by gated access. The assistant could have attempted to use the HF API with authentication tokens, but this would require user credentials and license acceptance — a multi-step process that might not succeed.
2. Guess based on similar DFlash drafters: The assistant had previously examined the config for z-lab/Qwen3-8B-DFlash-b16 (a different model) in [msg 6916]. That model had hidden_size=4096, 32 layers, and specific DFlash parameters. One could extrapolate from this, but the Qwen3.6-27B drafter could have different dimensions.
3. Inspect the weights directly: This is what the assistant chose. It's the most reliable method because the tensor shapes are ground truth — they cannot lie. The shapes directly encode the architectural parameters. The choice of tooling was also deliberate: using Python's safetensors library (the safe_open function) to iterate over all tensor keys and print their shapes. This is a lightweight operation that doesn't require loading the full 3.3GB file into memory — safe_open provides lazy tensor access via f.get_tensor(name), which reads individual tensors on demand. The script iterates over all keys, printing each tensor's name, shape, and dtype, producing a complete architectural fingerprint of the model.

Assumptions Made by the Assistant

Several assumptions underpin this message:

1. The safetensors file contains all model weights: The assistant assumes that the single .safetensors file is complete and contains every parameter of the DFlash drafter. This is a reasonable assumption for HuggingFace models, but some large models split weights across multiple safetensors files. If the drafter had been sharded, this inspection would only reveal part of the architecture.
2. The tensor naming convention follows the standard DFlash pattern: The assistant assumes that tensor names like fc.weight, hidden_norm.weight, layers.0.input_layernorm.weight, etc., follow the conventions of the DFlashDraftModel architecture. This is validated by the earlier research into the DFlash paper and the reference config for Qwen3-8B-DFlash-b16.
3. The model uses bfloat16 precision throughout: The output confirms all tensors are torch.bfloat16, which matches expectations for a modern LLM drafter.
4. The remote Python environment has safetensors installed: The command uses /root/ml-env/bin/python3, which is the ML environment set up earlier in the session. The assistant assumes this environment has the safetensors package available.
5. SSH access is available and working: The command executes via SSH to root@10.1.230.172, which is the CT129 LXC container. This assumes the SSH connection is stable and the remote Python interpreter is functional.

Mistakes or Incorrect Assumptions

The message itself doesn't contain obvious errors, but there are potential pitfalls worth examining:

1. The output is truncated in the conversation: The printed output shows only the first few tensors (fc.weight, hidden_norm.weight, and the start of layer 0's components). The full output would be much longer — potentially hundreds of lines for a 2B-parameter model. The conversation only captures the beginning of the output, which means the reader (and potentially the assistant in subsequent messages) doesn't see the complete architectural picture from this message alone. The assistant would need to either capture the full output or run additional targeted queries.
2. The inspection doesn't reveal DFlash-specific parameters: The tensor shapes tell us about the model's dimensions (hidden_size, intermediate_size, num_layers, vocab_size), but they don't directly reveal DFlash-specific configuration like block_size, mask_token_id, or target_layer_ids. These parameters are typically stored in the config.json, not encoded in weight shapes. The assistant would need to infer or guess these from the reference DFlash config or the model card.
3. No verification of tensor count: The script prints all tensor names and shapes but doesn't count them. A 2B-parameter model should have a specific number of parameter tensors (e.g., for each layer: input_layernorm, post_attention_layernorm, mlp.gate_proj, mlp.up_proj, mlp.down_proj, plus attention components). Without verifying the count, the assistant might miss missing components.
4. Assumption about the model architecture: The assistant assumes this is a standard DFlash draft model with a transformer architecture. However, the Qwen3.6 family uses GDN (Gated DeltaNet) hybrid attention, and the drafter might incorporate GDN components. The tensor names shown (input_layernorm, mlp.down_proj, etc.) suggest a standard transformer, but the full output might reveal GDN-specific layers.

Input Knowledge Required to Understand This Message

To fully grasp this message, one needs:

1. Understanding of safetensors format: The .safetensors file format is a safe serialization format for PyTorch tensors, designed as an alternative to pickle-based formats. The safe_open function provides dictionary-like access to tensor names.
2. Knowledge of transformer model architecture: Understanding that hidden_size determines the embedding dimension, intermediate_size (or ffn_hidden_size) determines MLP projection dimensions, num_layers determines the number of transformer blocks, and vocab_size determines the output projection dimension.
3. Familiarity with DFlash speculative decoding: DFlash (Draft model with FLASH attention) is a block diffusion method where the draft model predicts multiple future tokens simultaneously. The draft model is typically a smaller transformer with a special fc (fully connected) projection head that maps target model hidden states to draft model inputs.
4. Context about the deployment environment: The target machine is an LXC container (CT129) at IP 10.1.230.172, with a Python virtual environment at /root/ml-env/bin/python3. The SGLang service was running but consuming GPU memory, and the plan was to stop it and deploy vLLM instead.
5. The gated model situation: The z-lab/Qwen3.6-27B-DFlash model is gated on HuggingFace, meaning access requires authentication and license acceptance. The assistant obtained the safetensors file through an alternative path (local copy from the host machine) but lacks the accompanying configuration files.

Output Knowledge Created by This Message

The message produces critical architectural intelligence about the DFlash drafter:

1. Vocabulary size: The fc.weight tensor has shape [5120, 25600]. The first dimension (5120) is the hidden_size of the target model (Qwen3.6-27B), confirming the drafter takes target model hidden states as input. The second dimension (25600) is the draft model's hidden_size, suggesting a ~2B parameter model (25600 × 25600 × 4 for attention + MLP projections ≈ 2.6B parameters).
2. Hidden dimension: The hidden_norm.weight has shape [5120], which is the draft model's hidden_size. Wait — that's 5120, not 25600. Let me re-examine: fc.weight is [5120, 25600] — this is the feature connector (fc) that projects from target hidden_size (5120) to draft hidden_size (25600). The hidden_norm.weight being [5120] is the normalization after the fc projection... Actually, hidden_norm with shape [5120] suggests the draft model's hidden_size is 5120, and the fc layer projects from 5120 to 25600 as an intermediate expansion. Hmm, this needs careful interpretation. Actually, looking at the reference DFlash config for Qwen3-8B-DFlash-b16 (from [msg 6916]), the fc layer maps from target hidden_size to draft hidden_size. If target hidden_size=5120 (Qwen3.6-27B) and fc.weight is [5120, 25600], then the draft hidden_size is 25600? That seems too large for a 2B model. Wait, let me re-read: fc.weight [5120, 25600] — in PyTorch linear layers, the weight shape is [out_features, in_features]. So this maps from 25600 input features to 5120 output features. That means the draft model's hidden_size is 25600? No, that's enormous for a 2B model. Actually, in the DFlash architecture, the fc layer (feature connector) projects the target model's hidden states into the draft model's embedding space. If the target hidden_size is 5120 and the fc weight is [5120, 25600], this could mean: - Input: target hidden states of size 5120 - Output: projected to 25600 dimensions But 25600 hidden size is implausibly large for a 2B model. More likely, the weight shape notation might be [in_features, out_features] depending on how the safetensors were saved, or the interpretation is different. This ambiguity is exactly why the inspection is valuable — and why the truncated output is a limitation. The assistant would need to look at more tensors (like the attention projection weights) to determine the true hidden_size.
3. Layer structure: The output shows layer 0 has input_layernorm.weight [5120], mlp.down_proj.weight [5120, 17408], mlp.gate_proj.weight [...] (truncated). The intermediate size (17408) suggests the MLP expansion factor, which helps identify the model scale.
4. Precision: All tensors are bfloat16, confirming the model uses BF16 precision (not FP16 or FP32).
5. Model completeness: The fact that safe_open successfully opens the file and iterates over keys confirms the file is a valid safetensors archive with readable tensor metadata.

The Thinking Process Visible in the Reasoning

The reasoning preamble reveals the assistant's mental model:

"Now I need to figure out the right config for the Qwen3.6-27B drafter."

This frames the problem: the assistant needs to create a configuration file, and the current approach is to deduce it from the weights.

"The target model is Qwen3.6-27B which uses qwen3_5 architecture with 64 layers, hidden_size=5120, 24 attention heads, 4 KV heads."

The assistant references known facts about the target model to establish context. This knowledge comes from earlier research (the model card for Qwen3.6-27B). The assistant is mentally comparing the draft model's architecture to the target model's.

"Let me inspect the safetensors to understand the draft model's dimensions:"

This is the decision point: instead of guessing or searching for config files, the assistant chooses to inspect the weights directly. The thinking is: "The weights are the ground truth. If I read their shapes, I can reconstruct the architecture."

The choice of the safe_open API (rather than torch.load or safetensors.torch.load_file) is also deliberate. safe_open provides lazy, key-by-key access without loading the entire file into memory, which is ideal for inspection. The sorted(f.keys()) ensures deterministic ordering, and the f-string formatting with :60s and :>30s creates aligned output for readability.

Broader Significance

This message exemplifies a recurring pattern in machine learning engineering: deploying models in production often requires working with incomplete or non-standard artifacts. The clean, well-documented model on HuggingFace is the ideal case, but real-world deployment frequently involves raw weight files, partial checkpoints, or models from research repositories that lack proper packaging.

The assistant's approach — directly inspecting tensor shapes to reverse-engineer the architecture — is a practical skill that separates experienced ML engineers from novices. Rather than getting blocked by missing configuration files, the assistant treats the weights themselves as the authoritative source of architectural truth.

This message also highlights the gap between research code and production deployment. The DFlash drafter, labeled "still under training" by its authors, was shared as a raw safetensors file without the configuration scaffolding needed for serving frameworks. The assistant's investigation is the first step in bridging that gap — transforming a research artifact into a deployable component.

Conclusion

Message [msg 6917] captures a pivotal investigative step in the DFlash speculative decoding deployment. Faced with a gated model repository and missing configuration files, the assistant chose to reverse-engineer the model architecture by directly inspecting the weight tensors in the safetensors file. This decision reflects a pragmatic, data-driven approach to problem-solving: when documentation is unavailable, let the code (or in this case, the weights) speak for themselves.

The inspection revealed critical architectural parameters — the vocabulary projection dimensions, hidden size indicators, layer structure, and precision — that would inform the creation of a correct config.json. While the output was truncated and some DFlash-specific parameters (like block_size and target_layer_ids) remained unknown, the message laid the foundation for the next steps: constructing a complete configuration, deploying vLLM with DFlash, and ultimately benchmarking against the MTP baseline.

In the broader narrative of the session, this message represents the transition from deployment operations (copying files, installing software) to deep technical investigation (understanding model architecture, creating configuration from scratch). It's a reminder that in ML engineering, the most valuable tool is often not a framework or library, but the ability to read the raw data and infer its structure.