Reverse-Engineering a Speculative Decoding Drafter: Tensor Shapes as Configuration Blueprints

In the sprawling effort to deploy advanced speculative decoding on a Qwen3.6-27B model, one message stands out as a masterclass in model reverse-engineering. Message [msg 6918] is a brief but dense analytical moment where the assistant, faced with a raw safetensors file and no HuggingFace configuration, reconstructs the entire architecture of a DFlash drafter model from nothing but tensor shapes. This single message captures the essence of what it means to work at the frontier of open-source ML deployment: when the research artifacts are incomplete, the documentation is sparse, and the only reliable source of truth is the weights themselves.

The Broader Context: Racing Toward DDTree

To understand why this message matters, we need to step back into the larger narrative. The session had been pushing the boundaries of speculative decoding for days. The team had already achieved 73.5 tok/s with MTP (Multi-Token Prediction) speculation on Qwen3.6-27B using SGLang, but the goal was far more ambitious: DFlash (a block-diffusion-based drafter) and ultimately DDTree (tree-based verification that could theoretically double throughput again). The path was clear — deploy DFlash first as a linear-chain baseline, then patch DDTree's tree-construction logic on top — but the implementation required a specific drafter model: z-lab/Qwen3.6-27B-DFlash, a gated 2-billion-parameter draft model designed to work alongside the 27-billion-parameter target.

The problem was that this drafter was not available through standard HuggingFace download channels in the usual format. The user had provided the safetensors file directly (/tmp/dflash-q36-27b.safetensors), copied to the remote machine at /root/models/Qwen3.6-27B-DFlash/dflash-q36-27b.safetensors. But a safetensors file alone is not enough to load a model. You need a config.json that tells the framework the architecture: hidden sizes, number of layers, head dimensions, attention configuration, and — critically for DFlash — the target_layer_ids that specify which layers of the target model the drafter extracts hidden states from.

The Message: A Forensic Analysis of Tensor Shapes

The message opens with the assistant's synthesis of what the tensor inspection revealed:

Now I can derive the full config: - fc.weight: [5120, 25600] → hidden_size=5120, input is 55120=25600 → 5 target layers captured - 5 layers (0-4), each with full attention - q_proj: [4096, 5120] → 4096/128=32 Q heads - k_proj: [1024, 5120] → 1024/128=8 KV heads - o_proj: [5120, 4096] → output projects from 32128=4096 back to 5120 - mlp: intermediate_size=17408 (same as target model) - head_dim: 128 (from k_norm/q_norm weight size)

This is a remarkable piece of deductive reasoning. Let's unpack each inference.

The fc.weight tensor is the first clue. Its shape is [5120, 25600]. The assistant recognizes this as a projection from the concatenated hidden states of multiple target model layers. The drafter takes hidden states from specific layers of the target model, concatenates them, and feeds them into its own architecture. Since 25600 / 5120 = 5, the assistant deduces that the drafter captures hidden states from exactly 5 target layers. This is consistent with the DFlash architecture described in the research: the drafter doesn't process the entire target model's hidden states — it selectively extracts from a subset of layers to reduce computational overhead while preserving enough information about the target's internal representations.

The attention projection tensors reveal the head structure. The q_proj weight is [4096, 5120], meaning the query projection maps from hidden_size (5120) to 4096 dimensions. With a head_dim of 128 (deduced from the layer normalization weight sizes), this gives 4096 / 128 = 32 query heads. The k_proj is [1024, 5120], giving 1024 / 128 = 8 key-value heads — a 4:1 head ratio typical of grouped-query attention. The o_proj is [5120, 4096], projecting the concatenated attention outputs back to the model's hidden dimension.

The MLP intermediate size of 17408 matches the target model exactly, which makes sense: the drafter needs to model the target's behavior, so it uses the same feed-forward capacity.

The Critical Puzzle: Target Layer IDs

The most important unknown is the target_layer_ids — which specific layers of the 64-layer target model the drafter extracts from. The assistant reasons:

The target model has 64 layers, so 5 target layers would be evenly spaced. Let me check the Qwen3-8B DFlash pattern: target has 36 layers, captures [1,9,17,25,33] — roughly uniform spacing. For 64 layers with 5 captures: spacing is ~13, so [1, 14, 27, 40, 53] or similar.

This is a reasonable heuristic based on the published DFlash drafter for Qwen3-8B. The pattern in that model was approximately uniform spacing: layers 1, 9, 17, 25, 33 out of 36 total, which is roughly every 8-9 layers. For 64 layers with 5 captures, the spacing would be approximately 12-13 layers between each captured layer.

However, the assistant also notes a crucial caveat: "But since the model card says 'still under training', there might be a specific config." This acknowledges that the drafter is a work-in-progress, and the layer selection might not follow a simple uniform pattern. The model card's warning — "still under training" — is a recurring theme throughout this session, casting uncertainty over every assumption.

The Verification Step

The message concludes with a bash command to inspect the target model's configuration:

ssh root@10.1.230.172 '/root/ml-env/bin/python3 -c "
import json
with open(\"/root/models/Qwen3.6-27B/config.json\") as f:
    cfg = json.load(f)
tc = cfg[\"text_config\"]
print(\"hidden_size:\", tc[\"hidden_size\"])
print(\"num_hidden_layers:\", tc[\"num_hidden_layers\"])
...

The output confirms the target model's parameters: hidden_size=5120, 64 layers, 24 attention heads, 4 KV heads, head_dim=256, intermediate_size=17408, vocab_size=248320. This cross-validation is essential — the assistant needs to ensure the drafter's architecture is compatible with the target. Notably, the target's head_dim is 256, while the drafter's head_dim is 128. This is expected: the drafter is a smaller, more efficient model, and it doesn't need the same level of representational precision as the full target.

Input Knowledge Required

To fully understand this message, one needs:

  1. Knowledge of the DFlash architecture: Understanding that DFlash drafters extract hidden states from specific layers of the target model, concatenate them, and use them as input to a smaller transformer. The fc.weight tensor's shape directly encodes the number of captured layers.
  2. Transformer architecture fundamentals: Recognizing that q_proj, k_proj, v_proj, and o_proj tensors encode the attention head configuration, and that dividing the projection dimension by the head dimension yields the number of heads.
  3. The grouped-query attention pattern: Understanding why Q heads (32) and KV heads (8) differ, and that this is a deliberate design choice for efficiency.
  4. Context about the broader deployment: Knowing that this is part of a DFlash → DDTree pipeline, that the target is Qwen3.6-27B with 64 layers, and that the drafter is a gated model from Z Lab.
  5. Familiarity with HuggingFace model configurations: Understanding the text_config nested structure and the significance of parameters like rope_theta, max_position_embeddings, and bos_token_id.

Output Knowledge Created

This message produces several critical pieces of knowledge:

  1. The drafter's architecture is confirmed: 5 target layers, hidden_size=5120, 32 Q heads, 8 KV heads, head_dim=128, intermediate_size=17408. This is enough to construct a valid config.json.
  2. The target model's parameters are verified: 64 layers, 24 Q heads, 4 KV heads, head_dim=256, confirming the architectural relationship between drafter and target.
  3. The target_layer_ids remain uncertain: The assistant has a heuristic guess (uniform spacing around layers [1, 14, 27, 40, 53]) but acknowledges this might be wrong. This uncertainty will prove significant — later in the session, incorrect target_layer_ids will be identified as one of the root causes of DFlash's near-zero acceptance rate.
  4. The drafter's vocabulary and tokenizer compatibility: The target's vocab_size (248320) and token IDs (bos=248044, eos=248044) are confirmed, which the drafter must share.

Assumptions and Potential Mistakes

The most significant assumption is the uniform spacing heuristic for target_layer_ids. The assistant assumes that because the Qwen3-8B DFlash drafter used approximately uniform layer spacing, the Qwen3.6-27B drafter does too. This is a reasonable inference but not guaranteed — the model card's "still under training" warning suggests the configuration might be experimental or non-standard. In fact, later in the session ([msg 6919] onward), the assistant will discover that incorrect target_layer_ids contributed to the DFlash drafter's catastrophic 1.1% acceptance rate, requiring deep investigation into the vLLM DFlash proposer code and the discovery of multiple bugs.

Another assumption is that head_dim = 128 derived from the layer normalization weights. The assistant writes "head_dim: 128 (from k_norm/q_norm weight size)" — the k_norm and q_norm weights have size [128], which is consistent with head_dim=128. But this is an inference, not a direct configuration parameter. The target model uses head_dim=256, so the drafter's smaller head dimension is a design choice that could have implications for the quality of its predictions.

The assistant also assumes that all 5 layers have the same architecture — "5 layers (0-4), each with full attention." This is consistent with the tensor names (layers.0 through layers.4 all having the same projection tensors), but it's worth noting that the drafter could theoretically have heterogeneous layer types.

The Thinking Process

What makes this message remarkable is the forensic methodology the assistant employs. Rather than guessing or searching for documentation, the assistant lets the data speak. Each tensor shape is a clue, and the assistant systematically interprets each one:

  1. Start with the largest tensor (fc.weight) to understand the overall architecture — how many target layers are captured.
  2. Examine attention projections to determine head counts and dimensions.
  3. Check MLP dimensions to understand the feed-forward capacity.
  4. Cross-reference with the target model to ensure compatibility.
  5. Use published examples (Qwen3-8B DFlash) as a pattern for the unknown parameter. This is a pattern of reasoning that mirrors how experienced ML engineers work with undocumented models: the weights encode the architecture, and with enough knowledge of transformer design patterns, you can reconstruct the configuration from scratch. The message also demonstrates the importance of verification through independent data sources. The assistant doesn't just derive the drafter's config — it then reads the target model's official config to confirm compatibility. This dual-source verification is a hallmark of robust engineering practice.

Conclusion

Message [msg 6918] is a small but pivotal moment in a much larger journey. It represents the transition from having a raw model file to having a deployable artifact — from safetensors to configuration. The assistant's ability to reverse-engineer the DFlash drafter's architecture from tensor shapes alone is a testament to deep knowledge of transformer internals and the practical skills needed to work with research-stage models.

But the message also foreshadows the challenges ahead. The uncertainty around target_layer_ids will prove costly, and the "still under training" warning on the model card will echo through subsequent debugging sessions. In the end, the drafter's low acceptance rate will be traced not to the architecture derived here, but to integration bugs in the serving framework — bugs that no amount of tensor-shape analysis could have predicted. This message captures the moment before those difficulties emerged, when the path forward still seemed clear and the only obstacle was filling in a few configuration fields.