[None][feat] Support Inkling-NVFP4 model

[KleinBlueC]

Summary by CodeRabbit

• New Features

  • Added support for Inkling text and multimodal models, including image, audio, and video inputs.
  • Added NVFP4 checkpoint loading and optimized attention, caching, and hybrid reasoning support.
  • Added Inkling reasoning parsing for visible responses, reasoning content, tool calls, and streaming output.
  • Added Inkling evaluation support for GSM8K, MMLU, and MMMU benchmarks.

• Bug Fixes

  • Improved EOS token fallback and handling of unquantized configuration values.

• Documentation

  • Documented Inkling capabilities, supported modalities, checkpoint details, and known limitations.

Description

Adds PyTorch-backend support for Inkling (thinkingmachines/Inkling-NVFP4), a
RoPE-free hybrid-attention MoE reasoning model with vision and audio towers.
Registered as InklingForConditionalGeneration.

Architecture

• No RoPE — position comes from a learned relative-position bias added
  pre-softmax, per (query token, head, relative distance).
• Hybrid attention — 55 local sliding-window layers (window 512, 16 KV heads)
  interleaved with 11 global full-causal layers (8 KV heads).
• Four causal short convolutions per layer — two inside attention (K, V) and two
  on the residual stream, carrying per-request state across decode steps.
• Sigmoid-gated MoE — top-6 of 256 routed experts with an additive selection
  bias, then a log-sigmoid renorm spanning the routed and two shared-expert logits.
• muP + unpadded vocab — hidden states divided by logits_mup_width_multiplier
  before the head; logits sliced from 201024 down to 200058.
• Multimodal — hMLP vision tower and dMel audio codebook tower, both BF16, each
  fusing one row per patch / per frame into the text embedding stream. Video is
  multi-frame images; there is no separate video encoder.

NVFP4 covers the text decoder and routed experts (layers 3–65); layer-2 experts,
attention, shared experts, and both towers stay BF16 per the checkpoint's
hf_quant_config.json.

Implementation

Area	Notes
modeling_inkling.py	Decoder stack, MoE, short convolutions, multimodal wrapper
attention_backend/inkling_triton.py	Paged prefill + decode Triton kernels applying the relative bias as a score_mod. No existing fused backend exposes such a hook (context FMHA is disabled for kRELATIVE, trtllm-gen rejects a relative bias, FlashInfer has no additive per-token bias). rel_logits is static-shape, so the decode kernel is CUDA-graph capturable
modeling_inkling_multimodal.py	Vision tower, audio tower, video frame sampling, and the shared input processor that expands one <image> placeholder per patch and one <audio> per dMel frame. Pure numpy + torch preprocessing
configs/inkling.py	Config classes rebuilt from the checkpoint's nested dicts (transformers does not ship Inkling), plus the per-layer classification helpers used by both the model and KV-cache sizing
checkpoints/hf/inkling_weight_mapper.py	Checkpoint accounting (every key is consumed-text / deferred / unaccounted) and load-path remapping, including de-interleaving the gate/up projections
pyexecutor/*	Registers the short-conv state pool under a new CONV_STATE_MANAGER. Pool rows and attention decode metadata are published into stable CUDA buffers eagerly from input prep, so the captured decode forward does no host→device copy and each replay reads the current batch (same stable-pointer pattern as Mamba2Metadata)

KV cache: the per-layer KV-head split (local 16 / global 8) structurally requires
KVCacheManagerV2 — V1's unified pool would coerce it to one value and mis-size the
per-layer KV bytes, a correctness bug. The model defaults to V2 and the incompatible
path raises rather than silently downgrading.

Serving / eval plumbing (small, each needed end to end):

• --reasoning_parser inkling for Inkling's typed-content blocks, with streaming.
• needs_raw_special_tokens so a delimiter-based reasoning parser actually sees its
  markers (previously only the tool-parser path preserved them).
• end_id falls back to the config's eos_token_id — checkpoints whose terminator
  lives only in config.json otherwise never stop and run to max_tokens.
• hf_quant_config.json may spell "no quantization" as the string "none", which
  previously reached QuantAlgo("none") and raised.
• Warmup dummy ids are bounded to the head vocab; a vocab padded above the output
  vocab otherwise fails KV-cache estimation with "Token ID out of range".
• --post_process_fn inkling / inkling_mmmu for trtllm-eval.

Accuracy

Measured 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:

Task	TRT-LLM	SGLang	Δ
GSM8K (1319)	96.13 (1268/1319)	96.59 (1274/1319)	−0.45 pt
MMLU (14042, 5-shot)	85.35 (11985/14042)	85.48 (12003/14042)	−0.13 pt
MMMU (858)	79.72 (684/858)	78.32 (672/858)	+1.40 pt

Not supported in this release

MTP / speculative decoding (the checkpoint ships next-N draft weights; nothing
builds or loads them), LoRA, function calling, constrained/guided decoding, EPD
disaggregated serving, and multimodal-hash prefix caching (refused loudly per
modality, so multimodal requests still run — just uncached).

One workaround worth flagging for review: every Inkling all-reduce is rebuilt with
ONESHOT after construction. Under CUDA-graph capture a symmetric all-reduce
corrupts the run when its send buffer is unregistered while its recv buffer is a
registered NCCL window at a 12288 B message — which Inkling hits exactly (hidden
6144, bf16, one decode token), sending the first global-attention layer non-finite.
The all-reduces involved are built by generic modules (attention o_proj, MoE
down_proj), so pinning the strategy afterwards keeps the mitigation model-local.

Test Coverage

Accuracy (integration) — added to l0_b200 and llm_function_core:

• accuracy/test_llm_api_pytorch.py::TestInkling_NVFP4::test_nvfp4 — GSM8K + MMLU
  on the text decoder.
• accuracy/test_llm_api_pytorch_multimodal.py::TestInkling_NVFP4::test_nvfp4 —
  MMMU on the vision path.

Unit — CPU-only, no checkpoint / GPU / network needed, all well under a second:

• unittest/_torch/modeling/test_modeling_inkling.py — config parsing,
  registration, per-layer classification; plus checkpoint-gated weight accounting
  and tensor-shape checks that read only the safetensors index and skip cleanly
  when the checkpoint is absent.
• unittest/_torch/modeling/test_modeling_inkling_multimodal.py — the three media
  paths and the input processor on synthetic configs: hMLP scale plan and module
  tree, the fold's value preservation, dMel preprocessing, the audio codebook
  forward against a reference, frame sampling, and the fail-loud placeholder
  contract.
• unittest/llmapi/test_reasoning_parser.py — full-parse and streaming equivalence
  for the Inkling parser, including control tokens split across delta boundaries.
• unittest/llmapi/test_sampling_params.py — the end_id config fallback.
• unittest/others/test_lm_eval.py — the offline post-processing hook.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Inkling NVFP4 text and multimodal support is added to TensorRT-LLM. The changes include model configuration, checkpoint mapping, Triton attention, convolution state management, vision/audio processing, reasoning parsing, evaluation integration, serving updates, documentation, and accuracy coverage.

Inkling support

Layer / File(s)	Summary
Configuration and checkpoint loading tensorrt_llm/_torch/configs/*, tensorrt_llm/_torch/models/checkpoints/*, tensorrt_llm/_torch/model_config.py	Adds Inkling configuration types, registration, quantization normalization, checkpoint accounting, and NVFP4 weight remapping.
Vision, audio, video, and input assembly tensorrt_llm/_torch/models/modeling_inkling_multimodal.py, tensorrt_llm/_torch/models/modeling_inkling.py	Adds image, audio, and video processing. Multimodal inputs expand placeholders and fuse tower features into decoder embeddings.
Text model and execution runtime tensorrt_llm/_torch/attention_backend/inkling_triton.py, tensorrt_llm/_torch/models/modeling_inkling.py, tensorrt_llm/_torch/pyexecutor/*	Adds hybrid attention, short-convolution state, MoE layers, paged KV decode, CUDA-graph metadata, and executor resource management.
Reasoning, post-processing, and serving output tensorrt_llm/llmapi/*, tensorrt_llm/evaluate/*, tensorrt_llm/sampling_params.py, tensorrt_llm/serve/openai_server.py, tests/unittest/llmapi/*, tests/unittest/others/test_lm_eval.py	Adds Inkling typed-content parsing, streaming control-token handling, evaluation post-processors, EOS fallback behavior, and raw-special-token serving.
Documentation and accuracy integration docs/source/models/supported-models.md, tests/integration/defs/accuracy/*, tests/integration/test_lists/*	Documents Inkling support and registers text and multimodal NVFP4 accuracy tests and references.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Request
  participant InklingInputProcessor
  participant InklingForConditionalGeneration
  participant InklingAttention
  participant inkling_triton
  Request->>InklingInputProcessor: submit text and optional media
  InklingInputProcessor->>InklingForConditionalGeneration: assemble tokens and modality features
  InklingForConditionalGeneration->>InklingAttention: run hybrid decoder layers
  InklingAttention->>inkling_triton: execute prefill or paged decode
  inkling_triton-->>InklingAttention: return attention results
  InklingForConditionalGeneration-->>Request: return generated output

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16511: Adds related NVFP4 inference support and changes shared PyTorch quantization and model configuration paths.

Suggested reviewers: juney-nvidia

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 47.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly identifies the feature and the supported Inkling-NVFP4 model.
Description check	✅ Passed	The description explains the purpose, implementation, limitations, accuracy results, test coverage, and checklist items in detail.

✨ Finishing Touches 💡 1

⚔️ Resolve merge conflicts 💡

• Resolve merge conflict in branch inkling-nvfp4-bringup

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 15

🧹 Nitpick comments (21)

tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)

7815-7815: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use PascalCase test class names and update direct selectors.

Rename both classes to TestInklingNVFP4. Update the listed selectors to the new class name.

• tests/integration/defs/accuracy/test_llm_api_pytorch.py#L7815-L7815: rename TestInkling_NVFP4.
• tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py#L683-L683: rename TestInkling_NVFP4.
• tests/integration/test_lists/qa/llm_function_core.txt#L812-L813: update the two direct class selectors.
• tests/integration/test_lists/test-db/l0_b200.yml#L355-L356: update the two direct class selectors.

As per coding guidelines, use PascalCase for classes.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py` at line 7815, Rename
both TestInkling_NVFP4 classes to TestInklingNVFP4 in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:7815-7815 and
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py:683-683.
Update both direct class selectors to TestInklingNVFP4 in
tests/integration/test_lists/qa/llm_function_core.txt:812-813 and
tests/integration/test_lists/test-db/l0_b200.yml:355-356.

Source: Coding guidelines

tensorrt_llm/_torch/configs/inkling.py (2)

137-148: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate _local_ids with the element type and consider caching it.

_local_ids rebuilds a set on every is_local_layer call. num_kv_heads_per_layer() calls it once per layer. The cost is small, but a precomputed set is simpler and the annotation should be set[int] per the typing guideline.

♻️ Proposed refactor

-    `@property`
-    def _local_ids(self) -> set:
-        return set(self.local_layer_ids)
+    `@property`
+    def _local_ids(self) -> set[int]:
+        cached = getattr(self, "_local_ids_cache", None)
+        if cached is None or len(cached) != len(self.local_layer_ids):
+            cached = set(self.local_layer_ids)
+            self._local_ids_cache = cached
+        return cached

As per coding guidelines: "prefer built-in generic types and |".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/configs/inkling.py` around lines 137 - 148, The
_local_ids property should use the built-in generic annotation set[int] and
avoid rebuilding the set for every is_local_layer call. Precompute or cache the
converted local_layer_ids set at the configuration level, then have
is_local_layer reuse that cached set while preserving its current membership
behavior.

Source: Coding guidelines

---

195-248: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add type annotations to the constructor parameters and _as_config.

text_config, audio_config, vision_config, mtp_config, and _as_config(value) carry no annotations. The repository guideline requires annotating every function. Use dict[str, Any] | PretrainedConfig | None.

♻️ Proposed refactor

     def __init__(
         self,
-        text_config=None,
-        audio_config=None,
-        vision_config=None,
-        mtp_config=None,
+        text_config: "dict[str, Any] | InklingTextConfig | None" = None,
+        audio_config: "dict[str, Any] | PretrainedConfig | None" = None,
+        vision_config: "dict[str, Any] | PretrainedConfig | None" = None,
+        mtp_config: "dict[str, Any] | PretrainedConfig | None" = None,
         eos_token_id: int = 200006,

     `@staticmethod`
-    def _as_config(value):
+    def _as_config(
+        value: "dict[str, Any] | PretrainedConfig | None",
+    ) -> "PretrainedConfig | None":

As per coding guidelines: "Annotate every function, use None for non-returning functions".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/configs/inkling.py` around lines 195 - 248, Add type
annotations to InklingConfig.__init__ parameters text_config, audio_config,
vision_config, and mtp_config using dict[str, Any] | PretrainedConfig | None,
and annotate _as_config(value) with the same input type and an appropriate
return type. Add the required Any import, and annotate __init__ as returning
None while preserving the existing configuration handling.

Source: Coding guidelines

tests/unittest/_torch/modeling/test_modeling_inkling.py (2)

52-68: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Cache the parsed safetensors index.

_safetensors_shape calls _index(ckpt) for every tensor lookup, so the index JSON is re-read and re-parsed once per assertion. Use the module-scoped fixture pattern or functools.lru_cache to read it once.

♻️ Proposed refactor

+from functools import lru_cache
+
+@lru_cache(maxsize=None)
 def _index(ckpt: str) -> dict:
     with open(os.path.join(ckpt, "model.safetensors.index.json")) as f:
         return json.load(f)["weight_map"]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_modeling_inkling.py` around lines 52 -
68, Cache the parsed weight map used by _safetensors_shape so
model.safetensors.index.json is read and parsed only once per checkpoint. Apply
a module-scoped fixture or functools.lru_cache to _index, preserving
_safetensors_shape’s existing lookup behavior.

---

122-141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Coverage gap: the gate/up de-interleave is untested.

The tests cover accounting (inkling_account_checkpoint, inkling_nvfp4_expert_layers) and checkpoint shapes, but not the load-path remapping in inkling_weight_mapper.py. _split_interleaved_gate_up is correctness-critical: a contiguous chunk(2) instead of the strided even/odd split pairs the wrong SwiGLU channels and produces silently wrong output, with no shape error to catch it. That function is pure CPU and needs no checkpoint.

Coverage assessment: sufficient for config/accounting, insufficient for the mapper load path. Add a small CPU test in this file, for example:

💚 Suggested test

def test_split_interleaved_gate_up_takes_even_odd_rows():
    from tensorrt_llm._torch.models.checkpoints.hf.inkling_weight_mapper import (
        _split_interleaved_gate_up,
    )

    t = torch.arange(8).reshape(4, 2)  # rows [g0, u0, g1, u1]
    gate, up = _split_interleaved_gate_up(t, dim=0)
    assert torch.equal(gate, t[0::2])
    assert torch.equal(up, t[1::2])
    with pytest.raises(ValueError):
        _split_interleaved_gate_up(torch.zeros(3, 2), dim=0)

Do you want me to open an issue to track this?

As per path instructions: "suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_modeling_inkling.py` around lines 122 -
141, Add a CPU unit test in the test file for _split_interleaved_gate_up, using
an even-row tensor to verify gate receives even-indexed rows and up receives
odd-indexed rows, and assert that an odd-sized split dimension raises
ValueError. Import the mapper helper and pytest as needed; no checkpoint setup
is required.

Source: Path instructions

tensorrt_llm/_torch/models/modeling_inkling_multimodal.py (1)

1021-1028: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate the rejecting getters as non-returning.

get_num_tokens_per_image, get_num_tokens_per_audio, and get_num_tokens_per_video are annotated -> int but always raise. Use typing.NoReturn so type checkers do not report a missing return.

♻️ Proposed refactor

-    def get_num_tokens_per_image(self, *, image=None, **kwargs) -> int:
+    def get_num_tokens_per_image(self, *, image=None, **kwargs) -> NoReturn:
         self._reject_mm_hash_cache("image")
 
-    def get_num_tokens_per_audio(self, *, audio=None, **kwargs) -> int:
+    def get_num_tokens_per_audio(self, *, audio=None, **kwargs) -> NoReturn:
         self._reject_mm_hash_cache("audio")
 
-    def get_num_tokens_per_video(self, *, video=None, **kwargs) -> int:
+    def get_num_tokens_per_video(self, *, video=None, **kwargs) -> NoReturn:
         self._reject_mm_hash_cache("video")

As per coding guidelines: "Annotate every function ... use Literal, overload, TypeVar, or Protocol when appropriate."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_inkling_multimodal.py` around lines 1021
- 1028, Update the return annotations of get_num_tokens_per_image,
get_num_tokens_per_audio, and get_num_tokens_per_video from int to
typing.NoReturn, ensuring NoReturn is imported or referenced through the
existing typing conventions.

Source: Coding guidelines

tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py (3)

173-181: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Derive the expected bin geometry from the module constants.

Lines 177 and 179 hardcode 80 and 16. Those are DEFAULT_AUDIO_N_MELS and DEFAULT_AUDIO_NUM_DMEL_BINS. If a default changes, the test fails with an opaque number instead of pointing at the constant.

♻️ Proposed refactor

-    assert bins.ndim == 2 and bins.shape[1] == 80
+    assert bins.ndim == 2 and bins.shape[1] == DEFAULT_AUDIO_N_MELS
     assert bins.shape[0] >= 8, f"expected ~12 frames for 0.6s, got {bins.shape[0]}"
-    assert 0 <= int(bins.min()) and int(bins.max()) < 16
+    assert 0 <= int(bins.min()) and int(bins.max()) < DEFAULT_AUDIO_NUM_DMEL_BINS

The same applies to line 193 in test_dmel_multi_clip_concat_and_empty_clip.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py` around
lines 173 - 181, Update test_dmel_preprocess_shape_and_range and
test_dmel_multi_clip_concat_and_empty_clip to derive the expected mel-bin width
and value range from DEFAULT_AUDIO_N_MELS and DEFAULT_AUDIO_NUM_DMEL_BINS
instead of hardcoded 80 and 16. Keep the existing assertions and behavior
unchanged while referencing the module constants directly.

Source: Path instructions

---

124-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Coverage gap: image input coercion is untested.

Every vision test feeds a numpy array. _to_pil_rgb also accepts raw bytes, a PIL image, a file:// URI, and a local path, and it explicitly rejects http://, https://, and data: inputs. That rejection is a deliberate guard against the preprocessor reaching for the network; nothing pins it today.

Coverage assessment: sufficient for the tower geometry and forward, insufficient for input coercion. Add a short CPU test:

💚 Suggested test

def test_to_pil_rgb_rejects_remote_sources_and_accepts_bytes():
    from tensorrt_llm._torch.models.modeling_inkling_multimodal import _to_pil_rgb

    for url in ("http://x/y.png", "https://x/y.png", "data:image/png;base64,AA=="):
        with pytest.raises(ValueError, match="resolve"):
            _to_pil_rgb(url)
    with pytest.raises(TypeError):
        _to_pil_rgb(object())
    assert _to_pil_rgb(_image(8, 8)).mode == "RGB"

Do you want me to open an issue to track this?

As per path instructions: "suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR".

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py` around
lines 124 - 167, Add a CPU-focused test near the existing vision input tests
covering _to_pil_rgb: assert http://, https://, and data: inputs raise
ValueError matching “resolve”, unsupported objects raise TypeError, and a
generated image is accepted with RGB mode. This closes the input-coercion
coverage gap while preserving the existing tower tests.

Source: Path instructions

---

290-306: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

test_registration reaches into private registry internals.

Line 304 reads INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type and line 306 reads InklingInputProcessor._registered_model_type. Both are private. A rename in tensorrt_llm/inputs/registry.py breaks this test without any behavior change in Inkling. If the registry exposes a public lookup, prefer it.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py` around
lines 290 - 306, Update test_registration to validate processor registration
through the public lookup API exposed by INPUT_PROCESSOR_REGISTRY instead of
accessing _input_processors_cls_by_model_type. Also avoid asserting
InklingInputProcessor._registered_model_type directly; verify the expected
"inkling_mm_model" association through the public registry interface while
preserving the existing registration expectations.

Source: Path instructions

tensorrt_llm/_torch/models/modeling_inkling.py (4)

265-274: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Avoid the per-forward pageable allocation in write_state_indices.

torch.tensor(slots, dtype=torch.int32) allocates a pageable host tensor on every forward, then copies it into the pinned staging buffer. Write the slots into the pinned buffer directly.

♻️ Proposed refactor

         n = len(slots)
-        self.state_indices_cpu[:n].copy_(torch.tensor(slots, dtype=torch.int32))
+        for i, s in enumerate(slots):
+            self.state_indices_cpu[i] = s
         self.state_indices[:n].copy_(self.state_indices_cpu[:n], non_blocking=True)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_inkling.py` around lines 265 - 274,
Update write_state_indices to populate the pinned state_indices_cpu staging
buffer directly from slots, removing the per-forward torch.tensor allocation.
Preserve the existing n-based slicing, non-blocking device copy into
state_indices, and CUDA graph pointer validation.

---

1476-1504: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

The two decoder-layer paths diverge on **kwargs.

The pool path at line 1492 forwards **kwargs to self.attn. The stateless path at line 1479 does not. Any kwarg a future caller adds reaches the attention module on one path only. Forward **kwargs on both paths, or on neither.

♻️ Proposed fix

-            hidden_states = self.attn(position_ids, hidden_states, attn_metadata)
+            hidden_states = self.attn(position_ids, hidden_states, attn_metadata, **kwargs)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_inkling.py` around lines 1476 - 1504,
Make `**kwargs` handling consistent between the `conv_rt is None` and runtime
state-pool branches in the decoder layer: forward the same keyword arguments to
`self.attn` in the stateless path as in the pool path, or remove forwarding from
both paths. Preserve all existing attention and residual behavior.

---

1412-1427: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

One root cause: InklingGate.routing_method returns a new object per access, so the joint renorm cannot be shared. Because the property constructs a fresh InklingMoeRoutingMethod on every read, the instance create_moe holds cannot publish its shared_gammas back to InklingMoE.forward, which is why the renorm runs a second time per layer.

• tensorrt_llm/_torch/models/modeling_inkling.py#L1412-L1427: drop the second inkling_joint_renorm call and read shared_gammas from the stored routing-method instance.
• tensorrt_llm/_torch/models/modeling_inkling.py#L1292-L1301: build the InklingMoeRoutingMethod once in InklingGate.__init__ and return the stored instance from the property, so it can carry the per-forward shared_gammas.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_inkling.py` around lines 1412 - 1427,
Update tensorrt_llm/_torch/models/modeling_inkling.py lines 1412-1427 in
InklingMoE.forward to remove the second inkling_joint_renorm call and obtain
shared_gammas from the stored InklingMoeRoutingMethod instance. Update lines
1292-1301 in InklingGate.__init__ and its routing_method property to construct
InklingMoeRoutingMethod once and return that instance, allowing shared_gammas to
persist across the forward call.

---

969-991: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift

_build_rel_logits materializes a large fp32 tensor per layer at prefill.

The result is [T, local_heads, rel_extent] in fp32, plus a .contiguous() copy. For a global layer with T = 8192, 16 local heads and rel_extent = 1024, that is about 512 MiB live per layer during prefill, and the einsum output plus the contiguous copy can both be resident. Long-context prefill can therefore hit an out-of-memory condition that scales with the chunk size rather than with the model size.

Two mitigations to consider: emit the tensor in bf16 if the Triton kernel accepts it, and make einsum produce a contiguous result directly so the extra copy disappears.

♻️ Proposed change to drop the extra copy

-        rel = torch.einsum(
-            "thd,de->the", r.float(), self.rel_logits_proj.float()
-        )  # [T, H, rel_extent]
+        # matmul on the flattened [T*H, d_rel] view returns a contiguous result,
+        # so the trailing .contiguous() never copies.
+        rel = (r.float().reshape(-1, self.d_rel) @ self.rel_logits_proj.float()).view(
+            -1, self.local_num_heads, self.rel_extent
+        )

Please confirm the peak prefill memory against the configured max_num_tokens.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_inkling.py` around lines 969 - 991,
Update _build_rel_logits to avoid materializing both the einsum output and a
separate .contiguous() copy; make the einsum result contiguous directly where
supported, and emit bf16 when the consuming Triton kernels accept it while
preserving the required computation and scaling behavior. Verify peak prefill
memory using the configured max_num_tokens rather than only the example chunk
size.

tensorrt_llm/_torch/attention_backend/inkling_triton.py (1)

53-58: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Update the _NEG comment. _NEG is a module-level tl.constexpr, and Triton 3.6.0 supports this usage. Remove the claim that the value is inlined because @triton.jit cannot read module-level globals.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/attention_backend/inkling_triton.py` around lines 53 -
58, Update the comment above the module-level `_NEG` constant to remove the
inaccurate claim that Triton `@jit` functions cannot read module-level globals or
that the value is inlined for that reason. Preserve the existing explanation of
its finite, sufficiently negative value and its role in avoiding NaNs during
masked softmax bookkeeping.

tests/unittest/llmapi/test_sampling_params.py (1)

329-333: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add cases for a tuple EOS and an empty EOS list.

SamplingParams._setup branches on isinstance(config_eos, (list, tuple)) and then guards with if config_eos:. The tests cover only the non-empty list. The tuple path and the empty-sequence path are untested, and the empty-sequence guard is what prevents an IndexError on config_eos[0].

Coverage for the scalar, list, explicit-override, and missing-config paths is sufficient.

💚 Proposed tests

def test_setup_model_config_eos_tuple_is_handled():
    params = SamplingParams()
    params._setup(_FakeTokenizer(eos_token_id=None), _FakeConfig(eos_token_id=(7, 8)), None)
    assert params.end_id == 7
    assert params.stop_token_ids == [8]


def test_setup_model_config_empty_eos_list_leaves_end_id_none():
    params = SamplingParams()
    params._setup(_FakeTokenizer(eos_token_id=None), _FakeConfig(eos_token_id=[]), None)
    assert params.end_id is None

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/llmapi/test_sampling_params.py` around lines 329 - 333, Add
tests alongside test_setup_model_config_eos_list_sets_end_id_and_stop_tokens for
the tuple EOS path and empty EOS sequence path: verify tuple (7, 8) sets end_id
to 7 and stop_token_ids to [8], and verify an empty list leaves end_id as None
without raising an error.

tensorrt_llm/evaluate/lm_eval.py (1)

259-260: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Forward post_process_fn through super().__init__() for consistency.

keep_special_tokens is forwarded at Line 292, but post_process_fn is set later at Line 301 instead. The base constructor already accepts post_process_fn. The result is the same today, but the split assignment is easy to break during future edits.

♻️ Proposed consistency change

             preserve_caller_max_tokens=preserve_caller_max_tokens,
+            post_process_fn=post_process_fn,
             keep_special_tokens=keep_special_tokens,
         )

Then remove the duplicate assignment at Line 301.

Also applies to: 279-280, 292-292

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/evaluate/lm_eval.py` around lines 259 - 260, Update the
constructor around the visible post_process_fn and keep_special_tokens
parameters to pass post_process_fn through super().__init__() alongside
keep_special_tokens, then remove the later duplicate post_process_fn assignment.
Preserve the existing callback behavior while keeping all base-constructor
initialization together.

tensorrt_llm/serve/openai_server.py (1)

1514-1519: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Do not read the factory's private _parsers map from the server.

ReasoningParserFactory exposes only create_reasoning_parser() and keys() as public API. This code reaches into _parsers and depends on the entry being a tuple whose first element is the parser class. The adjacent tool-parser block uses the public ToolParserFactory.parsers.

If the factory later changes the stored entry shape, reasoning_entry[0] no longer yields the class. getattr(..., 'needs_raw_special_tokens', False) then returns False, skip_special_tokens stays True, and Inkling reasoning splitting silently stops working. The failure produces no error.

Add a public lookup on the factory and call it here.

♻️ Proposed public accessor

In tensorrt_llm/llmapi/reasoning_parser.py:

    `@classmethod`
    def needs_raw_special_tokens(cls, reasoning_parser: str) -> bool:
        """Return True if the named parser requires undecoded special tokens."""
        entry = cls._parsers.get(reasoning_parser.lower())
        if entry is None:
            return False
        return getattr(entry[0], "needs_raw_special_tokens", False)

In this file:

             if self.generator.args.reasoning_parser:
-                reasoning_entry = ReasoningParserFactory._parsers.get(
-                    self.generator.args.reasoning_parser.lower())
-                if reasoning_entry and getattr(
-                        reasoning_entry[0], 'needs_raw_special_tokens', False):
+                if ReasoningParserFactory.needs_raw_special_tokens(
+                        self.generator.args.reasoning_parser):
                     sampling_params.skip_special_tokens = False

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/serve/openai_server.py` around lines 1514 - 1519, Add the public
ReasoningParserFactory.needs_raw_special_tokens() accessor in
reasoning_parser.py, encapsulating the private _parsers lookup and parser-entry
inspection. Update the server logic around reasoning_parser to call this
accessor instead of reading _parsers directly, while preserving the existing
skip_special_tokens behavior.

tests/unittest/llmapi/test_reasoning_parser.py (1)

944-968: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add a test that reuses one parser across two streaming sessions.

Every streaming test builds a fresh parser or calls finish() once. No test drives parse_delta → finish → parse_delta on the same instance. That path leaves self._kind stale, as noted on InklingReasoningParser.finish in tensorrt_llm/llmapi/reasoning_parser.py.

Current coverage is otherwise sufficient for the batch and single-stream paths.

💚 Proposed test

def test_inkling_reasoning_parser_instance_reuse_after_finish():
    """A second stream on the same parser must not inherit the previous block kind."""
    parser = ReasoningParserFactory.create_reasoning_parser("inkling")
    parser.parse_delta(f"{INK_CH}first stream reasoning")
    parser.finish()

    r = parser.parse_delta("plain second stream answer")
    assert r.content == "plain second stream answer"
    assert r.reasoning_content == ""

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/llmapi/test_reasoning_parser.py` around lines 944 - 968, Add a
streaming reuse test alongside the existing Inkling parser tests that calls
parse_delta, then finish, then parse_delta again on the same parser instance.
Verify the second stream’s plain text is returned as content with empty
reasoning_content, confirming finish resets the parser state such as _kind
instead of carrying over the first stream’s reasoning block.

tests/unittest/others/test_lm_eval.py (1)

856-867: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add a case for the invalid --post_process_fn value.

The resolver test covers the three valid keys. It does not cover the else branch in _resolve_post_process_fn, which raises click.BadParameter. That branch is the only validation for the GSM8K option, which is declared as type=str rather than click.Choice. It should be pinned.

Coverage for the three supported processors is sufficient.

💚 Proposed test

def test_lm_eval_post_process_resolver_rejects_unknown_name():
    with pytest.raises(click.BadParameter):
        _resolve_post_process_fn("not_a_processor")


def test_lm_eval_post_process_resolver_passes_through_callables():
    fn, keep_special_tokens = _resolve_post_process_fn(extract_inkling_content)
    assert fn is extract_inkling_content
    assert keep_special_tokens is False

The second case pins the non-string passthrough at the top of the resolver, which LmEvalEvaluator relies on when a caller supplies a callable directly.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/others/test_lm_eval.py` around lines 856 - 867, Extend the
resolver tests around _resolve_post_process_fn with a
pytest.raises(click.BadParameter) case for an unknown string such as
"not_a_processor", covering the invalid-name branch. Also add a callable
passthrough case using extract_inkling_content and assert it is returned with
keep_special_tokens set to False.

tensorrt_llm/llmapi/reasoning_parser.py (2)

719-746: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy lift

The Inkling control-token format is implemented twice, and the two implementations disagree. InklingReasoningParser._consume and extract_inkling_content each define their own copy of the control-token alphabet and their own block state machine for the same wire format. Two divergences already exist: tool-invocation blocks route to visible content in the serving parser but are dropped in the offline extractor, and text before the first control token is kept as content in the serving parser but dropped in the offline extractor. The offline path is documented as "the offline analog" of the serving path, so these must agree. Neither divergence affects GSM8K, MMLU, or MMMU, because those generations contain no tool blocks and no leading unmarked text.

• tensorrt_llm/llmapi/reasoning_parser.py#L719-L746: extract the token alphabet and the block-classification state machine into one shared helper, then have _consume call it and map the resulting blocks onto content and reasoning. Keep the tool-block-to-content routing as the single agreed rule.
• tensorrt_llm/evaluate/post_processing.py#L164-L208: delete the duplicated _INK_* constants and _INK_CONTROL_RE, import the shared alphabet, and make extract_inkling_content reuse the shared state machine so tool blocks and leading unmarked text are classified identically to the serving path.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/llmapi/reasoning_parser.py` around lines 719 - 746, The
duplicated Inkling control-token alphabet and block state machine must be
unified. In tensorrt_llm/llmapi/reasoning_parser.py lines 719-746, extract the
shared alphabet and classification logic into a helper, then have
InklingReasoningParser._consume map its blocks to content or reasoning while
preserving tool blocks as visible content. In
tensorrt_llm/evaluate/post_processing.py lines 164-208, remove the local _INK_*
constants and _INK_CONTROL_RE and make extract_inkling_content reuse that helper
so leading unmarked text and tool blocks follow the same classification.

---

757-771: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Consider a precomputed prefix set for the streaming holdback.

_partial_control_length runs on every streaming delta. It performs up to _INKLING_MAX_CONTROL_LEN - 1 iterations, and each iteration scans all 18 control tokens. That is roughly 500 startswith calls per delta per request.

A module-level set of all strict prefixes turns the check into a hash lookup per length. Build it once next to _INKLING_CONTROL_TOKENS.

♻️ Proposed optimization

 _INKLING_MAX_CONTROL_LEN = max(len(t) for t in _INKLING_CONTROL_TOKENS)
+# Every strict prefix of every control token, for the streaming holdback.
+_INKLING_CONTROL_PREFIXES = frozenset(
+    t[:i] for t in _INKLING_CONTROL_TOKENS for i in range(1, len(t)))

         max_len = min(len(text), _INKLING_MAX_CONTROL_LEN - 1)
         for length in range(max_len, 0, -1):
-            suffix = text[-length:]
-            if any(
-                    len(suffix) < len(tok) and tok.startswith(suffix)
-                    for tok in _INKLING_CONTROL_TOKENS):
+            if text[-length:] in _INKLING_CONTROL_PREFIXES:
                 return length
         return 0

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/llmapi/reasoning_parser.py` around lines 757 - 771, Optimize
_partial_control_length by defining a module-level set containing every strict
prefix of the tokens in _INKLING_CONTROL_TOKENS, initialized alongside that
token collection. Replace the per-length any(...startswith...) scan with
membership checks against the precomputed prefix set while preserving the
longest-suffix and zero-result behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/source/models/supported-models.md`:
- Line 97: Update the Inkling footnote’s accuracy-gate description to state that
the text decoder is gated by both GSM8K and MMLU, while retaining MMMU as the
vision-path gate and leaving the remaining support details unchanged.

In `@tensorrt_llm/_torch/attention_backend/inkling_triton.py`:
- Around line 337-339: Update the affected attention paths, including
inkling_decode_attention, to avoid unpacking the unused total_tokens value and
validate that num_heads is evenly divisible by num_kv_heads before computing
kv_group_num. Raise the established appropriate error for invalid head
configurations, while preserving the existing mapping for divisible inputs.

In `@tensorrt_llm/_torch/configs/inkling.py`:
- Around line 121-124: Update the initializer around `moe_intermediate_size` to
read the explicitly supplied `moe_intermediate_size` kwarg before assigning its
default, following the existing `max_position_embeddings` handling; preserve the
checkpoint value when present and fall back to `intermediate_size` only when it
is absent.

In `@tensorrt_llm/_torch/model_config.py`:
- Around line 387-400: Update the MIXED_PRECISION branch’s kv_cache_quant_algo
re-derivation to call the existing _algo_or_none helper instead of constructing
QuantAlgo directly. Preserve the helper’s normalization of None, "none", "null",
and empty strings while retaining normal enum conversion for valid algorithm
values.

In `@tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py`:
- Around line 260-268: Update the mapper initialization or configuration
handling used by ModelLoader.load and
InklingForConditionalGeneration.load_weights so preprocess_weights and
_map_expert resolve the text sub-model configuration via
_text_sub_model_config(model_config). Ensure vocab_size and related text-model
fields come from that resolved config, including on the supplied-mapper path.

In `@tensorrt_llm/_torch/models/modeling_inkling_multimodal.py`:
- Around line 1046-1073: In __init__, after resolving image_token_id and
audio_token_id via _resolve_image_token_id and _resolve_audio_token_id, validate
that they are distinct and reject any collision with a clear configuration error
before placeholder processing can run. Keep _check_placeholder_count and the
multimodal expansion loop unchanged.
- Around line 396-408: Update plan_out_scales and its dependency configuration
so the linear_sum_assignment branch does not rely on an undeclared runtime scipy
import: either declare scipy in the runtime requirements consumed by setup.py,
or replace the assignment logic with an equivalent dependency-free
implementation. Preserve the existing scale-index selection and endpoint
behavior.

In `@tensorrt_llm/_torch/models/modeling_inkling.py`:
- Around line 783-795: Refactor the decode preparation around refresh and
prepare_inkling_attn_decode so seq_lens is computed and its host/device buffer
is allocated and copied once per generation step, then shared across all layers
instead of rebuilt for each layer. Replace the per-row page-table construction
loop and torch.tensor allocations with a single vectorized fill while preserving
padding and invalid-block handling. Benchmark decode-step latency before and
after the change.
- Around line 763-797: Bind the published ready state in refresh to a
per-forward token containing the current cache metadata, including
num_cached_tokens_per_seq, rather than only request IDs and batch size. Update
PyTorchModelEngine and InklingForCausalLM.forward so each forward invalidates or
validates that token before reading seq_lens or page_table, and ensure the eager
fallback cannot consume stale buffers when refresh is skipped.

In `@tensorrt_llm/evaluate/lm_eval.py`:
- Around line 805-812: Update the MMLU command in tensorrt_llm/evaluate/mmlu.py
to declare and propagate --post_process_fn, applying the selected processor to
outputs before compute_score(). Also change the GSM8K option declaration near
LmEvalEvaluator.command_harness to use click.Choice(["strip_thinking_mmmu",
"inkling", "inkling_mmmu"]) instead of type=str.

In `@tensorrt_llm/llmapi/reasoning_parser.py`:
- Around line 787-796: Update the finish() method to reset self._kind alongside
self._buffer before handling remaining content, matching parse() so reused
parser instances start each stream without stale block state.

In `@tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`:
- Around line 710-713: Update the validation description near the vision-path
reference to state that Inkling is validated against the reference using
aggregate accuracy, not item-for-item matching; preserve the existing
TensorRT-LLM/SGLang baseline context.

In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Line 7843: The test_nvfp4 methods in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:7843-7843 and
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py:707-707 both
require explicit return-type annotations. Update each test_nvfp4 definition to
include -> None.
- Around line 7836-7840: Annotate both shared EXTRA_EVALUATOR_KWARGS mappings
with ClassVar[...] in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:7836-7840 and
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py:700-704,
preserving their existing dictionary contents and evaluate() behavior.

In `@tests/unittest/llmapi/test_sampling_params.py`:
- Around line 342-345: Update
test_setup_without_any_eos_warns_and_leaves_end_id_none to match its actual
coverage: either configure capture for tensorrt_llm.logger and assert the
expected warning through caplog, or remove the unused caplog fixture and rename
the test to describe only the end_id-is-None behavior.

---

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/inkling_triton.py`:
- Around line 53-58: Update the comment above the module-level `_NEG` constant
to remove the inaccurate claim that Triton `@jit` functions cannot read
module-level globals or that the value is inlined for that reason. Preserve the
existing explanation of its finite, sufficiently negative value and its role in
avoiding NaNs during masked softmax bookkeeping.

In `@tensorrt_llm/_torch/configs/inkling.py`:
- Around line 137-148: The _local_ids property should use the built-in generic
annotation set[int] and avoid rebuilding the set for every is_local_layer call.
Precompute or cache the converted local_layer_ids set at the configuration
level, then have is_local_layer reuse that cached set while preserving its
current membership behavior.
- Around line 195-248: Add type annotations to InklingConfig.__init__ parameters
text_config, audio_config, vision_config, and mtp_config using dict[str, Any] |
PretrainedConfig | None, and annotate _as_config(value) with the same input type
and an appropriate return type. Add the required Any import, and annotate
__init__ as returning None while preserving the existing configuration handling.

In `@tensorrt_llm/_torch/models/modeling_inkling_multimodal.py`:
- Around line 1021-1028: Update the return annotations of
get_num_tokens_per_image, get_num_tokens_per_audio, and get_num_tokens_per_video
from int to typing.NoReturn, ensuring NoReturn is imported or referenced through
the existing typing conventions.

In `@tensorrt_llm/_torch/models/modeling_inkling.py`:
- Around line 265-274: Update write_state_indices to populate the pinned
state_indices_cpu staging buffer directly from slots, removing the per-forward
torch.tensor allocation. Preserve the existing n-based slicing, non-blocking
device copy into state_indices, and CUDA graph pointer validation.
- Around line 1476-1504: Make `**kwargs` handling consistent between the
`conv_rt is None` and runtime state-pool branches in the decoder layer: forward
the same keyword arguments to `self.attn` in the stateless path as in the pool
path, or remove forwarding from both paths. Preserve all existing attention and
residual behavior.
- Around line 1412-1427: Update tensorrt_llm/_torch/models/modeling_inkling.py
lines 1412-1427 in InklingMoE.forward to remove the second inkling_joint_renorm
call and obtain shared_gammas from the stored InklingMoeRoutingMethod instance.
Update lines 1292-1301 in InklingGate.__init__ and its routing_method property
to construct InklingMoeRoutingMethod once and return that instance, allowing
shared_gammas to persist across the forward call.
- Around line 969-991: Update _build_rel_logits to avoid materializing both the
einsum output and a separate .contiguous() copy; make the einsum result
contiguous directly where supported, and emit bf16 when the consuming Triton
kernels accept it while preserving the required computation and scaling
behavior. Verify peak prefill memory using the configured max_num_tokens rather
than only the example chunk size.

In `@tensorrt_llm/evaluate/lm_eval.py`:
- Around line 259-260: Update the constructor around the visible post_process_fn
and keep_special_tokens parameters to pass post_process_fn through
super().__init__() alongside keep_special_tokens, then remove the later
duplicate post_process_fn assignment. Preserve the existing callback behavior
while keeping all base-constructor initialization together.

In `@tensorrt_llm/llmapi/reasoning_parser.py`:
- Around line 719-746: The duplicated Inkling control-token alphabet and block
state machine must be unified. In tensorrt_llm/llmapi/reasoning_parser.py lines
719-746, extract the shared alphabet and classification logic into a helper,
then have InklingReasoningParser._consume map its blocks to content or reasoning
while preserving tool blocks as visible content. In
tensorrt_llm/evaluate/post_processing.py lines 164-208, remove the local _INK_*
constants and _INK_CONTROL_RE and make extract_inkling_content reuse that helper
so leading unmarked text and tool blocks follow the same classification.
- Around line 757-771: Optimize _partial_control_length by defining a
module-level set containing every strict prefix of the tokens in
_INKLING_CONTROL_TOKENS, initialized alongside that token collection. Replace
the per-length any(...startswith...) scan with membership checks against the
precomputed prefix set while preserving the longest-suffix and zero-result
behavior.

In `@tensorrt_llm/serve/openai_server.py`:
- Around line 1514-1519: Add the public
ReasoningParserFactory.needs_raw_special_tokens() accessor in
reasoning_parser.py, encapsulating the private _parsers lookup and parser-entry
inspection. Update the server logic around reasoning_parser to call this
accessor instead of reading _parsers directly, while preserving the existing
skip_special_tokens behavior.

In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Line 7815: Rename both TestInkling_NVFP4 classes to TestInklingNVFP4 in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:7815-7815 and
tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py:683-683.
Update both direct class selectors to TestInklingNVFP4 in
tests/integration/test_lists/qa/llm_function_core.txt:812-813 and
tests/integration/test_lists/test-db/l0_b200.yml:355-356.

In `@tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py`:
- Around line 173-181: Update test_dmel_preprocess_shape_and_range and
test_dmel_multi_clip_concat_and_empty_clip to derive the expected mel-bin width
and value range from DEFAULT_AUDIO_N_MELS and DEFAULT_AUDIO_NUM_DMEL_BINS
instead of hardcoded 80 and 16. Keep the existing assertions and behavior
unchanged while referencing the module constants directly.
- Around line 124-167: Add a CPU-focused test near the existing vision input
tests covering _to_pil_rgb: assert http://, https://, and data: inputs raise
ValueError matching “resolve”, unsupported objects raise TypeError, and a
generated image is accepted with RGB mode. This closes the input-coercion
coverage gap while preserving the existing tower tests.
- Around line 290-306: Update test_registration to validate processor
registration through the public lookup API exposed by INPUT_PROCESSOR_REGISTRY
instead of accessing _input_processors_cls_by_model_type. Also avoid asserting
InklingInputProcessor._registered_model_type directly; verify the expected
"inkling_mm_model" association through the public registry interface while
preserving the existing registration expectations.

In `@tests/unittest/_torch/modeling/test_modeling_inkling.py`:
- Around line 52-68: Cache the parsed weight map used by _safetensors_shape so
model.safetensors.index.json is read and parsed only once per checkpoint. Apply
a module-scoped fixture or functools.lru_cache to _index, preserving
_safetensors_shape’s existing lookup behavior.
- Around line 122-141: Add a CPU unit test in the test file for
_split_interleaved_gate_up, using an even-row tensor to verify gate receives
even-indexed rows and up receives odd-indexed rows, and assert that an odd-sized
split dimension raises ValueError. Import the mapper helper and pytest as
needed; no checkpoint setup is required.

In `@tests/unittest/llmapi/test_reasoning_parser.py`:
- Around line 944-968: Add a streaming reuse test alongside the existing Inkling
parser tests that calls parse_delta, then finish, then parse_delta again on the
same parser instance. Verify the second stream’s plain text is returned as
content with empty reasoning_content, confirming finish resets the parser state
such as _kind instead of carrying over the first stream’s reasoning block.

In `@tests/unittest/llmapi/test_sampling_params.py`:
- Around line 329-333: Add tests alongside
test_setup_model_config_eos_list_sets_end_id_and_stop_tokens for the tuple EOS
path and empty EOS sequence path: verify tuple (7, 8) sets end_id to 7 and
stop_token_ids to [8], and verify an empty list leaves end_id as None without
raising an error.

In `@tests/unittest/others/test_lm_eval.py`:
- Around line 856-867: Extend the resolver tests around _resolve_post_process_fn
with a pytest.raises(click.BadParameter) case for an unknown string such as
"not_a_processor", covering the invalid-name branch. Also add a callable
passthrough case using extract_inkling_content and assert it is returned with
keep_special_tokens set to False.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 43b4b5e1-4a1e-4209-adbf-ee6275b1e8ef

📥 Commits

Reviewing files that changed from the base of the PR and between 8e5a564 and ce9fa1f.

📒 Files selected for processing (33)

• docs/source/models/supported-models.md
• tensorrt_llm/_torch/attention_backend/inkling_triton.py
• tensorrt_llm/_torch/configs/__init__.py
• tensorrt_llm/_torch/configs/inkling.py
• tensorrt_llm/_torch/model_config.py
• tensorrt_llm/_torch/models/__init__.py
• tensorrt_llm/_torch/models/checkpoints/__init__.py
• tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py
• tensorrt_llm/_torch/models/modeling_inkling.py
• tensorrt_llm/_torch/models/modeling_inkling_multimodal.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/config_utils.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tensorrt_llm/evaluate/lm_eval.py
• tensorrt_llm/evaluate/post_processing.py
• tensorrt_llm/llmapi/llm_utils.py
• tensorrt_llm/llmapi/reasoning_parser.py
• tensorrt_llm/sampling_params.py
• tensorrt_llm/serve/openai_server.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/references/mmlu.yaml
• tests/integration/defs/accuracy/references/mmmu.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/modeling/test_modeling_inkling.py
• tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py
• tests/unittest/llmapi/test_reasoning_parser.py
• tests/unittest/llmapi/test_sampling_params.py
• tests/unittest/others/test_lm_eval.py

[KleinBlueC]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63406 [ run ] triggered by Bot. Commit: 3d079c0 Link to invocation

[tensorrt-cicd]

PR_Github #63406 [ run ] completed with state SUCCESS. Commit: 3d079c0
/LLM/main/L0_MergeRequest_PR pipeline #51386 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[KleinBlueC]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63448 [ run ] triggered by Bot. Commit: 3d079c0 Link to invocation

[Shixiaowei02]

Approving from a disaggregated-serving perspective.

[tensorrt-cicd]

PR_Github #63448 [ run ] completed with state FAILURE. Commit: 3d079c0
/LLM/main/L0_MergeRequest_PR pipeline #51419 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[KleinBlueC]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #63478 [ run ] triggered by Bot. Commit: 3d079c0 Link to invocation

[tensorrt-cicd]

PR_Github #63478 [ run ] completed with state FAILURE. Commit: 3d079c0
/LLM/main/L0_MergeRequest_PR pipeline #51449 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation