From 0e38652d095e3d6d512d85f367c2746d0daaba70 Mon Sep 17 00:00:00 2001 From: Zaheer Sheriff K Date: Wed, 9 Sep 2026 18:14:49 +0300 Subject: [PATCH 1/2] Fix negative offsets in GPUBiasingMultiModel.remove_model Removing a boosting model leaves every unused reserved slot holding a negative offset when the removed model started at arena offset 0. remove_model clears the departing model's own offsets to 0 and then applies the shift to the whole table with the predicate `offset < start_state`. Unused slots also hold 0, and when the removed model started at 0 that predicate is false for them, so each one is decremented to -num_states / -num_arcs. With the default reserve of 128 slots, removing the first of three registered models leaves 126 slots negative: offsets before removal : [0, 6, 11, 0, 0, ...] offsets after removal : [-6, 0, 5, -6, -6, ...] The root cause is that 0 means two different things in this table: a valid arena start address, and an empty slot that has never held a model. The shift cannot distinguish them, so it transforms slots that should not take part in it. This is currently latent. The Triton kernel returns early for a negative model_id, so the -1 sentinel never dereferences its slot, and callers gate on model2active, so no configuration I could construct produces a wrong score. Surviving models score bit-identically across a removal both before and after this change. But the invariant that an offset is a non-negative index into the arena is violated for almost every slot, and neither the kernel nor _advance_pytorch re-checks model2active before turning an offset into a pointer, so a stale or out-of-range positive id would find a negative base. The fix restricts the shift to ACTIVE models and clears the removed model's own offsets last. model2active[model_id] is already False at that point, so the removed model is excluded from the mask. Using an explicit active mask rather than only reordering the two statements makes the intent readable: inactive slots are not participants in the relocation transform. Adds a parametrized regression test covering removal from the front, middle and back of the arena, asserting that every offset stays non-negative, that inactive slots carry a canonical empty descriptor, that active extents stay inside the arena, and that a surviving model's scores are unchanged by the removal. The existing test only checked the survivor's shifted offset, which is why this was not caught. Verified on an NVIDIA RTX A6000 with Triton: negative offsets go from 126 to 0, surviving models remain bit-identical, slot reuse after removal still works, and the Triton and non-Triton paths continue to agree. Signed-off-by: Zaheer Sheriff K --- .../context_biasing/biasing_multi_model.py | 25 ++++---- .../asr/decoding/test_biasing_multi_model.py | 61 +++++++++++++++++++ 2 files changed, 72 insertions(+), 14 deletions(-) diff --git a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py index 694f6f190a24..d5af66da2df1 100644 --- a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py +++ b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py @@ -562,22 +562,19 @@ def remove_model(self, model_id: int): self.model2num_states[model_id] = 0 self.model2num_arcs[model_id] = 0 self.model2num_arcs_extended[model_id] = 0 - # shift model offsets + # Shift the offsets of the models that lived after the removed one. + # Only ACTIVE models take part in the shift. An unused slot holds offset 0, and 0 is also a + # legitimate arena address, so a shift applied to every slot underflows every unused slot to a + # negative offset whenever the removed model started at 0. `model2active[model_id]` is already + # False here, so the removed model is excluded from the mask and its own offsets are cleared last. + states_shift_mask = self.model2active & (self.model2states_offset >= start_state) + arcs_shift_mask = self.model2active & (self.model2arcs_offset >= start_arc) + self.model2states_offset[states_shift_mask] -= num_states + self.model2arcs_offset[arcs_shift_mask] -= num_arcs + + # clear the removed model's own offsets last self.model2states_offset[model_id] = 0 self.model2arcs_offset[model_id] = 0 - # shift states and arcs offsets - torch.where( - self.model2states_offset < start_state, - self.model2states_offset, - self.model2states_offset - num_states, - out=self.model2states_offset, - ) - torch.where( - self.model2arcs_offset < start_arc, - self.model2arcs_offset, - self.model2arcs_offset - num_arcs, - out=self.model2arcs_offset, - ) def get_init_states(self, batch_size: int, bos=True) -> torch.Tensor: """ diff --git a/tests/collections/asr/decoding/test_biasing_multi_model.py b/tests/collections/asr/decoding/test_biasing_multi_model.py index c4b0dd1b5cd8..6af9a41e3b1e 100644 --- a/tests/collections/asr/decoding/test_biasing_multi_model.py +++ b/tests/collections/asr/decoding/test_biasing_multi_model.py @@ -143,6 +143,67 @@ def test_add_then_remove_model(self, stt_en_conformer_transducer_small, device: assert multi_model.model2states_offset[model_id2].item() == 0 assert multi_model.model2arcs_offset[model_id2].item() == 0 + # Every offset must stay a valid, non-negative arena index, including the reserved slots that + # have never held a model. Before the ordering fix in remove_model, removing the model that + # started at offset 0 underflowed every zero-valued slot to -num_states / -num_arcs. + assert torch.all(multi_model.model2states_offset >= 0) + assert torch.all(multi_model.model2arcs_offset >= 0) + + # Inactive slots carry a canonical, empty descriptor. + inactive = ~multi_model.model2active + assert torch.all(multi_model.model2states_offset[inactive] == 0) + assert torch.all(multi_model.model2arcs_offset[inactive] == 0) + assert torch.all(multi_model.model2num_states[inactive] == 0) + assert torch.all(multi_model.model2num_arcs_extended[inactive] == 0) + + @pytest.mark.unit + @pytest.mark.with_downloads + @pytest.mark.parametrize("device", DEVICES) + @pytest.mark.parametrize("remove_index", [0, 1, 2]) + def test_offsets_stay_non_negative_for_any_removal_position( + self, stt_en_conformer_transducer_small, device: torch.device, remove_index: int + ): + """Removing a model from any position leaves every offset a valid arena index. + + Removing the model at the front of the arena is the interesting case: its start offset is 0, + which every unused reserved slot also holds, so a shift that does not exclude inactive slots + drives all of them negative. + """ + tokenizer = stt_en_conformer_transducer_small.tokenizer + vocab_size = tokenizer.vocab_size + + multi_model = GPUBiasingMultiModel(vocab_size=vocab_size).to(device) + phrase_lists = [["alpha", "beta"], ["gamma"], ["delta", "epsilon", "zeta"]] + model_ids = [ + multi_model.add_model(create_boosting_model(phrases, tokenizer, device), alpha=1.0) + for phrases in phrase_lists + ] + + # A surviving model's scores must be unchanged by an unrelated removal. + survivor = model_ids[(remove_index + 1) % len(model_ids)] + states = multi_model.get_init_states(batch_size=2, bos=True) + survivor_ids = torch.full((2,), survivor, dtype=torch.long, device=device) + scores_before, _ = multi_model.advance(states=states, model_ids=survivor_ids) + scores_before = scores_before.clone() + + multi_model.remove_model(model_ids[remove_index]) + + assert torch.all(multi_model.model2states_offset >= 0) + assert torch.all(multi_model.model2arcs_offset >= 0) + + inactive = ~multi_model.model2active + assert torch.all(multi_model.model2states_offset[inactive] == 0) + assert torch.all(multi_model.model2arcs_offset[inactive] == 0) + + active = multi_model.model2active + assert torch.all( + multi_model.model2states_offset[active] + multi_model.model2num_states[active] + <= multi_model.num_states_total + ) + + scores_after, _ = multi_model.advance(states=states, model_ids=survivor_ids) + assert torch.equal(scores_before, scores_after) + @pytest.mark.unit @pytest.mark.with_downloads @pytest.mark.parametrize("device", DEVICES) From ab257845a28d7f528f17e3b39c587ab7cf70e6ea Mon Sep 17 00:00:00 2001 From: Zaheer Sheriff K Date: Mon, 14 Sep 2026 18:34:21 +0300 Subject: [PATCH 2/2] Use torch.where(out=) for the offset shift instead of masked assignment Review feedback: the shift is a frequent operation on a small table, and a boolean-mask assignment forces a device synchronisation because the number of selected elements is only known on the host. torch.where with out= computes the same result at a fixed shape. Verified on an NVIDIA B300 that the two forms agree exactly -- identical offsets for removal from the front, middle and back of the arena, on CPU and CUDA, and both remove all 126 negative offsets the current code leaves behind -- and, under torch.cuda.set_sync_debug_mode(2), that the masked form synchronises and this one does not. Signed-off-by: Zaheer Sheriff K --- .../parts/context_biasing/biasing_multi_model.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py index d5af66da2df1..b46416fcd885 100644 --- a/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py +++ b/nemo/collections/asr/parts/context_biasing/biasing_multi_model.py @@ -567,10 +567,18 @@ def remove_model(self, model_id: int): # legitimate arena address, so a shift applied to every slot underflows every unused slot to a # negative offset whenever the removed model started at 0. `model2active[model_id]` is already # False here, so the removed model is excluded from the mask and its own offsets are cleared last. - states_shift_mask = self.model2active & (self.model2states_offset >= start_state) - arcs_shift_mask = self.model2active & (self.model2arcs_offset >= start_arc) - self.model2states_offset[states_shift_mask] -= num_states - self.model2arcs_offset[arcs_shift_mask] -= num_arcs + torch.where( + self.model2active & (self.model2states_offset >= start_state), + self.model2states_offset - num_states, + self.model2states_offset, + out=self.model2states_offset, + ) + torch.where( + self.model2active & (self.model2arcs_offset >= start_arc), + self.model2arcs_offset - num_arcs, + self.model2arcs_offset, + out=self.model2arcs_offset, + ) # clear the removed model's own offsets last self.model2states_offset[model_id] = 0