From 9fd875dbfea6d586d332c63d829b29b74be64b55 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 17:53:13 +0000 Subject: [PATCH] perf(rwkv): prefill text prompts in multi-token chunks Single-GPU unquantized bundles get a second engine (prefill.plan) that runs 1..64 prompt rows per launch. RWKV-4 decays do not depend on the token, so each row's WKV sums the incoming state, the earlier rows of the chunk and its own bonus under one max shift, as the one-row step does; the time-shift inputs are the previous rows. The runtime prefills in chunks on that engine and decodes on the unchanged one-token engine (decoder_engine_layout "split"). Signed-off-by: chaofengw --- families/rwkv/model.py | 244 +++++++++++++------- families/rwkv/runtime/pipeline.cpp | 52 ++++- families/rwkv/runtime/pipeline.h | 8 +- families/rwkv/runtime/plugin.cpp | 18 +- families/rwkv/tests/test_chunked_prefill.py | 146 ++++++++++++ families/rwkv/wkv_rows.py | 200 ++++++++++++++++ 6 files changed, 575 insertions(+), 93 deletions(-) create mode 100644 families/rwkv/tests/test_chunked_prefill.py create mode 100644 families/rwkv/wkv_rows.py 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]