From d27a33d2a08b0c56bff37ed4078c2a69a4d9af3a Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 16:21:42 +0000 Subject: [PATCH 1/2] perf(nemotron_h): prefill single-GPU prompts in multi-token chunks Single-GPU dense bundles get a second engine (prefill.plan) that runs 1..128 prompt rows per launch: Mamba-2 layers run each chunk as one SSD chunk and 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/nemotron_h/graph_blocks.py | 22 +- families/nemotron_h/model.py | 298 ++++++++++++++++-- families/nemotron_h/runtime/hybrid_state.cpp | 12 + families/nemotron_h/runtime/hybrid_state.h | 2 + families/nemotron_h/runtime/inference_state.h | 15 + families/nemotron_h/runtime/kv_cache.cpp | 13 + families/nemotron_h/runtime/kv_cache.h | 3 + families/nemotron_h/runtime/pipeline.cpp | 52 ++- families/nemotron_h/runtime/pipeline.h | 9 +- families/nemotron_h/runtime/plugin.cpp | 23 +- families/nemotron_h/ssd_rows.py | 228 ++++++++++++++ .../nemotron_h/tests/test_chunked_prefill.py | 168 ++++++++++ 12 files changed, 806 insertions(+), 39 deletions(-) create mode 100644 families/nemotron_h/ssd_rows.py create mode 100644 families/nemotron_h/tests/test_chunked_prefill.py diff --git a/families/nemotron_h/graph_blocks.py b/families/nemotron_h/graph_blocks.py index 8465ab4ee7..ad37d7898c 100644 --- a/families/nemotron_h/graph_blocks.py +++ b/families/nemotron_h/graph_blocks.py @@ -168,9 +168,13 @@ def add_attention_block( rotary_embedding_dim: int = 0, interleaved_rope: bool = False, decompose_attention: bool = False, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: """Pre-norm -> QKV -> RoPE -> cache concat -> attention -> output proj. + ``sequence_length`` is the query row count, or None for ``S`` runtime rows + (the mask then covers the cache plus the rows). + Returns {"normed": ..., "attn_out": ..., "present_k": ..., "present_v": ...}. Does NOT apply residual -- callers compose the residual pattern. @@ -181,7 +185,8 @@ def add_attention_block( native IAttention. """ matmul = _make_matmul_fn(network, dtype, quant_ctx) - attention_window = max_cache_length + 1 + if sequence_length is None and alibi_slopes_tensor is not None: + raise NotImplementedError("ALiBi attention runs one row per launch") if num_kv_heads is None: num_kv_heads = num_heads if kv_attention_size is None: @@ -213,12 +218,14 @@ def add_attention_block( 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, ) present_k = k present_v = v @@ -231,14 +238,15 @@ def add_attention_block( present_k, k, cache_k = _fp8_cache_qdq(network, k, cache_k, k_scale) present_v, v, cache_v = _fp8_cache_qdq(network, v, cache_v, v_scale) k_reshape = network.add_shuffle(k) - k_reshape.reshape_dims = (1, kv_attention_size) + rows = -1 if sequence_length is None else sequence_length + 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) all_k = network.add_concatenation([cache_k, k_reshape.get_output(0)]) all_k.axis = 0 all_v = network.add_concatenation([cache_v, v_reshape.get_output(0)]) all_v.axis = 0 - kv_seq = attention_window + kv_seq = None if sequence_length is None else max_cache_length + sequence_length if alibi_slopes_tensor is not None: if alibi_indices_tensor is None: raise ValueError("ALiBi attention requires cache position indices") @@ -260,7 +268,7 @@ def add_attention_block( num_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv_heads, - q_seq=1, + q_seq=sequence_length, kv_seq=kv_seq, causal=False, mask=mask_4d, diff --git a/families/nemotron_h/model.py b/families/nemotron_h/model.py index 50d18aa190..bb794c1563 100644 --- a/families/nemotron_h/model.py +++ b/families/nemotron_h/model.py @@ -68,10 +68,16 @@ ) from . import graph_ops from . import graph_blocks +from . import ssd_rows from .parallel import normalize_parallel_config from .prepared_quantization import PreparedPolicy, PreparedQuantContext, checkpoint_policy +# Prefill engine rows, bounded by the SSD chunk masks. +PREFILL_MAX_TOKENS = ssd_rows.MAX_ROWS +PREFILL_OPT_TOKENS = 64 + + def _parse_layer_types(pattern: str) -> list[str]: """Parse hybrid_override_pattern: M=mamba2, -=mlp, *=attention.""" mapping = {"M": "mamba2", "-": "mlp", "*": "attention", "E": "moe"} @@ -423,8 +429,27 @@ def build_engine( debug_layer_outputs: bool = False, parallel_config=None, target_feature_layers: tuple[int, ...] = (), + prefill_rows: int | None = None, + release_weights: bool = True, ) -> bytes: - """Build hybrid TRT engine, optionally exporting paired-draft features.""" + """Build hybrid TRT engine, optionally exporting paired-draft features. + + ``prefill_rows`` builds the prefill engine of a split bundle instead: + ``S`` token rows with one profile ``1..prefill_rows``, Mamba-2 as one SSD + chunk 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. + + ``release_weights`` drops the loaded checkpoint before the TensorRT + build; pass False to build another engine from the same weights. + """ + if prefill_rows is not None and (target_feature_layers or parallel_config is not None + and normalize_parallel_config(parallel_config).enabled): + raise NotImplementedError( + "Nemotron-H prefill engines are single-GPU without paired-draft features" + ) + if prefill_rows is not None and not 1 <= prefill_rows <= PREFILL_MAX_TOKENS: + raise ValueError(f"Nemotron-H prefill_rows must be in 1..{PREFILL_MAX_TOKENS}") if ( len(set(target_feature_layers)) != len(target_feature_layers) or any(type(index) is not int or not 0 <= index < config.num_hidden_layers @@ -510,9 +535,26 @@ def build_engine( trt_config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED # --- 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) conv_state_inputs = [] ssm_state_inputs = [] @@ -615,6 +657,7 @@ def build_engine( n_groups=n_groups, dtype=layer_np_dtype, quant_ctx=quant_ctx, + sequence_length=rows, ) hidden_state = result["hidden"] present_conv_outputs.append(result["present_conv"]) @@ -674,6 +717,7 @@ def build_engine( eps_tensor=layer_eps, dtype=layer_np_dtype, quant_ctx=quant_ctx, + sequence_length=rows, ) # add_attention_block does NOT apply residual residual = network.add_elementwise( @@ -689,6 +733,10 @@ def build_engine( if layer_idx in target_feature_layers and layer_idx != num_layers - 1: target_features[layer_idx] = hidden_state + if prefill_rows: + # Only the last row can affect the next token. + hidden_state = ssd_rows.last_row(network, hidden_state, hidden) + # --- Final norm --- if hidden_state.dtype != io_trt_dtype: hidden_state = network.add_cast(hidden_state, io_trt_dtype).get_output(0) @@ -771,7 +819,8 @@ def build_engine( # Every tensor the network needs is already embedded in it; drop the # resident checkpoint before the engine build, which runs its own # TensorRT workspace allocation on top of whatever weights hold. - weights.clear() + if release_weights: + weights.clear() plan = builder.build_serialized_network(network, trt_config) if plan is None: raise RuntimeError("TensorRT engine build failed") @@ -840,8 +889,9 @@ def _add_mamba2_layer( n_groups: int, dtype: np.dtype = np.float32, quant_ctx: NemotronHQuantContext | PreparedQuantContext | None = None, + sequence_length: int | None = 1, ) -> dict[str, trt.ITensor]: - """Add one Mamba-2 SSD layer (single-step decode). + """Add one Mamba-2 SSD layer over one row, or ``S`` rows when ``sequence_length`` is None. Mamba-2 in_proj splits: [gate(d_inner), hidden_B_C(conv_dim), dt(nheads)] Conv1d operates on hidden_B_C (d_inner + 2*n_groups*d_state channels). @@ -850,8 +900,6 @@ def _add_mamba2_layer( Returns: {hidden, present_conv, present_ssm} """ - groups_state_size = n_groups * d_state - # ===== 1. RMSNorm ===== normed = graph_ops.add_rms_norm( network, hidden, hidden_size, weights[f"{prefix}.input_norm"], eps_tensor, dtype=dtype @@ -868,7 +916,66 @@ def _add_mamba2_layer( f"{prefix}.mamba_in_proj", dtype=dtype, quant_ctx=quant_ctx, - ) # [1, proj_dim] + ) # [rows, proj_dim] + + mixer = _add_mamba2_mixer_rows if sequence_length is None else _add_mamba2_mixer_step + gated_tensor, present_conv, present_ssm = mixer( + network, + projected=projected, + conv_state_in=conv_state_in, + ssm_state_in=ssm_state_in, + weights=weights, + prefix=prefix, + d_inner=d_inner, + d_state=d_state, + d_conv=d_conv, + conv_dim=conv_dim, + mamba_num_heads=mamba_num_heads, + mamba_head_dim=mamba_head_dim, + n_groups=n_groups, + dtype=dtype, + ) + + # ===== 8. Output projection + residual ===== + out = _projection( + network, + gated_tensor, + d_inner, + hidden_size, + weights, + f"{prefix}.mamba_out_proj", + dtype=dtype, + quant_ctx=quant_ctx, + ) + + residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM) + + return { + "hidden": residual.get_output(0), + "present_conv": present_conv, + "present_ssm": present_ssm, + } + + +def _add_mamba2_mixer_step( + network: trt.INetworkDefinition, + *, + projected: trt.ITensor, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + weights: WeightDict, + prefix: str, + d_inner: int, + d_state: int, + d_conv: int, + conv_dim: int, + mamba_num_heads: int, + mamba_head_dim: int, + n_groups: int, + dtype: np.dtype, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]: + """Conv step, SSM update and gated norm for one token: ``(gated [1, d_inner], conv, ssm)``.""" + groups_state_size = n_groups * d_state # Split: gate [d_inner], hidden_B_C [conv_dim], dt [nheads] offset = 0 @@ -1120,25 +1227,145 @@ def _add_mamba2_layer( if gated_tensor.dtype != norm_output_dtype: gated_tensor = network.add_cast(gated_tensor, norm_output_dtype).get_output(0) - # ===== 8. Output projection + residual ===== - out = _projection( + return gated_tensor, present_conv, present_ssm + + +def _add_mamba2_mixer_rows( + network: trt.INetworkDefinition, + *, + projected: trt.ITensor, + conv_state_in: trt.ITensor, + ssm_state_in: trt.ITensor, + weights: WeightDict, + prefix: str, + d_inner: int, + d_state: int, + d_conv: int, + conv_dim: int, + mamba_num_heads: int, + mamba_head_dim: int, + n_groups: int, + dtype: np.dtype, +) -> tuple[trt.ITensor, trt.ITensor, trt.ITensor]: + """Conv, SSD and gated norm over ``S`` rows in one launch: ``(gated [S, d_inner], conv, ssm)``. + + Same arithmetic as :func:`_add_mamba2_mixer_step` applied token by token; + the SSM runs as one chunk (see :mod:`.ssd_rows`). + """ + groups_state_size = n_groups * d_state + rows = ssd_rows.row_count(network, projected) + + def columns(tensor, offset, width): + return ssd_rows.dynamic_slice(network, tensor, (0, offset), (rows, width)) + + def fp32(tensor): + if tensor.dtype == trt.float32: + return tensor + return network.add_cast(tensor, trt.float32).get_output(0) + + gate = columns(projected, 0, d_inner) + hidden_b_c = columns(projected, d_inner, conv_dim) + dt_raw = columns(projected, d_inner + conv_dim, mamba_num_heads) + + conv_w = graph_ops.add_constant( + network, (conv_dim, d_conv), weights[f"{prefix}.conv1d_weight"], dtype=dtype + ) + conv_columns, present_conv = ssd_rows.add_short_conv_rows( network, - gated_tensor, - d_inner, - hidden_size, - weights, - f"{prefix}.mamba_out_proj", + columns=ssd_rows.shuffle(network, hidden_b_c, 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, + ssd_rows.shuffle(network, conv_columns, transpose=(1, 0)), + conv_dim, + weights[f"{prefix}.conv1d_bias"], dtype=dtype, - quant_ctx=quant_ctx, + rank=2, ) + activated = graph_ops.add_activation(network, conv_out, "silu", dtype=dtype) + hidden_x = columns(activated, 0, d_inner) + heads_per_group = mamba_num_heads // n_groups - residual = network.add_elementwise(hidden, out, trt.ElementWiseOperation.SUM) + def heads_from_groups(tensor): + # [S, groups * N] -> head-major [H, S, N], head = group * heads_per_group + j. + grouped = ssd_rows.shuffle(network, tensor, reshape=(-1, n_groups, 1, d_state)) + if heads_per_group > 1: + tile = graph_ops.add_constant( + network, (1, 1, heads_per_group, 1), + np.ones((1, 1, heads_per_group, 1), dtype=dtype), dtype=dtype, + ) + grouped = network.add_elementwise(grouped, tile, trt.ElementWiseOperation.PROD).get_output(0) + heads = ssd_rows.shuffle(network, grouped, reshape=(-1, mamba_num_heads, d_state)) + return ssd_rows.shuffle(network, fp32(heads), transpose=(1, 0, 2)) - return { - "hidden": residual.get_output(0), - "present_conv": present_conv, - "present_ssm": present_ssm, - } + b_heads = heads_from_groups(columns(activated, d_inner, groups_state_size)) + c_heads = heads_from_groups(columns(activated, d_inner + groups_state_size, groups_state_size)) + x_heads = ssd_rows.shuffle( + network, + fp32(ssd_rows.shuffle(network, hidden_x, reshape=(-1, mamba_num_heads, mamba_head_dim))), + transpose=(1, 0, 2), + ) # [H, S, P] + + # dt = softplus(dt_raw + dt_bias), FP32 (dt_bias reaches 33.5). + dt_bias = graph_ops.add_constant( + network, (1, mamba_num_heads), weights[f"{prefix}.dt_bias"], dtype=dtype + ) + dt_biased = fp32(network.add_elementwise(dt_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) + dt = network.add_unary( + network.add_elementwise( + network.add_unary(dt_biased, trt.UnaryOperation.EXP).get_output(0), one, + trt.ElementWiseOperation.SUM, + ).get_output(0), + trt.UnaryOperation.LOG, + ).get_output(0) # [S, H] + + y_heads, present_ssm = ssd_rows.add_chunked_ssd( + network, + x=x_heads, + b=b_heads, + c=c_heads, + dt=ssd_rows.shuffle(network, dt, transpose=(1, 0)), + a=weights[f"{prefix}.A"], + d=weights[f"{prefix}.D"], + state=fp32(ssm_state_in), + num_heads=mamba_num_heads, + ) + y = ssd_rows.shuffle(network, y_heads, transpose=(1, 0, 2), reshape=(-1, d_inner)) + if y.dtype != gate.dtype: + y = network.add_cast(y, gate.dtype).get_output(0) + + # Gated group RMSNorm (norm_before_gate=False), as in the one-token mixer. + gate_activated = graph_ops.add_activation(network, gate, "silu", dtype=dtype) + y_gated = network.add_elementwise(y, gate_activated, trt.ElementWiseOperation.PROD).get_output(0) + group_size = d_inner // n_groups + norm_input = ssd_rows.shuffle(network, y_gated, reshape=(-1, n_groups, group_size)) + norm_output_dtype = norm_input.dtype + if dtype != np.float32: + norm_input = network.add_cast(norm_input, trt.float32).get_output(0) + eps_small = graph_ops.add_constant( + network, (1, 1, 1), np.array([[[1e-5]]], dtype=np.float32), dtype=np.float32 + ) + square = network.add_elementwise(norm_input, norm_input, trt.ElementWiseOperation.PROD) + mean = network.add_reduce(square.get_output(0), trt.ReduceOperation.AVG, 1 << 2, keep_dims=True) + denom = network.add_elementwise(mean.get_output(0), eps_small, trt.ElementWiseOperation.SUM) + recip = network.add_unary( + network.add_unary(denom.get_output(0), trt.UnaryOperation.SQRT).get_output(0), + trt.UnaryOperation.RECIP, + ) + normalized = network.add_elementwise(norm_input, recip.get_output(0), trt.ElementWiseOperation.PROD) + gamma_t = graph_ops.add_constant(network, (1, d_inner), weights[f"{prefix}.mamba_norm"], dtype=np.float32) + gated = network.add_elementwise( + ssd_rows.shuffle(network, normalized.get_output(0), reshape=(-1, d_inner)), gamma_t, + trt.ElementWiseOperation.PROD, + ).get_output(0) + if gated.dtype != norm_output_dtype: + gated = network.add_cast(gated, norm_output_dtype).get_output(0) + return gated, present_conv, present_ssm def _add_mlp_layer( @@ -1370,6 +1597,27 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter") -> None: writer.add_bytes(f"engine.rank{rank}.plan", plan) layout = "dual_profile" else: + # Single-GPU bundles without MoE layers add an engine that prefills + # the prompt in multi-token chunks; it is built first because the + # decoder build releases the weights. TensorRT 11.4 cannot lower + # NVFP4 dynamic activation quantization over a runtime row count, + # so NVFP4 bundles keep the token loop. + nvfp4 = (policy is not None and policy.request_format == "nvfp4") or getattr( + quant_ctx, "has_nvfp4_weights", False) + split = not nvfp4 and "moe" not in _config_layer_types(config.raw, config.num_hidden_layers) + prefill_plan = None + if split: + prefill_plan = model.build_engine( + config, + weights, + max_sequence_length, + precision=precision, + quant_ctx=quant_ctx, + verbose=bool(request.verbose), + parallel_config=parallel, + prefill_rows=PREFILL_MAX_TOKENS, + release_weights=False, + ) plan = model.build_engine( config, weights, @@ -1381,7 +1629,9 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter") -> None: parallel_config=parallel, ) writer.add_bytes("engine.plan", plan) - layout = "single" + if prefill_plan is not None: + writer.add_bytes("prefill.plan", prefill_plan) + layout = "split" if split else "single" writer.add_json( "runtime.json", _runtime_config( diff --git a/families/nemotron_h/runtime/hybrid_state.cpp b/families/nemotron_h/runtime/hybrid_state.cpp index ebc83658fd..444a3d8c66 100644 --- a/families/nemotron_h/runtime/hybrid_state.cpp +++ b/families/nemotron_h/runtime/hybrid_state.cpp @@ -32,6 +32,18 @@ void NemotronHHybridState::advance(int32_t n_tokens) { ssm_->advance(n_tokens); } +void NemotronHHybridState::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 NemotronHHybridState::commit_prefill(ITrtModule& module, int32_t rows) { + kv_->append_prefill_from(module, rows); + ssm_->advance(rows); +} + int32_t NemotronHHybridState::position() const { return kv_->position(); } diff --git a/families/nemotron_h/runtime/hybrid_state.h b/families/nemotron_h/runtime/hybrid_state.h index 60844af88c..59342a551f 100644 --- a/families/nemotron_h/runtime/hybrid_state.h +++ b/families/nemotron_h/runtime/hybrid_state.h @@ -22,6 +22,8 @@ class NemotronHHybridState final : public NemotronHInferenceState { 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/nemotron_h/runtime/inference_state.h b/families/nemotron_h/runtime/inference_state.h index 94390ecf70..ca55f05cc8 100644 --- a/families/nemotron_h/runtime/inference_state.h +++ b/families/nemotron_h/runtime/inference_state.h @@ -28,6 +28,7 @@ #include #include +#include namespace trtmc { @@ -69,6 +70,20 @@ class NemotronHInferenceState { // 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 Nemotron-H 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 Nemotron-H state does not support batched prefill"); + } + // --- Queries --- // Current sequence position (0 = empty, increments with advance()). diff --git a/families/nemotron_h/runtime/kv_cache.cpp b/families/nemotron_h/runtime/kv_cache.cpp index e4c215cbc0..abf8125473 100644 --- a/families/nemotron_h/runtime/kv_cache.cpp +++ b/families/nemotron_h/runtime/kv_cache.cpp @@ -254,6 +254,19 @@ void NemotronHKvCache::set_position(int32_t position) { position_ = std::max(0, std::min(position, max_length_)); } +void NemotronHKvCache::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("nemotron_h prefill context has no present K/V output"); + } + append_prefill_kv(keys, values, seq_len); +} + void NemotronHKvCache::advance(int32_t n_tokens) { // For now, only single-token advance is supported. // n_tokens > 1 reserved for future batched prefill (TASK-10). diff --git a/families/nemotron_h/runtime/kv_cache.h b/families/nemotron_h/runtime/kv_cache.h index 1180b4b5a4..cb708f164e 100644 --- a/families/nemotron_h/runtime/kv_cache.h +++ b/families/nemotron_h/runtime/kv_cache.h @@ -87,6 +87,9 @@ class NemotronHKvCache : public NemotronHInferenceState { // copies them via write_prefill_kv(). void bind_cache_inputs(ITrtModule& module); + // Append the first `seq_len` present K/V rows a prefill context produced. + void append_prefill_from(const ITrtModule& module, int32_t seq_len); + private: void write_position_input(TensorMap& inputs, int32_t seq_len); void write_batched_mask(TensorMap& inputs, int32_t seq_len); diff --git a/families/nemotron_h/runtime/pipeline.cpp b/families/nemotron_h/runtime/pipeline.cpp index 31b07c48bd..2c4b42374b 100644 --- a/families/nemotron_h/runtime/pipeline.cpp +++ b/families/nemotron_h/runtime/pipeline.cpp @@ -32,12 +32,23 @@ RecurrentPipeline::RecurrentPipeline(std::unique_ptr decoder, RecurrentGenConfig config, cudaStream_t stream, const char* name, std::shared_ptr tokenizer, std::string model_id_str, - std::unique_ptr paired) + std::unique_ptr paired, + 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)), - paired_(std::move(paired)) { + paired_(std::move(paired)), prefill_(std::move(prefill)) { if (!decoder_ || !decoder_->ok()) throw std::runtime_error(std::string(name_) + ": invalid decoder module"); + if (prefill_) { + if (!prefill_->ok() || paired_) + 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())); + } } RecurrentPipeline::~RecurrentPipeline() = default; @@ -111,9 +122,15 @@ std::vector RecurrentPipeline::generate_from_ids(const std::vector(input_ids.size()) <= state_->max_length()) { + run_prefill_chunks(input_ids, logits); + } else { + for (std::size_t i = 0; i < input_ids.size(); ++i) + run_step(input_ids[i], logits, + nemotron_h_recurrent::prefill_step_needs_logits(i, input_ids.size())); + } auto t_prefill_end = SteadyClock::now(); // ── Decode phase ── @@ -278,6 +295,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, bool copy_logits) { auto t0 = SteadyClock::now(); diff --git a/families/nemotron_h/runtime/pipeline.h b/families/nemotron_h/runtime/pipeline.h index 27569ddda3..79bb0a02eb 100644 --- a/families/nemotron_h/runtime/pipeline.h +++ b/families/nemotron_h/runtime/pipeline.h @@ -39,7 +39,8 @@ class RecurrentPipeline final : public ITextGeneration { cudaStream_t stream, const char* name, std::shared_ptr tokenizer = nullptr, std::string model_id_str = "", - std::unique_ptr paired = nullptr); + std::unique_ptr paired = nullptr, + std::unique_ptr prefill = nullptr); ~RecurrentPipeline() override; TextResult generate(const std::string& prompt, const TextGenerationConfig& cfg = {}) override; @@ -61,6 +62,9 @@ class RecurrentPipeline final : public ITextGeneration { std::shared_ptr tokenizer_; std::string model_id_; std::unique_ptr paired_; + // Multi-row prefill engine (single-GPU split bundles only); decoding stays on decoder_. + std::unique_ptr prefill_; + int32_t prefill_rows_{0}; std::vector generate_from_ids(const std::vector& input_ids, int32_t max_new_tokens, @@ -72,6 +76,9 @@ class RecurrentPipeline final : public ITextGeneration { std::vector& logits, int32_t& decode_steps); void run_step(int32_t token_id, std::vector& logits, bool copy_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/nemotron_h/runtime/plugin.cpp b/families/nemotron_h/runtime/plugin.cpp index 3883c4351a..c3c2fccb45 100644 --- a/families/nemotron_h/runtime/plugin.cpp +++ b/families/nemotron_h/runtime/plugin.cpp @@ -50,6 +50,8 @@ struct RuntimeConfig { std::string tensor_parallel_mode; std::string precision; std::string kv_cache_quantization{"none"}; + // Single-GPU split bundles carry a multi-row prefill engine. + bool split{false}; }; template @@ -82,7 +84,7 @@ RuntimeConfig parse_runtime_config(const BundleReader& bundle) { throw std::runtime_error("nemotron_h runtime.json has invalid quantization"); } const std::string layout = require_value(json, "decoder_engine_layout"); - if (layout != "single" && layout != "dual_profile") + if (layout != "single" && layout != "dual_profile" && layout != "split") throw std::runtime_error("nemotron_h runtime.json has invalid decoder_engine_layout"); RuntimeConfig config{ @@ -109,6 +111,11 @@ RuntimeConfig parse_runtime_config(const BundleReader& bundle) { require_value(json, "tensor_parallel_mode"), require_value(json, "precision"), }; + // Tensor-parallel bundles record dual_profile; their per-rank engines keep + // one-token inputs. Split bundles are single-GPU. + config.split = layout == "split"; + if (config.split && config.tensor_parallel_size != 1) + throw std::runtime_error("nemotron_h split bundles are single-GPU"); if (json.contains("kv_cache_quantization")) { config.kv_cache_quantization = require_value(json, "kv_cache_quantization"); if (config.kv_cache_quantization != "none" && config.kv_cache_quantization != "fp8") @@ -195,6 +202,18 @@ ITask* create(const FamilyContext& context) { if (decoder == nullptr || !decoder->ok()) throw std::runtime_error("nemotron_h failed to load decoder"); decoder->set_timing_label("nemotron_h 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 == nullptr || !prefill->ok()) + throw std::runtime_error("nemotron_h failed to load the prefill engine"); + prefill->set_timing_label("nemotron_h 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; @@ -227,7 +246,7 @@ ITask* create(const FamilyContext& context) { config.precision, config.tensor_parallel_size); return new RecurrentPipeline(std::move(decoder), std::move(state), std::move(generation), stream, "Nemotron-H", std::move(tokenizer), std::string{}, - std::move(paired)); + std::move(paired), std::move(prefill)); } } // namespace trtmc::nemotron_h diff --git a/families/nemotron_h/ssd_rows.py b/families/nemotron_h/ssd_rows.py new file mode 100644 index 0000000000..66fec200fb --- /dev/null +++ b/families/nemotron_h/ssd_rows.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mamba-2 (SSD) mixer over ``S`` rows in one launch. + +The single-token decoder applies, per head ``h`` and token ``t``:: + + S_t = exp(dt_t A) S_{t-1} + dt_t x_t B_t^T (state [P, N]) + y_t = S_t C_t + D x_t + +With ``a_t = dt_t A`` and ``gamma_t = sum_{i<=t} a_i`` a chunk of ``S`` tokens +is, without any sequential step:: + + y_t = exp(gamma_t) S_0 C_t + sum_{i<=t} exp(gamma_t - gamma_i) dt_i (C_t . B_i) x_i + D x_t + S_S = exp(gamma_S) S_0 + sum_i exp(gamma_S - gamma_i) dt_i x_i B_i^T + +Every exponent is ``<= 0``, so the chunk form stays finite in FP32. +""" + +from __future__ import annotations + +import numpy as np +import tensorrt as trt + +# Static size of the causal masks; it bounds the prefill rows. +MAX_ROWS = 128 + + +def _const(network, values: np.ndarray) -> trt.ITensor: + values = np.ascontiguousarray(values) + return network.add_constant(values.shape, trt.Weights(values)).get_output(0) + + +def _i64(network, values) -> trt.ITensor: + return _const(network, np.asarray(values, dtype=np.int64)) + + +def _ew(network, a, b, op) -> trt.ITensor: + return network.add_elementwise(a, b, op).get_output(0) + + +def _mm(network, a, b, transpose_a=False, transpose_b=False) -> trt.ITensor: + op_a = trt.MatrixOperation.TRANSPOSE if transpose_a else trt.MatrixOperation.NONE + op_b = trt.MatrixOperation.TRANSPOSE if transpose_b else trt.MatrixOperation.NONE + return network.add_matrix_multiply(a, op_a, b, op_b).get_output(0) + + +def shuffle(network, tensor, reshape=None, transpose=None) -> trt.ITensor: + layer = network.add_shuffle(tensor) + if transpose is not None: + layer.first_transpose = transpose + if reshape is not None: + layer.reshape_dims = reshape + return layer.get_output(0) + + +def row_count(network, tensor: trt.ITensor) -> trt.ITensor: + """Length of axis 0 as an int64 ``[1]`` shape tensor.""" + shape = network.add_shape(tensor).get_output(0) + return network.add_slice(shape, (0,), (1,), (1,)).get_output(0) + + +def _dims(network, parts) -> trt.ITensor: + tensors = [p if isinstance(p, trt.ITensor) else _i64(network, [p]) for p in parts] + joined = network.add_concatenation(tensors) + joined.axis = 0 + return joined.get_output(0) + + +def dynamic_slice(network, tensor, start, size) -> trt.ITensor: + """Slice with static or shape-tensor start and size entries.""" + rank = len(start) + layer = network.add_slice(tensor, (0,) * rank, (1,) * rank, (1,) * rank) + if any(isinstance(s, trt.ITensor) for s in start): + layer.set_input(1, _dims(network, start)) + else: + layer.start = tuple(start) + layer.set_input(2, _dims(network, size)) + return layer.get_output(0) + + +def last_row(network, tensor: trt.ITensor, width: int) -> trt.ITensor: + """``[S, width]`` -> ``[1, width]``, the last row.""" + last = _ew( + network, row_count(network, tensor), _i64(network, [1]), trt.ElementWiseOperation.SUB + ) + return dynamic_slice(network, tensor, (last, 0), (1, width)) + + +def add_short_conv_rows( + network: trt.INetworkDefinition, + *, + columns: trt.ITensor, + conv_state: trt.ITensor, + weight: trt.ITensor, + channels: int, + kernel: int, +) -> tuple[trt.ITensor, trt.ITensor]: + """Causal depthwise conv over ``columns`` ``[C, S]`` with a ``[C, kernel]`` history. + + The history holds the last ``kernel`` inputs (newest last). Output column + ``t`` is ``sum_j weight[:, j] * window[:, t + j]`` over + ``window = [history[:, 1:], columns]``; the new history is the last + ``kernel`` columns of the window. Returns ``(outputs [C, S], history)``. + """ + rows = row_count(network, shuffle(network, columns, transpose=(1, 0))) + if kernel > 1: + past = network.add_slice(conv_state, (0, 1), (channels, kernel - 1), (1, 1)).get_output(0) + window = network.add_concatenation([past, columns]) + window.axis = 1 + window = window.get_output(0) + else: + window = columns + total = None + for tap in range(kernel): + tap_weight = network.add_slice(weight, (0, tap), (channels, 1), (1, 1)).get_output(0) + shifted = dynamic_slice(network, window, (0, tap), (channels, rows)) + term = _ew(network, shifted, tap_weight, trt.ElementWiseOperation.PROD) + total = term if total is None else _ew(network, total, term, trt.ElementWiseOperation.SUM) + last = _ew(network, rows, _i64(network, [1]), trt.ElementWiseOperation.SUB) + history = dynamic_slice(network, window, (0, last), (channels, kernel)) + return total, history + + +def add_chunked_ssd( + network: trt.INetworkDefinition, + *, + x: trt.ITensor, + b: trt.ITensor, + c: trt.ITensor, + dt: trt.ITensor, + a: np.ndarray, + d: np.ndarray, + state: trt.ITensor, + num_heads: int, +) -> tuple[trt.ITensor, trt.ITensor]: + """SSD over one chunk, all tensors FP32. + + ``x`` is ``[H, S, P]``, ``b``/``c`` are ``[H, S, N]``, ``dt`` is ``[H, S]`` + (after softplus) and ``state`` is ``[H, P, N]``; ``a`` and ``d`` are the + per-head ``A`` (negative) and ``D`` values. Returns ``(y [H, S, P], state)``. + """ + sum_op, prod_op = trt.ElementWiseOperation.SUM, trt.ElementWiseOperation.PROD + rows = row_count(network, shuffle(network, dt, transpose=(1, 0))) + log_decay = _ew( + network, dt, _const(network, np.asarray(a, np.float32).reshape(num_heads, 1)), prod_op + ) # [H, S] + axis = network.add_constant((), np.array(1, dtype=np.int32)).get_output(0) + gamma = network.add_cumulative( + log_decay, axis, trt.CumulativeOperation.SUM, False, False + ).get_output(0) + gamma_rows = shuffle(network, gamma, reshape=(0, 0, 1)) # [H, S, 1] + diff = _ew( + network, + gamma_rows, + shuffle(network, gamma, reshape=(0, 1, -1)), + trt.ElementWiseOperation.SUB, + ) + diff = _ew( + network, + diff, + _const(network, np.zeros((1, 1, 1), np.float32)), + trt.ElementWiseOperation.MIN, + ) + causal = dynamic_slice( + network, + _const(network, np.tril(np.ones((MAX_ROWS, MAX_ROWS), np.float32))[None]), + (0, 0, 0), + (1, rows, rows), + ) + decay = _ew( + network, network.add_unary(diff, trt.UnaryOperation.EXP).get_output(0), causal, prod_op + ) + scores = _ew( + network, + _ew(network, decay, _mm(network, c, b, transpose_b=True), prod_op), + shuffle(network, dt, reshape=(0, 1, -1)), + prod_op, + ) # [H, t, i] + skip = _const(network, np.asarray(d, np.float32).reshape(num_heads, 1, 1)) + y = _ew( + network, + _ew( + network, + _ew( + network, + network.add_unary(gamma_rows, trt.UnaryOperation.EXP).get_output(0), + _mm(network, c, state, transpose_b=True), + prod_op, + ), + _mm(network, scores, x), + sum_op, + ), + _ew(network, skip, x, prod_op), + sum_op, + ) + + total = network.add_reduce( + log_decay, trt.ReduceOperation.SUM, 1 << 1, keep_dims=True + ).get_output(0) + weights = _ew( + network, + network.add_unary( + _ew(network, total, gamma, trt.ElementWiseOperation.SUB), trt.UnaryOperation.EXP + ).get_output(0), + dt, + prod_op, + ) # exp(gamma_S - gamma_i) dt_i: [H, S] + 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, + x, + _ew(network, b, shuffle(network, weights, reshape=(0, 0, 1)), prod_op), + transpose_a=True, + ), + sum_op, + ) + return y, final_state diff --git a/families/nemotron_h/tests/test_chunked_prefill.py b/families/nemotron_h/tests/test_chunked_prefill.py new file mode 100644 index 0000000000..e3f60b738b --- /dev/null +++ b/families/nemotron_h/tests/test_chunked_prefill.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Nemotron-H 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 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, fmt=None) + config = ModelConfig.from_dir(tmp_path) + instance = model._NemotronHModel() + + def build(**options): + weights = instance.load_weights(str(tmp_path), config) + return _engine( + instance.build_engine(config, weights, MAX_CACHE, precision="fp32", **options) + ) + + single = build() + prefill_engine = build(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_paired_draft_features(tmp_path: Path) -> None: + """Engines that export paired-draft features keep one row per launch.""" + _write_checkpoint(tmp_path, fmt=None) + config = ModelConfig.from_dir(tmp_path) + instance = model._NemotronHModel() + weights = instance.load_weights(str(tmp_path), config) + with pytest.raises(NotImplementedError, match="prefill engines"): + instance.build_engine( + config, weights, MAX_CACHE, prefill_rows=16, target_feature_layers=(0,) + ) + with pytest.raises(ValueError, match="prefill_rows"): + instance.build_engine(config, weights, MAX_CACHE, prefill_rows=model.PREFILL_MAX_TOKENS + 1) From 2f2cc0617a66a99bd7c1aad423dc4a6502d284b0 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Tue, 10 Nov 2026 19:46:44 +0000 Subject: [PATCH 2/2] fix(nemotron_h): take the chunk decay total from the cumulative sum A separate reduce-sum of the log decays can exceed the last cumulative sum by many ulps at large |dt A|, so exp(gamma_S - gamma_i) overflowed and the recurrent state went non-finite on some prompts and tactics. Read gamma_S from the cumsum and clamp the exponent at zero. Signed-off-by: chaofengw --- families/nemotron_h/ssd_rows.py | 20 ++++-- .../nemotron_h/tests/test_chunked_prefill.py | 61 +++++++++++++++++++ 2 files changed, 75 insertions(+), 6 deletions(-) diff --git a/families/nemotron_h/ssd_rows.py b/families/nemotron_h/ssd_rows.py index 66fec200fb..55c4051afe 100644 --- a/families/nemotron_h/ssd_rows.py +++ b/families/nemotron_h/ssd_rows.py @@ -195,14 +195,22 @@ def add_chunked_ssd( sum_op, ) - total = network.add_reduce( - log_decay, 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 |dt A|, and + # exp(gamma_S - gamma_i) then overflows into a non-finite state. 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, (0, last), (num_heads, 1)) # [H, 1] + tail = _ew( + network, + _ew(network, total, gamma, trt.ElementWiseOperation.SUB), + _const(network, np.zeros((1, 1), np.float32)), + trt.ElementWiseOperation.MIN, + ) weights = _ew( network, - network.add_unary( - _ew(network, total, gamma, trt.ElementWiseOperation.SUB), trt.UnaryOperation.EXP - ).get_output(0), + network.add_unary(tail, trt.UnaryOperation.EXP).get_output(0), dt, prod_op, ) # exp(gamma_S - gamma_i) dt_i: [H, S] diff --git a/families/nemotron_h/tests/test_chunked_prefill.py b/families/nemotron_h/tests/test_chunked_prefill.py index e3f60b738b..8f9495db46 100644 --- a/families/nemotron_h/tests/test_chunked_prefill.py +++ b/families/nemotron_h/tests/test_chunked_prefill.py @@ -166,3 +166,64 @@ def test_prefill_engine_rejects_paired_draft_features(tmp_path: Path) -> None: ) 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_chunked_ssd_state_stays_finite_at_large_decays() -> None: + """At |dt A| in the millions the chunk state matches the recurrence instead of overflowing.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + from .. import ssd_rows + + heads, rows, head_dim, state_dim = 2, ssd_rows.MAX_ROWS, 4, 3 + rng = np.random.default_rng(3) + x = rng.standard_normal((heads, rows, head_dim)).astype(np.float32) + b = rng.standard_normal((heads, rows, state_dim)).astype(np.float32) + c = rng.standard_normal((heads, rows, state_dim)).astype(np.float32) + dt = rng.uniform(1.0, 30.0, (heads, rows)).astype(np.float32) + a = np.array([-0.5, -5.0e4], np.float32) # the second head decays by ~1e6 per row + d = np.ones(heads, np.float32) + state = rng.standard_normal((heads, head_dim, state_dim)).astype(np.float32) + + builder = trt.Builder(trt.Logger(trt.Logger.WARNING)) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + inputs = { + name: network.add_input(name, trt.float32, value.shape) + for name, value in (("x", x), ("b", b), ("c", c), ("dt", dt), ("state", state)) + } + y, final = ssd_rows.add_chunked_ssd( + network, x=inputs["x"], b=inputs["b"], c=inputs["c"], dt=inputs["dt"], a=a, d=d, + state=inputs["state"], num_heads=heads, + ) + # gamma_S is read from the cumulative sum: a separate reduction sums in another order + # and can exceed gamma_i, which overflows exp(gamma_S - gamma_i) on some tactics. + layers = [network.get_layer(i).type for i in range(network.num_layers)] + assert trt.LayerType.REDUCE not in layers + y.name, final.name = "y", "final" + network.mark_output(y) + network.mark_output(final) + config = builder.create_builder_config() + config.clear_flag(trt.BuilderFlag.TF32) + engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine( + builder.build_serialized_network(network, config) + ) + context = engine.create_execution_context() + device = {name: torch.from_numpy(value).cuda() for name, value in + (("x", x), ("b", b), ("c", c), ("dt", dt), ("state", state))} + out = { + name: torch.empty(tuple(engine.get_tensor_shape(name)), device="cuda") + for name in ("y", "final") + } + for name, value in {**device, **out}.items(): + context.set_tensor_address(name, value.data_ptr()) + assert context.execute_async_v3(torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize() + + want = state.astype(np.float64) + for t in range(rows): # the one-token recurrence + decay = np.exp(dt[:, t] * a)[:, None, None] + want = decay * want + (dt[:, t, None, None] * x[:, t, :, None]) * b[:, t, None, :] + got = out["final"].cpu().numpy() + assert np.isfinite(got).all() + np.testing.assert_allclose(got, want, rtol=1e-3, atol=1e-3)