Merge pull request #114 from generative-computing/bugfix/51-52-audio-prompt-placeholder-invariants

Fixes three defects in the ASR audio cascade, plus the packaging gap that blocked testing them.

Author: aviv1ron1

.github/workflows/gpu-tests.yaml

@@ -16,7 +16,10 @@ jobs:
         with:
           enable-cache: true
 
-      - run: uv sync --frozen --group dev --extra hf --extra vllm --extra compose
+      # `--extra audio` is redundant with the dev group (which now includes it) but
+      # stated explicitly: the audio tests need soundfile/librosa at runtime, and a
+      # group refactor should not silently drop them again.
+      - run: uv sync --frozen --group dev --extra hf --extra vllm --extra compose --extra audio
 
       - name: Run GPU tests
         run: |

docs/AUDIO.md

@@ -9,6 +9,22 @@ intentionally simple and requires no training. The "proper" upgrade (feeding a
 trained projection of a speech encoder's embeddings straight into the LLM) reuses
 the same hooks — see [Design](#design) below.
 
+## Installing
+
+The audio path needs `soundfile` and `librosa` on top of the vLLM backend — they
+decode and resample the incoming waveform. They live in the `audio` extra, which is
+**not** part of `vllm`, so a plain `uv sync --extra vllm` gives you a checkpoint that
+fails on any non-16 kHz input:
+
+```bash
+# Serving an audio-enabled checkpoint
+uv sync --extra vllm --extra audio     # or --extra vllm20 --extra audio
+
+# Development / running the test suite (the dev groups include audio already)
+uv sync --group dev                    # vLLM 0.19.x
+uv sync --group dev-vllm20             # vLLM 0.20.x
+```
+
 ## Building an audio-enabled checkpoint
 
 Add `--enable-audio` when composing:
@@ -183,6 +199,12 @@ Per request, before the scheduler allocates KV cache:
 3. A `PromptReplacement` swaps the `<|audio|>` marker for those transcript token
    ids. The scheduler then sizes KV for the **real** length — the audio "window"
    is variable and decided at runtime, not reserved in advance.
+   A clip with no recognizable speech in it — silence, music, noise, or a clip
+   too short to hold a word — transcribes to the empty string. Since every audio
+   item has to occupy at least one prompt position (vLLM discards a zero-length
+   placeholder and then rejects the request), those clips are replaced with a
+   single space instead: the model sees an audio turn that said nothing, rather
+   than an error.
 4. The model's `embed_multimodal` supplies embeddings for those positions. In the
    alpha that is simply the transcript's own token embeddings (identical to
    embedding them as text). **This is the seam the future encoder reuses:** swap
@@ -220,6 +242,17 @@ adapter behaves identically to the text equivalent.
 
 ## Tests
 
+Everything on the audio path carries the `audio` marker, so the whole tier selects
+in one command regardless of where the tests live:
+
+```bash
+# All audio tests (13 of them need a GPU and a real checkpoint)
+pytest -m audio -v -s --tb=short
+
+# CPU tier only — runs in a few seconds
+pytest -m "audio and not gpu" -v -s --tb=short
+```
+
 - `tests/unit/test_asr.py` — CPU unit tests for the ASR backend (audio coercion,
   resampling, transcription with a mocked pipeline, pipeline-kwargs cache keying,
   and per-request decode-kwargs resolution). No GPU/vLLM required.

pyproject.toml

@@ -53,13 +53,17 @@ markers = [
     "slow: takes > 30s",
     "deep: expensive code-theory tests (m=8 / 256-dim); run with: pytest -m deep",
     "requires_model: needs a real model checkpoint",
+    "audio: exercises the audio/ASR path; run with: pytest -m audio",
 ]
 
 [dependency-groups]
 vllm19     = ["vllm>=0.19.1,<0.20.0"]
 vllm20     = ["vllm>=0.20.0,<0.21.0"]
-dev        = ["pytest", "pytest-cov", { include-group = "vllm19" }, "granite-switch[hf,compose]"]
-dev-vllm20 = ["pytest", "pytest-cov", { include-group = "vllm20" }, "granite-switch[hf,compose]"]
+# `audio` is included so the audio tests can actually run: the ASR path needs
+# soundfile/librosa at runtime, and no group pulled them in before (integration
+# tests failed with ModuleNotFoundError on a synced pod).
+dev        = ["pytest", "pytest-cov", { include-group = "vllm19" }, "granite-switch[hf,compose,audio]"]
+dev-vllm20 = ["pytest", "pytest-cov", { include-group = "vllm20" }, "granite-switch[hf,compose,audio]"]
 test       = ["pytest", "pytest-cov", "bitsandbytes", "optimum-quanto", { include-group = "dev" }]
 
 [tool.uv]

src/granite_switch/composer/compose_granite_switch.py

@@ -67,6 +67,7 @@
     configure_chat_template,
     get_alora_first_invocation_token_id,
 )
+from granite_switch.composer.validator import validate_control_lut
 from granite_switch.config import ASR_DTYPES
 
 # ---------------------------------------------------------------------------
@@ -871,7 +872,14 @@ def build():
         or args.asr_chunk_length_s is not None
         or args.asr_chunk_overlap_s is not None
     )
-    audio_token_id = add_audio_token(tokenizer) if audio_enabled else None
+    # The control tokens are re-passed so this call doesn't drop them from the
+    # tokenizer's additional-special-tokens list (add_special_tokens replaces
+    # that list rather than extending it).
+    audio_token_id = (
+        add_audio_token(tokenizer, keep_special_tokens=special_tokens)
+        if audio_enabled
+        else None
+    )
 
     # Configure chat template with adapter mappings (Granite models only).
     # Non-Granite models preserve the upstream template verbatim because
@@ -1011,6 +1019,23 @@ def build():
     new_embed_size = model.model.embed_tokens.weight.shape[0]
     print(f"Embeddings resized: {old_embed_size} -> {new_embed_size}")
 
+    # The switch sized its control->substitute table from the pre-resize
+    # config.vocab_size (copied from the base model), so the resize above leaves
+    # it short of the config this checkpoint will ship with. Re-derive it, then
+    # assert the two agree — see validate_control_lut for why a mismatch is not
+    # something the loader can recover from.
+    switch = getattr(model.model, "switch", None)
+    if switch is not None:
+        lut = getattr(switch, "control_to_substitute_lut", None)
+        if lut is not None and lut.numel() != model.config.vocab_size:
+            old_lut_size = lut.numel()
+            switch.rebuild_control_to_substitute_lut(model.config)
+            print(
+                "Switch control LUT rebuilt: "
+                f"{old_lut_size} -> {switch.control_to_substitute_lut.numel()}"
+            )
+    validate_control_lut(model)
+
     print(f"\nStep 3 complete in {time.time() - step_start:.2f}s")
 
     return (

src/granite_switch/composer/tokenizer_setup.py

@@ -94,21 +94,38 @@ def add_control_tokens(
     return adapter_token_ids, special_tokens
 
 
-def add_audio_token(tokenizer, marker: str = "<|audio|>") -> int:
+def add_audio_token(
+    tokenizer,
+    marker: str = "<|audio|>",
+    keep_special_tokens: list[str] | None = None,
+) -> int:
     """Add the audio placeholder marker token to the tokenizer.
 
     Used for the audio cascade: this single special token is placed in the
     prompt and the vLLM ASR processor replaces it with the transcript tokens at
     request time (see granite_switch.vllm.audio). Registering it as one special
     token keeps the processor's prompt-replacement match clean.
 
+    ``keep_special_tokens`` must list every token an earlier
+    ``add_special_tokens({"additional_special_tokens": ...})`` call registered —
+    in practice the adapter control tokens from :func:`add_control_tokens`.
+    That call *replaces* the additional-special-tokens list instead of appending
+    to it, and transformers exposes no way to read the current list back, so any
+    token not re-passed here silently drops out of ``all_special_tokens`` and
+    out of the saved ``tokenizer_config.json``. Re-passing an already-added
+    token is free: it keeps its id and does not grow the vocabulary.
+
     Must be called before the model's embedding resize so the new row is sized
     in. Returns the marker's token id.
     """
     print(f"\nAdding audio marker token: {marker}")
-    tokenizer.add_special_tokens({"additional_special_tokens": [marker]})
+    # Marker last so it takes the next free id and the kept tokens keep theirs.
+    kept = [t for t in (keep_special_tokens or []) if t != marker]
+    tokenizer.add_special_tokens({"additional_special_tokens": [*kept, marker]})
     token_id = tokenizer.convert_tokens_to_ids(marker)
     print(f"  {marker}: {token_id}")
+    if kept:
+        print(f"  (preserved {len(kept)} existing special token(s))")
     return token_id
 
 

src/granite_switch/composer/validator.py

@@ -13,6 +13,39 @@
 from .arch import ArchDescriptor
 
 
+def validate_control_lut(model) -> None:
+    """Check the switch's control->substitute table matches the shipped config.
+
+    The table is a persistent buffer sized from ``config.vocab_size``, so a
+    checkpoint whose buffer length disagrees with its own ``config.json`` is
+    internally inconsistent. Loading one is not a graceful degradation:
+    ``from_pretrained`` discards the mismatched tensor and leaves the buffer as
+    uninitialised memory, which turns every token into a "control" token and
+    sends out-of-range ids into the embedding gather — surfacing as an opaque
+    CUDA ``srcIndex < srcSelectDimSize`` device-side assert far from the cause.
+
+    Raises:
+        ValueError: if the table length differs from ``config.vocab_size``.
+    """
+    switch = getattr(getattr(model, "model", None), "switch", None)
+    lut = getattr(switch, "control_to_substitute_lut", None)
+    if lut is None:
+        return  # no token-exchange mapping on this model
+
+    expected = getattr(model.config, "vocab_size", None)
+    if expected is None or lut.numel() == expected:
+        return
+
+    raise ValueError(
+        f"control_to_substitute_lut has {lut.numel()} entries but "
+        f"config.vocab_size is {expected}. Saving this model would produce a "
+        f"checkpoint that cannot be loaded correctly. The table is derived from "
+        f"vocab_size and the adapter token ids, so rebuild it after any "
+        f"vocabulary change with "
+        f"switch.rebuild_control_to_substitute_lut(model.config)."
+    )
+
+
 def validate_all_parameters(
     model,
     arch: ArchDescriptor,

src/granite_switch/hf/switch/single.py

@@ -19,6 +19,35 @@
 from transformers.models.granite.modeling_granite import eager_attention_forward
 
 
+def build_control_to_substitute_lut(config) -> torch.Tensor | None:
+    """Derive the control->substitute lookup table from *config*.
+
+    Shape ``[max(vocab_size, max_ctrl_id + 1)]``: ``-1`` at every non-control id
+    and the substitute id at each control slot. The ``max`` keeps every control
+    id addressable even when ``vocab_size`` lags the tokenizer.
+
+    Returns ``None`` when *config* carries no token-exchange mapping, in which
+    case the switch leaves ``input_ids`` untouched.
+
+    Single source of truth for the sizing rule: the table is a pure function of
+    ``vocab_size``, ``adapter_token_ids`` and ``adapter_substitute_token_ids``,
+    so anything that changes those must re-derive it (see
+    :meth:`SingleSwitch.rebuild_control_to_substitute_lut`).
+    """
+    if config is None:
+        return None
+    ctrl_ids = getattr(config, "adapter_token_ids", None)
+    sub_ids = getattr(config, "adapter_substitute_token_ids", None)
+    if not ctrl_ids or not sub_ids:
+        return None
+
+    lut_size = max(getattr(config, "vocab_size", 0), max(ctrl_ids) + 1)
+    lut = torch.full((lut_size,), -1, dtype=torch.long)
+    for ctrl_id, sub_id in zip(ctrl_ids, sub_ids):
+        lut[ctrl_id] = sub_id
+    return lut
+
+
 class SingleSwitch(nn.Module):
     """Single-head attention-based switch for adapter selection.
 
@@ -89,22 +118,44 @@ def __init__(
         # control-token positions carry the substitute id by the time the
         # decoder embeds them. The decoder is then oblivious — it just calls
         # embed_tokens(input_ids) and gets the right result by construction.
-        if (
-            config is not None
-            and getattr(config, "adapter_token_ids", None) is not None
-            and getattr(config, "adapter_substitute_token_ids", None) is not None
-        ):
-            ctrl_ids = config.adapter_token_ids
-            sub_ids = config.adapter_substitute_token_ids
-            max_ctrl_id = max(ctrl_ids)
-            lut_size = max(getattr(config, "vocab_size", 0), max_ctrl_id + 1)
-            lut = torch.full((lut_size,), -1, dtype=torch.long)
-            for ctrl_id, sub_id in zip(ctrl_ids, sub_ids):
-                lut[ctrl_id] = sub_id
+        lut = build_control_to_substitute_lut(config)
+        if lut is not None:
             self.register_buffer("control_to_substitute_lut", lut)
         else:
             self.control_to_substitute_lut = None
 
+    def rebuild_control_to_substitute_lut(self, config=None) -> bool:
+        """Re-derive the control->substitute table after a vocabulary change.
+
+        ``__init__`` sizes the table from ``config.vocab_size``, so anything that
+        grows the vocabulary afterwards — notably
+        ``resize_token_embeddings`` when compose adds control and marker tokens —
+        leaves the buffer shorter than the config it will be saved alongside.
+
+        That matters because the buffer is persistent. On ``from_pretrained``,
+        a stored tensor whose shape disagrees with the freshly-constructed one is
+        discarded and the buffer is left as uninitialised memory (there is no
+        ``_init_weights`` rule for it), so every id reads as a control id and the
+        rewrite sends out-of-range ids into the embedding gather. Re-derive the
+        table before saving so the checkpoint and its config agree.
+
+        Returns ``True`` if a table was rebuilt, ``False`` if this switch has no
+        token-exchange mapping to rebuild.
+        """
+        lut = build_control_to_substitute_lut(
+            config if config is not None else self.config
+        )
+        if lut is None:
+            return False
+
+        existing = getattr(self, "control_to_substitute_lut", None)
+        if existing is not None:
+            lut = lut.to(device=existing.device)
+            self.control_to_substitute_lut = lut
+        else:
+            self.register_buffer("control_to_substitute_lut", lut)
+        return True
+
     @property
     def num_cache_layers(self) -> int:
         """Number of cache slots used."""

src/granite_switch/vllm/audio/processor.py

@@ -42,6 +42,10 @@
 
 AUDIO_MARKER = "<|audio|>"
 _TARGET_SR = 16_000
+# Stands in for a transcript with no speech in it. Must tokenize to at least one
+# token: vLLM discards a zero-length placeholder and then reports the item as
+# missing, so every audio item has to contribute something to the prompt.
+_EMPTY_TRANSCRIPT_TEXT = " "
 # Keeps the transcript budget finite if max_model_len cannot be read.
 _FALLBACK_CONTEXT_LEN = 8192
 _DUMMY_AUDIO_SECONDS = 5
@@ -159,7 +163,14 @@ def _transcribe(
         generate_kwargs: Mapping[str, object] | None = None,
     ) -> list[int]:
         """Transcribe one audio item to token ids. Never truncated here — an
-        oversized prompt is rejected by vLLM's own length check."""
+        oversized prompt is rejected by vLLM's own length check.
+
+        Guaranteed non-empty: silence, music, non-speech or a clip too short to
+        contain a word all transcribe to ``""``, and a zero-length replacement
+        would make vLLM drop the placeholder and reject the request. Those clips
+        get :data:`_EMPTY_TRANSCRIPT_TEXT` instead, so the model sees an audio
+        turn that simply said nothing.
+        """
         transcriber = get_transcriber(
             model_id=self.info._asr_model_id(),
             device=self.info._asr_device(),
@@ -176,7 +187,19 @@ def _transcribe(
             chunk_overlap_s=self.info._asr_chunk_overlap_s(),
         )
         tokenizer = self.info.get_tokenizer()
-        return tokenizer.encode(text, add_special_tokens=False)
+        if not text or not text.strip():
+            text = _EMPTY_TRANSCRIPT_TEXT
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        if not ids:
+            # Only reachable if the tokenizer drops _EMPTY_TRANSCRIPT_TEXT
+            # entirely. Fail loudly rather than emit a zero-length placeholder,
+            # which surfaces as an opaque "found 0 prompt placeholders".
+            raise ValueError(
+                f"Tokenizer produced no tokens for {text!r}; cannot build a "
+                "prompt placeholder for this audio item. Choose an "
+                "_EMPTY_TRANSCRIPT_TEXT this tokenizer encodes to >=1 token."
+            )
+        return ids
 
     def _call_hf_processor(
         self,
@@ -215,6 +238,31 @@ def _call_hf_processor(
             tensor_type="pt",
         )
 
+    def _hf_processor_applies_updates(
+        self,
+        prompt_text: str,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        tokenization_kwargs: Mapping[str, object],
+    ) -> bool:
+        """Always False: ``_call_hf_processor`` leaves the marker in place.
+
+        The base implementation returns True for raw (non-embedding) items,
+        which tells vLLM the processor already expanded the placeholder itself —
+        so vLLM skips applying our ``PromptReplacement`` and merely *searches*
+        the returned prompt for the transcript token ids. They are not there, and
+        it raises ``Expected there to be 1 audio prompt placeholders ... found 0``.
+
+        Unlike a real HF processor (e.g. Ultravox's), we tokenize the prompt with
+        the ``<|audio|>`` marker untouched and hand the transcript back out of
+        band, so the replacement must be applied by vLLM.
+
+        Only the uncached path consults this hook; the cached path already
+        hardcodes False, which is why audio works with the default
+        ``mm_processor_cache_gb=4`` and breaks under ``--mm-processor-cache-gb 0``.
+        """
+        return False
+
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,