Fix negative offsets in GPUBiasingMultiModel.remove_model

[kzos]

What

GPUBiasingMultiModel.remove_model leaves almost every unused reserved slot holding a negative offset whenever the removed model started at arena offset 0.

remove_model clears the departing model's own offsets to 0, then applies the shift across the whole table with the predicate offset < start_state. Unused slots also hold 0. 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, registering three boosting models and removing the first gives:

offsets before removal : [0, 6, 11, 0, 0, ...]
offsets after  removal : [-6, 0, 5, -6, -6, ...]
reserved slots holding a negative offset: 126 of 128

The underlying cause is that 0 carries two meanings in this table: a valid arena start address, and a slot that has never held a model. The shift cannot tell them apart, so it transforms slots that should not take part in it.

Impact, stated honestly

This is latent, not an active miscompute. I tried to make it produce a wrong answer and could not:

• the Triton kernel returns early for model_id < 0, so the -1 sentinel never dereferences its slot;
• rows carrying the sentinel come back all-zero on both the Triton and the non-Triton path, and the two agree bit for bit;
• a surviving model's scores are bit-identical before and after a removal;
• registering a new model after a removal reuses the freed id and scores correctly.

What is broken is the invariant that an offset is a non-negative index into the arena. Neither the Triton kernel nor _advance_pytorch re-checks model2active before turning an offset into a base pointer, so a stale or out-of-range positive id would find a negative base rather than failing closed. It is a hazard for the next consumer that does not gate, and it costs one masked subtraction to remove.

The change

Restrict the shift to active models, and clear 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.

Test

Extends test_add_then_remove_model to assert that every offset stays non-negative and that inactive slots carry a canonical empty descriptor.

Adds test_offsets_stay_non_negative_for_any_removal_position, parametrized over removal from the front, middle and back of the arena, asserting:

• every offset is non-negative;
• inactive slots hold a canonical empty descriptor;
• active extents stay inside the arena;
• a surviving model's advance() scores are unchanged by an unrelated removal.

The existing test checked only the survivor's shifted offset, which is why this was not caught.

Verification

Reproduced and verified on an NVIDIA RTX A6000 with Triton, using a standalone script that builds boosting trees from raw token id lists so it needs no checkpoint, tokenizer or audio. After the change: 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.

Two related things, deliberately not in this PR

Raising them here only so they are on the record; each deserves its own thread.

1. add_model fires reallocation_callbacks when the arena grows, and remove_model fires nothing, although both move data that a captured CUDA graph addresses. I could not demonstrate a wrong result from this, and the kernel loads the offsets through a pointer at run time, which suggests a replay sees the committed values. I would rather ask whether that is the intended contract than assert a bug.
2. There is no way to pre-size the arena. INIT_NUM_MODELS, INIT_NUM_ARCS and INIT_NUM_STATES are fixed and growth doubles, with no configurable maximum and no reserve(). A server that registers a phrase list per session cannot guarantee that no reallocation, and therefore no CUDA-graph reset, happens mid-stream for the other sessions in the batch. A capacity knob would let deployments trade memory for predictability.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[kzos]

Adding the local verification, since CI on NVIDIA runners needs a maintainer's /ok to test and this is the evidence that would otherwise cost runner time to obtain.

I ran the assertions from the new test against the unpatched and patched code on both devices. The parametrization matters: the defect fires only when the removed model sits at arena offset 0, which is exactly what the analysis predicts, so removals from the middle and back pass either way.

Unpatched:

remove_index=0: FAIL -> states_offset >= 0, arcs_offset >= 0,
                        inactive states_offset == 0, inactive arcs_offset == 0
remove_index=1: PASS
remove_index=2: PASS

Patched, CPU (_advance_pytorch) and CUDA (Triton):

remove_index=0: PASS
remove_index=1: PASS
remove_index=2: PASS

Both device paths, including the survivor-scores-unchanged assertion, which is the one that would catch a fix that corrected the offsets while breaking the compaction.

Separately, on the same hardware with a standalone reproducer: negative offsets go from 126 of 128 to 0, a surviving model's advance() output stays bit-identical across the removal, slot reuse after removal still works, and the Triton and non-Triton paths continue to agree with each other.

Happy to adjust the test's shape or placement if you would rather it lived somewhere else.