From 19fb41a41d586ec3f8955fc31af8655580848d83 Mon Sep 17 00:00:00 2001 From: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:28:38 -0700 Subject: [PATCH 01/15] [None][feat] Add Inkling NVFP4 model support Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com> --- docs/source/models/supported-models.md | 4 + .../sparse/inkling/__init__.py | 72 + .../sparse/inkling/backend.py | 293 +++ .../sparse/inkling/cache_manager.py | 498 ++++ .../sparse/inkling/conv_state.py | 374 +++ .../sparse/inkling/kernels.py | 570 +++++ .../sparse/inkling/metadata.py | 33 + .../sparse/inkling/page_table.py | 185 ++ .../sparse/inkling/params.py | 87 + .../attention_backend/sparse/registry.py | 17 + tensorrt_llm/_torch/configs/__init__.py | 5 + tensorrt_llm/_torch/configs/inkling.py | 230 ++ tensorrt_llm/_torch/model_config.py | 52 +- tensorrt_llm/_torch/models/__init__.py | 2 + tensorrt_llm/_torch/models/_arch_index.py | 4 + .../_torch/models/checkpoints/__init__.py | 3 +- .../checkpoints/hf/inkling_weight_mapper.py | 375 +++ .../_torch/models/modeling_inkling.py | 1334 ++++++++++ .../models/modeling_inkling_multimodal.py | 1126 +++++++++ tensorrt_llm/_torch/pyexecutor/_util.py | 60 +- .../_torch/pyexecutor/config_utils.py | 100 + .../_torch/pyexecutor/py_executor_creator.py | 4 +- tensorrt_llm/evaluate/lm_eval.py | 114 +- tensorrt_llm/evaluate/post_processing.py | 102 + tensorrt_llm/llmapi/llm_utils.py | 7 + tensorrt_llm/llmapi/reasoning_parser.py | 164 ++ tensorrt_llm/sampling_params.py | 40 + .../defs/accuracy/references/gsm8k.yaml | 31 + .../defs/accuracy/references/mmlu.yaml | 13 + .../defs/accuracy/references/mmmu.yaml | 34 + .../defs/accuracy/test_llm_api_pytorch.py | 150 ++ .../test_llm_api_pytorch_multimodal.py | 62 +- .../test_lists/qa/llm_function_core.txt | 4 + .../test_lists/test-db/l0_b200.yml | 4 + .../attention/test_inkling_chunked_prefill.py | 467 ++++ .../executor/test_mamba_cache_manager.py | 66 + .../_torch/modeling/test_modeling_inkling.py | 2143 +++++++++++++++++ .../test_modeling_inkling_multimodal.py | 466 ++++ .../references/trtllm_serve_cli.yaml | 2 +- .../unittest/llmapi/test_reasoning_parser.py | 164 ++ tests/unittest/llmapi/test_sampling_params.py | 80 + tests/unittest/others/test_lm_eval.py | 169 +- 42 files changed, 9661 insertions(+), 49 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/__init__.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/backend.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/cache_manager.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/conv_state.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/kernels.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/metadata.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/page_table.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/inkling/params.py create mode 100644 tensorrt_llm/_torch/configs/inkling.py create mode 100644 tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py create mode 100644 tensorrt_llm/_torch/models/modeling_inkling.py create mode 100644 tensorrt_llm/_torch/models/modeling_inkling_multimodal.py create mode 100644 tests/unittest/_torch/attention/test_inkling_chunked_prefill.py create mode 100644 tests/unittest/_torch/modeling/test_modeling_inkling.py create mode 100644 tests/unittest/_torch/modeling/test_modeling_inkling_multimodal.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 7e8fb78e6e3a..1a1935efcd53 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -30,6 +30,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`, `thinkingmachines/Inkling-Small-NVFP4` | | `InternLM3ForCausalLM` [^5] | InternLM3 | `internlm/internlm3-8b-instruct` | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | `moonshotai/Kimi-K2.5` | | `LagunaForCausalLM` | Laguna-XS | `poolside/laguna-XS.2` | @@ -79,6 +80,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 | No | Yes | No | Yes | No | No | No | No | Yes | Untested | Yes | 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 | @@ -99,6 +101,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]: NVFP4 checkpoint 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 in the `TRTLLM` family (the framework default). MTP, LoRA, function calling, constrained/guided decoding, disaggregated serving, and multimodal-hash prefix caching are not supported in this release. KV cache reuse additionally needs `kv_cache_config.mamba_state_config.periodic_snapshot_interval` set to a positive number of tokens, which snapshots the short-conv window a prefix hit has to restore; it applies to text prompts only. CUDA graph works at the default `moe_expert_parallel_size=1` but is broken with expert parallelism, hence No. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) @@ -127,6 +130,7 @@ complete encoder input must fit in the iteration token budget. | `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 | No | Yes | 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/sparse/inkling/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/__init__.py new file mode 100644 index 000000000000..2cca5d8c4dec --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/__init__.py @@ -0,0 +1,72 @@ +# 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, backend, cache manager, conv-state pool. + +Inkling's attention is dense (full causal, sliding window on local layers, plus a +learned relative-bias ``score_mod``), not sparse. It lives under ``sparse/`` +because that is where every model-specific attention backend already lives, and +selection goes through ``sparse/registry.py`` like its neighbours. + +The attention metadata adds one thing to the base: the conv pool's per-step slot +write, which has to run outside the CUDA-graph capture region. Everything else the +decode kernel needs is a view of the base's own buffers (``page_table.py``). +""" + +from .backend import InklingTritonAttention +from .cache_manager import InklingHybridCacheManager +from .conv_state import ( + CONV_ROLE, + InklingConvRuntime, + InklingConvState, + InklingConvStateCache, + InklingRole, + apply_short_conv, +) +from .kernels import ( + build_page_table, + inkling_decode_attention, + inkling_prefill_attention, + write_kv_cache_hnd, +) +from .metadata import InklingAttentionMetadata +from .page_table import gen_page_table, gen_seq_lens, page_div +from .params import ( + InklingBackendForwardArgs, + InklingSparseAttentionConfig, + InklingSparseParams, + inkling_forward_args, +) + +__all__ = [ + "InklingAttentionMetadata", + "InklingBackendForwardArgs", + "InklingSparseAttentionConfig", + "InklingSparseParams", + "InklingConvRuntime", + "InklingConvState", + "InklingConvStateCache", + "InklingRole", + "InklingHybridCacheManager", + "InklingTritonAttention", + "apply_short_conv", + "build_page_table", + "gen_page_table", + "gen_seq_lens", + "inkling_forward_args", + "inkling_decode_attention", + "inkling_prefill_attention", + "page_div", + "write_kv_cache_hnd", +] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/backend.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/backend.py new file mode 100644 index 000000000000..976a28c5822a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/backend.py @@ -0,0 +1,293 @@ +# 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 backend: KV write, page tables, prefill/decode dispatch.""" + +from typing import Optional + +import torch + +from ...interface import AttentionForwardArgs, merge_attention_forward_args +from ...trtllm import TrtllmAttention +from .kernels import ( + build_page_table, + inkling_decode_attention, + inkling_prefill_attention, + write_kv_cache_hnd, +) +from .metadata import InklingAttentionMetadata +from .page_table import gen_page_table, gen_seq_lens, page_div, validate_decode_layout +from .params import InklingBackendForwardArgs + + +class InklingTritonAttention(TrtllmAttention): + """Runs Inkling's Triton attention over the paged KV cache. + + Inkling's two private forward inputs (``rel_logits``, ``allow_mixed``) ride + ``AttentionForwardArgs.sparse_backend_args``; the decode kernel's seq lens and + page table are slices of buffers the metadata already keeps graph-stable (see + ``page_table.py``). ``sm_scale`` / ``rel_extent`` / ``window_left`` are + assigned by ``InklingAttention`` after construction, since + ``create_attention()`` has a fixed kwarg list. + """ + + Metadata = InklingAttentionMetadata + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: InklingAttentionMetadata, + forward_args: Optional[AttentionForwardArgs] = None, + **kwargs, + ) -> torch.Tensor: + """Dispatch prefill / decode over the paged cache. + + The runtime packs context requests first, then one-token generation + requests, so slicing at that boundary feeds the prefill and paged-decode + kernels respectively. ``forward`` is overridden rather than delegated to + keep Inkling off the sparse prediction hooks ``TrtllmAttention.forward`` + calls. + """ + forward_args = merge_attention_forward_args(forward_args, kwargs) + args = forward_args.sparse_backend_args + if not isinstance(args, InklingBackendForwardArgs): + raise TypeError( + "InklingTritonAttention.forward needs its rel_logits bias in " + "forward_args.sparse_backend_args as an " + f"InklingBackendForwardArgs, got {type(args).__name__}. Build it " + "with inkling_forward_args(); the shared AttentionForwardArgs " + "has no field that can carry a per-query-token bias." + ) + rel_logits = args.rel_logits + allow_mixed = args.allow_mixed + attn_metadata = metadata + # KVCacheManagerV2 takes the global layer index. + 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 batch needs the per-request short-conv state pool; the + # stateless path would convolve across the context/generation boundary. + if 0 < num_contexts < num_seqs and not allow_mixed: + raise NotImplementedError( + "InklingAttention: mixed context+generation batch needs the " + "short-conv state pool; the stateless short-conv path convolves " + "across the context/generation boundary of the packed batch. " + "Set allow_mixed=True on InklingBackendForwardArgs only when the " + "pool path is active (conv_rt is not None)." + ) + + 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) + # One prefill path, always over the pages. A request carrying cached + # history -- a later chunk, or a reused prefix -- attends to tokens it + # did not bring with it, and the write above has already put this + # chunk's K/V in the same pages, so no gather is needed. + max_total = max(int(c) + int(sl) for c, sl in zip(num_cached, seq_lens)) + max_pages = (max_total + page_size - 1) // page_size + page_table = build_page_table(block_ids, max_pages, device) + num_cached_dev = torch.tensor( + [int(c) for c in num_cached], dtype=torch.int32, device=device + ) + return inkling_prefill_attention( + q, + k_cache, + v_cache, + cu, + num_cached_dev, + page_table, + page_size, + 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 + # CUDA-graph-safe path: no host->device copy. Slice the base metadata's + # graph-stable kv_lens_cuda and kv_cache_block_offsets, then persist the + # new K/V with an in-graph scatter whose indices are derived on-GPU. + num_req = q.shape[0] + if ( + getattr(attn_metadata, "kv_cache_block_offsets", None) is not None + and attn_metadata.num_generations == num_req + ): + validate_decode_layout(attn_metadata, cache_layer, num_req) + sl = gen_seq_lens(attn_metadata, num_req) + pt = gen_page_table(attn_metadata, cache_layer)[:num_req] + div = page_div(attn_metadata) + 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 + # Divide after the gather: [num_req] elements, not [num_req, blocks]. + pages = pt.gather(1, page_row.unsqueeze(1)).squeeze(1).long() + if div > 1: + pages = torch.div(pages, div, rounding_mode="floor") + # Paired advanced indices select one (page, slot) per request -> + # [num_req, num_kv_heads, head_dim], matching the 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, + page_div=div, + ) + # Eager fallback: no paged page table to borrow, so build one here from + # the host block table. Illegal under CUDA graph. + if getattr(attn_metadata, "is_cuda_graph", False): + raise RuntimeError( + "Inkling decode metadata is unusable for a CUDA-graph batch: " + f"attn_metadata is {type(attn_metadata).__name__} with " + f"num_generations={getattr(attn_metadata, 'num_generations', None)} " + f"and kv_cache_block_offsets=" + f"{'set' if getattr(attn_metadata, 'kv_cache_block_offsets', None) is not None else 'None'}, " + f"expected a TrtllmAttentionMetadata with {num_req} generation " + "rows and an allocated page table. Inkling's backend is selected " + "through sparse/registry.py under the TRTLLM backend family; " + "remove any attn_backend override from --extra_llm_api_options / " + "LLM(attn_backend=...) so the default applies." + ) + 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) + # No page_div here: get_batch_cache_indices already divides by kv_factor, + # so these are block indices rather than page indices. + 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, + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/cache_manager.py new file mode 100644 index 000000000000..efe0f7171fa2 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/cache_manager.py @@ -0,0 +1,498 @@ +# 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. + +The short-conv state is registered with V2 as SSM layers rather than allocated on +the side: ``SsmLayerConfig`` is the framework's container for per-request +fixed-size state, and V2 sizes it with the right rule (one dedicated block per +request, never shared), so the pool's bytes enter V2's quota and come from the +pool allocator at stable device addresses. V2 owns the memory but not the +request-to-slot mapping, so :class:`InklingConvStateCache` still assigns rows, as +Mamba's manager does. + +The SSM layers are appended after the attention layers rather than interleaved, +which leaves every attention ``layer_id`` -- and therefore the paged-KV +addressing -- untouched. +""" + +import hashlib +from dataclasses import replace +from typing import List + +import torch + +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + DEFAULT_BEAM_INDEX, + BatchDesc, + BufferConfig, + DataRole, + KVCacheDesc, + LayerId, + PageIndexMode, + SsmLayerConfig, +) + +from .....logger import logger +from ....pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 +from ....pyexecutor.llm_request import LlmRequest +from ....pyexecutor.scheduler import ScheduledRequests +from .conv_state import CONV_ROLE, InklingConvState, InklingConvStateCache + + +def _resolve_conv_dtype(pretrained_config) -> torch.dtype: + """The compute dtype the short-conv pool holds. + + Not the manager's ``dtype`` argument: that is the KV cache dtype, which is + nvfp4/fp8 on quantized releases while this pool holds unquantized pre-conv + activations. An unresolvable value raises rather than defaulting, so a + wrong-dtype pool cannot surface far from its cause. + """ + config = getattr(pretrained_config, "text_config", pretrained_config) + dtype = getattr(config, "torch_dtype", None) + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, str): + resolved = getattr(torch, dtype, None) + if isinstance(resolved, torch.dtype): + return resolved + raise ValueError( + f"Inkling short-conv pool needs the model's compute dtype, but " + f"torch_dtype={dtype!r} on {type(config).__name__} is not a torch dtype" + ) + + +class InklingHybridCacheManager(KVCacheManagerV2): + """Paged KV (V2, per-layer geometry) + the short-conv state pool. + + Folding the pool into the manager -- ``CppMambaHybridCacheManager``'s shape + -- lets it reach the model through ``attn_metadata.kv_cache_manager`` and be + released by the manager's own ``free_resources``, so conv rows and KV blocks + cannot drift apart. + + Registering the conv window as V2 SSM layers is what makes block reuse + servable: V2 attaches it to the block committed at a snapshot ordinal and + hands it back on a prefix hit. + """ + + def __init__(self, *args, pretrained_config, mapping, max_batch_size, **kwargs): + # Resolved before super().__init__() because _build_cache_config runs + # inside it and the base does not keep pretrained_config. Everything else + # the conv sizing needs is on self by then. + self._conv_config = getattr(pretrained_config, "text_config", pretrained_config) + self._conv_dtype = _resolve_conv_dtype(pretrained_config) + # prepare_expect_snapshot_points needs the interval; the base keeps none. + self._kv_cache_config = args[0] if args else kwargs.get("kv_cache_config") + # A property of the model, so warn once per manager, not per request. + self._warned_multimodal_reuse = False + super().__init__( + *args, + pretrained_config=pretrained_config, + mapping=mapping, + max_batch_size=max_batch_size, + **kwargs, + ) + # Resolved on first use: `impl` is not guaranteed live at the end of + # the base constructor. From the first conv layer, as Mamba does. + self._conv_layer_group_id = None + self._conv_cache = InklingConvStateCache( + pretrained_config, + self._conv_tp_size, + self._num_conv_request_slots, + torch.device("cuda", torch.cuda.current_device()), + self._conv_dtype, + reserve_attention_dp_slot=self._reserve_attention_dp_slot, + max_draft_len=self.max_draft_len, + allocate=self._conv_state_buffer, + resolve_slot=self._conv_slot_for_request, + ) + logger.info( + f"Inkling short-conv state pool: {self._conv_cache.num_slots} rows " + f"({self._num_conv_request_slots} request + " + f"{self._num_reserved_conv_slots} reserved), " + f"{self._conv_cache.conv_state_bytes() / (1 << 20):.1f} MiB, " + "backed by V2 SSM layers" + ) + + def _conv_slot_for_request(self, request_id: int): + """V2's recurrent-state slot for ``request_id``, or None if it has none. + + V2 restores a reuse hit into the slot it assigned, so the pool has to + be indexed its way. None for CUDA-graph padding sentinels, which have no + cache entry and keep the reserved rows. + """ + kv_cache = self.kv_cache_map.get(request_id) + if kv_cache is None: + return None + if self._conv_layer_group_id is None: + self._conv_layer_group_id = self.impl.get_layer_group_id(self._conv_layer_id(0)) + slot = kv_cache.get_ssm_block_base_index(self._conv_layer_group_id) + return None if slot < 0 else int(slot) + + # ========================= KV cache block reuse ========================= + # Everything to "end KV cache block reuse" runs only under + # enable_block_reuse. + def prepare_expect_snapshot_points(self, requests) -> None: + """Where this batch's requests may snapshot their short-conv window. + + Found by ``py_executor`` through ``hasattr``; the scheduler then ends a + context chunk only here, which is the half a model cannot do for itself + -- a snapshot can only be taken where an iteration ends. The interval is + deliberately coarse: one snapshot costs the whole model's conv window, + so reuse lands only on multiples of it. + """ + state_config = getattr(self._kv_cache_config, "mamba_state_config", None) + interval = getattr(state_config, "periodic_snapshot_interval", 0) or 0 + for request in requests: + if not self.enable_block_reuse or not interval: + request.expect_snapshot_points = [] + continue + request.expect_snapshot_points = list(range(interval, request.prompt_len + 1, interval)) + + # The three methods below mirror MambaHybridCacheManagerV2's, copied rather + # than shared: lifting them into kv_cache_manager_v2.py edits a file every + # model depends on, and inheriting the Mamba manager drags in an SSM-state + # interface four short convs at two widths cannot satisfy. Every ``self`` + # here is a KVCacheManagerV2 member. A test compares them as ASTs. + def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of the Inkling V2 cache for " + f"request {request.py_request_id} to {history_length} tokens" + ) + + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + """Commit at each declared snapshot point, not only at the end. + + The base commits once as prefill advances, which is enough when + everything reusable lives in the KV pages. The conv window does not: V2 + attaches it to the block committed at a snapshot ordinal, so a commit + has to happen there or a later hit has no window to restore. + """ + should_block_reuse = ( + self.enable_block_reuse and not self.is_draft and not request.is_dummy_request + ) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = ( + min(max(snapshot_points), request.prompt_len) if snapshot_points else request.prompt_len + ) + commit_end = min(request.context_current_position, commit_limit) + if ( + request.context_current_position in request.expect_snapshot_points + and commit_end > kv_cache.num_committed_tokens + ): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = ( + self.enable_block_reuse and not self.is_draft and not request.is_dummy_request + ) + is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE + is_snapshot_boundary = ( + request.context_current_position in request.expect_snapshot_points + ) + has_pending_snapshot = any( + point > request.context_current_position for point in request.expect_snapshot_points + ) + should_resize = not should_block_reuse or ( + not is_all_reusable and not has_pending_snapshot + ) + should_commit = ( + is_all_reusable or is_snapshot_boundary or request.context_remaining_length == 0 + ) + + if should_resize and not kv_cache.resize(None, request.context_current_position): + raise ValueError( + "Failed to resize history length of the Inkling V2 cache " + f"for request {request.py_request_id} to " + f"{request.context_current_position} tokens at context update" + ) + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) + kv_cache.enable_swa_scratch_reuse = False + + def _augment_tokens_for_block_reuse(self, tokens, req, start=0, end=None): + """Keep multimodal requests out of the shared radix tree. + + The base distinguishes images by rewriting their spans into keys + carrying a content digest, but only for a request holding + ``multimodal_hashes`` -- which Inkling produces for none of them, so two + prompts with DIFFERENT images match as the same prefix: MMMU 81.3% with + reuse off against 31.3% with it on, same 32 items, no error anywhere. + Salting position 0 with the request id gives such a request a private + chain, leaving text requests alone. + """ + augmented = super()._augment_tokens_for_block_reuse(tokens, req, start, end) + if not self.enable_block_reuse: + return augmented + if getattr(req, "multimodal_hashes", None) is not None: + return augmented + if getattr(req, "py_multimodal_data", None) is None: + return augmented + # Only the chunk holding position 0 needs it: the tree chains each key + # through its parent, so a private first block privatises the rest. + if start != 0 or not len(augmented): + return augmented + if not self._warned_multimodal_reuse: + self._warned_multimodal_reuse = True + logger.warning( + "Inkling: KV cache block reuse is disabled for multimodal " + "requests. Their image/video/audio spans carry no content " + "digest, so a shared prefix would be matched on placeholder " + "token ids alone and serve one item's KV for another's. Text " + "requests are unaffected." + ) + salt = hashlib.sha256( + b"inkling-multimodal-no-digest-%d" % int(getattr(req, "request_id", 0)) + ).digest() + poisoned = list(augmented) + poisoned[0] = salt + return poisoned + + def prepare_context(self, req): + """Observation only: how much prefix was reused, and from whose slot. + + The framework's own ``KVCacheV2SsmSnapshotIterationStats`` would be the + right source, but a served deployment publishes no + ``kvCacheIterationStats*`` key at all, so there is nothing to read. A + reuse number without a gate reads the same whether reuse worked or never + ran. The slot source rides along because the two failure modes -- no + reuse, and reuse into a row the kernels do not read -- are + indistinguishable from accuracy. + """ + first = bool(getattr(req, "is_first_context_chunk", True)) + ok = super().prepare_context(req) + if not first: + return ok + d = getattr(self, "_reuse_dbg", None) + if d is None: + d = self._reuse_dbg = {"n": 0, "hits": 0, "best": 0, "v2_slots": 0} + d["n"] += 1 + pos = int(getattr(req, "context_current_position", 0) or 0) + if pos > 0: + d["hits"] += 1 + d["best"] = max(d["best"], pos) + if ( + self._conv_slot_for_request( + getattr(req, "py_request_id", getattr(req, "request_id", -1)) + ) + is not None + ): + d["v2_slots"] += 1 + if d["n"] % 64 == 0: + logger.info( + f"Inkling prefix reuse: {d['hits']}/{d['n']} requests, " + f"longest={d['best']} tokens, conv rows from V2: " + f"{d['v2_slots']}/{d['n']}" + ) + return ok + + # ======================= end KV cache block reuse ======================= + + # ---- conv geometry, all derived from what the base already resolved ----- + @property + def _conv_tp_size(self) -> int: + # The attention TP, not the global one: the k/v convs follow the kv-head + # split, as V2 does for the paged pool. + return 1 if self.mapping.enable_attention_dp else self.mapping.tp_size + + @property + def _reserve_attention_dp_slot(self) -> bool: + return bool(self.mapping.enable_attention_dp) + + @property + def _num_conv_request_slots(self) -> int: + # One row per resident sequence; each pipeline stage holds a microbatch. + return self.max_batch_size * self.mapping.pp_size + + @property + def _num_reserved_conv_slots(self) -> int: + # Asked of the pool rather than re-derived: the two counts must agree or + # slots_for indexes the V2 buffer out of bounds. + return InklingConvStateCache.reserved_slot_count( + reserve_attention_dp_slot=self._reserve_attention_dp_slot + ) + + def _conv_section_bytes(self, global_layer_idx: int) -> List[int]: + """Per-slot bytes of the layer's four conv sections, in pool order.""" + config = self._conv_config + kv_dim = ( + config.layer_num_kv_heads(global_layer_idx) * config.layer_head_dim(global_layer_idx) + ) // self._conv_tp_size + window = config.sconv_kernel_size - 1 + itemsize = torch.empty((), dtype=self._conv_dtype).element_size() + return [ + c * window * itemsize for c in (kv_dim, kv_dim, config.hidden_size, config.hidden_size) + ] + + def _conv_bytes_per_slot(self, global_layer_idx: int) -> int: + """Bytes one request occupies in the layer's conv state. + + One number, not four: the sections share a buffer. ``_conv_section_bytes`` + keeps their widths, which is what a TP-mismatched transfer splits on. + """ + return sum(self._conv_section_bytes(global_layer_idx)) + + # ---- V2 configuration ------------------------------------------------- + def _conv_layer_id(self, local_layer_idx: int) -> LayerId: + """Cache-layer id holding ``local_layer_idx``'s four conv states.""" + return LayerId(self.num_local_layers + local_layer_idx) + + def _build_cache_config(self, config): + """Append one SSM layer per decoder layer for the short-conv state. + + The attention layers ``_build_base_config`` produced are preserved + exactly: an Inkling layer needs both KV and conv state, and a + ``LayerConfig`` is one or the other, so appending is the only way. + """ + layers = list(config.layers) + num_attention_layers = len(layers) + # _conv_layer_id derives SSM ids arithmetically from this count. + assert num_attention_layers == self.num_local_layers, ( + num_attention_layers, + self.num_local_layers, + ) + for local_idx in range(num_attention_layers): + layers.append( + SsmLayerConfig( + layer_id=LayerId(num_attention_layers + local_idx), + buffers=[ + BufferConfig( + role=CONV_ROLE, + size=self._conv_bytes_per_slot(self.pp_layers[local_idx]), + ) + ], + ) + ) + + num_reserved = self._num_reserved_conv_slots + # Non-request rows as zero-capacity requests, like Mamba: no attention + # pages, one state slot each -- exactly their effect on the pool. + dummies = [KVCacheDesc(capacity=0, history_length=0) for _ in range(num_reserved)] + constraints = [ + replace(batch, kv_caches=[*batch.kv_caches, *dummies]) for batch in config.constraints + ] + # Length-independent floor: a conv row is fixed per resident sequence, so + # the base's token-scaled constraints do not bound the pool. + constraints.append( + BatchDesc( + [ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(self._num_conv_request_slots + num_reserved) + ] + ) + ) + # KVCacheManagerConfig asserts this whenever an SSM layer exists. + return replace(config, layers=layers, constraints=constraints, commit_min_snapshot=True) + + def _get_pool_roles(self, pool_id: int): + """Name a role that actually exists in ``pool_id``. + + The base returns ``Role.KEY`` unconditionally, but conv pools hold no KEY. + """ + first_layer = int(self.impl.layer_grouping[pool_id][0]) + if first_layer < self.num_local_layers: + return super()._get_pool_roles(pool_id) + if (pool_id, CONV_ROLE) in self._pool_layer_ids_by_role: + return CONV_ROLE, None + raise RuntimeError( + f"Inkling conv pool {pool_id} (first layer {first_layer}) does not " + f"hold {CONV_ROLE}; _build_cache_config and the pool packing " + "disagree about what landed where." + ) + + def _conv_state_buffer( + self, local_layer_idx: int, role: DataRole, state_shape: List[int] + ) -> torch.Tensor: + """A ``[num_slots, *state_shape]`` view of one conv state's pool memory. + + Mirrors ``MambaHybridCacheManagerV2._get_state_buffer``. The + ``as_strided`` step converts the raw page-indexed view into slot indexing, + since V2 packs ``page_index_scale`` pages per logical slot. + """ + layer_id = self._conv_layer_id(local_layer_idx) + addr = self.impl.get_mem_pool_base_address(layer_id, role, PageIndexMode.SHARED) + num_pages = self.impl.get_page_index_upper_bound(layer_id, role) + raw = convert_to_torch_tensor( + TensorWrapper(addr, self._conv_dtype, [num_pages] + list(state_shape)) + ) + scale = self.impl.get_page_index_scale(layer_id, role) + num_slots = (num_pages + scale - 1) // scale + return raw.as_strided( + [num_slots] + list(state_shape), + [raw.stride(0) * scale] + list(raw.stride()[1:]), + ) + + # ---- model-facing ----------------------------------------------------- + @property + def conv_state_cache(self) -> InklingConvStateCache: + """The short-conv state pool, for the metadata's per-step publication.""" + return self._conv_cache + + def get_conv_states(self, layer_idx: int) -> InklingConvState: + """The four short-conv state buffers of ``layer_idx``. + + Named after ``BaseMambaCacheManager.get_conv_states`` but deliberately not + implementing it: that hook returns one tensor per layer, which cannot + express four convs at two widths, and Inkling backs no SSM state at all. + """ + return self._conv_cache.layer_state(layer_idx) + + def get_state_indices(self) -> torch.Tensor: + """Pool rows of the current batch, in packed batch order.""" + return self._conv_cache.state_indices + + 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 together with the request's KV blocks, so a + leaked row cannot be reused with stale state.""" + 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/sparse/inkling/conv_state.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/conv_state.py new file mode 100644 index 000000000000..fbbf03c1bfbd --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/conv_state.py @@ -0,0 +1,374 @@ +# 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-request short-conv state for Inkling: the pool, the per-forward split, +and the helper that runs one short convolution through them. + +This is runtime state, not model weights: it has the lifetime of the paged KV +cache, is owned by :class:`~.cache_manager.InklingHybridCacheManager`, and its +rows are published each step from +:meth:`~.metadata.InklingAttentionMetadata.prepare`, before the captured region. +""" + +from collections import namedtuple +from dataclasses import dataclass +from typing import Callable, List, Optional + +import torch + +from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole + +from ....._utils import prefer_pinned + +# Per-request short-conv state of one decoder layer, carried across decode steps. +# Each field is a ``[num_req, channels, sconv_kernel_size - 1]`` window of the +# previous pre-conv inputs (oldest first): ``k``/``v`` for the attention k/v +# convs (TP-sharded), ``attn``/``mlp`` for the residual-stream convs (replicated). +InklingConvState = namedtuple("InklingConvState", ["k", "v", "attn", "mlp"]) + + +class InklingRole: + """V2 pool role for a decoder layer's short-conv state. + + One role, not four: the four convolutions share a window length, so they are + one buffer per layer holding ``[k | v | attn | mlp]`` along the channel axis. + That is also the shape a state transfer wants -- one pool with the section + widths beside it, as Mamba stores ``[x | B | C]``. + """ + + CONV_STATE = DataRole("inkling_conv") + + +CONV_ROLE = InklingRole.CONV_STATE +# Section order within the buffer; matches InklingConvState's fields. +CONV_SECTIONS = ("k", "v", "attn", "mlp") + + +class InklingConvStateCache: + """Runtime-owned per-request short-conv state pool for the whole decoder. + + Per layer it holds the four :class:`InklingConvState` buffers, each + ``[num_slots, channels, kernel_size - 1]``; the k/v conv channels follow the + fused-qkv k/v split (TP-sharded), the residual-stream convs are replicated. + The buffers come from the ``allocate`` callback, which + :class:`InklingHybridCacheManager` wires to V2 SSM-layer pool memory: this + class owns the request-to-row mapping, V2 owns the bytes. + + The row count is fixed at construction and every buffer keeps a stable device + address, mutated in place, so a captured graph replays cleanly. + """ + + @staticmethod + def reserved_slot_count(*, reserve_attention_dp_slot: bool) -> int: + """Rows sitting above the per-request ones: one shared by every CUDA-graph + padding sentinel, plus one for the attention-DP idle dummy when enabled. + + Static because the cache manager needs the same count to declare + V2's min-slots floor, and the two must agree exactly. + """ + return 1 + int(reserve_attention_dp_slot) + + def __init__( + self, + pretrained_config, + tp_size: int, + num_request_slots: int, + device: torch.device, + dtype: torch.dtype, + *, + reserve_attention_dp_slot: bool = False, + max_draft_len: int = 0, + allocate: Optional[Callable[[int, object, List[int]], torch.Tensor]] = None, + resolve_slot: Optional[Callable[[int], Optional[int]]] = None, + ): + # 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 + if num_request_slots <= 0: + raise ValueError(f"num_request_slots must be > 0, got {num_request_slots}") + self.num_request_slots = num_request_slots + # Reserved rows sit above the real ones so a real slot id is always a + # valid index into the first ``num_request_slots`` rows. + self._padding_slot = num_request_slots + self._attention_dp_dummy_slot = num_request_slots + 1 if reserve_attention_dp_slot else None + self._max_draft_len = max(0, int(max_draft_len)) + num_slots = num_request_slots + self.reserved_slot_count( + reserve_attention_dp_slot=reserve_attention_dp_slot + ) + self.num_slots = num_slots + self.kwin = kwin + + # ``allocate`` hands back V2 pool memory when the cache manager built us; + # standalone callers fall back to a private allocation of the same shapes. + # Pool memory arrives uninitialized, which is safe only because + # ``slots_for`` zeroes a row when it first hands it out. + if allocate is None: + + def allocate(_layer_idx, _role, state_shape): + return torch.zeros(num_slots, *state_shape, device=device, dtype=dtype) + + self._layers: List[InklingConvState] = [] + self._section_channels: List[List[int]] = [] + 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 + sections = [kv_dim, kv_dim, hidden, hidden] + self._section_channels.append(sections) + buf = allocate(i, CONV_ROLE, [sum(sections), kwin]) + if buf.shape[0] < num_slots: + raise RuntimeError( + f"Inkling conv pool layer {i}: allocator returned " + f"{buf.shape[0]} slots but the pool needs {num_slots} " + f"({num_request_slots} request + reserved). The V2 SSM " + "layer was sized from a different slot count than " + "InklingConvStateCache assumes." + ) + buf = buf[:num_slots] + # Views, not copies: the causal_conv1d ops read conv_state's + # strides off the tensor, so a section need not be contiguous. + offsets = [0] + for width in sections: + offsets.append(offsets[-1] + width) + self._layers.append( + InklingConvState(*(buf[:, a:b, :] for a, b in zip(offsets, offsets[1:]))) + ) + # Refreshed in place per forward so a captured graph aliases it and every + # replay sees the current batch. Indexed by batch position, not by slot. + self.state_indices = torch.arange(num_slots, dtype=torch.int32, device=device) + # Pinned host staging for that write: one async H2D copy per forward. + self.state_indices_cpu = torch.zeros( + num_slots, dtype=torch.int32, pin_memory=prefer_pinned() + ) + # Set by the cache manager to V2's slot for the request; see slots_for. + self._resolve_slot = resolve_slot + self._slot_of = {} + self._free = list(range(num_request_slots - 1, -1, -1)) + + def conv_state_bytes(self) -> int: + """Total device bytes this pool holds. Reported by the cache manager.""" + return sum(t.numel() * t.element_size() for st in self._layers for t in st) + + def section_bytes(self, layer_idx: int) -> List[int]: + """Per-slot bytes of each section, in ``CONV_SECTIONS`` order. + + What a state transfer needs to split one slot at its semantic + boundaries when the two sides disagree on TP -- the k/v sections are + sharded, the residual-stream ones are not. + """ + itemsize = self._layers[layer_idx].k.element_size() + return [width * self.kwin * itemsize for width in self._section_channels[layer_idx]] + + 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 _reserved_slot_for(self, request_id: int) -> Optional[int]: + """The reserved row ``request_id`` aliases, or None for a real request. + + ``cuda_graph_runner`` caches one dummy id per runtime draft length, so the + whole descending range maps here, not just the exact sentinel. + """ + from ....pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + from ....pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID + + if CUDA_GRAPH_DUMMY_REQUEST_ID - self._max_draft_len <= request_id: + return self._padding_slot + if request_id == ATTENTION_DP_DUMMY_REQUEST_ID: + return self._attention_dp_dummy_slot + return None + + def slots_for(self, request_ids: List[int]) -> List[int]: + """Map request ids to their pool rows. + + With a ``resolve_slot`` callback the row is V2's, and this class + allocates nothing: V2 restores a snapshot into the slot IT assigned, so + a second numbering over the same buffer would have the kernels read a + row nobody restored. Without the callback (standalone use, unit tests) + a private free list is self-consistent, since nothing else touches the + pool with reuse off. + """ + slots = [] + for r in request_ids: + if self._resolve_slot is not None: + slot = self._resolve_slot(r) + if slot is None: + slot = self._reserved_slot_for(r) + if slot is None: + raise RuntimeError( + f"Inkling short-conv pool: request {r} has no V2 " + "recurrent-state slot and is not a padding sentinel. " + "Every resident sequence gets one when its KV cache is " + "created; reaching here means the request was never " + "added to the cache manager." + ) + slots.append(slot) + continue + slot = self._slot_of.get(r) + if slot is None: + slot = self._reserved_slot_for(r) + if slot is None: + if not self._free: + raise RuntimeError( + "Inkling short-conv pool is out of rows " + f"({self.num_request_slots} request rows, all live). " + "Every resident sequence needs one row; a row is " + "returned only by the cache manager's free_resources." + ) + slot = self._free.pop() + self._slot_of[r] = slot + # Only the private path zeroes: a V2 slot may hold a window V2 + # just restored into it, and zeroing would throw that away. + for st in self._layers: + for t in st: + t[slot].zero_() + slots.append(slot) + return slots + + def write_state_indices(self, request_ids: List[int]) -> List[int]: + """Resolve ``request_ids`` to pool rows and publish them into the stable + ``state_indices`` CUDA buffer, in packed batch order (contexts first). + + A captured decode graph aliases ``state_indices``, so this must run every + forward from eager input-prep -- it is reached from the manager's + ``get_state_indices``, which the base ``prepare()`` calls. + """ + slots = self.slots_for(request_ids) + n = len(slots) + if n > self.num_slots: + raise RuntimeError( + f"Inkling short-conv batch has {n} rows but the pool publishes " + f"only {self.num_slots} state indices; the pool is sized to the " + "padded scheduler batch, so this signals a batch-shape mismatch." + ) + 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) + # Reserved rows are permanent: returning one to the free list would + # hand a real request a row the next padding batch overwrites. + if slot is not None and slot < self.num_request_slots: + self._free.append(slot) + + +# ---- chunked prefill: declare the window a preceding chunk left behind ------ +def _context_num_cached(attn_metadata, num_contexts): + """Per-context-request cached token counts, or None when unavailable. + + None on the cache-free unit-test path, which callers read as "all fresh". + """ + params = getattr(attn_metadata, "kv_cache_params", None) + if params is None: + return None + per_seq = getattr(params, "num_cached_tokens_per_seq", None) + if per_seq is None: + return None + return [int(c) for c in per_seq[:num_contexts]] + + +@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 short-conv seeds the pool for context requests and + updates it in place for generation requests, like the paged attention split. + """ + + 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 from_metadata(cls, attn_metadata) -> Optional["InklingConvRuntime"]: + """Build this forward's split, or ``None`` for the stateless conv path. + + Called from ``model.forward``, so the pool rows come back as a view of the + buffer ``prepare()`` already refreshed. The varlen offsets below do + allocate, but only for a batch carrying context requests, and a captured + batch never does -- decode graphs are generation-only. + """ + mgr = getattr(attn_metadata, "kv_cache_manager", None) + cache = getattr(mgr, "conv_state_cache", None) + if cache is None or attn_metadata.request_ids is None: + return None + num_contexts = attn_metadata.num_contexts + batch_size = len(attn_metadata.request_ids) + state_indices = cache.state_indices + device = state_indices.device + query_start_loc = has_initial_state = None + if num_contexts: + seq_lens = attn_metadata.seq_lens.tolist()[:num_contexts] + cu = torch.zeros(num_contexts + 1, dtype=torch.int32, device=device) + cu[1:] = torch.tensor(seq_lens, dtype=torch.int32, device=device).cumsum(0) + query_start_loc = cu + # The window a preceding chunk left in the pool row is consumed + # only if declared. Sufficient only because the prefill kernel also + # reads the cached KV back; alone it would be the subtler bug. + cached = _context_num_cached(attn_metadata, num_contexts) + if cached is None: + has_initial_state = torch.zeros(num_contexts, dtype=torch.bool, device=device) + else: + has_initial_state = torch.tensor( + [c > 0 for c in cached], dtype=torch.bool, device=device + ) + return cls( + num_ctx_tokens=sum(attn_metadata.seq_lens.tolist()[:num_contexts]), + ctx_indices=state_indices[:num_contexts] if num_contexts else None, + gen_indices=( + state_indices[num_contexts:batch_size] if num_contexts < batch_size else None + ), + query_start_loc=query_start_loc, + has_initial_state=has_initial_state, + ) + + +def apply_short_conv( + sconv, # models.modeling_inkling.InklingShortConv (duck-typed to avoid a cycle) + 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`` and the generation slice + updates it in place, then the two outputs are concatenated in packed order. + """ + 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) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/kernels.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/kernels.py new file mode 100644 index 000000000000..bc7910e47f92 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/kernels.py @@ -0,0 +1,570 @@ +# 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. + +Inkling adds a learned per-(query-token, head, relative-distance) bias inside +the attention score. No fused, CUDA-graph-safe TensorRT-LLM backend exposes a +``score_mod`` hook, which is why this pair of Triton kernels is the production +path. + +The bias is precomputed torch-side into ``rel_logits`` +``[num_query_tokens, num_heads, rel_extent]``; the kernels only gather+add: + + rel_dist = q_pos - k_pos + bias = rel_logits[q_idx, head, clamp(rel_dist, 0, rel_extent - 1)] + if 0 <= rel_dist < rel_extent else 0 + +Its shape is static, so the decode kernel is graph-capturable: fixed launch +grid, sequence lengths read from a GPU tensor. Both kernels read the paged KV +cache in HND layout 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, but finite, so a fully-masked tile +# does not poison the running max. Inlined as a literal inside the kernels, +# since 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_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, + PAGE_DIV: 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) + if PAGE_DIV > 1: + # Entries count pages; the K/V views count blocks. One integer + # division per KV tile, and PAGE_DIV is a constexpr power of two. + page_id = page_id // PAGE_DIV + + 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_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, + page_div: int = 1, +) -> 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 page ids, indexed by + absolute block ordinal (``k_pos // page_size``), scaled by + ``page_div``. + 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). + page_div: divisor on every ``page_table`` entry. 1 when the table already + holds indices into the ``[blocks, ...]`` K/V views; ``kv_factor`` + when it is a slice of ``kv_cache_block_offsets``, whose unit is + pages. Power of two, so the division lowers to a shift. + + Returns ``[batch, num_heads, head_dim]`` in q's dtype. + """ + assert page_div >= 1 and (page_div & (page_div - 1)) == 0, page_div + 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, + PAGE_DIV=page_div, + 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 + + +# --------------------------------------------------------------------------- +# Chunked-context prefill kernel: BLOCK_M query tiling like the prefill kernel, +# paged KV reads like the decode kernel, and absolute positions so a query can +# attend to history it did not bring with it. +# +# This is the path for any context request with cached history -- a later +# chunk of a chunked prefill, or a chunk sitting on a reused prefix. It needs +# no gather: InklingAttention._run_context writes this call's new K/V into the +# pages *before* attention, so prefix and new tokens are both already there and +# the kernel can read every key from the page table. +# +# Positions are absolute over the whole prompt (``num_cached + q_pos``), which +# is what makes the causal mask, the sliding window and the relative-bias +# distance come out right across a chunk boundary. ``rel_logits`` stays indexed +# by the PACKED query row: it is built per query token from that token's own +# hidden state, and only new tokens are queries. +# --------------------------------------------------------------------------- +# One prefill kernel, always over the paged cache. A packed variant used to sit +# beside it for the all-fresh case, on the theory that page indirection is worth +# avoiding there; measured at the serving level that cost did not show up above +# the baseline's own spread, and the second kernel cost a path that had to be +# proven bit-identical to the first. +@triton.jit +def _inkling_prefill_kernel( + Q, + K_Cache, + V_Cache, + Out, + RelLogits, + cu_seqlens, + num_cached, + page_table, + sm_scale, + stride_qt, + stride_qh, + stride_kp, + stride_kh, + stride_kt, + stride_vp, + stride_vh, + stride_vt, + stride_ot, + stride_oh, + stride_rt, + 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_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) + new_len = tl.load(cu_seqlens + cur_seq + 1) - seq_start + cached = tl.load(num_cached + cur_seq) + total_len = cached + new_len + + 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 # position within THIS chunk + mask_m = q_pos < new_len + q_abs = cached + q_pos # position within the whole prompt + + 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 in ABSOLUTE space: the last query row of this tile is + # cached + (cur_block_m + 1) * BLOCK_M - 1, so keys stop one past it. + end_n = tl.minimum(total_len, cached + (cur_block_m + 1) * BLOCK_M) + if WINDOW_LEFT >= 0: + lo = cached + 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): + k_pos = start_n + offs_n # ABSOLUTE key positions + mask_n = k_pos < total_len + + page_local = k_pos // page_size + tok_in_page = k_pos % page_size + page_id = tl.load(page_table + cur_seq * stride_ptb + page_local, mask=mask_n, other=0).to( + tl.int64 + ) + + # [BLOCK_DMODEL, BLOCK_N] for tl.dot, matching the packed prefill kernel. + 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) + qk = tl.dot(q, k, out_dtype=tl.float32) * sm_scale # [BLOCK_M, BLOCK_N] + + if HAS_REL: + rel_dist = q_abs[:, None] - k_pos[None, :] + rel_idx = tl.minimum(tl.maximum(rel_dist, 0), rel_extent - 1) + # Row is the PACKED query index: rel_logits covers new tokens only. + 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_abs[:, None] >= k_pos[None, :]) + if WINDOW_LEFT >= 0: + valid &= (q_abs[:, 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 = ( + 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) + 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, :]) + + +def inkling_prefill_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cu_seqlens: torch.Tensor, + num_cached: torch.Tensor, + page_table: torch.Tensor, + page_size: int, + 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 for requests that carry cached history. + + Same contract as :func:`inkling_prefill_attention`, except that K/V come + from the paged cache rather than a packed tensor, and each request declares + how many of its tokens are already there. + + Args: + q: ``[total_new_tokens, num_heads, head_dim]`` -- the chunk's queries. + k_cache, v_cache: ``[num_pages, num_kv_heads, page_size, head_dim]`` + HND pools. The caller must have written this chunk's K/V already. + cu_seqlens: ``[batch + 1]`` int32 cumulative NEW token counts. + num_cached: ``[batch]`` int32 tokens already in the cache per request. + All-zero is legal and must reproduce the packed kernel exactly. + page_table: ``[batch, max_pages]`` int32 physical page ids. + page_size: tokens per page. + max_seqlen: max NEW tokens for one request (host int; grid only). + rel_logits: ``[total_new_tokens, num_heads, rel_extent]`` fp32, indexed + by the packed query row, or None. + window_left: sliding-window radius (inclusive), -1 to disable. + + Returns ``[total_new_tokens, num_heads, head_dim]`` in q's dtype. + """ + q = q.contiguous() + _total_tokens, 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 = 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_cache, + v_cache, + o, + rel_arg, + cu_seqlens, + num_cached, + 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_st, + 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_M=BLOCK_M, + BLOCK_N=BLOCK_N, + Lk=head_dim, + num_warps=4, + num_stages=2, + ) + return o diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/metadata.py new file mode 100644 index 000000000000..d937f6389af1 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/metadata.py @@ -0,0 +1,33 @@ +# 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 metadata for Inkling: the base plus one pre-capture hook.""" + +from ...trtllm import TrtllmAttentionMetadata + + +class InklingAttentionMetadata(TrtllmAttentionMetadata): + """``TrtllmAttentionMetadata`` plus the short-conv pool's per-step slot write. + + Carries no fields of its own. The write is a host->device copy into a buffer + the captured decode graph aliases, so it has to run every step and outside + that region -- ``prepare()`` is the only hook that does both and still sees + the CUDA-graph padding rows. + """ + + def prepare(self) -> None: + super().prepare() + cache = getattr(self.kv_cache_manager, "conv_state_cache", None) + if cache is not None and self.request_ids is not None: + cache.write_state_indices(self.request_ids) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/page_table.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/page_table.py new file mode 100644 index 000000000000..890b405f867c --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/page_table.py @@ -0,0 +1,185 @@ +# 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. +"""Decode-time views into the attention metadata for Inkling's Triton kernels. + +The kernel needs each generation request's total KV length and page table through +fixed-pointer GPU buffers. Both are slices of ``kv_lens_cuda`` and +``kv_cache_block_offsets``, which ``prepare()`` already fills and keeps +graph-stable, so these are free functions rather than fields on a subclass. The +page table is laid out for the C++ attention op: plane 0 serves both K and V, but +its entries count pages while the kernel's views count blocks, hence +:func:`page_div`. +""" + +import os +from typing import List + +import torch + +from tensorrt_llm.logger import logger + +# Per-process latch: the crosscheck is silent on success, so announce it once. +_INK_XCHK_ANNOUNCED = False + + +def owned_layers(md) -> List[int]: + """Global decoder-layer indices this rank owns.""" + return list(getattr(md.kv_cache_manager, "pp_layers", [])) + + +def pt_row(md, layer: int) -> int: + """Row of ``kv_cache_block_offsets`` holding ``layer``'s page table. + + The tensor's leading axis is per pool by default but per attention-op layer + under ``enable_swa_scratch_reuse``; reading one as the other returns another + layer's pages. The scale check guards the same confusion: the C++ copy encodes + with the pool-level scale, which need not equal the layer's. + """ + mgr = md.kv_cache_manager + offset = mgr.layer_offsets[layer] + if getattr(mgr, "enable_swa_scratch_reuse", False): + return int(offset) + pool_id = int(mgr.layer_to_pool_mapping_dict[offset]) + layer_scale = int(mgr.get_layer_page_index_scale(layer)) + pool_scale = int(mgr.index_scales[pool_id]) + if layer_scale != pool_scale: + raise RuntimeError( + f"Inkling layer {layer} (offset {offset}) has page-index " + f"scale {layer_scale} but its pool {pool_id} is encoded with " + f"{pool_scale}; kv_cache_block_offsets cannot be reused as " + "this layer's page table because the recovered page address " + "would be wrong. Give the layer its own pool, or stage a " + "private table for it." + ) + return pool_id + + +def gen_page_table(md, layer: int) -> torch.Tensor: + """Generation-row page table for ``layer``: a view into the base + ``kv_cache_block_offsets`` K plane, entries scaled by :func:`page_div`.""" + row = pt_row(md, layer) + start = md.num_contexts + return md.kv_cache_block_offsets[row, start : start + md.num_generations, 0] + + +def page_div(md) -> int: + """Pages-to-blocks divisor: entries count pages (K and V separately), the + kernel's ``kv[:, 0]`` / ``kv[:, 1]`` views count blocks.""" + return int(getattr(md.kv_cache_manager, "kv_factor", 1)) + + +def gen_seq_lens(md, num_gen: int) -> torch.Tensor: + """Total-KV length per generation request, sliced off the base's + ``kv_lens_cuda``. Deliberately not ``kv_lens``: the host-side twin adds + ``num_extra_kv_tokens``, which this kernel must not see.""" + start = md.num_contexts + return md.kv_lens_cuda[start : start + num_gen] + + +def validate_decode_layout(md, layer: int, num_gen: int) -> None: + """Check the borrowed layout before the decode kernel reads it. + + Cheap enough to run per layer: a dict lookup plus three comparisons. + """ + num_contexts = md.num_contexts + # The captured kernel reads both borrowed buffers at a fixed generation + # offset, so it must be constant across replays -- 0 for a pure generation + # batch, which is what a decode graph is. + if md.is_cuda_graph and num_contexts: + raise RuntimeError( + f"Inkling got {num_contexts} context requests in a CUDA-graph " + "batch; the captured decode kernel reads seq lens at a fixed offset " + "into kv_lens_cuda, which only holds for pure generation batches" + ) + offsets = md.kv_cache_block_offsets + if offsets is None: + raise RuntimeError( + "Inkling has no kv_cache_block_offsets to borrow as its decode page " + "table. prepare() fills it for any batch that reaches a kernel, so " + "its absence is a setup error worth naming -- without this the next " + "line reports it as AttributeError on NoneType." + ) + row = pt_row(md, layer) + if row >= offsets.shape[0]: + raise RuntimeError( + f"Inkling needs row {row} of kv_cache_block_offsets but it has only " + f"{offsets.shape[0]} (num_attention_op_pools). The leading axis is " + "per-pool by default and per-attention-op-layer under " + "enable_swa_scratch_reuse; pt_row and the manager disagree about " + "which." + ) + if num_contexts + num_gen > offsets.shape[1]: + raise RuntimeError( + f"Inkling decode batch ({num_contexts} ctx + {num_gen} gen) exceeds " + f"kv_cache_block_offsets' {offsets.shape[1]} sequence rows " + "(max_num_sequences)." + ) + if os.environ.get("TLLM_INKLING_PT_CROSSCHECK") == "1": + layers = owned_layers(md) + # Once per step, not once per layer: the crosscheck walks every layer. + if layers and layer == layers[0]: + crosscheck_page_table(md, layers) + + +def crosscheck_page_table(md, layers: List[int]) -> None: + """Assert the borrowed table matches ``get_batch_cache_indices``. + + Gated on ``TLLM_INKLING_PT_CROSSCHECK=1`` because it rebuilds the host page + table every step; run it once on a real workload after touching the borrowed + layout. + """ + global _INK_XCHK_ANNOUNCED + mgr = md.kv_cache_manager + gen_ids = md.request_ids[md.num_contexts :] + div = page_div(md) + checked = 0 + for layer in layers: + borrowed = (gen_page_table(md, layer) // div).tolist() + for i, blocks in enumerate(mgr.get_batch_cache_indices(gen_ids, layer)): + want = [int(b) for b in blocks] + if any(b < 0 for b in want): + raise AssertionError( + f"layer {layer} request {gen_ids[i]}: get_batch_cache_indices " + f"returned an interior placeholder in {want}; the packed " + "representation this refactor replaced would have shifted " + "the survivors left, so the two are not equivalent here" + ) + got = borrowed[i][: len(want)] + if got != want: + raise AssertionError( + f"page-table mismatch at layer {layer} request " + f"{gen_ids[i]}: borrowed//{div}={got} vs " + f"get_batch_cache_indices={want}" + ) + # Reported, not asserted: the kernel masks every access by seq_len, + # so stale padding past the last block is harmless. + tail = borrowed[i][len(want) :] + if any(p != 0 for p in tail): + logger.warning( + f"Inkling page-table crosscheck: layer {layer} request " + f"{gen_ids[i]} carries stale values {tail[:8]} past its " + f"{len(want)} blocks. Harmless -- the kernel masks every " + "access by seq_len -- but the borrowed row is padded " + "differently from the staged one." + ) + checked += 1 + if not _INK_XCHK_ANNOUNCED: + _INK_XCHK_ANNOUNCED = True + logger.info( + f"INKLING_PT_CROSSCHECK ACTIVE: {checked} layer/request pairs " + f"agreed (borrowed kv_cache_block_offsets // {div} == " + f"get_batch_cache_indices) over {len(layers)} layers and " + f"{len(gen_ids)} generation rows" + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/inkling/params.py b/tensorrt_llm/_torch/attention_backend/sparse/inkling/params.py new file mode 100644 index 000000000000..954fd64f00e9 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/inkling/params.py @@ -0,0 +1,87 @@ +# 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 backend selection key and its per-forward backend inputs. + +Inkling is not sparse attention (see the package docstring); it reuses the +``sparse/`` seams because ``sparse/registry.py`` is where model-specific backends +and cache managers are selected, and ``sparse_backend_args`` is the registered +slot for a model-specific forward input. Consumers holding no ``ModelConfig`` +reach the backend only through the registry, so the config must be populated for +them to resolve anything but the family default. +""" + +from dataclasses import dataclass +from typing import Literal, Optional + +import torch + +from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig + +from ..params import SparseBackendForwardArgs, SparseParams + + +class InklingSparseParams(SparseParams): + """Selection key for Inkling's backend and cache manager in + ``sparse/registry.py``; the kernels' per-layer scalars live on the backend. + """ + + algorithm = "inkling" + + +class InklingSparseAttentionConfig(BaseSparseAttentionConfig): + """Architecture-derived config that carries :class:`InklingSparseParams`. + + Subclasses the interface but is deliberately not a member of the user-facing + ``SparseAttentionConfig`` union: there is no knob to expose, and the field is + injected by ``ModelConfig.from_pretrained`` keyed on architecture. + """ + + algorithm: Literal["inkling"] = "inkling" + + def to_sparse_params(self, pretrained_config=None, layer_idx=None) -> InklingSparseParams: + return InklingSparseParams() + + +@dataclass(kw_only=True, slots=True) +class InklingBackendForwardArgs(SparseBackendForwardArgs): + """Per-forward inputs ``InklingAttention`` hands to its backend, riding + ``AttentionForwardArgs.sparse_backend_args``. + + ``rel_logits`` is a ``[num_query_tokens, local_heads, rel_extent]`` fp32 + additive bias projected from the hidden states -- content-dependent, so the + shared ``relative_attention_bias`` (T5's position-indexed table) cannot carry + it. ``allow_mixed`` certifies the short-conv state pool is active, which is + what makes a mixed context+generation batch legal. + """ + + rel_logits: torch.Tensor + allow_mixed: bool = False + + +def inkling_forward_args( + rel_logits: torch.Tensor, + *, + allow_mixed: bool, + output: Optional[torch.Tensor] = None, +): + """Build the ``AttentionForwardArgs`` carrying Inkling's backend inputs.""" + from ...interface import AttentionForwardArgs + + return AttentionForwardArgs( + output=output, + sparse_backend_args=InklingBackendForwardArgs( + rel_logits=rel_logits, allow_mixed=allow_mixed + ), + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/registry.py b/tensorrt_llm/_torch/attention_backend/sparse/registry.py index 298462b0f409..40a4579167fc 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/registry.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/registry.py @@ -26,9 +26,17 @@ def get_sparse_attn_kv_cache_manager( from .deepseek_v4 import DeepseekV4CacheManager from .dsa import DSACacheManager, DSACacheManagerV2 + from .inkling import InklingHybridCacheManager from .minimax_m3 import MiniMaxM3KVCacheManagerV2 from .rocket import RocketKVCacheManager + if sparse_attention_config.algorithm == "inkling": + # Reached because Inkling populates `sparse_attention_config` (see + # `.inkling.params`), which makes `get_kv_cache_manager_cls` take its + # sparse branch before the `is_inkling` route in + # `_non_hybrid_kv_cache_manager_cls`. Without this entry that branch + # raises "Unsupported sparse attention algorithm" at engine startup. + return InklingHybridCacheManager if sparse_attention_config.algorithm == "rocket": return RocketKVCacheManager elif sparse_attention_config.algorithm == "dsa": @@ -85,8 +93,17 @@ def get_trtllm_sparse_attn_attention_backend( from .deepseek_v4 import DeepseekV4TrtllmAttention from .dsa import DSATrtllmAttention + from .inkling import InklingTritonAttention from .rocket import RocketTrtllmAttention + if sparse_params.algorithm == "inkling": + # Inkling is NOT a sparse-attention algorithm -- its attention is dense + # (full causal on global layers, a sliding window on local ones) with a + # learned relative-bias score_mod. It is dispatched here because this is + # where every model-specific backend is selected, and routing it through + # a branch in `get_attention_backend` instead left consumers that hold no + # ModelConfig resolving the family default. See `.inkling.params`. + return InklingTritonAttention if sparse_params.algorithm == "rocket": return RocketTrtllmAttention elif sparse_params.algorithm == "dsa": diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index 9110469e5fbb..ae59fb6865c5 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.kimi_k3 import KimiK3Config, KimiK3VisionConfig from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig from tensorrt_llm._torch.configs.laguna import LagunaConfig @@ -73,6 +74,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. @@ -96,6 +99,8 @@ def _register_custom_configs_with_transformers() -> None: "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", "Gemma4UnifiedVisionConfig", + "InklingConfig", + "InklingTextConfig", "KimiK3Config", "KimiK3VisionConfig", "KimiLinearConfig", diff --git a/tensorrt_llm/_torch/configs/inkling.py b/tensorrt_llm/_torch/configs/inkling.py new file mode 100644 index 000000000000..c853c9d3c99f --- /dev/null +++ b/tensorrt_llm/_torch/configs/inkling.py @@ -0,0 +1,230 @@ +# 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. + +Transformers does not ship Inkling, so these classes reconstruct the config from +the checkpoint's nested dicts without a transformers shim. The audio, vision and +MTP sub-configs are kept verbatim as ``PretrainedConfig`` blobs, which the towers +read their geometry off directly. + +Field names mirror the checkpoint ``text_config`` and the numeric defaults are the +real checkpoint values, overridable through ``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. + + ``KVCacheManagerV2`` takes this list as ``num_kv_heads`` so the paged pool + is sized per layer; ``head_dim`` is uniform, so only the 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, + # In-vocab chat-template placeholder ids, omitted from the checkpoint's + # config.json. They must be in-vocab: the executor rejects out-of-range + # request token ids. + 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 -- the draft weights stay deferred. + 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 b3ee33c3cca1..88d120bb1297 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -72,6 +72,15 @@ "MiniMaxM3SparseForConditionalGeneration", } +# Inkling is NOT a sparse-attention model. Populating `sparse_attention_config` +# for these architectures is what makes `sparse/registry.py` the single selection +# path for its backend and cache manager, including for consumers that hold no +# ModelConfig; see `attention_backend/sparse/inkling/params.py`. +_INKLING_ARCHITECTURES = { + "InklingForCausalLM", + "InklingForConditionalGeneration", +} + def _is_lock_infra_error(exc: BaseException) -> bool: """Whether exc indicates broken lock infrastructure (not mere contention).""" @@ -526,15 +535,23 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, quant_config = QuantConfig() layer_quant_config = None - # ``canonicalize_quant_algo`` is applied again here (and per layer - # below) because the ``quant_cfg.json`` overlay merged in for - # MIXED_PRECISION bypasses ``read_modelopt_quant_config``. - quant_config.quant_algo = ( - QuantAlgo(canonicalize_quant_algo(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, *, canonicalize=False): + # modelopt hf_quant_config.json may spell "no quantization" as JSON + # null or as the string "none"/"null"; both must map to None rather + # than QuantAlgo("none"), which is not a member. ``canonicalize`` is + # applied for quant_algo (and per layer below) because the + # ``quant_cfg.json`` overlay merged in for MIXED_PRECISION bypasses + # ``read_modelopt_quant_config``. + if value is None or (isinstance(value, str) and + value.strip().lower() in ("none", "null", "")): + return None + return QuantAlgo( + canonicalize_quant_algo(value) if canonicalize else value) + + quant_config.quant_algo = _algo_or_none( + json_quant_configs.get('quant_algo'), canonicalize=True) + 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) @@ -558,10 +575,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( @@ -1184,6 +1199,17 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): window_size=window_size, indexer_k_dtype=indexer_k_dtype, **indexer_config) + elif (pretrained_config.architectures + and pretrained_config.architectures[0] + in _INKLING_ARCHITECTURES + and kwargs.get('sparse_attention_config') is None): + # Architecture-derived, never user-supplied: Inkling has one + # correct backend and no alternative to choose between. + from .attention_backend.sparse.inkling import \ + InklingSparseAttentionConfig + kwargs[ + 'sparse_attention_config'] = InklingSparseAttentionConfig( + ) else: raise ValueError( "checkpoint_dir is None. Cannot load model config without a valid checkpoint directory." diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 4abb9565b7f8..d323ffaeaea2 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -48,6 +48,8 @@ "HCXVisionForCausalLM", "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", + "InklingForCausalLM", + "InklingForConditionalGeneration", "KimiK25ForConditionalGeneration", "KimiK3ForConditionalGeneration", "KimiLinearForCausalLM", diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index 8915af4c6475..becdc6dfb0d4 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -62,6 +62,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HCXVisionModel": "modeling_hyperclovax", "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "InklingForConditionalGeneration": "modeling_inkling", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", @@ -140,6 +141,8 @@ def is_builtin_zoo_module(module_name: str) -> bool: "HCXVisionForCausalLM": "modeling_hyperclovax", "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "InklingForCausalLM": "modeling_inkling", + "InklingForConditionalGeneration": "modeling_inkling", "KimiK25ForConditionalGeneration": "modeling_kimi_k25", "KimiK3ForConditionalGeneration": "modeling_kimi_k3_vl", "KimiLinearForCausalLM": "modeling_kimi_linear", @@ -203,6 +206,7 @@ def is_builtin_zoo_module(module_name: str) -> bool: "gemma4": "modeling_gemma4mm", "gemma4_unified": "modeling_gemma4_unified", "hyperclovax_vlm": "modeling_hyperclovax", + "inkling_mm_model": "modeling_inkling", "kimi_k25": "modeling_kimi_k25", "kimi_k3": "modeling_kimi_k3_vl", "llama4": "modeling_llama", 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..3378ee119079 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/inkling_weight_mapper.py @@ -0,0 +1,375 @@ +# 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. :func:`inkling_expected_text_keys` and + :func:`inkling_account_checkpoint` derive the exact set of ``model.llm.*`` keys + the text loader consumes and classify every checkpoint key as consumed-text / + deferred (audio, vision, MTP) / unaccounted. + +2. Name and layout remapping. :class:`InklingHfWeightMapper` renames the + checkpoint's keys to the TRT-LLM module tree, fuses q/k/v into ``qkv_proj``, + and unfuses the NVFP4 routed experts into the fused-MoE loader's layout. +""" + +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. ModelOpt stores the per-tensor activation +# ``input_scale`` as ``amax / (E2M1_MAX * E4M3_MAX)``; Inkling's checkpoint ships +# the raw ``.input_amax``, so ``_map_expert`` applies the conversion. +_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], quantized: bool = True) -> bool: + """Routed experts of an MoE layer are NVFP4 unless explicitly excluded. + + ``quantized=False`` for the BF16 release, where the exclusion list is empty + because nothing is quantized rather than because everything is. + """ + if not quantized: + return False + return f"model.llm.layers.{layer_idx}.mlp.experts" not in exclude_modules + + +def inkling_expected_text_keys( + config: InklingTextConfig, exclude_modules: Set[str], quantized: bool = True +) -> Set[str]: + """Exact set of ``model.llm.*`` checkpoint keys the text loader consumes. + + ``quantized=False`` for the BF16 release, which ships no scale sidecars. + """ + 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, quantized): + 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], + quantized: bool = True, +) -> 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 accounted for. + """ + expected = inkling_expected_text_keys(config, exclude_modules, quantized) + 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], quantized: bool = True +) -> 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, quantized) + ] + + +# --------------------------------------------------------------------------- +# 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; the fused ``qkv_proj`` loader +# collects and fuses them. Same for gate_up_proj <- gate_proj + up_proj. +_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 a gate/up-interleaved fused tensor into ``(gate, up)`` strided views + along ``dim``: gate = even indices, up = odd indices. + + The checkpoint interleaves the two projections along the output dim, so a + contiguous ``chunk(2)`` would pair the wrong channels in every SwiGLU. The + result is a view rather than a copy because the fused-MoE / gate_up loaders + call ``.contiguous()`` on the small per-rank shard themselves; de-interleaving + eagerly instead OOM-killed the TP=4 load. It reorders whole output rows, so it + holds for a packed NVFP4 weight and its per-block scale alike. + """ + 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`` of its own. + """ + 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 is padded to vocab_size while the + # tower emits logits only over the unpadded vocab. + # embed_tokens keeps the full matrix so input ids stay valid. + 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 along the output dim; split + # into gate (even rows) / up (odd rows) as strided views. + 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); the + # interleave is undone by the strided split in + # ``InklingSharedExperts.forward``. + + 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 gate/up-interleaved along the output dim, so ``w1`` + is the even rows and ``w3`` the odd ones (see + :func:`_split_interleaved_gate_up`); ``w2_weight[e]`` is the down + projection. NVFP4 sidecars map to the fused-MoE scale names, and + ``.original_shape`` is dropped as layout metadata. + """ + 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": + # The checkpoint stores a raw amax, but the fused-MoE loader expects + # ModelOpt's input_scale = amax / (E2M1_MAX * E4M3_MAX); without the + # conversion the activation block scales underflow e4m3. + 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, per-expert + # scalar weight_scale_2 (shared by 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": + # Reorders whole rows, so this holds for the packed fp4 + # weight and its per-block fp8 scale alike. + 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..b81598223f79 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_inkling.py @@ -0,0 +1,1334 @@ +# 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/sparse/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 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.sparse.inkling import ( + InklingConvRuntime, + InklingConvState, + apply_short_conv, + inkling_forward_args, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.hf.inkling_weight_mapper import InklingHfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import ( + DecoderModel, + DecoderModelForCausalLM, + MetaInitException, + 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 + +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, +) + + +def _module_excluded_from_quant(model_config: ModelConfig, name: str) -> bool: + """True if ``name`` (or an ancestor) is bf16, not NVFP4. + + The checkpoint lists its bf16 modules in ``quant_config.exclude_modules``, and + the lookup walks the dotted ancestry, so a listed ``...layers.5.attn`` covers + the projections under it. + """ + 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, which the stock routing methods cannot express. ``apply`` returns + only the routed pair the fused MoE needs; :class:`InklingMoE` recomputes the + shared gammas from the same 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. + + Prefill runs :func:`causal_conv1d_fn`, cached decode + :func:`causal_conv1d_update` against the per-request conv state; with no + ``conv_state`` the module falls back to a self-contained causal convolution. + + Under ``tp_shard`` the channels follow the fused qkv projection's kv-head + sharding and :meth:`load_weights` slices the rank's block out of the + checkpoint weight. The residual-stream convs are replicated. + """ + + 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. + + Replicated convs run with ``tp_size == 1`` and copy the full tensor; the + sharded k/v convs take the rank's contiguous, kv-head-aligned block. + """ + 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 branch runs in fp32 (per the source); the fused cached + branches run in the input dtype, since the ``causal_conv1d`` ops require + ``weight.dtype == x.dtype``. ``conv_state`` is updated in place. + """ + in_dtype = x.dtype + residual = x + w = self.weight.squeeze(1).to(x.dtype) # [channels, kernel] + if conv_state is not None and is_decode: + # causal_conv1d_update writes in place into its x argument, so pass a + # copy -- otherwise the internal residual becomes conv(x) + conv(x). + 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 learned relative bias is a per-(query, head, + relative-distance) additive ``score_mod`` that no fused backend exposes, so it + is precomputed here into ``rel_logits`` and gathered+added by the Triton + kernels. Local layers apply the sliding window natively in the kernel; global + layers fold the log-scaling ``tau`` into ``rel_logits``. + + The KV write, page-table construction and prefill/decode dispatch run in + :meth:`InklingTritonAttention.forward`, reached through the standard backend + contract. ``forward`` overrides the base because the k/v short-convs must run + between ``split_qkv`` and ``apply_qk_norm``, and ``rel_logits`` is projected + from ``hidden_states``, which ``forward_impl`` never receives. + """ + + 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 is bf16, not NVFP4: hand the base a shallow ModelConfig copy + # with an empty quant_config so qkv_proj/o_proj are built unquantized. + 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: forward runs the Inkling Triton attention directly, so + # this keeps the base from building an unused RotaryEmbedding. + 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. + # The backend uses 1/(sqrt(head_dim) * q_scaling). + q_scaling=float(head_dim) ** 0.5, + skip_rope=True, + fuse_qk_norm_rope=False, + is_qk_norm=True, + ) + # head_dim differs from hidden_size // num_heads, so the base Attention + # must have picked it up from config.head_dim (there is no kwarg for it). + assert self.head_dim == head_dim, (self.head_dim, head_dim) + + # Applied directly by the Triton kernels. The sliding window uses an + # inclusive radius: 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 + # These three differ between local and global layers, and + # create_attention() has no passthrough for model-specific kwargs. + self.attn.sm_scale = self.sm_scale + self.attn.rel_extent = self.rel_extent + self.attn.window_left = self.window_left + + # Attention-scoped TP: the Inkling-only tensors below (r_proj, k/v sconv) + # hang off the same head/kv-head split as qkv_proj, so they must follow + # the attention TP rather than the global one. + tp_size = 1 if model_config.mapping.enable_attention_dp else model_config.mapping.tp_size + 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, sharded by head like q and not + # gathered (consumed locally to build the bias). Replicated 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), + # so this must use ``self.rel_extent``, not ``config.rel_extent``. + self.rel_logits_proj = nn.Parameter(torch.empty(self.d_rel, self.rel_extent)) + # k/v short convs act on the k/v stream of the fused qkv projection and + # are sharded by kv-head like it; InklingShortConv slices at load. + 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 + + 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`` + ``conv_rt`` the + k/v short-convs run through the runtime state pool; 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_short_conv(self.k_sconv, k, pool_k, conv_rt) + v = apply_short_conv(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``. Global layers fold in the + per-query-token log-scaling ``tau``. The Triton kernels index this by + ``clamp(q_pos - k_pos, 0, rel_extent - 1)`` and zero it outside the range. + """ + 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 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`` + ``conv_rt`` drive the k/v short-convs through the + runtime state pool; without them they run stateless over the sequence. + """ + num_tokens = hidden_states.shape[0] + # The pre-attention RMSNorm can emit fp32 while the attention/r + # projections are bf16, so cast once here. + 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) + # Standard backend contract; rel_logits and the mixed-batch certificate + # ride AttentionForwardArgs.sparse_backend_args (see inkling/params.py). + attn_out = self.attn.forward( + q, + k, + v, + attn_metadata, + forward_args=inkling_forward_args(rel_logits, 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: keeping the + # column/row split would all-reduce partials from different requests. + 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 fp32 and promotes the output; cast back so the + # residual stream stays in the input dtype. + 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. Model dtype: these run as raw bmms and load unquantized. + 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]; gammas: [T, n_shared] fp32, applied after + # the (linear) down projection where it commutes. + 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 with gate/up interleaved along its output dim, so + # gate = even channels and up = odd; chunk(2) would pair the wrong ones. + 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: under TP each rank holds a shard of the experts and + # produces only a partial routed sum, so it must be all-reduced before the + # replicated shared-expert output is added on top. + 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, in + which case an empty ``QuantConfig`` gives ``create_moe`` a bf16 MoE.""" + 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 this step's + per-rank token count, which ``FusedMoE`` needs to gather under DP.""" + 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) + # fp32 scales in the routed/shared paths can promote the sum; keep the + # residual-stream dtype. + 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 pool + buffers and each short-conv runs through them; without it they 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_short_conv(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_short_conv(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 short-conv pool comes from the cache manager and its + per-forward split from the state metadata the base ``prepare()`` built; a + metadata without one keeps the stateless conv. + + ``inputs_embeds_prenormed`` is set on the multimodal path, where the + wrapper has already applied ``embed_norm`` to the text embeddings before + scattering the raw tower rows in.""" + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + # Views only, no copy: the slot write happened in prepare(), outside the + # captured region. conv_rt is the gate for the pool path. + conv_rt = InklingConvRuntime.from_metadata(attn_metadata) + conv_cache = ( + getattr(attn_metadata.kv_cache_manager, "conv_state_cache", None) + if conv_rt is not None + else None + ) + # Per-rank token counts for this step, set on attn_metadata only under + # attention DP; FusedMoE reads it to pad and gather across ranks. + 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 attention backend family was overridden. + + Inkling's backend is registered in ``sparse/registry.py`` under the + ``TRTLLM`` family only; the ``VANILLA`` / ``FLASHINFER`` registries do not + know the ``"inkling"`` algorithm and fail with a message naming neither + Inkling nor the setting responsible. + """ + backend = getattr(model_config, "attn_backend", None) + if backend is not None and str(backend).upper() != "TRTLLM": + raise ValueError( + f"Inkling requires the TRTLLM attention backend family (got " + f"{backend!r}). Its backend is registered in " + "attention_backend/sparse/registry.py under that family only, " + "and its Triton decode kernel reads per-step seq_lens and page " + "tables off TrtllmAttentionMetadata. Remove the attn_backend " + "override from --extra_llm_api_options / LLM(attn_backend=...)." + ) + + @staticmethod + def _assert_inkling_moe_parallel(model_config) -> None: + """Reject an expert-parallel layout the MoE backend cannot serve. + + The CUTLASS backend -- the only routed-expert backend Inkling ships -- + does not opt into ``_supports_non_divisible_ep``, so a non-divisible + ``moe_expert_parallel_size`` fails inside expert-slot bookkeeping rather + than at load. + """ + 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 + # Report "more ranks than experts" first: it subsumes non-divisibility, + # and "pick a divisor" is not useful advice in that case. + 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}." + ) + # Pure EP (moe_tp_size 1) reproduces the TP-only accuracy per item but + # segfaults during CUDA-graph capture; root cause not yet found, so + # reject only that combination rather than the whole 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; Inkling's decode message size hits that case and decode collapses + to a repeated token. The affected all-reduces are built by generic + modules, so rebuilding them here keeps the mitigation model-local. + """ + for mod in self.modules(): + old = getattr(mod, "all_reduce", None) + if not isinstance(old, AllReduce): + continue + # Carry the module's own mapping and dtype over so the rebuild + # 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 reaches the decoder through attn_metadata, so + # there are no conv kwargs and no ResourceManager lookup here. + 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. + + Returns ``[feats]`` of shape ``(sum_patches, decoder_dmodel)`` -- the shape + ``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. + + Returns ``[feats]`` of shape ``(sum_frames, decoder_dmodel)``, the same shape + contract as the image encoder, or ``[]`` when there is no audio.""" + 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)] + + +def _has_meta_tensors(module: nn.Module) -> bool: + return any(getattr(p, "is_meta", False) for p in module.parameters()) or any( + getattr(b, "is_meta", False) for b in module.buffers() + ) + + +def _build_replicated_bf16_tower(tower_cls, tower_config): + """Build a replicated bf16 media tower, or defer it under ``MetaInitMode``. + + Returns ``(tower, deferred_config)``: exactly one is ``None``. A non-``None`` + ``deferred_config`` means ``load_weights`` must rebuild the tower before + loading into it. + + The deferral is what keeps meta init working: ``MetaInitMode`` permits only + random-init ops on meta parameters, so the tower's ``.to(torch.bfloat16)`` + raises ``MetaInitException``. The loader wraps that mode around the entire + ``from_config`` call, so one tower raising there silently costs the whole text + stack its meta init. ``modeling_kimi_k3_vl`` defers its encoder for the same + reason. + """ + if tower_config is None or not getattr(tower_config, "decoder_dmodel", None): + return None, None + try: + tower = tower_cls(tower_config) + except MetaInitException: + return None, tower_config + # Constructing may succeed while leaving meta parameters behind; the + # `.to(bfloat16)` below is what would raise on them, so check first. + if _has_meta_tensors(tower): + return None, tower_config + return tower.to(torch.bfloat16), None + + +@register_auto_model("InklingForConditionalGeneration") +@register_input_processor( + InklingInputProcessor, + model_type="inkling_mm_model", + placeholder_metadata=MultimodalPlaceholderMetadata( + # Image and audio are distinct modalities; video has no separate token or + # tower -- its frames are served as ordinary ```` placeholders. + placeholder_map={"image": "", "audio": "