Repository navigation
[TRTLLM-16460][feat] modeling_v2: DeepSeek-V4.1-Flash-NVFP4 target (sm_103, unified parallelism) - #20192
Draft
Wanli-Jiang wants to merge 53 commits into
Draft
[TRTLLM-16460][feat] modeling_v2: DeepSeek-V4.1-Flash-NVFP4 target (sm_103, unified parallelism)#20192Wanli-Jiang wants to merge 53 commits into
Wanli-Jiang wants to merge 53 commits into
Conversation
…Wrapper Each catalog entry becomes an OpWrapper: raw_call is the certified surface (takes everything explicitly, reads nothing off self), beside the entry's reference, is_valid, compare and CELLS. The declared name lists go, the step contract is gated behind TRTLLM_MODELING_V2_VALIDATE, and torch is called directly rather than through mirror modules.
…rapper The nineteen entries, converted in families: the first seven, four more, the two collectives, the two paged-cache MLA entries (with a gate for cells), thop_attention, and the last three. The unused_ argument prefix goes; each op keeps its own names.
A validating() fixture that rejects a cell outside the entry's CELLS, and reducescatter's dtypes stated as an explicit allowlist.
A core is what a model statically is; a target is one way of running a step over it. The predicate is num_contexts == 0 and nothing else, so a mixed batch routes to prefill and decode may assume it has no context rows. The base class holds two rules the tree has already paid for: a target reads the core rather than snapshotting it, and the phase is read per routable unit rather than cached. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…de target The core's forward becomes a dispatcher: it runs the phase-independent contract check -- which is also what regrows the rope table -- and hands the step to one of two targets built after weights load. gpt_oss has no phase branch to divide, so the two targets share the step body and differ only where they genuinely can: decode states the two counts its routing guarantees instead of reading them back off the metadata. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Two checks, read off disk like the rest of this file. A decode target may not read back the count it was routed on -- inside a captured graph that value is frozen at whatever the capturing batch carried. Prefill stays exempt: it holds the mixed batch and needs the split. The second pins the class names the first depends on, with an explicit list of cores not yet migrated so that finishing the migration is a visible diff. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
… runs test_a_decode_target_never_reads_the_phase_back keyed on the DecodeTarget class name. Decode's body is inherited from _GptOssTarget (where decode actually executes) and its step_args delegates to the module-level _build_step_args -- neither is a DecodeTarget node, so the gate guarded exactly one line: the literal call in DecodeTarget.step_args. Invert the rule instead of enumerating more class names, since enumeration goes stale the same way: an attribute read of num_contexts/num_ctx_tokens anywhere in a core's module is now a violation unless it falls inside PrefillTarget's body (it holds the mixed batch and needs the split) or _check_step_contract's body (phase-independent metadata validation that legitimately mirrors the same projection). Verified the widened gate catches the newly-covered case: inserting a read into _GptOssTarget.forward turns it red naming that line, and reverting turns it green again. Also drop test_gpt_oss_core_declares_both_phases, which hardcoded the gpt_oss path; test_every_core_ships_both_targets derives the same paths from the routing tables and already covers it. Correct two docstrings that claimed this tree regrows `_rope` on an admitted longer max_seq_len -- gpt_oss has no `_rope` attribute at all; that claim describes deepseek and was copied in by mistake. State the reason that actually applies here: the contract check is phase-independent, so it belongs in the dispatcher rather than duplicated into both targets. Name the attribute each load-order assert actually guards (_targets vs _layers/_call_tensors) instead of a shared, ambiguous message, and drop the redundant attn_metadata type assert from the shared step body -- the dispatcher is the only way in and has already checked it. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…nd a binding layer An entry's arguments arrive at three different times -- engine constants and op-owned fixtures at construction, the batch state once a forward, the tensors at the call -- but the signature admitted only the last. This adds the two earlier stages and makes raw_call the surface underneath them. raw_call takes everything explicitly and reads no bound state, so CELLS keeps certifying the whole input space whatever a target binds, and reference and is_valid go on mirroring it. The validation harness patches raw_call for the same reason: is_valid still sees the merged argument set. No behaviour change. Every entry is constructed with nothing bound, so the default __call__ is a passthrough and all 56 call sites are untouched. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
d95f77e renamed each catalog entry's __call__ to raw_call and made __call__ a new base-class merge layer, but left the prose describing the old shape. Two module docstrings had come to state the opposite of the code directly beneath them: _op.py's contract table still listed __call__ as "the op call a target makes" and said the entry's test drives cells through __call__, and _validating.py's module docstring said the wrapper's __call__ goes straight to the kernel and is what gets patched -- both now true of raw_call instead. About a dozen subclass docstrings across catalog/*/*.py that said reference or is_valid "mirrors __call__" get the same correction, including the collectives' documented exception to mirroring, which stays an exception, now to raw_call. The base class's own new docstrings for __call__ as the three-stage merge layer are accurate and untouched. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The binding layer is only safe because raw_call takes every argument explicitly. One gate reads every entry's raw_call and refuses any read of self; the other pins is_valid to accepting everything raw_call accepts -- either by naming the same parameters or by absorbing the remainder with **kwargs -- rather than demanding a parameter-for-parameter mirror, which several entries never had and do not need (reference is cell-driven, not forward-driven, and is exempt). Both gates assert the first parameter is self, which a previous mirror check assumed rather than checked, and so missed an entry where both signatures had dropped it. Also adds a per-step staleness check, live only inside the validation harness, that catches a forward which binds step state without rebinding it on a later step. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
A target is not a second line of defence against routing: the router guarantees it picked the right target from the architecture and key parameters, so re-validating a checkpoint's full configuration inside every target does not scale. Remove all 35 assert statements from the gpt-oss-120b / sm_103 / tp1 target, along with the expressions and comments whose sole purpose was to feed them. Two of the removed asserts guarded runtime inputs this target does not implement (lora_params, spec_metadata) rather than routing decisions; those inputs are now silently ignored rather than rejected. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…class
An entry's __init__ now also accepts layered=, a {argument: per-layer
sequence} mapping distinct from the flat bound= kwargs. __call__ gains a
layer= keyword that, when given, projects one row of each layered table
onto the merged call after bound and step -- so a target can bind a
per-layer weight table once, at construction, and index it with layer=i
at each call instead of unpacking a parallel tuple itself.
An entry that never binds layered= has an empty table, so a call that
never passes layer= runs the exact same three-stage merge it always did;
this is what keeps every existing call site -- including deepseek's,
which binds nothing -- byte-for-byte unchanged.
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…rote thop_attention, fused_qk_norm_rope and flashinfer_fused_add_rmsnorm wrote their output by mutating an argument and returned None, which hid what a call site produced. Each raw_call now returns the tensor(s) it wrote -- output, qkv, and (x, residual) respectively -- so a caller can rebind to the result instead of reading it back off the argument it passed in. The mutation is unchanged: none of these allocate, so a caller that reuses a preallocated buffer across many calls (gpt_oss's attn_out, one per layer for 36 layers) still gets the same object back, just with a name attached to it at the call site. reference already mirrors this shape for all three, so no reference changes. deepseek's call sites (thop_attention x4, flashinfer_fused_add_rmsnorm x3) ignore the return value, as they did before; the two catalog tests that asserted "is None" now assert identity against the arguments instead. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…talog entries _GptOssTarget gains an __init__: it builds one bound instance per op this target calls, with per-layer weight tables (w_qkv/b_qkv, w_o/b_o, w_n2, w_rt/b_rt, the MoE expert operands, attention_sinks, the next-norm table) in layered=, and config scalars (eps, heads_q, heads_kv, head_dim, rotary_dim, theta, the four yarn_*, topk, num_experts, hidden, inter, inter_pad, _CALL_CONSTANTS) in bound=. fused_qk_norm_rope's inert q/k norm weight and the MoE runner's per-expert alpha/beta/limit vectors move from core._call_tensors (now gone) into this same construction, as the op-owned fixtures they are. forward's per-layer body collapses from eight positional op calls carrying ~90 lines of arguments to ten bound-instance calls indexed by layer=i. attention_window_size stays a per-forward, per-call computation rather than joining attention_sinks in the layered table: its full-attention value is attn_metadata.max_seq_len, which resolves through the KV cache manager's *actual* sized max_seq_len -- not known until after weight load, when _GptOssTarget.__init__ runs, and not guaranteed to equal the engine's requested max_seq_len (resource_manager.py's window validation can clamp it down under memory pressure). Baking a construction-time value in would silently diverge from the real one on a memory-constrained run. core._layers and core._call_tensors are deleted outright -- nothing reads per-layer state through the core anymore. The eleven cfg-derived scalars core copied for forward to read (eps, heads_q, heads_kv, head_dim, rotary_dim, theta, the four yarn_*, topk, num_experts, hidden, inter, inter_pad) stay on core: weights.py's load() reads several of them directly, and build_layer_views (via _GptOssTarget.__init__) now binds all of them into the op instances it constructs -- forward itself is simply no longer one of their readers. core.window and core.sliding are the one exception that still reads from forward, for the reason above. advance_step_generation() plus self._attn.bind_step(**step) now run once per forward, before the loop, replacing the **step unpacked at every thop_attention call site. Unit suite: 281 passed, 2 failed (known MPI-slot comm tests), 1 skipped -- no regressions. gpt_oss gsm8k accuracy: 89.992 (reference 90.300), matching the pre-refactor spread on this cluster. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…all sites Four follow-ups from review of the previous commit, none changing the computed values: - eps, topk and the GLU clamp limit were core attributes whose only remaining reader was _GptOssTarget.__init__ itself -- weights.py never touches them, unlike heads_q/heads_kv/head_dim/hidden/inter/num_experts/ the three padded widths, which stay attributes because weights.py reads them directly, at a separate time, through core.X. Deleted; the target now reads cfg.rms_norm_eps / cfg.num_experts_per_tok / cfg.swiglu_limit straight off core.model_config.pretrained_config. theta and the four yarn_* stay attributes even though the target is their only other reader too: recomputing them means re-deriving ~20 lines of YaRN math with real failure modes (KeyError, math.log domain errors) a second time, either duplicated or deferred past a multi-GB weight load: worse on both counts than one cached, early-validated copy. rotary_dim is gone outright -- it was already just self.head_dim under a second name. - self._moe's call site passed routing_bias as a bare positional None. It is now bound to None at construction (next to the comment explaining why: the router bias rides the router GEMM's own epilogue instead), and the call site names every argument. - self._attn's call site passed local_layer_idx=i beside layer=i, the same loop variable under two names. local_layer_idx is now a layered table (a trivial range(num_layers) one, since the op's own argument and the binding index happen to coincide for every call this target makes) so the call site states i once. - core._layers and core._call_tensors are confirmed fully dissolved: zero dotted attribute accesses remain, grep-verified separately from the `num_layers` substring collisions that made a cruder grep overcount them. The two mentions left are build_layer_views' and __init__'s own docstrings explaining that they used to exist and don't anymore; reworded away from the attribute spelling so they stop looking like hits. Re-verified unchanged: unit suite 281 passed / 2 known failures / 1 skipped (job 948502), gpt_oss gsm8k accuracy 89.992 (job 948586) -- both identical to the pre-fix run, as expected since no bound value changed. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…nd_layered A catalog entry's binding stages used to encode when a value became known (construction vs. per-step) rather than whether it varies by layer, which is the axis the entry actually cares about. OpWrapper now exposes two methods instead of three construction-time styles: bind_const(**values) for a value that is the same for every layer, and bind_layered(**tables) for one value per layer selected by layer= at the call. Both are callable from anywhere a caller finds a value knowable, and both update rather than replace prior bindings for other keys. __init__ no longer takes binding arguments; it only sets up empty state. Dropping the construction/step split means the base class can no longer tell a construction-time constant from a volatile per-forward value by name. The staleness check that used to compare self._step_generation now stamps each bound key with the step generation in effect when it was last bound, in a per-key self._generation dict, and flags under validating() only a key whose stamp is >= 1 (bound during some forward) and behind the engine's current generation -- a key that was once bound mid-forward and was not rebound this one. A key bound before the first forward keeps generation 0 forever and can never be flagged. gpt_oss's call sites are mechanically converted to the new methods with no behavioral change in this commit; binding attention_window_size through bind_layered follows separately. deepseek is unaffected: its module-level singletons never call bind_const or bind_layered, so their _const and _layered stay empty and __call__ remains a pure passthrough. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
… bind_layered attention_window_size used to be the one argument thop_attention's binding couldn't absorb: it is per-layer, but not knowable at construction (it derives from attn_metadata.max_seq_len, which the KV cache manager does not produce until after weights load), and not per-step either since it never changes again once known. With bind_const/bind_layered callable from anywhere, it now has a home: bound once with bind_layered on the target's first forward, guarded by a one-shot _window_pending flag (mirroring ModelingV2Core._contract_pending) so the check costs one branch per target's lifetime rather than one per layer per step. attn_metadata.max_seq_len is a plain field (interface.py), no attention backend reassigns it per step, and the only writes are in KVCacheManager.__init__ and engine construction -- both already done by the time a target's first forward runs, so the value read there is final. The per-layer attention call drops the attention_window_size kwarg entirely: self._attn(q=qkv, output=attn_out, layer=i). Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
A target is not a second line of defence against routing. The router picks a target from the architecture and key parameters, and re-validating a checkpoint's whole configuration inside every target does not scale -- this file carried 93 such assertions. Removes all 93, along with the expressions left dead behind them. Among the 93 were guards on runtime inputs this target does not implement rather than on configuration, so those inputs are now silently ignored rather than rejected. Verified on the rebuilt 26.08 container: three dep4 gsm8k cases pass (95.148 / 95.072 / 94.845 / 94.238) and the unit suite is unchanged at 283 passed. Note there is no verified pre-rebuild baseline to compare those numbers against; the argument that this is safe is that a passing assert cannot alter a computation. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…eek accuracy gate mv2_accuracy_deepseek_19345.sbatch is a copy of modeling_v2_accuracy_deepseek.sbatch repointed at the 26.08/19345 container (the tree this branch is built against), used to verify the assert-removal change against a real dep4 gsm8k run. Running it surfaced two pre-existing script defects, independent of that change: - Each leg trusted pytest's raw exit code. pytest can exit 4 (USAGE_ERROR) on an all-pass run because of a pre-existing nvidia_cutlass_dsl sys.path guard that fires during collection teardown; the script already greps the summary line for the expected pass count to detect a real failure, but then still reported the step as failed via the untouched raw code. Under `srun --mpi=pmix --ntasks=4`, a rank exiting nonzero also makes srun force-terminate the other three ranks, so a run observed one leg's reported exit flip to 143/SIGTERM after every test had already printed "passed" -- a symptom of this bug, not a wall-clock timeout (that run used 23m30s of its 4-hour budget). Fixed by pinning rc to exactly 0 or 1 off the grep result instead of the raw exit code. - /tmp/mv2_require.out and /tmp/mv2_off.out were written inside the container by the pytest step, but read by the summary section after the container had exited -- different filesystem namespaces, so both greps silently matched nothing. Moved both files under $REPO, which is bind-mounted the same way inside and outside the container. --time is left at the existing 04:00:00: the run that hit the 143 symptom used well under a tenth of that budget, so the time limit was never the actual constraint. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
… own its fixtures A target that needed a second instance of an entry had to reach its type through the singleton -- `type(flashinfer_rmsnorm)()` -- because each module exported only the instance and kept the class private. The class is now the exported name; the singleton stays, since deepseek and the entry tests still call it. bind_layered now takes the layer it is binding rather than a whole table. A model whose layers are not uniform can then bind only the layers that have a given operand, which deepseek will need: its first three layers are dense and the rest are MoE. The mx MoE runner builds its own alpha and beta from the expert count and a device, instead of receiving tensors the model manufactured from two constants that only this op ever read. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The core carried ten `self.X = cfg.Y` assignments that existed only so the forward could read them. The values are bound into the entries now, so the readers take them from `cfg` directly and the attributes are gone. What remains is derived rather than copied -- the padded widths, the YaRN quantities -- plus `window`, which the forward pairs with the derived `sliding` list per layer. attention_window_size is rebound every forward rather than once behind a pending flag. It cannot be known at construction, since the KV cache manager does not exist until after weights load, and it does not change afterwards -- so rebuilding a 36-entry table each step buys nothing except the removal of a state bit and a branch, which is the trade taken here. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The core carried 19 `self.X = cfg.Y` assignments -- num_layers, hidden, eps, vocab, heads, nope, rope, v_dim, kv_lora, q_lora, max_pos, dense_layers, num_experts, topk, moe_inter, dense_inter, n_group, topk_group, routed_scale -- that existed only so derive_after_load(), forward() and weights.py could read them back. Every reader now takes the value from `core.model_config.pretrained_config` directly, mirroring gpt_oss's 7c18232: a local alias at the top of __init__/derive_after_load/forward (and in MTPLayer's methods, off `core.model_config.pretrained_config`), not a second name stored on the core. weights.py's six functions that read the removed `core.X` attributes do the same. qk_dim and lat_dim (`self.nope + self.rope`, `self.kv_lora + self.rope`) and shared_inter (`moe_intermediate_size * n_shared_experts`) are not `cfg.Y` copies, but the same reasoning applies -- they are recomputed locally at each of their few call sites rather than kept as a third attribute. What stays: theta, rope_factor, rope_orig_max, beta_fast, beta_slow, rope_amplitude and q_scaling (the YaRN derivation has real failure modes -- KeyError, math.log domain errors -- worth catching once, before the weight load, rather than duplicating); local_experts and expert_offset (combine cfg with the parallel mapping, and weights.py reads them at a separate time); mtp_layers and mtp_enabled (derived with real fallback/resolution logic, also read by weights.py); quant_mode and the parallelism fields (rank, ep_rank, ep_size, dp_group, dp_size) are not cfg copies at all and are untouched. No call-site shapes change in this commit -- every op still takes the same arguments it did before, just read from a local instead of an attribute. Binding those call sites through the catalog entries' bind_const/bind_layered is the next two commits (trunk, then MTPLayer). Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Targets in this tree share no code with each other. _GptOssTarget existed only because prefill's body and decode's happened to coincide -- gpt_oss is not MLA, so there is no phase branch to divide -- and a base class justified by a coincidence is one that breaks the moment the coincidence ends. PrefillTarget and DecodeTarget now derive from Target directly, each with its own __init__, step_args and forward. The one real difference stays where it was: prefill reads both counts off the metadata because it holds the mixed batch, decode states them as zero because its routing guarantees that. This costs about 230 duplicated lines. It buys each target being readable and changeable without reference to the other. The phase gate gets simpler as a result. It had been widened to follow into the shared base and the module-level helpers, because scanning DecodeTarget alone reached a single line; with the body in its own class the direct form suffices. Verified by planting a read inside DecodeTarget.forward and watching the gate name it. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…he manifest from it A weight's shape, its checkpoint source and its load-time transform lived in three places: the model's declaration block, the manifest, and a function far from either. The contents were different facts, but the seventeen key names were written twice, and an assert in load() existed only to police that. WEIGHTS states each weight once -- shape, dtype, source, transform -- and the manifest unrolls it rather than repeating the keys. shape is a callable because every shape here is derived: fc1_k_pad is a padded width, q_width a product. Sources that assemble one parameter from several checkpoint tensors, as qkv does from q/k/v, carry a slice per source. The declaration block in modeling.py still stands; this does not switch over. A test drives the table and a transcription of that block and asserts they produce identical keys, shapes and dtypes, which is what the switchover will stand on. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…alog entries Fifty-six call sites across three forwards carried shape scalars, call constants and per-layer weights that never change once the engine exists. Those move into the entries -- bind_const for what every layer shares, bind_layered for what each layer owns -- and the calls keep the tensors and the layer index. Instances live on the core rather than on a target, because this model has not been split into prefill and decode targets yet. The phase branch inside the layer body is untouched for the same reason: dividing it is the Target layer's job, and the two paths run genuinely different algorithms. bind_layered taking a layer id is what makes this model work at all -- its first three layers are dense and the rest are MoE, so only the layers that have a given operand are bound. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Three things, all at the same entry point.
The contract check's position_ids parameter went unused once the dtype assert
it guarded was removed; the gpt_oss call site reached it through
kwargs.get("position_ids"), which would have silently passed None had a
caller ever switched to positional arguments. Both are gone.
The pending flag moves inside the check, so the caller writes one line rather
than a branch around a call. That costs a Python call per forward even when
validation is off -- per forward, not per layer, against a step measured in
milliseconds.
advance_step_generation moves to the dispatcher. "A new engine step has begun"
is the dispatcher's knowledge, not a target's, and keeping it there had left
two copies in gpt_oss after the targets stopped sharing a body. It is
deliberately not folded into bind_const: it must fire once per forward while
bind_const fires once per op, and deepseek binds step state to two entries in
one forward. The same placement is what keeps deepseek's MTPLayer correct
without anyone remembering to exclude it -- the layer runs inside one engine
forward and must not start a generation of its own.
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…ss model The model allocated the parameters and the loader said where their bytes came from, naming all seventeen keys twice. The table names them once and the model asks it for the dict; the declaration block and its local P() helper are gone. The drift assert goes with the duplication it policed -- the manifest and the parameters now unroll the same tuple, so a key in one and not the other is unreachable. The coverage assert stays: that every checkpoint key is consumed is a separate claim and still worth making. The equivalence test keeps its transcription of the deleted block rather than importing the model. That was the reason for transcribing it: the table stays held to what this target used to declare, where an import would now have the test agree with itself. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
gpt_oss and deepseek each defined their own class literally named ModelingV2Core, deriving straight from DecoderModel. The only thing genuinely common between the two -- the opt-in step-contract check (_check_step_contract plus the _probe_step_surface/_after_contract_check hooks it calls) and the _contract_pending flag's initialisation -- now lives in a real shared base, tensorrt_llm/_torch/_experimental/modeling_v2/_core.py. Everything model-specific (__init__, forward, build_layer_views, derive_after_load) stays on each model's own class, renamed to GptOssModelingV2 and DeepseekV3ModelingV2 respectively, so the two still share no code with each other. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
… phase-keyed dict self._targets: dict[Phase, Target] baked the reason a target is chosen into the data structure. Replace it with named self._prefill/self._decode attributes and a _select_target(attn_metadata) method, so a future selection criterion beyond phase has one place to grow into instead of the shape of a lookup table. _select_target delegates to _target.py's phase_of rather than restating num_contexts == 0 inline, keeping that predicate in the one place that already states it in full. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…check into its targets _build_step_args was a module-level function both targets called -- the same shared logic the deleted _GptOssTarget base had held, in another form. Mixins, helpers and partials were ruled out when that base went; a module-level function was not, so it survived. Each target now writes out its own projection. Prefill reads both counts off the metadata because it holds the mixed batch; decode states them as zero because its routing guarantees that. The contract check moves with it, into Target: it probes the metadata surface step_args reads, so it belongs beside it. Both targets running it is correct rather than wasteful -- their projections differ, so probing one proves nothing about the other. deepseek has no targets yet, so its check stays on its core and _core.py keeps the skeleton for it. The two models are in different states until deepseek is split; there is a comment saying so, because the alternative is someone tidying away code deepseek still needs. The four yarn_* values become locals passed to both target constructors. They were computed in __init__ and read only when the targets bound fused_qk_norm_rope -- nothing in between touched them. Verified on 4xGB300: modeling_v2 unit suite 287 passed / 1 skipped with the two pre-existing comm op-matrix failures unchanged, and gpt_oss accuracy 89.765. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Comments and docstrings were 38% of gpt_oss's modeling.py and 33% of deepseek's. Most of that was the refactor narrating itself -- why a value is a local and not an attribute, why a name is not duplicated, why a helper was not extracted. AGENTS.md puts that in the PR description, not in the code, and a reader a year from now wants the code explained as it stands rather than a record of how it got here. The rest restated the line below it. What survives is one category: a trap where getting the line wrong produces a wrong answer or a crash the reader could not have inferred. Most of them are silent -- the MoE op ignoring routing_bias on this routing method, a swizzled scale layout whose byte count happens to match whenever num_tokens is a multiple of 128 (every decode CUDA-graph batch), a kv_norm weight that folds a normalization this target already applied. Each is now one or two lines stating the trap and nothing else. gpt_oss keeps 5, deepseek 14. Docstrings keep their first sentence; the later sentences stay only where they state a contract a caller must obey that the signature does not show. gpt_oss 1006 -> 658 lines, deepseek 1917 -> 1381. Proved comment-only by mv2_comments_gate.py, which parses each file before and after, strips docstrings from both, and compares the executable AST. It is red-green tested: unchanged file OK, deleted comment block OK, a one-token edit to routing_bias FAIL. Both files report OK, so no test run is needed to know the behaviour is untouched. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…it build the manifest ModelWeights is to weights.py what ModelingV2Core is to modeling.py: the table's three passes -- declare, manifest, load -- are the same loop over the same table for any model, so they live on a base and each model supplies only its WEIGHTS and its dims(). The manifest was the point. It did not unroll WEIGHTS; it wrote the checkpoint key, destination index and transform of all seventeen roles out a second time by hand, which is exactly what the table exists to stop. A comment above the loader claimed both sides unrolled the table and used that to justify dropping the drift check between them -- the claim was not true of the code it sat in. manifest() now generates those rows, so a role has one source of truth again. Two more copies go with it. dims() is now the only place a derived width or padded size is worked out: q_width and kv_width used to be computed in the core's __init__ and again in _manifest, and the three padded sizes in __init__ while the transforms read them back off the core. And _FC1_K_ALIGN was a literal in modeling.py, for the alignment the quantizer pads to, while the table independently decided the declared width it has to match; it is one constant in weights.py now, imported by modeling.py. Dropped on the way: _pad_up and the dtype local in modeling.py, both left without readers once the padding moved. _manifest stays, dead, for one commit. The equivalence test compares the generated manifest against the real function rather than a transcription of it, and that is only possible while both exist. It is already green. That test also found a fixture bug it was not looking for: _DIMS said fc1_k_pad was 2944, padded up to 128 like its neighbours, where the model pads to FC1_K_ALIGN and gets 3072. The declaration tests feed both sides the same bundle, so the wrong number agreed with itself and had been proving the table at a shape -- fc1_w 1472 columns wide instead of 1536 -- that the model never allocates. Corrected here. Not yet verified: the unit suite and the gpt_oss accuracy run are in flight as job 1009393. The weight-table tests passed; declare, dims and manifest equivalence are green. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…, and drop the probe indirection build_layer_views said what it did on its first day and has not since: it builds the targets, derives the GEMM-side views, and computes the YaRN ramp. What it actually is is the model's post-load hook, and DecoderModel already has one by that name -- so this was a second name for an existing concept, and the shell had to know the private one to call it. The core now overrides post_load_weights, calls super(), and the shell's line reads self.model.post_load_weights(). Checked before renaming: the base's post_load_weights runs setup_aliases, transform_weights and cache_derived_state, which are empty on this path and do not descend into submodules, so calling super() here adds three no-ops rather than re-running a transform over the core's subtree. _probe_step_surface was a one-line method that existed only so a generic _check_step_contract could reach a model-specific projection. The check now calls the projection directly. For gpt_oss that was already true -- Target calls its own step_args -- so this only concerned deepseek, whose check moves onto its core beside the _build_step_args it probes. That empties ModelingV2Core: gpt_oss set _contract_pending there and never read it (its targets own one each), and deepseek now owns both the flag and the check. A base class holding nothing is worse than no base class, so _core.py is deleted and both cores derive from DecoderModel. deepseek keeps _after_contract_check as a named method -- it has a real body, growing the rope table to the engine's resolved max_seq_len. _build_step_args stays a module-level function: MTPLayer.forward calls it too, so it is not the core's method to own. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Folds in the follow-up fix: the core's post_load_weights is not an override (DecoderModel has no such method, so there is no super() to call), and the weight table is a class attribute of its ModelWeights subclass.
…decode target deepseek is the model the target layer was built for. MLA's TrtllmAttention serves context_only or generation_only and refuses a mixed batch, so the phase split already existed -- as two attention instances and an `if tc:` / `if gen:` branch inside one 233-line forward. This moves it up to where gpt_oss keeps it. DecodeTarget binds 23 ops against the prefill target's 27: the context attention instance, the kv_b projection and the two paged-KV helpers all serve a path its routing guarantees never runs. Its forward loses the context block and, per layer, the two torch.split calls that separated rows it cannot have -- 61 layers' worth. forward 233 -> 177 lines, __init__ 289 -> 236. The two targets share no code, including their copy of _dense_mlp, which binds target-owned ops and so cannot stay on the core. Both bodies were cut from the original source rather than retyped, and the result is checked structurally: PrefillTarget constructs the same 27 ops and the same 24 layered bindings the core did, and DecodeTarget's set differs by exactly those four names. Separately, a bug this exposed. _after_contract_check hung off the contract probe, which is gated on TRTLLM_MODELING_V2_VALIDATE -- set by both accuracy tests and by nothing in a served engine. But it does not check anything: it grows the rope table to the engine's resolved max_seq_len and decides _cached_ctx. With the flag off, production kept a table sized to max_position_embeddings -- an unchecked out-of-bounds read once speculative_config inflates max_seq_len past it -- and _cached_ctx stuck at False, making the cached-context path dead code. It is now _derive_first_step on the core, unconditional. The probe itself stays opt-in, on each target, where step_args is what it calls. The rope tables, MTP operands and side stream stay on the core: MTPLayer reads the first two, and a CUDA stream made inside a graph capture is not capturable. derive_after_load becomes post_load_weights, matching gpt_oss. Not yet verified: deepseek accuracy is in flight. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…ore class name DeepseekV3ModelingV2 said which checkpoint family it served and nothing about how it is parallelised, but the parallelism is not incidental here: the forward allgathers before the router and reduce-scatters after it, the expert window is this rank's slice of 256, and the MoE is the only thing that needs a collective. A second strategy over the same checkpoint is a different core, not a flag on this one, so the name carries the tag: DeepseekV3Dep4ModelingV2. dep4, not tep4: this repo treats them as different strategies -- deepseek_v3's routing.py even has a "tep4_to_dep4" case -- and this target's config is tensor_parallel_size=4, moe_expert_parallel_size=4, enable_attention_dp=True, which is attention DP, not attention TP. The shell already carried the full tag (ModelingV2DeepseekR10528Nvfp4Sm103Dep4); this brings the core into line with it. gpt_oss has the same gap -- its core is GptOssModelingV2 against a ModelingV2GptOss120bSm103Tp1 shell -- and is left alone here rather than folded into a deepseek change. Name-only: no behaviour, and the accuracy run that cleared the target split (gsm8k 94.69 require / 93.93 off, acceptance 2.920 / 3.037) stands. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
MTPLayer never got the op-wrapper conversion the trunk had. It made 28 module-level calls -- thop_attention(...), mla_rope_generation(...), cublas_mm(...) -- each passing every argument explicitly, including forty on one of them, while the trunk next to it called 30 bound instances. That is why "move MTP into the targets" is not a relocation: a target holds bound ops, and these were not ops yet. Now 30 bound instances and zero raw calls. The module is a single layer, so every binding is bind_const and there is no layered table. The positional calls were converted against each op's raw_call signature rather than by eye: the three MLA entries take 21, 16 and 47 parameters and the call sites passed 21, 16 and 40 in order, the last seven taking their defaults. Each argument kept its value and gained its name. Binding happens in a new MTPLayer.post_load_weights, driven by the shell after the core's own hook: this module is built in the shell's __init__, while every parameter is still meta, and its operands come from core._mtp, which that hook derives. The step projection and the rope tables are rebound on every call rather than once. The worker rewrites attn_metadata in place between draft step 0 and step 1+, so a value read on the first of the N calls in one engine forward is wrong on the rest -- and frozen wrong under capture. core._rope can also have grown on the engine's first step. This is D2a. D2b splits the forward by phase into the two targets; it needs the ops to exist first, which is what this does. Not yet verified: deepseek mtp3 accuracy is in flight. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…to carry it The deepseek target now states each weight once, as a W, like gpt_oss: a 75-line hand declaration block in the core's __init__ becomes one declare() call, the hand-written _manifest goes, and __init__ loses the sixteen local dims whose only reader was that block -- dims() derives them now, once. W keeps its field count; two existing fields widen instead. `src` gains a callable form `(d, i) -> [(key, index), ...]`, whose two old forms are special cases: it exists for sources that are not a fixed template -- this rank's expert window (the checkpoint key contains expert_offset), a prefix that differs between layer families, and multi-key rows whose tensors feed one transform together. `per_layer` becomes `layers`: ALL_LAYERS, None, or a callable saying which layers carry the weight, because deepseek's dense and MoE layers do not carry the same set. declare and manifest now share _targets(), so they cannot disagree about whether a weight exists. Two findings the table forced into the open: The MTP module is layer 61. Eleven of its weights match a trunk layer's exactly -- shape, key template, transform -- and the checkpoint names them model.layers.61.*. The trunk entries now extend one layer when MTP is live, and the mtp_* parameter keys become l61_*; only enorm/hnorm/eh/head_norm and the bf16 expert family are the module's own. The dense MLP and the shared expert share a parameter key and nothing else: different intermediate width, different checkpoint prefix. Two table entries over disjoint layer ranges, same name, rather than one entry pretending they are one weight. deepseek overrides load: multi-key rows hand their transform several tensors, the 0-dim NVFP4 scalars reject `[:]`, lm_head is asserted whole (attention DP replicates it), and coverage is asserted in both directions, with the expected leftover -- this rank's off-window experts plus the MTP layers -- supplied through the expected_unconsumed hook. Proven before the hand-written code went: with both still present, the generated manifest was compared against _manifest per key and per source, MTP on and off, and declare against a transcription of the init block -- 12 passed (SLURM 1081689), after which _manifest, its helpers and the comparison tests were deleted together and the suite rerun green (1081889, 293 passed with the two pre-existing comm failures). The declare transcriptions stay in the tests. gpt_oss: per_layer=False becomes layers=None on two entries; its own dead _manifest and the test that had proven it go the same way. Not yet verified: both models' accuracy runs are in flight. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The MTP module's forward carried the same `if tc:` / `if gen:` branch the trunk had before its target split, and the same reason to lose it: MLA serves context_only or generation_only, never mixed. Each target now carries its own mtp_forward and binds the module's ops in its own __init__, from core._mtp, under the same _mtp_-prefixed names on both. DecodeTarget's copy binds 23 of the prefill side's 27 -- the context attention instance, the kv_b projection and the two paged-KV helpers serve a path its routing guarantees never runs -- and its body loses the context block, the three per-call splits and the attn_ctx rebind: 118 lines against prefill's 174. Every draft step past the first in an engine forward takes this shorter path. MTPLayer stays, as the object the spec worker holds across calls, and shrinks to the two things that are genuinely its own: shared_head with its norm -- that entry point is called directly by the worker and reads no phase -- and a forward that is now one routing line. The routing runs per call, not per engine forward, because the worker rewrites attn_metadata in place between draft step 0 and step 1+: one engine step's N draft calls can cross from prefill to decode, and each one asks _select_target again. Both bodies were cut from the D2a source rather than retyped, with the same statement-level tooling as the trunk split, and checked the same way: the prefill op set equals D2a's bindings minus the head norm, decode's differs by exactly the four context names, and each tc-elimination rewrite asserted a hit. Verified on 4xGB300: modeling_v2 suite 293 passed with the two pre-existing comm op-matrix failures unchanged; deepseek accuracy require leg 2 passed (gsm8k 95.07, acceptance 2.912 -- the mtp3 case drives both targets' mtp_forward, draft step 0 through prefill and steps 1+ through decode), stock leg 1 passed (95.07, acceptance 3.018). Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…re class name too The same change deepseek's core already had: tp1 is not incidental to this target -- every expert is local, the quantizer pads for a single rank's operands, and there is no collective anywhere in its forward. A second strategy over the same checkpoint would be a different core, so the name carries the tag, matching the shell (ModelingV2GptOss120bSm103Tp1) and the directory (gpt_oss_120b__sm_103__tp1). Name-only; five references, one file. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
3fafe99 committed the deepseek accuracy gate's SLURM script to the repo root while fixing its exit-code handling. It is developer scaffolding -- the in-tree gates are the tests it drives -- and does not belong in the package. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The target split removed deepseek's phase branch by routing: DecodeTarget's body lost `if tc:` because its routing guarantees no context rows. The prefill side kept its body verbatim -- including the now always-true guard. A target is not a second line of defence against its own routing: `phase_of` sends a batch here exactly when num_contexts >= 1, so tc >= 1, and the guard goes, along with the two `if tc else 0` expressions it justified. Both forward and mtp_forward, since the draft step routes the same way per call. `if gen:` stays in both, and is not a phase branch: prefill is the general case -- context rows and possibly generation rows beside them -- so gen may be zero, and that is batch composition, not routing. Verified on 4xGB300: modeling_v2 suite 295 passed / 0 failed, deepseek gsm8k identity 94.69 / mtp3 95.00 (reference 94.24), acceptance 2.954 against stock 2.998. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…he layer loop The trunk forwards passed seven metadata projections to _mla_gen at every layer -- sequence lengths, host lengths, the three KV-pool handles, num_contexts -- values that cannot change within a step. They join the per-step bind_const at the top of forward, as an explicit subset rather than **step because the MLA helpers' raw_calls do not take the projection's engine-bookkeeping keys. Prefill's _mla_append and _mla_load get the same treatment, guarded by the cached-ctx flavor whose ctx_* fields are the only place those metadata attributes exist. The attention scratch moves out of the loop: eight tensors (nine on the prefill side, which also owns the context K buffer) were allocated per layer, 488 empties per decode step for buffers whose sizes depend only on step constants. One set now serves all layers -- safe because each layer's producer fully rewrites a buffer before that layer's consumer reads it and the whole attention path runs on the main stream, so layer i+1 cannot overwrite what layer i still reads. Six of them ride the per-step bind of _mla_gen and _attn_gen together with predicted_tokens_per_seq. The call sites keep only the data flow: _mla_gen(fused_q=..., q_pe=..., latent_cache=..., layer=i). The two mtp_forwards are deliberately untouched: the draft step runs a single layer per call, so there is no repetition to hoist, and the same change there would scatter one call across three binding sites for no reuse. Verified on 4xGB300: modeling_v2 suite 295 passed / 0 failed; deepseek gsm8k identity 94.69 / mtp3 94.69 (reference 94.24), acceptance 2.942 against stock 3.018. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…ut = op(in)
gpt_oss's targets follow one form -- a call site always binds what the op
wrote, even when it also passes the buffer (attn_out = self._attn(q=...,
output=attn_out)). deepseek still had op(in, out) statements, which read as if
their tensors were only consumed when some were being overwritten.
The root cause sat in the catalog, not the call sites: three deepseek-only
in-place entries -- bmm_out, mla_rope_generation, mla_rope_append_paged_kv_
assign_q -- had raw_calls that returned nothing, so there was nothing to bind.
They now return what they write (out, fused_q, q), the same move the three
norm entries got when gpt_oss adopted the form.
All four forwards convert, both trunk and draft:
fq_nope_t = self._absorb(a=torch.transpose(q_nope, 0, 1), out=fq_nope_t, layer=i)
fused_q = self._mla_gen(fused_q=fused_q, q_pe=q_pe, latent_cache=latent_gen, layer=i)
lat_out = self._attn_gen(..., output=lat_out, layer=i)
attn_gen_t = self._expand(a=lat_out_t, out=attn_gen_t, layer=i)
q_ctx = self._mla_append(q=q_ctx, latent_cache=latent_ctx, layer=i)
attn_ctx = self._attn_ctx(..., output=attn_ctx, layer=i)
o, residual = self._norm2(o, residual, layer=i)
mlp_out, residual = self._next_norm(mlp_out, residual, layer=i)
The transposed scratch views (fq_nope_t, lat_out_t, attn_gen_t) become named
step constants beside the buffers they view, which also collapses the absorb
and expand call sites from nested-transpose blocks to one line each. Checked
structurally: zero bare-expression op calls remain in any target forward of
either model.
bmm_out's unit test pinned the old contract -- `assert op(a, w, out) is None,
"the op writes in place"`. It now asserts the stronger version of the same
fact: `is out` proves the op wrote this very buffer and returned it, not a
copy that happens to match.
Verified on 4xGB300: modeling_v2 suite 295 passed / 0 failed; deepseek gsm8k
identity 94.77 / mtp3 94.77 (reference 94.24), acceptance 2.930 against stock
3.019.
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
…log into the model
The catalog is meant to be the shared vocabulary; an op only one target ever
calls does not belong there. Eleven move into an `ops/` subtree under their
owning model, keeping the category layout (attention/moe/gemm/quantization):
deepseek (8): load_paged_kv_cache_for_mla, mla_rope_append_paged_kv_assign_q,
mla_rope_generation, bmm_out, fp4_block_scale_moe_runner, fused_moe,
noaux_tc_op, fp4_quantize
gpt_oss (3): fused_qk_norm_rope, mxe4m3_mxe2m1_block_scale_moe_runner,
mxfp8_quantize
Eight stay shared: thop_attention, cublas_mm, flashinfer_rmsnorm,
flashinfer_fused_add_rmsnorm (multi-model), plus flashinfer_silu_and_mul,
allgather, reducescatter, nvfp4_gemm (deepseek-only today but kept in catalog
for generality -- SiLU is a ubiquitous activation, the two collectives are
distributed infrastructure any TP/EP model needs, and nvfp4_gemm is the generic
NVFP4 GEMM beside cublas_mm).
Each moved op is wrapper + GPU test (no .md contracts exist despite the index's
aspiration). Both move via git mv, preserving history. The wrapper's one import
-- `from .._op import ...` -- becomes the absolute catalog path, since _op
stays in the catalog; modeling.py imports the op through a relative `.ops.<cat>`
path. The test keeps its `__extra_import_path__` mechanism, widened to reach
_validating from two directories deeper.
index.yaml drops the eleven moved entries (now 8). The CI lists collect
`unittest/_torch/modeling_v2` by directory, and the moved tests stay under that
tree (models/<model>/<cat>/), so they are still collected without a list change.
Import closure checked statically across all model files; ruff clean. Accuracy
verification in flight.
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
The allgather and reducescatter op matrices spawned their own 4-rank mpirun, which inside a SLURM-launched container picks the slurm launcher (shells out to srun, absent in the container) and sizes its slot budget from the outer allocation (refuses -n 4). The tests were correct -- CI runs them -- but they could not run anywhere else without three OMPI env vars, and they looked like known failures for weeks. They now spawn one torch.multiprocessing worker per device and rendezvous through torch.distributed's env://, so they run on any 4-GPU box with a plain pytest and no environment at all. Each worker scrubs the inherited PMIx/PMI/ OMPI namespace from its children, or the first hard exit trips srun's pmix errhandler and SIGKILLs the whole step. Running under torch.distributed means the ops take the ProcessGroup path the product uses when mpi_disabled(): the entries gain a `process_group` argument that selects the `_pg` op symbol (None keeps the MPI-session symbol), and the matrices bind a boxed c10d group once, so every call site keeps its certified shape. A rank subset is now a real sub-ProcessGroup, not just a shorter group list. The host-side engine-sync check uses a gloo group rather than constructing the engine's TorchDist, which needs ray importable. Three kinds of check changed because the PG path's behaviour differs from the MPI path's, none of them the wrapper's own contract: - the wrapper guards (non-contiguous, 0-d, wrong-length sizes, uncovered split, float8) drop the "the bare op silently corrupts" demonstration and keep only the assertion that the wrapper rejects the input -- the PG op raises instead of corrupting, and the guard is what the test protects. - the cold-first-capture check inverts: the PG communicator exists before any op runs, so a group's first call captures into a graph straight away rather than failing; replay correctness is what it now certifies. - the two checks that measured the MPI op's corruption-on-misuse (byte-pattern float8 sum and poison / mispaired) are removed: they asserted MPI implementation detail, the PG path does not reproduce it, and neither carries a wrapper contract. The launcher's wedge sub-job and its helpers go with them. The per-batch-size graph capture copies each op output into a persistent buffer outside the shared graph pool: the ops' pooled outputs alias across graphs, so a burst of replays with no sync between them left them all holding the last writer's result. The copy, ordered after the op on the one stream, reads the right value before the next graph overwrites the pool, which is what lets the burst pass test the collective's ordering rather than die on pool aliasing. Verified on 4xGB300: allgather 21 checks, reducescatter 28 checks, all ranks, with no mpirun and no OMPI env; the modeling_v2 unit suite is green. Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
… contiguous-side offset in 64 bits token_idx * head_dim was evaluated in 32 bits and wrapped once token_idx * 64 reached 2^31, so the gather wrote rows at index >= 2^25 out of place without an error. The same patch is proposed separately against main. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…he OpenAI server's pending media loads A message parsed by the OpenAI server carries its pending media loads (coroutines) under "media"; deep-copying it raised and every image request through trtllm-serve returned HTTP 400. The renderer never reads them, so they are left out of the copy. The same patch is proposed separately against main. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…s a template parameter and adds (128, 3) DeepSeek-V4.1-Flash routes 128 experts with top-3. The custom MoE routing kernel and its thop binding take top-k as a template parameter, on top of the image-routing variant already in the tree. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…-V4 sparse backend and the runtime hooks the modeling_v2 target uses The in-tree DeepSeek-V4.1 runs on the CSA2 sparse backend and is unchanged. The modeling_v2 target keeps the DeepSeek-V4 backend in a V4.1 layout mode, entered only through v41_layout.use_deepseek_v41_layout_mode(model_config): 128-token pages, the Engram token role (DeepSeek-V4's disaggregated layer ids are unchanged), encoder/decoder replay and native / FP4 / FP8 KV storage. Adds Triton custom ops for NVFP4/MXFP8 KV fake quantization, read-back and E4M3 transcode, trtllm::dsv4_local_to_global_indices, and small hooks: autotuner fallback_tactic(), per-request last-chunk/context-logits flags in attention metadata, admission refusal in resource managers, profiling of a multimodal encoder that runs inside the forward, and a model-class opt-in for lazy safetensors. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…target calls, and its certified cells on them Promotes bmm_out, fp4_block_scale_moe_runner and fp4_quantize (from deepseek_v3's ops/) and mxe4m3_mxe2m1_block_scale_moe_runner and mxfp8_quantize (from gpt_oss's ops/) back to the shared catalog, since a second target now calls them; the owning targets import them from there. Adds the cells the DeepSeek-V4.1 target calls shared entries at, and makes the claims gates scan every target's ops/ as well as the catalog. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…rapper entries The kernels only the DeepSeek-V4.1 target calls, each an OpWrapper whose CELLS, ARCHS and note are its certified contract, with a test that drives every cell: sparse MLA (bf16 and FP8 KV), the compressor and indexer kernels, DeepGEMM FP8xFP4 MQA logits, Engram hash/gather/gate, mHC fusions, MXFP8 GEMM and quantization, NVFP4/MXFP8 KV conversions, gate_forward, allreduce and pipeline send/receive. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…get (sm_103, unified parallelism) A ModelingV2Core with PrefillTarget and DecodeTarget, ModelWeights tables, the Engram host table and the target's own input processor, routed by TRTLLM_MODELING_V2 for DeepseekV41ForConditionalGeneration checkpoints. Text-only DeepseekV41ForCausalLM and TRTLLM_MODELING_V2=off stay on the built-in model. create_input_processor asks the modeling_v2 router first; under auto, a configuration no target serves is refused rather than pairing the target's processor with the built-in model. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…s and quickstart support for the modeling_v2 DeepSeek-V4.1 target GSM8K and MMMU gates and encoder/decoder replay bitwise gates for the target on GB300, keyed apart from the in-tree DeepSeek-V4.1's references; test-db entries sized from measured durations (the target's ops tests post-merge, its multi-GPU tests on GB300); CODEOWNERS; quickstart_advanced accepts --spec_decode_algo DSPARK. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
This branch has not been deployed
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.
@coderabbitai summary
Description
This adds a modeling_v2 target for DeepSeek-V4.1-Flash-NVFP4 on GB300 (sm_103), with unified parallelism: tep4, dep4, tp4, plus tep8/dep8 and PP2 members.
It is opt-in. With
TRTLLM_MODELING_V2=off(the default) nothing changes, and DeepSeek-V4.1 runs on the in-tree implementation from #19727. WithTRTLLM_MODELING_V2=require|auto, aDeepseekV41ForConditionalGenerationcheckpoint routes to this target:DeepseekV41ForCausalLMstays on the built-in model.The target was developed and gated on GB300 before #19977 and #19727 landed. This PR ports it onto #19977's layout and onto main as it is after #19727.
What changed (one commit per area)
indexerKCacheGather/Scattercompute the contiguous-side offset in 64 bits. Also proposed alone against main: #20190.trtllm-serveno longer fail with HTTP 400. Also proposed alone against main: #20191.v41_layout.use_deepseek_v41_layout_mode(model_config): 128-token pages, an Engram token role, replay cursors, native/FP4/FP8 KV, and the long-decode pool constraint.trtllm::dsv4_local_to_global_indices.fallback_tactic(), per-request last-chunk/context-logits flags, admission refusal, profiling of an in-forward multimodal encoder, a lazy-safetensors opt-in.bmm_out,fp4_block_scale_moe_runner,fp4_quantize,mxe4m3_mxe2m1_block_scale_moe_runnerandmxfp8_quantizemove back to the shared catalog, since a second target calls them.ops/.models/deepseek_v41/deepseek_v4_1_flash_nvfp4__sm_103__unified/: core, PrefillTarget/DecodeTarget, ModelWeights tables, the Engram host table, and its own input processor._router_indexentries.create_input_processorasks the modeling_v2 router first. Withoff, it returns nothing.auto, a configuration no target serves is refused instead of pairing this target's input processor with the built-in model.quickstart_advanced.py --spec_decode_algo DSPARK.In-tree behaviour is unchanged.
__post_init__takes the V4 path.Test Coverage
All results are on GB300 (4× GB300 per node, sm_103), with a native build of this branch (
103-real) in the image pinned byjenkins/current_image_tags.properties. They were measured on the integration branch, whose final tree equals this PR's. Multi-rank runs usempirun … trtllm-llmapi-launch, because the container forbidsMPI_Comm_spawn.Model level (
TRTLLM_MODELING_V2=require):max_seq_len=1048576with chunked prefill. The pass key is exact at 30,680, 245,720 and 998,360 prompt tokens; the ~1M prompt takes 35 s.TRTLLM_MODELING_V2=auto: tp2, which no target serves, is refused with the=offmessage; tep4 selects the target.Unit and op tests:
auto_fallback: 136 passed. The target's own tests: 95 passed.test_modeling_deepseekv4.py: 46 passed.py_executor: 226 passed.csa2plustest_modeling_deepseekv41.pygive 747 passed, plus 218 with the release checkpoint.Not run here:
test_moe_routingcases and beam-searchmulti_process, which needMPI_Comm_spawn.csa2 test_two_rank_module_projection_parityfails here with an NCCL "invalid usage" ingetComm. No file on its path differs from the base, and it was not re-run on a base build.Lint: the repo's pre-commit, run twice over every changed file; pass 2 rewrites nothing. The one failing hook,
validate-test-lists, reports the three stalewaives.txtentries (lines 258–260) that #19977 leaves behind. This PR movesfp4_quantize's andmxfp8_quantize's tests back to the paths those waives name, but the waived functions no longer exist.Known limitations and follow-ups
mpi_disabledpath (PPCommTorch) is not.B300-PyTorch-Post-Merge-1is a single split, and it may need another.deepseek_v4tokenizer that its reference was measured with.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.