[feat] export the front-end and serving contract for genrec

[tiankongdeguiji]

Summary

The HF export of genrec_causal_lm_model wrote only the backbone weights, the plain backbone config.json and a bare tokenizer.json, and refused any prompt with a PROJECTED slot. A serving engine therefore had to reimplement the prompt assembler or not serve the model.

A genrec model now exports like every other tzrec model. export() wraps the served half of the model, GenRecFrontEnd, and hands it to the same export_model every model goes through; the LM rides beside it as HuggingFace weights for the engine that decodes. One flat directory:

export_dir/
  scripted_model.pt           # GenRecFrontEnd: features -> input_ids / hole_positions / slot_embeds / hole_keys / hole_slot_counts
  fg.json  pipeline.config  model_acc.json
  dense_meta.json  sparse/    # USE_DISTRIBUTED_EMBEDDING=1, exactly as for any model
  config.json                 # composite: architectures=GenRecForCausalLM, model_type=genrec, backbone under text_config
  model.safetensors           # unchanged (dcp_to_hf)
  tokenizer.json  tokenizer_config.json   # extended tokenizer beside the config, so no --tokenizer-path

• tzrec/prompt/assembler.py: one PromptAssembler, a TorchScript-able tensor-op walk with two call sites (RFC 0001 §8.1): the collator and ScriptWrapper both run it on the parsed dict and prefix its streams into additional_infos. It is prompt structure only and trusts the parsed batch as every scripted module does: no schema, width, band, length or batch-size checks (each cost a device sync at serving; the data contract belongs to the parser and to the request), holes indexed by projected occurrence in emission order, the canonical parsed ABI (.values / .lengths / .key_lengths) and nothing else, one index_copy_ for every segment. Its streams sit in additional_infos under their own names; nothing else there shares them, so there is no prompt_ prefix.
• tzrec/prompt/hole_keys.py: the integer hole_keys fold of RFC 0002 §8 as its own module, HoleKeyBuilder, composed by ScriptWrapper beside the assembler at serving only; training never computes a key it discards. It iterates projected_slots, the same order project_slots and the assembler's holes use, so the three hole-indexed streams stay row-aligned; a dense member folds its float32 bit pattern per row, for DEEP and sequence slots alike, and each hole's accumulator starts at its slot salt so a hole no value reaches is still that slot's hole. FoldConstants is gone: the multipliers are class constants baked into the scripted front-end, nothing in sglang reads them, and no plan salt enters a key (a cache shared across model versions already matches static tokens alike, so only the engine can namespace it).
• FX cannot trace the walk or the fold (they branch on tensor values and size buffers from tensor sums), so both are FX leaves the way ComputeJTDictToKJT already is; TorchScript compiles them whole inside the traced graph.
• tzrec/models/genrec_model.py::GenRecFrontEnd: shares the model's embedding_group and projections (checkpoint names unchanged under ScriptWrapper, LM never loaded at export); predict returns the streams plus slot_embeds. It owns no export logic: export() runs export_model on it -- freeing the LM as soon as the front-end is built, since only the HuggingFace weights serve it -- and then, on rank 0, hf_export_util.export_hf_assets writes the HF weights, the extended tokenizer and one composite config.json, written once from the backbone config dcp_to_hf returns (written locally in a staging dir dropped however the export ends, and uploaded when the dir is remote). dcp_to_hf maps the checkpoint onto the architecture on the DCP metadata's names, so it reads the backbone's tensors and not the sparse tables beside them. The composite config is what sglang keys on: model_type: genrec selects GenRecConfig, GenRecForCausalLM resolves the backbone from text_config.architectures. Nothing tzrec-specific is in it: the server never reads the SID space (decode runs on the npz index, the processor needs no sentinel id, generation reads the standard eos_token_id / pad_token_id / vocab_size), and the offline index builder derives the token base from <|sid_0|>'s id in tokenizer.json and the bands and bundle identity from the bundle manifest.
• Because the front-end is an ordinary tzrec export, QUANT_EMB, INPUT_TILE and USE_DISTRIBUTED_EMBEDDING apply unchanged: the distributed-embedding export splits the slot tables into sparse/*.npz and makes the front-end the processor's dense stage with dense_meta.json, with no genrec-specific writer.
• The export records no bundle identity; compile_prompt checks the codebook against the manifest at export, and the server checks the npz's bundle_uuid against --sid-bundle-uuid.
• _get_sparse_embedding_tensor skips .weight entries that are not sparse tables.
• vocab_hash, plan_hash, the prompt_digests a checkpoint recorded and the restore guard (check_prompt_assets) are removed. Nothing in the export or in serving read them, and the guard refused routine warm starts: a SID bundle refresh with the same codebook, a renamed sentinel, any pre-existing checkpoint. Prompt-config compatibility with a checkpoint is the user's responsibility, as for every other tzrec config, recorded by the pipeline.config training saves. hf_export_meta.json keeps only the backbone prefix.

Paths considered

• Where the constraint-index builder lives. Kept out of tzrec: the npz format and its invariants are owned by sglang's ConstraintCsr, so the builder sits beside the loader there and reads the export's tokenizer.json, the bundle's manifest.json and its sid_to_items as data; tzrec's integration test asserts the token base the builder derives from.
• A bespoke front-end vs the standard export. The first two commits built the front-end by hand (SlotTables copied out of EmbeddingGroup, a hand-written dense_meta.json/sparse/ writer, a frontend/ subdirectory, a class-specific block in export()). That re-implemented what export_model already does; the third commit removes all of it in favour of a TowerWrapper-style wrapper and one isinstance case beside MatchModel/TDM.
• One walk vs two. A first cut added a second scripted assembler beside the collator's numpy one and pinned them with a parity test; that duplicated the one thing RFC 0001 A.13 exists to prevent, so the existing class was refactored into the scripted module instead.
• Hole keys inside the walk or beside it. RFC 0002 A.8 keeps the fold in the assembler so that training and serving do not run two behaviours of one module. A separate module over projected_slots keeps one walk, stops the dataloader hashing what training discards, and takes hashlib, the fold constants and a second input ABI out of the assembler; the hole order it shares with the assembler is the contract project_slots already relies on.
• Keeping the digests. A digest over repr(sid_space) and the tokenizer JSON catches a codebook re-split with the same total, which shapes would not, but it also fails every bundle refresh and cosmetic tokenizer change. Comparing the recorded SID space field by field would give an actionable error; nothing asked for it yet, so the guard goes rather than grows.

Notes

• Genrec is spelled GenRec throughout (classes and proto messages; config field and file names unchanged), matching DeepFM / DlrmHSTU and sglang's GenRecForCausalLM.
• Checkpoints written before this change carry digests that are now ignored.
• The front-end is TorchScript-only: AOT/TRT/RTP export is refused with a clear error (data-dependent shapes in the walk).
• Under USE_DISTRIBUTED_EMBEDDING=1 the processor runs input-tiled, so a request is one user and batch_size counts candidates (RFC 0002 §12); the front-end also still needs every slot member's raw .values / .lengths / .key_lengths beside the looked-up embeddings, because the walk and the fold read the canonical parsed dict only: the processor-side pass-through RFC 0002 §5.7 lists.

Out of scope, listed for follow-up: drop_if_empty, the raw-id pass-through in TorchEasyRec Processor (RFC 0002 §5.7), text-slot vocab_file injection into fg.json. User-facing docs wait until the export contract settles.

Test Plan

Env: torch 2.13.0 / torchrec 1.8.0 / transformers 5.16.1.

• tzrec/prompt/assembler_test.py: every case of the numpy walk on tensor inputs, the key_lengths layout walking like the flat one, sentinel/hole/hole_slot_counts invariants, max_seqlen, the batch sized from a dense or a jagged anchor, output-key constants, scripting round trip.
• tzrec/prompt/hole_keys_test.py: the mix64 host reference; the fold cases (equal content, changed content, same id in another slot, permuted item, swapped members, dense bit pattern, dense sequence per item, CPU/GPU bit-identity); keys follow the assembler's hole order across two slots; the builder scripts and is an FX leaf; two empty holes in different slots do not collide.
• Test fixtures live in tzrec/utils/test_util.py beside create_tiny_causal_lm: create_genrec_test_tokenizer and create_genrec_test_model (tiny LM, tokenizer, compiled prompt and model in one call); tzrec/tests/prompt_test_util.py and its base class are gone, each test is a plain TestCase.
• tzrec/models/genrec_model_test.py: ScriptWrapper(GenRecFrontEnd) returns the walk's streams, hole_keys from the builder and slot_embeds equal to the model's own projection on the same batch; it shares the checkpoint's parameter names and holds no LM; it FX-traces and scripts (integer streams compared exactly, floats with a tolerance); two projected slots concatenate occurrence by occurrence, split by hole_slot_counts; plus the model-level cases (vocab resize, finite loss into the backbone, FX trace of the training forward).
• tzrec/tests/genrec_integration_test.py, built like rank_integration_test.py: train_eval -> eval -> export in torchrun subprocesses on mock parquet with a tiny LM, INLINE + PROJECTED slots and a SID manifest; then the export is read back the way the sglang stack does it: flat layout, the scripted front-end (called with and without a device) equals PromptAssembler + HoleKeyBuilder on a parsed batch of the eval data, slot_embeds shape, composite config, tokenizer dir decoding a SID atom, config.json shape (no SID block) and <|sid_0|> at the token base. GPU-gated: the USE_DISTRIBUTED_EMBEDDING=1 export writes sparse/, dense_meta.json, model_acc.json, and a processor simulator (numpy lookup in the exported npz per sparse_features.json, fed per dense_meta.json, input-tiled) equals the default artifact. tzrec.predict is not run: the front-end emits token streams, not per-row columns.
• tzrec/utils/hf_export_util_test.py also pins that dcp_to_hf asks the checkpoint for the backbone keys alone.
• Existing suites still green: prompt.compile_test, models.genrec_causal_lm_model_test, datasets.dataset_test, utils.hf_export_util_test, utils.export_util_test, main_test. prompt/persist_test.py tested only the removed guard and is deleted.
• sglang (unpushed builder branch): the constraint-index builder takes --model-path, --manifest, --catalog; its unit test and an import smoke of the renamed GenRecForCausalLM / GenRecConfig / GenRecProcessor modules pass in the sglang test image.
• pre-commit run --files <changed> and pyrefly check clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01MB8aCo5mgF5ohAcBpWYDaN

[github-actions]

Code review — PR #662 (static review; no tests/builds run)

Five review passes ran over the change (code quality, performance, test coverage, documentation accuracy, security/multi-process). Overall this is a well-structured change: routing genrec through the standard export_model path, collapsing to one scripted walk with two call sites, and keeping the fold serving-only are all the right shapes. The Genrec→GenRec rename and the digest/guard removal are clean — no stale references to persist, vocab_hash, plan_hash, check_prompt_assets, prompt_test_util, or the old PROMPT_* keys remain anywhere in tzrec/ or docs/, config field names are unchanged (text-format compatible), and the mix64 fold checks out against SplitMix64 semantics (masked arithmetic shifts, wrapping int64, order-independent index_add_).

Inline comments cover the noteworthy findings; the ones I'd weight highest:

• hole_keys.py:113 — empty holes fold to key 0 in every slot, so two different projected slots that are both empty for a request collide silently in the engine's prefix cache; a one-line salt-seeded accumulator fixes it.
• main.py genrec export branch — the full LM is now materialized on every rank and discarded (the old path refused dense-EMA/missing-checkpoint before building any model); the discarded backbone also stays resident through export_model's DMP shard.
• hf_export_util.py — dcp_to_hf loads the entire checkpoint (all sparse tables) on rank 0 to keep only backbone keys, now on every genrec export; and the composite config merges the checkpoint-era vocab_size with freshly compiled SID ids with no comparison — exactly the silent "plausible output rather than an error" failure the deleted guard documented.
• assembler.py:248 — the band/width/max_length validations are gone for the training collator too, where the serving-sync rationale doesn't apply (CPU, eager); types.py:145 and prompt.proto still promise "an over-long row is an error".
• Test gaps: the front-end is only ever exercised with one projected slot, the scripted round trip compares int64 key streams through allclose on .float(), and projected=False (Pattern-I export) is dead in the integration test.

Smaller notes, not worth separate threads:

• export_hf_assets leaks the mkdtemp dir (holding full LM weights) if dcp_to_hf/compile_prompt/upload raises — a finally: shutil.rmtree would do; a partial remote upload leaves a half-populated export dir (matches the pre-existing export_model pattern, so noted only).
• batch_device returns whichever tensor is first in the dict; a mixed-device request (host batch_size scalar first) would build the streams on the wrong device. No test feeds a CUDA batch to the front-end — the GPU integration test loads with map_location="cpu" and feeds CPU tensors.
• Both integration tests pin TEST_NPROC_PER_NODE=1, so the rank-0-only export_hf_assets beside the all-rank export_model is never exercised multi-rank (traced statically and it looks race-free — uploads are disjoint and rank-0-sequential, dcp_to_hf loads non-distributed — just untested). _train_eval_export also only asserts eval_result.txt exists; sid_integration_test.py parses it and asserts finiteness.
• Exposing compiled_prompt on BaseGenRecModel (not just the front-end) makes every ScriptWrapper around a genrec model build the serving assembler with include_response=False; harmless today (main.py:1227's wrapper is discarded), but a footgun for tools like online_dense_export, whose build_dense_graph_module doesn't register the FX leaves.
• HoleKeyBuilder iterates projected_slots (body and response) while the serving assembler counts occurrences over body segments only; the two agree solely because compile forces response slots to INLINE — worth one docstring sentence somewhere, since a projected response slot would silently misalign hole_keys against hole_positions.
• _add_sparse_table's new skip comment says "such as a genrec backbone's", but the distributed-embedding export wraps the front-end, which holds no backbone — the dense .weight entries it filters are projection heads (and the sentence trails off). Similarly, _save_tokenizer_dir's docstring says it names "the two special ids"; it writes the token strings.

Nice touches worth calling out: the FX-leaf list shared correctly across both tracers' matching semantics (class-name vs path), the output_keys-vs-literals pinning test, and the CPU/GPU bit-identity test for the fold.