diff --git a/families/mamba/model.py b/families/mamba/model.py index 439a471dd6..1a1d2b7ea9 100644 --- a/families/mamba/model.py +++ b/families/mamba/model.py @@ -53,6 +53,7 @@ _transpose_2d, ) from . import graph_ops +from . import scan_rows from .parallel import normalize_parallel_config @@ -202,9 +203,15 @@ def build_engine( verbose: bool = False, debug_layer_outputs: bool = False, parallel_config=None, + prefill_rows: int | None = None, ) -> bytes: """Build TRT engine for Mamba SSM. + ``prefill_rows`` builds the prefill engine of a split bundle instead: + ``token_id`` is ``[S]`` with one profile ``1..prefill_rows``, every layer + runs its selective scan over the ``S`` rows, the logits come from the + last row and the present states are those after the last row. + max_cache_length is accepted for the shared build signature but is not used by Mamba (SSM state is constant size regardless of sequence length). @@ -219,6 +226,10 @@ def build_engine( present_ssm_0..N: float32 [1, d_inner * state_size] """ parallel = normalize_parallel_config(parallel_config) + if prefill_rows is not None and (parallel.enabled or debug_layer_outputs): + raise NotImplementedError("Mamba prefill engines are single-GPU without debug outputs") + if prefill_rows is not None and not 1 <= prefill_rows <= scan_rows.MAX_ROWS: + raise ValueError(f"Mamba prefill_rows must be in 1..{scan_rows.MAX_ROWS}") if parallel.enabled: parallel.validate() from .tp_builder import build_mamba_tp_engine @@ -257,7 +268,14 @@ def build_engine( # ----------------------------------------------------------- # Inputs # ----------------------------------------------------------- - token_id = network.add_input("token_id", trt.int32, (1,)) + rows = 1 if prefill_rows is None else None + token_id = network.add_input("token_id", trt.int32, (1 if rows else -1,)) + if prefill_rows is not None: + profile = builder.create_optimization_profile() + profile.set_shape( + "token_id", (1,), (min(PREFILL_OPT_TOKENS, prefill_rows),), (prefill_rows,) + ) + trt_config.add_optimization_profile(profile) conv_state_inputs = [] ssm_state_inputs = [] @@ -323,6 +341,7 @@ def build_engine( conv_kernel=conv_kernel, dt_rank=dt_rank, dtype=work_np_dtype, + sequence_length=rows, ) hidden_state = result["hidden"] @@ -332,6 +351,10 @@ def build_engine( if debug_layer_outputs: _mark_debug_output(network, hidden_state, f"debug_hidden_{layer_idx}") + if rows is None: + # Only the last row can affect the next token. + hidden_state = scan_rows.last_row(network, hidden_state, hidden) + # ----------------------------------------------------------- # Final norm # ----------------------------------------------------------- @@ -418,8 +441,9 @@ def _add_mamba_layer( conv_kernel: int, dt_rank: int, dtype: np.dtype = np.float32, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: - """Add one Mamba SSM layer. + """Add one Mamba SSM layer over one row, or ``S`` rows when ``sequence_length`` is None. Returns: {hidden, present_conv, present_ssm} """ @@ -437,6 +461,38 @@ def _add_mamba_layer( z = graph_ops.add_matmul_rhs_constant( network, normed, hidden_size, d_inner, weights[f"{prefix}.w_in_z"], dtype=dtype ) + if sequence_length is None: + y, present_conv, present_ssm = _add_mamba_mixer_rows( + network, + x, + weights=weights, + prefix=prefix, + conv_state_in=conv_state_in, + ssm_state_in=ssm_state_in, + d_inner=d_inner, + state_size=state_size, + conv_kernel=conv_kernel, + dt_rank=dt_rank, + dtype=dtype, + ) + if y.dtype != hidden.dtype: + y = network.add_cast(y, hidden.dtype).get_output(0) + z_activated = graph_ops.add_activation(network, z, "silu", dtype=dtype) + gated = network.add_elementwise(y, z_activated, trt.ElementWiseOperation.PROD) + out = graph_ops.add_matmul_rhs_constant( + network, + gated.get_output(0), + d_inner, + hidden_size, + weights[f"{prefix}.w_out"], + dtype=dtype, + ) + residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM) + return { + "hidden": residual.get_output(0), + "present_conv": present_conv, + "present_ssm": present_ssm, + } # ===== 3. Conv1d step ===== # conv_state_in: [d_inner, conv_kernel] @@ -632,6 +688,113 @@ def _add_mamba_layer( } +def _add_mamba_mixer_rows( + network: trt.INetworkDefinition, + x: trt.ITensor, + *, + weights: WeightDict, + prefix: str, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + d_inner: int, + state_size: int, + conv_kernel: int, + dt_rank: int, + dtype: np.dtype, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]: + """Conv and selective scan over ``x [S, d_inner]``: ``(y [S, d_inner] FP32, conv, ssm)``. + + Same arithmetic as the one-row layer applied token by token; the scan + combines the rows as described in :mod:`.scan_rows`. + """ + prod_op = trt.ElementWiseOperation.PROD + + def fp32(tensor): + return ( + tensor + if tensor.dtype == trt.float32 + else network.add_cast(tensor, trt.float32).get_output(0) + ) + + conv_w = graph_ops.add_constant( + network, (d_inner, conv_kernel), weights[f"{prefix}.conv1d_weight"], dtype=dtype + ) + conv_columns, present_conv = scan_rows.add_short_conv_rows( + network, + columns=scan_rows.shuffle(network, x, transpose=(1, 0)), + conv_state=conv_state_in, + weight=conv_w, + channels=d_inner, + kernel=conv_kernel, + ) + conv_out = graph_ops.add_bias_sum( + network, + scan_rows.shuffle(network, conv_columns, transpose=(1, 0)), + d_inner, + weights[f"{prefix}.conv1d_bias"], + dtype=dtype, + ) + activated = graph_ops.add_activation(network, conv_out, "silu", dtype=dtype) # [S, d] + + dt_in = graph_ops.add_matmul_rhs_constant( + network, activated, d_inner, dt_rank, weights[f"{prefix}.w_dt_in"], dtype=dtype + ) + b = graph_ops.add_matmul_rhs_constant( + network, activated, d_inner, state_size, weights[f"{prefix}.w_B"], dtype=dtype + ) + c = graph_ops.add_matmul_rhs_constant( + network, activated, d_inner, state_size, weights[f"{prefix}.w_C"], dtype=dtype + ) + dt = graph_ops.add_matmul_rhs_constant( + network, dt_in, dt_rank, d_inner, weights[f"{prefix}.w_dt_out"], dtype=dtype + ) + dt = fp32( + graph_ops.add_bias_sum(network, dt, d_inner, weights[f"{prefix}.dt_proj_bias"], dtype=dtype) + ) + one = graph_ops.add_constant(network, (1, 1), np.array([1.0], dtype=np.float32)) + dt = network.add_unary( + network.add_elementwise( + network.add_unary(dt, trt.UnaryOperation.EXP).get_output(0), + one, + trt.ElementWiseOperation.SUM, + ).get_output(0), + trt.UnaryOperation.LOG, + ).get_output(0) # softplus, [S, d] + + dt_3d = scan_rows.shuffle(network, dt, reshape=(-1, d_inner, 1)) + a_const = graph_ops.add_constant( + network, (1, d_inner, state_size), weights[f"{prefix}.A"].reshape(1, d_inner, state_size) + ) + decay = network.add_unary( + network.add_elementwise(dt_3d, a_const, prod_op).get_output(0), trt.UnaryOperation.EXP + ).get_output(0) # exp(dt A): [S, d, n] + x_3d = scan_rows.shuffle(network, fp32(activated), reshape=(-1, d_inner, 1)) + b_3d = scan_rows.shuffle(network, fp32(b), reshape=(-1, 1, state_size)) + update = network.add_elementwise( + network.add_elementwise(dt_3d, b_3d, prod_op).get_output(0), x_3d, prod_op + ).get_output(0) # dt B x: [S, d, n] + states, present_ssm = scan_rows.add_selective_scan_rows( + network, decay=decay, update=update, state=ssm_state_in, width=(d_inner, state_size) + ) + c_3d = scan_rows.shuffle(network, fp32(c), reshape=(-1, 1, state_size)) + y = network.add_reduce( + network.add_elementwise(states, c_3d, prod_op).get_output(0), + trt.ReduceOperation.SUM, + 1 << 2, + keep_dims=False, + ).get_output(0) # [S, d] + d_const = graph_ops.add_constant(network, (1, d_inner), weights[f"{prefix}.D"]) + y = network.add_elementwise( + y, + network.add_elementwise(d_const, fp32(activated), prod_op).get_output(0), + trt.ElementWiseOperation.SUM, + ).get_output(0) + return y, present_conv, present_ssm + + +# Optimum rows of the prefill engine. +PREFILL_OPT_TOKENS = 32 + _BUNDLE_FILES = ( "tokenizer.json", "tokenizer_config.json", @@ -757,7 +920,19 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: parallel_config=parallel, ) writer.add_bytes("engine.plan", plan) - layout = "single" + # A second engine prefills the prompt in chunks; decode keeps the one-row engine. + writer.add_bytes( + "prefill.plan", + model.build_engine( + config, + weights, + max_sequence_length, + precision=precision, + verbose=bool(request.verbose), + prefill_rows=scan_rows.MAX_ROWS, + ), + ) + layout = "split" writer.add_json( "runtime.json", _runtime_config( diff --git a/families/mamba/runtime/pipeline.cpp b/families/mamba/runtime/pipeline.cpp index 8d5871f30d..ba4f23ce26 100644 --- a/families/mamba/runtime/pipeline.cpp +++ b/families/mamba/runtime/pipeline.cpp @@ -29,11 +29,21 @@ RecurrentPipeline::RecurrentPipeline(std::unique_ptr decoder, std::unique_ptr state, RecurrentGenConfig config, cudaStream_t stream, const char* name, std::shared_ptr tokenizer, - std::string model_id_str) - : decoder_(std::move(decoder)), state_(std::move(state)), config_(config), stream_(stream), - name_(name), tokenizer_(std::move(tokenizer)), model_id_(std::move(model_id_str)) { + std::string model_id_str, std::unique_ptr prefill) + : decoder_(std::move(decoder)), state_(std::move(state)), prefill_(std::move(prefill)), + config_(config), stream_(stream), name_(name), tokenizer_(std::move(tokenizer)), + model_id_(std::move(model_id_str)) { if (!decoder_ || !decoder_->ok()) throw std::runtime_error(std::string(name_) + ": invalid decoder module"); + if (prefill_) { + if (!prefill_->ok()) + throw std::runtime_error(std::string(name_) + ": invalid prefill module"); + const auto rows = prefill_->input_profile_shape("token_id", prefill_->profile_idx(), + ProfileShapeSelector::kMax); + if (rows.size() != 1 || rows[0] < 1) + throw std::runtime_error(std::string(name_) + ": invalid prefill token profile"); + prefill_rows_ = static_cast(rows[0]); + } } static std::vector encode_prompt(const ITokenizer& tokenizer, @@ -103,10 +113,12 @@ std::vector RecurrentPipeline::generate_from_ids(const std::vector& input_ids, + std::vector& logits) { + // The prefill engine reads and writes the same recurrent state buffers. + state_->bind_to(*prefill_); + const auto total = static_cast(input_ids.size()); + for (int32_t start = 0; start < total; start += prefill_rows_) { + const int32_t rows = std::min(prefill_rows_, total - start); + TensorMap inputs; + inputs["token_id"] = + Tensor{const_cast(input_ids.data() + start), {rows}, DType::kInt32}; + state_->prepare_step(inputs, rows); + prefill_->forward_async(inputs); + state_->advance(rows); + } + prefill_->sync(); + const auto* device_logits = prefill_->device_ptr("logits"); + if (device_logits == nullptr) + throw std::runtime_error(std::string(name_) + ": prefill engine has no logits output"); + logits.resize(static_cast(config_.vocab_size)); + const cudaError_t copied = cudaMemcpy(logits.data(), device_logits, + logits.size() * sizeof(float), cudaMemcpyDeviceToHost); + if (copied != cudaSuccess) + throw std::runtime_error(std::string(name_) + + ": failed to copy prefill logits: " + cudaGetErrorString(copied)); +} + void RecurrentPipeline::run_step(int32_t token_id, std::vector& logits) { auto t0 = SteadyClock::now(); diff --git a/families/mamba/runtime/pipeline.h b/families/mamba/runtime/pipeline.h index 84736ba395..450894b30b 100644 --- a/families/mamba/runtime/pipeline.h +++ b/families/mamba/runtime/pipeline.h @@ -36,7 +36,7 @@ class RecurrentPipeline final : public ITextGeneration { std::unique_ptr state, RecurrentGenConfig config, cudaStream_t stream, const char* name, std::shared_ptr tokenizer = nullptr, - std::string model_id_str = ""); + std::string model_id_str = "", std::unique_ptr prefill = nullptr); TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg = {}) override; int32_t default_max_new_tokens() const override { return 128; } @@ -51,6 +51,9 @@ class RecurrentPipeline final : public ITextGeneration { private: std::unique_ptr decoder_; std::unique_ptr state_; + // Multi-row prefill engine of a split bundle; decode stays on decoder_. + std::unique_ptr prefill_; + int32_t prefill_rows_{0}; RecurrentGenConfig config_; cudaStream_t stream_; const char* name_; @@ -62,6 +65,9 @@ class RecurrentPipeline final : public ITextGeneration { const MambaSamplingParams& params); void run_step(int32_t token_id, std::vector& logits); + // Feed the prompt to prefill_ in chunks of up to prefill_rows_ tokens and + // return the logits of the last token. + void run_prefill_chunks(const std::vector& input_ids, std::vector& logits); using SteadyClock = std::chrono::steady_clock; void report_timing(SteadyClock::time_point t_prefill_start, diff --git a/families/mamba/runtime/plugin.cpp b/families/mamba/runtime/plugin.cpp index 1fafc43e39..654026d6ea 100644 --- a/families/mamba/runtime/plugin.cpp +++ b/families/mamba/runtime/plugin.cpp @@ -77,7 +77,8 @@ ITask* create(const FamilyContext& context) { if (hidden_size <= 0 || num_layers <= 0 || vocab_size <= 0 || inner <= 0 || state_size <= 0 || conv_kernel <= 1 || tp_size <= 0 || max_cache_length <= 0 || num_attention_heads <= 0 || num_key_value_heads <= 0 || head_dim <= 0 || expand <= 0 || - (layout != "single" && layout != "dual_profile") || + (layout != "single" && layout != "dual_profile" && layout != "split") || + (layout == "split" && tp_size != 1) || tp_mode != (tp_size > 1 ? "tensor_parallel" : "single") || (precision != "fp16" && precision != "bf16" && precision != "fp32")) { throw std::runtime_error("mamba runtime.json contains invalid geometry"); @@ -97,6 +98,18 @@ ITask* create(const FamilyContext& context) { throw std::runtime_error("mamba failed to load decoder"); decoder->set_timing_label("mamba decoder"); const cudaStream_t stream = decoder->stream(); + std::unique_ptr prefill; + if (layout == "split") { + // The prefill engine shares the decoder's stream so state copies stay ordered. + ModuleCreateOptions prefill_options = options; + prefill_options.stream = stream; + const auto& prefill_plan = require_section(context.reader, "prefill.plan"); + prefill = context.backend.create_module(prefill_plan.data(), prefill_plan.size(), + prefill_options); + if (prefill == nullptr || !prefill->ok()) + throw std::runtime_error("mamba failed to load the prefill engine"); + prefill->set_timing_label("mamba prefill"); + } std::vector specs = { {"conv_state", {inner, conv_kernel}, "present_conv"}, @@ -114,7 +127,8 @@ ITask* create(const FamilyContext& context) { generation.chat_template_format = mamba_detect_chat_template_format(chat_template(context.reader)); return new RecurrentPipeline(std::move(decoder), std::move(state), std::move(generation), - stream, "Mamba", create_tokenizer(context.reader), std::string{}); + stream, "Mamba", create_tokenizer(context.reader), std::string{}, + std::move(prefill)); } } // namespace trtmc::mamba diff --git a/families/mamba/scan_rows.py b/families/mamba/scan_rows.py new file mode 100644 index 0000000000..2dabee3e66 --- /dev/null +++ b/families/mamba/scan_rows.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mamba selective scan over ``S`` rows in one launch. + +The single-token decoder applies, per channel ``d`` and state ``n``:: + + h_t = exp(dt_t[d] A[d, n]) h_{t-1} + dt_t[d] B_t[n] x_t[d] + y_t[d] = sum_n C_t[n] h_t[d, n] + D[d] x_t[d] + +With ``h_t = a_t h_{t-1} + b_t`` the rows are combined by a Hillis-Steele scan: +``ceil(log2(S))`` rounds of ``(a, b)_t <- (a_t a_{t-k}, a_t b_{t-k} + b_t)`` +for ``k = 1, 2, 4, ...``. Afterwards ``a_t`` is the decay product since the +chunk start and ``b_t`` the zero-state recurrence, so ``h_t = a_t h_0 + b_t``. +Every factor is a decay in ``(0, 1]``, so the scan stays finite in FP32. +""" + +from __future__ import annotations + +import numpy as np +import tensorrt as trt + +# Prefill rows of one launch; the scan runs ceil(log2(MAX_ROWS)) rounds. +MAX_ROWS = 64 + + +def _const(network, values: np.ndarray) -> trt.ITensor: + values = np.ascontiguousarray(values) + return network.add_constant(values.shape, trt.Weights(values)).get_output(0) + + +def _i64(network, values) -> trt.ITensor: + return _const(network, np.asarray(values, dtype=np.int64)) + + +def _ew(network, a, b, op) -> trt.ITensor: + return network.add_elementwise(a, b, op).get_output(0) + + +def _mm(network, a, b, transpose_a=False, transpose_b=False) -> trt.ITensor: + op_a = trt.MatrixOperation.TRANSPOSE if transpose_a else trt.MatrixOperation.NONE + op_b = trt.MatrixOperation.TRANSPOSE if transpose_b else trt.MatrixOperation.NONE + return network.add_matrix_multiply(a, op_a, b, op_b).get_output(0) + + +def shuffle(network, tensor, reshape=None, transpose=None) -> trt.ITensor: + layer = network.add_shuffle(tensor) + if transpose is not None: + layer.first_transpose = transpose + if reshape is not None: + layer.reshape_dims = reshape + return layer.get_output(0) + + +def row_count(network, tensor: trt.ITensor) -> trt.ITensor: + """Length of axis 0 as an int64 ``[1]`` shape tensor.""" + shape = network.add_shape(tensor).get_output(0) + return network.add_slice(shape, (0,), (1,), (1,)).get_output(0) + + +def _dims(network, parts) -> trt.ITensor: + tensors = [p if isinstance(p, trt.ITensor) else _i64(network, [p]) for p in parts] + joined = network.add_concatenation(tensors) + joined.axis = 0 + return joined.get_output(0) + + +def dynamic_slice(network, tensor, start, size) -> trt.ITensor: + """Slice with static or shape-tensor start and size entries.""" + rank = len(start) + layer = network.add_slice(tensor, (0,) * rank, (1,) * rank, (1,) * rank) + if any(isinstance(s, trt.ITensor) for s in start): + layer.set_input(1, _dims(network, start)) + else: + layer.start = tuple(start) + layer.set_input(2, _dims(network, size)) + return layer.get_output(0) + + +def last_row(network, tensor: trt.ITensor, width: int) -> trt.ITensor: + """``[S, width]`` -> ``[1, width]``, the last row.""" + last = _ew( + network, row_count(network, tensor), _i64(network, [1]), trt.ElementWiseOperation.SUB + ) + return dynamic_slice(network, tensor, (last, 0), (1, width)) + + +def add_short_conv_rows( + network: trt.INetworkDefinition, + *, + columns: trt.ITensor, + conv_state: trt.ITensor, + weight: trt.ITensor, + channels: int, + kernel: int, +) -> tuple[trt.ITensor, trt.ITensor]: + """Causal depthwise conv over ``columns`` ``[C, S]`` with a ``[C, kernel]`` history. + + The history holds the last ``kernel`` inputs (newest last). Output column + ``t`` is ``sum_j weight[:, j] * window[:, t + j]`` over + ``window = [history[:, 1:], columns]``; the new history is the last + ``kernel`` columns of the window. Returns ``(outputs [C, S], history)``. + """ + rows = row_count(network, shuffle(network, columns, transpose=(1, 0))) + if kernel > 1: + past = network.add_slice(conv_state, (0, 1), (channels, kernel - 1), (1, 1)).get_output(0) + window = network.add_concatenation([past, columns]) + window.axis = 1 + window = window.get_output(0) + else: + window = columns + total = None + for tap in range(kernel): + tap_weight = network.add_slice(weight, (0, tap), (channels, 1), (1, 1)).get_output(0) + shifted = dynamic_slice(network, window, (0, tap), (channels, rows)) + term = _ew(network, shifted, tap_weight, trt.ElementWiseOperation.PROD) + total = term if total is None else _ew(network, total, term, trt.ElementWiseOperation.SUM) + last = _ew(network, rows, _i64(network, [1]), trt.ElementWiseOperation.SUB) + history = dynamic_slice(network, window, (0, last), (channels, kernel)) + return total, history + + +def _shift_rows(network, tensor, rows, shift: int, fill: float, width: tuple[int, int]): + """``out[t] = tensor[t - shift]`` for ``t >= shift``, else ``fill`` (``[S, d, n]``).""" + pad = _ew( + network, + _const(network, np.full((shift, 1, 1), fill, np.float32)), + _const(network, np.zeros((1, *width), np.float32)), + trt.ElementWiseOperation.SUM, + ) + joined = network.add_concatenation([pad, tensor]) + joined.axis = 0 + return dynamic_slice(network, joined.get_output(0), (0, 0, 0), (rows, *width)) + + +def add_selective_scan_rows( + network: trt.INetworkDefinition, + *, + decay: trt.ITensor, + update: trt.ITensor, + state: trt.ITensor, + width: tuple[int, int], +) -> tuple[trt.ITensor, trt.ITensor]: + """``h_t = decay_t h_{t-1} + update_t`` over ``[S, d, n]`` FP32 rows from ``state [d, n]``. + + Returns ``(h [S, d, n], h_last [d, n])``. + """ + prod_op, sum_op = trt.ElementWiseOperation.PROD, trt.ElementWiseOperation.SUM + rows = row_count(network, decay) + shift = 1 + while shift < MAX_ROWS: + prior_decay = _shift_rows(network, decay, rows, shift, 1.0, width) + prior_update = _shift_rows(network, update, rows, shift, 0.0, width) + update = _ew(network, _ew(network, decay, prior_update, prod_op), update, sum_op) + decay = _ew(network, decay, prior_decay, prod_op) + shift *= 2 + start = shuffle(network, state, reshape=(1, *width)) + states = _ew(network, _ew(network, decay, start, prod_op), update, sum_op) + last = _ew(network, rows, _i64(network, [1]), trt.ElementWiseOperation.SUB) + final = dynamic_slice(network, states, (last, 0, 0), (1, *width)) + return states, shuffle(network, final, reshape=width) diff --git a/families/mamba/tests/test_chunked_prefill.py b/families/mamba/tests/test_chunked_prefill.py new file mode 100644 index 0000000000..7225de24fc --- /dev/null +++ b/families/mamba/tests/test_chunked_prefill.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mamba prefill engine: chunked selective scan against the one-token engine.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") +safetensors_numpy = pytest.importorskip("safetensors.numpy") + +from .. import model, scan_rows # noqa: E402 +from ..config import ModelConfig # noqa: E402 + +VOCAB, HIDDEN, LAYERS, STATE, KERNEL, DT_RANK = 48, 64, 2, 16, 4, 8 +INNER = 2 * HIDDEN + + +def _write_checkpoint(path: Path) -> ModelConfig: + rng = np.random.default_rng(0) + + def w(*shape, scale=0.1): + return (rng.standard_normal(shape) * scale).astype(np.float32) + + tensors = { + "backbone.embeddings.weight": w(VOCAB, HIDDEN, scale=1.0), + "backbone.norm_f.weight": np.ones(HIDDEN, np.float32), + "lm_head.weight": w(VOCAB, HIDDEN), + } + for layer in range(LAYERS): + base = f"backbone.layers.{layer}" + tensors.update( + { + f"{base}.norm.weight": np.ones(HIDDEN, np.float32), + f"{base}.mixer.in_proj.weight": w(2 * INNER, HIDDEN), + f"{base}.mixer.conv1d.weight": w(INNER, 1, KERNEL, scale=0.5), + f"{base}.mixer.conv1d.bias": w(INNER), + f"{base}.mixer.x_proj.weight": w(DT_RANK + 2 * STATE, INNER), + f"{base}.mixer.dt_proj.weight": w(INNER, DT_RANK), + f"{base}.mixer.dt_proj.bias": w(INNER, scale=1.0), + # Decays from slow (A ~ -0.1) to fast (A ~ -16), as in real checkpoints. + f"{base}.mixer.A_log": np.log( + np.tile(np.linspace(0.1, 16.0, STATE), (INNER, 1)) + ).astype(np.float32), + f"{base}.mixer.D": np.ones(INNER, np.float32), + f"{base}.mixer.out_proj.weight": w(HIDDEN, INNER), + } + ) + path.mkdir(parents=True, exist_ok=True) + safetensors_numpy.save_file(tensors, str(path / "model.safetensors")) + (path / "config.json").write_text( + json.dumps( + { + "model_type": "mamba", + "architectures": ["MambaForCausalLM"], + "vocab_size": VOCAB, + "hidden_size": HIDDEN, + "intermediate_size": INNER, + "num_hidden_layers": LAYERS, + "state_size": STATE, + "conv_kernel": KERNEL, + "time_step_rank": DT_RANK, + "expand": 2, + "layer_norm_epsilon": 1e-5, + } + ) + ) + return ModelConfig.from_dir(path) + + +def _run(engine, state: dict, tokens: list[int]) -> np.ndarray: + context = engine.create_execution_context() + stream = torch.cuda.current_stream().cuda_stream + token = torch.tensor(tokens, dtype=torch.int32).cuda() + if -1 in tuple(engine.get_tensor_shape("token_id")): + assert context.set_input_shape("token_id", (len(tokens),)) + context.set_tensor_address("token_id", token.data_ptr()) + outputs = {} + for i in range(engine.num_io_tensors): + name = engine.get_tensor_name(i) + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + if name in state: + context.set_tensor_address(name, state[name].data_ptr()) + continue + outputs[name] = torch.empty( + tuple(context.get_tensor_shape(name)), dtype=torch.float32, device="cuda" + ) + context.set_tensor_address(name, outputs[name].data_ptr()) + assert context.execute_async_v3(stream) + torch.cuda.synchronize() + for name, value in outputs.items(): + if name.startswith("present_"): + state[ + name.replace("present_conv", "conv_state").replace("present_ssm", "ssm_state") + ].copy_(value) + return outputs["logits"][-1].cpu().numpy() + + +def _state(engine) -> dict: + return { + engine.get_tensor_name(i): torch.zeros( + tuple(engine.get_tensor_shape(engine.get_tensor_name(i))), + dtype=torch.float32, + device="cuda", + ) + for i in range(engine.num_io_tensors) + if engine.get_tensor_name(i).startswith(("conv_state", "ssm_state")) + } + + +@pytest.mark.gpu +@pytest.mark.trt +def test_prefill_engine_matches_one_token_engine(tmp_path: Path) -> None: + """Prefill chunks leave the logits and recurrent state of token-by-token steps.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + config = _write_checkpoint(tmp_path) + instance = model._MambaModel() + weights = instance.load_weights(str(tmp_path), config) + runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING)) + decoder = runtime.deserialize_cuda_engine(instance.build_engine(config, weights, 64)) + prefill = runtime.deserialize_cuda_engine( + instance.build_engine(config, weights, 64, prefill_rows=scan_rows.MAX_ROWS) + ) + assert [tuple(s) for s in prefill.get_tensor_profile_shape("token_id", 0)] == [ + (1,), + (min(model.PREFILL_OPT_TOKENS, scan_rows.MAX_ROWS),), + (scan_rows.MAX_ROWS,), + ] + + rng = np.random.default_rng(1) + prompt = rng.integers(0, VOCAB, 75).tolist() # two chunks: 64 + 11 rows + want_state = _state(decoder) + for token in prompt: + want = _run(decoder, want_state, [token]) + got_state = _state(decoder) + got = None + for start in range(0, len(prompt), scan_rows.MAX_ROWS): + got = _run(prefill, got_state, prompt[start : start + scan_rows.MAX_ROWS]) + + errors = {"logits": np.abs(got - want).max() / np.abs(want).max()} + for name, reference in want_state.items(): + reference = reference.cpu().numpy() + errors[name] = np.abs(got_state[name].cpu().numpy() - reference).max() / max( + np.abs(reference).max(), 1e-6 + ) + assert np.argmax(got) == np.argmax(want) + # FP32 engines run matmuls in TF32; one row and several rows take different kernels. + assert max(errors.values()) < 5e-3, errors + + +@pytest.mark.trt +def test_prefill_engine_rejects_unsupported_builds(tmp_path: Path) -> None: + """Prefill engines are single-GPU, bounded by the scan size and without debug outputs.""" + config = _write_checkpoint(tmp_path) + instance = model._MambaModel() + weights = instance.load_weights(str(tmp_path), config) + with pytest.raises(ValueError, match="prefill_rows"): + instance.build_engine(config, weights, 64, prefill_rows=scan_rows.MAX_ROWS + 1) + with pytest.raises(NotImplementedError, match="prefill"): + instance.build_engine(config, weights, 64, prefill_rows=8, debug_layer_outputs=True)