diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 21ee5daf3433..48482c3f28e3 100644 --- a/docs/source/models/supported-models.md +++ b/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 | diff --git a/tensorrt_llm/_torch/attention_backend/inkling/__init__.py b/tensorrt_llm/_torch/attention_backend/inkling/__init__.py new file mode 100644 index 000000000000..b62b06b28aa9 --- /dev/null +++ b/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", +] diff --git a/tensorrt_llm/_torch/attention_backend/inkling/backend.py b/tensorrt_llm/_torch/attention_backend/inkling/backend.py new file mode 100644 index 000000000000..9eec17b25bd1 --- /dev/null +++ b/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 diff --git a/tensorrt_llm/_torch/attention_backend/inkling/cache_manager.py b/tensorrt_llm/_torch/attention_backend/inkling/cache_manager.py new file mode 100644 index 000000000000..73b4af94eacb --- /dev/null +++ b/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) diff --git a/tensorrt_llm/_torch/attention_backend/inkling/kernels.py b/tensorrt_llm/_torch/attention_backend/inkling/kernels.py new file mode 100644 index 000000000000..8d93c6d0370a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/inkling/kernels.py @@ -0,0 +1,524 @@ +# 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 Triton attention: paged prefill + decode with a learned relative-bias +``score_mod`` and native sliding window. + +Why this exists +--------------- +Inkling attention adds a learned per-(query-token, head, relative-distance) +additive bias INSIDE the attention score, and windows local layers separately. +No fused/CUDA-graph-safe TensorRT-LLM backend exposes a ``score_mod`` hook: +``attentionOp.cpp`` disables context FMHA for ``kRELATIVE`` position embedding, +and the trtllm-gen decode kernel rejects a relative bias. So the production +attention path for Inkling is a pair of Triton kernels that apply the bias as an +aux-tensor ``score_mod``. + +The bias is precomputed on the torch side as a contiguous ``rel_logits`` aux +tensor ``[num_query_tokens, num_heads, rel_extent]`` (``einsum('thd,de->the', r, +proj)`` with the global-layer ``tau`` folded in). The kernels only gather+add: + + rel_dist = q_pos - k_pos + rel_idx = clamp(rel_dist, 0, rel_extent - 1) + bias = rel_logits[q_idx, head, rel_idx] if 0 <= rel_dist < rel_extent else 0 + qk += bias + +This keeps ``rel_logits`` a *static-shape* tensor (``num_query_tokens`` == batch +in the decode phase), so the decode kernel is CUDA-graph capturable: the launch +grid ``(batch, num_heads)`` is fixed, and per-request sequence lengths are read +from a GPU tensor inside the kernel (no host sync, no ``.item()``). + +Both kernels read the paged KV cache in the ``KVCacheManagerV2`` HND layout +(``[num_pages, num_kv_heads, page_size, head_dim]`` after selecting K or V from +the ``[num_pages, 2, ...]`` pool), addressed through a per-request page table. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +# Additive value used to drop a masked key from the softmax. Large enough that +# ``exp(qk - max)`` underflows to 0 in fp32, finite so online-softmax bookkeeping +# never sees a NaN. (float("-inf") would poison the running max on the first, +# fully-masked tile of a windowed row.) Inlined as a literal inside the kernels +# because Triton @jit functions cannot read non-constexpr module globals. +_NEG = tl.constexpr(-1.0e30) + + +# --------------------------------------------------------------------------- +# Prefill (context) kernel: contiguous varlen Q/K/V, causal + optional window, +# optional relative-bias score_mod. One fresh context has no cached prefix, so +# K/V are read from the packed extend tensors directly. +# --------------------------------------------------------------------------- +@triton.jit +def _inkling_prefill_kernel( + Q, + K, + V, + Out, + RelLogits, + cu_seqlens, + sm_scale, + stride_qt, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_ot, + stride_oh, + stride_rt, + stride_rh, + kv_group_num, + rel_extent: tl.constexpr, + HAS_REL: tl.constexpr, + WINDOW_LEFT: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + Lk: tl.constexpr, +): + cur_seq = tl.program_id(0) + cur_head = tl.program_id(1) + cur_block_m = tl.program_id(2) + cur_kv_head = cur_head // kv_group_num + + seq_start = tl.load(cu_seqlens + cur_seq) + seq_len = tl.load(cu_seqlens + cur_seq + 1) - seq_start + + offs_m = tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + mask_d = offs_d < Lk + + q_pos = cur_block_m * BLOCK_M + offs_m # [BLOCK_M], position within sequence + mask_m = q_pos < seq_len + + q_ptrs = (seq_start + q_pos)[:, None] * stride_qt + cur_head * stride_qh + offs_d[None, :] + q = tl.load(Q + q_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0) + + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + e_sum = tl.zeros([BLOCK_M], dtype=tl.float32) + + # Causal: query block cur_block_m only attends to keys <= its last row. + end_n = tl.minimum(seq_len, (cur_block_m + 1) * BLOCK_M) + # Sliding window: skip whole key tiles older than the window low bound. + if WINDOW_LEFT >= 0: + lo = cur_block_m * BLOCK_M - WINDOW_LEFT + if lo < 0: + lo = 0 + lo = (lo // BLOCK_N) * BLOCK_N + else: + lo = 0 + + for start_n in range(lo, end_n, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + k_pos = start_n + offs_n # [BLOCK_N] + mask_n = k_pos < seq_len + + k_ptrs = ( + (seq_start + k_pos)[None, :] * stride_kt + cur_kv_head * stride_kh + offs_d[:, None] + ) + k = tl.load(K + k_ptrs, mask=mask_n[None, :] & mask_d[:, None], other=0.0) + qk = tl.dot(q, k, out_dtype=tl.float32) * sm_scale # [BLOCK_M, BLOCK_N] + + if HAS_REL: + rel_dist = q_pos[:, None] - k_pos[None, :] + rel_idx = tl.minimum(tl.maximum(rel_dist, 0), rel_extent - 1) + rel_ptrs = (seq_start + q_pos)[:, None] * stride_rt + cur_head * stride_rh + rel_idx + rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent) + bias = tl.load( + RelLogits + rel_ptrs, mask=mask_m[:, None] & mask_n[None, :] & rel_valid, other=0.0 + ) + qk += bias + + valid = mask_m[:, None] & mask_n[None, :] & (q_pos[:, None] >= k_pos[None, :]) + if WINDOW_LEFT >= 0: + valid &= (q_pos[:, None] - k_pos[None, :]) <= WINDOW_LEFT + qk = tl.where(valid, qk, _NEG) + + row_max = tl.max(qk, 1) + n_e_max = tl.maximum(e_max, row_max) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max[:, None]) + e_sum = e_sum * re_scale + tl.sum(p, 1) + + v_ptrs = ( + (seq_start + k_pos)[:, None] * stride_vt + cur_kv_head * stride_vh + offs_d[None, :] + ) + v = tl.load(V + v_ptrs, mask=mask_n[:, None] & mask_d[None, :], other=0.0) + acc = acc * re_scale[:, None] + tl.dot(p.to(v.dtype), v, out_dtype=tl.float32) + e_max = n_e_max + + acc = acc / e_sum[:, None] + o_ptrs = (seq_start + q_pos)[:, None] * stride_ot + cur_head * stride_oh + offs_d[None, :] + tl.store(Out + o_ptrs, acc.to(Out.dtype.element_ty), mask=mask_m[:, None] & mask_d[None, :]) + + +# --------------------------------------------------------------------------- +# Decode (generation) kernel: one query token per request, paged KV read, +# causal + optional window, optional relative-bias score_mod. CUDA-graph safe: +# static grid (batch, num_heads); seq lengths and the page table are read from +# GPU tensors, no host sync. +# --------------------------------------------------------------------------- +@triton.jit +def _inkling_decode_kernel( + Q, + K_Cache, + V_Cache, + Out, + RelLogits, + seq_lens, + page_table, + sm_scale, + stride_qb, + stride_qh, + stride_kp, + stride_kh, + stride_kt, + stride_vp, + stride_vh, + stride_vt, + stride_ob, + stride_oh, + stride_rb, + stride_rh, + stride_ptb, + kv_group_num, + page_size: tl.constexpr, + rel_extent: tl.constexpr, + HAS_REL: tl.constexpr, + WINDOW_LEFT: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + BLOCK_N: tl.constexpr, + Lk: tl.constexpr, +): + cur_batch = tl.program_id(0).to(tl.int64) + cur_head = tl.program_id(1) + cur_kv_head = cur_head // kv_group_num + + seq_len = tl.load(seq_lens + cur_batch) + q_pos = seq_len - 1 # decode query sits at the last cached position + + offs_d = tl.arange(0, BLOCK_DMODEL) + offs_n = tl.arange(0, BLOCK_N) + mask_d = offs_d < Lk + + q = tl.load( + Q + cur_batch * stride_qb + cur_head * stride_qh + offs_d, mask=mask_d, other=0.0 + ).to(tl.float32) # [BLOCK_DMODEL] + + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + e_max = -float("inf") + e_sum = 0.0 + + if WINDOW_LEFT >= 0: + lo = q_pos - WINDOW_LEFT + if lo < 0: + lo = 0 + lo = (lo // BLOCK_N) * BLOCK_N + else: + lo = 0 + + for start_n in range(lo, seq_len, BLOCK_N): + k_pos = start_n + offs_n # [BLOCK_N] + mask_n = k_pos < seq_len + + page_local = k_pos // page_size + tok_in_page = k_pos % page_size + page_id = tl.load( + page_table + cur_batch * stride_ptb + page_local, mask=mask_n, other=0 + ).to(tl.int64) + + k_ptrs = ( + page_id[:, None] * stride_kp + + cur_kv_head * stride_kh + + tok_in_page[:, None] * stride_kt + + offs_d[None, :] + ) + k = tl.load(K_Cache + k_ptrs, mask=mask_n[:, None] & mask_d[None, :], other=0.0).to( + tl.float32 + ) + qk = tl.sum(q[None, :] * k, 1) * sm_scale # [BLOCK_N] + + if HAS_REL: + rel_dist = q_pos - k_pos + rel_idx = tl.minimum(tl.maximum(rel_dist, 0), rel_extent - 1) + rel_ptrs = cur_batch * stride_rb + cur_head * stride_rh + rel_idx + rel_valid = (rel_dist >= 0) & (rel_dist < rel_extent) + bias = tl.load(RelLogits + rel_ptrs, mask=mask_n & rel_valid, other=0.0) + qk += bias + + valid = mask_n & (k_pos <= q_pos) + if WINDOW_LEFT >= 0: + valid &= (q_pos - k_pos) <= WINDOW_LEFT + qk = tl.where(valid, qk, _NEG) + + n_e_max = tl.maximum(e_max, tl.max(qk, 0)) + re_scale = tl.exp(e_max - n_e_max) + p = tl.exp(qk - n_e_max) # [BLOCK_N] + e_sum = e_sum * re_scale + tl.sum(p, 0) + + v_ptrs = ( + page_id[:, None] * stride_vp + + cur_kv_head * stride_vh + + tok_in_page[:, None] * stride_vt + + offs_d[None, :] + ) + v = tl.load(V_Cache + v_ptrs, mask=mask_n[:, None] & mask_d[None, :], other=0.0).to( + tl.float32 + ) + acc = acc * re_scale + tl.sum(p[:, None] * v, 0) + e_max = n_e_max + + o = acc / e_sum + tl.store( + Out + cur_batch * stride_ob + cur_head * stride_oh + offs_d, + o.to(Out.dtype.element_ty), + mask=mask_d, + ) + + +# --------------------------------------------------------------------------- +# Python wrappers +# --------------------------------------------------------------------------- +def _block_dmodel(head_dim: int) -> int: + return triton.next_power_of_2(head_dim) + + +def inkling_prefill_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + sm_scale: float, + rel_logits: Optional[torch.Tensor] = None, + rel_extent: int = 0, + window_left: int = -1, +) -> torch.Tensor: + """Context-phase attention over packed varlen Q/K/V. + + Args: + q: ``[total_tokens, num_heads, head_dim]`` + k, v: ``[total_tokens, num_kv_heads, head_dim]`` + cu_seqlens: ``[batch + 1]`` int32 cumulative token counts. + max_seqlen: max per-request length (host int; used for the grid). + sm_scale: softmax scale (``1 / head_dim`` for Inkling). + rel_logits: ``[total_tokens, num_heads, rel_extent]`` fp32 aux bias, or + None to skip the score_mod. + rel_extent: relative-bias extent (profile width). + window_left: sliding-window radius (inclusive), -1 to disable. + + Returns ``[total_tokens, num_heads, head_dim]`` in q's dtype. + """ + # The kernels index the head_dim with an implicit stride-1 last axis, so the + # inputs must be contiguous. ``v`` in particular reaches here non-contiguous: + # it is the fused-qkv v slice run through the short conv, and (unlike ``k``) + # never passes through ``apply_qk_norm``'s reshape, so it keeps the qkv row + # stride. ``.contiguous()`` is a no-op for already-contiguous q/k. + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + _total_tokens, num_heads, head_dim = q.shape + num_kv_heads = k.shape[1] + # The kernel maps a query head to its KV head as ``cur_head // kv_group_num``; + # a non-divisible pair would silently mis-map instead of failing. + assert num_heads % num_kv_heads == 0, (num_heads, num_kv_heads) + kv_group_num = num_heads // num_kv_heads + o = torch.empty_like(q) + + has_rel = rel_logits is not None + if has_rel: + assert rel_logits.is_contiguous() and rel_logits.shape[-1] == rel_extent + r_st, r_sh = rel_logits.stride(0), rel_logits.stride(1) + rel_arg = rel_logits + else: + r_st = r_sh = 0 + rel_arg = q # unused placeholder pointer + + BLOCK_DMODEL = _block_dmodel(head_dim) + BLOCK_M = 64 + BLOCK_N = 64 + batch = cu_seqlens.shape[0] - 1 + grid = (batch, num_heads, triton.cdiv(max_seqlen, BLOCK_M)) + + _inkling_prefill_kernel[grid]( + q, + k, + v, + o, + rel_arg, + cu_seqlens, + sm_scale, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + r_st, + r_sh, + kv_group_num, + rel_extent=rel_extent if has_rel else 1, + HAS_REL=has_rel, + WINDOW_LEFT=window_left, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + Lk=head_dim, + num_warps=4, + num_stages=2, + ) + return o + + +def inkling_decode_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + sm_scale: float, + rel_logits: Optional[torch.Tensor] = None, + rel_extent: int = 0, + window_left: int = -1, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Generation-phase attention: one query per request over paged KV. + + Args: + q: ``[batch, num_heads, head_dim]`` + k_cache, v_cache: ``[num_pages, num_kv_heads, page_size, head_dim]`` HND + views (K/V selected from the ``[num_pages, 2, ...]`` pool). + seq_lens: ``[batch]`` int32 GPU total-KV length per request. + page_table: ``[batch, max_pages]`` int32 GPU physical page ids. + page_size: tokens per page. + sm_scale: softmax scale (``1 / head_dim``). + rel_logits: ``[batch, num_heads, rel_extent]`` fp32 aux bias, or None. + rel_extent: relative-bias extent. + window_left: sliding-window radius (inclusive), -1 to disable. + out: optional pre-allocated ``[batch, num_heads, head_dim]`` output (for + CUDA-graph static buffers). + + Returns ``[batch, num_heads, head_dim]`` in q's dtype. + """ + q = q.contiguous() # kernel indexes head_dim as the stride-1 axis + batch, num_heads, head_dim = q.shape + num_kv_heads = k_cache.shape[1] + assert num_heads % num_kv_heads == 0, (num_heads, num_kv_heads) + kv_group_num = num_heads // num_kv_heads + o = out if out is not None else torch.empty_like(q) + + has_rel = rel_logits is not None + if has_rel: + assert rel_logits.is_contiguous() and rel_logits.shape[-1] == rel_extent + r_sb, r_sh = rel_logits.stride(0), rel_logits.stride(1) + rel_arg = rel_logits + else: + r_sb = r_sh = 0 + rel_arg = q + + BLOCK_DMODEL = _block_dmodel(head_dim) + BLOCK_N = 64 + grid = (batch, num_heads) + + _inkling_decode_kernel[grid]( + q, + k_cache, + v_cache, + o, + rel_arg, + seq_lens, + page_table, + sm_scale, + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + o.stride(0), + o.stride(1), + r_sb, + r_sh, + page_table.stride(0), + kv_group_num, + page_size=page_size, + rel_extent=rel_extent if has_rel else 1, + HAS_REL=has_rel, + WINDOW_LEFT=window_left, + BLOCK_DMODEL=BLOCK_DMODEL, + BLOCK_N=BLOCK_N, + Lk=head_dim, + num_warps=4, + num_stages=2, + ) + return o + + +def build_page_table(block_ids_per_seq, max_pages: int, device) -> torch.Tensor: + """Pack a ragged ``block_ids_per_seq`` (from + ``KVCacheManagerV2.get_batch_cache_indices``) into a dense + ``[batch, max_pages]`` int32 page table, padding short rows with 0 (never + read: the decode kernel bounds every access by the per-request ``seq_len``). + """ + batch = len(block_ids_per_seq) + pt = torch.zeros((batch, max_pages), dtype=torch.int32, device=device) + for i, blocks in enumerate(block_ids_per_seq): + valid = [int(b) for b in blocks if int(b) >= 0] + if valid: + pt[i, : len(valid)] = torch.tensor(valid, dtype=torch.int32, device=device) + return pt + + +def write_kv_cache_hnd( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + new_k: torch.Tensor, + new_v: torch.Tensor, + block_ids, + start_slot: int, + page_size: int, +) -> None: + """Write ``new_k``/``new_v`` (``[n, num_kv_heads, head_dim]``) for ONE + request into the paged HND cache starting at logical position ``start_slot``. + + ``k_cache``/``v_cache`` are ``[num_pages, num_kv_heads, page_size, + head_dim]`` views. ``block_ids`` is the request's physical page list. Used at + prefill/decode to populate the cache before attention reads it. + """ + valid_blocks = [int(b) for b in block_ids if int(b) >= 0] + n = new_k.shape[0] + written = 0 + while written < n: + pos = start_slot + written + page = valid_blocks[pos // page_size] + off = pos % page_size + take = min(page_size - off, n - written) + k_cache[page, :, off : off + take, :] = ( + new_k[written : written + take].transpose(0, 1).to(k_cache.dtype) + ) + v_cache[page, :, off : off + take, :] = ( + new_v[written : written + take].transpose(0, 1).to(v_cache.dtype) + ) + written += take diff --git a/tensorrt_llm/_torch/attention_backend/inkling/metadata.py b/tensorrt_llm/_torch/attention_backend/inkling/metadata.py new file mode 100644 index 000000000000..8fa14eb96ffe --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/inkling/metadata.py @@ -0,0 +1,228 @@ +# 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. +"""Per-step decode metadata for the Inkling Triton attention path.""" + +from typing import Dict, List, Optional + +import torch + +from ...._utils import prefer_pinned +from ..trtllm import TrtllmAttentionMetadata + + +class InklingAttentionMetadata(TrtllmAttentionMetadata): + """Per-step decode metadata for the Inkling Triton kernels. + + The decode kernel needs, per generation request, the total KV length + (``num_cached + 1``) and the physical page table. Building those from host + lists inside ``model.forward`` raises ``Cannot copy between CPU and CUDA + tensors during CUDA graph capture``, so they live in fixed-pointer GPU + buffers that :meth:`prepare` overwrites each step -- the same shape as the + base metadata's ``seq_lens_cuda`` in-place ``copy_``, and the same reason + :meth:`AttentionMetadata.create_cuda_graph_metadata` exists. + + ``prepare()`` is the framework's documented "before the forward step of the + model" hook and runs on all three input-preparation paths including the + steady-generation fast path, which is exactly the set of places this used to + be published from inside ``PyTorchModelEngine``. It runs after the padded + batch is assembled (``CUDAGraphRunner.pad_batch`` wraps ``_prepare_inputs``) + and after ``super().prepare()`` has re-clamped + ``kv_cache_params.num_cached_tokens_per_seq``, so it sees exactly the data + the model-engine hook saw. + + Freshness needs no epoch counter: the buffers belong to the metadata object + for the step being prepared, and ``ink_num_gen`` is reset at the top of + every ``prepare()``, so a step that publishes nothing cannot leave the + previous step's page table readable. + + ``TrtllmAttentionMetadata`` is the base because that is the backend Inkling + ran on before this class existed; Inkling reads only base fields, but + inheriting keeps every other consumer's expectations intact. + """ + + def __post_init__(self) -> None: + super().__post_init__() + self.kv_layout = "HND" + # Number of generation rows published this step. 0 means the decode + # buffers hold nothing valid for this forward. + self.ink_num_gen: int = 0 + self.ink_max_pages: Optional[int] = None + self.ink_cap: int = 0 + self.ink_seq_lens: Optional[torch.Tensor] = None # [cap] int32 total-KV + self.ink_page_table: Dict[int, torch.Tensor] = {} # layer -> [cap, pages] + self._ink_sl_host: Optional[torch.Tensor] = None + self._ink_pt_host: Optional[torch.Tensor] = None + # Short-conv pool + this forward's context/generation split, published + # alongside the attention metadata because both are per-step host work + # that must land in stable buffers before CUDA-graph capture. + self.ink_conv_cache = None + self.ink_conv_rt = None + + def _ink_layers(self) -> List[int]: + """Global decoder-layer indices this rank owns. + + ``get_batch_cache_indices`` is per-layer (per pool_id / index_scale), so + the page table is too. ``pp_layers`` is already the local slice, and the + model addresses the cache by global layer index. + """ + return list(getattr(self.kv_cache_manager, "pp_layers", [])) + + def _ink_ensure(self, num_gen: int) -> None: + """Size the stable buffers, refusing to grow them under CUDA graph.""" + mgr = self.kv_cache_manager + if self.ink_max_pages is None: + self.ink_max_pages = max(1, int(mgr.max_blocks_per_seq)) + if self.ink_seq_lens is not None and num_gen <= self.ink_cap: + return + if self.is_cuda_graph and self.ink_seq_lens is not None: + raise RuntimeError( + f"InklingAttentionMetadata would grow its stable decode buffers " + f"during CUDA graph capture/replay (num_gen={num_gen} > " + f"cap={self.ink_cap}); the buffers are sized to the padded " + f"scheduler batch, so this signals a capture-shape mismatch" + ) + self.ink_cap = max(num_gen, self.ink_cap) + device = self.seq_lens_cuda.device + self.ink_seq_lens = torch.ones(self.ink_cap, dtype=torch.int32, device=device) + self.ink_page_table = { + layer: torch.zeros((self.ink_cap, self.ink_max_pages), dtype=torch.int32, device=device) + for layer in self._ink_layers() + } + + def prepare(self) -> None: + super().prepare() + self._prepare_inkling_conv() + self._prepare_inkling_decode() + + def _prepare_inkling_conv(self) -> None: + """Publish the short-conv pool rows for this batch. + + Runs here rather than from PyTorchModelEngine because prepare() is the + framework's pre-forward hook and is already called on every input-prep + path, so the host->device slot write stays outside the captured region. + + Type-tests the concrete manager rather than an abstract protocol. A + protocol would have to be satisfied by exactly one class -- both its + useful methods return Inkling's own pool and runtime types -- so it + would abstract nothing while putting an Inkling file under + ``pyexecutor``. The framework rule this looks like it breaks ("test the + capability, not the model", the way _prepare_mamba_metadata tests + BaseMambaCacheManager) is about not branching on model identity in + SHARED code; this class only ever serves Inkling by construction. + + If a second short-conv model ever appears, the right move is to widen + the framework's existing hook -- ``BaseMambaCacheManager`` already + declares ``get_conv_states(layer_idx)`` and three models implement it -- + not to invent a parallel protocol beside it. + """ + from .cache_manager import InklingHybridCacheManager + + mgr = self.kv_cache_manager + if not isinstance(mgr, InklingHybridCacheManager) or self.request_ids is None: + self.ink_conv_cache = self.ink_conv_rt = None + return + self.ink_conv_cache, self.ink_conv_rt = mgr.prepare_conv_runtime(self) + + def _prepare_inkling_decode(self) -> None: + # Reset first: a step that returns early below must not leave the + # previous step's buffers advertised as current. 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. + self.ink_num_gen = 0 + mgr = self.kv_cache_manager + if mgr is None or self.request_ids is None or self.kv_cache_params is None: + return + num_contexts = self.num_contexts + num_gen = len(self.request_ids) - num_contexts + if num_gen <= 0: + return + layers = self._ink_layers() + if not layers: + return + self._ink_ensure(num_gen) + + # Total-KV lengths are layer-independent: staged once per step. + num_cached = self.kv_cache_params.num_cached_tokens_per_seq[num_contexts:] + if self._ink_sl_host is None or self._ink_sl_host.shape[0] < num_gen: + self._ink_sl_host = torch.empty(num_gen, dtype=torch.int32, pin_memory=prefer_pinned()) + sl_host = self._ink_sl_host[:num_gen] + sl_np = sl_host.numpy() + for i in range(num_gen): + sl_np[i] = int(num_cached[i]) + 1 + self.ink_seq_lens[:num_gen].copy_(sl_host, non_blocking=True) + + # Reused pinned staging, ONE ROW PER LAYER. It must not be a single + # [cap, max_pages] buffer refilled per layer: the copies below are + # non_blocking, so the next layer's fill would overwrite the host bytes + # while the previous layer's H2D copy is still in flight, and every + # layer past the first would land a torn page table -- attention then + # reads the wrong KV pages and decode collapses to repeated tokens. + # A fresh pinned allocation per layer is not an option either (66 + # cudaHostAllocs a token), hence one 3-D buffer indexed by layer. + n_layers = len(layers) + if ( + self._ink_pt_host is None + or self._ink_pt_host.shape[0] < n_layers + or self._ink_pt_host.shape[1] < num_gen + ): + self._ink_pt_host = torch.zeros( + (n_layers, max(num_gen, self.ink_cap), self.ink_max_pages), + dtype=torch.int32, + pin_memory=prefer_pinned(), + ) + gen_ids = self.request_ids[num_contexts:] + pt_np = self._ink_pt_host.numpy() + pt_np[:n_layers, :num_gen].fill(0) + for li, layer in enumerate(layers): + block_ids = mgr.get_batch_cache_indices(gen_ids, layer) + # Clamp to num_gen: the staging row is sized to the generation + # slice, so a manager returning more rows than request ids would + # otherwise write past it. + for i, blocks in enumerate(block_ids[:num_gen]): + valid = [b for b in map(int, blocks) if b >= 0][: self.ink_max_pages] + if valid: + pt_np[li, i, : len(valid)] = valid + # Issue the copies only after every row is filled, so no in-flight copy + # can alias a row this loop still has to write. + for li, layer in enumerate(layers): + self.ink_page_table[layer][:num_gen].copy_( + self._ink_pt_host[li, :num_gen], non_blocking=True + ) + self.ink_num_gen = num_gen + + def create_cuda_graph_metadata( + self, max_batch_size: int, *args, **kwargs + ) -> "InklingAttentionMetadata": + md = super().create_cuda_graph_metadata(max_batch_size, *args, **kwargs) + if md is self or md.kv_cache_manager is None: + return md + # Same treatment interface.py gives block_ids_per_seq under + # enable_flash_mla: create_cuda_graph_metadata is a SHALLOW copy, so the + # graph metadata would otherwise share -- and then resize -- the eager + # metadata's buffers, stranding the captured pointers. Allocate at the + # padded batch size up front so _ink_ensure never grows under capture. + md.ink_max_pages = max(1, int(md.kv_cache_manager.max_blocks_per_seq)) + md.ink_cap = max_batch_size + md.ink_num_gen = 0 + md._ink_sl_host = None + md._ink_pt_host = None + md.ink_conv_cache = None + md.ink_conv_rt = None + md.ink_seq_lens = torch.ones(max_batch_size, dtype=torch.int32, device="cuda") + md.ink_page_table = { + layer: torch.zeros((max_batch_size, md.ink_max_pages), dtype=torch.int32, device="cuda") + for layer in md._ink_layers() + } + return md diff --git a/tensorrt_llm/_torch/attention_backend/utils.py b/tensorrt_llm/_torch/attention_backend/utils.py index ef83c99159ed..d243dde6dd5f 100644 --- a/tensorrt_llm/_torch/attention_backend/utils.py +++ b/tensorrt_llm/_torch/attention_backend/utils.py @@ -28,6 +28,14 @@ def get_attention_backend( if sparse_params is not None: return get_trtllm_sparse_attn_attention_backend(sparse_params) return TrtllmAttention + elif backend_name == "INKLING": + # Intentionally NOT added to the attn_backend telemetry categorical list + # in llm_args.py: touching that file trips the API-stability label gate + # and stales the golden manifest, and "INKLING" is never user-supplied + # anyway -- Inkling's get_model_defaults selects it, and an explicit + # override is rejected at load. + from .inkling import InklingTritonAttention + return InklingTritonAttention elif backend_name == "FLASHINFER" and IS_FLASHINFER_AVAILABLE: from .flashinfer import FlashInferAttention if sparse_params is not None: diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index c5893c21bfae..482e683e644a 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -23,6 +23,7 @@ Gemma4UnifiedTextConfig, Gemma4UnifiedVisionConfig, ) +from tensorrt_llm._torch.configs.inkling import InklingConfig, InklingTextConfig from tensorrt_llm._torch.configs.laguna import LagunaConfig from tensorrt_llm._torch.configs.minicpmv4_6 import MiniCPMV4_6Config, MiniCPMV4_6VisionConfig @@ -63,6 +64,8 @@ def _register_custom_configs_with_transformers() -> None: "gemma4_unified_text": Gemma4UnifiedTextConfig, "gemma4_unified_vision": Gemma4UnifiedVisionConfig, "gemma4_unified_audio": Gemma4UnifiedAudioConfig, + "inkling_mm_model": InklingConfig, + "inkling_text": InklingTextConfig, } # Cosmos3Config resolves vision sub-configs via ``qwen3_vl_vision``; that # alias is only present in newer transformers releases. @@ -86,6 +89,8 @@ def _register_custom_configs_with_transformers() -> None: "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", "Gemma4UnifiedVisionConfig", + "InklingConfig", + "InklingTextConfig", "LagunaConfig", "MiniCPMV4_6Config", "MiniCPMV4_6VisionConfig", diff --git a/tensorrt_llm/_torch/configs/inkling.py b/tensorrt_llm/_torch/configs/inkling.py new file mode 100644 index 000000000000..e7d517ad278b --- /dev/null +++ b/tensorrt_llm/_torch/configs/inkling.py @@ -0,0 +1,250 @@ +# 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. +"""Config classes for the Inkling multimodal checkpoint. + +The checkpoint publishes ``config.json`` with ``model_type == +"inkling_mm_model"`` and ``architectures == +["InklingForConditionalGeneration"]``. Transformers does not ship Inkling, so +these classes reconstruct the config from the checkpoint's nested dicts without +any transformers shim. + +The audio, vision, and MTP sub-configs are kept verbatim (as +``PretrainedConfig`` blobs); the vision and audio towers read their geometry +directly off them (see ``models/modeling_inkling_multimodal.py``). + +Field names mirror the checkpoint ``text_config``. All numeric defaults are the +real checkpoint values, but a ``config.json`` that spells a field out overrides +the default via ``from_dict``. +""" + +from transformers.configuration_utils import PretrainedConfig + + +class InklingTextConfig(PretrainedConfig): + """Text-tower sub-config (``InklingCausalLLM``). + + A RoPE-free hybrid-attention decoder: per-head q/k RMSNorm, learned + relative-position bias, four short convolutions per layer, sigmoid-gated MoE + with two shared experts, muP logit scaling, and an unpadded vocab slice. + """ + + model_type = "inkling_text" + + def __init__( + self, + vocab_size: int = 201024, + unpadded_vocab_size: int = 200058, + hidden_size: int = 6144, + num_hidden_layers: int = 66, + num_attention_heads: int = 64, + num_key_value_heads: int = 8, + head_dim: int = 128, + rms_norm_eps: float = 1e-6, + model_max_length: int = 1048576, + logits_mup_width_multiplier: float = 24.0, + use_embed_norm: bool = True, + tie_word_embeddings: bool = False, + # hybrid attention geometry + local_layer_ids: list[int] | None = None, + sliding_window_size: int = 512, + swa_num_attention_heads: int = 64, + swa_num_key_value_heads: int = 16, + swa_head_dim: int = 128, + # relative-bias / log-scaling + d_rel: int = 16, + rel_extent: int = 1024, + log_scaling_n_floor: int = 128000, + log_scaling_alpha: float = 0.1, + # short conv + use_sconv: bool = True, + sconv_kernel_size: int = 4, + # dense MLP / MoE + dense_mlp_idx: int = 2, + intermediate_size: int = 3072, + dense_intermediate_size: int = 24576, + n_routed_experts: int = 256, + num_experts_per_tok: int = 6, + n_shared_experts: int = 2, + shared_expert_sink: bool = True, + route_scale: float = 8.0, + use_gate_bias: bool = True, + gate_activation: str = "sigmoid", + norm_after_topk: bool = True, + use_global_scale: bool = True, + hidden_act: str = "silu", + attention_dropout: float = 0.0, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.vocab_size = vocab_size + self.unpadded_vocab_size = unpadded_vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.rms_norm_eps = rms_norm_eps + self.model_max_length = model_max_length + # `max_position_embeddings` is read by several TRT-LLM code paths + # (Attention, RopeParams probing); Inkling has no RoPE but keep the 1M + # context window available so nothing clamps sequence length. + self.max_position_embeddings = kwargs.get("max_position_embeddings", model_max_length) + self.logits_mup_width_multiplier = logits_mup_width_multiplier + self.use_embed_norm = use_embed_norm + + self.local_layer_ids = list(local_layer_ids) if local_layer_ids else [] + self.sliding_window_size = sliding_window_size + self.swa_num_attention_heads = swa_num_attention_heads + self.swa_num_key_value_heads = swa_num_key_value_heads + self.swa_head_dim = swa_head_dim + + self.d_rel = d_rel + self.rel_extent = rel_extent + self.log_scaling_n_floor = log_scaling_n_floor + self.log_scaling_alpha = log_scaling_alpha + + self.use_sconv = use_sconv + self.sconv_kernel_size = sconv_kernel_size + + self.dense_mlp_idx = dense_mlp_idx + self.intermediate_size = intermediate_size + # A config.json may spell this out; PretrainedConfig.__init__ has already + # stored it from **kwargs, so read it back before defaulting. + self.moe_intermediate_size = kwargs.get("moe_intermediate_size", intermediate_size) + self.dense_intermediate_size = dense_intermediate_size + self.n_routed_experts = n_routed_experts + self.num_experts_per_tok = num_experts_per_tok + self.n_shared_experts = n_shared_experts + self.shared_expert_sink = shared_expert_sink + self.route_scale = route_scale + self.use_gate_bias = use_gate_bias + self.gate_activation = gate_activation + self.norm_after_topk = norm_after_topk + self.use_global_scale = use_global_scale + self.hidden_act = hidden_act + self.attention_dropout = attention_dropout + + # ---- per-layer classification helpers (single source of truth) ---- + @property + def _local_ids(self) -> set: + return set(self.local_layer_ids) + + def is_dense_layer(self, layer_idx: int) -> bool: + """Dense MLP layers are the ones with index < ``dense_mlp_idx``.""" + return layer_idx < self.dense_mlp_idx + + def is_local_layer(self, layer_idx: int) -> bool: + """Local (sliding-window) layers are listed in ``local_layer_ids``.""" + return layer_idx in self._local_ids + + def layer_num_kv_heads(self, layer_idx: int) -> int: + return ( + self.swa_num_key_value_heads + if self.is_local_layer(layer_idx) + else self.num_key_value_heads + ) + + def layer_num_heads(self, layer_idx: int) -> int: + return ( + self.swa_num_attention_heads + if self.is_local_layer(layer_idx) + else self.num_attention_heads + ) + + def layer_head_dim(self, layer_idx: int) -> int: + return self.swa_head_dim if self.is_local_layer(layer_idx) else self.head_dim + + def layer_window(self, layer_idx: int) -> int | None: + """Sliding-window size for local layers; ``None`` for global layers.""" + return self.sliding_window_size if self.is_local_layer(layer_idx) else None + + def num_kv_heads_per_layer(self) -> list[int]: + """Per-layer KV-head counts for the hybrid attention geometry. + + Local (sliding-window) layers use ``swa_num_key_value_heads`` (16) and + global layers use ``num_key_value_heads`` (8). ``KVCacheManagerV2`` + accepts this ``List[int]`` as ``num_kv_heads`` (it divides each by + ``tp_size``), so the paged KV cache allocates the right per-layer head + count instead of a single uniform value. ``head_dim`` is uniform (128) + across local and global layers, so only the KV-head count varies. + """ + return [self.layer_num_kv_heads(i) for i in range(self.num_hidden_layers)] + + +class InklingConfig(PretrainedConfig): + """Top-level Inkling multimodal config (``inkling_mm_model``). + + Reconstructs ``text_config`` with :class:`InklingTextConfig`; + ``audio_config``, ``vision_config`` and ``mtp_config`` are retained as plain + ``PretrainedConfig`` blobs so the multimodal checkpoint round-trips. + """ + + model_type = "inkling_mm_model" + sub_configs = {"text_config": InklingTextConfig} + + def __init__( + self, + text_config=None, + audio_config=None, + vision_config=None, + mtp_config=None, + eos_token_id: int = 200006, + # Multimodal placeholder ids. One appears per media item in the + # pre-rendered token stream; the input processor expands it to one token + # per patch (image) / audio frame and the vision/audio fusion overwrites + # those positions. Both placeholders are IN-VOCAB chat-template tokens: + # image ``<|unused_200054|>`` (id 200054) and audio ``<|unused_200053|>`` + # (id 200053) -- exactly the tokens the Inkling chat template renders for + # an image / audio content part + # (``<|content_image|><|unused_200054|>...`` and + # ``<|content_audio_input|><|unused_200053|><|audio_end|>...``). + # They MUST be in-vocab: TensorRT-LLM's executor validates request token + # ids and rejects an out-of-range id. The checkpoint's config.json omits + # these, so the defaults apply unless a config spells them out. + image_token_id: int = 200054, + audio_token_id: int = 200053, + tie_word_embeddings: bool = False, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.eos_token_id = eos_token_id + self.image_token_id = image_token_id + self.audio_token_id = audio_token_id + + if text_config is None: + self.text_config = InklingTextConfig() + elif isinstance(text_config, dict): + self.text_config = InklingTextConfig(**text_config) + else: + self.text_config = text_config + + # Retained verbatim; interpreted only in the Phase-3 multimodal stage. + self.audio_config = self._as_config(audio_config) + self.vision_config = self._as_config(vision_config) + # Retained verbatim so the checkpoint round-trips. MTP / next-N draft + # decoding is not supported: nothing builds or loads these layers, and + # the draft weights stay in the weight mapper's deferred set. + self.mtp_config = self._as_config(mtp_config) + + @staticmethod + def _as_config(value): + if value is None or isinstance(value, PretrainedConfig): + return value + if isinstance(value, dict): + cfg = PretrainedConfig() + for k, v in value.items(): + setattr(cfg, k, v) + return cfg + return value diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 08f6f325329c..882ed829d324 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -391,12 +391,20 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, quant_config = QuantConfig() layer_quant_config = None - quant_config.quant_algo = (QuantAlgo(json_quant_configs['quant_algo']) - if json_quant_configs.get('quant_algo') - is not None else None) - quant_config.kv_cache_quant_algo = ( - QuantAlgo(json_quant_configs['kv_cache_quant_algo']) if - json_quant_configs.get('kv_cache_quant_algo') is not None else None) + def _algo_or_none(value): + # modelopt hf_quant_config.json may spell "no quantization" as JSON + # null (-> None) OR as the string "none"/"null" (e.g. the Inkling + # NVFP4 checkpoint uses ``"kv_cache_quant_algo": "none"``); both must + # map to None rather than QuantAlgo("none"), which is not a member. + if value is None or (isinstance(value, str) and + value.strip().lower() in ("none", "null", "")): + return None + return QuantAlgo(value) + + quant_config.quant_algo = _algo_or_none( + json_quant_configs.get('quant_algo')) + quant_config.kv_cache_quant_algo = _algo_or_none( + json_quant_configs.get('kv_cache_quant_algo')) quant_config.group_size = json_quant_configs.get('group_size', None) quant_config.exclude_modules = json_quant_configs.get( 'exclude_modules', None) @@ -420,10 +428,8 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, ) json_quant_configs.update(json_extended_quant_configs) # kv_cache_quant_algo is global regardless of MIXED_PRECISION - kv_cache_quant_algo = (QuantAlgo( - json_quant_configs['kv_cache_quant_algo']) if - json_quant_configs.get('kv_cache_quant_algo') - is not None else None) + kv_cache_quant_algo = _algo_or_none( + json_quant_configs.get('kv_cache_quant_algo')) mixed_quant_configs = json_quant_configs.get( 'quantized_layers', None) if (kv_quant_lhs := json_extended_quant_configs.get( diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 4f4905d83086..cc519e1b6b2c 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -29,6 +29,8 @@ from .modeling_hunyuan_dense import HunYuanDenseV1ForCausalLM from .modeling_hunyuan_moe import HunYuanMoEV1ForCausalLM from .modeling_hyperclovax import HCXVisionForCausalLM +from .modeling_inkling import (InklingForCausalLM, + InklingForConditionalGeneration) from .modeling_kimi_k25 import KimiK25ForConditionalGeneration from .modeling_laguna import LagunaForCausalLM from .modeling_llama import LlamaForCausalLM @@ -85,6 +87,8 @@ "Gemma4ForConditionalGeneration", "Gemma4UnifiedForConditionalGeneration", "HCXVisionForCausalLM", + "InklingForCausalLM", + "InklingForConditionalGeneration", "LagunaForCausalLM", "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", diff --git a/tensorrt_llm/_torch/models/checkpoints/__init__.py b/tensorrt_llm/_torch/models/checkpoints/__init__.py index ab2f322ab51d..33f52db851ba 100644 --- a/tensorrt_llm/_torch/models/checkpoints/__init__.py +++ b/tensorrt_llm/_torch/models/checkpoints/__init__.py @@ -5,6 +5,7 @@ from .hf.cosmos3_weight_mapper import Cosmos3HfWeightMapper from .hf.gemma3_weight_mapper import Gemma3HfWeightMapper from .hf.gemma4_weight_mapper import Gemma4HfWeightMapper +from .hf.inkling_weight_mapper import InklingHfWeightMapper from .hf.llama4_weight_mapper import Llama4HfWeightMapper from .hf.llava_next_weight_mapper import LlavaNextHfWeightMapper from .hf.mixtral_weight_mapper import MixtralHfWeightMapper @@ -33,7 +34,7 @@ "MixtralHfWeightMapper", "Llama4HfWeightMapper", "Qwen2MoeHfWeightMapper", "Qwen3MoeHfWeightMapper", "Qwen2VLHfWeightMapper", "Qwen3_5MoeHfWeightMapper", "Qwen3NextHfWeightMapper", - "Gemma4HfWeightMapper", "LlavaNextHfWeightMapper", + "Gemma4HfWeightMapper", "InklingHfWeightMapper", "LlavaNextHfWeightMapper", "MistralLarge3CheckpointLoader", "MistralLarge3WeightMapper", "MXCheckpointLoader", "Qwen3VLHfWeightMapper", "Cosmos3HfWeightMapper" ] diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py new file mode 100644 index 000000000000..0dfc7f045b11 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py @@ -0,0 +1,400 @@ +# 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. +"""HF -> TensorRT-LLM weight mapping for the Inkling text tower. + +Two responsibilities: + +1. **Accounting (authoritative, CPU-testable).** :func:`inkling_expected_text_keys` + and :func:`inkling_account_checkpoint` derive the exact set of ``model.llm.*`` + checkpoint keys the text loader consumes, and classify every checkpoint key as + consumed-text / intentionally-deferred (audio, vision, MTP) / unaccounted. + Pinned by ``tests/unittest/_torch/modeling/test_modeling_inkling.py`` against + the real checkpoint index (no GPU). It guarantees no missing q/k-norm, + rel-bias, short-conv, route/global-scale or unpadded-logit tensor can hide. + +2. **Name/layout remapping (the load path).** :class:`InklingHfWeightMapper` + renames the checkpoint's keys (``wq_du``, ``w13_weight`` …) to + the TRT-LLM module tree, fuses q/k/v into the attention ``qkv_proj``, and + unfuses the NVFP4 routed experts (``w13_weight`` -> per-expert ``w1``/``w3`` + with their block scales) into the layout the fused-MoE loader expects. +""" + +from __future__ import annotations + +import re +from typing import Dict, List, Set, Tuple + +import torch + +from tensorrt_llm._torch.configs.inkling import InklingTextConfig +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper + +# NVFP4 two-level-scale maxima: E2M1 element max (6.0) and E4M3 block-scale max +# (448.0). ModelOpt stores the per-tensor activation ``input_scale`` as +# ``amax / (E2M1_MAX * E4M3_MAX)``; Inkling's checkpoint instead ships the raw +# ``.input_amax``, so the mapper must apply this conversion (see ``_map_expert``). +_NVFP4_E2M1_MAX = 6.0 +_NVFP4_E4M3_MAX = 448.0 + +# Prefixes the text loader does not consume: the vision / audio towers load +# themselves (see modeling_inkling_multimodal.py) and MTP is not implemented. +INKLING_DEFERRED_PREFIXES: Tuple[str, ...] = ( + "model.audio.", + "model.visual.", + "model.mtp.", +) + +# Per-layer checkpoint keys (relative to ``model.llm.layers.N.``), present in +# every one of the 66 decoder layers. +_ATTN_AND_NORM_KEYS: Tuple[str, ...] = ( + "attn.wq_du.weight", + "attn.wk_dv.weight", + "attn.wv_dv.weight", + "attn.wr_du.weight", + "attn.wo_ud.weight", + "attn.q_norm.weight", + "attn.k_norm.weight", + "attn.k_sconv.weight", + "attn.v_sconv.weight", + "attn.rel_logits_proj.proj", + "attn_norm.weight", + "mlp_norm.weight", + "attn_sconv.weight", + "mlp_sconv.weight", +) + +# Dense MLP (layers 0, 1). +_DENSE_MLP_KEYS: Tuple[str, ...] = ( + "mlp.w13_dn.weight", + "mlp.w2_md.weight", + "mlp.global_scale", +) + +# MoE common (all MoE layers, bf16 or NVFP4). +_MOE_COMMON_KEYS: Tuple[str, ...] = ( + "mlp.experts.w13_weight", + "mlp.experts.w2_weight", + "mlp.gate.weight", + "mlp.gate.bias", + "mlp.gate.global_scale", + "mlp.shared_experts.shared_w13_weight", + "mlp.shared_experts.shared_w2_weight", +) + +# NVFP4 sidecars attached to each routed-expert weight tensor (layers 3..65). +_NVFP4_SIDECARS: Tuple[str, ...] = (".input_amax", ".original_shape", ".scale", ".scale2") +_NVFP4_QUANTIZED_EXPERT_TENSORS: Tuple[str, ...] = ( + "mlp.experts.w13_weight", + "mlp.experts.w2_weight", +) + +_NON_LAYER_TEXT_KEYS: Tuple[str, ...] = ( + "model.llm.embed.weight", + "model.llm.embed_norm.weight", + "model.llm.norm.weight", + "model.llm.unembed.weight", +) + + +def _experts_are_nvfp4(layer_idx: int, exclude_modules: Set[str]) -> bool: + """Routed experts of an MoE layer are NVFP4 unless explicitly excluded.""" + return f"model.llm.layers.{layer_idx}.mlp.experts" not in exclude_modules + + +def inkling_expected_text_keys(config: InklingTextConfig, exclude_modules: Set[str]) -> Set[str]: + """Exact set of ``model.llm.*`` checkpoint keys the text loader consumes.""" + keys: Set[str] = set(_NON_LAYER_TEXT_KEYS) + for n in range(config.num_hidden_layers): + pfx = f"model.llm.layers.{n}." + for k in _ATTN_AND_NORM_KEYS: + keys.add(pfx + k) + if config.is_dense_layer(n): + for k in _DENSE_MLP_KEYS: + keys.add(pfx + k) + else: + for k in _MOE_COMMON_KEYS: + keys.add(pfx + k) + if _experts_are_nvfp4(n, exclude_modules): + for base in _NVFP4_QUANTIZED_EXPERT_TENSORS: + for side in _NVFP4_SIDECARS: + keys.add(pfx + base + side) + return keys + + +def inkling_account_checkpoint( + all_keys: Set[str], config: InklingTextConfig, exclude_modules: Set[str] +) -> Dict[str, Set[str]]: + """Classify every checkpoint key into consumed-text / deferred / + unaccounted. + + ``unaccounted`` and ``missing`` must both be empty for the checkpoint to be + fully and exactly accounted. The audio, vision and MTP blocks are carried by + the checkpoint but not loaded by the text tower, so their keys land in + ``deferred`` (see ``INKLING_DEFERRED_PREFIXES``). + """ + expected = inkling_expected_text_keys(config, exclude_modules) + consumed_text = all_keys & expected + deferred = {k for k in all_keys if k.startswith(INKLING_DEFERRED_PREFIXES)} + missing = expected - all_keys + + unaccounted = all_keys - consumed_text - deferred + return { + "consumed_text": consumed_text, + "deferred": deferred, + "unaccounted": unaccounted, + "missing": missing, + } + + +def inkling_nvfp4_expert_layers(config: InklingTextConfig, exclude_modules: Set[str]) -> List[int]: + """Layers whose routed experts are stored as NVFP4 (expected: 3..65).""" + return [ + n + for n in range(config.num_hidden_layers) + if not config.is_dense_layer(n) and _experts_are_nvfp4(n, exclude_modules) + ] + + +# --------------------------------------------------------------------------- +# Load path +# --------------------------------------------------------------------------- +# Simple 1:1 renames from the (``model.llm.`` stripped) checkpoint name to the +# TRT-LLM module tree. +_SIMPLE_RENAMES = { + "embed.weight": "model.embed_tokens.weight", + "embed_norm.weight": "model.embed_norm.weight", + "norm.weight": "model.norm.weight", + "unembed.weight": "lm_head.weight", +} + +# Per-layer renames (regex on the ``layers.N.`` tail -> TRT name tail). +# q/k/v map to the standard separate HF names at the ``attn.`` level; the fused +# ``qkv_proj`` Linear's loader collects attn.q_proj/k_proj/v_proj via its +# special-handling callback and fuses them. Same for gate_up_proj <- gate_proj + +# up_proj (the dense w13_dn tensor is pre-fused and is split in _map_dense_w13). +_LAYER_RENAMES = { + "attn.wq_du.weight": "attn.q_proj.weight", + "attn.wk_dv.weight": "attn.k_proj.weight", + "attn.wv_dv.weight": "attn.v_proj.weight", + "attn.wo_ud.weight": "attn.o_proj.weight", + "attn.wr_du.weight": "attn.r_proj.weight", + "attn.q_norm.weight": "attn.q_norm.weight", + "attn.k_norm.weight": "attn.k_norm.weight", + "attn.k_sconv.weight": "attn.k_sconv.weight", + "attn.v_sconv.weight": "attn.v_sconv.weight", + "attn.rel_logits_proj.proj": "attn.rel_logits_proj", + "attn_norm.weight": "attn_norm.weight", + "mlp_norm.weight": "mlp_norm.weight", + "attn_sconv.weight": "attn_sconv.weight", + "mlp_sconv.weight": "mlp_sconv.weight", + # dense (w13_dn is split in _map_dense_w13; w2_md -> down_proj) + "mlp.w2_md.weight": "mlp.down_proj.weight", + "mlp.global_scale": "mlp.global_scale", + # moe (non-expert) + "mlp.gate.weight": "mlp.gate.weight", + "mlp.gate.bias": "mlp.gate.bias", + "mlp.gate.global_scale": "mlp.gate.global_scale", + "mlp.shared_experts.shared_w13_weight": "mlp.shared_experts.shared_w13", + "mlp.shared_experts.shared_w2_weight": "mlp.shared_experts.shared_w2", +} + +_EXPERT_RE = re.compile(r"layers\.(\d+)\.mlp\.experts\.(w13_weight|w2_weight)(\.\w+)?$") +_DENSE_W13_RE = re.compile(r"layers\.(\d+)\.mlp\.w13_dn\.weight$") + + +def _split_interleaved_gate_up(t: torch.Tensor, dim: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Split an Inkling gate/up-INTERLEAVED fused tensor into ``(gate, up)`` STRIDED + VIEWS (no copy) along ``dim``: gate = even indices, up = odd indices. + + The Inkling checkpoint stores every fused gate+up weight with the two + projections INTERLEAVED along the output (``2*inter``) dim: + ``[g0, u0, g1, u1, ...]``, i.e. the reference SwiGLU reads it as + ``silu(z[..., ::2]) * z[..., 1::2]``. TRT-LLM's fused gate_up / fused-MoE + loaders instead want separate gate/up; a plain contiguous ``chunk(2)`` + (``[first half | second half]``) would pair the WRONG gate/up channels in + every dense-MLP, routed-expert and shared-expert SwiGLU. + + Returns STRIDED VIEWS rather than a contiguous copy on purpose: the fused-MoE + / gate_up loaders shard each rank's slice then call ``.contiguous()`` on that + small shard (see quantization.py ``load_expert_w3_w1_weight``), so no + full-tensor host copy is needed. A contiguous de-interleave here instead + materialized a private per-rank copy of the ~hundreds-of-GiB fused w13, + doubling host memory and OOM-killing the TP=4 load. Reorders whole output + rows only -> valid for a packed NVFP4 weight and its per-block fp8 scale. + """ + dim = dim % t.dim() + if t.shape[dim] % 2 != 0: + raise ValueError(f"cannot split odd gate/up dim {dim}: {tuple(t.shape)}") + even = [slice(None)] * t.dim() + odd = [slice(None)] * t.dim() + even[dim] = slice(0, None, 2) + odd[dim] = slice(1, None, 2) + return t[tuple(even)], t[tuple(odd)] + + +@register_mapper("HF", "InklingForConditionalGeneration") +class InklingHfWeightMapper(HfWeightMapper): + """Renames Inkling checkpoint keys to the TRT-LLM module tree. + + Runs after ``filter_weights("model.llm", ...)`` in the model's + ``load_weights`` (so incoming keys start at ``layers.N.…`` / ``embed.weight`` + …). The NVFP4 routed experts are unfused from the checkpoint's stacked, + gate+up-fused ``w13_weight [E, 2*inter, hidden/2]`` into the per-expert + ``w1``/``w3`` layout (plus block ``weight_scale``, per-expert + ``weight_scale_2`` and ``input_scale``) that the fused-MoE loader consumes. + """ + + @property + def _text_config(self) -> InklingTextConfig: + """The text sub-config the mapped weights actually describe. + + ``ModelLoader.load`` initializes the mapper with the TOP-LEVEL + ``InklingConfig``, which carries the decoder geometry under + ``text_config`` and has no ``vocab_size`` / ``n_routed_experts`` of its + own. Resolving here covers both entry points: the loader-supplied mapper + and the one ``InklingForConditionalGeneration.load_weights`` builds. + """ + cfg = self.config.pretrained_config + return getattr(cfg, "text_config", cfg) + + def preprocess_weights(self, weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + new_weights: dict[str, torch.Tensor] = {} + text_config = self._text_config + unpadded_vocab = int( + getattr( + text_config, + "unpadded_vocab_size", + text_config.vocab_size, + ) + ) + for name, tensor in weights.items(): + if name in _SIMPLE_RENAMES: + if name == "unembed.weight" and tensor.shape[0] > unpadded_vocab: + # The checkpoint LM-head matrix is padded to vocab_size + # (201024); the text tower emits logits only over the + # unpadded vocab (200058). Dropping the padding rows here is + # exactly the required "slice logits to unpadded" (logit[i] + # = h @ unembed[i]), and lets LMHead(num_embeddings=200058) + # load without a shape mismatch. embed_tokens keeps the full + # matrix (built at vocab_size), so input ids stay in range. + tensor = tensor[:unpadded_vocab] + new_weights[_SIMPLE_RENAMES[name]] = tensor + continue + + expert_match = _EXPERT_RE.search(name) + if expert_match is not None: + self._map_expert(name, tensor, expert_match, new_weights) + continue + + dense_match = _DENSE_W13_RE.search(name) + if dense_match is not None: + # Dense w13_dn is gate/up-INTERLEAVED [g0,u0,...] along the output + # (2*inter) dim; split into gate (even rows) / up (odd rows) for + # the fused gate_up_proj loader (strided views, no copy). + layer_idx = dense_match.group(1) + gate, up = _split_interleaved_gate_up(tensor, dim=0) + new_weights[f"model.layers.{layer_idx}.mlp.gate_proj.weight"] = gate + new_weights[f"model.layers.{layer_idx}.mlp.up_proj.weight"] = up + continue + + # shared_experts.shared_w13_weight loads RAW (interleaved) via + # _LAYER_RENAMES; the gate/up interleave is undone by the strided split + # in InklingSharedExperts.forward (zero-copy, param materialized once). + + m = re.match(r"layers\.(\d+)\.(.*)$", name) + if m is not None: + layer_idx, tail = m.group(1), m.group(2) + trt_tail = _LAYER_RENAMES.get(tail, tail) + new_weights[f"model.layers.{layer_idx}.{trt_tail}"] = tensor + continue + + # Unknown key: keep as-is so any mismatch surfaces loudly at load. + new_weights[name] = tensor + return new_weights + + def _map_expert( + self, + name: str, + tensor: torch.Tensor, + match: re.Match[str], + out: dict[str, torch.Tensor], + ) -> None: + """Unfuse a stacked expert tensor into per-expert fused-MoE keys. + + ``w13_weight[e]`` is ``[2*inter, hidden]`` (gate rows first, up rows + second, per HF ``InklingExperts``); split into ``w1`` (gate) and ``w3`` + (up). ``w2_weight[e]`` is the down projection. NVFP4 sidecars map to the + fused-MoE scale names: ``.scale`` -> ``weight_scale`` (block), + ``.scale2`` -> ``weight_scale_2`` (per-expert), ``.input_amax`` -> + ``input_scale``. ``.original_shape`` is metadata and is dropped. + """ + layer_idx, which, sidecar = match.group(1), match.group(2), match.group(3) + prefix = f"model.layers.{layer_idx}.mlp.experts" + + scale_name = { + None: "weight", + ".scale": "weight_scale", + ".scale2": "weight_scale_2", + ".input_amax": "input_scale", + }.get(sidecar) + if scale_name is None: # .original_shape -> drop (layout metadata) + return + + if sidecar == ".input_amax": + # Inkling's NVFP4 checkpoint stores the routed-expert activation + # calibration as a RAW amax (``.input_amax``). The fused-MoE loader + # (``NVFP4FusedMoEMethod.process_weights_after_loading`` -> + # ``fc31_input_scale = 1 / max_e(input_scale)``) and + # ``torch.ops.trtllm.fp4_quantize`` expect the ModelOpt per-tensor + # activation ``input_scale = amax / (E2M1_MAX * E4M3_MAX)``, so the + # activation global scale ``1 / max_e(input_scale)`` lands the e4m3 + # activation block scales in range. Without this conversion the global + # scale is (E2M1_MAX*E4M3_MAX)=2688x too small, the activation fp4 + # block scales underflow e4m3, and every routed expert diverges from + # the bf16 ground truth. Weight / block-scale / scale2 layout is + # already correct, so ONLY the activation input scale needs this. + tensor = tensor.to(torch.float32) / (_NVFP4_E2M1_MAX * _NVFP4_E4M3_MAX) + + n_experts = int(getattr(self._text_config, "n_routed_experts", tensor.shape[0])) + projs = ("w1", "w3") if which == "w13_weight" else ("w2",) + + def _assign(e: int, vals: tuple[torch.Tensor, ...]) -> None: + for proj, val in zip(projs, vals): + out[f"{prefix}.{e}.{proj}.{scale_name}"] = val + + # Three sidecar shapes: per-expert multi-dim weight/block-scale (chunk + # w13 into gate/up along the out dim), per-expert scalar weight_scale_2 + # (same value for gate and up), and a single global input_amax scalar + # broadcast to every expert/proj. + if tensor.dim() >= 2 and tensor.shape[0] == n_experts: + for e in range(n_experts): + if which == "w13_weight": + # w13 (packed fp4 weight AND its per-block fp8 scale) is + # gate/up-INTERLEAVED [g0,u0,...] along the per-expert output + # (2*inter) dim; split into w1 (gate = even rows) / w3 (up = + # odd rows) as strided views (no copy). Reorders whole rows, so + # it is correct for both the uint8 weight and the fp8 scale. + per = _split_interleaved_gate_up(tensor[e], dim=0) + else: + per = (tensor[e],) + _assign(e, per) + elif tensor.dim() >= 1 and tensor.shape[0] == n_experts: + for e in range(n_experts): + _assign(e, (tensor[e],) * len(projs)) + else: # global scalar (input_amax [1]) -> broadcast to all experts + val = tensor.reshape(-1)[0] + for e in range(n_experts): + _assign(e, (val,) * len(projs)) diff --git a/tensorrt_llm/_torch/models/modeling_inkling.py b/tensorrt_llm/_torch/models/modeling_inkling.py new file mode 100644 index 000000000000..cdf0f7ab82c8 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_inkling.py @@ -0,0 +1,2073 @@ +# 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 text decoder for the TensorRT-LLM PyTorch backend. + +The multimodal towers and their input processor live in +``modeling_inkling_multimodal.py``; see ``configs/inkling.py`` for the config +classes and ``checkpoints/hf/inkling_weight_mapper.py`` for the HF -> TRT weight +mapping. MTP is not implemented (its weights are accounted as unused). + +Architecture summary: + * RoPE-free attention with per-head q/k RMSNorm and score scale ``1/head_dim``. + * Learned relative-position bias (``RelLogitsProj``), added pre-softmax as a + ``score_mod`` inside the Inkling Triton attention kernels (prefill + paged + decode); see ``attention_backend/inkling/``. + * Hybrid layers: 55 local sliding-window (win=512, 16 kv-heads) + 11 global + full-causal (8 kv-heads). Global layers apply log-scaling tau (a no-op below + 128k tokens, still implemented for correctness). + * Four causal short convolutions per layer (k, v inside attention before the + k/q norm; one post-attention and one post-MLP on the residual stream). + * Sigmoid-gated MoE, top-6 of 256 routed experts with an additive selection + bias, log-sigmoid renorm over the selected-routed *plus* two shared logits, + scaled by ``route_scale * global_scale``. Layers 0/1 are dense MLP. + * Routed experts for layers 3..65 are NVFP4; layer-2 experts and everything + else are bf16. + * muP: divide hidden states by ``logits_mup_width_multiplier`` before the head; + slice logits to ``unpadded_vocab_size``. ``embed_norm`` folds onto embeddings. +""" + +import copy +from collections import namedtuple +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Optional, Tuple + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +import torch +from torch import nn + +from tensorrt_llm._torch.attention_backend import AttentionMetadata +from tensorrt_llm._torch.attention_backend.inkling import ( + build_page_table, + inkling_decode_attention, + inkling_prefill_attention, + write_kv_cache_hnd, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_utils import ( + DecoderModel, + DecoderModelForCausalLM, + filter_weights, + register_auto_model, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.fused_moe import ( + BaseMoeRoutingMethod, + RoutingMethodType, + create_moe, +) +from tensorrt_llm._torch.modules.linear import ( + Linear, + TensorParallelMode, + WeightMode, + WeightsLoadingConfig, +) +from tensorrt_llm._torch.modules.mamba.causal_conv1d import causal_conv1d_fn, causal_conv1d_update +from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +# Protocol only -- no BaseResourceManager / KV-cache import, so this stays free +# of an import cycle back through pyexecutor at model-load time. +from tensorrt_llm._utils import prefer_pinned + +from ...inputs import ( + ContentFormat, + MultimodalPlaceholderMetadata, + MultimodalPlaceholderPlacement, + register_input_processor, +) +from ..configs.inkling import InklingConfig, InklingTextConfig +from .modeling_inkling_multimodal import ( + DEFAULT_AUDIO_TOKEN_ID, + DEFAULT_IMAGE_TOKEN_ID, + InklingAudioModel, + InklingInputProcessor, + InklingVisionModel, +) +from .modeling_multimodal_utils import ( + filter_mm_token_from_input_ids, + find_input_mm_embeds, + fuse_input_embeds, +) + +# Per-layer, per-request short-conv state carried across decode steps: the four +# causal short convolutions of one Inkling decoder layer. Each field is a +# ``[num_req, channels, sconv_kernel_size - 1]`` buffer holding the previous +# ``kernel_size - 1`` pre-conv inputs (oldest first): ``k``/``v`` for the +# attention k/v short-convs (channels = num_kv_heads * head_dim, TP-sharded) and +# ``attn``/``mlp`` for the post-attention / post-MLP residual-stream short-convs +# (channels = hidden_size, replicated). The generation phase reads these, +# convolves the one new token, and rolls the window forward IN PLACE, so the +# buffers keep stable addresses across decode steps and CUDA-graph replay. +InklingConvState = namedtuple("InklingConvState", ["k", "v", "attn", "mlp"]) + + +class InklingConvStateCache: + """Runtime-owned per-request short-conv state pool for the whole decoder. + + Carries the four causal short-convs of every decoder layer per request + across decode steps, with the same lifetime as the paged KV cache. + + Per layer it allocates the four short-conv state buffers + (:class:`InklingConvState`), each ``[max_batch, channels, kernel_size - 1]`` + holding the previous ``kernel_size - 1`` pre-conv inputs (oldest first). The + k/v conv channels follow the fused-qkv k/v split (TP-sharded like + ``InklingShortConv(tp_shard=True)``); the post-attention / post-MLP convs run + on the full (all-reduced) hidden stream and are replicated. The buffers keep + stable device addresses for their whole lifetime -- the fused + ``causal_conv1d_update`` / ``causal_conv1d_fn`` ops mutate them IN PLACE at + the per-request ``state_indices`` slots, so a captured CUDA graph replays + cleanly (no realloc, no gather/scatter). ``state_indices`` is a single stable + ``[max_batch]`` int32 CUDA buffer written in place per forward (the + Mamba2Metadata stable-pointer pattern), so the runtime can alias it under + graph capture. + """ + + def __init__( + self, + pretrained_config, + tp_size: int, + max_batch_size: int, + device: torch.device, + dtype: torch.dtype = torch.bfloat16, + ): + # Takes the pretrained config + tp_size rather than a ``ModelConfig`` so + # the KV cache manager can build the pool from the arguments + # ``_create_kv_cache_manager`` already passes it, without a ModelConfig. + # Accept either the text config or the top-level multimodal one. + config = getattr(pretrained_config, "text_config", pretrained_config) + kwin = config.sconv_kernel_size - 1 + self.max_batch_size = max_batch_size + self.kwin = kwin + + def buf(channels): + return torch.zeros(max_batch_size, channels, kwin, device=device, dtype=dtype) + + self._layers: List[InklingConvState] = [] + for i in range(config.num_hidden_layers): + kv_dim = (config.layer_num_kv_heads(i) * config.layer_head_dim(i)) // tp_size + hidden = config.hidden_size + self._layers.append( + InklingConvState(k=buf(kv_dim), v=buf(kv_dim), attn=buf(hidden), mlp=buf(hidden)) + ) + # Stable per-request slot-index buffer (int32, CUDA). Refreshed in place + # per forward -- EAGERLY, from input preparation, before CUDA-graph + # capture/replay (see :meth:`write_state_indices`) -- so a captured + # decode graph aliases it and every replay reads the current batch's + # rows (Mamba2Metadata stable-pointer pattern). + self.state_indices = torch.arange(max_batch_size, dtype=torch.int32, device=device) + # Pinned host staging for that per-forward write: the eager input-prep + # phase fills this and issues ONE async H2D copy into ``state_indices``. + # Pinned so the copy is cheap and legal even under graph capture; kept in + # lock-step size with ``state_indices`` across :meth:`_grow`. + self.state_indices_cpu = torch.zeros( + max_batch_size, dtype=torch.int32, pin_memory=prefer_pinned() + ) + self._slot_of = {} + self._free = list(range(max_batch_size - 1, -1, -1)) + + def layer_state(self, layer_idx: int) -> InklingConvState: + """The four short-conv state buffers for ``layer_idx`` (pool views).""" + return self._layers[layer_idx] + + def slots_for(self, request_ids: List[int]) -> List[int]: + """Map request ids to their (stable) pool rows, allocating new ones. + + Fresh requests get a zero-initialised slot; existing requests keep their + row so their carried short-conv windows persist across decode steps. + + If a single forward presents more *fresh* requests than the pool has + free rows, the pool grows to fit (see :meth:`_grow`). Steady-state + serving is bounded by ``max_batch_size`` (+1 CUDA-graph pad row) and + never triggers growth, but the one-time KV-cache estimation forward can + exceed it: that dummy batch is sized to saturate ``max_num_tokens`` (and + is replicated ``x tp_size`` under attention DP), independent of + ``max_batch_size``. Growing there (instead of ``IndexError`` on an empty + free list) lets estimation profile memory correctly, and because growth + only happens in that eager estimation/warmup window the buffers a later + CUDA graph captures are the final, pointer-stable ones. + """ + num_new = sum(1 for r in request_ids if r not in self._slot_of) + if num_new > len(self._free): + self._grow(num_new - len(self._free)) + slots = [] + for r in request_ids: + if r not in self._slot_of: + slot = self._free.pop() + self._slot_of[r] = slot + for st in self._layers: + for t in st: + t[slot].zero_() + slots.append(self._slot_of[r]) + return slots + + def _grow(self, extra: int): + """Append ``extra`` fresh (zeroed) rows to every per-request buffer. + + Reallocates each layer's four short-conv state tensors and the shared + ``state_indices`` scratch to ``max_batch_size + extra`` rows, copying the + existing rows forward so any in-flight request keeps its carried window, + and returns the new rows to the free list. Called only from + :meth:`slots_for` when a batch needs more rows than the pool owns; see + there for why that happens (KV-cache estimation / attention-DP), and why + it is safe w.r.t. CUDA-graph pointer stability. + """ + old = self.max_batch_size + new = old + extra + for i, st in enumerate(self._layers): + grown = [] + for t in st: + buf = torch.zeros(new, t.shape[1], t.shape[2], device=t.device, dtype=t.dtype) + buf[:old].copy_(t) + grown.append(buf) + self._layers[i] = InklingConvState(*grown) + self.state_indices = torch.arange(new, dtype=torch.int32, device=self.state_indices.device) + # Keep the pinned host-staging buffer sized in lock-step, else the eager + # H2D write in write_state_indices would index past its end. + self.state_indices_cpu = torch.zeros(new, dtype=torch.int32, pin_memory=prefer_pinned()) + # New rows old..new-1 join the free list, popped ascending like __init__. + self._free = list(range(new - 1, old - 1, -1)) + self._free + self.max_batch_size = new + + def write_state_indices(self, request_ids: List[int], is_graph: bool) -> List[int]: + """Resolve ``request_ids`` to pool rows and publish them into the stable + ``state_indices`` CUDA buffer -- the EAGER, pre-capture slot write. + + Returns the resolved slot list (context requests first, then + generation, matching the packed batch order). The host->device copy goes + through the pinned ``state_indices_cpu`` staging buffer so it is legal + under CUDA-graph capture and non-blocking. Because a captured decode + graph aliases ``state_indices`` (via the ``gen_indices`` view built in + :meth:`InklingConvRuntime.build`), this MUST run every forward from eager + input-prep -- NOT inside the captured ``model.forward`` -- so each replay + reads the current batch's rows rather than the stale capture-time ones. + + ``is_graph`` (``attn_metadata.is_cuda_graph``) guards pool-pointer + stability. Growth reallocates ``state_indices`` and would strand a + captured graph's aliased pointer, so it may only happen in the eager + estimation/warmup window (``is_graph`` False). The pool is sized + ``max_batch_size + 1`` >= any graph batch, so a graph forward never needs + to grow; assert it to turn a latent pointer bug into a loud failure + instead of silent decode corruption. + """ + before = self.state_indices.data_ptr() + slots = self.slots_for(request_ids) + if is_graph and self.state_indices.data_ptr() != before: + raise RuntimeError( + "Inkling short-conv pool grew during CUDA graph capture/replay; " + "the pool must be sized to the max graph batch up front (a grown " + "pool strands the captured state_indices pointer)." + ) + n = len(slots) + self.state_indices_cpu[:n].copy_(torch.tensor(slots, dtype=torch.int32)) + self.state_indices[:n].copy_(self.state_indices_cpu[:n], non_blocking=True) + return slots + + def free(self, request_ids: List[int]): + for r in request_ids: + slot = self._slot_of.pop(r, None) + if slot is not None: + self._free.append(slot) + + +@dataclass +class InklingConvRuntime: + """Per-forward short-conv plumbing for the pool path (all layers share it). + + Splits the packed ``[context tokens | one-token generation]`` batch at the + context boundary so each of the four short-convs seeds the pool for context + requests (varlen ``causal_conv1d_fn``) and updates it in place for generation + requests (``causal_conv1d_update``), exactly like the paged attention split + in :meth:`InklingAttention._attention`. + """ + + num_ctx_tokens: int + ctx_indices: Optional[torch.Tensor] # int32 pool slots, context requests + gen_indices: Optional[torch.Tensor] # int32 pool slots, generation requests + query_start_loc: Optional[torch.Tensor] # int32 [n_ctx+1] varlen offsets + has_initial_state: Optional[torch.Tensor] # bool [n_ctx] + + @classmethod + def build(cls, attn_metadata, cache: InklingConvStateCache) -> "InklingConvRuntime": + """Publish this batch's pool rows, then build the context/generation split. + + Resolves the batch's request ids to pool rows and writes them into the + stable ``state_indices`` buffer + (:meth:`InklingConvStateCache.write_state_indices`), then slices the + context/generation views of that buffer. The split mirrors the attention + split: context requests first (each with its full new-token span), then + one-token generation requests. Prefill-only tensors + (``query_start_loc`` / ``has_initial_state``) are built only when + ``num_contexts > 0`` -- never during decode-graph capture. Reached from + ``InklingAttentionMetadata.prepare()`` via + :meth:`InklingHybridCacheManager.prepare_conv_runtime`, so the + host->device slot write lands outside the captured ``model.forward``. + """ + is_graph = bool(getattr(attn_metadata, "is_cuda_graph", False)) + slots = cache.write_state_indices(list(attn_metadata.request_ids), is_graph) + seq_lens = attn_metadata.seq_lens.tolist() + num_contexts = attn_metadata.num_contexts + state_indices = cache.state_indices + device = state_indices.device + num_ctx_tokens = sum(seq_lens[:num_contexts]) + ctx_indices = state_indices[:num_contexts] if num_contexts else None + gen_indices = ( + state_indices[num_contexts : len(slots)] if num_contexts < len(slots) else None + ) + query_start_loc = has_initial_state = None + if num_contexts: + cu = torch.zeros(num_contexts + 1, dtype=torch.int32, device=device) + cu[1:] = torch.tensor(seq_lens[:num_contexts], dtype=torch.int32, device=device).cumsum( + 0 + ) + query_start_loc = cu + # Fresh prefill carries no prior conv window. This is only correct + # because Inkling defaults ``enable_block_reuse`` off + # (``get_model_defaults``): a reused prefix would need this set per + # request from ``num_cached_tokens_per_seq``, with the preceding + # ``kernel_size - 1`` activations restored into the pool. Do not + # re-enable block reuse without implementing both. + has_initial_state = torch.zeros(num_contexts, dtype=torch.bool, device=device) + return cls( + num_ctx_tokens=num_ctx_tokens, + ctx_indices=ctx_indices, + gen_indices=gen_indices, + query_start_loc=query_start_loc, + has_initial_state=has_initial_state, + ) + + +def _apply_sconv( + sconv: "InklingShortConv", + x: torch.Tensor, + pool_buf: Optional[torch.Tensor], + rt: Optional[InklingConvRuntime], +) -> torch.Tensor: + """Run one short-conv over a (possibly mixed) batch through the state pool. + + ``rt is None`` -> stateless full-sequence causal conv (no pool registered). + Otherwise the context slice seeds ``pool_buf`` (varlen prefill) and the + generation slice updates it in place at ``rt.gen_indices`` (decode), then the + two outputs are concatenated in packed order. ``pool_buf`` is this conv's + ``[max_batch, channels, kernel-1]`` state buffer from + :class:`InklingConvStateCache`. + """ + if rt is None: + return sconv(x) + parts = [] + nctx = rt.num_ctx_tokens + if nctx > 0: + parts.append( + sconv.forward( + x[:nctx], + conv_state=pool_buf, + cache_indices=rt.ctx_indices, + query_start_loc=rt.query_start_loc, + has_initial_state=rt.has_initial_state, + is_decode=False, + ) + ) + if x.shape[0] > nctx: + parts.append( + sconv.forward( + x[nctx:], conv_state=pool_buf, cache_indices=rt.gen_indices, is_decode=True + ) + ) + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0) + + +def _module_excluded_from_quant(model_config: ModelConfig, name: str) -> bool: + """True if ``name`` (or an ancestor) is bf16, not NVFP4. + + This plain-NVFP4 checkpoint lists its bf16 modules in + ``hf_quant_config.json`` ``quantization.exclude_modules`` (read into + ``quant_config.exclude_modules`` by ``from_pretrained``) rather than in + ``per_layer_quant_configs`` (only populated for MIXED_PRECISION checkpoints). + ``QuantConfig.is_module_excluded_from_quantization`` walks the dotted + ancestry, so a listed ``model.llm.layers.5.attn`` covers the qkv/o + projections under it. Used to build attention (all ``.attn`` excluded) and + layer-2 routed experts (``.mlp.experts`` excluded) as bf16. + """ + qc = model_config.quant_config + return ( + qc is not None + and qc.exclude_modules is not None + and qc.is_module_excluded_from_quantization(name) + ) + + +# ---------------------------------------------------------------------------- +# Routing method +# ---------------------------------------------------------------------------- +class InklingMoeRoutingMethod(BaseMoeRoutingMethod): + """Sigmoid gate + additive-bias top-k selection + log-sigmoid renorm. + + The renorm denominator spans the selected routed logits *and* the shared + logits together (``shared_expert_sink``), so this cannot be expressed by the + stock sigmoid/MiniMax routing methods. ``apply`` returns only the routed + ``(topk_ids, topk_weights)`` needed by the fused MoE; the shared gammas come + from the same joint renorm and are recomputed in :class:`InklingMoE` for the + shared-expert branch (see :func:`inkling_joint_renorm`). + """ + + def __init__( + self, + top_k: int, + num_experts: int, + n_shared_experts: int, + callable_gate_bias, + callable_global_scale, + route_scale: float, + ): + super().__init__() + self.top_k = top_k + self.num_experts = num_experts + self.n_shared_experts = n_shared_experts + self._callable_gate_bias = callable_gate_bias + self._callable_global_scale = callable_global_scale + self.route_scale = route_scale + + def apply( + self, router_logits: torch.Tensor, input_ids=None + ) -> tuple[torch.Tensor, torch.Tensor]: + # router_logits: [num_tokens, num_experts + n_shared] in fp32. + routed_w, topk_idx, _ = inkling_joint_renorm( + router_logits.float(), + gate_bias=self._callable_gate_bias(), + global_scale=self._callable_global_scale(), + route_scale=self.route_scale, + top_k=self.top_k, + num_routed=self.num_experts, + n_shared=self.n_shared_experts, + ) + return topk_idx.to(torch.int32), routed_w.to(torch.float32) + + @property + def routing_method_type(self): + # CUTLASS computes the dispatch torch-side via :meth:`apply`, so the + # kernel needs no routing enum of its own. + return RoutingMethodType.Unspecified + + +def inkling_joint_renorm( + router_logits: torch.Tensor, + gate_bias: torch.Tensor, + global_scale: torch.Tensor, + route_scale: float, + top_k: int, + num_routed: int, + n_shared: int, +): + """Exact Inkling router math (fp32). Mirrors HF ``InklingTopkRouter``. + + Returns ``(routed_weights [T, top_k], topk_idx [T, top_k], shared_gammas + [T, n_shared])``. Selection uses ``sigmoid(routed) + bias``; the weights are + a softmax over ``logsigmoid`` of the selected-routed-plus-shared *logits*, + scaled by ``route_scale * global_scale``. + """ + routed_logits = router_logits[..., :num_routed] + shared_logits = router_logits[..., num_routed : num_routed + n_shared] + + scores = routed_logits.sigmoid() + scores_for_choice = scores + gate_bias + topk_idx = torch.topk(scores_for_choice, top_k, dim=-1, sorted=False)[1] + + topk_logits = torch.cat([routed_logits.gather(-1, topk_idx), shared_logits], dim=-1) + topk_log_probs = torch.nn.functional.logsigmoid(topk_logits) + weights = torch.exp(topk_log_probs - torch.logsumexp(topk_log_probs, dim=-1, keepdim=True)) + weights = weights * route_scale * global_scale + + routed_weights = weights[..., :top_k].contiguous() + shared_gammas = weights[..., top_k : top_k + n_shared].contiguous() + return routed_weights, topk_idx, shared_gammas + + +# ---------------------------------------------------------------------------- +# Short convolution (four per layer) +# ---------------------------------------------------------------------------- +class InklingShortConv(nn.Module): + """Causal depthwise short convolution (kernel 4) with an internal residual. + + The weight matches the checkpoint layout ``[channels, 1, kernel]``. At + prefill this runs :func:`causal_conv1d_fn`; at cached decode it runs + :func:`causal_conv1d_update` against the per-request conv state carried by + the state cache manager. ``conv_state`` (and the runtime metadata that + selects the per-request slot) is threaded in by the caller; when it is + ``None`` the module falls back to a self-contained causal convolution over + the provided sequence. + + TP sharding (``tp_shard=True``): the k/v short convs act on the per-rank + slice of the k/v stream produced by the fused qkv projection, so their + channels are sharded by kv-head exactly like that projection. The checkpoint + stores the *full* (unsharded) conv weight, so :meth:`load_weights` slices the + rank's contiguous channel block -- the same pattern as the mamba mixer, which + stores its depthwise conv in a column-parallel ``Linear``. The + post-attention / post-MLP convs run on the full (all-reduced) hidden stream + and are replicated (``tp_shard=False``). + """ + + def __init__(self, channels: int, kernel_size: int, mapping=None, tp_shard: bool = False): + super().__init__() + self.kernel_size = kernel_size + self.tp_size = mapping.tp_size if (mapping is not None and tp_shard) else 1 + self.tp_rank = mapping.tp_rank if (mapping is not None and tp_shard) else 0 + assert channels % self.tp_size == 0, (channels, self.tp_size) + self.channels_full = channels + # Local (this rank's) channel count -- what the forward actually sees. + self.channels = channels // self.tp_size + # Depthwise conv weight, one filter per (local) channel: [channels,1,kernel]. + self.weight = nn.Parameter(torch.empty(self.channels, 1, kernel_size)) + self.register_parameter("bias", None) + + def load_weights(self, weights, allow_partial_loading: bool = False): + """Copy the (full) checkpoint conv weight, slicing this rank's channels. + + The loader routes here (``hasattr(module, 'load_weights')``) with a + one-element list of ``{'weight': [channels_full, 1, kernel]}``. For the + replicated post-attn/post-MLP convs ``tp_size == 1`` and the full tensor + is copied; for the sharded k/v convs the rank's contiguous channel block + is taken (kv-head aligned, matching the fused qkv k/v split). + """ + w = weights[0]["weight"] + if self.tp_size > 1: + w = w.chunk(self.tp_size, dim=0)[self.tp_rank] + self.weight.data.copy_(w[:]) + + def forward( + self, + x: torch.Tensor, + conv_state: Optional[torch.Tensor] = None, + cache_indices: Optional[torch.Tensor] = None, + query_start_loc: Optional[torch.Tensor] = None, + has_initial_state: Optional[torch.Tensor] = None, + is_decode: bool = False, + ) -> torch.Tensor: + """x: [num_tokens, channels]; internal residual ``y = conv(x) + x``. + + The stateless (no-cache) branch runs the conv in fp32 (per the source); + the fused cached branches run in the input dtype (the ``causal_conv1d`` + ops require ``weight.dtype == x.dtype``, so the fp32 conv Parameter is + cast to ``x.dtype`` and ``conv_state`` -- the bf16 state pool -- matches). + Output is cast back to the input dtype. ``conv_state`` is updated in place + by the fused ops. + """ + in_dtype = x.dtype + residual = x + # Fused ops need weight and state in the input dtype (bf16); the fp32 + # conv Parameter is cast here (the stateless branch below uses fp32). + w = self.weight.squeeze(1).to(x.dtype) # [channels, kernel] + if conv_state is not None and is_decode: + # Cached single/short-step decode: [num_tokens, channels] -> op. + # ``causal_conv1d_update`` writes its output IN PLACE into its ``x`` + # argument (and returns that same tensor), so it must be given a + # COPY -- otherwise it clobbers ``residual`` (which aliases ``x``) + # and the internal residual becomes ``conv(x) + conv(x)`` instead of + # ``conv(x) + x``. (The prefill branch is safe: ``transpose(). + # contiguous()`` already copies.) + y = causal_conv1d_update( + x.clone(), + conv_state, + w, + self.bias, + activation=None, + conv_state_indices=cache_indices, + ) + elif conv_state is not None: + # Prefill with cache: varlen [channels, total_tokens]. + xt = x.transpose(0, 1).contiguous() + y = causal_conv1d_fn( + xt, + w, + self.bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_state, + conv_states=conv_state, + activation=None, + ) + y = y.transpose(0, 1).contiguous() + else: + # No cache: self-contained causal depthwise conv over the sequence. + xt = x.float().transpose(0, 1).unsqueeze(0) # [1, channels, T] + y = torch.nn.functional.conv1d( + xt, + self.weight.float(), + bias=None, + padding=self.kernel_size - 1, + groups=self.channels, + ) + y = y[..., : x.shape[0]].squeeze(0).transpose(0, 1) + return (y.to(in_dtype) + residual).to(in_dtype) + + +# ---------------------------------------------------------------------------- +# Attention +# ---------------------------------------------------------------------------- +class InklingAttention(QKNormRoPEAttention): + """RoPE-free attention with per-head q/k RMSNorm, k/v short-conv, and a + learned relative-position bias applied as a Triton ``score_mod``. + + Reuses :class:`QKNormRoPEAttention` for the fused qkv/o projections and + per-head q/k RMSNorm (``skip_rope=True`` gives qk-norm without RoPE), and + owns the extra ``r`` projection, the k/v short convolutions, and the + relative-logit projection. The attention *compute* itself runs through the + Inkling Triton attention path (``attention_backend/inkling/``) + rather than the base backend, because Inkling's learned relative bias is a + per-(query,head,relative-distance) additive ``score_mod`` that no fused, + CUDA-graph-safe TensorRT-LLM backend exposes: + * ``cpp/.../common/attentionOp.cpp`` disables context FMHA for + ``position_embedding_type == kRELATIVE`` (unfused MHA fallback); + * the trtllm-gen decode kernel rejects a relative attention bias; + * FlashInfer has no additive per-token bias hook. + The bias is + precomputed on the torch side as a contiguous ``rel_logits`` aux tensor + ``[num_query_tokens, local_heads, rel_extent]`` (``einsum('thd,de->the', r, + proj)`` with the global-layer ``tau`` folded in), and the Triton prefill / + paged-decode kernels gather+add it: ``bias = rel_logits[q_idx, head, + clamp(q_pos-k_pos, 0, rel_extent-1)]`` where ``0 <= q_pos-k_pos < + rel_extent``. Because ``rel_logits`` is a static-shape tensor (its first dim + equals the batch in the decode phase), the paged-decode kernel captures and + replays cleanly under CUDA graph -- the launch grid ``(batch, heads)`` is + fixed and per-request sequence lengths are read from a GPU tensor. Local + layers apply the sliding window natively inside the kernel + (``window_left = sliding_window_size - 1``); global layers apply the + log-scaling ``tau`` folded into ``rel_logits`` (a no-op below + ``log_scaling_n_floor`` = 128k positions). + + KV read/write goes through ``KVCacheManagerV2`` in the HND paged layout: the + context phase writes new K/V to the cache (for later reuse) and attends over + the contiguous extend tensors; the generation phase writes the one new + token's K/V and attends over the paged cache. ``self.attn`` (the base + backend) is built but unused -- only its runtime-assigned ``local_layer_idx`` + (the KV-cache layer offset) is read here. + """ + + def __init__(self, model_config: ModelConfig[InklingTextConfig], layer_idx: int): + config = model_config.pretrained_config + self.is_local = config.is_local_layer(layer_idx) + head_dim = config.layer_head_dim(layer_idx) + num_heads = config.layer_num_heads(layer_idx) + num_kv_heads = config.layer_num_kv_heads(layer_idx) + self.attention_window_size = config.layer_window(layer_idx) + self.d_rel = config.d_rel + self.rel_extent = config.sliding_window_size if self.is_local else config.rel_extent + self.log_scaling_n_floor = None if self.is_local else config.log_scaling_n_floor + self.log_scaling_alpha = config.log_scaling_alpha + + # Attention (q/k/v/o projections + KV cache) is bf16, not NVFP4: the + # checkpoint excludes ``model.llm.layers.{i}.attn``. The base Attention + # builds qkv_proj/o_proj from ``config.get_quant_config()`` (the global + # NVFP4 config, which packs the input dim to hidden/2 and demands scale + # sidecars the bf16 checkpoint does not have), so hand it a shallow + # ModelConfig copy whose ``quant_config`` is empty for this layer. r_proj + # below is already unquantized (no quant_config passed). + # (ModelConfig.__setattr__ whitelists ``quant_config`` for exactly this + # per-module-quant override, so the shallow copy needs no unfreeze.) + attn_model_config = model_config + if _module_excluded_from_quant(model_config, f"model.llm.layers.{layer_idx}.attn"): + from tensorrt_llm.models.modeling_utils import QuantConfig + + attn_model_config = copy.copy(model_config) + attn_model_config.quant_config = QuantConfig() + + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=num_heads, + num_key_value_heads=num_kv_heads, + max_position_embeddings=config.max_position_embeddings, + bias=False, + # No RoPE: this model overrides forward to run the Inkling Triton + # attention (qk-norm + sconv + relative-bias score_mod) directly, so + # pos_embd_params=None keeps the base from building an unused + # RotaryEmbedding. The base backend ``self.attn`` is still + # constructed but unused for compute (only its runtime-assigned + # ``local_layer_idx`` -- the KV-cache layer offset -- is read). + pos_embd_params=None, + layer_idx=layer_idx, + dtype=config.torch_dtype, + config=attn_model_config, + # q/k are per-head RMS-normalized, so the score scale is 1/head_dim + # rather than 1/sqrt(head_dim). The backend uses + # 1/(sqrt(head_dim) * q_scaling); q_scaling = sqrt(head_dim) yields + # the required 1/head_dim. + q_scaling=float(head_dim) ** 0.5, + skip_rope=True, + fuse_qk_norm_rope=False, + is_qk_norm=True, + ) + # head_dim is uniform (128) across local/global layers and differs from + # hidden_size // num_heads (96), so the base Attention must read it from + # config.head_dim (QKNormRoPEAttention does not accept a head_dim kwarg). + assert self.head_dim == head_dim, (self.head_dim, head_dim) + + # Inkling score scale is 1/head_dim (per-head q/k RMSNorm replaces the + # usual 1/sqrt(head_dim)), applied directly by the Triton kernels. The + # sliding window is applied natively inside the kernel for local layers + # (inclusive radius = window - 1: query p attends to keys [p-(w-1), p]). + self.sm_scale = 1.0 / float(head_dim) + self.window_left = (self.attention_window_size - 1) if self.is_local else -1 + + # Attention-scoped TP. Under attention DP every rank runs the FULL head + # set over its own requests, so the base Attention above built qkv_proj / + # o_proj from an internal ``tp_size=1`` mapping (modules/attention.py). + # The three Inkling-only tensors below (r_proj, k/v sconv) hang off that + # same head/kv-head split, so they must follow the attention TP, not the + # global one -- sharding them by mapping.tp_size while the base keeps + # full heads is a silent shape mismatch, not a slowdown. + tp_size = 1 if model_config.mapping.enable_attention_dp else model_config.mapping.tp_size + # Cross-check against the base rather than trusting two copies of the + # rule to stay in step: if modules/attention.py ever changes how it + # scopes attention TP, fail here at load instead of silently building + # r_proj and the sconvs for a different head count than qkv_proj. + assert self.num_heads == num_heads // tp_size, ( + f"attention TP disagrees with the base Attention: base kept " + f"{self.num_heads} of {num_heads} heads, this rule expects " + f"{num_heads // tp_size} (enable_attention_dp=" + f"{model_config.mapping.enable_attention_dp}, " + f"mapping.tp_size={model_config.mapping.tp_size})" + ) + # r projection: per-head relative states (num_heads * d_rel), sharded by + # head like q. Output is not gathered (consumed locally to build bias). + # Under attention DP it is replicated: no mapping / no TP mode, matching + # how DeepSeek-V3 builds its non-expert Linears under ADP. + self.r_proj = Linear( + config.hidden_size, + num_heads * self.d_rel, + bias=False, + dtype=config.torch_dtype, + mapping=None if model_config.mapping.enable_attention_dp else model_config.mapping, + tensor_parallel_mode=( + None if model_config.mapping.enable_attention_dp else TensorParallelMode.COLUMN + ), + gather_output=False, + ) + # Learned relative-logit profiles, replicated across TP ranks. The + # profile length is per-layer: local layers store only the + # sliding-window extent (512), global layers the full rel_extent (1024) + # -- so the parameter must use ``self.rel_extent``, not the global + # ``config.rel_extent`` (mismatch here is the 1024-vs-512 load crash). + self.rel_logits_proj = nn.Parameter(torch.empty(self.d_rel, self.rel_extent)) + # k/v short convs act on the k/v stream from the fused qkv projection, + # so they are sharded by kv-head like that projection. Pass the FULL + # channel count and let InklingShortConv slice this rank's block at load. + # Under attention DP that projection keeps every kv head, so the convs + # must keep every channel: tp_shard follows the attention TP. + full_kv_dim = num_kv_heads * head_dim + sconv_tp_shard = not model_config.mapping.enable_attention_dp + self.k_sconv = InklingShortConv( + full_kv_dim, + config.sconv_kernel_size, + mapping=model_config.mapping, + tp_shard=sconv_tp_shard, + ) + self.v_sconv = InklingShortConv( + full_kv_dim, + config.sconv_kernel_size, + mapping=model_config.mapping, + tp_shard=sconv_tp_shard, + ) + self.local_num_heads = num_heads // tp_size + # Stable GPU buffers for the CUDA-graph-safe runtime decode metadata, + # refreshed eagerly (before capture/replay) by the model engine via + + def _project(self, hidden_states, conv_pool_kv=None, conv_rt=None): + """Fused qkv projection -> split -> k/v short-conv -> per-head qk RMSNorm. + + Returns ``(q, k, v)`` shaped ``[T, local_heads, head_dim]`` / + ``[T, local_kv_heads, head_dim]``. With ``conv_pool_kv=(pool_k, pool_v)`` + + ``conv_rt`` the k/v short-convs run through the runtime state pool + (seed for context tokens, in-place update at the per-request slots for + generation tokens; fused ops, CUDA-graph safe, mixed-batch capable); + without them they run the stateless full-sequence causal conv. + """ + D = self.head_dim + num_tokens = hidden_states.shape[0] + qkv = self.qkv_proj(hidden_states) + q, k, v = self.split_qkv(qkv, None, None) + # k/v short convolution before the q/k norm (source order). + if conv_pool_kv is not None: + pool_k, pool_v = conv_pool_kv + k = _apply_sconv(self.k_sconv, k, pool_k, conv_rt) + v = _apply_sconv(self.v_sconv, v, pool_v, conv_rt) + else: + k = self.k_sconv(k) + v = self.v_sconv(v) + q, k = self.apply_qk_norm(q, k) + nh = self.q_size // D + nkv = self.kv_size // D + return ( + q.view(num_tokens, nh, D), + k.view(num_tokens, nkv, D), + v.view(num_tokens, nkv, D), + ) + + def _build_rel_logits( + self, hidden_states: torch.Tensor, position_ids: Optional[torch.Tensor] + ) -> torch.Tensor: + """Contiguous relative-bias aux tensor ``[T, local_heads, rel_extent]``. + + ``rel_logits[t, h, e] = sum_d r[t, h, d] * proj[d, e]`` (fp32), mirroring + the reference ``InklingRelativeLogits``. For global layers + the log-scaling ``tau`` (a no-op below ``log_scaling_n_floor`` = 128k) is + folded in per query token. The + Triton kernels index this by ``clamp(q_pos-k_pos, 0, rel_extent-1)`` and + zero it outside ``[0, rel_extent)`` -- the exact source score_mod. + """ + r = self.r_proj(hidden_states).view(-1, self.local_num_heads, self.d_rel) + rel = torch.einsum( + "thd,de->the", r.float(), self.rel_logits_proj.float() + ) # [T, H, rel_extent] + if self.log_scaling_n_floor is not None and position_ids is not None: + pos = position_ids.reshape(-1).float() + tau = 1.0 + self.log_scaling_alpha * torch.log( + ((pos + 1.0) / self.log_scaling_n_floor).clamp(min=1.0) + ) + rel = rel * tau[:, None, None] + return rel.contiguous() + + def _attention(self, q, k, v, rel_logits, attn_metadata, *, allow_mixed=False): + """Dispatch prefill / decode over the paged cache, supporting mixed + context+generation batches. + + The runtime packs context requests first (each with its full new-token + span) then one-token generation requests (``seq_lens == 1``). We slice + the packed q/k/v/rel_logits + per-request metadata at that boundary and + run the context slice through the prefill kernel and the generation + slice through the paged-decode kernel, concatenating the outputs. Pure + context (``num_contexts == num_seqs``) and pure generation + (``num_contexts == 0``) fall out as the single-slice cases. + """ + # ``KVCacheManagerV2.get_buffers`` / ``get_batch_cache_indices`` take the + # GLOBAL layer index and map it through ``layer_offsets`` themselves + # (identity for single-node TP-only, the pp-local offset under PP). Use + # ``self.layer_idx`` (the model's global decoder layer index) directly: + # the base backend's ``self.attn.local_layer_idx`` is only primed inside + # the base attention forward, which Inkling bypasses, so it stays ``None``. + cache_layer = self.layer_idx + kv = attn_metadata.kv_cache_manager.get_buffers(cache_layer, kv_layout="HND") + # kv: [num_pages, 2, num_kv_heads, page_size, head_dim] + k_cache, v_cache = kv[:, 0], kv[:, 1] + page_size = kv.shape[3] + mgr = attn_metadata.kv_cache_manager + request_ids = attn_metadata.request_ids + num_cached = attn_metadata.kv_cache_params.num_cached_tokens_per_seq + seq_lens = attn_metadata.seq_lens.tolist() + num_contexts = attn_metadata.num_contexts + num_seqs = len(seq_lens) + ctx_tokens = sum(seq_lens[:num_contexts]) + + # A mixed context+generation batch needs ``_project`` to apply the + # prefill short-conv to the context tokens and the decode short-conv to + # the generation tokens -- which only the per-request short-conv state + # pool (``InklingConvStateCache``, the ``conv_rt`` runtime path) does + # correctly. The stateless path convolves one token group at a time, so a + # mixed batch there would convolve across the context/generation + # boundary; refuse it explicitly unless the pool path is active + # (``allow_mixed``). + if 0 < num_contexts < num_seqs and not allow_mixed: + raise NotImplementedError( + "InklingAttention: mixed context+generation batch needs the " + "short-conv state pool (pass conv_cache/conv_rt); the stateless " + "short-conv path cannot mix a batch" + ) + + outs = [] + if num_contexts > 0: + outs.append( + self._run_context( + q[:ctx_tokens], + k[:ctx_tokens], + v[:ctx_tokens], + rel_logits[:ctx_tokens], + seq_lens[:num_contexts], + num_cached[:num_contexts], + request_ids[:num_contexts], + mgr, + cache_layer, + k_cache, + v_cache, + page_size, + ) + ) + if num_contexts < num_seqs: + outs.append( + self._run_generation( + q[ctx_tokens:], + k[ctx_tokens:], + v[ctx_tokens:], + rel_logits[ctx_tokens:], + num_cached[num_contexts:], + request_ids[num_contexts:], + mgr, + cache_layer, + k_cache, + v_cache, + page_size, + attn_metadata, + ) + ) + return outs[0] if len(outs) == 1 else torch.cat(outs, dim=0) + + def _run_context( + self, + q, + k, + v, + rel_logits, + seq_lens, + num_cached, + request_ids, + mgr, + cache_layer, + k_cache, + v_cache, + page_size, + ): + device = q.device + # Persist new K/V to the paged cache for later generation reuse. + block_ids = mgr.get_batch_cache_indices(request_ids, cache_layer) + off = 0 + for i, sl in enumerate(seq_lens): + write_kv_cache_hnd( + k_cache, + v_cache, + k[off : off + sl], + v[off : off + sl], + block_ids[i], + int(num_cached[i]), + page_size, + ) + off += sl + cu = torch.zeros(len(seq_lens) + 1, dtype=torch.int32, device=device) + cu[1:] = torch.tensor(seq_lens, dtype=torch.int32, device=device).cumsum(0) + max_seqlen = max(seq_lens) + return inkling_prefill_attention( + q, k, v, cu, max_seqlen, self.sm_scale, rel_logits, self.rel_extent, self.window_left + ) + + def _run_generation( + self, + q, + k, + v, + rel_logits, + num_cached, + request_ids, + mgr, + cache_layer, + k_cache, + v_cache, + page_size, + attn_metadata, + ): + device = q.device + # --- Runtime CUDA-graph-safe path. --------------------------------- + # ``InklingAttentionMetadata.prepare()`` published this batch's decode + # metadata into the metadata object's stable GPU buffers, so the + # captured forward performs ZERO host->device copy: it reads the + # stable ``ink_seq_lens``/``ink_page_table`` buffers and persists the new token's + # K/V into the paged cache with an in-graph GPU scatter whose (page, + # offset) indices are derived on-GPU from those buffers. This replaces the + # host ``write_kv_cache_hnd`` loop + ``torch.tensor(..., device=cuda)`` + # build that raised ``Cannot copy between CPU and CUDA tensors during CUDA + # graph capture``. Padding rows carry their own (dummy) registered request + # slots -- ``attn_metadata.request_ids`` is padded after ``prepare()`` -- + # so the scatter never corrupts a real request's page 0. + num_req = q.shape[0] + if getattr(attn_metadata, "ink_num_gen", 0) == num_req: + sl = attn_metadata.ink_seq_lens[:num_req] + pt = attn_metadata.ink_page_table[cache_layer][:num_req] + pos = (sl - 1).long() # write slot = total_kv_len - 1 = num_cached + page_row = torch.div(pos, page_size, rounding_mode="floor") + offs = pos - page_row * page_size + pages = pt.gather(1, page_row.unsqueeze(1)).squeeze(1).long() + # HND paged cache: [num_pages, num_kv_heads, page_size, head_dim]; + # paired advanced indices (pages, offs) select one (page, slot) per + # request -> [num_req, num_kv_heads, head_dim], matching new k/v. + k_cache[pages, :, offs, :] = k.to(k_cache.dtype) + v_cache[pages, :, offs, :] = v.to(v_cache.dtype) + return inkling_decode_attention( + q, + k_cache, + v_cache, + sl, + pt, + page_size, + self.sm_scale, + rel_logits, + self.rel_extent, + self.window_left, + ) + # Eager fallback (never captured): the decode metadata was not published, + # so build it here from the host block table, like the context path. + # + # Under CUDA graph this path is ILLEGAL -- the torch.tensor() below is a + # host->device copy from unpinned memory -- and the error it raises names + # a tensor, not the cause. The realistic way to arrive here is an + # ``attn_backend`` override: get_model_defaults selects "INKLING" so that + # attn_metadata carries the decode publish, but an explicit + # ``attn_backend: TRTLLM`` in --extra_llm_api_options wins the model + # defaults deep-merge and silently takes it away. Say that instead. + if getattr(attn_metadata, "is_cuda_graph", False): + raise RuntimeError( + "Inkling decode metadata was not published for a CUDA-graph " + f"batch (ink_num_gen=" + f"{getattr(attn_metadata, 'ink_num_gen', None)}, expected " + f"{num_req}); attn_metadata is " + f"{type(attn_metadata).__name__}, not InklingAttentionMetadata. " + "Inkling requires attn_backend='INKLING'; remove any " + "attn_backend override from --extra_llm_api_options / " + "LLM(attn_backend=...) and let the model default apply." + ) + num_req = len(request_ids) + block_ids = mgr.get_batch_cache_indices(request_ids, cache_layer) + for i in range(num_req): + write_kv_cache_hnd( + k_cache, + v_cache, + k[i : i + 1], + v[i : i + 1], + block_ids[i], + int(num_cached[i]), + page_size, + ) + total = [int(num_cached[i]) + 1 for i in range(num_req)] + decode_seq_lens = torch.tensor(total, dtype=torch.int32, device=device) + max_pages = max(len(b) for b in block_ids) + decode_page_table = build_page_table(block_ids, max_pages, device) + return inkling_decode_attention( + q, + k_cache, + v_cache, + decode_seq_lens, + decode_page_table, + page_size, + self.sm_scale, + rel_logits, + self.rel_extent, + self.window_left, + ) + + def forward( + self, + position_ids: Optional[torch.IntTensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + *, + conv_pool_kv=None, + conv_rt=None, + **kwargs, + ): + """Inkling attention through the Triton score_mod path. + + ``conv_pool_kv=(pool_k, pool_v)`` + ``conv_rt`` drive the k/v short-convs + through the runtime state pool (seed on context, in-place update at the + per-request slots on generation, mixed-batch capable, CUDA-graph safe); + without them the short-convs run stateless over the whole sequence. + """ + num_tokens = hidden_states.shape[0] + # The pre-attention RMSNorm can emit fp32 (the residual-stream norm + # path), but the attention/r projections are bf16 (``.attn`` is excluded + # from NVFP4). Cast once so the decoder-layer forward is robust to the + # norm's output dtype. + hidden_states = hidden_states.to(self.qkv_proj.weight.dtype) + q, k, v = self._project(hidden_states, conv_pool_kv, conv_rt) + rel_logits = self._build_rel_logits(hidden_states, position_ids) + attn_out = self._attention( + q, k, v, rel_logits, attn_metadata, allow_mixed=conv_rt is not None + ) + attn_out = attn_out.reshape(num_tokens, self.q_size) + return self.o_proj(attn_out) + + +# ---------------------------------------------------------------------------- +# Dense MLP (layers 0, 1) and MoE (layers 2..65) +# ---------------------------------------------------------------------------- +class InklingDenseMLP(nn.Module): + """SwiGLU MLP with a learned scalar ``global_scale`` (layers 0, 1). + + Fused gate+up (``w13_dn``) column-parallel, down (``w2_md``) row-parallel. + """ + + def __init__(self, model_config: ModelConfig[InklingTextConfig]): + super().__init__() + config = model_config.pretrained_config + inter = config.dense_intermediate_size + # Under attention DP the dense MLP goes data-parallel too: each rank + # holds the full weight and runs it over its OWN tokens. Keeping the + # column/row split here would be a correctness bug, not just a + # different partitioning -- the row-parallel down_proj all-reduces its + # partial sum across the TP group, and under ADP the peers' partials + # belong to DIFFERENT requests, so the reduce would add unrelated + # tokens together. This mirrors DeepSeek-V3, whose + # ``_compute_mlp_tp_size`` returns 1 under ADP for the same reason: + # only the routed experts stay sharded, and FusedMoE re-joins the ranks + # explicitly via all_rank_num_tokens. + dp = model_config.mapping.enable_attention_dp + mlp_mapping = None if dp else model_config.mapping + self.gate_up_proj = Linear( + config.hidden_size, + 2 * inter, + bias=False, + dtype=config.torch_dtype, + mapping=mlp_mapping, + tensor_parallel_mode=None if dp else TensorParallelMode.COLUMN, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_GATE_UP_LINEAR + ), + ) + self.down_proj = Linear( + inter, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + mapping=mlp_mapping, + tensor_parallel_mode=None if dp else TensorParallelMode.ROW, + ) + self.global_scale = nn.Parameter(torch.ones(1)) + self.act_fn = torch.nn.functional.silu + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + # ``global_scale`` is an fp32 scalar Parameter; multiplying promotes the + # output to fp32. Cast back to the input dtype so the bf16 residual + # stream (and the next layer's bf16 projections) stay bf16. + out = self.down_proj(self.act_fn(gate) * up) * self.global_scale + return out.to(x.dtype) + + +class InklingGate(nn.Module): + """fp32 router: logits over 256 routed + 2 shared experts, plus the additive + selection bias and the learned global scale. Feeds + :class:`InklingMoeRoutingMethod`. + """ + + def __init__(self, config: InklingTextConfig): + super().__init__() + self.num_routed = config.n_routed_experts + self.n_shared = config.n_shared_experts + self.top_k = config.num_experts_per_tok + self.route_scale = config.route_scale + n_total = self.num_routed + self.n_shared + self.weight = nn.Parameter(torch.empty(n_total, config.hidden_size, dtype=torch.float32)) + self.bias = nn.Parameter(torch.empty(self.num_routed, dtype=torch.float32)) + self.global_scale = nn.Parameter(torch.ones(1, dtype=torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.linear(hidden_states.float(), self.weight) + + @property + def routing_method(self) -> InklingMoeRoutingMethod: + return InklingMoeRoutingMethod( + top_k=self.top_k, + num_experts=self.num_routed, + n_shared_experts=self.n_shared, + callable_gate_bias=lambda: self.bias, + callable_global_scale=lambda: self.global_scale, + route_scale=self.route_scale, + ) + + +class InklingSharedExperts(nn.Module): + """Two shared SwiGLU experts, each weighted by a per-token gamma and summed. + + Reference: HF ``InklingSharedExperts`` (batched 2-expert SwiGLU, fp32 sum). + """ + + def __init__(self, config: InklingTextConfig): + super().__init__() + self.n_shared = config.n_shared_experts + inter = config.intermediate_size + hidden = config.hidden_size + # [n_shared, 2*inter, hidden] fused gate+up; [n_shared, hidden, inter] down. + # Must be created in the model dtype (bf16): the shared experts run as raw + # bmms against the bf16 hidden stream, so an untyped (default-fp32) param + # dtype-mismatches the bmm ("expected BFloat16 but found Float"). The + # checkpoint stores these bf16 (shared_experts is in exclude_modules). + self.shared_w13 = nn.Parameter( + torch.empty(self.n_shared, 2 * inter, hidden, dtype=config.torch_dtype) + ) + self.shared_w2 = nn.Parameter( + torch.empty(self.n_shared, hidden, inter, dtype=config.torch_dtype) + ) + self.act_fn = torch.nn.functional.silu + + def forward(self, hidden_states: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + # hidden_states: [T, hidden] (bf16); gammas: [T, n_shared] fp32 (from the + # joint renorm). Keep both bmms in the activation dtype and apply the + # per-token gamma in fp32 AFTER the (linear) down projection, where it + # commutes: gamma * (act @ w2) == (act * gamma) @ w2. This avoids + # upcasting the down-proj bmm's LHS to fp32 (which mismatches the bf16 + # shared_w2 -- an "expected BFloat16 but found Float" bmm error) while + # keeping gamma at full fp32 precision and matching the source fp32 + # gamma-weighted sum. + x = hidden_states.unsqueeze(0).expand(self.n_shared, -1, -1) + gate_up = torch.bmm(x, self.shared_w13.transpose(1, 2)) + # ``shared_w13`` loads RAW: gate/up are Inkling-INTERLEAVED [g0,u0,...] + # along its 2*inter output dim, so the + # bmm output channels are interleaved -- gate = even, up = odd. A + # contiguous chunk(2) here would pair the wrong channels (silu(mix)*mix). + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] + activated = self.act_fn(gate) * up + out = torch.bmm(activated, self.shared_w2.transpose(1, 2)) # [S, T, hidden] + out = out.float() * gammas.transpose(0, 1).unsqueeze(-1).float() + return out.sum(dim=0).to(hidden_states.dtype) + + +class InklingMoE(nn.Module): + """Router + routed experts (fused MoE) + two shared experts. + + Routed experts run through :func:`create_moe` (NVFP4 for layers 3..65, bf16 + for layer 2 via a per-layer quant override). Shared experts and the router + stay bf16/fp32. The routed output already reduces over the top-6 experts; the + gamma-weighted shared output is added on top (source ``h + shared``). + """ + + def __init__(self, model_config: ModelConfig[InklingTextConfig], layer_idx: int): + super().__init__() + config = model_config.pretrained_config + self.gate = InklingGate(config) + self.num_routed = config.n_routed_experts + self.n_shared = config.n_shared_experts + self.top_k = config.num_experts_per_tok + self.route_scale = config.route_scale + + experts_quant_config = self._experts_quant_config(model_config, layer_idx) + # reduce_results=True: all-reduce the routed-expert output across the TP + # group. Under TP each rank holds a shard of the 256 experts and produces + # only a PARTIAL routed sum, so the full routed output is the sum across + # ranks. Without this all-reduce the TP=4 runtime adds a per-rank partial + # routed output to the (replicated, full) shared-expert output and the + # whole model produces garbage from the first token (the dense layers 0/1 + # are correct because their row-parallel down_proj already all-reduces). + # The shared experts stay replicated and are added AFTER this reduce + # (full + full), so they are not double-counted. + self.experts = create_moe( + routing_method=self.gate.routing_method, + num_experts=self.num_routed, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + dtype=config.torch_dtype, + reduce_results=True, + model_config=model_config, + override_quant_config=experts_quant_config, + layer_idx=layer_idx, + ) + self.shared_experts = InklingSharedExperts(config) + + @staticmethod + def _experts_quant_config(model_config: ModelConfig, layer_idx: int): + """Per-layer expert quant: NVFP4 unless the checkpoint excludes it. + + The checkpoint lists its bf16 modules in ``hf_quant_config.json`` + ``quantization.exclude_modules`` (read into + ``quant_config.exclude_modules`` by ``from_pretrained``). Layer-2 routed + experts are excluded (bf16 MoE) while layers 3..65 routed experts are + NVFP4. ``quant_config_dict`` / ``per_layer_quant_configs`` are only + populated for MIXED_PRECISION checkpoints, so for this plain-NVFP4 + checkpoint the authoritative per-layer signal is ``exclude_modules``. + Return an empty (no-quant) ``QuantConfig`` for an excluded expert module + so ``create_moe`` builds an unquantized bf16 MoE; otherwise the NVFP4 + base config. + """ + if _module_excluded_from_quant(model_config, f"model.llm.layers.{layer_idx}.mlp.experts"): + from tensorrt_llm.models.modeling_utils import QuantConfig + + return QuantConfig() + return model_config.quant_config + + def forward( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> torch.Tensor: + """Routed + shared experts. + + ``all_rank_num_tokens`` is the per-rank token count this step, taken + from ``attn_metadata``. ``FusedMoE`` sets ``use_dp`` from + ``mapping.enable_attention_dp`` and needs the list to pad and gather + across ranks; without it a DP or EP-with-DP layout cannot know how much + each peer contributed. ``None`` is the non-DP case and leaves the + expert call exactly as it was. + """ + router_logits = self.gate(hidden_states) # [T, 258] fp32 + routed = self.experts( + hidden_states, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + _, _, shared_gammas = inkling_joint_renorm( + router_logits, + gate_bias=self.gate.bias, + global_scale=self.gate.global_scale, + route_scale=self.route_scale, + top_k=self.top_k, + num_routed=self.num_routed, + n_shared=self.n_shared, + ) + shared = self.shared_experts(hidden_states, shared_gammas) + # Keep the bf16 residual-stream dtype (fp32 scales in the routed/shared + # paths can promote the sum) so the next layer's projections stay bf16. + return (routed + shared).to(hidden_states.dtype) + + +# ---------------------------------------------------------------------------- +# Decoder layer / model / causal LM +# ---------------------------------------------------------------------------- +class InklingDecoderLayer(nn.Module): + """Pre-norm attention + MLP, each followed by a short-conv with an internal + residual, then the residual add (HF ``InklingDecoderLayer`` order). + """ + + def __init__(self, model_config: ModelConfig[InklingTextConfig], layer_idx: int): + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + self.attn_norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.attn = InklingAttention(model_config, layer_idx) + self.attn_sconv = InklingShortConv(config.hidden_size, config.sconv_kernel_size) + self.mlp_norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + if config.is_dense_layer(layer_idx): + self.mlp = InklingDenseMLP(model_config) + else: + self.mlp = InklingMoE(model_config, layer_idx) + self.mlp_sconv = InklingShortConv(config.hidden_size, config.sconv_kernel_size) + + def _run_mlp( + self, + hidden_states: torch.Tensor, + all_rank_num_tokens: Optional[List[int]], + ) -> torch.Tensor: + """Dense layers 0/1 take only the activations; MoE layers also take the + per-rank token counts the fused kernel needs to gather across ranks.""" + if isinstance(self.mlp, InklingMoE): + return self.mlp(hidden_states, all_rank_num_tokens=all_rank_num_tokens) + return self.mlp(hidden_states) + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + *, + conv_state: Optional[InklingConvState] = None, + conv_rt: Optional[InklingConvRuntime] = None, + all_rank_num_tokens: Optional[List[int]] = None, + **kwargs, + ) -> torch.Tensor: + """Pre-norm attention + MLP, each followed by a short-conv (internal + residual), then the residual add. + + With ``conv_rt`` given, ``conv_state`` holds this layer's four + ``[max_batch, C, K-1]`` pool buffers + (:meth:`InklingConvStateCache.layer_state`) and each short-conv seeds the + pool for context tokens and updates it in place at the per-request slots + for generation tokens (fused ops, mixed-batch + CUDA-graph safe). + Without it the short-convs run stateless over the whole sequence. + """ + if conv_rt is None: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states = self.attn(position_ids, hidden_states, attn_metadata) + hidden_states = self.attn_sconv(hidden_states) # internal residual + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self._run_mlp(hidden_states, all_rank_num_tokens) + hidden_states = self.mlp_sconv(hidden_states) # internal residual + return residual + hidden_states + + # --- Runtime state-pool path (prefill-seed / decode / mixed). --- + residual = hidden_states + h = self.attn_norm(hidden_states) + h = self.attn( + position_ids, + h, + attn_metadata, + conv_pool_kv=(conv_state.k, conv_state.v), + conv_rt=conv_rt, + **kwargs, + ) + h = residual + _apply_sconv(self.attn_sconv, h, conv_state.attn, conv_rt) + + residual = h + hm = self._run_mlp(self.mlp_norm(h), all_rank_num_tokens) + return residual + _apply_sconv(self.mlp_sconv, hm, conv_state.mlp, conv_rt) + + +class InklingModel(DecoderModel): + """The Inkling text decoder stack. ``embed_norm`` folds onto the token + embeddings before the layers (``use_embed_norm``).""" + + def __init__(self, model_config: ModelConfig[InklingTextConfig]): + super().__init__(model_config) + config = model_config.pretrained_config + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + self.embed_norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.layers = nn.ModuleList( + [InklingDecoderLayer(model_config, i) for i in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + inputs_embeds_prenormed: bool = False, + **kwargs, + ) -> torch.Tensor: + """Decoder stack. The runtime short-conv state pool and this forward's + context/generation split come from ``attn_metadata`` -- published by + ``InklingAttentionMetadata.prepare()`` from the cache manager, which owns + the pool. Each layer reads its own four ``[max_batch, C, K-1]`` buffers + and the shared split, so the four short-convs of every layer carry + per-request state across decode steps exactly like the paged KV cache. + A metadata without them (no conv-capable cache manager) keeps the + stateless conv. + + ``inputs_embeds_prenormed``: on the multimodal path the wrapper has + ALREADY applied ``embed_norm`` to the text embeddings and scattered the + RAW tower rows in AFTER the norm, so the fused stream must NOT be + re-normed here -- the tower rows carry their own final norm, and pushing + them through an extra RMSNorm would corrupt them so the decoder could not + read the media. Text-only callers pass raw ``inputs_embeds`` and keep the + norm (default ``False``).""" + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + conv_cache = getattr(attn_metadata, "ink_conv_cache", None) + conv_rt = getattr(attn_metadata, "ink_conv_rt", None) + # Per-rank token counts for this step. The model engine fills this on + # attn_metadata only when attention DP is on; FusedMoE reads it to pad + # and gather across ranks. None everywhere else, which is the + # single-rank / pure-TP path and behaves exactly as before. + all_rank_num_tokens = getattr(attn_metadata, "all_rank_num_tokens", None) + hidden_states = inputs_embeds if inputs_embeds_prenormed else self.embed_norm(inputs_embeds) + for i, layer in enumerate(self.layers): + layer_state = conv_cache.layer_state(i) if conv_cache is not None else None + hidden_states = layer( + position_ids, + hidden_states, + attn_metadata, + conv_state=layer_state, + conv_rt=conv_rt, + all_rank_num_tokens=all_rank_num_tokens, + ) + return self.norm(hidden_states) + + +class InklingForCausalLM(DecoderModelForCausalLM[InklingModel, InklingTextConfig]): + """Text CausalLM: muP logit scaling + unpadded-vocab slice. + + ``embed`` and ``unembed`` are separate checkpoint tensors (never tied). The + ``LMHead`` is built at the unpadded vocab size so its forward slices off the + padding automatically; hidden states are divided by + ``logits_mup_width_multiplier`` before the head (accuracy-critical). + """ + + def __init__(self, model_config: ModelConfig[InklingTextConfig]): + config = model_config.pretrained_config + self.mup_multiplier = float(config.logits_mup_width_multiplier) + super().__init__( + InklingModel(model_config), + config=model_config, + hidden_size=config.hidden_size, + vocab_size=config.unpadded_vocab_size, + ) + self._assert_inkling_attn_backend(model_config) + self._assert_inkling_moe_parallel(model_config) + self._apply_allreduce_strategy() + + @staticmethod + def _assert_inkling_attn_backend(model_config) -> None: + """Fail at load if the Inkling attention backend was overridden. + + ``get_model_defaults`` selects ``attn_backend='INKLING'`` because + ``InklingAttentionMetadata`` is what publishes the decode seq_lens and + page table into CUDA-graph-stable buffers. Model defaults are a + deep-merge in which an explicit user value wins, so an + ``attn_backend: TRTLLM`` left in ``--extra_llm_api_options`` -- a very + easy thing to carry over from another model's serve config -- silently + removes that publish. The run then dies deep in CUDA-graph capture with + "Cannot copy between CPU and CUDA tensors", which names neither Inkling + nor the setting responsible. + """ + backend = getattr(model_config, "attn_backend", None) + if backend is not None and str(backend).upper() != "INKLING": + raise ValueError( + f"Inkling requires attn_backend='INKLING' (got {backend!r}). " + "The Triton decode kernel reads its per-step seq_lens and page " + "table from InklingAttentionMetadata, which only the INKLING " + "backend supplies. Remove the attn_backend override from " + "--extra_llm_api_options / LLM(attn_backend=...) so the model " + "default applies." + ) + + @staticmethod + def _assert_inkling_moe_parallel(model_config) -> None: + """Reject an expert-parallel layout the MoE backend cannot serve. + + Inkling's routed experts go through the generic ``create_moe`` factory, + so expert parallelism needs no Inkling-specific code: ``Mapping`` + derives ``moe_ep_size`` / ``moe_tp_size``, ``FusedMoE`` slices the 256 + experts with ``_compute_ep_partition``, and CutlassFusedMoE remaps the + NVFP4 per-expert scales onto the local slice. What it does NOT have is + a check that the requested split is one the backend supports. + + ``FusedMoE._supports_non_divisible_ep`` is opt-in and the CUTLASS + backend -- the only routed-expert backend Inkling ships -- does not opt + in, so a non-divisible ``moe_expert_parallel_size`` fails somewhere + inside expert-slot bookkeeping rather than at load. 256 divides evenly + by every power of two, so this only bites on values like 3, 5 or 6. + + Note this deliberately does NOT constrain moe_tp_size: with + ``moe_ep_size = 1`` (the default) the experts are TP-sharded, which is + what every Inkling accuracy run to date measured. + """ + mapping = getattr(model_config, "mapping", None) + if mapping is None: + return + ep_size = getattr(mapping, "moe_ep_size", 1) or 1 + if ep_size <= 1: + return + config = model_config.pretrained_config + config = getattr(config, "text_config", config) + num_experts = getattr(config, "n_routed_experts", None) + if num_experts is None: + return + # Order matters: more ranks than experts is the more fundamental + # problem and subsumes non-divisibility, so report it first. Checking + # divisibility first would tell a user with 8 ranks and 4 experts to + # "pick a divisor of 4", which is not the advice they need. + if num_experts < ep_size: + raise ValueError( + f"moe_expert_parallel_size={ep_size} exceeds Inkling's " + f"{num_experts} routed experts; ranks with zero experts are " + f"not supported by any MoE backend." + ) + if num_experts % ep_size != 0: + raise ValueError( + f"Inkling has {num_experts} routed experts, which " + f"moe_expert_parallel_size={ep_size} does not divide evenly. " + f"The CUTLASS MoE backend does not opt into non-divisible " + f"expert parallelism, so the uneven split would fail inside " + f"expert-slot bookkeeping instead of here. Pick a divisor of " + f"{num_experts}." + ) + # Measured, not theoretical. On 4 GPUs, against the golden GSM8K run: + # + # ep 1 / 2, cuda_graph on acc 0.9667, zero score flips + # ep 4, cuda_graph OFF acc 0.9667, zero score flips + # ep 4, cuda_graph on SIGSEGV in warmup, all four ranks + # + # So the expert split itself is sound at every size tried, including + # moe_tp_size 1 -- pure EP reproduces the TP-only result per item. What + # breaks is ep_size 4 together with CUDA-graph capture; changing + # max_batch_size / max_num_tokens does not move it, which rules out the + # expert GEMM shape. Root cause not yet found. + # + # Ruled out by experiment, so nobody repeats them: + # * the expert GEMM shape -- varying max_batch_size / max_num_tokens + # does not move the crash + # * the ONESHOT all-reduce pin that works around the captured + # symmetric-all-reduce defect -- letting the framework pick the + # strategy (AUTO) instead still crashes, so the two are unrelated + # * MoE workspaces not being re-allocated between + # MoERunner.clear_all_workspaces() and capture -- enabling the + # autotuner, which re-allocates them, still crashes + # * alltoall -- CutlassFusedMoE.enable_alltoall is always False, so + # ep 2 and ep 4 use the same collective path + # + # Reject only the combination that was observed to crash, and point at + # the configuration that works, rather than removing a usable layout. + moe_tp_size = getattr(mapping, "moe_tp_size", None) + use_cuda_graph = getattr(model_config, "use_cuda_graph", False) + if moe_tp_size is not None and moe_tp_size < 2 and use_cuda_graph: + raise ValueError( + f"moe_expert_parallel_size={ep_size} leaves moe_tp_size=" + f"{moe_tp_size}, which segfaults during CUDA-graph capture " + f"for Inkling. The same layout runs correctly with CUDA " + f"graphs disabled (cuda_graph_config=None), reproducing the " + f"TP-only accuracy per item. Either disable CUDA graphs or " + f"use moe_expert_parallel_size <= {max(1, ep_size // 2)}." + ) + + def _apply_allreduce_strategy(self) -> None: + """Keep Inkling's all-reduces off the NCCL_SYMMETRIC tactic. + + 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. Inkling meets 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. + + The all-reduces that hit this are built by generic modules (attention + ``o_proj``, MoE ``down_proj``), not by this file, so the strategy cannot + be passed at construction without editing shared code. Rebuilding each + ``AllReduce`` after the fact keeps the whole mitigation model-local. + Pinning ONESHOT also drops the window requirement -- ``AllReduce`` only + takes an NCCL window under NCCL_SYMMETRIC/NCCL/AUTO -- so two of the five + trigger conditions go away, not just one. + + Costs symmetric on every Inkling all-reduce, eager included: roughly a + third of captured decode all-reduces pick it today. + """ + for mod in self.modules(): + old = getattr(mod, "all_reduce", None) + # ``None`` means the module reduces nothing (no TP, or DP handles it); + # giving it an AllReduce would add a collective, not remove a tactic. + if not isinstance(old, AllReduce): + continue + # Carry the module's own mapping and dtype over so the rebuilt + # instance differs from the original in strategy alone. + mod.all_reduce = AllReduce( + mapping=old.mapping, + strategy=AllReduceStrategy.ONESHOT, + dtype=getattr(mod, "dtype", None), + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + inputs_embeds_prenormed: bool = False, + **kwargs, + ) -> torch.Tensor: + # The short-conv state pool is owned by InklingHybridCacheManager, so it + # shares the KV cache's request lifetime and reaches the decoder through + # attn_metadata (published by InklingAttentionMetadata.prepare(), outside + # the captured region). No conv kwargs on the main path, and no + # ResourceManager lookup from inside forward. + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + inputs_embeds_prenormed=inputs_embeds_prenormed, + ) + hidden_states = hidden_states / self.mup_multiplier + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + + +def _encode_inkling_image_embeds( + visual: InklingVisionModel, multimodal_params: list +) -> List[torch.Tensor]: + """Run the hMLP vision tower over the context requests' patch features. + + Reads ``multimodal_data['image']['vision_patches_bthwc']`` (the tensor the + :class:`InklingInputProcessor` attaches) from each context + ``MultimodalParams``, concatenates them, and runs the tower on the tower's + device/dtype. Returns a single-element list ``[feats]`` with ``feats`` of + shape ``(sum_patches, decoder_dmodel)`` -- the same shape + ``get_multimodal_embeddings`` returns and ``find_input_mm_embeds`` slices -- + or ``[]`` when no context request carries image features.""" + patches = [] + for param in multimodal_params: + data = getattr(param, "multimodal_data", None) or {} + image = data.get("image") or {} + vp = image.get("vision_patches_bthwc") + if vp is not None: + patches.append(vp) + if not patches: + return [] + p = next(visual.parameters()) + x = torch.cat([vp.to(device=p.device, dtype=p.dtype) for vp in patches], dim=0) + return [visual(x)] + + +def _encode_inkling_audio_embeds( + audio_tower: InklingAudioModel, multimodal_params: list +) -> List[torch.Tensor]: + """Run the dMel audio tower over the context requests' audio features. + + Reads ``multimodal_data['audio']['dmel_bins']`` (the tensor the + :class:`InklingInputProcessor` attaches) from each context + ``MultimodalParams``, concatenates them, and runs the tower. Returns a + single-element list ``[feats]`` with ``feats`` of shape + ``(sum_frames, decoder_dmodel)`` -- the same shape contract as the image + encoder -- or ``[]`` when no context request carries audio features.""" + frames = [] + for param in multimodal_params: + data = getattr(param, "multimodal_data", None) or {} + audio = data.get("audio") or {} + db = audio.get("dmel_bins") + if db is not None: + frames.append(db) + if not frames: + return [] + dev = audio_tower.encoder.weight.device + # dMel bins are integer codebook indices; keep them integral (the tower casts + # to long internally), only moving them onto the tower's device. + x = torch.cat([f.to(device=dev) for f in frames], dim=0) + return [audio_tower(x)] + + +@register_auto_model("InklingForConditionalGeneration") +@register_input_processor( + InklingInputProcessor, + model_type="inkling_mm_model", + placeholder_metadata=MultimodalPlaceholderMetadata( + # Image and audio are distinct registered modalities (```` -> one + # token per vision patch; ``