[feat] prompt-native generative recommendation

[WhiteSwan1]

Ready for review. Multi-rank is still untested and two items from the
design remain unbuilt (see Not done below); both are called out rather than
silently pending.

Replaces the SID-feature-based generative stack with a prompt-native one: a new
prompt_config owns the template, slots, SID space and tokenizer, and
model_config keeps only what belongs to the LM.

Why

The per-level SID offset used to be applied twice — once in the SID tool's
sid_offset_codes, once again in SidFeature._parse — with no hash covering the
pair. The collision tool now emits an offset_codebook column, so the offset
arrives already applied and that second implementation can go.

What remains inside tzrec is a uniform + base_vocab, which is not
position-dependent and therefore needs no bespoke feature type. SidFeature is
deleted; an INLINE slot is an ordinary sequence_raw_feature and a PROJECTED one
is a sequence_id_feature with an injected num_buckets.

The offset form is also self-validating: a raw code 100 is legal under codebook
4096 and 8192 alike, so a tool/config mismatch passes a range check, whereas
offset codes carry cumsum(codebook[:l]) inside the value and throw every level

= 1 out of band.

What is here

area
tzrec/protos/prompt.proto	PromptConfig, PromptSlot, PromptProjection, SidSpace
tzrec/prompt/compile.py	compiles the config into a ResolvedSidSpace, a PromptPlan, a ProjectionPlan and an extended tokenizer
tzrec/prompt/assembler.py	builds the packed varlen token stream in the dataloader worker
tzrec/prompt/persist.py	writes the contract beside the weights, checks it on restore
tzrec/models/prompt_generative_qwen.py	the model: inputs_embeds forward, index_copy scatter, band-constrained decode
tzrec/modules/prompt_projection.py	slot width -> LM hidden size
docs/source/models/prompt_generative_qwen.md	operator manual (Chinese)

Removed: sid_feature.py, generative_model.py, generative_qwen.py,
generative_model.proto and their tests. Nothing is kept for compatibility.

Two properties worth reviewing specifically:

• The compiler resolves no dimension. No artifact stores in_dim/out_dim;
  the model resolves both at __init__ from group_total_dim and the backbone
  config, so one compiled plan is valid across model sizes.
• No shape is derived from device tensors. The collator computes
  max_seqlen on the host, so the model never calls lengths.max(). It is
  carried in additional_infos, which Batch.to() migrates, so reading it back
  does cost one host sync per step — a known performance item, not a correctness
  one, tracked in the review triage.

Test Plan

Unit: 55 tests across compile, assemble, projection, persistence and the model.
pre-commit run -a and python scripts/pyre_check.py both clean.

End to end, on a two-layer Qwen saved locally (no download), both an INLINE-only
prompt and one with an added PROJECTED slot:

step	result
tzrec.train_eval	3 steps, ce_loss 4.79 against ln(128) = 4.85 — near-random from a fresh init, as expected
eval	ce_loss computed and written to train_eval_result_v2.txt
checkpoint	model.ckpt-N/prompt/ carries sid_space.json, prompt_plan.json, prompt_hashes.json, tokenizer/
--continue_train	restores and continues
hash guard	changing the codebook to 8,8,8 is refused at restore
tzrec.predict	64 rows in, 64 out; generated_sids of shape (num_return, num_levels); every code inside [0, codebook), so the bands held and detokenize inverted both shifts
tzrec.export (HF)	loads with AutoModelForCausalLM.from_pretrained; export dir carries the prompt contract

Eight defects were found by running the pipeline rather than by unit tests —
four missed _create_model call sites, a (total, value_dim) shape the
assembler mis-read, the absent loss/metric hooks, an FX Proxy in decode, and
the un-derived feature groups a projected slot needs. All are fixed and covered.

Review round

codex-review raised 12 comments; all are triaged. Fixed: plan_hash now
covers projection bodies and routing (slot_to_module), a checkpoint with no
recorded prompt assets is fatal rather than a warning, and predict builds
embeddings outside both FX leaves so the inference tracer sees the sharded
lookup. Earlier in the round: sid_space and response are required at compile,
a token_format with no {i} is rejected, spliced projections take the LM
dtype, and two unreachable validation branches were removed.

Deferred with reasons, not dismissed: the incomplete HF export of a
projected-slot model, SID placement relative to the backbone's vocabulary (a
migration — it changes vocab_hash), the per-step host sync above, and two
beam-search memory items. Three were declined as guards for states the pipeline
cannot produce. The full reasoning is in the PR comment.

Not done

• Multi-rank is untested. Everything above is --nproc-per-node=1, so
  sharding, DMP and the planner have not been exercised. The projected config is
  the one with a sparse table to shard.
• The scripted serving front end and cache_ids (§11 of the design). Without
  cache_ids, a stock radix prefix cache keyed on token ids will reuse KV across
  users at a sentinel position — a correctness bug, so it must land before any
  projected slot is served.
• TorchScript export is refused with a message pointing at export_format: HF;
  the model's input is a stream the dataloader assembles, which an export-time
  dummy batch cannot supply.

Open design question

design_v2.md §8.1 has an "INLINE text" branch, but TokenizeFeature extends
IdFeature and is therefore always sparse, hence always PROJECTED. No existing
type can be an INLINE text slot. I implemented INLINE as unconditionally meaning
SID. Either the branch goes, or TokenizeFeature gains a no-embedding mode.

[github-actions]

[P1] Preserve or reject projected-slot state during HF export. This prefix filter intentionally keeps only backbone keys, so learned embedding_group and projections parameters are silently dropped while projected prompts are still exportable. The result loads as an HF LM but cannot reproduce checkpoint inference; either serialize/load the non-backbone state or fail export when projected slots exist.

[github-actions]

[P1] Make the accepted response plan match fixed banded decode. _generate() immediately forces one SID token per level, but this validator accepts static-only responses, prefixes such as Answer: {{answer}}, and multiple slots. Training then supervises a first token/sequence inference can never emit; restrict the response to the supported INLINE SID shape or make decode walk the compiled response plan.

[github-actions]

[P1] Enforce the compiled response width per row. Response slots are planned as exactly num_levels, but _inline_tokens() accepts zero or any whole multiple; an empty answer can yield an all-ignored/NaN loss, while two items exceed logits_suffix_len and silently train only the tail. Validate the actual response length against the plan here.

[github-actions]

[P1] Include projection routing in the persisted hash. slot_to_module is omitted, so swapping two slots' projection_name values when their bodies/shapes match leaves this hash unchanged; checkpoint keys still load, but each slot receives the other's learned weights. Hash the stable mapping and treat routing mismatches as incompatible.

[github-actions]

[P1] Preserve batch size for static-only prediction prompts. When the body has no slot and the answer exists only in the omitted response, there is no included feature to infer from, so a nonempty input batch is assembled as zero rows (cu_seqlens == [0]). Pass the raw record-batch row count explicitly instead of defaulting to zero.

[github-actions]

[P1] Validate the base tokenizer against the backbone before resizing. base_vocab_size defines every SID/sentinel id, but nothing checks it equals the pretrained LM vocabulary size; with a mismatched tokenizer those ids can overlap existing text rows and silently use the wrong embeddings/bands. Fail before resize (and verify tokenizer identity where available).

[github-actions]

[P1] Avoid materializing full-vocabulary FP32 logits for every live beam. At later levels there are batch * beam_width rows; for a 152K-token Qwen, batch 32/width 100 is about 486M logits, plus this FP32 copy, before slicing a small SID band. Preserve global-vocabulary scoring with chunked/fused lm_head + log-sum-exp, or enforce a memory-safe schedule.

[github-actions]

[P1] Do not physically duplicate the whole prompt KV cache for every first-level beam. With the documented 100+ widths, cache memory and copy bandwidth scale as batch * width * prompt_length * layers, which can OOM before subsequent levels run. Use prefix-sharing/paged cache, or a chunked and explicitly memory-bounded beam strategy.

[github-actions]

[P2] Keep max_seqlen on the host. Batch.to() moves every additional_infos tensor to CUDA, so converting this scalar to int performs a device-to-host read and synchronizes every forward, defeating sparse-pipeline copy/compute overlap. Carry it as host/Python metadata that Batch.to() does not migrate.

[github-actions]

[P2] Keep projected embedding lookup visible to the inference tracer. Wrapping _generate(self, batch) also hides its build_input() call, so TorchRec sees no sharded-module node and projected-slot distribution/lookup runs synchronously inside the leaf. Mirror the loss path: build embeddings before the wrapper and wrap only padding/decode.

[github-actions]

[P1] Do not fail open when a prompt checkpoint lacks hashes. CheckpointManager.save() swallows rank-0 asset-write failures, and this branch then restores with only a warning, disabling the vocabulary guard and allowing plausible but incorrect SID bands. Propagate asset-write status after required collectives and make missing prompt assets fatal, with any legacy bypass explicit.

[github-actions]

[P2] Validate finite integer-valued SID inputs before band checks. INLINE raw features are commonly floating-point; fractional values pass these comparisons and are truncated by astype(int64), while NaN makes both comparisons false and converts to an invalid id. Reject non-finite or non-integral values before conversion.

[github-actions]

Static review complete. The compiler/assembler/model separation is clear, and I left inline comments only for issues that affect correctness, export fidelity, or realistic decode performance.

One general gap: the PR body references docs/source/models/prompt_generative_qwen.md, but that guide is absent and the new public model/HF export mode is not indexed or documented. Before merge, please add the Chinese operator guide and HF export documentation covering the required response/sid_space, offset-SID input contract, tokenizer/backbone pairing, projected-slot export limitations, and beam-memory constraints.

Per review request, I reviewed statically and did not run tests or builds.

[WhiteSwan1]

Thanks — this was a useful pass. Triage of all 12 comments below: 3 fixed, 6 deferred, 3 not taken. Where the behaviour was checkable I verified it against the code rather than judging by reading; the evidence is in each item below.

Fixed

• plan_hash omitted projection routing (compile.py) — sorted(projection_plan.projections) iterates a dict, so the digest covered only the module ids. plan_hash now hashes both projections.items() and slot_to_module.items(). Without the latter, two slots could swap projection_name with matching bodies and still hash the same, so the keys load and each slot receives the other's learned weights. Two regression tests, both of which fail if either mapping is dropped.
• Missing prompt assets failed open (persist.py) — agreed, and it is worse in combination with CheckpointManager.save() swallowing rank-0 asset-write failures. An absent prompt_hashes.json is now fatal. No legacy bypass, since this stack has never been released.
• The inference tracer could not see the projected lookup (prompt_generative_qwen.py) — correct, and the asymmetry was unintentional: the training path already built embeddings outside the leaf. predict now does so on both paths and passes them in, leaving only the padding and the LM hidden.

Deferred, with reasons

• HF export drops projected-slot state — real, and tracked. It affects only a projected-slot model; the INLINE-only SID configuration this PR targets exports correctly, which is what the smoke runs and the checked-in mock config exercise. Taking the "fail export when projected slots exist" half separately is attractive, since dcp_to_hf currently writes a valid model.safetensors containing only the backbone, so the incompleteness is silent.
• Response plan vs. banded decode — accepted as a real train/serve mismatch, but nothing is broken for the shape in use: a response holding one INLINE SID slot trains and decodes consistently. Restricting the plan now would forbid shapes nobody has requested; better resolved alongside the extension that wants them.
• SID ids placed without knowing the backbone's table — the sharpest comment here, and it goes deeper than stated. Verified: SID ids land inside the backbone's existing table while the resize simultaneously adds rows above it that no token addresses, and a small sum(codebook) can make target_vocab_size round below config.vocab_size and shrink the table. A base_vocab_size == config.vocab_size check is not the fix — HF pads above the tokenizer (Qwen2.5 ships vocab_size 151936 against ~151.6k tokens), so that would reject every legitimate config. The correct invariant is max(tokenizer_vocab_size, lm.config.vocab_size), but compile_prompt has no access to the backbone config and the bands are already hashed into vocab_hash by the time the model could compare. That makes it a migration, not a patch, and it wants validating against a real backbone rather than a 64-vocab fixture.
• max_seqlen read back off the device — confirmed: Batch.to() migrates every additional_infos tensor, so this is a device-to-host sync per step, and it contradicts a design property the surrounding comments asserted. Performance only, though — the width is the same number either way. A fix was written and reverted: it carried the width as a plain int field on Batch alongside tile_size, which worked, but adds a prompt-specific field to a dataclass every model touches. Deferred pending a GPU measurement to size the win. The stale docstring claiming the model never reads it off the device has been corrected regardless.
• Full-vocabulary fp32 logits per beam and per-beam KV cache duplication — both real and both want a design rather than a patch (chunked or fused lm_head + log-sum-exp; prefix-sharing or a paged cache). Worth noting the current _band_logp is already a deliberate trade — normalize then slice, rather than a full-vocab log_softmax — so these ask for the next step rather than correct an oversight. The KV duplication is likely the largest single inference-memory item.

Not taken

These three each guard a state the pipeline does not produce:

• Response width unchecked per row — every route was checked. A codebook mismatch is already caught by the offset validation (_inline_tokens rejects values that do not carry their level offset — the self-validating property the offset encoding exists for). An answer holding two items requires an upstream labeling bug. An empty answer cannot come from the SID pipeline: resolve_sid_collisions never drops an item — one it cannot relocate keeps its original SID and overflows the bucket.
• Non-finite / fractional SID values — same guarantee. SID codes arrive from the collision tool's offset_codebook column as integers.
• Static-only prediction prompt assembles zero rows — reproduced, but a prompt whose body carries no slot has no per-request content at all; every request would produce the same static prefix. Reaching that state means the prompt needs rethinking, not that the assembler needs a row count.

The first two rest on one upstream guarantee — the offset_codebook column supplies exactly one complete SID per row, as integers — which is real but stated nowhere in the prompt stack. Noted as an assumption to make explicit.