[feat] SID: in-model candidate-set matching (top-k last-layer candidate SIDs)

[WhiteSwan1]

What

Adds inference-time candidate-set matching to the SID (semantic-ID) generation models. Alongside the greedy SID, the residual quantizers emit the top-k last-layer candidate SID tuples + scores, gated by a new SidCandidateOutputConfig (strategy "last_layer_knn"): keep the greedy prefix and vary only the last layer across its nearest codes, best-first (slot 0 == the greedy SID). These candidates feed downstream codebook-collision handling. Opt-in and inference-only.

Changes

• proto: SidCandidateOutputConfig (enabled / topk / strategy) on SidRqvae and SidRqkmeans.
• QuantizeLayer: nearest_neighbors + a shared _topk_output builder (via the base lookup()), so both backends (VQ / K-Means) return the greedy pick plus the top-k ids/scores in eval from one place.
• ResidualQuantizer (now a BaseModule, to pick up the framework is_inference mode flag via set_is_inference) owns the candidate walk (_build_code_candidates) and a ResidualPassOutput result type. Candidate emission is decided internally by _should_output_candidates() = inference mode + config-enabled + a backend _candidates_available() hook (K-Means -> is_fitted), so an unfit codebook emits no candidates instead of crashing.
• SID models expose candidates via predict() at inference only.
• Tidies the quantizer subsystem: unified _quantize_layer, a single ResidualPassOutput walk result, one combined topk range guard, and argmin on the raw cdist in the K-Means training path (dropped a needless .pow(2)).

Testing

• Unit coverage for every piece (quantizers, models, losses); FX-symbolic-trace and export-trace tests cover the candidate path. 118 SID unit tests pass; ruff check + format clean.
• Tests are kept DRY: a shared faiss_unavailable in test_util + @unittest.skipIf, parameterized validation/shape tests, and shared construction/assert helpers.

[github-actions]

Review summary

Overall a clean, well-tested refactor. The shared _residual_pass → ResidualPassOutput / _residual_output factoring is nicely done, candidate output is correctly gated on is_inference + enabled + _candidates_available(), candidate_codes[:, 0] provably equals the greedy SID, and latents is correctly skipped on the inference path. A security/DDP pass found no real issues: topk is double-bounded (construction + nearest_neighbors) so there's no OOB into a codebook, the K-Means unfit path degrades to "no candidates" instead of crashing, and the candidate walk is inference-only so it introduces no new collective op / rank divergence.

Posted a few inline notes (none blocking):

• Two docstrings gone stale after the refactor — SidRqvae.predict ("codes only") and the ResidualQuantizer class doc (_quantize_layer is no longer a subclass primitive).
• A minor training-path perf regression: the STE distance matrix is now built with autograd tracking and discarded each step.
• A fragility in _init_candidate_output_config (candidate attrs left unset when disabled — safe only via short-circuit).
• The new @skipIf(*faiss_unavailable) decorators don't actually make the K-Means test files faiss-optional (top-level faiss import fails collection first).

Two smaller items, not worth an inline comment:

• Docs: the new candidate_output_config proto field isn't mentioned in docs/source/models/sid_rqvae.md / sid_rqkmeans.md, which otherwise enumerate every config field.
• Coverage: the candidate FX/export-trace path is tested for RQ-VAE but not for the K-Means backend, whose eval branch differs (KMeansQuantizeLayer.quantize has its own if self.training split + distances.pow(2)). Since export is how this feature is consumed offline, a K-Means export-trace test would close that gap.

[github-actions]

This skipIf(*faiss_unavailable) gate looks spurious: the test loads centroids directly via load_centroids_ and runs no FAISS fit, and since import faiss.contrib.torch_utils was moved into train_offline, construction/load/forward no longer need faiss. This is the only place the K-Means residual last-layer candidate ordering ([2, 3, 1]) and greedy-prefix invariance are asserted, so on a faiss-less CI runner the K-Means-specific candidate contract goes unverified. Recommend dropping the decorator so it runs unconditionally.

[github-actions]

Worth documenting the candidate_scores semantics more precisely, since this is a public inference output whose stated purpose is offline collision prevention: (1) the values are distances (lower = better, ascending), not similarity scores; and (2) they are the last-layer residual distance only and backend-scoped (squared-L2 for K-Means / L2-VQ, negative cosine for cosine-VQ), so they rank candidates within one item but are not directly comparable across items. A one-line caveat here (and/or in the proto) would prevent a downstream consumer from treating them as cross-item comparable scores.

[github-actions]

reserved 4, 5; in a brand-new message with no prior released version reads as leftover from development iterations. Either add a short comment noting what 4/5 were (to prevent accidental reuse) or drop the reservation — as-is it will puzzle future readers.

[github-actions]

Code review summary

Reviewed across code quality, performance, tests, docs, and safety. This is a solid, well-structured PR — the ResidualQuantizer base-class refactor is coherent, the candidate walk cleanly reuses the existing residual pass, and the feature is correctly gated (inference-only + config-enabled + backend _candidates_available). Verified specifically:

• No training-path regression. _should_output_candidates() short-circuits on the plain-Python is_inference flag before any candidate work; train/eval keep topk=1 and skip the candidate block entirely. Inference also correctly skips the latents stack.
• No correctness/safety issues. Config validation is double-layered (construction-time _init_candidate_output_config + call-time nearest_neighbors guard), fail-safe on unset config, and the unfit-K-Means path degrades to no-candidates rather than crashing. Checkpoint-integrity and DDP behavior are unaffected (candidate path is purely per-sample, no collectives).
• Docs are accurate — shapes, defaults, and the "slot 0 == greedy" contract all match the implementation.

A few noteworthy items (inline comments posted for the specific-line ones):

Test coverage gaps (all inference-only, low risk but worth closing):

• n_layers == 1 candidate path is untested — every candidate test uses n_layers=2/3. This is exactly the empty-prefix zero-width cat/expand case in _build_code_candidates where a shape/FX-trace regression would hide.
• SidRqkmeans has no FX/export-trace test with candidates enabled (the CPU-only inference model), while SidRqvae does (test_rqvae_export_trace_keeps_candidate_output).
• topk == n_embed (the exact upper boundary of the topk > n_embed guard) is never tested as a valid value.

Minor code quality: the refactor hoisted the encode primitive _quantize_layer into a concrete base method, but its decode mirror _lookup_code stays abstract with two byte-identical subclass overrides (return self.layers[layer_idx].lookup(code_idx)). It could be made concrete in the base too, symmetrizing the pair and removing the duplication.

Nothing here is blocking. Nice work on the DRY tests and the thorough validation.