From 385855c85aee551db2b78604ae58814b1a20993e Mon Sep 17 00:00:00 2001 From: Michael Williams Date: Tue, 15 Sep 2026 18:38:49 -0700 Subject: [PATCH] Add Kimi-K3 text prefill exporter --- tools/hf/exporters/models/__init__.py | 1 + tools/hf/exporters/models/kimi/__init__.py | 1 + tools/hf/exporters/models/kimi/helpers.py | 78 ++++ tools/hf/exporters/models/kimi/patches.py | 350 ++++++++++++++++++ tools/hf/exporters/models/kimi/spec.py | 155 ++++++++ tools/hf/exporters/plugin/kimi_kda.py | 49 +++ tools/hf/exporters/plugin/plugin_converter.py | 32 ++ tools/hf/exporters/plugin/plugin_utils.py | 2 + tools/hf/run_kimi_export.py | 66 ++++ 9 files changed, 734 insertions(+) create mode 100644 tools/hf/exporters/models/kimi/__init__.py create mode 100644 tools/hf/exporters/models/kimi/helpers.py create mode 100644 tools/hf/exporters/models/kimi/patches.py create mode 100644 tools/hf/exporters/models/kimi/spec.py create mode 100644 tools/hf/exporters/plugin/kimi_kda.py create mode 100644 tools/hf/run_kimi_export.py diff --git a/tools/hf/exporters/models/__init__.py b/tools/hf/exporters/models/__init__.py index eaebd1daeb8..182bbd3efa4 100644 --- a/tools/hf/exporters/models/__init__.py +++ b/tools/hf/exporters/models/__init__.py @@ -1,5 +1,6 @@ """Model families. Importing this package registers EdgeSpecs.""" from .groot import spec as _groot # noqa: F401 +from .kimi import spec as _kimi # noqa: F401 from .nemotron import spec as _nemotron # noqa: F401 from .pi05 import spec as _pi05 # noqa: F401 diff --git a/tools/hf/exporters/models/kimi/__init__.py b/tools/hf/exporters/models/kimi/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tools/hf/exporters/models/kimi/__init__.py @@ -0,0 +1 @@ + diff --git a/tools/hf/exporters/models/kimi/helpers.py b/tools/hf/exporters/models/kimi/helpers.py new file mode 100644 index 00000000000..d28e9f43933 --- /dev/null +++ b/tools/hf/exporters/models/kimi/helpers.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn + + +def language_model(model: nn.Module) -> nn.Module: + language = getattr(model, "language_model", None) + if not isinstance(language, nn.Module): + raise AttributeError(f"{type(model).__name__} has no language_model") + return language + + +def decoder_model(model: nn.Module) -> nn.Module: + decoder = getattr(language_model(model), "model", None) + if not isinstance(decoder, nn.Module) or not hasattr(decoder, "layers"): + raise AttributeError(f"{type(model).__name__} has no Kimi decoder layers") + return decoder + + +def kda_layer_indices(model: nn.Module) -> list[int]: + return [ + index + for index, layer in enumerate(decoder_model(model).layers) + if bool(getattr(layer, "is_linear_attn", False)) + ] + + +def allocate_kda_states( + model: nn.Module, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[tuple[torch.Tensor, ...], list[str]]: + """Allocate Q/K/V convolution state plus V-first recurrent state per KDA layer.""" + states: list[torch.Tensor] = [] + names: list[str] = [] + decoder = decoder_model(model) + + for layer_index in kda_layer_indices(model): + attention = decoder.layers[layer_index].self_attn + num_heads = int(attention.num_heads) + head_dim = int(attention.head_dim) + projection_size = num_heads * head_dim + conv_size = int(attention.conv_size) + + for suffix in ("q", "k", "v"): + states.append( + torch.zeros( + batch_size, + projection_size, + conv_size, + device=device, + dtype=dtype, + ) + ) + names.append(f"kda_conv_{suffix}_{layer_index}") + + states.append( + torch.zeros( + batch_size, + num_heads, + head_dim, + head_dim, + device=device, + dtype=torch.float32, + ) + ) + names.append(f"kda_recurrent_{layer_index}") + + return tuple(states), names + + +def text_config(model: nn.Module) -> Any: + return language_model(model).config diff --git a/tools/hf/exporters/models/kimi/patches.py b/tools/hf/exporters/models/kimi/patches.py new file mode 100644 index 00000000000..346fa209d13 --- /dev/null +++ b/tools/hf/exporters/models/kimi/patches.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +from contextlib import ExitStack, contextmanager +from typing import Any, Callable, Iterator + +import torch +import torch.nn.functional as F + +from ...plugin.attn_patches import patch_attribute +from .helpers import decoder_model, language_model + + +def _conv_bias(conv, projected: torch.Tensor) -> torch.Tensor: + bias = getattr(conv, "bias", None) + if bias is not None: + return bias.to(device=projected.device, dtype=projected.dtype) + return torch.zeros( + int(conv.weight.shape[0]), + device=projected.device, + dtype=projected.dtype, + ) + + +def _plugin_convolution( + conv, + projected: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kernel_size = int(conv.weight.shape[-1]) + output, next_state = torch.ops.trt.causal_conv1d.default( + projected, + conv.weight, + _conv_bias(conv, projected), + state, + context_lengths, + 1, + kernel_size - 1, + 1, + int(projected.shape[-1]), + ) + return F.silu(output), next_state + + +def _patch_kda_attention(original: Callable) -> Callable: + def forward( + self, + hidden_states, + attention_mask=None, + cache_params=None, + *, + context_lengths=None, + conv_state_q=None, + conv_state_k=None, + conv_state_v=None, + recurrent_state=None, + **kwargs, + ): + if recurrent_state is None: + return original( + self, + hidden_states, + attention_mask=attention_mask, + cache_params=cache_params, + **kwargs, + ) + + del attention_mask, cache_params, kwargs + if context_lengths is None: + raise ValueError("Kimi KDA export requires context_lengths") + if conv_state_q is None or conv_state_k is None or conv_state_v is None: + raise ValueError("Kimi KDA export requires Q/K/V convolution states") + + q, _ = _plugin_convolution( + self.q_conv1d, + self.q_proj(hidden_states), + conv_state_q, + context_lengths, + ) + k, _ = _plugin_convolution( + self.k_conv1d, + self.k_proj(hidden_states), + conv_state_k, + context_lengths, + ) + v, _ = _plugin_convolution( + self.v_conv1d, + self.v_proj(hidden_states), + conv_state_v, + context_lengths, + ) + + batch_size, seq_len, _ = hidden_states.shape + num_heads = int(self.num_heads) + head_dim = int(self.head_dim) + q = q.reshape(batch_size, seq_len, num_heads, head_dim) + k = k.reshape(batch_size, seq_len, num_heads, head_dim) + v = v.reshape(batch_size, seq_len, num_heads, head_dim) + gate = self.f_b_proj(self.f_a_proj(hidden_states)).reshape( + batch_size, seq_len, num_heads, head_dim + ) + beta = self.b_proj(hidden_states).to(dtype=hidden_states.dtype) + lower_bound = getattr(self, "gate_lower_bound", None) + + output, _ = torch.ops.trt.kimi_kda_plugin.default( + q, + k, + v, + gate, + beta, + self.A_log.float(), + self.dt_bias.reshape(num_heads, head_dim).to(hidden_states.dtype), + recurrent_state, + context_lengths, + float(lower_bound if lower_bound is not None else -5.0), + lower_bound is not None, + ) + + output_gate = ( + self.g_proj(hidden_states) + if self.use_full_rank_gate + else self.g_b_proj(self.g_a_proj(hidden_states)) + ).reshape(batch_size, seq_len, num_heads, head_dim) + + output_float = output.float() + eps = float( + getattr( + self.o_norm, + "variance_epsilon", + getattr(self.o_norm, "eps", 1e-5), + ) + ) + output = ( + output_float + * torch.rsqrt(output_float.square().mean(-1, keepdim=True) + eps) + ).to(hidden_states.dtype) + output = output * self.o_norm.weight * output_gate.sigmoid() + return self.o_proj(output.reshape(batch_size, seq_len, -1)) + + return forward + + +def _patch_mla_attention(original: Callable) -> Callable: + def forward( + self, + hidden_states, + attention_mask=None, + position_ids=None, + past_key_values=None, + **kwargs, + ): + del position_ids, past_key_values, kwargs + batch_size, seq_len, _ = hidden_states.shape + query_shape = (batch_size, seq_len, -1, self.q_head_dim) + key_shape = ( + batch_size, + seq_len, + -1, + self.qk_nope_head_dim + self.v_head_dim, + ) + + if self.q_lora_rank is not None: + query = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + else: + query = self.q_proj(hidden_states) + query = query.reshape(query_shape).transpose(1, 2) + query_nope, query_position = torch.split( + query, + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=-1, + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + key_latent, key_position = torch.split( + compressed_kv, + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + key_value = self.kv_b_proj(self.kv_a_layernorm(key_latent)) + key_value = key_value.reshape(key_shape).transpose(1, 2) + key_nope, value = torch.split( + key_value, + [self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + key_position = key_position.reshape( + batch_size, 1, seq_len, self.qk_rope_head_dim + ).expand(*key_nope.shape[:-1], -1) + query = torch.cat((query_nope, query_position), dim=-1) + key = torch.cat((key_nope, key_position), dim=-1) + + scores = torch.matmul(query, key.transpose(-2, -1)) * float(self.scaling) + if attention_mask is not None: + scores = scores + attention_mask[..., : key.shape[-2]] + probabilities = scores.softmax(dim=-1, dtype=torch.float32).to(query.dtype) + output = torch.matmul(probabilities, value).transpose(1, 2).contiguous() + output = output.reshape(batch_size, seq_len, -1) + if self.use_output_gate: + output = output * self.g_proj(hidden_states).sigmoid() + return self.o_proj(output) + + return forward + + +def _patch_sparse_moe(original: Callable) -> Callable: + def forward(self, hidden_states): + identity = hidden_states + original_shape = hidden_states.shape + topk_indices, topk_weights = self.gate(hidden_states) + flattened = hidden_states.reshape(-1, hidden_states.shape[-1]) + + if self.use_latent_moe: + flattened = self.routed_expert_down_proj(flattened) + + expert_outputs = torch.stack( + [expert(flattened) for expert in self.experts], + dim=1, + ) + routing = F.one_hot( + topk_indices, + num_classes=int(self.num_experts), + ).to(topk_weights.dtype) + routing = (routing * topk_weights.unsqueeze(-1)).sum(dim=1) + output = (expert_outputs * routing.unsqueeze(-1)).sum(dim=1) + output = output.to(flattened.dtype) + + if self.use_latent_moe: + if self.latent_moe_use_norm: + output = self.routed_expert_norm(output) + output = self.routed_expert_up_proj(output) + output = output.reshape(original_shape) + if self.config.num_shared_experts is not None: + output = output + self.shared_experts(identity).to(output.dtype) + return output + + return forward + + +def _causal_mask( + context_lengths: torch.Tensor, + seq_len: int, + dtype: torch.dtype, +) -> torch.Tensor: + positions = torch.arange(seq_len, device=context_lengths.device) + causal = positions.unsqueeze(0) <= positions.unsqueeze(1) + valid_keys = positions.unsqueeze(0) < context_lengths.unsqueeze(1) + allowed = causal.unsqueeze(0) & valid_keys.unsqueeze(1) + mask = torch.zeros( + context_lengths.shape[0], + 1, + seq_len, + seq_len, + device=context_lengths.device, + dtype=dtype, + ) + return mask.masked_fill(~allowed.unsqueeze(1), torch.finfo(dtype).min) + + +def _patch_language_forward(original: Callable) -> Callable: + def forward( + self, + inputs_embeds=None, + context_lengths=None, + last_token_ids=None, + *kda_states, + **kwargs, + ): + if context_lengths is None: + return original(self, inputs_embeds=inputs_embeds, **kwargs) + if inputs_embeds is None or last_token_ids is None: + raise ValueError("Kimi export requires embeddings and last-token indices") + + decoder = self.model + hidden = inputs_embeds.to(dtype=next(decoder.parameters()).dtype) + batch_size, seq_len, hidden_size = hidden.shape + causal_mask = _causal_mask(context_lengths, seq_len, hidden.dtype) + block_residual = hidden.new_zeros(batch_size * seq_len, 0, hidden_size) + state_index = 0 + + for layer in decoder.layers: + layer_kwargs: dict[str, Any] = {} + layer_mask = causal_mask + if layer.is_linear_attn: + if state_index + 4 > len(kda_states): + raise ValueError("Missing flattened Kimi KDA states") + layer_kwargs = { + "context_lengths": context_lengths, + "conv_state_q": kda_states[state_index], + "conv_state_k": kda_states[state_index + 1], + "conv_state_v": kda_states[state_index + 2], + "recurrent_state": kda_states[state_index + 3], + } + state_index += 4 + layer_mask = None + + hidden, block_residual = layer( + hidden, + attention_mask=layer_mask, + block_residual=block_residual, + **layer_kwargs, + ) + + hidden = decoder._apply_output_attn_res(hidden, block_residual) + hidden = decoder.norm(hidden) + indices = ( + last_token_ids if last_token_ids.ndim == 1 else last_token_ids.squeeze(-1) + ) + selected = hidden[ + torch.arange(batch_size, device=hidden.device, dtype=torch.long), + indices, + ] + return self.lm_head(selected).float() + + return forward + + +@contextmanager +def apply_kimik3_patches(model: Any | None = None) -> Iterator[None]: + if model is None: + yield + return + + language = language_model(model) + decoder = decoder_model(model) + kda_class = mla_class = moe_class = None + for layer in decoder.layers: + if layer.is_linear_attn: + kda_class = type(layer.self_attn) + else: + mla_class = type(layer.self_attn) + if hasattr(layer, "block_sparse_moe"): + moe_class = type(layer.block_sparse_moe) + + with ExitStack() as stack: + stack.enter_context( + patch_attribute(type(language), "forward", _patch_language_forward) + ) + if kda_class is not None: + stack.enter_context( + patch_attribute(kda_class, "forward", _patch_kda_attention) + ) + if mla_class is not None: + stack.enter_context( + patch_attribute(mla_class, "forward", _patch_mla_attention) + ) + if moe_class is not None: + stack.enter_context( + patch_attribute(moe_class, "forward", _patch_sparse_moe) + ) + yield diff --git a/tools/hf/exporters/models/kimi/spec.py b/tools/hf/exporters/models/kimi/spec.py new file mode 100644 index 00000000000..8fd310f0f26 --- /dev/null +++ b/tools/hf/exporters/models/kimi/spec.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from typing import Any + +import torch +import torch.nn as nn + +from ...ops import call_engine +from ...spec import ComponentBundle, EdgeSpec, register_edge_spec +from .helpers import allocate_kda_states, language_model +from .patches import apply_kimik3_patches + + +@register_edge_spec("kimi", "kimi-k3", "kimi_k3") +class KimiK3Spec(EdgeSpec): # type: ignore[misc] + def apply_patches(self, model=None): + return apply_kimik3_patches(model) + + def prepare_sample_inputs( + self, + model: nn.Module, + raw: Mapping[str, Any], + config: Any, + ) -> MutableMapping[str, Any]: + del config + if "inputs_embeds" in raw: + inputs_embeds = raw["inputs_embeds"] + elif "input_ids" in raw: + embedding = model.get_input_embeddings() + input_ids = raw["input_ids"].to(embedding.weight.device) + with torch.no_grad(): + inputs_embeds = embedding(input_ids) + else: + raise KeyError("Kimi K3 requires input_ids or inputs_embeds") + + attention_mask = raw.get("attention_mask") + if attention_mask is None: + attention_mask = torch.ones( + inputs_embeds.shape[:2], + device=inputs_embeds.device, + dtype=torch.long, + ) + else: + attention_mask = attention_mask.to(inputs_embeds.device) + return { + "inputs_embeds": inputs_embeds.contiguous(), + "attention_mask": attention_mask.contiguous(), + } + + def capture_eager_outputs( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + bench: dict[str, float] | None = None, + ) -> dict[str, torch.Tensor]: + del config + from ...measure import cuda_ms + + kwargs = { + "inputs_embeds": sample["inputs_embeds"], + "attention_mask": sample["attention_mask"], + "use_cache": False, + "return_dict": True, + } + with torch.no_grad(): + logits = language_model(model)(**kwargs).logits + + context_lengths = sample["attention_mask"].sum(dim=-1, dtype=torch.int64) + last_token_ids = context_lengths - 1 + batch = torch.arange(logits.shape[0], device=logits.device) + selected_logits = logits[batch, last_token_ids] + + if bench is not None: + bench["language"] = cuda_ms(lambda: language_model(model)(**kwargs).logits) + return {"language": selected_logits} + + def prepare( + self, + model: nn.Module, + sample: MutableMapping[str, Any], + config: Any, + ) -> dict[str, ComponentBundle]: + del config + embeddings = sample["inputs_embeds"] + batch_size = int(embeddings.shape[0]) + context_lengths = sample["attention_mask"].sum(dim=-1, dtype=torch.int32) + if bool((context_lengths <= 0).any()): + raise ValueError("Kimi attention_mask contains an empty sequence") + last_token_ids = (context_lengths - 1).to(torch.int64).unsqueeze(-1) + states, state_names = allocate_kda_states( + model, + batch_size=batch_size, + device=embeddings.device, + dtype=embeddings.dtype, + ) + args = (embeddings, context_lengths, last_token_ids, *states) + input_names = [ + "inputs_embeds", + "context_lengths", + "last_token_ids", + *state_names, + ] + sample.update(dict(zip(input_names, args))) + + return { + "language": ComponentBundle( + module=language_model(model).eval(), + trace_args=args, + save_args=args, + input_names=input_names, + output_names=["logits"], + parity_output="logits", + model_type="kimi_k3", + engine_file="language.engine", + trt_settings={ + "disable_tf32": True, + "use_fp32_acc": True, + "use_explicit_typing": True, + "decompose_attention": True, + "assume_dynamic_shape_support": True, + }, + ) + } + + def run(self, engines: Mapping[str, str], sample: Mapping[str, Any]) -> Any: + state_names = sorted( + ( + name + for name in sample + if name.startswith("kda_conv_") or name.startswith("kda_recurrent_") + ), + key=_state_sort_key, + ) + return call_engine( + engines["language"], + "language", + sample["inputs_embeds"], + sample["context_lengths"], + sample["last_token_ids"], + *(sample[name] for name in state_names), + ) + + +def _state_sort_key(name: str) -> tuple[int, int]: + layer_index = int(name.rsplit("_", 1)[1]) + kind = name[: name.rfind("_")] + order = { + "kda_conv_q": 0, + "kda_conv_k": 1, + "kda_conv_v": 2, + "kda_recurrent": 3, + } + return layer_index, order[kind] diff --git a/tools/hf/exporters/plugin/kimi_kda.py b/tools/hf/exporters/plugin/kimi_kda.py new file mode 100644 index 00000000000..6642d1b3d29 --- /dev/null +++ b/tools/hf/exporters/plugin/kimi_kda.py @@ -0,0 +1,49 @@ +"""Kimi KDA custom op lowered to the Edge-LLM ``kimi_kda`` plugin.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def register_kimi_kda_plugin_op() -> None: + """Register ``trt::kimi_kda_plugin`` for Dynamo export.""" + if hasattr(torch.ops, "trt") and hasattr(torch.ops.trt, "kimi_kda_plugin"): + return + + @torch.library.custom_op("trt::kimi_kda_plugin", mutates_args=()) + def kimi_kda_plugin( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + context_lengths: torch.Tensor, + lower_bound: float, + use_lower_bound: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + del q, k, gate, beta, a_log, dt_bias + del context_lengths, lower_bound, use_lower_bound + return torch.empty_like(v), torch.empty_like(state) + + @kimi_kda_plugin.register_fake + def _( + q, + k, + v, + gate, + beta, + a_log, + dt_bias, + state, + context_lengths, + lower_bound, + use_lower_bound, + ): + del q, k, gate, beta, a_log, dt_bias + del context_lengths, lower_bound, use_lower_bound + return torch.empty_like(v), torch.empty_like(state) diff --git a/tools/hf/exporters/plugin/plugin_converter.py b/tools/hf/exporters/plugin/plugin_converter.py index efb6db8ebdd..dbec1886833 100644 --- a/tools/hf/exporters/plugin/plugin_converter.py +++ b/tools/hf/exporters/plugin/plugin_converter.py @@ -296,6 +296,38 @@ def convert_update_ssm_state(ctx: ConversionContext, target, args, kwargs, name) return layer.get_output(0), layer.get_output(1) +@dynamo_tensorrt_converter( + torch.ops.trt.kimi_kda_plugin.default, + supports_dynamic_shapes=True, + priority=ConverterPriority.HIGH, +) +def convert_kimi_kda_plugin(ctx: ConversionContext, target, args, kwargs, name): + del target, kwargs + args = list(args) + tensors = args[:9] + lower_bound = float(args[9]) + use_lower_bound = int(bool(args[10])) + + creator = get_trt_plugin_creator("kimi_kda", "1", "") + if creator is None: + raise RuntimeError("kimi_kda not found in TensorRT plugin registry") + + plugin = _create_trt_plugin( + creator, + name, + [ + _float_field("lower_bound", lower_bound), + _int_field("use_lower_bound", use_lower_bound), + ], + ) + if plugin is None: + raise RuntimeError("Failed to create kimi_kda plugin") + + inputs = _as_plugin_inputs(ctx, tensors, name) + layer = _add_plugin_layer(ctx, inputs, plugin, name) + return layer.get_output(0), layer.get_output(1) + + def _convert_nvfp4_moe(ctx: ConversionContext, args, name: str, plugin_name: str): args = list(args) tensors = args[:11] diff --git a/tools/hf/exporters/plugin/plugin_utils.py b/tools/hf/exporters/plugin/plugin_utils.py index 31a3c58a82c..e983f178cf6 100644 --- a/tools/hf/exporters/plugin/plugin_utils.py +++ b/tools/hf/exporters/plugin/plugin_utils.py @@ -340,6 +340,7 @@ def load_plugin(): def load_plugins_for_trt(): + from .kimi_kda import register_kimi_kda_plugin_op from .mamba import register_mamba_plugin_ops from .moe import register_moe_plugin_ops @@ -347,6 +348,7 @@ def load_plugins_for_trt(): _register_vit_attention_plugin_op() register_mamba_plugin_ops() register_moe_plugin_ops() + register_kimi_kda_plugin_op() from . import attn_patches as _attn_patches # noqa: F401,E402 load_plugin() diff --git a/tools/hf/run_kimi_export.py b/tools/hf/run_kimi_export.py new file mode 100644 index 00000000000..3c756a46267 --- /dev/null +++ b/tools/hf/run_kimi_export.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""EdgeExporter on kimi-k3 model.""" + +from __future__ import annotations + +import torch +import torch_tensorrt +from exporters import EdgeConfig, EdgeExporter +from exporters.measure import print_bench +from exporters.plugin.plugin_utils import load_plugins_for_trt +from transformers import AutoModelForCausalLM, AutoTokenizer + +checkpoint = "inference-optimization/Kimi-K3-0.40B" + + +def load_policy(device: torch.device, dtype: torch.dtype): + model = ( + AutoModelForCausalLM.from_pretrained( + checkpoint, + trust_remote_code=True, + dtype=dtype, + ) + .to(device=device, dtype=dtype) + .eval() + ) + tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) + # if pad token is not set, set it to eos token + if tokenizer.pad_token_id is None: + tokenizer.pad_token = tokenizer.eos_token + return model, tokenizer + + +def main() -> None: + load_plugins_for_trt() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.float16 + + model, tokenizer = load_policy(device, dtype) + encoded = tokenizer("Hello, how are you?", return_tensors="pt") + sample_inputs = { + "input_ids": encoded["input_ids"].to(device), + "attention_mask": encoded["attention_mask"].to(device), + } + + exporter = EdgeExporter() + config = EdgeConfig( + model_type="kimi", + engine_dir="/tmp/kimi", + max_seq_len=1024, + ) + + program = exporter.export(model, sample_inputs, config=config) + + print("engines:", exporter.engines) + + with torch.no_grad(): + out = program.module()(**exporter.sample) + + logits = out[0] if isinstance(out, (tuple, list)) else out + print("logits", tuple(logits.shape), "mean", float(logits.float().mean())) + print_bench(exporter.bench) + + +if __name__ == "__main__": + main()