25Hz tokenizer: qkv_attention_manual never masks padded keys (bool masked_fill), short attention windows get wrong outputs

Where: qwen_tts/core/tokenizer_25hz/vq/whisper_encoder.py, MultiHeadAttention.qkv_attention_manual, lines 245-252 (unchanged since the initial commit d8146d5).

attn_mask = torch.arange(max_seqlen, device=q.device)[None, :] < torch.tensor(seqlens, device=q.device)[:, None]
attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask == 0, -torch.finfo(q.dtype).max)
attn_scores = torch.matmul(q_padded, k_padded.transpose(-2, -1)) * scale
attn_scores = attn_scores + attn_mask

What happens: attn_mask is a torch.bool tensor, and masked_fill keeps the dtype of self, so -3.4e38 is cast to True. After that line the mask is all True, and attn_scores + attn_mask adds 1.0 to every score, which softmax ignores. Padded keys are never masked. Because k_padded / v_padded are zero-initialised, every padded slot receives weight exp(1)/Σ in the softmax, which shrinks the weights of the real keys. No error or warning is raised. The cast behaves the same on every PyTorch release I checked (1.5 to 2.14; the CPU and CUDA kernels both do value.to<bool>()).

When it is reached: MultiHeadAttention.forward uses the manual path when flash_attn is not importable (any CPU / macOS install), and also on GPU whenever q.dtype is not fp16/bf16 (it then sets self.use_flash_attention = False for the rest of the session). Loading without dtype= gives fp32 under the pinned transformers 4.57.3, and the README tokenizer example passes no dtype, so a default GPU install hits it as well. Only flash-attn + fp16/bf16 avoids it.

When it changes results: whenever a batch contains attention windows of unequal length, i.e. any clip longer than 2 s (one window = n_window = 100 post-CNN frames = 32,000 samples) whose length is not a multiple of 2 s, or several clips whose tail windows differ. Only the tokens of the shorter windows are wrong; full windows in the same batch are exact.

Repro (after pip install qwen-tts):

import torch, torch.nn.functional as F
from qwen_tts.core.tokenizer_25hz.vq.whisper_encoder import MultiHeadAttention

torch.manual_seed(0)
mha = MultiHeadAttention(n_state=16, n_head=2).eval()

def reference(q, k, v, seqlens):  # attention computed window by window
    outs, s = [], 0
    for L in seqlens:
        qi, ki, vi = (t[s:s + L].view(L, mha.n_head, -1).transpose(0, 1) for t in (q, k, v))
        outs.append(F.scaled_dot_product_attention(qi, ki, vi).transpose(0, 1).reshape(L, -1)); s += L
    return torch.cat(outs)

for seqlens in ([5, 3], [5, 5]):
    q, k, v = torch.randn(3, sum(seqlens), 16)
    cu = torch.tensor([0] + list(torch.tensor(seqlens).cumsum(0)), dtype=torch.int32)
    got, want = mha.qkv_attention_manual(q, k, v, cu), reference(q, k, v, seqlens)
    print(f"seqlens={seqlens}: max|diff| = {(got - want).abs().max().item():.3e}")

seqlens=[5, 3]: max|diff| = 4.043e-01
seqlens=[5, 5]: max|diff| = 3.576e-07

With the real encoder width (n_state=1280, 20 heads) and windows [100, 100, 50], the two full windows match to 7e-7 and the 50-token window is off by 4.1e-1. Through WhisperEncoderVQ.forward with a 76-frame and a 48-frame mel (n_window=4), only the two tokens of the short tail window differ (4.6e-1); every other position matches to 3e-7.

Fix: build the additive mask as a float tensor:

attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype).masked_fill(~attn_mask, -torch.finfo(q.dtype).max)

With this change all cases above agree to fp32 round-off (fp16/bf16 to 2e-3 / 2e-2). PR with the one-line fix: #370.

Scope: only the 25 Hz tokenizer encoder uses this class. The 12 Hz tokenizer and the TTS model do not import it, so the published 12 Hz checkpoints are unaffected. The 25 Hz weights are not released yet (#34), so today this matters for whoever runs or ports Qwen3TTSTokenizerV1Model from this code. Found while porting the 25 Hz tokenizer to Hugging Face transformers (huggingface/transformers#44517), where the port's per-window attention diverged from this fallback on exactly the tail-window tokens.