Conversation
_apply_admission picked the rows for the caller's table initializer with the two branches swapped. When the strategy reported it had written the non-admitted rows, the caller re-initialized all missing rows and overwrote them, so the strategy's initializer_args had no effect. When the strategy had no initializer (the default), the caller initialized only the admitted rows and the non-admitted ones reached the forward holding whatever storage.find's torch.empty value buffer contained. Swap the branches: the table initializer covers the admitted rows when the strategy took the rest, and every missing row otherwise. AdmissionStrategy.initialize_non_admitted_embeddings was annotated -> None while the caller branches on its bool; say bool and write down the fallback contract. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…efetch paths initialize_non_admitted_embeddings had exactly one caller, _apply_admission, which only runs on the generic forward path. _prefetch_cache_path and _prefetch_hbm_direct_path recorded non_admitted_positions and left the forward to fill those rows with the table's own initializer, so the same FrequencyAdmissionStrategy(initializer_args=...) silently stopped taking effect once a table used a cache or went HBM-direct. Give the strategy first refusal there too, falling back to the table initializer when it declines -- the contract _apply_admission already uses. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…d with The admission tests only ever asserted which keys reached the table, never what a rejected key's lookup returned, which is why both initialization bugs went unnoticed. They also always gave the strategy an initializer, leaving the fallback-to-table-initializer path uncovered -- and that initializer was DynamicEmbInitializerArgs(value=0.0), whose default mode is UNIFORM, so the value was ignored and the rows were random anyway. Add --expect-all-rejected: with a threshold no key can reach, every value the forward returns has to be the constant written by whichever initializer owns non-admitted rows. Run it over both owners (--non-admitted-init-value and --no-strategy-initializer) and all three storage configurations, since each reaches admission through a different code path. Name the tables' own initializer constant TABLE_INITIALIZER_VALUE so the fallback case has something to compare against, and make the strategy's initializer explicitly CONSTANT. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
DynamicEmbInitializerArgs.__eq__ branched on self.mode and then compared only that mode's fields, never checking that other was in the same mode. CONSTANT(value=0.0) == NORMAL(mean=0.0, std_dev=1.0) came out True, because NORMAL leaves value at its 0.0 default. Both methods also returned the NotImplementedError class rather than the NotImplemented singleton for a foreign operand. A class object is truthy, so args == 5 was True and args != 5 was False. Return NotImplemented and let Python fall back to identity. __ne__ only reimplemented what Python 3 already derives from __eq__, and reimplemented it with the same defect; drop it. Nothing compares these args yet -- this is groundwork for deciding whether the tables of a fused module share one initializer. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… initializer A fused module built one initializer per table but every consumer took initializers[0], so tables 1..N-1 were initialized with table 0's distribution, in training and in eval alike. Not a corner case: tables are grouped without regard to initializer_args, and the planner resolves an unbounded UNIFORM to +/-sqrt(1 / num_embeddings), so two tables of different row counts get different parameters without anyone asking. The tests miss it because all four of their tables are the same size and explicitly CONSTANT. Give each mode a second kernel that reads a [num_tables, num_params] tensor with the table id of the row it is writing, and a MultiTableInitializer that decides once which of the two a module needs. Tables that agree, and a module with one table, take the same scalar path as before: the two are separate entry points, so the common path carries neither pointer nor branch of the other's. That only works if tables may differ in initializer parameters and still fuse, so grouping now asks each object what has to match -- an initializer answers with its mode, a counter with its bucket layout and key type, a strategy with its threshold -- rather than comparing whole objects. One side effect: equal but separately constructed strategies now fuse where identity comparison kept them apart, which changes sharding plans. Checkpoints address tables by name, so they are unaffected. The same collapse ran through admission, so it gets the same treatment. An AdmissionStrategy now owns what deciding takes: its counter, which it adds to and erases from itself instead of having three call sites do it, and the initializer for the rows it rejects, which it declares and the module runs. materialize_for_tables turns a module's per-table configurations into the one strategy it runs, and is the only place any of it touches the device -- a strategy as written is inert and shareable. DynamicEmbTableOptions.admission_counter is deprecated accordingly: it warns, and __post_init__ folds it into a copy of the strategy so that capacities configured per table still are. initialize_non_admitted_embeddings is gone. It asked per batch a question settled at construction -- who writes a rejected row -- and getting that handshake backwards is the bug the first two commits on this branch fix. Nothing here has been compiled or run; there is no CUDA where it was written. docs/admission_strategy_design.md records the design, what is deliberate, and the tests still to write. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A strategy that tosses a coin rather than keeping a tally, and the first one to exercise the shape the previous commit put in place: state() stays None so no counter is ever built for it, materialize_for_tables has only the rejected-row initializer to open, and admit is a comparison. Nothing in the framework changed to make room for it. Admission is consulted only for a key that is missing, so a key gets a fresh toss every time it turns up and is not yet in the table: it takes 1 / p appearances on average to get in. That filters by frequency without counting anything. Deciding by hash(key) instead would be reproducible and rank-stable, but it would settle each key's fate forever -- at p = 0.1 nine keys in ten could never get in however hot they are -- and rank stability buys nothing when row-wise sharding already gives a key to one rank. A batch holding a key k times deserves k tosses. Tossing k times and taking any success is exactly one toss against 1 - (1 - p)^k, so that is what the draw is compared to, one draw per key. frequencies is where k comes from, and is 1 wherever the module is not counting occurrences. Three things are easy to get backwards here and are commented where they are decided. torch.rand draws from [0, 1), so the test is strict, the opposite of the initializers' curand_uniform on (0, 1]. log1p(-1) has no value, so both ends of the range short-circuit. And 1 - p loses a small p outright in float32 -- at p = 1e-8 it rounds to 1, leaving 1 - (1 - p)^k at zero and admitting nothing ever -- so the threshold goes through log1p/expm1, which never puts the two next to each other. The tests check what this code decides rather than what torch.rand or an inherited default does: the rate, both ends, repeats compounding, the grouping key, and a probability small enough that computing (1 - p)^k directly would admit nothing, which separates the two forms without needing a large sample. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…shell driver The driver ran the same twelve-line torchrun invocation from seven places, three of them the same triple loop copied over so that two storage flags could differ. Of its 214 lines about 150 were parameters being passed along, and adding a dimension meant copying the block again. The suite is now a table of named cases in the test itself, generated by the loops the shell used to spell out three times. --list-cases prints the names, optionally only those worth a given process count, and --case runs one by supplying its options as click defaults -- so anything also given on the command line still wins and every existing invocation still works. That leaves the shell with the one thing it has to decide, since torchrun decides it: how many processes. 18 lines, no parameter. The 70 runs are the same 70 runs: each case carries exactly the options its torchrun line used to, and the count and spot-checked parameters were compared against the old script. A failure now names the combination that failed -- cache-adam-lfu, reject-all-table-hybrid -- which is also what you pass to --case to run it again alone. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three files out of the twenty-odd in unit_tests are about admission, and unit_tests already keeps incremental_dump, retain_evicted_keys and table_operation this way. One wrinkle the other three do not have: test_embedding_dump_load is the fixture module several tests share -- test_lfu_scores and test_hybrid_storage_export import it too -- so it stays where it is and the moved test puts its directory on the path, with a comment saying why. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
test_probabilistic_admission was registered beside the rest of unit_test.sh while the suite it belongs to ran from test_embedding_admission.sh, so the two halves of admission were entered from different places. Run it from that driver instead, first, since it needs one GPU and a few seconds and will say what is wrong before the 70 runs behind it do. It moves group as a result: unit_test.sh fwd_bwd no longer reaches it, unit_test.sh load_dump does, along with everything else about admission. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
jiashuy
force-pushed
the
feat/dynamicemb-prob-admission
branch
from
September 18, 2026 06:53
a89c225 to
d9966b3
Compare
… with what Two things about the same interface. An initializer took (buffer, indices, keys, table_ids), with keys and table_ids addressed by buffer row while indices selected rows -- two conventions the signature did nothing to distinguish. A caller that lined something up with indices instead read past the end of it, taking some other table's id and then some other table's parameters, silently. Order them so the signature says it: everything before indices runs alongside buffer, one entry per row of it, and indices comes last because it is the mask over those rows, not another thing to line up with them. The launchers check that length against the buffer's row count, so getting it wrong is refused rather than read. And the second kernel of each mode was named _per_table, which reads as "one call per table"; _multi_table would have been worse, implying the first is not, when both write the same fused buffer and only the parameters differ. Name it for what it takes, <mode>_init_table_params, beside the table_params tensor it reads, which C++ and Python now spell the same way rather than table_args and table_parameters. The rows that tensor is built from keep a name of their own, table_param_rows, so the list a subclass writes down and the tensor the kernel indexes stop sharing one. InitParams<kPerTable, kNumParams> stays: there the flag qualifies the parameters, which is what it means. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…y meets The Counter and AdmissionStrategy sections described an interface nobody could implement against. Some of it never matched: get_initializer_args was documented as an abstract method and has never existed, Counter.add was given a frequencies/inplace pair instead of keys/table_ids/frequencies, Counter.create does not exist, and KVCounter was shown inheriting Counter when it is configuration and does not. The rest this branch made stale -- admit takes table_ids now, a strategy keeps its own counter rather than being handed accumulated frequencies, and materialize_for_tables, state and non_admitted_initializer went undocumented, as did ProbabilisticAdmissionStrategy. Write down what the classes are, with the flow the right way round: a strategy is configuration until a module materializes it, and the counter belongs to the strategy that counts, the framework only carrying it for accounting and checkpoints. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
isort and black on the files this branch touches, and clang-format 18 on the CUDA sources, which pre-commit does not cover but STYLE_GUIDE.md asks for -- the renames earlier on this branch left every continuation line in initializer.cu aligned to a name that is no longer there. One edit is not theirs: isort hoists a comment above the import group it sorts, away from the aliased import it explains, so the comment is worded to read correctly from there. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jiashuy
force-pushed
the
feat/dynamicemb-prob-admission
branch
from
September 18, 2026 08:09
d9966b3 to
c23e8e8
Compare
The shim that folds DynamicEmbTableOptions.admission_counter into the strategy had nothing holding it, so the two things easy to get wrong about it were only right by inspection: whose object gets written to, and which table's capacity survives when one strategy is handed to every table, as it usually is. Cover that it still reaches the strategy and still warns, that each table keeps the capacity it was given while the caller's own object stays as written, and that a strategy carrying its own counter is left as is. None of it needs a device. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…t runs AdmissionStrategy was both: the thing a caller writes on a table, and the thing a module runs. So a configuration carried an admit that only raised, a state that was always None and an initializer that was never there -- half an object, right only because nobody called that half. Two classes now. AdmissionStrategy stays the configuration: inert, shared across tables, and answering get_grouped_key. MultiTableAdmitter is what a fused module runs, and holds everything deciding takes -- admit, the counter behind state(), the initializer for the rows it rejects. create_admitter takes a module's per-table configurations and builds one, which lets each configuration pick the admitter it needs rather than have a base class enumerate them. That is four classes where one did less, and the reason to prefer them is that none of the eight members between them is dead. Nothing on either class is optional-in-practice any more, the module attribute is renamed for what it now holds, and both documents say so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68562836 -- failed
Result: 3/5 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68580485 -- failed
Result: 14/17 jobs passed |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Checklist