Skip to content

Fix shape mismatch in KyutaiSpeechToText generate() last window#46952

Merged
Rocketknight1 merged 4 commits into
huggingface:mainfrom
jiqing-feng:kyutai
Jul 21, 2026
Merged

Fix shape mismatch in KyutaiSpeechToText generate() last window#46952
Rocketknight1 merged 4 commits into
huggingface:mainfrom
jiqing-feng:kyutai

Conversation

@jiqing-feng

@jiqing-feng jiqing-feng commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

CI

Summary

Generating with a KyutaiSpeechToTextForConditionalGeneration model crashes when the
streaming audio_window_size is larger than the number of encoded audio frames:

RuntimeError: The size of tensor a (200) must match the size of tensor b (118) at non-singleton dimension 1

Root cause

In prepare_inputs_for_generation, the audio_tokens buffer is pre-allocated with
audio_window_size frames. Each window the Mimi codec re-encodes the corresponding
audio slice and writes it back via audio_tokens.copy_(new_audio_tokens). When the
last (or only) window has fewer frames than audio_window_size, the codec emits a
shorter tensor than the buffer, so the exact-size copy_ fails.

Fix

Copy only the overlapping frames so the buffer and codec output stay aligned:

n = min(audio_tokens.shape[1], new_audio_tokens.shape[1])
audio_tokens[:, :n].copy_(new_audio_tokens[:, :n])

Single-shot generation is unchanged; repeated generation no longer crashes.

With the default audio_window_size (estimated from the audio length) the same
mismatch shows up as an off-by-one when the estimate exceeds the codec frame count
by one; the repro sets a large window to make the mismatch deterministic.

Reproduction

import numpy as np
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

model_id = "kyutai/stt-2.6b-en-trfs"
processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, dtype=torch.bfloat16).eval()
model.generation_config.audio_window_size = 200  # > encoded frames -> buffer/codec mismatch

audio = np.random.randn(24000 * 5).astype(np.float32) * 0.01  # ~118 frames
inputs = processor(audio=audio, return_tensors="pt").to(model.device)

out = model.generate(**inputs)
print(f"generated {out.shape[-1]} tokens")

Before the fix it aborts with the shape mismatch above; after the fix it prints the
token count.


Signed-off-by: jiqing-feng <[email protected]>
Signed-off-by: jiqing-feng <[email protected]>
@jiqing-feng
jiqing-feng marked this pull request as ready for review June 29, 2026 08:44
@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: kyutai_speech_to_text

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 29792117293:2
Result: success | Jobs: 5 | Tests: 327 | Failures: 0 | Duration: 2m 12s

@Rocketknight1 Rocketknight1 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, makes sense to me, and should be safe!

@Rocketknight1
Rocketknight1 enabled auto-merge July 21, 2026 13:21
@Rocketknight1
Rocketknight1 added this pull request to the merge queue Jul 21, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Merged via the queue into huggingface:main with commit ac884e7 Jul 21, 2026
36 checks passed
SangbumChoi added a commit to SangbumChoi/transformers that referenced this pull request Jul 22, 2026
* upstream/main: (39 commits)
  Remove deprecated training args and `is_fast` property (huggingface#46917)
  Consistent output shape from `get_image_features` (huggingface#46405)
  Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423)
  fix failed test cases for qwen3_omni_moe model (huggingface#47449)
  Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233)
  Move `value` padding into the attention interfaces that need it (huggingface#47451)
  Simplify function dispatch for linear attention (huggingface#47450)
  [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447)
  Fix double-shifted training loss in GitForCausalLM (huggingface#47395)
  Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895)
  Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379)
  Update bug report list (huggingface#46607)
  Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952)
  Optimize flash attention max seqlen computation in vision attention (huggingface#47170)
  fix: remove unreachable return in special token builder (huggingface#47420)
  Add Harry to slow CI (huggingface#47454)
  BLT: vectorize patch length processing (huggingface#47385)
  Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935)
  [Kimi] add integration tests (huggingface#47383)
  Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974)
  ...
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants