The One-Position Bug: Tracing a Subtle Attention Mask Mismatch in DFlash Drafter Evaluation

The Message

[assistant] Now fix the context slicing in the main loop -- strictly BEFORE anchor:
[edit] /data/dflash/scripts/eval_drafter.py
Edit applied successfully.

At first glance, this appears to be a routine edit commit—a single-line change to an evaluation script, accompanied by a brief explanatory comment. But this message, <msg id=8952>, represents the culmination of an extended debugging odyssey spanning dozens of messages, multiple machines, and a cascade of architectural discoveries. It is the moment when a subtle but critical mismatch between training-time and inference-time attention semantics was identified and corrected, closing one chapter of a larger investigation into why the DFlash drafter was producing garbled, repetitive output.

The Context: Building a Correct Evaluation Harness

To understand why this message was written, one must understand the broader investigation unfolding across the conversation. The assistant had been tasked with evaluating the DFlash drafter—a speculative decoding model trained to predict multiple tokens in parallel using a draft-then-verify architecture. The drafter was trained on hidden states extracted from a Qwen3.6-27B target model, using a sophisticated flex attention mask that controlled which positions each token could attend to during training.

The evaluation harness, running on a separate machine (CT129), needed to replicate the exact same attention semantics that the model experienced during training. Any mismatch—even a single position offset—could cause the drafter to receive different information at inference time than it was conditioned on during training, leading to degraded or nonsensical output.

The assistant had already discovered several major bugs through this evaluation work: the fc projection was using only 4 of 5 target layers, the noise schedule was corrupting target logits, and the loss function was a soft KL divergence instead of the paper's hard cross-entropy. But before those architectural issues could even be diagnosed, the evaluation harness itself needed to be trustworthy—and the garbled output suggested it wasn't.

The Discovery: Strict Inequality in the Training Mask

The critical insight came in the preceding message, <msg id=8951>, where the assistant read the training code's flex attention mask function and traced through its logic. The training mask's base_prefix_mod function contained this condition:

before_anchor = kv_base_pos < q_anchor  # STRICTLY before, NOT including anchor

This strict inequality meant that block tokens could only attend to context positions strictly before the anchor position. The anchor token itself—the token whose hidden state served as the conditioning signal for the block—was excluded from the context K/V cache. Its information entered the block's computation only through the block's own position 0 embedding, not through the attention mechanism's key-value lookup.

The assistant's evaluation code, however, was using fc_output[:, :anchor_pos + 1, :] to slice the context, which included the anchor position. This one-position discrepancy meant the drafter was seeing information at inference time that it never saw during training—a subtle form of data leakage that could distort attention distributions and produce unreliable draft tokens.

The Reasoning: Tracing Through a Concrete Example

The assistant's reasoning in <msg id=8951> works through a concrete example with anchor_pos = 24 to verify the analysis. During training, position IDs run from 1 to N, so at index 24 the position ID is 25. The get_base_indices function, called with anchor position 24 and block size 16, generates indices [25, 26, ..., 40]. The flex attention mask then allows the block to access context positions [0, 1, ..., 23] with their corresponding position IDs [1, 2, ..., 24].

In the evaluation setup, using anchor_pos = 24, the context should be extracted from the first 24 positions, giving ctx_len = 24. The context slice should be fc_output[:, :24, :] (i.e., [:anchor_pos]), not fc_output[:, :25, :] (i.e., [:anchor_pos + 1]). The fix is to change the slice endpoint from anchor_pos + 1 to anchor_pos.

This also requires adjusting the position ID ranges: the context spans positions 1 through anchor_pos, while the block starts at anchor_pos + 1 and extends through anchor_pos + block_size. The assistant also recognized the need to handle the edge case where anchor_pos = 0, which would result in an empty context and only block-to-block attention—a scenario guarded in the subsequent message <msg id=8953>.

Why This Bug Matters

The one-position offset might seem trivial—what difference could a single token make in a 27-billion-parameter model processing sequences of thousands of tokens? But in the DFlash architecture, the anchor position is special: it is the boundary between the conditioning context and the generated block. The anchor token's hidden state serves dual purposes: it is both the last token of the context that the block can condition on and the first token of the block's own generation. The training mask carefully separates these two roles by excluding the anchor from the context K/V and letting it enter only through the block's position embedding.

By including the anchor in the context, the evaluation harness was effectively giving the drafter access to the anchor's hidden state through two paths simultaneously: once through the context K/V cache (as a key that block queries could attend to) and once through the block's own position 0 embedding (as a query that could attend to itself and other block positions). This double-counting could distort attention distributions, potentially explaining the garbled, repetitive output patterns the assistant had observed—sequences like "FizzFizzFizzBuzz" and "elif elif elif" that suggested the drafter was stuck in loops.

Input Knowledge Required

Understanding this message requires familiarity with several concepts: the DFlash anchor block mechanism, where a drafter predicts multiple tokens in parallel conditioned on a prefix; flex attention masks, which use custom masking functions to control token-to-token visibility; the distinction between context positions (which provide conditioning information) and block positions (which are being generated); and the importance of strict vs. non-strict inequality in attention masking. One must also understand that the evaluation harness reimplements the attention mechanism using standard PyTorch operations (since flex_attention requires CUDA and the evaluation setup may not support it), making it vulnerable to subtle semantic mismatches.

Output Knowledge Created

This message produces a corrected evaluation harness that properly excludes the anchor position from the context K/V cache, aligning inference-time attention semantics with training-time behavior. This is one of several fixes needed to make the evaluation trustworthy—alongside the hidden state extraction fix (using fla instead of PyTorch's CPU fallback), the architecture fix (expanding fc to use all 5 target layers), and the loss function fix (switching from soft KL to hard cross-entropy). Each fix builds on the previous one, and the evaluation harness must be correct before any of the other fixes can be meaningfully assessed.

The Broader Significance

This message exemplifies a pattern that recurs throughout the conversation: the assistant methodically traces through training code to identify subtle mismatches between training-time and inference-time behavior, then corrects the evaluation harness to match. The one-position bug is not the most consequential discovery of this session—the three training bugs (noise corrupting target logits, fc shortcut including target layer, loss function mismatch) are far more impactful. But it is a necessary prerequisite: without a correct evaluation harness, none of those bugs could have been reliably diagnosed, and any attempted fixes would have been evaluated against a broken baseline.

The message also illustrates the importance of precise attention semantics in speculative decoding architectures. The DFlash paper's flex attention mask is carefully designed to prevent information leakage between blocks while allowing full bidirectional attention within each block. Replicating these semantics in a different framework (standard attention vs. flex attention) requires exact fidelity to every inequality and boundary condition. A single off-by-one error can invalidate the entire evaluation.