dflash_generate: final mask cleanup is value-based and drops legitimate output tokens equal to mask_token_id

[shaun0927]

Summary

dflash_generate's final cleanup at dflash/model.py:150-151 removes any output token that happens to equal mask_token_id, regardless of whether it was generated by the model or just leftover unfilled buffer.

output_ids = output_ids[:, :max_length]
output_ids = output_ids[:, output_ids[0] != mask_token_id]

The intent is to drop the trailing mask-padded positions left over when generation stops early. The actual effect is value-based, so any legitimate token whose id equals mask_token_id is silently dropped from the middle of the sequence.

Evidence (no model required)

import torch
mask_token_id = 0
buf = torch.tensor([[11, 12, 13,
                     101, 102, 0, 103, 104, 0, 105, 106, 107, 108,
                     0, 0, 0]])  # last 3 = unfilled mask
stripped = buf[:, buf[0] != mask_token_id]
print(stripped.tolist())
# [[11, 12, 13, 101, 102, 103, 104, 105, 106, 107, 108]]   <- two zeros gone
# expected: [[11, 12, 13, 101, 102, 0, 103, 104, 0, 105, 106, 107, 108]]

Why this matters in practice

The published draft configs choose mask_token_id from each tokenizer's reserved/special-token zone (e.g. Llama-3.1: 128002 = <|reserved_special_token_0|>, Qwen3.5-4B/9B: 248070, Qwen3-8B-b16: 151669, gpt-oss-20b: 200000), which avoids the worst case under the supported HF chat templates. But the cleanup is still fragile by construction:

• A user-supplied or future config that chooses an in-vocabulary id (or 0) silently truncates output.
• If the model ever emits the special token (e.g. wrong tokenizer revision, raw / non-chat usage), that token vanishes with no error and the reported num_output_tokens is wrong.

Suggested fix

Track the number of valid tokens written during the loop and slice by length, rather than filtering by value. The start variable already tracks this:

output_ids = output_ids[:, :start]
# stop-token trim still applied as today

If you'd like, I can prepare a small PR with the change plus a unit test that asserts the legitimate-token case.

Environment

• repo at HEAD 6e0c951
• python 3.12, torch 2.11