[None][feat] Support Inkling-NVFP4 model

docs/source/models/supported-models.md

@@ -29,6 +29,7 @@ The following is a table of supported models for the PyTorch backend:
 | `GptOssForCausalLM`                  | GPT-OSS                            | `openai/gpt-oss-20b`, `openai/gpt-oss-120b`  |
 | `HunYuanDenseForCausalLM` [^5]       | Hunyuan Dense                      | `tencent/Hunyuan-7B-Instruct`                |
 | `HunYuanMoEForCausalLM` [^5]         | Hunyuan MoE                        | `tencent/Hunyuan-A13B-Instruct`              |
+| `InklingForConditionalGeneration` [^15] | Inkling                         | `thinkingmachines/Inkling-NVFP4`             |
 | `InternLM3ForCausalLM` [^5]          | InternLM3                          | `internlm/internlm3-8b-instruct`             |
 | `KimiK25ForConditionalGeneration`    | Kimi-K2.5                          | `moonshotai/Kimi-K2.5`                       |
 | `LagunaForCausalLM`                  | Laguna-XS                          | `poolside/laguna-XS.2`                       |
@@ -75,6 +76,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Qwen3_5MoeForCausalLM`          | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Untested         | Yes            | N/A                      | Untested              | Untested        |
 | `Llama4ForConditionalGeneration` | Yes               | Yes        | Yes                        | Yes                   | Yes             | No  | Yes              | Yes               | No     | Yes           | Yes              | Untested       | N/A                      | Yes                   | Yes             |
 | `GptOssForCausalLM`              | Yes               | Yes        | Yes                        | Yes                   | Yes             | No  | Yes              | No                | Yes    | Yes           | Yes              | Yes            | N/A                      | Yes                   | Yes             |
+| `InklingForConditionalGeneration` [^15] | Yes       | Yes        | Yes                        | No                    | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Yes                   | No              |
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | Yes | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
@@ -95,6 +97,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
 [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`).
+[^15]: Inkling is an NVFP4 checkpoint (`thinkingmachines/Inkling-NVFP4`) with a RoPE-free hybrid-attention text decoder and BF16 vision and audio towers; video is handled as multi-frame images. Use `--reasoning_parser inkling` for its typed-content thinking blocks. Enforced at load: `moe_expert_parallel_size` must divide 256, and `attn_backend` must stay at the default `INKLING`. KV cache reuse, MTP, LoRA, function calling, constrained/guided decoding, EPD disaggregated serving, and multimodal-hash prefix caching are not supported in this release.
 
 # Multimodal Feature Support Matrix (PyTorch Backend)
 
@@ -105,6 +108,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Gemma4ForConditionalGeneration`     | Untested          | Yes        | Yes             | Yes           | Untested         | No             | Untested              | No                        | L + I + V + A [^9] |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested        | Yes           | Untested         | No             | Untested              | No                        | L + I + A |
 | `HCXVisionForCausalLM`               | Yes               | Yes        | No              | Yes           | Yes              | Yes            | Yes                   | No                        | L + I     |
+| `InklingForConditionalGeneration`    | Yes               | Yes        | Untested        | Yes           | Untested         | No             | Yes                   | No                        | L + I + V + A [^15] |
 | `LlavaLlamaModel (VILA)`             | Yes               | Yes        | No              | Yes           | Yes              | No             | Yes                   | No                        | L + I + V |
 | `LlavaNextForConditionalGeneration`  | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | Yes                       | L + I     |
 | `Llama4ForConditionalGeneration`     | Yes               | Yes        | No              | Yes           | Yes              | No             | Yes                   | No                        | L + I     |

tensorrt_llm/_torch/attention_backend/inkling/__init__.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inkling attention: Triton kernels, per-step metadata, backend, cache manager.
+
+Split the way ``sparse/minimax_m3`` is -- kernels, metadata, backend and cache
+manager in their own modules -- rather than one flat file. Inkling is NOT under
+``sparse/``: that package is gated on ``sparse_attention_config`` /
+``SparseParams`` and its machinery (index caches, top-k block masks, per-sparse
+-layer pools) assumes only part of the KV is scored. Inkling's attention is
+dense -- full causal on global layers, a 512-token sliding window on local ones
+-- with a learned relative-bias ``score_mod``.
+"""
+
+from .backend import InklingTritonAttention
+from .cache_manager import InklingHybridCacheManager
+from .kernels import (
+    build_page_table,
+    inkling_decode_attention,
+    inkling_prefill_attention,
+    write_kv_cache_hnd,
+)
+from .metadata import InklingAttentionMetadata
+
+__all__ = [
+    "InklingAttentionMetadata",
+    "InklingHybridCacheManager",
+    "InklingTritonAttention",
+    "build_page_table",
+    "inkling_decode_attention",
+    "inkling_prefill_attention",
+    "write_kv_cache_hnd",
+]

tensorrt_llm/_torch/attention_backend/inkling/backend.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Attention backend that carries :class:`InklingAttentionMetadata`."""
+
+from ..trtllm import TrtllmAttention
+from .metadata import InklingAttentionMetadata
+
+
+class InklingTritonAttention(TrtllmAttention):
+    """Carries :class:`InklingAttentionMetadata`.
+
+    Inkling never routes attention through a backend ``forward``:
+    ``InklingAttention.forward`` overrides the base module entirely and calls
+    the Triton kernels above. The backend object exists so the model engine
+    picks the right ``Metadata`` class (``metadata_cls = attn_backend.Metadata``)
+    and so the base module can assign ``local_layer_idx``. Subclassing
+    ``TrtllmAttention`` rather than ``AttentionBackend`` keeps construction and
+    every non-Inkling code path byte-identical to the TRTLLM backend Inkling
+    used before.
+    """
+
+    Metadata = InklingAttentionMetadata

tensorrt_llm/_torch/attention_backend/inkling/cache_manager.py

@@ -0,0 +1,129 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Inkling's KV cache manager: paged KV plus the short-conv state pool.
+
+Lives with the model's attention package rather than under ``pyexecutor``,
+matching ``sparse/minimax_m3/cache_manager.py``. Nothing about it is installed
+into shared framework directories: ``_util`` selects this class the same way it
+selects ``MiniMaxM3KVCacheManagerV2``, and ``InklingAttentionMetadata``
+type-tests it directly.
+
+There is deliberately no ``BaseConvStateManager`` protocol. Per-request
+short-conv state is NOT new -- ``BaseMambaCacheManager`` already declares
+``get_conv_states(layer_idx)`` and nemotron_h / qwen3_next / qwen3_5 implement
+it -- but that protocol also mandates ``get_ssm_states``, ``is_speculative``,
+``mamba_layer_cache`` and replay metadata, none of which Inkling can back, and
+its one-tensor-per-layer accessor cannot express Inkling's four convs per layer
+at two different widths (k/v follow the TP-sharded kv split; the post-attention
+and post-MLP convs run replicated on the full residual stream).
+
+A parallel protocol was tried and removed: both of its useful methods returned
+Inkling's own pool and runtime types, so it abstracted nothing while putting an
+Inkling-specific file under ``pyexecutor``. If a second short-conv model ever
+appears, widen the framework's existing hook rather than inventing another one
+beside it.
+"""
+
+import torch
+
+from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
+
+
+class InklingHybridCacheManager(KVCacheManagerV2):
+    """Paged KV (V2, per-layer geometry) + the short-conv state pool.
+
+    The pool used to be a separate ResourceManager registered under its own
+    ResourceManagerType, published from three call sites inside
+    ``PyTorchModelEngine``. Making it part of the cache manager -- the shape
+    ``CppMambaHybridCacheManager`` uses for mamba conv/SSM state -- removes all
+    of that: the pool reaches the model through
+    ``attn_metadata.kv_cache_manager``, which is a standard AttentionMetadata
+    field, and rows are released by the manager's own ``free_resources``, which
+    every caller (including the warmup/estimation dummy-batch path) already
+    invokes.
+
+    It also removes a whole class of bug rather than just some code. A pool that
+    lives beside the cache manager can disagree with it about block reuse or
+    request lifetime and nothing forces the two views back together; here the
+    conv rows are freed by the same call that frees the request's KV blocks, so
+    they cannot drift apart.
+
+    Note: because the manager owns the pool, the pool is now also allocated for
+    the throwaway manager built during KV-cache size estimation. That is roughly
+    66 layers x 4 convs x (max_batch+1) rows -- tens of MB, freed with the
+    estimation manager -- and it buys the lifetime coupling above.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Imported here, not at module scope: modeling_inkling imports from
+        # _torch.attention_backend and _torch.modules, and a top-level import
+        # would close a cycle back through pyexecutor at model-load time.
+        from ...models.modeling_inkling import InklingConvStateCache
+
+        pretrained_config = kwargs["pretrained_config"]
+        mapping = kwargs["mapping"]
+        max_batch_size = kwargs["max_batch_size"]
+        # NOT kwargs["dtype"]: that is the KV cache dtype, a C++
+        # ``tensorrt_llm.bindings.DataType``, and torch.zeros rejects it. The
+        # conv pool holds pre-conv activations, so it takes the model's compute
+        # dtype from the (text) config.
+        text_config = getattr(pretrained_config, "text_config", pretrained_config)
+        conv_dtype = getattr(text_config, "torch_dtype", None)
+        if not isinstance(conv_dtype, torch.dtype):
+            conv_dtype = torch.bfloat16
+        # The conv pool's k/v width follows the attention kv-head split, so it
+        # takes the ATTENTION TP, not the global one. Under attention DP every
+        # rank keeps the full kv-head set for its own requests -- the same rule
+        # KVCacheManagerV2 already applies to the paged pool
+        # (``tp_size = 1 if mapping.enable_attention_dp``) and that
+        # InklingAttention applies to the k/v short convs themselves. Dividing
+        # by the global tp_size here would allocate quarter-width conv rows for
+        # full-width convs.
+        attn_tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size
+        # +1 row for the CUDA-graph padding / dummy-request slot (the mamba
+        # pattern): a padded decode batch admits up to max_batch_size real
+        # requests plus a shared dummy row.
+        self._conv_cache = InklingConvStateCache(
+            pretrained_config,
+            attn_tp_size,
+            max_batch_size + 1,
+            torch.device("cuda", torch.cuda.current_device()),
+            conv_dtype,
+        )
+
+    # ---- model-facing -----------------------------------------------------
+    def prepare_conv_runtime(self, attn_metadata):
+        from ...models.modeling_inkling import InklingConvRuntime
+
+        return self._conv_cache, InklingConvRuntime.build(attn_metadata, self._conv_cache)
+
+    def free_conv_state(self, request_ids) -> None:
+        self._conv_cache.free(list(request_ids))
+
+    # ---- KVCacheManagerV2 -----------------------------------------------------
+    def free_resources(self, request, *args, **kwargs):
+        """Release the conv row with the request's KV blocks.
+
+        This is what lets the model engine's warmup/estimation dummy-batch
+        cleanup drop its Inkling-specific branch: it already calls
+        ``kv_cache_manager.free_resources(req)`` for every dummy request, and a
+        leaked conv row would later be reused, with stale state, by a real
+        request whose id collides with a dummy id.
+        """
+        rid = getattr(request, "py_request_id", None)
+        if rid is not None:
+            self.free_conv_state([rid])
+        return super().free_resources(request, *args, **kwargs)