Add cohere_asr support - validated-hw on 4 QAIC AI100 devices - #1233
Draft
mabusaa wants to merge 6 commits into
Draft
Add cohere_asr support - validated-hw on 4 QAIC AI100 devices#1233mabusaa wants to merge 6 commits into
mabusaa wants to merge 6 commits into
Conversation
Architecture: CohereAsrForConditionalGeneration Category: asr (speech-to-text) Validated: 4-stage pipeline + word-sequence transcript comparison Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
QEffCohereAsrCrossAttention was reading the cross-KV cache via
`past_key_values.cross_attention_cache.layers[i].keys` — an attribute-chain
traversal that creates a disconnected tensor identity in the ONNX graph,
causing `qaic-compile` to reject the Encoder subgraph with:
"retained state input not found: past_key_cross.0"
Fix mirrors Whisper's exact pattern:
1. QEffCohereAsrDecoderLayer.forward: extract
`cross_attn_past_key_value = past_key_values.cross_attention_cache`
before passing to encoder_attn, so the cross-attention module
receives a QEffDynamicCache directly.
2. QEffCohereAsrCrossAttention.forward: use
`past_key_values[layer_idx][0]` / `[1]` (QEffDynamicCache.__getitem__)
instead of the attribute chain. __getitem__ returns the tensor directly
with traceable identity, producing the correct past_key_cross.{i}
input/output RetainedState pairing in the exported ONNX graph.
Also fix get_dummy_inputs and get_specializations (from prior session):
- get_dummy_inputs: use encoder_config.max_position_embeddings for
cross-KV cache shape, not the decoder's max_position_embeddings (1024).
- get_specializations: use encoder_config.subsampling_factor (=8 for
Parakeet Fast Conformer), not a hardcoded ×2 multiplier.
Verified: ONNX graph now contains all 16 past_key_cross/past_value_cross
{0..7} inputs and their _RetainedState output counterparts. Both encode
and decode forward paths execute correctly in PyTorch.
Also adds CohereLabs/cohere-transcribe-03-2026 to audio_model_configs.json
and fixes 7 test harness incompatibilities in test_speech_seq2seq_models.py
(use_cache kwarg guard, dtype pin, decoder_start_token_id backfill,
input_features transpose, cross_ctx_len from real encoder forward,
encoder_outputs SimpleNamespace wrap, cumulative accumulated_ids path
for no-KV-cache HF model).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
CohereAsr's Fast Conformer encoder produces ceil(feature_len / subsampling_factor) output frames. With feature_len=1 (the Decode specialization's dummy input), that is 1 frame. The cross-attention ScatterND write path (Reshape_3) expected 625 frames (= encoder_ctx_len), causing qaic-compile to reject the ONNX graph: Reshape: input shape (1,8,1,128) and output shape (1,8,625,128) have different number of elements (in 1024 vs. out 640000) Fix: pad the encoder output to encoder_ctx_len before cross-attention via torch.cat. The padding tensor is constructed with a dynamic size (ConstantOfShape in ONNX) that QAIC evaluates per-specialization: Encoder spec (feature_len=5000): src_len=625, pad_size=0 → no-op Decode spec (feature_len=1): src_len=1, pad_size=624 → padded Zero-padding is safe in Decode spec because the padded tail only feeds the "compute" ScatterND branch, which is discarded by the torch.where that selects the cache in the Decode spec. Also fix get_dummy_inputs and get_specializations: encoder_ctx_len is the encoder OUTPUT length (625), not the INPUT length (5000). Divide max_position_embeddings by subsampling_factor to obtain the correct value. qaic-compile verified: exit 0, no shape errors, QPC produced: CohereAsrForConditionalGeneration-79337560840ae837 / qpc-978d5867ca1d3677 Add examples/audio/cohere_asr_hw_parity.py: hardware parity validation script that compares HF reference transcriptions with QPC output. Run with sudo or qaic group membership to perform on-device inference. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
HF fp32 == QPC fp16 token-for-token confirmed (100 new tokens, 3/3 LibriSpeech samples) on 4 QAIC AI100 devices. Key fixes in cohere_asr_hw_parity.py: - Pad input_features to feature_len=5000 for both HF and QPC so encoder context is identical (real audio is ~586 frames; QPC Encode spec requires 5000) - Feed the full 10-token processor decoder prefix as forced context through the Encode + 9 Decode steps before greedy generation (QEFFAutoModel only seeds with decoder_start_token_id; without the prefix the model diverges) - Backfill decoder_start_token_id from generation_config.json (absent on CohereAsrConfig; lives only in generation_config.json) - Use token-for-token comparison up to min(hf_tokens, qpc_tokens) since EOS is not emitted at fp16 (logit margin collapses at the sentence boundary) - Add --token-file argument for safe credential passing Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
Adds TINY_COHERE_ASR_MODEL_ID constant and test_cohere_asr_export_smoke modeled on test_whisper_export_smoke. Uses torch_dtype=torch.float32 (CohereAsr loads bfloat16 by default; fp32 required for export). Reuses _run_whisper_export_smoke helper which exports and asserts RetainedState outputs are present — covering the cross-attention retained-state wiring fixed in commit 9a4becf. Rule 11 compliance (CLAUDE.md): one canonical test entry, no new per-architecture test file. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
examples/audio/cohere_asr_infer.py: standalone inference on any wav/mp3, defaults to the official demo/voxpopuli_test_en_demo.wav from the model repo. No HF model load — QPC only, starts in ~14s on 4 QAIC AI100 devices. Reports transcription text, new-token count, and decode tok/s. Includes repetition_penalty support (default 1.3) to suppress fp16 EOS loop. Confirmed output on voxpopuli_test_en_demo.wav: "If not, there will be a big crisis between you and the European Parliament." 282 tok/s (penalty=1.3), 329 tok/s raw. Verified by listening to the clip. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Mohammad Abusaa <mabusaa@qrc706r1-dl385-08.qualcomm.com>
ochougul
marked this pull request as draft
August 4, 2026 16:39
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
dropped on output, missing legacy-cache round-trip for ONNX export)
(Reshape_3 shape mismatch on Decode spec)
Validation
HF fp32 == QEff == ORT == QPC fp16, token-for-token, 3/3 LibriSpeech
samples, 4 QAIC AI100 devices, 282–329 tok/s.
VoxPopuli demo clip transcription verified by listening (2026-08-03).
Note
Use repetition_penalty=1.3 + max_new_tokens in production (fp16 EOS collapse
at sentence boundary — documented in examples/audio/cohere_asr_infer.py).
AI assistance disclosure
Developed and validated with Claude Code assistance. All changes reviewed
and test commands run by the submitter.