diff --git a/families/rwkv/model.py b/families/rwkv/model.py index 9c365ffebc..29a70dc8ab 100644 --- a/families/rwkv/model.py +++ b/families/rwkv/model.py @@ -66,6 +66,7 @@ _transpose_2d, ) from . import graph_ops +from . import wkv_rows from .parallel import normalize_parallel_config @@ -236,9 +237,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 RWKV linear attention. + ``prefill_rows`` builds the prefill engine of a split bundle instead: + ``token_id`` is ``[S]`` with one profile ``1..prefill_rows``, every layer + mixes the ``S`` rows in one launch (see :mod:`.wkv_rows`), the logits + come from the last row and the present states are those after it. + max_cache_length is accepted for the shared build signature but is not used by RWKV (recurrent state is constant size regardless of sequence length). @@ -259,6 +266,14 @@ def build_engine( present_max_0..N: float32 [1, hidden_size] """ parallel = normalize_parallel_config(parallel_config) + if prefill_rows is not None and ( + parallel.enabled or debug_layer_outputs or quant_ctx is not None + ): + raise NotImplementedError( + "RWKV prefill engines are single-GPU, unquantized and without debug outputs" + ) + if prefill_rows is not None and not 1 <= prefill_rows <= wkv_rows.MAX_ROWS: + raise ValueError(f"RWKV prefill_rows must be in 1..{wkv_rows.MAX_ROWS}") if parallel.enabled: parallel.validate() if quant_ctx is not None: @@ -298,7 +313,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) attn_state_inputs = [] ff_state_inputs = [] @@ -402,6 +424,7 @@ def build_engine( intermediate_size=intermediate, dtype=work_np_dtype, quant_ctx=quant_ctx, + sequence_length=rows, ) hidden_state = result["hidden"] @@ -414,6 +437,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 = wkv_rows.last_row(network, hidden_state, hidden) + # ----------------------------------------------------------- # Final LayerNorm # ----------------------------------------------------------- @@ -520,8 +547,10 @@ def _add_rwkv_layer( intermediate_size: int, dtype: np.dtype = np.float32, quant_ctx=None, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: - """Add one RWKV layer (time-mixing + channel-mixing). + """Add one RWKV layer (time-mixing + channel-mixing) over one row, or ``S`` rows + when ``sequence_length`` is None. ``quant_ctx`` (FP8 builds only) routes the projections it quantizes through Q/DQ; every other projection is the unchanged 16/32-bit matmul. @@ -555,6 +584,11 @@ def _project(inp, in_width, out_width, weight_key, module): # The normed output before time-shift becomes present_attn for next step present_attn = normed_attn + attn_prev = attn_state_in + if sequence_length is None: + # Row t blends with row t - 1; row 0 with the incoming state. + present_attn = wkv_rows.last_row(network, normed_attn, hidden_size) + attn_prev = wkv_rows.shifted_rows(network, normed_attn, attn_state_in, hidden_size) # ===== 2. Time-shift blending for R, K, V ===== # For each of R, K, V: blended = mix * normed + (1 - mix) * attn_state_in @@ -574,9 +608,9 @@ def _time_shift_blend(normed, prev_state, mix_weights_key): ) return blended.get_output(0) - xk = _time_shift_blend(normed_attn, attn_state_in, f"{prefix}.time_mix_key") - xv = _time_shift_blend(normed_attn, attn_state_in, f"{prefix}.time_mix_value") - xr = _time_shift_blend(normed_attn, attn_state_in, f"{prefix}.time_mix_receptance") + xk = _time_shift_blend(normed_attn, attn_prev, f"{prefix}.time_mix_key") + xv = _time_shift_blend(normed_attn, attn_prev, f"{prefix}.time_mix_value") + xr = _time_shift_blend(normed_attn, attn_prev, f"{prefix}.time_mix_receptance") # ===== 3. Projections ===== # R = sigmoid(xr @ w_attn_r) @@ -603,95 +637,118 @@ def _time_shift_blend(normed, prev_state, mix_weights_key): time_decay = graph_ops.add_constant(network, (1, hidden_size), weights[f"{prefix}.time_decay"]) time_first = graph_ops.add_constant(network, (1, hidden_size), weights[f"{prefix}.time_first"]) - # decay_plus_max = max_state_in + time_decay - decay_plus_max = network.add_elementwise(max_state_in, time_decay, trt.ElementWiseOperation.SUM) + def _wkv_step(): + """One row: ``(wkv, num, den, max)`` with the max-shifted exponentials.""" + # decay_plus_max = max_state_in + time_decay + decay_plus_max = network.add_elementwise(max_state_in, time_decay, trt.ElementWiseOperation.SUM) - # tf_plus_k = time_first + K - tf_plus_k = network.add_elementwise(time_first, k_recur, trt.ElementWiseOperation.SUM) + # tf_plus_k = time_first + K + tf_plus_k = network.add_elementwise(time_first, k_recur, trt.ElementWiseOperation.SUM) - # ---- WKV output (uses time_first as bonus, NO decay) ---- - # HF: max_for_output = max(max_state, key + time_first) - # time_decay is ONLY used in the state update, not the output. - q_out = network.add_elementwise( - tf_plus_k.get_output(0), max_state_in, trt.ElementWiseOperation.MAX - ) + # ---- WKV output (uses time_first as bonus, NO decay) ---- + # HF: max_for_output = max(max_state, key + time_first) + # time_decay is ONLY used in the state update, not the output. + q_out = network.add_elementwise( + tf_plus_k.get_output(0), max_state_in, trt.ElementWiseOperation.MAX + ) - # e2 = exp(key + time_first - q) - tf_k_minus_q = network.add_elementwise( - tf_plus_k.get_output(0), q_out.get_output(0), trt.ElementWiseOperation.SUB - ) - exp_tf_k = network.add_unary(tf_k_minus_q.get_output(0), trt.UnaryOperation.EXP) + # e2 = exp(key + time_first - q) + tf_k_minus_q = network.add_elementwise( + tf_plus_k.get_output(0), q_out.get_output(0), trt.ElementWiseOperation.SUB + ) + exp_tf_k = network.add_unary(tf_k_minus_q.get_output(0), trt.UnaryOperation.EXP) - # e1 = exp(max_state - q) (NO time_decay here) - ms_minus_q = network.add_elementwise( - max_state_in, q_out.get_output(0), trt.ElementWiseOperation.SUB - ) - exp_dpm = network.add_unary(ms_minus_q.get_output(0), trt.UnaryOperation.EXP) + # e1 = exp(max_state - q) (NO time_decay here) + ms_minus_q = network.add_elementwise( + max_state_in, q_out.get_output(0), trt.ElementWiseOperation.SUB + ) + exp_dpm = network.add_unary(ms_minus_q.get_output(0), trt.UnaryOperation.EXP) - # wkv_num = exp(tf+k-q) * V + exp(dpm-q) * num_state - term1_num = network.add_elementwise( - exp_tf_k.get_output(0), v_recur, trt.ElementWiseOperation.PROD - ) - term2_num = network.add_elementwise( - exp_dpm.get_output(0), num_state_in, trt.ElementWiseOperation.PROD - ) - wkv_num = network.add_elementwise( - term1_num.get_output(0), term2_num.get_output(0), trt.ElementWiseOperation.SUM - ) + # wkv_num = exp(tf+k-q) * V + exp(dpm-q) * num_state + term1_num = network.add_elementwise( + exp_tf_k.get_output(0), v_recur, trt.ElementWiseOperation.PROD + ) + term2_num = network.add_elementwise( + exp_dpm.get_output(0), num_state_in, trt.ElementWiseOperation.PROD + ) + wkv_num = network.add_elementwise( + term1_num.get_output(0), term2_num.get_output(0), trt.ElementWiseOperation.SUM + ) - # wkv_den = exp(tf+k-q) + exp(dpm-q) * den_state - term2_den = network.add_elementwise( - exp_dpm.get_output(0), den_state_in, trt.ElementWiseOperation.PROD - ) - wkv_den = network.add_elementwise( - exp_tf_k.get_output(0), term2_den.get_output(0), trt.ElementWiseOperation.SUM - ) + # wkv_den = exp(tf+k-q) + exp(dpm-q) * den_state + term2_den = network.add_elementwise( + exp_dpm.get_output(0), den_state_in, trt.ElementWiseOperation.PROD + ) + wkv_den = network.add_elementwise( + exp_tf_k.get_output(0), term2_den.get_output(0), trt.ElementWiseOperation.SUM + ) - # wkv = wkv_num / wkv_den - wkv = network.add_elementwise( - wkv_num.get_output(0), wkv_den.get_output(0), trt.ElementWiseOperation.DIV - ) + # wkv = wkv_num / wkv_den + wkv = network.add_elementwise( + wkv_num.get_output(0), wkv_den.get_output(0), trt.ElementWiseOperation.DIV + ) - # ---- State update (does NOT use time_first) ---- - # q2 = max(K, decay_plus_max) - q2 = network.add_elementwise( - k_recur, decay_plus_max.get_output(0), trt.ElementWiseOperation.MAX - ) + # ---- State update (does NOT use time_first) ---- + # q2 = max(K, decay_plus_max) + q2 = network.add_elementwise( + k_recur, decay_plus_max.get_output(0), trt.ElementWiseOperation.MAX + ) - # exp(K - q2) - k_minus_q2 = network.add_elementwise(k_recur, q2.get_output(0), trt.ElementWiseOperation.SUB) - exp_k_q2 = network.add_unary(k_minus_q2.get_output(0), trt.UnaryOperation.EXP) + # exp(K - q2) + k_minus_q2 = network.add_elementwise(k_recur, q2.get_output(0), trt.ElementWiseOperation.SUB) + exp_k_q2 = network.add_unary(k_minus_q2.get_output(0), trt.UnaryOperation.EXP) - # exp(decay_plus_max - q2) - dpm_minus_q2 = network.add_elementwise( - decay_plus_max.get_output(0), q2.get_output(0), trt.ElementWiseOperation.SUB - ) - exp_dpm_q2 = network.add_unary(dpm_minus_q2.get_output(0), trt.UnaryOperation.EXP) + # exp(decay_plus_max - q2) + dpm_minus_q2 = network.add_elementwise( + decay_plus_max.get_output(0), q2.get_output(0), trt.ElementWiseOperation.SUB + ) + exp_dpm_q2 = network.add_unary(dpm_minus_q2.get_output(0), trt.UnaryOperation.EXP) - # present_num = exp(K-q2)*V + exp(dpm-q2)*num_state - st_term1 = network.add_elementwise( - exp_k_q2.get_output(0), v_recur, trt.ElementWiseOperation.PROD - ) - st_term2 = network.add_elementwise( - exp_dpm_q2.get_output(0), num_state_in, trt.ElementWiseOperation.PROD - ) - present_num = network.add_elementwise( - st_term1.get_output(0), st_term2.get_output(0), trt.ElementWiseOperation.SUM - ) + # present_num = exp(K-q2)*V + exp(dpm-q2)*num_state + st_term1 = network.add_elementwise( + exp_k_q2.get_output(0), v_recur, trt.ElementWiseOperation.PROD + ) + st_term2 = network.add_elementwise( + exp_dpm_q2.get_output(0), num_state_in, trt.ElementWiseOperation.PROD + ) + present_num = network.add_elementwise( + st_term1.get_output(0), st_term2.get_output(0), trt.ElementWiseOperation.SUM + ) - # present_den = exp(K-q2) + exp(dpm-q2)*den_state - st_den_term2 = network.add_elementwise( - exp_dpm_q2.get_output(0), den_state_in, trt.ElementWiseOperation.PROD - ) - present_den = network.add_elementwise( - exp_k_q2.get_output(0), st_den_term2.get_output(0), trt.ElementWiseOperation.SUM - ) + # present_den = exp(K-q2) + exp(dpm-q2)*den_state + st_den_term2 = network.add_elementwise( + exp_dpm_q2.get_output(0), den_state_in, trt.ElementWiseOperation.PROD + ) + present_den = network.add_elementwise( + exp_k_q2.get_output(0), st_den_term2.get_output(0), trt.ElementWiseOperation.SUM + ) + + # present_max = q2 + present_max = q2.get_output(0) + return ( + wkv.get_output(0), + present_num.get_output(0), + present_den.get_output(0), + present_max, + ) - # present_max = q2 - present_max = q2.get_output(0) + if sequence_length is None: + wkv_out, present_num, present_den, present_max = wkv_rows.add_wkv_rows( + network, + k=k_recur, + v=v_recur, + time_decay=time_decay, + time_first=time_first, + num_state=num_state_in, + den_state=den_state_in, + max_state=max_state_in, + width=hidden_size, + ) + else: + wkv_out, present_num, present_den, present_max = _wkv_step() # ===== 5. Gated output + residual ===== - wkv_for_gate = wkv.get_output(0) + wkv_for_gate = wkv_out if wkv_for_gate.dtype != r_gate.get_output(0).dtype: wkv_for_gate = network.add_cast(wkv_for_gate, r_gate.get_output(0).dtype).get_output(0) gated = network.add_elementwise( @@ -724,10 +781,14 @@ def _time_shift_blend(normed, prev_state, mix_weights_key): # The normed output before time-shift becomes present_ff for next step present_ff = normed_ffn + ff_prev = ff_state_in + if sequence_length is None: + present_ff = wkv_rows.last_row(network, normed_ffn, hidden_size) + ff_prev = wkv_rows.shifted_rows(network, normed_ffn, ff_state_in, hidden_size) # ===== 2. Time-shift for key and receptance ===== - xk_ffn = _time_shift_blend(normed_ffn, ff_state_in, f"{prefix}.time_mix_ffn_key") - xr_ffn = _time_shift_blend(normed_ffn, ff_state_in, f"{prefix}.time_mix_ffn_receptance") + xk_ffn = _time_shift_blend(normed_ffn, ff_prev, f"{prefix}.time_mix_ffn_key") + xr_ffn = _time_shift_blend(normed_ffn, ff_prev, f"{prefix}.time_mix_ffn_receptance") # ===== 3. Key projection + squared ReLU ===== k_ffn = _project(xk_ffn, hidden_size, intermediate_size, "w_ffn_k", "feed_forward.key") @@ -752,12 +813,15 @@ def _time_shift_blend(normed, prev_state, mix_weights_key): "hidden": residual_ffn.get_output(0), "present_attn": present_attn, "present_ff": present_ff, - "present_num": present_num.get_output(0), - "present_den": present_den.get_output(0), + "present_num": present_num, + "present_den": present_den, "present_max": present_max, } +# Optimum rows of the prefill engine. +PREFILL_OPT_TOKENS = 32 + _BUNDLE_FILES = ( "tokenizer.json", "tokenizer_config.json", @@ -899,6 +963,20 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: ) writer.add_bytes("engine.plan", plan) layout = "single" + if quant_ctx is None: + # 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=wkv_rows.MAX_ROWS, + ), + ) + layout = "split" writer.add_json( "runtime.json", _runtime_config( diff --git a/families/rwkv/runtime/pipeline.cpp b/families/rwkv/runtime/pipeline.cpp index b66486f08d..96c371e875 100644 --- a/families/rwkv/runtime/pipeline.cpp +++ b/families/rwkv/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/rwkv/runtime/pipeline.h b/families/rwkv/runtime/pipeline.h index 7a081bb396..325a1c67f9 100644 --- a/families/rwkv/runtime/pipeline.h +++ b/families/rwkv/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 RwkvSamplingParams& 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/rwkv/runtime/plugin.cpp b/families/rwkv/runtime/plugin.cpp index b231eacf1a..ba608601af 100644 --- a/families/rwkv/runtime/plugin.cpp +++ b/families/rwkv/runtime/plugin.cpp @@ -71,7 +71,8 @@ ITask* create(const FamilyContext& context) { const std::string layout = require_value(json, "decoder_engine_layout"); if (hidden_size <= 0 || num_layers <= 0 || vocab_size <= 0 || tp_size <= 0 || max_cache_length <= 0 || num_attention_heads <= 0 || num_key_value_heads <= 0 || - head_dim <= 0 || (layout != "single" && layout != "dual_profile") || + head_dim <= 0 || (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("rwkv runtime.json contains invalid geometry"); @@ -91,6 +92,18 @@ ITask* create(const FamilyContext& context) { throw std::runtime_error("rwkv failed to load decoder"); decoder->set_timing_label("rwkv 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("rwkv failed to load the prefill engine"); + prefill->set_timing_label("rwkv prefill"); + } std::vector specs = { {"attn_state", {hidden_size}, "present_attn"}, {"ff_state", {hidden_size}, "present_ff"}, @@ -109,7 +122,8 @@ ITask* create(const FamilyContext& context) { generation.chat_template_format = rwkv_detect_chat_template_format(chat_template(context.reader)); return new RecurrentPipeline(std::move(decoder), std::move(state), std::move(generation), - stream, "RWKV", create_tokenizer(context.reader), std::string{}); + stream, "RWKV", create_tokenizer(context.reader), std::string{}, + std::move(prefill)); } } // namespace trtmc::rwkv diff --git a/families/rwkv/tests/test_chunked_prefill.py b/families/rwkv/tests/test_chunked_prefill.py new file mode 100644 index 0000000000..63e2a85993 --- /dev/null +++ b/families/rwkv/tests/test_chunked_prefill.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""RWKV prefill engine: multi-row chunks against token-by-token steps of the one-token engine.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +trt = pytest.importorskip("tensorrt") +if not hasattr(trt, "Builder"): # another test module put a stand-in module in sys.modules + pytest.skip("needs TensorRT", allow_module_level=True) + +from .. import model, wkv_rows # noqa: E402 +from ..config import ModelConfig # noqa: E402 +from .test_quantization import HIDDEN, LAYERS, VOCAB, _write_tiny_rwkv # noqa: E402 + +_STATES = ("attn", "ff", "num", "den", "max") + + +def _wide_mixing(seed: int) -> dict: + """Decays and bonuses over the range of trained checkpoints, and keys large enough to need the shift.""" + rng = np.random.default_rng(seed) + extra = {} + for layer in range(LAYERS): + prefix = f"rwkv.blocks.{layer}.attention" + extra[f"{prefix}.time_decay"] = rng.uniform(-5.0, 2.0, HIDDEN).astype(np.float32) + extra[f"{prefix}.time_first"] = rng.uniform(-2.0, 3.0, HIDDEN).astype(np.float32) + extra[f"{prefix}.key.weight"] = (rng.standard_normal((HIDDEN, HIDDEN)) * 0.5).astype( + np.float32 + ) + return extra + + +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).split("_state")[0] in _STATES + and "_state_" in engine.get_tensor_name(i) + } + + +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_"): + kind, layer = name[len("present_") :].rsplit("_", 1) + state[f"{kind}_state_{layer}"].copy_(value) + return outputs["logits"][-1].cpu().numpy() + + +@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") + _write_tiny_rwkv(tmp_path, seed=3, extra=_wide_mixing(4)) + config = ModelConfig.from_dir(tmp_path) + instance = model._RwkvModel() + 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, 128)) + prefill = runtime.deserialize_cuda_engine( + instance.build_engine(config, weights, 128, prefill_rows=wkv_rows.MAX_ROWS) + ) + assert [tuple(s) for s in prefill.get_tensor_profile_shape("token_id", 0)] == [ + (1,), + (model.PREFILL_OPT_TOKENS,), + (wkv_rows.MAX_ROWS,), + ] + + rng = np.random.default_rng(1) + prompt = rng.integers(0, VOCAB, 75).tolist() # two chunks: 64 + 11 rows + following = rng.integers(0, VOCAB, 3).tolist() + want_state = _state(decoder) + want = [_run(decoder, want_state, [token]) for token in prompt][-1:] + got_state = _state(decoder) + got = None + for start in range(0, len(prompt), wkv_rows.MAX_ROWS): + got = _run(prefill, got_state, prompt[start : start + wkv_rows.MAX_ROWS]) + got = [got] + # Decode on the one-token engine from both states. + for token in following: + want.append(_run(decoder, want_state, [token])) + got.append(_run(decoder, got_state, [token])) + + def relative(value, reference): + return float(np.abs(value - reference).max() / max(np.abs(reference).max(), 1e-6)) + + errors = {} + for step, (value, reference) in enumerate(zip(got, want)): + assert np.argmax(value) == np.argmax(reference), step + errors[f"logits[{step}]"] = relative(value, reference) + for layer in range(LAYERS): + for kind in ("attn", "ff"): + name = f"{kind}_state_{layer}" + errors[name] = relative(got_state[name].cpu().numpy(), want_state[name].cpu().numpy()) + # num and den may carry another scale (max); their ratio may not. + ratio = { + key: (state[f"num_state_{layer}"] / state[f"den_state_{layer}"]).cpu().numpy() + for key, state in (("got", got_state), ("want", want_state)) + } + errors[f"num/den_{layer}"] = relative(ratio["got"], ratio["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, unquantized, bounded by the distance table, without debug outputs.""" + _write_tiny_rwkv(tmp_path) + config = ModelConfig.from_dir(tmp_path) + instance = model._RwkvModel() + weights = instance.load_weights(str(tmp_path), config) + with pytest.raises(ValueError, match="prefill_rows"): + instance.build_engine(config, weights, 128, prefill_rows=wkv_rows.MAX_ROWS + 1) + with pytest.raises(NotImplementedError, match="prefill"): + instance.build_engine(config, weights, 128, prefill_rows=8, debug_layer_outputs=True) diff --git a/families/rwkv/wkv_rows.py b/families/rwkv/wkv_rows.py new file mode 100644 index 0000000000..d292d9f528 --- /dev/null +++ b/families/rwkv/wkv_rows.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""RWKV-4 time mixing over ``S`` rows in one launch. + +The single-token decoder keeps, per channel, ``num * exp(max)`` and +``den * exp(max)`` of + + A_t = exp(w) A_{t-1} + exp(k_t) v_t, B_t = exp(w) B_{t-1} + exp(k_t) + +and outputs ``wkv_t = (A_{t-1} + exp(u + k_t) v_t) / (B_{t-1} + exp(u + k_t))`` +(``w`` is the negative ``time_decay``, ``u`` the ``time_first`` bonus). The +decay does not depend on the token, so for a chunk with incoming state +``(num_0, den_0, max_0)`` row ``t`` sums three kinds of terms, with exponents + + max_0 + t w (the incoming state) + k_i + (t - 1 - i) w, i < t (earlier rows of the chunk) + u + k_t (the bonus of row t) + +Each row subtracts the largest exponent before ``exp``, as the decoder does. +Row ``S`` (no bonus) is the state after the chunk, kept as ``(num, den, max)`` +with ``max`` that largest exponent: the same ``A`` and ``B`` at a possibly +different scale than token-by-token steps would keep. +""" + +from __future__ import annotations + +import numpy as np +import tensorrt as trt + +# Rows of one launch; the distance table below is sized for it. +MAX_ROWS = 64 +# Exponent of terms that do not exist; exp() of it minus any real maximum is 0. +_ABSENT = -1.0e30 + + +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 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 _plus_one(network, rows: trt.ITensor) -> trt.ITensor: + return _ew(network, rows, _i64(network, [1]), trt.ElementWiseOperation.SUM) + + +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 shifted_rows(network, rows_in: trt.ITensor, previous: trt.ITensor, width: int) -> trt.ITensor: + """Row ``t`` of the result is row ``t - 1`` of ``rows_in`` (``[S, width]``); row 0 is ``previous``.""" + joined = network.add_concatenation([previous, rows_in]) + joined.axis = 0 + return dynamic_slice( + network, joined.get_output(0), (0, 0), (row_count(network, rows_in), width) + ) + + +def _append_row(network, rows_in: trt.ITensor, value: float, width: int) -> trt.ITensor: + """``[S, width]`` -> ``[S + 1, width]`` with a constant last row.""" + joined = network.add_concatenation( + [rows_in, _const(network, np.full((1, width), value, np.float32))] + ) + joined.axis = 0 + return joined.get_output(0) + + +def add_wkv_rows( + network: trt.INetworkDefinition, + *, + k: trt.ITensor, + v: trt.ITensor, + time_decay: trt.ITensor, + time_first: trt.ITensor, + num_state: trt.ITensor, + den_state: trt.ITensor, + max_state: trt.ITensor, + width: int, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor, trt.ITensor]: + """WKV of ``S`` rows (FP32 ``[S, width]`` ``k`` and ``v``, ``[1, width]`` decay, bonus and state). + + Returns ``(wkv [S, width], num, den, max)``, the new state ``[1, width]`` each. + """ + op = trt.ElementWiseOperation + rows = row_count(network, k) + rows1 = _plus_one(network, rows) + + # distance[t, i] = t - 1 - i: the decay steps from row i to the state read by row t. + t_index = np.arange(MAX_ROWS + 1, dtype=np.float32) + distance = dynamic_slice( + network, + _const(network, t_index[:, None] - 1.0 - np.arange(MAX_ROWS, dtype=np.float32)[None, :]), + (0, 0), + (rows1, rows), + ) + steps = dynamic_slice(network, _const(network, t_index[:, None]), (0, 0), (rows1, 1)) + + # Earlier rows: [S + 1, S, width]; rows i >= t do not exist. + distance3 = shuffle(network, distance, reshape=(0, 0, 1)) + decay3 = shuffle(network, time_decay, reshape=(1, 1, width)) + earlier = _ew( + network, + shuffle(network, k, reshape=(1, -1, width)), + _ew(network, distance3, decay3, op.PROD), + op.SUM, + ) + exists = _ew( + network, distance3, _const(network, np.full((1, 1, 1), -0.5, np.float32)), op.GREATER + ) + earlier = network.add_select( + exists, earlier, _const(network, np.full((1, 1, 1), _ABSENT, np.float32)) + ).get_output(0) + + # Incoming state and bonus: [S + 1, width]; the state row S has no bonus. + incoming = _ew(network, max_state, _ew(network, steps, time_decay, op.PROD), op.SUM) + bonus = _append_row(network, _ew(network, time_first, k, op.SUM), _ABSENT, width) + bonus_value = _append_row(network, v, 0.0, width) + + top = network.add_reduce(earlier, trt.ReduceOperation.MAX, 1 << 1, keep_dims=False).get_output( + 0 + ) + top = _ew(network, top, _ew(network, incoming, bonus, op.MAX), op.MAX) + + weight = network.add_unary( + _ew(network, earlier, shuffle(network, top, reshape=(0, 1, width)), op.SUB), + trt.UnaryOperation.EXP, + ).get_output(0) + value3 = shuffle(network, v, reshape=(1, -1, width)) + num = network.add_reduce( + _ew(network, weight, value3, op.PROD), trt.ReduceOperation.SUM, 1 << 1, keep_dims=False + ).get_output(0) + den = network.add_reduce(weight, trt.ReduceOperation.SUM, 1 << 1, keep_dims=False).get_output(0) + + exp_bonus = network.add_unary( + _ew(network, bonus, top, op.SUB), trt.UnaryOperation.EXP + ).get_output(0) + exp_state = network.add_unary( + _ew(network, incoming, top, op.SUB), trt.UnaryOperation.EXP + ).get_output(0) + num = _ew(network, num, _ew(network, exp_bonus, bonus_value, op.PROD), op.SUM) + num = _ew(network, num, _ew(network, exp_state, num_state, op.PROD), op.SUM) + den = _ew(network, den, exp_bonus, op.SUM) + den = _ew(network, den, _ew(network, exp_state, den_state, op.PROD), op.SUM) + + wkv = _ew( + network, + dynamic_slice(network, num, (0, 0), (rows, width)), + dynamic_slice(network, den, (0, 0), (rows, width)), + op.DIV, + ) + new_state = [ + dynamic_slice(network, tensor, (rows, 0), (1, width)) for tensor in (num, den, top) + ] + return wkv, new_state[0], new_state[1], new_state[2]