Repository navigation
Inkling chunked prefill reuse v2 - #18050
Draft
KleinBlueC wants to merge 122 commits into
Draft
KleinBlueC wants to merge 122 commits into
KleinBlueC wants to merge 122 commits into
Conversation
Add the Inkling NVFP4 model to the _torch stack: modeling_inkling.py, a Triton score_mod attention backend (SWA + global + relative-position bias) over KVCacheManagerV2, HF NVFP4 weight mapper/configs, trtllm-gen blockScaleMoe runner with sink-renorm routing, reasoning-parser/effort rendering, lm_eval post-processing, and the inkling_* unittest suite. Progress (working snapshot): - Component validation passes in isolation on the TP=4 NVFP4 / trtllm-gen MoE stack: weight load & accounting, attention source-activation replay, MoE replay, and full-model source-logit replay. - Baseline (cuda_graph=off, overlap=off) accuracy vs the SGLang reference: GSM8K 0.916 vs 0.972 (-5.6pt), full MMLU 82.22 vs 85.66 (-3.44pt). The gap is dominated by runaway / non-terminating generation on hard prompts: 76% of GSM8K errors are 7k-8.6k-token spirals where the model reaches the answer but never emits EOS, while SGLang commits. Per-layer localization is exhausted; the residual is a diffuse fp4 / kernel-family divergence (bf16 Triton attention + trtllm-gen fp4 MoE vs SGLang flashinfer), not a single fixable layer bug. - Enabled (cuda_graph=on) is blocked by a decode collapse (B2) localized to the global-attention block under CUDA-graph capture/replay at TP=4: h_attn goes non-finite at the first global-attention layer once decode crosses a KV-page boundary; a reduced-model TP=2 harness reproduces it. Next: decode-side termination fix for the baseline runaway (finish_reason-based detection + EOS/stop handling), and resolve B2 in the global-attention-under- graph path. Excludes build artifacts (libtensorrt_llm.so), locks, and the ext/ submodule. Signed-off-by: kleinc <kleinc@nvidia.com>
…(6*448)) Root cause of the baseline NVFP4 accuracy gap vs SGLang. The Inkling checkpoint ships routed-expert activation calibration as a RAW `.input_amax`, but the fused-MoE loader / trtllm `fp4_quantize` expect the ModelOpt per-tensor `input_scale = amax / (E2M1_MAX * E4M3_MAX) = amax / (6*448)`. The mapper renamed `.input_amax` -> `input_scale` WITHOUT the conversion, making the activation global scale 2688x too small: every routed expert's fp4 output was ~0.62 rel_rms from the bf16 ground truth (7.6x SGLang's 0.082). Mirrors sglang inkling.py:1222 / inkling_common/dense_mlp.py:497. Positional bisection confirmed the weight/block-scale/gate-up-interleave layout was already element-wise correct (24/24 L3 experts, 18.8M elems, 0 diverged) -- the defect was ONLY the activation input scale, shared by both TRT MoE backends. Fix: `inkling_weight_mapper._map_expert` divides `.input_amax` by (6*448). Result (baseline: cuda_graph=off, overlap=off, TP=4, CUTLASS MoE, fp4 fix active): - L3 routed-expert vs bf16 truth: rel_rms 0.624 -> 0.081 (== SGLang 0.082) - GSM8K paired 5x100: TRT 0.916 -> 0.968 vs SGLang 0.974; mean_delta -0.056 -> -0.006 (Gate 2 PASS, within +/-0.02); runaway canary passes; collapse 0 - MMLU paired 6x100: TRT ~0.822 -> 0.862 vs SGLang 0.872; gap -3.44pt -> -1.0pt (Gate 3 accuracy PASS); B-bias error-is-B 65.9% -> 40%, TRT-B 33.2% -> 29.8% Residual (not this fix): TRT (CUTLASS) vs SGLang (flashinfer) fp4 kernel-family near-tie noise -- accuracy-neutral, not bit-reproducible at batch>1. The strict within-2pp-of-gold MMLU B-bias criterion is model-inherent (SGLang itself over-picks B) and is left to human adjudication, not a TRT defect. Also adds env-gated per-layer/per-module dump instrumentation (modeling_inkling.py dump_sink; inkling_perlayer_localize_test.py) and the trtllm-gen MoE backend path used to localize the bug. Signed-off-by: kleinc <kleinc@nvidia.com>
…ring graph capture
Root cause of the Inkling TP=4 enabled-runtime (cuda_graph=on + overlap) decode collapse
("B2"): the AutoTuner-selected all-reduce (`tunable_allreduce`, AUTO strategy) is not
CUDA-graph-capture-safe. Frozen into a decode graph, the tuned tactic produces a non-finite
result on replay, so decode goes NaN from the first global-attention layer and collapses to
a token-0 repeat ("Paris!!!!"). Eager is finite at identical metadata; the fault is baked
into the captured graph.
Localized by single-variable determinism isolation (autotuner ON -> collapse, autotuner OFF
-> clean), NOT by op-fingerprints -- any in-graph probe shifts the graph memory pool / tactic
selection and suppresses the bug (a Heisenbug).
Fix (`tensorrt_llm/_torch/distributed/ops.py`, AllReduce.forward): during CUDA-graph capture
(`torch.cuda.is_current_stream_capturing()`), skip `tunable_allreduce` and fall back to the
static, graph-safe `all_reduce_op(AUTO)`. Warm-up / eager (not capturing) still autotune, so
steady-state performance is unchanged. This is a mainline TRT-LLM robustness fix (any TP
model that captures an AUTO all-reduce benefits), not Inkling-specific.
Confirmed: enabled generation_parity is BIT-IDENTICAL to the baseline (cuda_graph=off) --
tf_mismatch/neartie/confident=16/11/5, freerun_collapse=0, identical logit_checksum; enabled
5x100 GSM8K (0.966 vs SGLang 0.974) and MMLU (0.875 vs 0.872) within +/-0.02, zero errors,
no regression vs the accepted CUTLASS baseline.
Also includes env-gated per-layer/per-op B2 localization instrumentation
(`modeling_inkling.py` dump_sink / INKLING_FP*; `inkling_fp_localize_test.py`), zero-cost
when the INKLING_FP* env is unset.
Signed-off-by: kleinc <kleinc@nvidia.com>
…MLP tower, image fusion (WIP) WIP -- Stage-1 progress snapshot on the Inkling NVFP4 multimodal tower, on top of the accepted text tower. Not finished: the vision path is verified clean but the MMMU Accounting gap is still open (decode-side), and audio / MTP are still deferred. Committed to record current progress, not as a complete feature. Text decode/attention/MoE paths are untouched. Model / config: - configs/inkling.py: add image_token_id (200054, the in-vocab chat-template <|unused_200054|>) and audio_token_id. The SGLang-internal -101 sentinel is rejected by TensorRT-LLM's executor token-id validation; the two ids are interchangeable for parity since both are overwritten by vision embeddings. - modeling_inkling_vision.py (new): hMLP vision tower InklingVisionModel and InklingInputProcessor, which expands the <image> placeholder to one token per vision patch and attaches vision_patches_bthwc features. - modeling_inkling.py: InklingForConditionalGeneration registers the input processor, builds the vision tower as a replicated bf16 submodule, and fuses per-patch embeddings into the text stream via fuse_input_embeds with explicit text/mm indices (OOV-safe). Fixes the "image not visible" hallucination: the fused stream must NOT be re-normed, since SGLang scatters raw vision rows in after embed_norm; the extra RMSNorm corrupted the image rows. Tests / diagnostics (tests/unittest/_torch/modeling/, 25 new files): MMMU harness alignment against the SGLang scorer, input-processor and vision-tower unit checks, image e2e / fusion / logit-replay / generation-parity drivers, and the localization probes used to isolate the vision-vs-decode split (vision verified bitwise-clean; the residual Accounting gap is decode-side). These require GPU + the NVFP4 checkpoint (TP=4) and were not run for this commit. Signed-off-by: kleinc <kleinc@nvidia.com>
…s_token SamplingParams._setup() looked for the end-of-generation token in only two places: tokenizer.eos_token_id, then generation_config.eos_token_id. When a checkpoint provides neither, end_id stayed None, nothing could terminate a request, and every generation silently ran to max_tokens. That combination is reachable. Multi-part chat formats have no single terminator -- a message end, an end-of-sampling marker and a document separator are distinct tokens -- so such checkpoints register their control tokens under extra_special_tokens / additional_special_tokens, neither of which populates tokenizer.eos_token_id, and declare the real stop token as eos_token_id in config.json. Those checkpoints also tend to ship no generation_config.json. _setup() already received hf_model_config but never consulted it. Observed on a checkpoint of that shape: responses ran to the token limit while emitting the configured eos_token_id up to 441 times in a single response. Fall back to hf_model_config.eos_token_id when the first two sources yield nothing. A list value sets end_id from its first entry and appends the rest to stop_token_ids, mirroring the existing generation_config path. Priority is otherwise unchanged: an explicit SamplingParams(end_id=...) still wins and tokenizer.eos_token_id still takes precedence over config.json, so models that already resolved an end_id see no behavioural change. Warn when all three sources come up empty. Generating to the cap on every request with no diagnostic is the part that makes this expensive to find. Add unit tests for tokenizer priority, the config fallback, the list form, an explicit end_id not being overridden, the all-empty case, and a missing hf_model_config. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Adds the audio and video modalities alongside the existing vision path.
Audio:
* InklingAudioPreprocessor -- dMel feature extraction (mel basis, hz<->mel,
per-frame bin quantization) producing int32 [N, 80] bins, one audio token
per frame.
* InklingAudioModel -- codebook encoder (bin m occupies codebook rows
[m*V, (m+1)*V), summed over bins) plus optional final norm, bf16, loaded
strictly from the real checkpoint's audio tensors.
* Placeholder expansion and fail-loud count checks in the input processor.
Video:
* sample_video_frames / sample_video_as_images / DecodedVideo -- frame
sampling ported to match SGLang's sample_video_frames semantics.
* The <image>-per-frame path: every sampled frame becomes its own image
span through the existing vision tower.
Audio and video land together because both wire into the same
InklingInputProcessor.assemble dispatch; splitting them would leave an
intermediate commit whose processor references helpers that do not exist yet.
Tests (all GPU-verified on TP=4):
* inkling_audio_tower_test.py -- 10 passed, incl. real-weight CUDA forward
(AUDIO_TOWER_CUDA_OK, out=(n_frames, 6144) bf16 finite) and a
reference-math allclose(atol=1e-5) check of the codebook sum.
* inkling_video_utils_test.py -- 12 passed, incl. a port of SGLang's
test_video_utils.py::test_sample_video_frames_lengths (same 4 cases and
the same expected frame indices) and a real-weight multi-frame CUDA
forward (VIDEO_TOWER_CUDA_OK, out=(total_patches, 6144)).
* inkling_audio_e2e_test.py / inkling_video_e2e_test.py -- strict TP=4
end-to-end smokes: every prompt must be finite, non-empty and
non-collapsed. Green both baseline (5/5) and with cuda_graph+overlap
enabled (5/5).
Scope note: these prove the modalities run and are shape/dtype/finiteness
correct; they are not a cross-stack numerical parity check against SGLang.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
`logits_rows` is a view of the same slice that is written back, so the assignment self-overlaps and torch rejects it with "... refer to a single memory location. Please clone()...". The processors edit in place, so the write-back is redundant anyway; cloning the source makes it overlap-safe. Only reached for requests carrying a py_logits_post_processor, so normal requests are unaffected. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
… parity runs evaluate/interface.py gains summarize_generation_stats/log_generation_stats, and lm_eval emits a greppable GEN_STATS marker per batch. Without it a runaway / no-EOS regression hides behind a parseable answer buried in a wall of repeated text -- the failure mode that cost several bring-up iterations. The MMMU harness/runner changes carry the sharded union runs used for the TRT-vs-SGLang comparison (shard plan, incremental atomic per-item writes so a wall-killed shard keeps everything it scored, cap/max_seq plumbing), plus unit coverage of the answer parser. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Adds streaming-vs-batch equivalence over arbitrary split points, tool-call and repetition segmentation, and end-tokens split across deltas. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Standalone probes used to localize the TRT-vs-SGLang vision divergence: prefill logits, transformers-reference decode, termination behaviour, and a per-layer activation dump. Diagnostics only -- not part of any test suite. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Two in-progress pieces, neither reached by a runtime path yet. MTP static tier: InklingMTPConfig carries the checkpoint's num_nextn_predict_layers / chain_hidden_post_norm / local_layer_ids plus the per-depth banded-attention geometry injected from the text tower, so draft depths in local_layer_ids get SWA head geometry and the rest stay global. The weight mapper gains inkling_expected_mtp_keys() and accounts model.mtp.* as consumed rather than deferred when an mtp_config is supplied; the default mtp_config=None leaves the text-tower accounting byte-identical. Unit coverage for config parse, weight accounting, the BF16/unquantized requirement, and per-depth banding against the checkpoint shapes. MMMU harness: INKLING_MMMU_TEXT_ONLY reruns the same items as pure text (image placeholder stripped, no image attached) so the shared decoder is exercised at the identical bs / cap / overlap regime without the vision path. Default-off — the vision scoring path is unchanged when unset. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…del code The cuda-graph decode collapse is a defect in symmetric all-reduce: a captured NCCL_SYMMETRIC reduce whose send buffer is unregistered while its recv buffer is a registered NCCL window corrupts the run at a 12288 B message. Inkling hits that size exactly -- hidden 6144, bf16, one decode token -- so the first global-attention layer goes non-finite and decode collapses to a repeated token 0. Revert the shared-code mitigation in distributed/ops.py: that file is now byte-identical to its pre-Inkling state, so no other model's all-reduce path changes on our account. Mitigate in modeling_inkling.py instead. The all-reduces that trigger this are built by generic modules -- attention o_proj, MoE down_proj -- so the strategy cannot be passed at construction without editing shared code; rebuilding each AllReduce after super().__init__() keeps the mitigation model-local. Each rebuilt instance carries the module's own mapping and dtype over, so strategy is the only delta. Pinning ONESHOT also drops the window requirement, since AllReduce only takes an NCCL window under NCCL_SYMMETRIC/NCCL/AUTO -- two of the five trigger conditions go away, not just one. Active by default; INKLING_ALLREDUCE_STRATEGY=AUTO restores stock behaviour, and the defect with it, for A/B runs. Measured on job 5728192, alternating arms in one job against one binary: default 0/3 collapse, AUTO 3/3, 331 modules swapped. Cost: symmetric is disabled on every Inkling all-reduce, eager included, and roughly a third of captured decode all-reduces pick it today. The performance impact is unmeasured and should be measured before this is treated as final. This is containment, not a root-cause fix -- the defect remains for any other model that meets all five conditions. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Removes the bring-up's debug scaffolding and completes the deliverables the Inkling NVFP4 change was missing. No model behaviour changes. Deleted 41 debug/localization test files (divergence probes, dump/isolate helpers, teacher-forcing and graph-capture localizers) that existed to find the bring-up's defects and have no role now that it works. Nothing imports them: every deleted basename was grepped across the tree, zero references remain. The two environment variables the model still reads are the ones worth keeping — INKLING_ALLREDUCE_STRATEGY (the escape hatch for the ONESHOT all-reduce mitigation) and INKLING_MOE_BACKEND (the trtllm-gen MoE kernel select) — and both are documented; the 19 debug-only INKLING_* knobs are gone. Completed the deliverables: - docs/source/models/supported-models.md — architecture row, multimodal feature-matrix row, and footnote [^14] covering modality coverage, the unsupported set (MTP, LoRA, function calling, constrained decoding, EPD, mm-hash caching), and the all-reduce mitigation plus its escape hatch. - TestInkling_NVFP4::test_nvfp4 added to the shared multimodal accuracy file rather than a standalone per-model test, with a sourced MMMU reference. - Registered that id in test-db/l0_b200.yml and qa/llm_function_core.txt. - Dropped the orphaned inkling_vision_tower_artifact.json; the surviving regression test consumes the generated artifact instead. Static checks: git diff --check clean, every modified Python file compiles, no stale references to the deleted files. Runtime coverage as of the last completed suite: unit tiers green (vision 43, audio 10, video 12, text 31, collect 95 with zero residual debug files) and GSM8K cg0ov0 parity at delta=0.0. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
The MTP / next-N draft work never reached a runtime path: nothing builds the draft layers and nothing consumes them. Only the config plumbing and the load-accounting existed, so remove them rather than ship dead code. - configs/inkling.py: drop InklingMTPConfig, InklingConfig.has_mtp and _as_mtp_config. mtp_config goes back to a plain retained blob, so the checkpoint still round-trips. - inkling_weight_mapper.py: drop inkling_expected_mtp_keys and the consumed_mtp bucket; inkling_account_checkpoint loses its mtp_config parameter. - test_modeling_inkling.py: drop the four Stage-9 MTP tests. - supported-models.md: the footnote no longer claims the draft weights are weight-accounted. Checkpoint accounting is unaffected: "model.mtp." stays in INKLING_DEFERRED_PREFIXES, so the draft weights are classified as deferred exactly like the audio and vision blocks, and `unaccounted` stays empty. Verified: every touched file compiles, and no reference to any removed symbol remains anywhere in the tree. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…ences MMMU already had a reference entry; the two text benchmarks the bring-up actually measured did not, so the numbers lived only in run logs. Both references are the cached SGLang NVFP4 measurement under the same harness, matching how the MMMU entry was sourced: - GSM8K 95.53 (full set, 0.9553). TRT-Inkling was validated against it with the paired 5x100 protocol: flexible-parse mean 0.968, and all four cuda-graph x overlap-scheduler corners scored 0.98-0.99 on the 100 paired items where SGLang scored 0.99. - MMLU 85.66 (full Hendrycks, 14042 samples, weighted_accuracy 85.6573). The TRT side was measured with the 5-seed text-regression protocol on the harness' 114-item subset (83.33 / 84.21 / 85.09 / 85.96 / 87.72, mean ~85.3). The comment says so explicitly: it tracks the reference, but no full-set TRT-LLM MMLU run has been recorded yet. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…that matters
The bring-up left 22 Inkling files under tests/unittest/_torch/modeling. Most
were per-defect regressions or SGLang-alignment scaffolding whose outcome is
now covered end-to-end by the MMMU/GSM8K accuracy tests. Keep one test per
thing that nothing else covers, drop the rest.
Kept:
- test_modeling_inkling.py config parse, layer classification, weight
accounting over the real checkpoint
- inkling_vision_tower_test.py vision tower
- inkling_audio_tower_test.py audio tower (no accuracy benchmark covers it)
- inkling_video_utils_test.py video utils (no accuracy benchmark covers it)
- inkling_input_processor_test.py multimodal input processor
- inkling_moe_backend_select_test.py the INKLING_MOE_BACKEND kernel select
inkling_mmmu_real_align_test.py was doing double duty: the vision-tower and
input-processor tests import its MMMU item fetch/cache and its importlib
loader. Split that half out as inkling_mmmu_fixtures.py (a fixture module, no
tests) and drop the SGLang-alignment machinery with the rest.
Deleted (16): the mmmu align/harness/run/parser set, image_prompts, the three
per-modality e2e smokes, image_fusion, image_norm_fix, attn_decode_meta,
gate_up_deinterleave, kv_manager_v2, generation_parity and
source_logit_replay.
Every kept file compiles and no reference to a deleted module remains in the
tree.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
The multimodal file already carried TestInkling_NVFP4::test_nvfp4 for MMMU, so the vision path had an integration test but the shared text decoder -- what GSM8K and MMLU actually exercise -- had none. Adds TestInkling_NVFP4 to test_llm_api_pytorch.py running both text benchmarks against the references recorded earlier. It mirrors the multimodal class: NVFP4 assert, 16384-token budget for the long chain of thought, and extract_inkling_content as the post-processor so the <|content_thinking|> channel is dropped and only the visible answer is scored. Registered the new id in both lists that already carry the multimodal one: test-db/l0_b200.yml and qa/llm_function_core.txt. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…keep CUTLASS The Inkling routed experts now run only the default CUTLASS MoE backend. The trtllm-gen blockScaleMoe path was an opt-in experiment behind INKLING_MOE_BACKEND=TRTLLM, never the default, and it needed a dedicated routing enum plumbed all the way into the CUDA runner to work. Restored to their pre-bring-up state (the additions were Inkling-only, so these files are now byte-identical to 44f0521): - cpp/.../trtllmGenKernels/blockScaleMoe/runner.h the InklingSinkRenorm = 9 enum value and its name case - cpp/.../trtllmGenKernels/blockScaleMoe/runner.cu the precomputed-routing dispatch branch - _torch/modules/fused_moe/routing.py the matching Python enum value and its autotuner-dummy mapping Removed from the model: - _inkling_trtllm_moe_backend / _moe_config_with_trtllm_backend and the frozen-config copy they needed to retarget moe_backend - InklingMoeRoutingMethod._trtllm_backend, so routing_method_type is always Unspecified and requires_separated_routing goes back to the default - the per-layer INKLING_MOE_SELECT backend-introspection log Also deletes inkling_moe_backend_select_test.py, which existed only to cover that knob, and the stale comment that pointed at it as the fix for the fused combine's cross-row non-determinism. INKLING_ALLREDUCE_STRATEGY is now the only environment variable the model reads. Everything compiles and git diff --check is clean. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
The MMLU comment added in 5f28628 claimed no full-set TRT-LLM run had been recorded. That is wrong. One exists: 2026-07-20, all 14042 Hendrycks samples, weighted_accuracy 82.2176 -- 3.44 points below the 85.66 reference, and it failed the bring-up's 2-point gate. The gap was later traced to an fp4 expert-GEMM family difference (CUTLASS/trtllm-gen vs SGLang's flashinfer) plus non-terminating generation on hard prompts. What is true is narrower: the full set has not been re-measured since those fixes. Post-fix MMLU evidence is subset-only -- a 570-item stratified canary at 84.21 (-1.45, inside the gate) and a 5-seed 114-item regression averaging ~85.3. The comment now says exactly that, including that meeting 85.66 at full scale is unverified. GSM8K's comment was not wrong but was too vague about coverage. TRT-LLM has never been measured on the full 1319-item set: the evidence is the paired 5x100 protocol (flexible mean 0.968), the four cuda-graph x overlap corners at 0.98-0.99 on 100 paired items, and one early full-set attempt that completed only its first 120-item chunk (0.925 vs 0.975). Spelled out. Both files still record the SGLang reference as the accuracy value, so TestInkling_NVFP4::test_nvfp4 grades against a bar the model has not been shown to clear at full scale. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…timodal
The module already carried the image, audio and video paths, so the "vision"
name no longer described it. Rename it and reorganize the contents into four
labelled sections -- vision tower, audio tower, video tower, and the shared
multimodal input processor -- so each modality is easy to locate.
Behavior-preserving. Along the way:
* factor the duplicated tower weight-loading into ``_load_tower_weights``
and the duplicated RMSNorm into a single shared ``InklingRMSNorm``
(was ``InklingVisionRMSNorm``, used by both towers);
* factor the input processor's repeated media-list coercion and its
placeholder/feature-row count checks into small helpers, dropping a
tautological per-item check (``num_tokens`` is built from ``num_patches``);
* trim the docstrings to what the code needs, dropping the development-time
stage/goal references and reference-implementation file paths.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Bring-up leftovers that do not belong in the production model:
* cuda_graph_runner: drop the "[cuda-graph] CAPTURED/REPLAYED" evidence
logging and its once-only gate flag. This was pure instrumentation added
to prove the runtime really captured and replayed a graph, and it sat in
shared (non-Inkling) code.
* modeling_inkling: drop the INKLING_ALLREDUCE_STRATEGY environment knob.
The ONESHOT pin is a correctness mitigation, not a tuning option, so it
is applied unconditionally instead of via an A/B toggle (and the log line
that reported it is gone).
* modeling_inkling: drop the explicit-window short-conv decode path
(InklingShortConv.forward_decode, the decoder layer's third branch, and
the attention's conv_states/return_conv_state arguments). Only the
bring-up replay harness ever passed those; the runtime always drives the
short convs through the per-request state pool.
* modeling_inkling: drop the attention's decode_seq_lens / decode_page_table
/ skip_kv_write arguments. The runtime publishes decode metadata into the
layer's stable GPU buffers before capture, so the pre-supplied static
tensors had no caller left; the eager fallback that builds them from the
host block table stays.
* modeling_inkling: drop InklingConvStateCache.reset (unused) and merge
InklingConvRuntime.from_metadata into build (the split existed only so the
replay harness could publish slots itself).
Also rewrite the comments and docstrings across the Inkling _torch files to
describe the code as it stands: no development stage/goal numbers, job ids,
absolute paths into local reference checkouts, or references to the deleted
replay harness. The stale "audio / vision / MTP are deferred" module docstring
now reflects that only MTP is unimplemented.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
The Inkling unit tests were five files carrying a bring-up harness rather than a test suite: they loaded the SGLang reference from an absolute path in a local checkout, downloaded MMMU rows over the network into a gitignored cache, wrote JSON artifacts, hardcoded a personal /lustre checkpoint path, and shipped __main__ runners that printed comparison tables. None of that can run in CI. Fold everything multimodal into test_modeling_inkling_multimodal.py -- one file covering the vision, audio and video paths plus the shared input processor -- with small synthetic configs and inputs: no checkpoint, no GPU, no network, no SGLang import, no artifacts. 27 tests, well under a second. Slim test_modeling_inkling.py the same way. The config and layer-classification tests now build their config explicitly instead of requiring the checkpoint, so they actually run; the weight-accounting and tensor-shape tests keep the checkpoint (index JSON only, no weights) and resolve it through the standard llm_models_root() with an INKLING_CHECKPOINT override, so they skip cleanly instead of always skipping on a path that only existed on one machine. Removed: inkling_vision_tower_test.py, inkling_input_processor_test.py, inkling_audio_tower_test.py, inkling_video_utils_test.py (folded in) inkling_mmmu_fixtures.py (MMMU downloader + SGLang loader; no longer used) Also drop the .gitignore entries for the deleted caches and artifacts. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…nkling PR Run the repo's pre-commit formatters (isort, yapf, ruff, ruff-format) over the files this PR touches, and fix the eleven ruff-legacy D205/D209 docstring regressions it flags in tensorrt_llm/evaluate/interface.py and tests/unittest/llmapi/test_reasoning_parser.py. Formatting and docstring wording only; no behavior change. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Two bring-up leftovers outside _torch:
* evaluate: drop summarize_generation_stats / log_generation_stats and both
call sites (Evaluator.evaluate and LmEvalWrapper). This logged a greppable
GEN_STATS line with the finish_reason mix and generated-token distribution
so a runaway/no-EOS regression would be visible while bringing the model up.
It is observability scaffolding, not part of the eval contract: it never
influenced what was scored, nothing consumed the marker, and no test
covered it.
* docs: drop "The text decoder is also usable standalone (text-only) via the
InklingForCausalLM architecture" from the Inkling footnote. It is not true.
InklingForCausalLM carries no @register_auto_model, so it is not in
MODEL_CLASS_MAPPING and no checkpoint can select it; the config registry has
no inkling_text entry either. The class is the base that
InklingForConditionalGeneration derives from, and the inkling_text handling
in config_utils/_util exists for that nested sub-config, not for standalone
loading. The published checkpoint declares InklingForConditionalGeneration
and the text-only accuracy test loads it through that same architecture, so
registering the class would advertise a path with no checkpoint to exercise
it and no test coverage.
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
All three references were previously the SGLang measurement. GSM8K and MMLU carried a long note that TRT-LLM had never been confirmed at full scale, and MMMU claimed TRT-LLM matched SGLang item-for-item, which is not what the data shows. GSM8K and MMLU have now been run on the complete datasets, TRT-LLM and SGLang side by side, at TP=4 with CUDA graph and the overlap scheduler on, batch 8, driven by the same client so prompt rendering and scoring are shared code. MMMU was already measured on the full 858-item pool during the bring-up: GSM8K TRT-LLM 96.13 (1268/1319) SGLang 96.59 (1274/1319) -0.45 pt MMLU TRT-LLM 85.35 (11985/14042) SGLang 85.48 (12003/14042) -0.13 pt MMMU TRT-LLM 79.72 (684/858) SGLang 78.32 (672/858) +1.40 pt The references now hold the TRT-LLM numbers, which is what CI reproduces, and the comments record the SGLang cross-check. The earlier -3.44 pt MMLU gap is gone now that the routed experts run on CUTLASS only and the non-terminating generations are fixed. The MMMU comment also drops the item-for-item claim and records what the pool actually shows: the two stacks agree on aggregate accuracy but 74/858 (8.6%) of items score differently in both directions, which is inherent to two different NVFP4 implementations resolving near-ties differently deep in a long-CoT decode. Aggregate accuracy, not per-item identity, is the metric to gate on. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Brings the branch up to 93e7c91 (402 upstream commits since 44f0521). Six files conflicted; all resolutions keep both sides: * supported-models.md upstream took footnote [^14] for MiniCPM-V, so the Inkling note moves to [^15] (3 table rows renumbered) * _util.py union of the import blocks; the new is_hybrid_linear V2 guard sits alongside the Inkling one * model_engine.py upstream's steady-gen fast path and the Inkling runtime publish are independent additions; upstream also fixed the logits self-assignment defect by dropping the redundant write-back, which supersedes the clone() this branch carried * py_executor_creator union of the config_utils imports * lm_eval.py six additive-kwarg conflicts, union throughout * test_lm_eval.py both sides appended an independent test section One resolution is not textual. Upstream's new _apply_steady_gen_fast_prepare returns early from _prepare_tp_inputs, bypassing the tail where Inkling publishes its short-conv pool rows and per-layer attention decode metadata. git reports no conflict there -- the two sides touch different lines -- but an unchanged generation-only batch is exactly Inkling's steady-state decode, so a plain merge would have left the captured graph reading capture-time page tables. The fast path now publishes before returning; it is a no-op for models without a registered CONV_STATE_MANAGER. Also initializes the new 3rdparty/MSA submodule (upstream NVIDIA#16291), which setup.py now requires. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…d it Two review findings on InklingDecodeMeta, both on the decode critical path. **Stale publication.** `ready` was a plain flag that only `refresh` cleared, so it described "some forward published these buffers", not "this forward did". A forward that reached the decode kernel without the engine having prepared -- an eager fallback, or a prepare path that returns early -- read the previous step's seq_lens and page table. The same requests advance num_cached_tokens_per_seq every step, so a page table one step out of date silently drops a newly allocated page. Request ids and the batch size cannot detect that. `prepare_inkling_attn_decode` now bumps a shared per-forward counter before doing anything else, and `ready` is a property comparing the layer's recorded epoch against it. A layer that is not refreshed this forward reports False by construction, so the attention takes its host fallback instead of reading stale buffers. This is the same contract the upstream merge broke once already (a new fast path bypassed the publish site); making the invariant self-checking means the next such refactor fails loudly rather than silently. **Per-layer work on every token.** `refresh` ran for all 66 layers per decode step, and each call allocated two pinned host tensors and built the page table with one `torch.tensor` per row -- 132 cudaHostAllocs and 66*num_gen tensor constructions per token. The total-KV lengths do not depend on the layer, so they are now built once per step and shared; the per-layer pinned page-table staging is allocated once and reused, and filled through its numpy view with a vectorized row assignment. Behavior is unchanged when the engine publishes every forward, which is the runtime path. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Three review findings on the way a checkpoint is read. The weight mapper resolved its geometry from whatever config it was handed. ``ModelLoader.load`` initializes it with the TOP-LEVEL ``InklingConfig``, which keeps the decoder geometry under ``text_config`` and has no ``vocab_size`` or ``n_routed_experts`` of its own, so ``preprocess_weights`` raised ``AttributeError`` on the unpadded-vocab lookup and ``_map_expert`` silently fell back to ``tensor.shape[0]`` for the expert count. The mapper now resolves ``text_config`` itself, which covers both entry points: the loader-supplied mapper and the one ``InklingForConditionalGeneration.load_weights`` builds. ``InklingTextConfig`` unconditionally set ``moe_intermediate_size`` from ``intermediate_size``, discarding the value a ``config.json`` had already stored through ``**kwargs``. It now reads the kwarg back first, matching how ``max_position_embeddings`` is handled two lines above. ``load_modelopt_quant_config`` gained ``_algo_or_none`` to map both JSON null and the strings "none"/"null" onto ``None``, but the MIXED_PRECISION branch still built ``QuantAlgo`` directly, so a mixed-precision checkpoint spelling ``kv_cache_quant_algo: "none"`` would still raise. That branch uses the helper now. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Each of these turned a bad input into wrong output rather than an error. The Triton kernels map a query head to its KV head as ``cur_head // kv_group_num``. Both entry points computed ``kv_group_num = num_heads // num_kv_heads`` without checking divisibility, so a non-divisible pair produced a wrong mapping inside the kernel with no diagnostic. Both now assert. The prefill entry also unpacked ``total_tokens`` without using it (ruff RUF059); renamed to ``_total_tokens``. ``InklingInputProcessor.assemble`` dispatches on the token value, and both placeholder ids are config-overridable. If a config gave them the same value every occurrence was consumed as an image, the audio branch never ran, and the placeholder-count check reported the mismatch against the wrong modality. The constructor now rejects the collision. ``InklingReasoningParser.finish`` cleared ``_buffer`` but left ``_kind`` at the last block kind, so reusing an instance for a second stream would route that stream's first segment by the previous stream's channel. ``parse`` already resets both. Not reachable today (the factory builds one parser per request), but the asymmetry is the kind that survives a refactor. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
The supported-models footnote named GSM8K and MMMU as the accuracy gates, but TestInkling_NVFP4 also evaluates MMLU on the text decoder. Name both. The MMMU test docstring claimed Inkling was validated against SGLang "item-for-item". It is not: 74 of the 858 pool items score differently, in both directions, and the harness gates the aggregate score. Say what is actually validated, and why per-item identity is the wrong metric here -- two different NVFP4 implementations resolve near-ties differently deep in a long-CoT decode. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
Second pass. Comment+docstring share 44% -> 36% overall across the package, 2295 -> 2010 lines. This round cut the long module and class docstrings (params, cache_manager, metadata, conv_state, backend, kernels, __init__) and the multi-paragraph inline blocks in cache_manager and conv_state, keeping the reasoning that cannot be read off the code and dropping restatement. Behaviour-neutral: with every docstring stripped, all seven modules are node-for-node identical to HEAD as ASTs. Signed-off-by: kleinc <kleinc@nvidia.com>
89 -> 74 tests, 2131 -> 1867 lines. Two categories removed. Ten asserted on source text or message wording via inspect.getsource rather than on behaviour: they break when an implementation is rephrased and pass when it is rewritten wrongly but keeps the string. Where the property mattered it already has behavioural coverage -- the global-tp guard, for instance, is covered by four tests that check actual head/channel shapes under and without attention DP. Five were tautological or subsumed: a config default compared against a constant declared in the same file, the absence of an enum member never added on this branch, an 'INKLING not in llm_args' check that test_the_fake_config_stays_out_of_the_user_facing_schema already makes more strongly, and two page-table/conv-pool tests whose assertions are a subset of a neighbour's. Two deletions leave a property with no other test, both weakly guarded to begin with: the MoE reduce-scatter preference and the expert-sharding delegation were only ever asserted as substrings of the implementation. test_modeling_inkling_multimodal.py is left alone: 23 tests over 466 lines, all asserting behaviour, with no redundancy found -- cutting there would remove coverage rather than noise. Signed-off-by: kleinc <kleinc@nvidia.com>
…assertion Two unrelated things the CI run surfaced. Formatting: isort, ruff and ruff-format rewrote import blocks and line wrapping in five files. Applied as the hooks produce them. test_serving_and_profiler_interface_contract asserted that the input processor exposes get_dummy_mm_data_for_tokens. That name no longer exists anywhere -- upstream renamed it to get_dummy_mm_data in b3e6813, which reached this branch through the main merge at 9fea415. The test never went red because the local runner only ran test_modeling_inkling.py; adding the multimodal file to that invocation is what exposed it. Pre-existing, unrelated to the attention work. Signed-off-by: kleinc <kleinc@nvidia.com>
74 -> 68 tests, 1868 -> 1745 lines. Same standard as the first pass: assertions on source text rather than behaviour, where the property has behavioural cover. - the KV-cache-creator guard: three tests already call the guard and expect the raise, so asserting the function name appears in the caller's source adds nothing; - 'INKLING not in the family selector': the positive routing is asserted by test_backend_and_cache_manager_both_resolve_from_the_sparse_registry; - the pyexecutor file-naming scan and the package-docstring wording check: conventions, not behaviour; - the framework's own KvCacheConfig default, which is not Inkling's to pin. test_the_decode_fast_path_is_guarded_by_metadata_type goes too, and it was mine from earlier today -- keeping a source-grep of my own while deleting nine of someone else's would not be a standard. What that guard actually protects is exercised end to end by the TLLM_INKLING_PT_CROSSCHECK run, not by a substring. 103 passed / 3 skipped across both Inkling test files (job 6282255). Signed-off-by: kleinc <kleinc@nvidia.com>
All three describe behaviour this branch changed, and all three would mislead a user into a configuration the code now rejects or accepts differently. attn_backend is no longer pinned to the name INKLING. Selection moved to sparse/registry.py, so what _assert_inkling_attn_backend enforces is the TRTLLM *family* -- which is the framework default, so the constraint is weaker than the footnote claimed, not stronger. Chunked Prefill read Untested in both tables. It is not untested: the KV-cache guard raises on it, so No is the accurate value. Disaggregated serving was listed as EPD only. The guard refuses the regular kind too -- it keys on cache_transceiver_config, not on the EPD path -- so the footnote now names both, and says why these are refusals rather than gaps: the four per-layer short convolutions hold their window as per-request state outside the KV cache, so a reused prefix has no window to restore and a transferred request arrives without one. Docs only; no code change. Signed-off-by: kleinc <kleinc@nvidia.com>
Signed-off-by: kleinc <kleinc@nvidia.com> # Conflicts: # tests/integration/test_lists/qa/llm_function_core.txt
- registry: group the Inkling imports with the other backend imports - metadata: drop the redundant kv_layout assignment (HND is the base default) and the create_cuda_graph_metadata override (the base already re-runs __post_init__ on the copy) - metadata: compute the page-table row per layer on demand instead of caching a whole-model row map - model_config: move the sparse_attention_config injection next to the other architecture-derived sparse config updates - _util / config_utils: drop comments that explain code which is not there Signed-off-by: kleinc <kleinc@nvidia.com>
Review feedback: the Inkling files carried far more prose than the code around them. Comment + docstring lines are down from 34% to 25% of the Inkling sources (the sparse/ neighbours run ~24%), keeping the reason for each non-obvious decision and dropping the surrounding essay. Signed-off-by: kleinc <kleinc@nvidia.com>
The metadata mirrored the pool into ink_conv_cache purely so the decoder could reach it; the manager is its only owner, so read it from there and keep only the per-forward split (ink_conv_rt) on the metadata. That split still has to be published from prepare(): building it performs a pinned host->device write into the graph-stable state_indices buffer, which must not be captured. Signed-off-by: kleinc <kleinc@nvidia.com>
… geometry class Two review comments, one change because they share the same tests. The metadata subclass now carries no fields and overrides only prepare(), for the one thing that cannot live anywhere else: the conv pool's slot write is a host->device copy into a buffer the captured decode graph aliases, so it has to run every step, outside that region, and see the CUDA-graph padding rows. Everything else moves out -- the decode seq lens, page-table row and page divisor become free functions in page_table.py over the base metadata, and the per-forward conv split is sliced in model.forward, views only. The backend picks its graph-safe path on the borrowed page table being allocated rather than on the metadata type. _InklingConvGeometry is gone. Only pretrained_config has to be resolved before super().__init__(), because the base does not keep it; mapping, max_batch_size, pp_layers and max_draft_len are all on self by the time _build_cache_config runs, so the rest is derived there -- max_draft_len was duplicating the base outright. The pool dtype still comes from the model config rather than binding_to_torch_dtype(self.dtype): that one is the KV cache dtype, which is fp8/nvfp4 once a checkpoint asks for a quantized KV cache, while this pool holds unquantized pre-conv activations. Signed-off-by: kleinc <kleinc@nvidia.com>
Line-wrap the conv byte-size comprehension, collapse an over-wrapped assert, and drop an import left unused when the registry assertion moved to its own test. Signed-off-by: kleinc <kleinc@nvidia.com>
The metadata rework introduced fresh prose that did not match what it replaced. Trim it back: the four conv-geometry properties take one-line comments instead of docstrings, and the metadata, page-table, conv-runtime and backend docstrings lose the restatement. Inkling's attention_backend files are back at 24% comment and docstring lines, which is where the sparse/ neighbours sit. Signed-off-by: kleinc <kleinc@nvidia.com>
_create_kv_cache_manager never sees the top-level Inkling config, so unwrapping it was dead code. Every path that reaches here passes model_engine, and model.model_config.pretrained_config is already InklingTextConfig because InklingForConditionalGeneration hands the text sub-config to super().__init__; the one call site that passes model_config explicitly resolves it to self._model_engine.model.model_config for MTP, and Inkling has no MTP. What is left is the per-layer local/global KV-head split, which is geometry rather than config nesting -- the same shape as the is_gemma4_hybrid branch above it -- so the identity test is inlined at its only remaining use. Signed-off-by: kleinc <kleinc@nvidia.com>
Comments only, no code change. The refusal helper's docstring restated all three raise messages at length (45 lines for a 35-line body); the rest were paragraphs where the neighbouring branch gets a line or two -- the Gemma4 block right beside the per-layer KV-head branch carries no comment at all. Signed-off-by: kleinc <kleinc@nvidia.com>
modeling_inkling.load_weights imported the weight mapper from the middle of its body; the mapper pulls in nothing from this module, so it belongs at the top like modeling_nemotron_h's. The audio preprocessor also re-imported stdlib io under an alias when the module already imports it; only the soundfile import stays local, since that one is an optional dependency a decoded-waveform caller never needs. Signed-off-by: kleinc <kleinc@nvidia.com>
Two halves that only work together: - a chunked-context prefill Triton kernel -- BLOCK_M query tiling like the prefill kernel, paged KV reads like the decode kernel, and absolute positions (num_cached + q_pos) so the causal mask, the sliding window and the relative-bias distance stay right across a chunk boundary. rel_logits remains indexed by the packed query row, since only new tokens are queries. No gather is needed: _run_context writes this chunk's K/V into the pages first. - has_initial_state derived from num_cached_tokens_per_seq, so the conv window a preceding chunk left in the pool row is declared and therefore consumed. Either alone is wrong: the conv fix without the kernel leaves attention losing history while the convs keep it, which is a subtler wrong answer than the original refusal. _needs_chunked_context routes per batch; all-fresh batches keep the packed kernel and pay no page indirection. INKLING_FORCE_CHUNKED_ATTN=1 forces the chunked path for A/B measurement, since this model's sampled output is not reproducible between two generate() calls unless the autotuner is off. reject_unsupported_inkling_kv_cache_features drops its chunked-prefill raise and keeps the block-reuse and disaggregated-serving ones. Validated on Inkling-small-NVFP4: the chunked kernel is bit-identical to one-shot prefill at the layer level, 63 GPU parity tests including realistic page layouts, and GSM8K over the full 1319 questions -- delta -0.08 pt with exact McNemar p=1.000 against a same-code noise floor that itself moves 0.15 pt. Signed-off-by: kleinc <kleinc@nvidia.com>
The conv state is already registered as V2 SSM layers, so V2 can attach a short-conv window to the block committed at a snapshot ordinal and hand it back with that block on a prefix hit. What was missing is the policy: where snapshots may be taken, and a guarantee that iterations actually end there. prepare_expect_snapshot_points declares the boundaries. py_executor finds it by hasattr and the scheduler will not end a context chunk anywhere else -- which is the half a model cannot do for itself. A snapshot can only be taken where an iteration ends, so without a say in where chunks end, capture depends on the operator having picked a max_num_tokens smaller than the shared prefix; nothing states that and nothing checks it. needs_block_aligned_context_chunks drives FORCE_CHUNK and is named for the property rather than the model family. Inkling is deliberately not folded into is_hybrid_linear, which would also route it through extract_mamba_kv_cache_params and the Mamba conv-state layouts it cannot satisfy. ReusableStateSnapshotMixin lifts commit-at-snapshot-points out of MambaHybridCacheManagerV2, until now its only implementation. Every self reference in those three methods was already a base-class member, so the move is behaviour-preserving; both managers inherit it. The interval reuses kv_cache_config.mamba_state_config.periodic_snapshot_interval rather than adding a second user-facing knob for the same quantity. Coarse on purpose: a snapshot costs the whole model's conv window, so one per block over an 8k prompt would be hundreds of MiB for a single request. _augment_tokens_for_block_reuse keeps multimodal requests out of the shared radix tree. The base rewrites image spans into cache-key tokens carrying the item's content digest precisely because a placeholder token id is the same whatever image is behind it -- and it can only do that for a request carrying multimodal_hashes, which Inkling produces for none of them. Without the override, two prompts holding different images match as the same prefix and one is served the other's KV: measured on MMMU over the same 32 items, 81.3% with reuse off against 31.3% with reuse on, restored to 81.3% by salting position 0 with the request id. A blanket refusal would be the wrong shape, since it would also remove the text reuse that measures good. Block reuse is no longer refused outright, only when no snapshot policy is configured -- which is the case the refusal was really protecting against. The reuse counter is kept, and counts the FIRST context chunk only. A continuation chunk also arrives with context_current_position > 0, and counting those made a reuse-disabled arm report hits. It stays because this feature can be completely inert while the server is healthy, every request succeeds, and the score matches the baseline -- there is no other signal that separates "reuse works" from "reuse never ran". Verified on Inkling-small-NVFP4, GSM8K 100 samples, chunk 256, interval 256, two runs per arm: reuse on 97/97 with 63/64 requests reused (longest 512 tokens), reuse off 95/98 with 0 reused. The off arm is the negative control the counter had to pass before the on-arm number meant anything, and its own 3 pt spread is the floor those numbers are read against. Signed-off-by: kleinc <kleinc@nvidia.com>
…pported Chunked Prefill moves to Yes in both matrices. KV Cache Reuse moves to Yes in the text matrix and stays No in the multimodal one, because a request carrying image, video or audio input never reuses by design. The footnote states the two settings reuse takes -- enable_block_reuse plus a positive mamba_state_config.periodic_snapshot_interval -- because reuse with only the first is still refused, and a reader who tries it gets a NotImplementedError with no hint in the table that a second knob exists. Signed-off-by: kleinc <kleinc@nvidia.com>
The chunked-prefill and block-reuse commits came over from a branch written before the comment-density pass on this one, and carried its prose with them. Comment and docstring lines are down from 27% to 21% of what the two features add, which is where the sparse/ and pyexecutor/ neighbours sit. What is kept is the reason for each non-obvious decision -- why the conv fix is necessary but not sufficient without the kernel, why a snapshot interval is coarse on purpose, why multimodal requests get a private reuse chain, why the warning is said at startup rather than per request. What went is the retelling of history around them: job ids, superseded readings, and paragraphs restating what the next three lines of code already say. Signed-off-by: kleinc <kleinc@nvidia.com>
At 328 words it was six times the longest other footnote in the table (55) and had turned into a design document: why the conv window cannot be restored, how the reuse tree is keyed, what the chunked kernel reads. That belongs in the code, which now says it. What is kept is what a reader of the table cannot get anywhere else: what the checkpoint is, the reasoning parser, the two constraints enforced at load, the unsupported list, the extra knob KV cache reuse needs and its text-only scope, and why CUDA graph reads No. Signed-off-by: kleinc <kleinc@nvidia.com>
…agers Three changes, all in service of one rule: the two features must not edit code other models depend on. kv_cache_manager_v2.py and mamba_cache_manager.py go back to their upstream state. The reuse commit had lifted commit-at-snapshot-points out of MambaHybridCacheManagerV2 into a ReusableStateSnapshotMixin so both managers could inherit one copy; that is the better factoring in the abstract, but it edits the file every model imports. The three methods now live on InklingHybridCacheManager itself. They are byte-for-byte the Mamba versions modulo the error strings, and a test compares them as ASTs, so the copy cannot drift silently. Why they exist at all, since registering the conv window as V2 SSM layers is supposed to be enough: it is enough for the storage and for the snapshot itself -- _KVCache.commit() captures the SSM state at the block holding the last committed token. What it is not enough for is the contract that commit asserts, "commit_min_snapshot requires commit() to start or end at history_length". The base try_commit_blocks commits as prefill advances and never sets history_length, so with SSM layers present it does not hold. These three methods are what maintain it. The hand-rolled reuse counter is gone. The framework already reports exactly what it counted, per SSM lifecycle and over /metrics: KVCacheV2SsmSnapshotIterationStats carries iterSnapshotLookups / Hits / Misses / ReusedTokens / Aligned / Unaligned. Reading those also gives the offline harness a gate it never had, since worker logs do not surface there. Both features' code is now fenced off by section banners in each file it touches, so what belongs to chunked prefill and what belongs to block reuse can be read -- and reverted -- without tracing call sites. Signed-off-by: kleinc <kleinc@nvidia.com>
…ling The refusal claimed the short-conv windows sit "outside the paged KV cache", which stopped being true when they became KV cache manager V2 SSM layers: they are pool memory with a life cycle, like Mamba's. The real reason is one level up. The cache-transfer path enumerates state pools by cache-manager class -- kv_extractor.py branches on isinstance(manager, MambaHybridCacheManagerV2) -- and Inkling's manager is not that class, so its SSM layers are never described as state and fall through to the attention branch. Nothing errors; the generation instance simply resumes from a window nobody transferred. Worth being precise about because it is also the shape of the fix: a Mamba layer group for Inkling, not a new transfer mechanism. Signed-off-by: kleinc <kleinc@nvidia.com>
The pool is V2 SSM-layer memory, but the row a request used was picked by a private free list in InklingConvStateCache. Two allocators over one buffer is harmless only while nothing outside the model touches it -- true until block reuse, which is exactly when V2 starts writing there: _KVCache.commit() copies a restored recurrent state into the slot V2 assigned (storage.slot_address(..., new_slot.slot_id, ...)), and the model is meant to read that slot back through get_ssm_block_base_index(). Mamba does; Inkling published its own row instead. So a reuse hit restored the conv window into a row the kernels never read, and the row they did read may belong to another live request. Nothing raises, and GSM8K cannot see a three-tap depthwise window being wrong -- the same shape as the multimodal key defect this branch already fixed. The pool now asks the manager, which asks V2. The private free list stays for standalone construction (the unit tests), where nothing else touches the pool. Its zeroing does NOT move with it: a V2 slot may hold a window V2 has just restored, and zeroing on first sight would throw it away. Padding sentinels keep the reserved rows -- they have no cache entry, and V2 reports a negative index for a request whose state is not resident, so both read as "no slot" rather than as row -1. Signed-off-by: kleinc <kleinc@nvidia.com>
…buffers Two things, because the second is what made the first visible. The geometry block -- _conv_tp_size, _reserve_attention_dp_slot, _num_conv_request_slots, _num_reserved_conv_slots, _conv_bytes_per_slot -- was deleted by the refactor commit on this branch. The edit that removed the reuse counter took everything between it and the "V2 configuration" banner, and those properties sat in between. Nothing caught it: __init__ reads them through self, so it is an AttributeError at construction, not an import error, and the unit job that would have failed was still queued. Restored, and the same class now gets a structural check in the test file. The four per-layer conv buffers become one holding [k | v | attn | mlp] along the channel axis, with the section widths kept beside it. The four windows share a length, so this is a pure layout change: the sections are views, and the causal_conv1d ops read conv_state's strides off the tensor (causalConv1dOp.cpp:246) rather than assuming contiguity, so a strided section is fine. It is also the shape the rest of the framework already expects. A state transfer describes one pool per layer group plus the byte widths to split a slot on when the two sides disagree on TP -- MambaLayerGroup.conv_section_bytes, which Mamba fills with [x | B | C]. With four pools there was nothing to map to; _get_pool_roles' search over four roles goes away with them. Signed-off-by: kleinc <kleinc@nvidia.com>
validate_decode_layout reads kv_cache_block_offsets and goes straight to .shape, so a batch that reaches a kernel without one reports "'NoneType' object has no attribute 'shape'" from inside a helper, with nothing naming the field or saying who was supposed to fill it. The check existed before the page-table helpers moved into page_table.py and came back as free functions; the test that pins it (test_metadata_rejects_a_missing_block_offsets_tensor) has been failing since. 131 passed / 3 skipped (job 6358837), against 130 passed / 1 failed before. Signed-off-by: kleinc <kleinc@nvidia.com>
First GPU run of the branch: 3 failed, 6 errors, 207 passed. Both causes were introduced by edits that only ever saw py_compile and pre-commit. The manager read self.impl in __init__ to resolve its conv layer group. That attribute is not guaranteed to exist by the end of the base constructor, and a manager built for a unit test has none, so every test that constructs one died with AttributeError. It is now resolved on first use, where impl is certainly live -- which is also the only place the value is needed. The multimodal reuse tests lost their @pytest.fixture decorator when they were ported across the refactor: the port extracted top-level blocks by the name of the def they contain, and a decorated fixture is two lines with the name on the second. Six tests then errored at setup rather than running. 216 passed / 3 skipped (job 6364346), against 207 passed / 3 failed / 6 errors. Signed-off-by: kleinc <kleinc@nvidia.com>
…urce too The framework has these counters -- KVCacheV2SsmSnapshotIterationStats, per SSM lifecycle -- and they are the right source, so this counter was removed in favour of them. Measured since: a served deployment's /metrics carries no kvCacheIterationStats* key at all, over 1001 iterations, with enable_iter_perf_stats on and iteration_stats_interval at its default of 1 (jobs 6381795, 6382310). There is nothing there to read yet. Without a gate, an accuracy pair says nothing: reuse that worked and reuse that never ran both look like the reuse-off number. That has already cost this campaign two rounds. The line also reports how many requests got their conv row from V2 rather than from the pool's private free list, because the two ways reuse can be silently wrong are indistinguishable by accuracy: no reuse at all, or reuse restored into a row the kernels never read. First run of the branch to produce evidence rather than a number: reuse on 63/64 requests reused, longest 512 tokens, conv rows from V2 64/64 reuse off 0/64 conv rows from V2 64/64 Both arms 97/100 on GSM8K n=100 (jobs 6382688 / 6382689). Signed-off-by: kleinc <kleinc@nvidia.com>
There were three: a packed prefill for all-fresh requests, a paged one for requests carrying cached history, and decode. The packed one existed on the theory that page indirection is worth avoiding in the common case, and it cost a standing obligation -- the two prefill kernels had to stay bit-identical at num_cached == 0, which is a property nothing enforces outside a test. That theory did not survive measurement. At the serving level the all-paged build (chunked off, so every request all-fresh) came in at 183.3 output tok/s against the packed build's 194.7 -- but the packed build re-measured at 176.2, so the difference is inside the baseline's own spread. Accuracy was already settled: GSM8K over the full 1319, -0.08 pt with exact McNemar p=1.000 against a noise floor of -0.15 pt. So the paged kernel becomes the prefill kernel. _needs_chunked_context and the INKLING_FORCE_CHUNKED_ATTN knob go with the second path -- with one kernel there is nothing to route between, and nothing to force for an A/B. The parity suite keeps its shape: its oracle was the packed kernel and is now the same kernel run one-shot over a fresh cache at num_cached == 0, which is what "splitting a prompt must not change the answer" was always measured against. 215 passed / 3 skipped (job 6390415). Net -199 lines. Signed-off-by: kleinc <kleinc@nvidia.com>
Open
1 task done
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
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
Test Coverage
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.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.