From ea2e659db7a01661689cac72d210493b25733002 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 16:21:42 +0000 Subject: [PATCH 1/3] perf(qwen3_5): prefill text prompts in multi-token chunks Text-only single-GPU bundles get a second engine (prefill.plan) that runs 1..64 prompt rows per launch: Gated DeltaNet layers run each chunk as one chunked delta rule (the unit lower-triangular solve is a batched block inversion) and full-attention layers attend over the cache plus the chunk. The runtime prefills in chunks on that engine and decodes on the unchanged one-token engine (decoder_engine_layout "split"). NVFP4 bundles keep the token loop. Signed-off-by: chaofengw --- families/qwen3_5/deltanet_rows.py | 308 +++++++++++++ families/qwen3_5/model.py | 431 ++++++++++++++---- families/qwen3_5/runtime/hybrid_state.cpp | 12 + families/qwen3_5/runtime/hybrid_state.h | 2 + families/qwen3_5/runtime/inference_state.h | 15 + families/qwen3_5/runtime/kv_cache.cpp | 13 + families/qwen3_5/runtime/kv_cache.h | 3 + families/qwen3_5/runtime/pipeline.cpp | 53 ++- families/qwen3_5/runtime/pipeline.h | 9 +- families/qwen3_5/runtime/plugin.cpp | 24 +- .../qwen3_5/tests/test_chunked_prefill.py | 205 +++++++++ 11 files changed, 975 insertions(+), 100 deletions(-) create mode 100644 families/qwen3_5/deltanet_rows.py create mode 100644 families/qwen3_5/tests/test_chunked_prefill.py diff --git a/families/qwen3_5/deltanet_rows.py b/families/qwen3_5/deltanet_rows.py new file mode 100644 index 0000000000..735b94493b --- /dev/null +++ b/families/qwen3_5/deltanet_rows.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Gated DeltaNet mixer over ``S`` rows in one launch (chunked delta rule). + +The single-token decoder applies, per head and token ``t``:: + + S_t = exp(g_t) S_{t-1} + k_t u_t^T, u_t = beta_t (v_t - exp(g_t) S_{t-1}^T k_t) + o_t = S_t^T q_t + +For a chunk of ``S`` tokens with cumulative decay ``gamma_t = sum_{i<=t} g_i`` +and ``Gamma_t = exp(gamma_t)`` this is equivalent to:: + + (I + A) U = beta V - diag(beta Gamma) K S_0, + A_ij = beta_i exp(gamma_i - gamma_j) k_i.k_j for j < i, else 0 + O = diag(Gamma) Q S_0 + (exp(gamma_t - gamma_i) q_t.k_i)_{i<=t} U + S_S = exp(gamma_S) S_0 + (K * exp(gamma_S - gamma))^T U + +``I + A`` is unit lower triangular. It is inverted by 2x2 block inversion on +a fixed ``SOLVE_ROWS`` padding (the padded rows are identity), batched over +the blocks of each level. This stays accurate where a single Neumann product +over 32 or more rows does not (nearly identical keys with beta close to one). +""" + +from __future__ import annotations + +import math + +import numpy as np +import tensorrt as trt + +# Static size of the padded triangular solve; it bounds the prefill rows. +SOLVE_ROWS = 64 +# Diagonal blocks of this size are inverted with a Neumann product. +_SOLVE_LEAF = 8 + + +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 _block_rows(network, matrix, heads: int, blocks: int, size: int, pairs: np.ndarray): + """Blocks ``(r, c)`` of ``[heads, blocks*size, blocks*size]`` with ``pairs[r, c] = 1``, summed over c. + + Returns ``[heads, blocks, size, size]``: row block ``r`` holds the selected block of that row. + """ + split = _shuffle(network, matrix, reshape=(heads, blocks, size, blocks, size)) + mask = _const(network, pairs.astype(np.float32).reshape(1, blocks, 1, blocks, 1)) + selected = _ew(network, split, mask, trt.ElementWiseOperation.PROD) + return network.add_reduce( + selected, trt.ReduceOperation.SUM, 1 << 3, keep_dims=False + ).get_output(0) + + +def _pair_half(network, tensor, heads: int, pairs: int, size: int, index: int) -> trt.ITensor: + """``[heads, 2*pairs, size, size]`` -> block ``index`` (0 or 1) of each pair.""" + grouped = _shuffle(network, tensor, reshape=(heads, pairs, 2, size, size)) + half = network.add_slice( + grouped, (0, 0, index, 0, 0), (heads, pairs, 1, size, size), (1, 1, 1, 1, 1) + ).get_output(0) + return _shuffle(network, half, reshape=(heads, pairs, size, size)) + + +def _inverse_unit_lower(network, lower: trt.ITensor, heads: int, size: int) -> trt.ITensor: + """Inverse of a static ``[heads, size, size]`` unit lower-triangular matrix. + + All ``_SOLVE_LEAF`` diagonal blocks are inverted at once with a Neumann + product, then each level merges pairs of blocks with + ``[[T11, 0], [-T22 L21 T11, T22]]``, batched over the pairs, so the + graph has a handful of matmuls per level rather than per block. + """ + sum_op = trt.ElementWiseOperation.SUM + leaf = min(_SOLVE_LEAF, size) + if size % leaf or (size // leaf) & (size // leaf - 1): + raise ValueError("the triangular solve size must be a power-of-two multiple of the leaf") + blocks = size // leaf + diagonal = _block_rows(network, lower, heads, blocks, leaf, np.eye(blocks)) + identity = _const(network, np.eye(leaf, dtype=np.float32)[None, None]) + # (I - N)^-1 = (I + N)(I + N^2)(I + N^4)... with N = I - D nilpotent. + nilpotent = _ew(network, identity, diagonal, trt.ElementWiseOperation.SUB) + inverse = _ew(network, identity, nilpotent, sum_op) + power = nilpotent + for _ in range(max(0, math.ceil(math.log2(leaf)) - 1)): + power = _mm(network, power, power) + inverse = _mm(network, inverse, _ew(network, identity, power, sum_op)) + width = leaf + while blocks > 1: + pairs = blocks // 2 + below = np.zeros((blocks, blocks)) + below[np.arange(1, blocks, 2), np.arange(0, blocks, 2)] = 1.0 + lower_left = _pair_half( + network, + _block_rows(network, lower, heads, blocks, width, below), + heads, + pairs, + width, + 1, + ) + upper = _pair_half(network, inverse, heads, pairs, width, 0) + bottom_right = _pair_half(network, inverse, heads, pairs, width, 1) + corner = network.add_unary( + _mm(network, _mm(network, bottom_right, lower_left), upper), trt.UnaryOperation.NEG + ).get_output(0) + zeros = _const(network, np.zeros((heads, pairs, width, width), dtype=np.float32)) + top = network.add_concatenation([upper, zeros]) + top.axis = 3 + bottom = network.add_concatenation([corner, bottom_right]) + bottom.axis = 3 + merged = network.add_concatenation([top.get_output(0), bottom.get_output(0)]) + merged.axis = 2 + inverse = merged.get_output(0) + blocks, width = pairs, width * 2 + return _shuffle(network, inverse, reshape=(heads, size, size)) + + +def add_chunked_delta_rule( + network: trt.INetworkDefinition, + *, + q: trt.ITensor, + k: trt.ITensor, + v: trt.ITensor, + g: trt.ITensor, + beta: trt.ITensor, + state: trt.ITensor, + num_heads: int, + head_dim: int, +) -> tuple[trt.ITensor, trt.ITensor]: + """Gated delta rule over one chunk, all tensors FP32. + + ``q``/``k``/``v`` are head-major ``[H, S, D]`` (``q`` already scaled), + ``g``/``beta`` are ``[H, S]`` and ``state`` is ``[H, D_k, D_v]``. + Returns ``(outputs [H, S, D_v], final state [H, D_k, D_v])``. + """ + sum_op, prod_op = trt.ElementWiseOperation.SUM, trt.ElementWiseOperation.PROD + rows = row_count(network, _shuffle(network, g, transpose=(1, 0))) + + axis = network.add_constant((), np.array(1, dtype=np.int32)).get_output(0) + gamma = network.add_cumulative(g, axis, trt.CumulativeOperation.SUM, False, False).get_output(0) + gamma_rows = _shuffle(network, gamma, reshape=(0, 0, 1)) # [H, S, 1] + gamma_cols = _shuffle(network, gamma, reshape=(0, 1, -1)) # [H, 1, S] + # exp(gamma_i - gamma_j) for j <= i; the exponent is clamped at zero so the + # masked upper triangle cannot overflow before it is multiplied by zero. + diff = _ew(network, gamma_rows, gamma_cols, trt.ElementWiseOperation.SUB) + diff = _ew( + network, + diff, + _const(network, np.zeros((1, 1, 1), np.float32)), + trt.ElementWiseOperation.MIN, + ) + decay = network.add_unary(diff, trt.UnaryOperation.EXP).get_output(0) + + def mask(offset: int) -> trt.ITensor: + full = np.tril(np.ones((SOLVE_ROWS, SOLVE_ROWS), np.float32), offset)[None] + return dynamic_slice(network, _const(network, full), (0, 0, 0), (1, rows, rows)) + + decay_strict = _ew(network, decay, mask(-1), prod_op) + decay_inclusive = _ew(network, decay, mask(0), prod_op) + + beta_rows = _shuffle(network, beta, reshape=(0, 0, 1)) # [H, S, 1] + a_matrix = _ew( + network, + _ew(network, beta_rows, decay_strict, prod_op), + _mm(network, k, k, transpose_b=True), + prod_op, + ) + padded = network.add_slice(a_matrix, (0, 0, 0), (num_heads, SOLVE_ROWS, SOLVE_ROWS), (1, 1, 1)) + padded.mode = trt.SampleMode.FILL + padded.set_input(4, network.add_constant((), np.array(0.0, dtype=np.float32)).get_output(0)) + lower = _ew( + network, + _const(network, np.eye(SOLVE_ROWS, dtype=np.float32)[None]), + padded.get_output(0), + sum_op, + ) + inverse = dynamic_slice( + network, + _inverse_unit_lower(network, lower, num_heads, SOLVE_ROWS), + (0, 0, 0), + (num_heads, rows, rows), + ) + + gamma_exp = network.add_unary(gamma_rows, trt.UnaryOperation.EXP).get_output(0) # [H, S, 1] + history = _mm(network, k, state) # [H, S, D_v]: S_0^T k_i + rhs = _ew( + network, + _ew(network, beta_rows, v, prod_op), + _ew(network, _ew(network, beta_rows, gamma_exp, prod_op), history, prod_op), + trt.ElementWiseOperation.SUB, + ) + updates = _mm(network, inverse, rhs) # U: [H, S, D_v] + + scores = _ew(network, decay_inclusive, _mm(network, q, k, transpose_b=True), prod_op) + outputs = _ew( + network, + _ew(network, gamma_exp, _mm(network, q, state), prod_op), + _mm(network, scores, updates), + sum_op, + ) + + total = network.add_reduce(g, trt.ReduceOperation.SUM, 1 << 1, keep_dims=True).get_output(0) + tail = network.add_unary( + _ew( + network, + _shuffle(network, total, reshape=(0, 0, 1)), + gamma_rows, + trt.ElementWiseOperation.SUB, + ), + trt.UnaryOperation.EXP, + ).get_output(0) # exp(gamma_S - gamma_i): [H, S, 1] + carried = _ew( + network, + network.add_unary( + _shuffle(network, total, reshape=(0, 1, 1)), trt.UnaryOperation.EXP + ).get_output(0), + state, + prod_op, + ) + final_state = _ew( + network, + carried, + _mm(network, _ew(network, k, tail, prod_op), updates, transpose_a=True), + sum_op, + ) + return outputs, final_state + + +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, (_i64(network, [0]), last), (channels, kernel)) + return total, history diff --git a/families/qwen3_5/model.py b/families/qwen3_5/model.py index 8a10c927a1..6339d7dd55 100644 --- a/families/qwen3_5/model.py +++ b/families/qwen3_5/model.py @@ -84,12 +84,18 @@ ) from . import graph_ops from . import graph_blocks +from . import deltanet_rows from .parallel import add_all_reduce_sum, broadcast_from_rank0 from .prepared_quantization import ( PreparedLinear, load_fp8_kv_scales, load_projection, prepared_format, prepare_fp8_cache, ) +# Prefill engine rows; the chunked delta rule pads its solve to this size. +PREFILL_MAX_TOKENS = deltanet_rows.SOLVE_ROWS +PREFILL_OPT_TOKENS = 32 + + def _projection_matmul(network, dtype): """Use the same original/prepared projections in target and draft graphs.""" def matmul(lhs, lhs_width, rhs_width, rhs_weights, _name): @@ -532,8 +538,23 @@ def build_engine( media_inputs: bool = False, export_hidden_state: bool = False, feature_layers: tuple[int, ...] = (), + prefill_rows: int | None = None, ) -> bytes: - """Build hybrid TRT engine with DeltaNet + attention layers.""" + """Build hybrid TRT engine with DeltaNet + attention layers. + + ``prefill_rows`` builds the prefill engine of a split bundle instead: + ``S`` token rows with one profile ``1..prefill_rows``, DeltaNet as one + chunked delta rule and attention over the cache plus the chunk. Its + logits come from the last row; present K/V cover all rows and the + recurrent outputs hold the state after the last row. Decoding stays on + the one-token engine. + """ + if prefill_rows is not None and (media_inputs or export_hidden_state or feature_layers): + raise NotImplementedError( + "Qwen3.5 prefill engines are text-only without MTP or draft features" + ) + if prefill_rows is not None and not 1 <= prefill_rows <= PREFILL_MAX_TOKENS: + raise ValueError(f"Qwen3.5 prefill_rows must be in 1..{PREFILL_MAX_TOKENS}") hidden = config.hidden_size vocab = config.vocab_size num_layers = config.num_hidden_layers @@ -596,9 +617,26 @@ def build_engine( trt_config.build_route = "-peep:match_dual_gemm=off" # --- Inputs --- - token_id = network.add_input("token_id", trt.int32, (1,)) - position_id = network.add_input("position_id", trt.int32, (1,)) - attention_mask = network.add_input("attention_mask", trt.float32, (1, attention_window)) + rows = None if prefill_rows else 1 + row_dim = -1 if prefill_rows else 1 + token_id = network.add_input("token_id", trt.int32, (row_dim,)) + position_id = network.add_input("position_id", trt.int32, (row_dim,)) + attention_mask = network.add_input( + "attention_mask", trt.float32, (row_dim, -1 if prefill_rows else attention_window) + ) + if prefill_rows: + max_rows = min(prefill_rows, max_cache_length) + opt_rows = min(PREFILL_OPT_TOKENS, max_rows) + profile = builder.create_optimization_profile() + for name in ("token_id", "position_id"): + profile.set_shape(name, (1,), (opt_rows,), (max_rows,)) + profile.set_shape( + "attention_mask", + (1, max_cache_length + 1), + (opt_rows, max_cache_length + opt_rows), + (max_rows, max_cache_length + max_rows), + ) + trt_config.add_optimization_profile(profile) # DeltaNet state inputs (conv + ssm per DeltaNet layer) conv_state_inputs = [] @@ -754,6 +792,7 @@ def layer_cast(tensor): mlp_size=mlp_size, dtype=layer_np_dtype, quant_ctx=quant_ctx, + sequence_length=rows, ) hidden_state = result["hidden"] present_conv_outputs.append(result["present_conv"]) @@ -787,6 +826,7 @@ def layer_cast(tensor): mlp_size=mlp_size, dtype=layer_np_dtype, quant_ctx=quant_ctx, + sequence_length=rows, ) hidden_state = result["hidden"] present_k_outputs.append(result["present_k"]) @@ -810,6 +850,17 @@ def layer_cast(tensor): target_features.name = "target_hidden" network.mark_output(target_features) + if prefill_rows: + # Only the last row can affect the next token. + last = network.add_elementwise( + deltanet_rows.row_count(network, hidden_state), + graph_ops.add_constant(network, (1,), np.array([1], dtype=np.int64), dtype=np.int64), + trt.ElementWiseOperation.SUB, + ).get_output(0) + hidden_state = deltanet_rows.dynamic_slice( + network, hidden_state, (last, 0), (1, hidden) + ) + # --- Final norm --- if hidden_state.dtype != work_trt_dtype: hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0) @@ -1585,8 +1636,9 @@ def _add_deltanet_layer( mlp_size: int, dtype: np.dtype = np.float32, quant_ctx: Qwen35QuantContext | None = None, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: - """Add one Gated DeltaNet layer (single-step decode). + """Add one Gated DeltaNet layer over one row, or ``S`` rows when ``sequence_length`` is None. DeltaNet uses delta-rule linear attention with: - Conv1d on QKV projection output @@ -1598,8 +1650,6 @@ def _add_deltanet_layer( Returns: {hidden, present_conv, present_ssm} """ - qk_dim = num_kv_heads * head_dim # Q and K dimension before Compact GQA/MQA K/V - # ===== 1. RMSNorm ===== normed = graph_ops.add_rms_norm( network, hidden, hidden_size, weights[f"{prefix}.input_norm"], eps_tensor, dtype=dtype @@ -1626,6 +1676,146 @@ def _add_deltanet_layer( network, normed, hidden_size, num_heads, weights, f"{prefix}.deltanet_b_proj", dtype=dtype, quant_ctx=quant_ctx ) + mixer = _add_deltanet_mixer_rows if sequence_length is None else _add_deltanet_mixer_step + mixer_output, present_conv, present_ssm = mixer( + network, + qkv=qkv, + a_raw=a_raw, + b_raw=b_raw, + conv_state_in=conv_state_in, + ssm_state_in=ssm_state_in, + weights=weights, + prefix=prefix, + d_inner=d_inner, + d_conv=d_conv, + conv_dim=conv_dim, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + dtype=dtype, + ) + + # ===== 10. Gated RMSNorm per-head: weight * norm(output) * silu(z) ===== + # The reference recurrent kernel returns the attention output in the model + # storage dtype before Qwen3_5RMSNormGated casts it back to FP32. + recurrent_output = mixer_output + rows = 1 if sequence_length is not None else -1 + head_shape = (num_heads, head_dim) if sequence_length is not None else (-1, num_heads, head_dim) + if recurrent_output.dtype != hidden.dtype: + recurrent_output = network.add_cast(recurrent_output, hidden.dtype).get_output(0) + + # HF norm operates per head_v_dim: reshape to [num_heads, head_dim], norm, reshape back + deltanet_norm_w = weights[f"{prefix}.deltanet_norm"] + # Use same eps as HF Qwen3_5RMSNormGated (config.rms_norm_eps = 1e-6) + eps_small = graph_ops.add_constant( + network, (1, 1), np.array([1e-6], dtype=np.float32), dtype=np.float32 + ) + + # Reshape output and z to [num_heads, head_dim] for per-head norm + output_heads = network.add_shuffle(recurrent_output) + output_heads.reshape_dims = head_shape + norm_input = output_heads.get_output(0) + norm_output_dtype = norm_input.dtype + if dtype != np.float32: + norm_input = network.add_cast(norm_input, trt.float32).get_output(0) + + # Per-head RMSNorm: norm each head independently + sq = network.add_elementwise(norm_input, norm_input, trt.ElementWiseOperation.PROD) + head_axis = len(head_shape) - 1 + mean = network.add_reduce( + sq.get_output(0), trt.ReduceOperation.AVG, 1 << head_axis, keep_dims=True + ) + if head_axis == 2: + eps_small = network.add_shuffle(eps_small) + eps_small.reshape_dims = (1, 1, 1) + eps_small = eps_small.get_output(0) + denom_in = network.add_elementwise(mean.get_output(0), eps_small, trt.ElementWiseOperation.SUM) + sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) + recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) + normalized = network.add_elementwise( + norm_input, recip.get_output(0), trt.ElementWiseOperation.PROD + ) + + # Reshape back and apply weight + norm_flat = network.add_shuffle(normalized.get_output(0)) + norm_flat.reshape_dims = (rows, d_inner) + gamma_t = graph_ops.add_constant(network, (1, d_inner), deltanet_norm_w, dtype=np.float32) + normed_output = network.add_elementwise( + norm_flat.get_output(0), gamma_t, trt.ElementWiseOperation.PROD + ) + normed_output_tensor = normed_output.get_output(0) + if normed_output_tensor.dtype != norm_output_dtype: + normed_output_tensor = network.add_cast(normed_output_tensor, norm_output_dtype).get_output( + 0 + ) + + # Gate: multiply by silu(z) + z_activated = graph_ops.add_activation(network, z, "silu", dtype=dtype) + gated = network.add_elementwise( + normed_output_tensor, z_activated, trt.ElementWiseOperation.PROD + ) + + # ===== 11. Output projection + residual ===== + out = _projection( + network, gated.get_output(0), d_inner, hidden_size, weights, f"{prefix}.deltanet_out_proj", dtype=dtype, quant_ctx=quant_ctx + ) + + out = add_all_reduce_sum(network, out, weights) + residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM) + hidden_after_attn = residual.get_output(0) + + # ===== 12. Post-attention norm + SwiGLU MLP + residual ===== + post_normed = graph_ops.add_rms_norm( + network, + hidden_after_attn, + hidden_size, + weights[f"{prefix}.post_attn_norm"], + eps_tensor, + dtype=dtype, + ) + + mlp_out = graph_blocks.add_swiglu_mlp( + network, + post_normed, + weights=weights, + prefix=prefix, + hidden_size=hidden_size, + mlp_size=mlp_size, + dtype=dtype, + quant_ctx=quant_ctx, + ) + + mlp_out = add_all_reduce_sum(network, mlp_out, weights) + mlp_residual = network.add_elementwise(hidden_after_attn, mlp_out, trt.ElementWiseOperation.SUM) + + return { + "hidden": mlp_residual.get_output(0), + "present_conv": present_conv, + "present_ssm": present_ssm, + } + + +def _add_deltanet_mixer_step( + network: trt.INetworkDefinition, + *, + qkv: trt.ITensor, + a_raw: trt.ITensor, + b_raw: trt.ITensor, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + weights: WeightDict, + prefix: str, + d_inner: int, + d_conv: int, + conv_dim: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + dtype: np.dtype, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]: + """Conv step and delta-rule update for one token: ``(output [1, d_inner], conv, ssm)``.""" + qk_dim = num_kv_heads * head_dim # Q and K dimension before Compact GQA/MQA K/V + # ===== 3. Conv1d step on QKV ===== # conv_state_in: [conv_dim, d_conv] # qkv: [1, conv_dim] -> [conv_dim, 1] @@ -1836,95 +2026,133 @@ def recurrent_cast(tensor: trt.ITensor) -> trt.ITensor: output_flat = network.add_shuffle(output_3d.get_output(0)) output_flat.reshape_dims = (1, d_inner) - # ===== 10. Gated RMSNorm per-head: weight * norm(output) * silu(z) ===== - # The reference recurrent kernel returns the attention output in the model - # storage dtype before Qwen3_5RMSNormGated casts it back to FP32. - recurrent_output = output_flat.get_output(0) - if recurrent_output.dtype != hidden.dtype: - recurrent_output = network.add_cast(recurrent_output, hidden.dtype).get_output(0) - # HF norm operates per head_v_dim: reshape to [num_heads, head_dim], norm, reshape back - deltanet_norm_w = weights[f"{prefix}.deltanet_norm"] - # Use same eps as HF Qwen3_5RMSNormGated (config.rms_norm_eps = 1e-6) - eps_small = graph_ops.add_constant( - network, (1, 1), np.array([1e-6], dtype=np.float32), dtype=np.float32 - ) + return output_flat.get_output(0), present_conv, present_ssm - # Reshape output and z to [num_heads, head_dim] for per-head norm - output_heads = network.add_shuffle(recurrent_output) - output_heads.reshape_dims = (num_heads, head_dim) - norm_input = output_heads.get_output(0) - norm_output_dtype = norm_input.dtype - if dtype != np.float32: - norm_input = network.add_cast(norm_input, trt.float32).get_output(0) - # Per-head RMSNorm: norm each head independently - sq = network.add_elementwise(norm_input, norm_input, trt.ElementWiseOperation.PROD) - mean = network.add_reduce(sq.get_output(0), trt.ReduceOperation.AVG, 1 << 1, keep_dims=True) - denom_in = network.add_elementwise(mean.get_output(0), eps_small, trt.ElementWiseOperation.SUM) - sqrt_l = network.add_unary(denom_in.get_output(0), trt.UnaryOperation.SQRT) - recip = network.add_unary(sqrt_l.get_output(0), trt.UnaryOperation.RECIP) - normalized = network.add_elementwise( - norm_input, recip.get_output(0), trt.ElementWiseOperation.PROD - ) - - # Reshape back and apply weight - norm_flat = network.add_shuffle(normalized.get_output(0)) - norm_flat.reshape_dims = (1, d_inner) - gamma_t = graph_ops.add_constant(network, (1, d_inner), deltanet_norm_w, dtype=np.float32) - normed_output = network.add_elementwise( - norm_flat.get_output(0), gamma_t, trt.ElementWiseOperation.PROD - ) - normed_output_tensor = normed_output.get_output(0) - if normed_output_tensor.dtype != norm_output_dtype: - normed_output_tensor = network.add_cast(normed_output_tensor, norm_output_dtype).get_output( - 0 - ) +def _add_deltanet_mixer_rows( + network: trt.INetworkDefinition, + *, + qkv: trt.ITensor, + a_raw: trt.ITensor, + b_raw: trt.ITensor, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + weights: WeightDict, + prefix: str, + d_inner: int, + d_conv: int, + conv_dim: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + dtype: np.dtype, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]: + """Conv and delta rule over ``S`` rows in one launch: ``(output [S, d_inner], conv, ssm)``. - # Gate: multiply by silu(z) - z_activated = graph_ops.add_activation(network, z, "silu", dtype=dtype) - gated = network.add_elementwise( - normed_output_tensor, z_activated, trt.ElementWiseOperation.PROD - ) + Same arithmetic as :func:`_add_deltanet_mixer_step` applied token by token; + the recurrence runs as one chunk (see :mod:`.deltanet_rows`). + """ + qk_dim = num_kv_heads * head_dim - # ===== 11. Output projection + residual ===== - out = _projection( - network, gated.get_output(0), d_inner, hidden_size, weights, f"{prefix}.deltanet_out_proj", dtype=dtype, quant_ctx=quant_ctx - ) + def shuffle(tensor, reshape=None, transpose=None): + 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) - out = add_all_reduce_sum(network, out, weights) - residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM) - hidden_after_attn = residual.get_output(0) + def fp32(tensor): + return tensor if tensor.dtype == trt.float32 else network.add_cast(tensor, trt.float32).get_output(0) - # ===== 12. Post-attention norm + SwiGLU MLP + residual ===== - post_normed = graph_ops.add_rms_norm( + conv_w = graph_ops.add_constant( + network, (conv_dim, d_conv), weights[f"{prefix}.conv1d_weight"], dtype=dtype + ) + conv_columns, present_conv = deltanet_rows.add_short_conv_rows( network, - hidden_after_attn, - hidden_size, - weights[f"{prefix}.post_attn_norm"], - eps_tensor, - dtype=dtype, + columns=shuffle(qkv, transpose=(1, 0)), + conv_state=conv_state_in, + weight=conv_w, + channels=conv_dim, + kernel=d_conv, ) + conv_out = graph_ops.add_bias_sum( + network, shuffle(conv_columns, transpose=(1, 0)), conv_dim, + weights[f"{prefix}.conv1d_bias"], dtype=dtype, + ) + activated = graph_ops.add_activation(network, conv_out, "silu", dtype=dtype) + rows = deltanet_rows.row_count(network, activated) - mlp_out = graph_blocks.add_swiglu_mlp( - network, - post_normed, - weights=weights, - prefix=prefix, - hidden_size=hidden_size, - mlp_size=mlp_size, - dtype=dtype, - quant_ctx=quant_ctx, + def columns(offset, width): + return deltanet_rows.dynamic_slice(network, activated, (0, offset), (rows, width)) + + q_normed = graph_ops.add_l2_norm( + network, shuffle(columns(0, qk_dim), reshape=(-1, num_kv_heads, head_dim)), 2, + eps=1e-6, dtype=dtype, + ) + k_normed = graph_ops.add_l2_norm( + network, shuffle(columns(qk_dim, qk_dim), reshape=(-1, num_kv_heads, head_dim)), 2, + eps=1e-6, dtype=dtype, ) + heads_per_group = num_heads // num_kv_heads - mlp_out = add_all_reduce_sum(network, mlp_out, weights) - mlp_residual = network.add_elementwise(hidden_after_attn, mlp_out, trt.ElementWiseOperation.SUM) + def expand_heads(tensor): + # [S, kv_heads, D] -> [S, heads, D], head = kv_head * heads_per_group + g. + if heads_per_group == 1: + return tensor + tile = graph_ops.add_constant( + network, (1, 1, heads_per_group, 1), + np.ones((1, 1, heads_per_group, 1), dtype=dtype), dtype=dtype, + ) + tiled = network.add_elementwise( + shuffle(tensor, reshape=(-1, num_kv_heads, 1, head_dim)), tile, + trt.ElementWiseOperation.PROD, + ).get_output(0) + return shuffle(tiled, reshape=(-1, num_heads, head_dim)) - return { - "hidden": mlp_residual.get_output(0), - "present_conv": present_conv, - "present_ssm": present_ssm, - } + def head_major(tensor): + return shuffle(fp32(tensor), transpose=(1, 0, 2)) # [H, S, D] + + q_scale = graph_ops.add_constant( + network, (1, 1, 1), np.array([[[1.0 / np.sqrt(head_dim)]]], dtype=np.float32), + dtype=np.float32, + ) + q_heads = network.add_elementwise( + head_major(expand_heads(q_normed)), q_scale, trt.ElementWiseOperation.PROD + ).get_output(0) + k_heads = head_major(expand_heads(k_normed)) + v_heads = head_major(shuffle(columns(2 * qk_dim, d_inner), reshape=(-1, num_heads, head_dim))) + + # Decay g = A * softplus(a + dt_bias) and beta = sigmoid(b), FP32, [H, S]. + a_const = graph_ops.add_constant(network, (1, num_heads), weights[f"{prefix}.A"], dtype=np.float32) + dt_bias = graph_ops.add_constant( + network, (1, num_heads), weights[f"{prefix}.dt_bias"], dtype=np.float32 + ) + biased = network.add_elementwise(fp32(a_raw), dt_bias, trt.ElementWiseOperation.SUM).get_output(0) + one = graph_ops.add_constant(network, (1, 1), np.array([1.0], dtype=np.float32), dtype=np.float32) + softplus = network.add_unary( + network.add_elementwise( + network.add_unary(biased, trt.UnaryOperation.EXP).get_output(0), one, + trt.ElementWiseOperation.SUM, + ).get_output(0), + trt.UnaryOperation.LOG, + ).get_output(0) + decay = network.add_elementwise(a_const, softplus, trt.ElementWiseOperation.PROD).get_output(0) + beta = fp32(network.add_activation(b_raw, trt.ActivationType.SIGMOID).get_output(0)) + + outputs, present_ssm = deltanet_rows.add_chunked_delta_rule( + network, + q=q_heads, + k=k_heads, + v=v_heads, + g=shuffle(decay, transpose=(1, 0)), + beta=shuffle(beta, transpose=(1, 0)), + state=fp32(ssm_state_in), + num_heads=num_heads, + head_dim=head_dim, + ) + return shuffle(outputs, transpose=(1, 0, 2), reshape=(-1, d_inner)), present_conv, present_ssm def _add_media_attention(network, q, k, v, mask, num_heads, kv_heads, head_dim, window): @@ -1990,6 +2218,7 @@ def _add_full_attention_layer( media_attention: bool = False, dtype: np.dtype = np.float32, quant_ctx: Qwen35QuantContext | None = None, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: """Add one full self-attention layer with output gating. @@ -1999,9 +2228,16 @@ def _add_full_attention_layer( - Output gating: context * sigmoid(gate) BEFORE o_proj - SwiGLU MLP after attention + ``sequence_length`` is the query row count, or None for ``S`` runtime rows + attending over the cache plus their own causal prefix. + Returns: {hidden, present_k, present_v} """ + if sequence_length is None and media_attention: + raise NotImplementedError("Qwen3.5 media attention runs one row per launch") attention_window = max_cache_length + 1 + rows = -1 if sequence_length is None else sequence_length + kv_rows = None if sequence_length is None else max_cache_length + sequence_length # Pre-attention norm normed = graph_blocks.apply_norm( @@ -2030,12 +2266,14 @@ def _add_full_attention_layer( q_norm = weights.get(f"{prefix}.q_norm") if q_norm is not None: q = graph_ops.add_rms_norm_per_head( - network, q, num_heads, head_dim, q_norm, eps_tensor, dtype=dtype + network, q, num_heads, head_dim, q_norm, eps_tensor, dtype=dtype, + sequence_length=sequence_length, ) k_norm = weights.get(f"{prefix}.k_norm") if k_norm is not None: k = graph_ops.add_rms_norm_per_head( - network, k, num_kv_heads, head_dim, k_norm, eps_tensor, dtype=dtype + network, k, num_kv_heads, head_dim, k_norm, eps_tensor, dtype=dtype, + sequence_length=sequence_length, ) # Native RoPE @@ -2048,6 +2286,7 @@ def _add_full_attention_layer( sin_half_tensor, position_id, rotary_embedding_dim, + sequence_length=sequence_length, ) k = graph_ops.add_apply_rope_native( network, @@ -2058,6 +2297,7 @@ def _add_full_attention_layer( sin_half_tensor, position_id, rotary_embedding_dim, + sequence_length=sequence_length, ) # Q stays in compute precision for one-token decoding. K/V are stored @@ -2071,9 +2311,9 @@ def _add_full_attention_layer( # Reshape K, V for concatenation k_reshape = network.add_shuffle(k) - k_reshape.reshape_dims = (1, kv_attention_size) + k_reshape.reshape_dims = (rows, kv_attention_size) v_reshape = network.add_shuffle(v) - v_reshape.reshape_dims = (1, kv_attention_size) + v_reshape.reshape_dims = (rows, kv_attention_size) # Concatenate with cache all_k = network.add_concatenation([cache_k, k_reshape.get_output(0)]) @@ -2094,8 +2334,8 @@ def _add_full_attention_layer( num_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv_heads, - q_seq=1, - kv_seq=attention_window, + q_seq=sequence_length, + kv_seq=kv_rows, mask=mask_4d, ) @@ -2337,6 +2577,11 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter", *, require_mt return if require_mtp: model._load_mtp_weights(readers, weights, config) + # Text-only bundles add an engine that prefills the prompt in multi-token chunks. + # TensorRT 11.4 cannot lower NVFP4 dynamic activation quantization over a + # runtime row count, so NVFP4 bundles keep the token loop. + nvfp4 = quantization == "nvfp4" or getattr(quant_ctx, "has_nvfp4_weights", False) + split = not visual and not require_mtp and draft_config is None and not nvfp4 plan = model.build_engine( config, weights, @@ -2351,6 +2596,12 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter", *, require_mt ) writer.add_bytes("engine.plan", plan) + if split: + writer.add_bytes("prefill.plan", model.build_engine( + config, weights, max_sequence_length, precision=precision, + verbose=bool(request.verbose), quant_ctx=quant_ctx, + prefill_rows=PREFILL_MAX_TOKENS, + )) if require_mtp: # Preserve the admitted three draft steps plus one target anchor. writer.add_bytes("mtp_engine.plan", model.build_mtp_engine( @@ -2395,7 +2646,7 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter", *, require_mt model, precision=precision, max_cache_length=max_sequence_length, - decoder_engine_layout="single", + decoder_engine_layout="split" if split else "single", **({"has_mtp": True, "mtp_num_draft_tokens": 3} if require_mtp else {}), **({"kv_cache_quantization": "fp8"} if weights.get("_fp8_kv_cache", False) else {}), ), diff --git a/families/qwen3_5/runtime/hybrid_state.cpp b/families/qwen3_5/runtime/hybrid_state.cpp index 1d412f3f8d..81399c9ee9 100644 --- a/families/qwen3_5/runtime/hybrid_state.cpp +++ b/families/qwen3_5/runtime/hybrid_state.cpp @@ -32,6 +32,18 @@ void Qwen35HybridState::advance(int32_t n_tokens) { ssm_->advance(n_tokens); } +void Qwen35HybridState::bind_prefill(ITrtModule& module) { + // The prefill context reads the shared caches and recurrent state; its + // present K/V rows stay in its own buffers until commit_prefill(). + kv_->bind_cache_inputs(module); + ssm_->bind_to(module); +} + +void Qwen35HybridState::commit_prefill(ITrtModule& module, int32_t rows) { + kv_->append_prefill_from(module, rows); + ssm_->advance(rows); +} + int32_t Qwen35HybridState::position() const { return kv_->position(); } diff --git a/families/qwen3_5/runtime/hybrid_state.h b/families/qwen3_5/runtime/hybrid_state.h index 03fd41d16b..d8e2e87731 100644 --- a/families/qwen3_5/runtime/hybrid_state.h +++ b/families/qwen3_5/runtime/hybrid_state.h @@ -21,6 +21,8 @@ class Qwen35HybridState final : public Qwen35InferenceState { void bind_to(ITrtModule& module) override; void prepare_step(TensorMap& inputs, int32_t seq_len = 1) override; void advance(int32_t n_tokens = 1) override; + void bind_prefill(ITrtModule& module) override; + void commit_prefill(ITrtModule& module, int32_t rows) override; int32_t position() const override; int32_t max_length() const override; int32_t num_layers() const override; diff --git a/families/qwen3_5/runtime/inference_state.h b/families/qwen3_5/runtime/inference_state.h index 5bf0f09f51..293626e29d 100644 --- a/families/qwen3_5/runtime/inference_state.h +++ b/families/qwen3_5/runtime/inference_state.h @@ -28,6 +28,7 @@ #include #include +#include namespace trtmc { @@ -69,6 +70,20 @@ class Qwen35InferenceState { // State types that do not distinguish the phases can ignore this. virtual void mark_prefill_complete() {} + // Bind a multi-row prefill context (a second profile of the decoder + // engine) to the same caches and recurrent state as bind_to(). + virtual void bind_prefill(ITrtModule& module) { + (void)module; + throw std::runtime_error("this Qwen3.5 state does not support batched prefill"); + } + + // Commit the `rows` tokens a bind_prefill() context just processed. + virtual void commit_prefill(ITrtModule& module, int32_t rows) { + (void)module; + (void)rows; + throw std::runtime_error("this Qwen3.5 state does not support batched prefill"); + } + // --- Queries --- // Current sequence position (0 = empty, increments with advance()). diff --git a/families/qwen3_5/runtime/kv_cache.cpp b/families/qwen3_5/runtime/kv_cache.cpp index d485babfd6..bbf392ca42 100644 --- a/families/qwen3_5/runtime/kv_cache.cpp +++ b/families/qwen3_5/runtime/kv_cache.cpp @@ -247,6 +247,19 @@ void Qwen35KvCache::append_prefill_kv(const std::vector& prefill_k, position_ += seq_len; } +void Qwen35KvCache::append_prefill_from(const ITrtModule& module, int32_t seq_len) { + std::vector keys(static_cast(num_layers_)); + std::vector values(static_cast(num_layers_)); + for (int32_t i = 0; i < num_layers_; ++i) { + const auto layer = static_cast(i); + keys[layer] = module.device_ptr(names_.present_k[layer]); + values[layer] = module.device_ptr(names_.present_v[layer]); + if (keys[layer] == nullptr || values[layer] == nullptr) + throw std::runtime_error("qwen3_5 prefill context has no present K/V output"); + } + append_prefill_kv(keys, values, seq_len); +} + void Qwen35KvCache::set_position(int32_t position) { position_ = std::max(0, std::min(position, max_length_)); } diff --git a/families/qwen3_5/runtime/kv_cache.h b/families/qwen3_5/runtime/kv_cache.h index df555b4abe..03199304a5 100644 --- a/families/qwen3_5/runtime/kv_cache.h +++ b/families/qwen3_5/runtime/kv_cache.h @@ -80,6 +80,9 @@ class Qwen35KvCache : public Qwen35InferenceState { void append_prefill_kv(const std::vector& prefill_k, const std::vector& prefill_v, int32_t seq_len); + // Append the first `seq_len` present K/V rows a prefill context produced. + void append_prefill_from(const ITrtModule& module, int32_t seq_len); + // Move the logical cache length without touching device memory. Stale rows // remain masked out by subsequent prepare_step calls. void set_position(int32_t position); diff --git a/families/qwen3_5/runtime/pipeline.cpp b/families/qwen3_5/runtime/pipeline.cpp index 54ec838892..f70fe739d7 100644 --- a/families/qwen3_5/runtime/pipeline.cpp +++ b/families/qwen3_5/runtime/pipeline.cpp @@ -31,14 +31,26 @@ RecurrentPipeline::RecurrentPipeline(std::unique_ptr decoder, const char* name, std::shared_ptr tokenizer, std::string model_id_str, std::unique_ptr mtp_scheduler, - std::unique_ptr paired_draft) + std::unique_ptr paired_draft, + std::unique_ptr prefill) : 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)), - mtp_scheduler_(std::move(mtp_scheduler)), paired_draft_(std::move(paired_draft)) { + mtp_scheduler_(std::move(mtp_scheduler)), paired_draft_(std::move(paired_draft)), + prefill_(std::move(prefill)) { if (mtp_scheduler_ && paired_draft_) throw std::runtime_error("Qwen3.5 cannot combine MTP and external draft schedulers"); if (!decoder_ || !decoder_->ok()) throw std::runtime_error(std::string(name_) + ": invalid decoder module"); + if (prefill_) { + if (!prefill_->ok() || mtp_scheduler_ || paired_draft_) + throw std::runtime_error(std::string(name_) + + ": batched prefill needs a plain text decoder"); + 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(std::min(rows[0], state_->max_length())); + } // Inspect declared inputs rather than optional Edge metadata. Older text // bundles remain loadable and never receive media-only bindings. int media_bindings = 0; @@ -205,10 +217,14 @@ std::vector RecurrentPipeline::generate_from_ids(const std::vector(input_ids.size()) <= state_->max_length()) { + run_prefill_chunks(input_ids, logits); + } else { + for (const auto token : input_ids) + run_step(token, logits); + } auto t_prefill_end = SteadyClock::now(); // ── Decode phase ── @@ -449,6 +465,31 @@ void RecurrentPipeline::report_timing(SteadyClock::time_point t_prefill_start, } } +void RecurrentPipeline::run_prefill_chunks(const std::vector& input_ids, + std::vector& logits) { + const auto total = static_cast(input_ids.size()); + state_->bind_prefill(*prefill_); + 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_->commit_prefill(*prefill_, rows); + } + prefill_->sync(); + const auto* device_logits = prefill_->device_ptr("logits"); + if (device_logits == nullptr) + throw std::runtime_error(std::string(name_) + ": prefill context 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, std::vector* hidden_state_out) { auto t0 = SteadyClock::now(); diff --git a/families/qwen3_5/runtime/pipeline.h b/families/qwen3_5/runtime/pipeline.h index e893d3b6bb..391392f728 100644 --- a/families/qwen3_5/runtime/pipeline.h +++ b/families/qwen3_5/runtime/pipeline.h @@ -53,7 +53,8 @@ class RecurrentPipeline final : public ITextGeneration { std::shared_ptr tokenizer = nullptr, std::string model_id_str = "", std::unique_ptr mtp_scheduler = nullptr, - std::unique_ptr paired_draft = nullptr); + std::unique_ptr paired_draft = nullptr, + 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; } @@ -78,6 +79,9 @@ class RecurrentPipeline final : public ITextGeneration { std::string model_id_; std::unique_ptr mtp_scheduler_; std::unique_ptr paired_draft_; + // Multi-row prefill engine (split bundles only); decoding stays on decoder_. + std::unique_ptr prefill_; + int32_t prefill_rows_{0}; bool media_inputs_{false}; int32_t media_hidden_size_{0}; DType media_dtype_{DType::kFloat16}; @@ -97,6 +101,9 @@ class RecurrentPipeline final : public ITextGeneration { const Qwen35SamplingParams& params); void run_step(int32_t token_id, std::vector& logits, std::vector* hidden_state_out = nullptr); + // 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); void* hidden_state_device_ptr_{nullptr}; std::size_t hidden_state_numel_{0}; diff --git a/families/qwen3_5/runtime/plugin.cpp b/families/qwen3_5/runtime/plugin.cpp index f2664f948a..fbd506207d 100644 --- a/families/qwen3_5/runtime/plugin.cpp +++ b/families/qwen3_5/runtime/plugin.cpp @@ -49,6 +49,8 @@ struct RuntimeConfig { bool has_mtp; std::int32_t mtp_num_draft_tokens; std::int32_t tensor_parallel_size; + // A second engine prefills prompt chunks; decoding stays on engine.plan. + bool split; }; template @@ -90,8 +92,9 @@ RuntimeConfig parse_runtime_config(const BundleReader& bundle) { (json.contains("has_mtp") ? 2 : 0) + static_cast(json.contains("tensor_parallel_size"))) throw std::runtime_error("qwen3_5 runtime.json has an unexpected field set"); - if (require_value(json, "decoder_engine_layout") != "single") - throw std::runtime_error("qwen3_5 requires decoder_engine_layout='single'"); + const auto layout = require_value(json, "decoder_engine_layout"); + if (layout != "single" && layout != "split") + throw std::runtime_error("qwen3_5 requires decoder_engine_layout 'single' or 'split'"); RuntimeConfig config{ require_value(json, "hidden_size"), @@ -120,7 +123,10 @@ RuntimeConfig parse_runtime_config(const BundleReader& bundle) { json.contains("tensor_parallel_size") ? require_value(json, "tensor_parallel_size") : 1, + layout == "split", }; + if (config.split && (config.tensor_parallel_size != 1 || config.has_mtp)) + throw std::runtime_error("qwen3_5 split bundles are single-GPU without MTP"); const auto layer_types = require_value>(json, "layer_types"); if (config.hidden_size <= 0 || config.num_layers <= 0 || config.num_attention_heads <= 0 || config.num_key_value_heads <= 0 || config.head_dim <= 0 || config.vocab_size <= 0 || @@ -173,6 +179,18 @@ ITask* create_native(const FamilyContext& context, const RuntimeConfig& config, if (!decoder || !decoder->ok()) throw std::runtime_error("Qwen3.5 native decoder failed to load"); decoder->set_timing_label("qwen3_5 decoder"); + std::unique_ptr prefill; + if (config.split) { + // The prefill engine shares the decoder's stream so state copies stay ordered. + ModuleCreateOptions prefill_options = options; + prefill_options.stream = decoder->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 || !prefill->ok()) + throw std::runtime_error("Qwen3.5 native prefill engine failed to load"); + prefill->set_timing_label("qwen3_5 prefill"); + } const cudaStream_t stream = decoder->stream(); const std::int32_t kv_dim = config.num_key_value_heads / config.tensor_parallel_size * config.head_dim; @@ -233,7 +251,7 @@ ITask* create_native(const FamilyContext& context, const RuntimeConfig& config, auto tokenizer = create_tokenizer(context.reader); auto pipeline = std::make_unique( std::move(decoder), std::move(state), std::move(generation), stream, "Qwen3.5", tokenizer, - std::string{}, std::move(mtp_scheduler), std::move(paired_draft)); + std::string{}, std::move(mtp_scheduler), std::move(paired_draft), std::move(prefill)); if (context.reader.find_section("native_visual.json")) { auto visual = load_engine( context.backend, diff --git a/families/qwen3_5/tests/test_chunked_prefill.py b/families/qwen3_5/tests/test_chunked_prefill.py new file mode 100644 index 0000000000..702d2559af --- /dev/null +++ b/families/qwen3_5/tests/test_chunked_prefill.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3.5 prefill engine: chunked prefill against the one-token engine.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors.torch") + +from .. import deltanet_rows, model # noqa: E402 +from ..config import ModelConfig # noqa: E402 +from .test_quantization import _write_checkpoint # noqa: E402 + +MAX_CACHE = 32 +MASKED = -1.0e4 + + +class _Stepper: + """Runs one engine context with host-managed KV caches and recurrent state.""" + + def __init__(self, engine, profile: int, state: dict) -> None: + self.engine = engine + self.context = engine.create_execution_context() + self.stream = torch.cuda.current_stream().cuda_stream + if engine.num_optimization_profiles > 1: + assert self.context.set_optimization_profile_async(profile, self.stream) + self.state = state + + def run(self, tokens: list[int], position: int) -> np.ndarray: + rows = len(tokens) + width = MAX_CACHE + rows + mask = torch.full((rows, width), MASKED) + mask[:, :position] = 0.0 + for row in range(rows): + mask[row, MAX_CACHE : MAX_CACHE + row + 1] = 0.0 + inputs = { + "token_id": torch.tensor(tokens, dtype=torch.int32), + "position_id": torch.arange(position, position + rows, dtype=torch.int32), + "attention_mask": mask, + } + bound = {} + for name, value in inputs.items(): + value = value.cuda() + bound[name] = value + if -1 in tuple(self.engine.get_tensor_shape(name)): + assert self.context.set_input_shape(name, tuple(value.shape)) + self.context.set_tensor_address(name, value.data_ptr()) + outputs = {} + for i in range(self.engine.num_io_tensors): + name = self.engine.get_tensor_name(i) + if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + if name in self.state: + self.context.set_tensor_address(name, self.state[name].data_ptr()) + continue + shape = tuple(self.context.get_tensor_shape(name)) + dtype = ( + torch.float16 + if self.engine.get_tensor_dtype(name) == trt.float16 + else torch.float32 + ) + outputs[name] = torch.empty(shape, dtype=dtype, device="cuda") + self.context.set_tensor_address(name, outputs[name].data_ptr()) + assert self.context.execute_async_v3(self.stream) + torch.cuda.synchronize() + for name, value in outputs.items(): + if name.startswith(("present_conv", "present_ssm")): + target = self.state[ + name.replace("present_conv", "conv_state").replace("present_ssm", "ssm_state") + ] + target.copy_(value.reshape(target.shape)) + elif name.startswith("present_"): + cache = self.state[name.replace("present_", "cache_")] + cache[position : position + rows] = value + return outputs["logits"][-1].float().cpu().numpy() + + +def _state(engine) -> dict: + state = {} + for i in range(engine.num_io_tensors): + name = engine.get_tensor_name(i) + if name.startswith(("cache_", "conv_state", "ssm_state")): + dtype = torch.float16 if engine.get_tensor_dtype(name) == trt.float16 else torch.float32 + state[name] = torch.zeros( + tuple(engine.get_tensor_shape(name)), dtype=dtype, device="cuda" + ) + return state + + +def _engine(plan: bytes): + return trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(plan) + + +@pytest.mark.gpu +@pytest.mark.trt +def test_chunked_prefill_matches_one_token_engine(tmp_path: Path) -> None: + """Prefill chunks leave the logits, KV rows and recurrent state of token-by-token steps.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + _write_checkpoint(tmp_path, quantized=False) + config = ModelConfig.from_dir(tmp_path) + instance = model._Qwen35Model() + weights = instance.load_weights(str(tmp_path), config) + single = _engine(instance.build_engine(config, weights, MAX_CACHE, precision="fp32")) + prefill_engine = _engine( + instance.build_engine( + config, weights, MAX_CACHE, precision="fp32", prefill_rows=model.PREFILL_MAX_TOKENS + ) + ) + assert prefill_engine.num_optimization_profiles == 1 + assert [tuple(s) for s in prefill_engine.get_tensor_profile_shape("token_id", 0)] == [ + (1,), + (min(model.PREFILL_OPT_TOKENS, MAX_CACHE),), + (MAX_CACHE,), + ] + + prompt, next_token = [3, 17, 5, 9, 30, 2, 11, 21, 7, 14, 26], 8 + want_state = _state(single) + stepper = _Stepper(single, 0, want_state) + want = [stepper.run([token], position) for position, token in enumerate(prompt)] + want.append(stepper.run([next_token], len(prompt))) + + # Prefill chunks on the prefill engine, then decode on the one-token engine. + got_state = _state(single) + prefill = _Stepper(prefill_engine, 0, got_state) + decode = _Stepper(single, 0, got_state) + got = [ + prefill.run(prompt[:7], 0), + prefill.run(prompt[7:], 7), + decode.run([next_token], len(prompt)), + ] + + # FP32 engines still run matmuls in TF32, and one row and several rows take + # different kernels, so compare at a fraction of each tensor's scale. + errors = {} + for step, (logits, reference) in enumerate(zip(got, (want[6], want[10], want[11]))): + assert np.argmax(logits) == np.argmax(reference), step + errors[f"logits[{step}]"] = np.abs(logits - reference).max() / np.abs(reference).max() + for name, reference in want_state.items(): + reference = reference.float().cpu().numpy() + value = got_state[name].float().cpu().numpy() + if name.startswith("cache_"): + reference, value = reference[: len(prompt) + 1], value[: len(prompt) + 1] + errors[name] = np.abs(value - reference).max() / max(np.abs(reference).max(), 1e-6) + assert max(errors.values()) < 5e-3, errors + + +@pytest.mark.trt +def test_prefill_engine_rejects_media_and_speculative_builds(tmp_path: Path) -> None: + """Media, MTP hidden-state and draft-feature engines keep one row per launch.""" + _write_checkpoint(tmp_path, quantized=False) + config = ModelConfig.from_dir(tmp_path) + instance = model._Qwen35Model() + weights = instance.load_weights(str(tmp_path), config) + for options in ( + {"media_inputs": True}, + {"export_hidden_state": True}, + {"feature_layers": (0,)}, + ): + with pytest.raises(NotImplementedError, match="prefill engines"): + instance.build_engine(config, weights, MAX_CACHE, prefill_rows=16, **options) + with pytest.raises(ValueError, match="prefill_rows"): + instance.build_engine(config, weights, MAX_CACHE, prefill_rows=model.PREFILL_MAX_TOKENS + 1) + + +@pytest.mark.gpu +@pytest.mark.trt +def test_triangular_solve_matches_numpy_on_nearly_parallel_keys() -> None: + """The batched block inverse stays exact where a Neumann series over 64 rows breaks down.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + heads, size = 2, deltanet_rows.SOLVE_ROWS + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + lower = network.add_input("lower", trt.float32, (heads, size, size)) + inverse = deltanet_rows._inverse_unit_lower(network, lower, heads, size) + inverse.name = "inverse" + network.mark_output(inverse) + config = builder.create_builder_config() + # Check the algorithm in full FP32; engines may run these matmuls in TF32. + config.clear_flag(trt.BuilderFlag.TF32) + engine = trt.Runtime(logger).deserialize_cuda_engine( + builder.build_serialized_network(network, config) + ) + context = engine.create_execution_context() + rng = np.random.default_rng(0) + # beta ~ 1 and nearly identical keys: strictly lower entries close to one. + strict = np.tril(rng.uniform(0.9, 1.0, (heads, size, size)), -1) + matrix = (np.eye(size) + strict).astype(np.float32) + source = torch.from_numpy(matrix).cuda() + result = torch.empty_like(source) + context.set_tensor_address("lower", source.data_ptr()) + context.set_tensor_address("inverse", result.data_ptr()) + assert context.execute_async_v3(torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize() + want = np.linalg.inv(matrix.astype(np.float64)) + got = result.cpu().numpy() + assert np.abs(got - want).max() <= 1e-4 * np.abs(want).max() From 50e18ddddb526d922cca7f027e2cfecf06c16718 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 19:37:32 +0000 Subject: [PATCH 2/3] fix(qwen3_5): take the chunk's total decay from the cumulative sum The chunk state's exp(gamma_S - gamma_i) used a separate reduction for gamma_S. Summed in another order it can exceed gamma_i by many ulps at large decays and overflow; the same pattern made the Nemotron-H SSD chunk state non-finite on some prompts. gamma_S is now the last cumulative sum, and the exponent is clamped at zero. Signed-off-by: chaofengw --- families/qwen3_5/deltanet_rows.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/families/qwen3_5/deltanet_rows.py b/families/qwen3_5/deltanet_rows.py index 735b94493b..23fb322e37 100644 --- a/families/qwen3_5/deltanet_rows.py +++ b/families/qwen3_5/deltanet_rows.py @@ -246,13 +246,24 @@ def mask(offset: int) -> trt.ITensor: sum_op, ) - total = network.add_reduce(g, trt.ReduceOperation.SUM, 1 << 1, keep_dims=True).get_output(0) + # gamma_S is the last cumulative sum, not a separate reduction: summed in + # another order it can exceed gamma_i by many ulps at large |g|, and + # exp(gamma_S - gamma_i) would then overflow. A float running sum of + # non-positive terms never increases, and the clamp keeps the exponent at + # most zero regardless. + last = _ew(network, rows, _i64(network, [1]), trt.ElementWiseOperation.SUB) + total = dynamic_slice(network, gamma, (_i64(network, [0]), last), (num_heads, 1)) # [H, 1] tail = network.add_unary( _ew( network, - _shuffle(network, total, reshape=(0, 0, 1)), - gamma_rows, - trt.ElementWiseOperation.SUB, + _ew( + network, + _shuffle(network, total, reshape=(0, 0, 1)), + gamma_rows, + trt.ElementWiseOperation.SUB, + ), + _const(network, np.zeros((1, 1, 1), np.float32)), + trt.ElementWiseOperation.MIN, ), trt.UnaryOperation.EXP, ).get_output(0) # exp(gamma_S - gamma_i): [H, S, 1] From 306b97b6a5fb9d4828c2eca33d80ec9aa95bfe10 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Wed, 11 Nov 2026 04:51:26 +0000 Subject: [PATCH 3/3] fix(qwen3_5): keep the bundle loadable and drain prefill on errors The prefill engine carries its own copy of every weight; build it only when both engines' weights fit in 80% of GPU memory, and run prompts through the decode engine otherwise. Drain the prefill stream before rethrowing an error from the chunk loop, so queued chunks never read a released request. Signed-off-by: chaofengw --- families/qwen3_5/model.py | 35 ++++++++++++++++++- families/qwen3_5/runtime/pipeline.cpp | 26 +++++++++----- .../qwen3_5/tests/test_chunked_prefill.py | 10 ++++++ 3 files changed, 62 insertions(+), 9 deletions(-) diff --git a/families/qwen3_5/model.py b/families/qwen3_5/model.py index 6339d7dd55..e9efd96cb9 100644 --- a/families/qwen3_5/model.py +++ b/families/qwen3_5/model.py @@ -2456,6 +2456,38 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: dispatch_build(request, writer, _build_native) +# Share of GPU memory that the decode and prefill engines' weights may take together. +_TWO_ENGINE_MEMORY_FRACTION = 0.8 + + +def _two_engines_fit(weight_bytes: int, gpu_bytes: int | None) -> bool: + """Whether the prefill engine, which carries its own copy of every weight, fits beside the decode engine. + + Without it, prompts run through the decode engine one token at a time (the bundle stays loadable). + """ + if gpu_bytes is None or 2 * weight_bytes <= gpu_bytes * _TWO_ENGINE_MEMORY_FRACTION: + return True + print(f"Skipping the prefill engine: the decode and prefill engines need about {2 * weight_bytes / 1024**3:.0f} " + f"GiB of weights and the GPU has {gpu_bytes / 1024**3:.0f} GiB; prompts run through the decode engine", + file=sys.stderr) + return False + + +def _host_weight_bytes(weights) -> int: + return sum(value.nbytes for value in weights.values() if isinstance(value, np.ndarray)) + + +def _gpu_memory_bytes() -> int | None: + """Total memory of the current CUDA device, or None when no device can be queried.""" + try: + import torch + except ImportError: + return None + if not torch.cuda.is_available(): + return None + return int(torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory) + + def _build_native(request: "BuildRequest", writer: "BundleWriter", *, require_mtp: bool = False, draft_dir: Path | None = None, draft_variant: str = "dflash") -> None: """Build one Qwen3.5 hybrid bundle through family-owned code only.""" @@ -2581,7 +2613,8 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter", *, require_mt # TensorRT 11.4 cannot lower NVFP4 dynamic activation quantization over a # runtime row count, so NVFP4 bundles keep the token loop. nvfp4 = quantization == "nvfp4" or getattr(quant_ctx, "has_nvfp4_weights", False) - split = not visual and not require_mtp and draft_config is None and not nvfp4 + split = (not visual and not require_mtp and draft_config is None and not nvfp4 + and _two_engines_fit(_host_weight_bytes(weights), _gpu_memory_bytes())) plan = model.build_engine( config, weights, diff --git a/families/qwen3_5/runtime/pipeline.cpp b/families/qwen3_5/runtime/pipeline.cpp index f70fe739d7..703d5f0eaa 100644 --- a/families/qwen3_5/runtime/pipeline.cpp +++ b/families/qwen3_5/runtime/pipeline.cpp @@ -469,14 +469,24 @@ void RecurrentPipeline::run_prefill_chunks(const std::vector& input_ids std::vector& logits) { const auto total = static_cast(input_ids.size()); state_->bind_prefill(*prefill_); - 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_->commit_prefill(*prefill_, rows); + try { + 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_->commit_prefill(*prefill_, rows); + } + } catch (...) { + // Queued chunks may still read the request's tokens: drain them before the error + // releases the request. A failure of the drain must not replace the original error. + try { + prefill_->sync(); + } catch (...) { + } + throw; } prefill_->sync(); const auto* device_logits = prefill_->device_ptr("logits"); diff --git a/families/qwen3_5/tests/test_chunked_prefill.py b/families/qwen3_5/tests/test_chunked_prefill.py index 702d2559af..0acfdf1506 100644 --- a/families/qwen3_5/tests/test_chunked_prefill.py +++ b/families/qwen3_5/tests/test_chunked_prefill.py @@ -203,3 +203,13 @@ def test_triangular_solve_matches_numpy_on_nearly_parallel_keys() -> None: want = np.linalg.inv(matrix.astype(np.float64)) got = result.cpu().numpy() assert np.abs(got - want).max() <= 1e-4 * np.abs(want).max() + + +def test_the_prefill_engine_is_built_only_when_a_second_weight_copy_fits() -> None: + gib = 1024**3 + assert model._two_engines_fit(18 * gib, 121 * gib) # 9B FP16 on a 128 GB unified-memory GPU + assert not model._two_engines_fit(54 * gib, 121 * gib) # 27B FP16: the decode engine alone + assert model._two_engines_fit(54 * gib, 275 * gib) + assert model._two_engines_fit(54 * gib, None) # no device to ask: keep the prefill engine + weights = {"w": np.zeros((4, 8), dtype=np.float16), "_rows": 4} + assert model._host_weight_bytes(weights) == 64