diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md index 686829541385..c8e3483a9d35 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/README.md +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/README.md @@ -379,3 +379,11 @@ That boundary may move as models need more customization. Model-specific runtime policy can be split from shared services, but a change to common runtime behavior still has to work for every other model. How to specialize the runtime further is still being designed. + +### Beam diagnostic validation + +For Beam, `TRTLLM_MODELING_V2_VALIDATE=1` additionally enables repeated per-layer +metadata checks and per-step, all-rank routing agreement using device-to-host copies +and SHA-256 hashes. The serving adapter also enables checkpoint digests in this mode. +This is expensive diagnostic behavior, not only a first-forward contract check. +Use the same setting on every rank and leave it disabled for performance measurements. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 1d714f7dd885..309b0ccb0d1a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -80,15 +80,19 @@ def registered_routes() -> Dict[str, Callable[..., type]]: #: construction, and a check that can only ever hold is a check whose failure #: mode is its own bug taking down a working deployment. #: -#: An environment variable rather than an LLM arg while modeling_v2 is -#: experimental; it has to reach the ranks, so a test passes it through +#: Beam additionally enables checkpoint digests, per-layer metadata rechecks and per-step +#: routing agreement (device-to-host copies, hashing and collectives). Keep it off for +#: performance measurements; all Beam ranks must use the same setting. +#: +#: An environment variable while modeling_v2 is experimental. It must reach all ranks, +#: so a test passes it through #: ``env_overrides``. When this stops being experimental it belongs in the #: LLM args, where a caller can set it per engine rather than per process. MODELING_V2_VALIDATE_ENV = "TRTLLM_MODELING_V2_VALIDATE" def step_contract_enabled() -> bool: - """Whether targets should check their step contract on the first forward. + """Whether contract diagnostics are enabled (Beam also checks every step). Read once per model instance at construction rather than per forward: the answer cannot change within a process, and the forward is the hot path. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py index 3e67f8d6690b..a7ac57303355 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/flashinfer_silu_and_mul.py @@ -21,9 +21,8 @@ class FlashinferSiluAndMul(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # deepseek-r1 is the only shipped caller. It reaches this op from two - # widths -- the routed/shared expert intermediate and the dense layers' - # -- at decode and prefill token counts. + # Exercise decode and prefill across narrow and wide intermediates. + # Model-specific shapes live in the test driver. CELLS: tuple[Cell, ...] = ( Cell( why="decode-shaped, 2 x moe_intermediate (2048)", diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py index 0e65c6823a2c..a04be7d6f938 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/gemm/cublas_mm.py @@ -20,8 +20,8 @@ class CublasMm(OpWrapper): ARCHS = frozenset({Arch.SM_100, Arch.SM_103}) - # Both targets use this for every dense projection. This entry's test - # carries their three argument shapes that differ in kind rather than in size: fused + # Targets use this for dense projections. This entry's test + # carries argument shapes that differ in kind rather than in size: fused # bias (gpt-oss carries one on qkv, o and the router), no bias (deepseek's # projections), and an out_dtype that is not the input's (deepseek's router # asks for fp32 logits). @@ -34,7 +34,7 @@ class CublasMm(OpWrapper): note = """ `output_buffer_kind` and `group` are the op's multi-stream and collective - arguments. Both shipped targets leave them at their defaults, so neither is + arguments. The covered targets leave them at their defaults, so neither is a cell axis and neither is certified here. """ diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py new file mode 100644 index 000000000000..51c47b7ad054 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/catalog/residual/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Residual entries.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md new file mode 100644 index 000000000000..9c04ef9a75df --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/README.md @@ -0,0 +1,36 @@ + + + +# Beam StairCase components + +This directory implements Beam components using the experimental ModelingV2/StairCase +architecture: family-scoped operators and a 52-layer NVFP4-expert component for attention TP4 and MoE TP4/EP1 on four GB300 GPUs. + +The full component implements sliding-window and full attention with a BF16 KV cache. +Its loader streams caller-provided weights into rank-local storage; no checkpoint is +included. Router projection and selection use FP32, and packed NVFP4 experts use the +native GEMM path. + +These are directly constructed components, not a registered `LLM()` or `trtllm-serve` +target. Engine routing, the engine adapter, checkpoint-loader registration and a named +whole-model accuracy gate belong to the serving integration. Operator certification and +CPU component-contract tests accompany these components; broader component tests are +reviewed separately. Speculative decoding, FP8 KV cache, disaggregated serving, CUDA +graphs and overlap scheduling are outside this component baseline. + +Beam-only operators live under `ops/`, with tests in the matching Beam test tree. +Component-specific cases for shared catalog operators reuse the catalog test drivers, +including TP1 shapes and the per-rank TP4 projection, attention and expert shapes. + +## Checkout-only bring-up helpers + +The six TP1 bring-up components and their generation controller live with their tests in +`tests/unittest/_torch/modeling_v2/models/beam/_components/`. They are not installed in +the TensorRT-LLM wheel. The runtime does not expose intermediate-tensor capture hooks. +Normal runtime validation and rank-consistency diagnostics remain available. + +The TP4 checkpoint layout, expert provider and full-model execution remain in this +library. Public serving integration uses `LLM()` and `trtllm-serve`. + +Shared BF16 TP4 layout, loading and component base modules live in `_tp4/`; the +`full_model_nvfp4__sm_103__tp4/` package adds NVFP4 expert storage and execution. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py new file mode 100644 index 000000000000..890bdb17925a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BeamForCausalLM components.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py new file mode 100644 index 000000000000..579fa8b62c94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_attention.py @@ -0,0 +1,453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The attention sublayer of the Beam decoder block components, one request step. + +Shared by the dense and sparse block components. It holds no parameters, cache or +configuration: each block passes its own tensors and geometry, its layer's position +encoding and window, and the prepared attention metadata of the call. The computation, +before the block's post-attention norm:: + + q, k, v, g = h @ q_proj^T, h @ k_proj^T, h @ v_proj^T, h @ gate_proj^T + q, k = rmsnorm(q, q_norm), rmsnorm(k, k_norm) per head, rounded to bf16 + sliding layer: q, k = rope(q, k) fp32 NeoX rotation, rounded + full layer: q = q * exp(0.1 * log1p((p + 1) / 2048)) fp32 scale, rounded; k as is + a = causal attention(q, k, v), scale 1/sqrt(128), 6 q heads per kv head, + window 512 on a sliding layer, the whole context on a full one + a = a * sigmoid(g + gate_bias) per head, bf16 steps + out = a @ o_proj^T + +A full-attention layer has no position embedding: neither q nor k is rotated, and its +only position dependence is the per-token query scale above. + +A step of one request: a context call (``num_contexts == 1``) appends its T new tokens +after the ``past`` tokens the cache already holds for the request -- none for a fresh +context, a reused or earlier-computed prefix otherwise -- and a generation call +(``num_contexts == 0``) appends exactly one token. Every position is absolute: the new +tokens sit at ``past .. past + T - 1``, where ``past + T`` is the request's KV length in +the metadata, and the positions the block receives must be exactly those. A context call +over a cached prefix must use paged-context FMHA, the path that reads the prefix's keys and +values from the pool; without it the context kernel would attend to the new rows only. +The keys and values the op appends are the transformed k (normalized, rotated on a sliding +layer) and the projected v of the new tokens, written to the request's own pages. + +A step of several requests (``check_batch_metadata``) is the native runtime's batch: its +context requests first, then its generation requests, each with its own ``past``, KV length +and pages, their new tokens concatenated in that order. Every operation but the attention +op is per token; the op reads each request's own metadata row. + +Logical layer and cache slot. A block's ``layer_idx`` is its depth in the model: it fixes +the block's semantics (its residual coefficients, its attention and feed-forward kind). +The KV cache it appends to is a ``KVCacheManagerV2`` that holds that layer, so its storage +slot is the manager's local index for it -- ``layer_offsets[layer_idx]``: 0 for the one +layer a ``layer_mask`` selects, the layer's own position in a manager holding several. +The attention op and the manager's per-layer arrays take that local index; manager APIs +such as ``get_buffers`` take the logical one and translate it. The slot is resolved from +each call's metadata, never stored, and a manager that does not hold the block's logical +layer is rejected before any kernel runs. + +Host synchronization: the one-request position check (``check_step``) compares the +positions with ``arange(past, past + T)`` on the device and reads the answer back, once per +call. The metadata checks read host tensors only. +""" + +from __future__ import annotations + +from typing import Mapping + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + STANDARD, + thop_attention, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_neox_rope import ( + beam_neox_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_nope_query_scale import ( # noqa: E501 + beam_nope_query_scale, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_output_gate import ( + beam_output_gate, +) +from tensorrt_llm._torch.attention.backends.fmha.phased import get_spec_decoding_position_offsets +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm.bindings import DataType + +from ._validation import check_positions, metadata_values + +#: thop_attention arguments that do not change with the step on these blocks: packed q|k|v, +#: append-and-attend, causal, bf16 pool, no quantization, and any position transform +#: applied before the call -- the rope_* values are inert placeholders the op requires, +#: not this model's RoPE. +CALL_CONSTANTS = dict( + k=None, + v=None, + output_sf=None, + is_fused_qkv=True, + update_kv_cache=True, + predicted_tokens_per_seq=1, + attention_input_type=STANDARD, + is_mla_enable=False, + mask_type=1, + q_scaling=1.0, + quant_mode=0, + kv_scale_orig_quant=None, + kv_scale_quant_orig=None, + attention_sinks=None, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + rope_dim=0, + rope_base=10000.0, + rope_scale_type=0, + rope_scale=1.0, + rope_short_m_scale=1.0, + rope_long_m_scale=1.0, + rope_max_positions=1024, + rope_original_max_positions=1024, + out_scale=None, + latent_cache=None, + q_pe=None, + q_lora_rank=None, + kv_lora_rank=None, + qk_nope_head_dim=None, + qk_rope_head_dim=None, + v_head_dim=None, + rope_append=None, + chunked_prefill_buffer_batch_size=1, + attention_chunk_size=None, + softmax_stats_tensor=None, + sparse_kv_indices=None, + sparse_kv_offsets=None, + sparse_attn_indices=None, + sparse_attn_offsets=None, + sparse_attn_indices_block_size=0, + mrope_rotary_cos_sin=None, + mrope_position_deltas=None, + helix_position_offsets=None, + helix_is_inactive_rank=None, + cross_kv=None, + relative_attention_bias=None, + relative_attention_max_distance=0, + kv_norm_weight=None, + kv_norm_eps=1e-6, + skip_correction_threshold=0.0, +) + + +def local_layer_index(md: TrtllmAttentionMetadata, layer_idx: int) -> int: + """The cache slot of logical layer ``layer_idx`` in the call's cache manager. + + Raises ValueError unless the manager holds ``layer_idx`` at an integer local index + within its local layers. Reads the manager's host-side layer map only. + """ + mgr = md.kv_cache_manager + offsets = getattr(mgr, "layer_offsets", None) + if not isinstance(offsets, Mapping) or layer_idx not in offsets: + held = sorted(offsets) if isinstance(offsets, Mapping) else None + raise ValueError( + f"the KV cache manager does not hold logical layer {layer_idx}; it holds {held}. " + f"Select the layer with a layer_mask rather than reusing another layer's cache" + ) + num_local = getattr(mgr, "num_local_layers", None) + local = offsets[layer_idx] + if ( + type(num_local) is not int + or type(local) is not int + or len(offsets) != num_local + or not 0 <= local < num_local + ): + raise ValueError( + f"the KV cache manager maps logical layer {layer_idx} to slot {local!r} of " + f"{num_local!r} local layers holding {len(offsets)} logical layers" + ) + return local + + +def check_step( + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + md: TrtllmAttentionMetadata, + *, + layer_idx: int, + hidden_size: int, + num_kv_heads: int, + head_dim: int, + device: torch.device, + max_tokens: int, +) -> None: + """Raise ValueError on any call input a block would otherwise compute wrongly. + + Accepted: ``hidden_states`` contiguous ``[1, T, hidden_size]`` bf16 on ``device``, + 1 <= T <= ``max_tokens``; integer ``position_ids`` ``[1, T]`` equal to + ``arange(past, past + T)`` on that device; a prepared ``TrtllmAttentionMetadata`` + holding one step of one request on a BF16 ``KVCacheManagerV2`` that holds logical + layer ``layer_idx`` with ``num_kv_heads`` kv heads of ``head_dim`` (``check_metadata`` + states the steps and returns ``past``). + """ + if not isinstance(hidden_states, torch.Tensor) or hidden_states.dtype != torch.bfloat16: + raise ValueError("hidden_states must be a torch.bfloat16 tensor") + if ( + hidden_states.dim() != 3 + or hidden_states.shape[0] != 1 + or hidden_states.shape[2] != hidden_size + ): + raise ValueError( + f"hidden_states must be [1, T, {hidden_size}], got {tuple(hidden_states.shape)}" + ) + if hidden_states.device != device or not hidden_states.is_contiguous(): + raise ValueError("hidden_states must be contiguous on the block's device") + num_tokens = hidden_states.shape[1] + if not 1 <= num_tokens <= max_tokens: + raise ValueError(f"supported token counts per call are 1..{max_tokens}, got {num_tokens}") + check_positions(position_ids, (1, num_tokens), device) + past = check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + expected = torch.arange( + past, past + num_tokens, dtype=position_ids.dtype, device=position_ids.device + ) + if not torch.equal(position_ids[0], expected): + raise ValueError( + f"the call appends tokens {past}..{past + num_tokens - 1} after {past} cached " + "tokens; its positions must be exactly those" + ) + + +def check_metadata( + md: TrtllmAttentionMetadata, + num_tokens: int, + *, + layer_idx: int, + num_kv_heads: int, + head_dim: int, +) -> int: + """Raise ValueError unless ``md`` is one prepared step of one request; return ``past``. + + ``check_batch_metadata`` restricted to a batch of exactly one request: a context step of + ``num_tokens`` new tokens after ``past`` cached ones, or a generation step of one token. + ``past`` is the absolute position of the step's first new token. Reads host-side state + only. + """ + spans = check_batch_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=num_kv_heads, head_dim=head_dim + ) + if len(spans) != 1: + raise ValueError(f"one request per call is supported; metadata has {len(spans)}") + return spans[0][0] + + +def check_batch_metadata( + md: TrtllmAttentionMetadata, + num_tokens: int, + *, + layer_idx: int, + num_kv_heads: int, + head_dim: int, +) -> list: + """Raise ValueError unless ``md`` is one prepared step of its requests; return their spans. + + The native runtime orders a step's ``num_seqs`` requests with its ``num_contexts`` + context requests first, and their new tokens fill the call's ``num_tokens`` rows in that + order. Accepted on a BF16 ``KVCacheManagerV2`` holding logical layer ``layer_idx`` with + ``num_kv_heads`` kv heads of ``head_dim``, beam width 1, no cross attention or + speculative decoding, where each request ``i`` with KV length ``L_i`` (at most the + manager's ``max_seq_len``) is + + * a context request of ``n_i >= 1`` new tokens: context length ``n_i``, + ``past_i = L_i - n_i >= 0`` tokens already cached, and paged-context FMHA whenever + ``past_i > 0``; or + * a generation request: one new token after ``past_i = L_i - 1 >= 0`` cached tokens + (the native runtime's warmup issues generation steps over a cache that holds only the + new token). + + Returns ``[(past_i, n_i), ...]`` in row order: request ``i``'s new tokens sit at the + absolute positions ``past_i .. past_i + n_i - 1``. Reads host-side state only. + """ + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError(f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}") + mgr = md.kv_cache_manager + if mgr is None: + raise ValueError("attn_metadata needs a KV cache manager for the paged append") + if md.is_cross or md.is_spec_decoding_enabled or md.beam_width != 1: + raise ValueError("cross attention, speculative decoding and beams are not supported") + num_seqs, num_contexts = md.num_seqs, md.num_contexts + if num_seqs < 1 or not 0 <= num_contexts <= num_seqs: + raise ValueError( + f"the metadata holds {num_seqs} requests of which {num_contexts} are contexts" + ) + if md.kv_cache_params is None or md.kv_cache_params.num_extra_kv_tokens: + raise ValueError("the metadata must use the KV cache without extra KV tokens") + if num_tokens > md.max_num_tokens: + raise ValueError( + f"the call has {num_tokens} tokens; the metadata holds {md.max_num_tokens} per call" + ) + new_tokens = metadata_values(md.seq_lens, num_seqs, "seq_lens") + kv_lens = metadata_values(md.kv_lens_runtime, num_seqs, "kv_lens_runtime") + request_types = metadata_values( + md.host_request_types_runtime, num_seqs, "host_request_types_runtime" + ) + context_lens = metadata_values( + md.prompt_lens_cpu_runtime, num_contexts, "prompt_lens_cpu_runtime" + ) + if sum(new_tokens) != num_tokens or md.num_ctx_tokens != sum(new_tokens[:num_contexts]): + raise ValueError( + f"the requests' new tokens {new_tokens} ({md.num_ctx_tokens} in contexts) do not " + f"fill the call's {num_tokens} tokens" + ) + spans = [] + for i, (n, kv_len, request_type) in enumerate(zip(new_tokens, kv_lens, request_types)): + past = kv_len - n + if i < num_contexts: + if n < 1 or context_lens[i] != n or request_type != 0: + raise ValueError( + f"context request {i} of {n} new tokens needs context length {n}; " + f"metadata has context length {context_lens[i]} and request type " + f"{request_type}" + ) + if past < 0: + raise ValueError( + f"request {i}'s KV length {kv_len} is shorter than its {n} new tokens" + ) + if past and not md.use_paged_context_fmha: + raise ValueError( + f"a context step after {past} cached tokens must read them from the pool: " + "paged-context FMHA is required" + ) + else: + if n != 1 or request_type != 1: + raise ValueError( + f"generation request {i} appends exactly one token, got {n} (request type " + f"{request_type})" + ) + if past < 0: + raise ValueError( + f"request {i}'s KV length {kv_len} is shorter than its generation token" + ) + spans.append((past, n)) + if max(kv_lens) > mgr.max_seq_len or md.max_seq_len != mgr.max_seq_len: + raise ValueError( + f"a request reaches {max(kv_lens)} tokens; the cache holds {mgr.max_seq_len} per " + f"sequence and the metadata states {md.max_seq_len}" + ) + local = local_layer_index(md, layer_idx) + heads = getattr(mgr, "num_kv_heads_per_layer", None) or [] + local_heads = heads[local] if len(heads) == mgr.num_local_layers else None + if mgr.dtype != DataType.BF16 or local_heads != num_kv_heads or mgr.head_dim != head_dim: + raise ValueError( + f"the KV cache must be BF16 with {num_kv_heads} kv heads of {head_dim}; " + f"got {mgr.dtype}, {heads}, head_dim {mgr.head_dim}" + ) + return spans + + +def step_args(md: TrtllmAttentionMetadata, layer_idx: int) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state. + + Includes ``local_layer_idx``, the cache slot of logical layer ``layer_idx`` resolved + from this call's manager. + """ + return dict( + local_layer_idx=local_layer_index(md, layer_idx), + sequence_length=md.kv_lens_cuda_runtime, + host_past_key_value_lengths=md.kv_lens_runtime, + host_total_kv_lens=md.host_total_kv_lens, + context_lengths=md.prompt_lens_cuda_runtime, + host_context_lengths=md.prompt_lens_cpu_runtime, + host_request_types=md.host_request_types_runtime, + kv_cache_block_offsets=md.kv_cache_block_offsets, + host_kv_cache_pool_pointers=md.host_kv_cache_pool_pointers, + host_kv_cache_pool_mapping=md.host_kv_cache_pool_mapping, + workspace=md.effective_workspace, + tokens_per_block=md.tokens_per_block, + max_num_requests=md.max_num_requests, + max_context_length=md.max_context_length, + max_seq_len=md.max_seq_len, + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, + use_paged_context_fmha=md.use_paged_context_fmha, + beam_width=md.effective_beam_width, + cache_indirection=md.cache_indirection, + block_ids_per_seq=md.block_ids_per_seq, + max_context_q_len_override=md.max_context_q_len_override, + is_cross=md.is_cross, + is_spec_decoding_enabled=md.is_spec_decoding_enabled, + use_spec_decoding=md.use_spec_decoding, + is_spec_dec_tree=md.is_spec_dec_tree, + spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, + spec_decoding_position_offsets=get_spec_decoding_position_offsets(md), + spec_decoding_packed_mask=md.spec_decoding_packed_mask, + spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, + spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, + spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + num_sparse_topk=md.num_sparse_topk, + flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, + flash_mla_num_splits=md.flash_mla_num_splits, + max_num_sequences=md.max_num_sequences, + force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + ) + + +def attention( + h: torch.Tensor, + positions: torch.Tensor, + step: dict, + *, + q_proj: torch.Tensor, + k_proj: torch.Tensor, + v_proj: torch.Tensor, + gate_proj: torch.Tensor, + o_proj: torch.Tensor, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + gate_bias: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + eps: float, + rope_theta: float | None, + window: int, +) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm. + + ``step`` is ``step_args`` of the call's metadata; the projections are ``[out, in]``. + ``rope_theta`` is the sliding layer's RoPE base; ``None`` is a full-attention layer + without position embedding, whose q takes the per-token query scale instead. ``window`` + is the attention window the op applies: the sliding window, or the metadata's + ``max_seq_len`` for a full layer. + """ + T = h.shape[0] + heads, kv_heads, d = num_heads, num_kv_heads, head_dim + q = cublas_mm(h, q_proj.t()) + k = cublas_mm(h, k_proj.t()) + v = cublas_mm(h, v_proj.t()) + g = cublas_mm(h, gate_proj.t()) + qn = flashinfer_rmsnorm(q.view(T * heads, d), q_norm, eps) + kn = flashinfer_rmsnorm(k.view(T * kv_heads, d), k_norm, eps) + qn, kn = qn.view(T, heads, d), kn.view(T, kv_heads, d) + sliding = rope_theta is not None + if sliding: + qa, ka = beam_neox_rope(qn, kn, positions, rope_theta, 1.0) + else: + qa, ka = beam_nope_query_scale(qn, positions), kn + qkv = torch.cat((qa.reshape(T, heads * d), ka.reshape(T, kv_heads * d), v), dim=-1) + core = torch.empty(T, heads * d, dtype=h.dtype, device=h.device) + core = thop_attention( + q=qkv, + output=core, + num_heads=heads, + num_kv_heads=kv_heads, + head_size=d, + attention_window_size=window, + **step, + **CALL_CONSTANTS, + ) + gated = beam_output_gate(core.view(T, heads, d), g.view(T, heads, d), gate_bias) + part = cublas_mm(gated.reshape(T, heads * d), o_proj.t()) + return part diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_constants.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_constants.py new file mode 100644 index 000000000000..e2a79548d92c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_constants.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Supported Beam model configuration and numerical constants shared across components.""" + +ROUTER_EPS = 1e-20 + +SM = (10, 3) + +WORLD_SIZE = 4 + +#: The configuration this component computes, as the checkpoint's config.json states it. +SUPPORTED_CONFIG = { + "model_type": "beam", + "hidden_size": 5120, + "intermediate_size": 20480, + "moe_intermediate_size": 2560, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "num_experts": 256, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "norm_topk_prob": True, + "vocab_size": 151936, + "num_hidden_layers": 52, + "rms_norm_eps": 1e-06, + "routed_scaling_factor": 2.5, + "final_logit_softcapping": 30.0, + "sliding_window": 512, + "hidden_act": "silu", + "tie_word_embeddings": False, + "attention_bias": False, +} +SLIDING_ROPE_THETA = 10000.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_load.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_load.py new file mode 100644 index 000000000000..47ecf617c829 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_load.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collect rank-local checkpoint failures before any rank leaves a load phase.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from tensorrt_llm.mapping import Mapping + + +def load_failure_details(mapping: Mapping, ok: bool, error: BaseException | None) -> str | None: + """Return the same failing-rank summary on every rank, or None after complete success.""" + outcome = (ok, repr(error) if error is not None else "incomplete tensor coverage") + outcomes = [outcome] + if mapping.world_size > 1: + from tensorrt_llm._torch.distributed.communicator import Distributed + + outcomes = Distributed.get(mapping).allgather(outcome) + failures = [ + f"rank {rank}: {cause}" for rank, (passed, cause) in enumerate(outcomes) if not passed + ] + return "; ".join(failures) or None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/__init__.py new file mode 100644 index 000000000000..fd1a4e95309b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared BF16 TP4 layout, checkpoint reader and component base for Beam.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/component_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/component_bf16.py new file mode 100644 index 000000000000..a87385287350 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/component_bf16.py @@ -0,0 +1,610 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam BF16 checkpoint-layout component / sm_103 / tp4: rank-local real weights and blocks. + +Owned by the routed NVFP4 target (this directory): the NVFP4 component (``component``) builds +on it and the full model (``model``) on that. Constructed on its own it is a component, not a +target: a caller on each of four ranks constructs this class with the checkpoint's +configuration and its rank's ``Mapping`` (attention TP4, MoE TP4/EP1, PP1), loads a +caller-supplied checkpoint directory into the rank's slices, and runs representative +computations. It does not run a full-depth forward, prefill/decode loop or sampler. + +Storage is rank-local and final: each logical tensor's slice (``layout_bf16``) lives in its own +parameter, except that each sparse layer keeps its 256 experts' slices in three contiguous +banks. Loading streams bounded tiles straight into that storage (``loader``) and is a +transaction: the component is usable only after every rank loaded every selected tensor. + +Per call, rank r computes its 12 query heads / 2 kv heads, its dense or expert intermediate +slice and its vocabulary rows; the enclosing method applies exactly one collective at each +boundary, on the TRT-LLM NCCL all-reduce or all-gather:: + + embed masked local lookup, one all-reduce (sum) + attention local q/k/v/gate, q/k norms, RoPE or NoPE scale, TRTLLM attention on the + rank's kv heads, output gate, local o_proj partial, one all-reduce + dense FFN local gate/up, fp32 SwiGLU, local down partial, one all-reduce + MoE replicated fp32 router (global expert ids), in diagnostic mode an + all-rank agreement on the router input, ids and weights, rank-local + expert slices combined in ascending global expert order (bf16 routes, + fp32 sum, x2.5, bf16), one all-reduce + head final norm, local vocabulary logits, one all-gather along the + vocabulary in rank order, fp32 softcap once + +""" + +from __future__ import annotations + +import hashlib +import json +import os +from collections import OrderedDict +from typing import Any, Mapping, Optional, Sequence + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + MODELING_V2_VALIDATE_ENV, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.beam_logit_softcap import ( + beam_logit_softcap, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.swiglu_fp32 import ( + swiglu_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.moe.beam_grouped_topk_router import ( + beam_grouped_topk_router, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.norm.beam_rmsnorm_fp32 import ( + beam_rmsnorm_fp32, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.residual.beam_depth_residual import ( + beam_depth_residual, + depth_scales, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.distributed import AllReduce, allgather +from tensorrt_llm.functional import AllReduceStrategy +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from .. import _attention, dispatch_combine +from .._constants import ROUTER_EPS +from .._constants import SM as _SM +from .._load import load_failure_details +from .._validation import check_positions, explicit_device +from .layout_bf16 import ( + EMBED, + FINAL_NORM, + HEAD, + WORLD_SIZE, + TensorSpec, + check_config, + check_partition, + check_rank, + tensor_specs, +) +from .loader_bf16 import CheckpointFormatError, SafetensorsCatalog, TileReader, read_slice_into + +#: Default tokens per call; the representative cases use at most 129. A caller that runs longer +#: calls (a full-model prefill) passes its own ``max_tokens`` at construction. +MAX_TOKENS = 1024 + + +class LoadError(RuntimeError): + """A load failed on at least one rank; the component is not usable on any rank.""" + + +def check_max_tokens(max_tokens: int) -> int: + """Return ``max_tokens`` if it is a positive int (the most tokens one call accepts).""" + if type(max_tokens) is not int or max_tokens < 1: + raise ValueError(f"max_tokens must be a positive int, got {max_tokens!r}") + return max_tokens + + +def _check_mapping(mapping: ParallelMapping) -> None: + check_rank(mapping.rank, mapping.world_size) + expected = dict(tp_size=4, pp_size=1, moe_tp_size=4, moe_ep_size=1, cp_size=1) + for field, value in expected.items(): + actual = getattr(mapping, field) + if actual != value: + raise ValueError(f"this layout needs Mapping {field}={value}, got {actual}") + if mapping.enable_attention_dp: + raise ValueError("this layout does not use attention data parallelism") + if list(mapping.tp_group) != list(range(WORLD_SIZE)): + raise ValueError(f"the TP group must be ranks 0..3, got {mapping.tp_group}") + + +class BeamCheckpointLayoutBf16Sm103Tp4(nn.Module): + """Rank-local real-checkpoint weights of the Beam model on one of four TP ranks. + + ``config`` is the checkpoint's config.json (a mapping); ``mapping`` this rank's + ``tensorrt_llm.mapping.Mapping`` with world size 4, TP4, PP1, MoE TP4/EP1; ``layers`` + the logical layers to hold (default all 52; the embedding, final norm and head are + always held); ``max_tokens`` the most tokens one call accepts (default ``MAX_TOKENS``). + Construction allocates uninitialized rank-local storage on the current CUDA device and + validates the configuration and layout; it reads no checkpoint data. + """ + + _tensor_specs = staticmethod(tensor_specs) + _check_partition = staticmethod(check_partition) + + def __init__( + self, + config: Mapping[str, Any], + mapping: ParallelMapping, + layers: Optional[Sequence[int]] = None, + device: Optional[torch.device] = None, + *, + max_tokens: int = MAX_TOKENS, + ) -> None: + super().__init__() + self.max_tokens = check_max_tokens(max_tokens) + self.schedule = check_config(config) + _check_mapping(mapping) + device = explicit_device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this component is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.config = dict(config) + self.mapping = mapping + self.rank = mapping.rank + self.device = device + self.layers = tuple(sorted(range(self.schedule.num_layers) if layers is None else layers)) + self.specs = self._tensor_specs(config, self.layers) + self._check_partition(self.specs) + self.hidden = config["hidden_size"] + self.head_dim = config["head_dim"] + self.local_heads = config["num_attention_heads"] // WORLD_SIZE + self.local_kv_heads = config["num_key_value_heads"] // WORLD_SIZE + self.eps = config["rms_norm_eps"] + self.window = config["sliding_window"] + self.rope_theta = config["rope_parameters"]["sliding_attention"]["rope_theta"] + self.softcap = config["final_logit_softcapping"] + self.vocab_interval = self.specs[EMBED].interval(self.rank)[1:] + self._allocate_storage() + self._allreduce = AllReduce(mapping=mapping, strategy=AllReduceStrategy.NCCL) + self.usable = False + self.load_record: Optional[dict] = None + self._init_diagnostics() + + def _allocate_storage(self) -> None: + #: Storage is always resolved from the currently registered parameters, never from + #: references taken here, so a replaced parameter or bank is what every call uses. + self.bank_names: list = [] + for spec in self.specs.values(): + if spec.bank is not None: + if spec.bank not in self.bank_names: + local = spec.local_shape(self.rank) + bank = torch.empty( + (self.config["num_experts"], *local), + dtype=torch.bfloat16, + device=self.device, + ) + self.register_parameter( + _attr(spec.bank), nn.Parameter(bank, requires_grad=False) + ) + self.bank_names.append(spec.bank) + continue + dtype = torch.bfloat16 if spec.effective_dtype == "BF16" else torch.float32 + param = nn.Parameter( + torch.empty(spec.local_shape(self.rank), dtype=dtype, device=self.device), + requires_grad=False, + ) + self.register_parameter(_attr(spec.name), param) + + # ------------------------------------------------------------------ diagnostics + + def _init_diagnostics(self) -> None: + """Diagnostic validation mode: ``step_contract_enabled()``, read once per instance. + + On, each MoE layer also proves the four ranks routed identical bytes (an all-gather of + host digests) and callers re-check per-call facts their own validation already fixed. + Off (the default), none of that host work or diagnostic collective runs. Every rank + must run in the same mode, since the agreement is a collective; ``load`` checks that + before the component can run. + """ + self.diagnostics = step_contract_enabled() + + def _check_diagnostics_agree(self) -> None: + """Raise on every rank unless all four ranks run the same diagnostic mode (collective).""" + flag = torch.tensor([1.0 if self.diagnostics else 0.0], device=self.device) + on = int(self._allreduce(flag).item()) + if on not in (0, WORLD_SIZE): + raise LoadError( + f"rank {self.rank}: diagnostic validation is on for {on} of {WORLD_SIZE} ranks; " + f"set {MODELING_V2_VALIDATE_ENV} alike on every rank" + ) + + # ------------------------------------------------------------------ storage + + def tensor(self, name: str) -> torch.Tensor: + """The rank-local storage of logical tensor ``name`` (an expert's view of its bank slot).""" + spec = self.specs[name] + if spec.bank is not None: + return self.bank(spec.bank)[spec.expert] + return getattr(self, _attr(name)) + + def bank(self, bank: str) -> torch.Tensor: + """The current ``[256, ...]`` storage of one sparse layer's expert bank.""" + return getattr(self, _attr(bank)) + + # ------------------------------------------------------------------ loading + + @torch.no_grad() + def load(self, checkpoint_dir: str, *, record_digests: bool = True) -> dict: + """Stream the selected tensors of a caller-supplied checkpoint into rank-local storage. + + Collective: every rank calls it. Metadata is validated against the whole model + (every expected tensor and scale, nothing unsupported, exact dtypes and shapes) and + the directory's config.json against the constructed configuration before any + storage changes. The component is unusable from the first write until every rank + reports a complete load; a failure on any rank raises ``LoadError`` on all ranks + and leaves the component unusable. Returns the load record (coverage, source + digests, bytes read). + """ + self._check_diagnostics_agree() + error: Optional[BaseException] = None + catalog = None + try: + with open(os.path.join(checkpoint_dir, "config.json")) as f: + stored = json.load(f) + if stored != self.config: + differing = sorted( + k for k in set(stored) | set(self.config) if stored.get(k) != self.config.get(k) + ) + raise CheckpointFormatError( + f"config.json differs from the constructed configuration: {differing[:8]}" + ) + catalog = SafetensorsCatalog(checkpoint_dir) + metadata = self._validate_catalog(checkpoint_dir, catalog, stored) + except Exception as exc: # Keep rank-local failures inside the collective protocol. + error = exc + if failure := load_failure_details(self.mapping, error is None, error): + raise LoadError( + f"rank {self.rank}: checkpoint rejected before loading: {failure}" + ) from error + + self.usable = False + self.load_record = None + self._reset_loaded_state() + coverage = OrderedDict() + reader = None + bytes_read = reads = 0 + try: + reader = TileReader() + order = sorted( + self.specs.values(), + key=lambda s: ( + catalog.entries[self._entry_name(s)].path, + catalog.entries[self._entry_name(s)].begin, + ), + ) + for spec in order: + digest = hashlib.sha256() if record_digests else None + coverage[spec.name] = self._load_tensor(catalog, reader, spec, digest) + extra = self._load_extra(catalog, reader) + except Exception as exc: # Every rank must reach the completion agreement. + error = exc + finally: + if reader is not None: + bytes_read, reads = reader.bytes_read, reader.reads + try: + reader.close() + except Exception as exc: # Cleanup failure must also reach the agreement. + if error is None: + error = exc + complete = error is None and list(coverage) == [s.name for s in order] + if failure := load_failure_details(self.mapping, complete, error): + self._reset_loaded_state() + raise LoadError( + f"rank {self.rank}: load did not complete on every rank: {failure}" + ) from error + self.load_record = { + "rank": self.rank, + "checkpoint_dir": os.fspath(checkpoint_dir), + "layers": list(self.layers), + "tensors": len(coverage), + "bytes_read": bytes_read, + "reads": reads, + "coverage": coverage, + **metadata, + **extra, + } + self.usable = True + return self.load_record + + def _validate_catalog( + self, checkpoint_dir: str, catalog: SafetensorsCatalog, stored: dict + ) -> dict: + catalog.validate(self._tensor_specs(self.config), complete=True) + return {} + + def _reset_loaded_state(self) -> None: + """Clear format-specific state after storage changes or a failed transaction.""" + + @staticmethod + def _entry_name(spec: TensorSpec) -> str: + return spec.name + + def _load_tensor( + self, + catalog: SafetensorsCatalog, + reader: TileReader, + spec: TensorSpec, + digest: Optional["hashlib._Hash"], + ) -> dict: + dst = self.tensor(spec.name) + read_slice_into(catalog, reader, spec, self.rank, dst, digest=digest) + return { + "interval": spec.interval(self.rank), + "local_shape": list(dst.shape), + "dtype": str(dst.dtype), + "source_sha256": digest.hexdigest() if digest is not None else None, + "bank": spec.bank, + "slot": spec.expert, + } + + def _load_extra(self, catalog: SafetensorsCatalog, reader: TileReader) -> dict: + return {} + + # ------------------------------------------------------------------ helpers + + def _check_usable(self, layer: Optional[int] = None) -> None: + if not self.usable: + raise RuntimeError("the component is not usable: no complete load on every rank") + if layer is not None and layer not in self.layers: + raise ValueError(f"logical layer {layer} is not held by this component ({self.layers})") + + def _layer(self, layer: int, name: str) -> torch.Tensor: + return self.tensor(f"model.layers.{layer}.{name}") + + def _bank(self, layer: int, which: str) -> torch.Tensor: + return self.bank(f"layers.{layer}.experts.{which}") + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor: + """The one sum over the four ranks at a partial-result boundary (NCCL).""" + return self._allreduce(tensor) + + def _digests_agree(self, tensors: Sequence[torch.Tensor]) -> list: + """Raise unless every rank holds byte-identical ``tensors``; return this rank's digests. + + Diagnostic (``self.diagnostics``): host hashing of device copies and one all-gather, + called on every rank alike or on none. + """ + digests = [] + for t in tensors: + h = hashlib.sha256(f"{t.dtype}{list(t.shape)}".encode()) + h.update(memoryview(t.contiguous().cpu().view(torch.uint8).numpy().reshape(-1))) + digests.append(h.digest()) + local = torch.frombuffer(bytearray(b"".join(digests)), dtype=torch.uint8).to(self.device) + gathered = allgather(local, self.mapping, dim=0).view(WORLD_SIZE, -1).cpu() + if not all(torch.equal(gathered[0], gathered[r]) for r in range(1, WORLD_SIZE)): + raise RuntimeError( + "replicated MoE routing differs between ranks (input, ids or weights)" + ) + return [d.hex() for d in digests] + + # ------------------------------------------------------------------ computations + + @torch.no_grad() + def embed(self, token_ids: torch.Tensor) -> torch.Tensor: + """Global token ids ``[T]`` int64 -> replicated embeddings ``[T, hidden]`` bf16.""" + self._check_usable() + _check_ids(token_ids, self.config["vocab_size"], self.device, self.max_tokens) + lo, hi = self.vocab_interval + inside = (token_ids >= lo) & (token_ids < hi) + local = nn.functional.embedding( + torch.where(inside, token_ids - lo, torch.zeros_like(token_ids)), self.tensor(EMBED) + ) + part = torch.where(inside.unsqueeze(-1), local, torch.zeros_like(local)) + out = self.all_reduce(part) + return out + + @torch.no_grad() + def attention( + self, + layer: int, + h: torch.Tensor, + positions: torch.Tensor, + md: TrtllmAttentionMetadata, + *, + call_checked: bool = False, + ) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> attention output after the one o_proj all-reduce. + + ``md`` is one prepared step of one or more requests (``_attention.check_batch_metadata`` + with one row per generation request) whose new tokens are the T rows, contexts first; ``positions`` are their + absolute positions. ``call_checked`` states that the caller validated this call's + metadata and its cache manager's every layer already (a full-depth call does so once); + the per-layer re-check then runs only in diagnostic mode. + """ + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + T = h.shape[0] + check_positions(positions, (T,), self.device) + if not call_checked or self.diagnostics: + _attention.check_batch_metadata( + md, + T, + layer_idx=layer, + num_kv_heads=self.local_kv_heads, + head_dim=self.head_dim, + ) + step = _attention.step_args(md, layer) + sliding = self.schedule.is_sliding(layer) + part = _attention.attention( + h, + positions, + step, + q_proj=self._layer(layer, "self_attn.q_proj.weight"), + k_proj=self._layer(layer, "self_attn.k_proj.weight"), + v_proj=self._layer(layer, "self_attn.v_proj.weight"), + gate_proj=self._layer(layer, "self_attn.gate_proj.weight"), + o_proj=self._layer(layer, "self_attn.o_proj.weight"), + q_norm=self._layer(layer, "self_attn.q_norm.weight"), + k_norm=self._layer(layer, "self_attn.k_norm.weight"), + gate_bias=self._layer(layer, "self_attn.gate_bias"), + num_heads=self.local_heads, + num_kv_heads=self.local_kv_heads, + head_dim=self.head_dim, + eps=self.eps, + rope_theta=self.rope_theta if sliding else None, + window=self.window if sliding else step["max_seq_len"], + ) + out = self.all_reduce(part) + return out + + @torch.no_grad() + def dense_ffn(self, layer: int, h: torch.Tensor) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> dense SwiGLU output after the one down all-reduce.""" + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + if self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a sparse feed-forward") + gate = cublas_mm(h, self._layer(layer, "mlp.gate_proj.weight").t()) + up = cublas_mm(h, self._layer(layer, "mlp.up_proj.weight").t()) + act = swiglu_fp32(gate, up) + activation = act["act"] + del gate, up, act + part = cublas_mm(activation, self._layer(layer, "mlp.down_proj.weight").t()) + out = self.all_reduce(part) + return out + + @torch.no_grad() + def route(self, layer: int, h: torch.Tensor) -> dict: + """The replicated FP32 router on ``h`` ``[T, hidden]``: every routing intermediate.""" + self._check_usable(layer) + _check_hidden(h, self.hidden, self.device, self.max_tokens) + if not self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a dense feed-forward") + return beam_grouped_topk_router( + h, + self._layer(layer, "mlp.gate.weight"), + self._layer(layer, "mlp.gate.e_score_correction_bias"), + n_group=self.config["n_group"], + topk_group=self.config["topk_group"], + top_k=self.config["num_experts_per_tok"], + eps=ROUTER_EPS, + ) + + @torch.no_grad() + def moe(self, layer: int, h: torch.Tensor) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> MoE output after the one all-reduce. + + In diagnostic mode, before any expert runs, every rank's router input bytes, ordered + expert ids and FP32 weights must be identical (one all-gather of their digests). + """ + routing = self.route(layer, h) + ids, weights = routing["ids"], routing["weights"] + if self.diagnostics: + self._digests_agree((h, ids, weights)) + gate_bank, up_bank, down_bank = (self._bank(layer, w) for w in ("gate", "up", "down")) + + def local_expert(expert_id: int, rows: torch.Tensor) -> torch.Tensor: + g = cublas_mm(rows, gate_bank[expert_id].t()) + u = cublas_mm(rows, up_bank[expert_id].t()) + act = swiglu_fp32(g, u) + y = cublas_mm(act["act"], down_bank[expert_id].t()) + return y + + local = dispatch_combine.dispatch_and_combine(h, ids, weights, local_expert) + out = self.all_reduce(local) + return out + + @torch.no_grad() + def block( + self, + layer: int, + x: torch.Tensor, + positions: torch.Tensor, + md: TrtllmAttentionMetadata, + *, + call_checked: bool = False, + ) -> torch.Tensor: + """One decoder block of logical layer ``layer`` on ``x`` ``[T, hidden]``, one engine step. + + The input norm runs the source model's torch op order (``beam_rmsnorm_fp32``), so the + q/k/v/gate projections read the source's normalized bytes for the same block input. + The other norms use the flashinfer kernel. ``call_checked`` states that the + full-model caller already validated the step metadata (see ``attention``). + """ + self._check_usable(layer) + p = f"model.layers.{layer}." + h = beam_rmsnorm_fp32(x, self.tensor(p + "input_layernorm.weight"), self.eps) + a = self.attention( + layer, + h, + positions, + md, + call_checked=call_checked, + ) + pa = flashinfer_rmsnorm(a, self.tensor(p + "post_attention_layernorm.weight"), self.eps) + x1 = beam_depth_residual(pa, x, *depth_scales(layer, is_attention=True)) + pf = flashinfer_rmsnorm(x1, self.tensor(p + "pre_feedforward_layernorm.weight"), self.eps) + m = self.moe(layer, pf) if self.schedule.is_sparse(layer) else self.dense_ffn(layer, pf) + pm = flashinfer_rmsnorm(m, self.tensor(p + "post_feedforward_layernorm.weight"), self.eps) + out = beam_depth_residual(pm, x1, *depth_scales(layer, is_attention=False)) + return out + + @torch.no_grad() + def final_norm(self, h: torch.Tensor) -> torch.Tensor: + self._check_usable() + _check_hidden(h, self.hidden, self.device, self.max_tokens) + out = flashinfer_rmsnorm(h, self.tensor(FINAL_NORM), self.eps) + return out + + @torch.no_grad() + def head(self, n: torch.Tensor) -> torch.Tensor: + """Final-norm output ``[T, hidden]`` -> softcapped fp32 logits over the full vocabulary. + + The rank's vocabulary rows give local logits; one all-gather along the vocabulary + axis restores global order (rank 0's rows first), and the softcap is applied once. + """ + self._check_usable() + _check_hidden(n, self.hidden, self.device, self.max_tokens) + local = cublas_mm(n, self.tensor(HEAD).t()) + gathered = allgather(local, self.mapping, dim=-1) + logits = beam_logit_softcap(gathered, self.softcap) + return logits + + +def _attr(name: str) -> str: + return name.replace(".", "__") + + +def _check_hidden( + h: torch.Tensor, hidden: int, device: torch.device, max_tokens: int = MAX_TOKENS +) -> None: + if not isinstance(h, torch.Tensor) or h.dtype != torch.bfloat16: + raise ValueError("hidden states must be a torch.bfloat16 tensor") + if h.dim() != 2 or h.shape[1] != hidden or not 1 <= h.shape[0] <= max_tokens: + raise ValueError( + f"hidden states must be [T, {hidden}] with 1 <= T <= {max_tokens}, got {list(h.shape)}" + ) + if h.device != device or not h.is_contiguous(): + raise ValueError("hidden states must be contiguous on the component's device") + + +def _check_ids( + ids: torch.Tensor, vocab: int, device: torch.device, max_tokens: int = MAX_TOKENS +) -> None: + if not isinstance(ids, torch.Tensor) or ids.dtype != torch.int64 or ids.dim() != 1: + raise ValueError("token ids must be a 1-D torch.int64 tensor") + if ids.device != device or not 1 <= ids.shape[0] <= max_tokens: + raise ValueError(f"token ids must hold 1..{max_tokens} ids on the component's device") + if bool(((ids < 0) | (ids >= vocab)).any()): + raise ValueError(f"token ids must be in [0, {vocab})") + + +__all__ = [ + "BeamCheckpointLayoutBf16Sm103Tp4", + "LoadError", + "MAX_TOKENS", + "check_max_tokens", + "TensorSpec", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/layout_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/layout_bf16.py new file mode 100644 index 000000000000..228c57da173e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/layout_bf16.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration, schedule and four-rank ownership of the Beam checkpoint layout component. + +The layout is one fixed four-rank configuration: attention TP4 and MoE TP4 (EP1) on the +same four ranks, PP1, no attention data parallelism. Every logical checkpoint tensor is +either split along one axis into four contiguous, gap-free, non-overlapping intervals in +global checkpoint coordinates, or replicated identically on every rank:: + + embed_tokens, lm_head [151936, 5120] vocabulary rows [37984 r, 37984 (r + 1)) + q_proj, self_attn.gate_proj [6144, 5120] query-head rows [1536 r, 1536 (r + 1)) + k_proj, v_proj [1024, 5120] kv-head rows [256 r, 256 (r + 1)) + o_proj [5120, 6144] query-head columns [1536 r, 1536 (r + 1)) + gate_bias [48] query heads [12 r, 12 (r + 1)) + dense gate_proj, up_proj [20480, 5120] intermediate rows [5120 r, 5120 (r + 1)) + dense down_proj [5120, 20480] intermediate cols [5120 r, 5120 (r + 1)) + expert gate_proj, up_proj [2560, 5120] intermediate rows [640 r, 640 (r + 1)) + expert down_proj [5120, 2560] intermediate cols [640 r, 640 (r + 1)) + norms, q_norm, k_norm replicated BF16 + router weight, correction replicated FP32; expert ids stay global + +Rank r therefore owns query heads [12 r, 12 r + 12) with their two kv heads [2 r, 2 r + 2), +keeping the 6:1 grouped-query ratio, and an intermediate slice of every expert. A row-split +FP8 matrix takes the matching rows of its per-output-channel scale; a column-split one keeps +all of its output scales. Each sparse layer stores its experts' slices in three contiguous +banks, ``gate`` and ``up`` ``[256, 640, 5120]`` and ``down`` ``[256, 5120, 640]``; expert e +is slot e of each bank. +""" + +from __future__ import annotations + +import dataclasses +from collections import OrderedDict +from typing import Any, Mapping, Optional + +from .._constants import SLIDING_ROPE_THETA, SUPPORTED_CONFIG, WORLD_SIZE + +#: Stored dtype names, as safetensors headers spell them. +FP8 = "F8_E4M3" +BF16 = "BF16" +F32 = "F32" + +ATTENTION_NORMS = ( + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", +) +EMBED = "model.embed_tokens.weight" +FINAL_NORM = "model.norm.weight" +HEAD = "lm_head.weight" + + +@dataclasses.dataclass(frozen=True) +class Schedule: + """The per-layer attention and feed-forward kinds the checkpoint's configuration states.""" + + attention: tuple # "sliding_attention" | "full_attention", one per logical layer + mlp: tuple # "dense" | "sparse", one per logical layer + + @property + def num_layers(self) -> int: + return len(self.attention) + + def is_sparse(self, layer: int) -> bool: + return self.mlp[layer] == "sparse" + + def is_sliding(self, layer: int) -> bool: + return self.attention[layer] == "sliding_attention" + + +@dataclasses.dataclass(frozen=True) +class TensorSpec: + """One logical checkpoint tensor: global geometry, stored dtype and four-rank rule.""" + + name: str + shape: tuple + stored_dtype: str + family: str + layer: Optional[int] + expert: Optional[int] + split_dim: Optional[int] # None: replicated on every rank + + @property + def scale_name(self) -> Optional[str]: + """The consumed per-output-channel scale of an FP8 matrix, ``[out, 1]`` BF16.""" + return self.name + "_scale" if self.stored_dtype == FP8 else None + + def interval(self, rank: int, world_size: int = WORLD_SIZE) -> Optional[tuple]: + """``(dim, start, stop)`` of rank ``rank``'s slice, or None for a replicated tensor.""" + check_rank(rank, world_size) + if self.split_dim is None: + return None + extent = self.shape[self.split_dim] + if extent % world_size: + raise ValueError( + f"{self.name}: dim {self.split_dim} of {extent} does not split by {world_size}" + ) + size = extent // world_size + return self.split_dim, rank * size, (rank + 1) * size + + def local_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + interval = self.interval(rank, world_size) + if interval is None: + return tuple(self.shape) + dim, start, stop = interval + shape = list(self.shape) + shape[dim] = stop - start + return tuple(shape) + + @property + def effective_dtype(self) -> str: + """BF16 for FP8 matrices (decoded and scaled) and BF16 tensors; F32 router tensors.""" + return BF16 if self.stored_dtype in (FP8, BF16) else F32 + + @property + def bank(self) -> Optional[str]: + """The sparse layer's bank holding this expert slice, or None for a parameter.""" + if self.expert is None: + return None + return f"layers.{self.layer}.experts.{self.family.removeprefix('expert_')}" + + +def check_rank(rank: int, world_size: int = WORLD_SIZE) -> None: + """Raise ValueError on any rank/world other than one of the four ranks of this layout.""" + if type(world_size) is not int or world_size != WORLD_SIZE: + raise ValueError( + f"this layout runs on exactly {WORLD_SIZE} ranks, got world_size={world_size!r}" + ) + if type(rank) is not int or not 0 <= rank < world_size: + raise ValueError(f"rank must be an int in [0, {world_size}), got {rank!r}") + + +def _get(config: Any, field: str) -> Any: + if isinstance(config, Mapping): + return config.get(field) + return getattr(config, field, None) + + +def check_config(config: Any) -> Schedule: + """Raise ValueError on the first field outside what this layout computes; return the schedule. + + ``config`` is the checkpoint's config.json as a mapping (or an object with those + attributes). The schedule comes from its ``layer_types``/``mlp_layer_types``: every layer + is sliding or full attention and dense or sparse, the dense layers precede the sparse + ones, and only sliding layers use RoPE (``rope_parameters``). + """ + for field, expected in SUPPORTED_CONFIG.items(): + actual = _get(config, field) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this layout computes {expected!r}") + for field, expected in (("partial_rotary_factor", 1.0), ("use_sliding_window", True)): + actual = _get(config, field) + if actual is not None and actual != expected: + raise ValueError(f"unsupported {field}={actual!r}; this layout computes {expected!r}") + layers = SUPPORTED_CONFIG["num_hidden_layers"] + for field in ("layer_types", "mlp_layer_types"): + values = _get(config, field) + if not isinstance(values, (list, tuple)) or not all(isinstance(v, str) for v in values): + raise ValueError(f"{field} must be a list of layer names") + attention = tuple(_get(config, "layer_types")) + mlp = tuple(_get(config, "mlp_layer_types")) + if len(attention) != layers or set(attention) - {"sliding_attention", "full_attention"}: + raise ValueError(f"layer_types must name {layers} sliding/full layers, got {attention!r}") + if len(mlp) != layers or set(mlp) - {"dense", "sparse"}: + raise ValueError(f"mlp_layer_types must name {layers} dense/sparse layers, got {mlp!r}") + if ( + "dense" in mlp + and "sparse" in mlp + and mlp.index("sparse") < len(mlp) - mlp[::-1].index("dense") + ): + raise ValueError("dense feed-forward layers must precede the sparse ones") + rope = _get(config, "rope_parameters") + sliding = (rope or {}).get("sliding_attention") if isinstance(rope, Mapping) else None + if ( + not isinstance(rope, Mapping) + or set(rope) - {"sliding_attention", "full_attention"} + or rope.get("full_attention") is not None + or not isinstance(sliding, Mapping) + or set(sliding) - {"rope_type", "rope_theta", "partial_rotary_factor"} + or sliding.get("rope_type") != "default" + or sliding.get("rope_theta") != SLIDING_ROPE_THETA + or sliding.get("partial_rotary_factor", 1.0) != 1.0 + ): + raise ValueError( + "rope_parameters must give sliding layers default RoPE with theta " + f"{SLIDING_ROPE_THETA} and full layers none, got {rope!r}" + ) + return Schedule(attention=attention, mlp=mlp) + + +def tensor_specs(config: Any, layers: Optional[tuple] = None) -> "OrderedDict[str, TensorSpec]": + """Every logical tensor of the selected layers and the outer tensors, by checkpoint name. + + ``layers`` selects logical layers (default: all of them); the embedding, final norm and + head are always included. Consumed FP8 scales are not logical tensors; see + ``TensorSpec.scale_name``. + """ + schedule = check_config(config) + hidden = _get(config, "hidden_size") + heads = _get(config, "num_attention_heads") + kv_heads = _get(config, "num_key_value_heads") + head_dim = _get(config, "head_dim") + vocab = _get(config, "vocab_size") + dense_inter = _get(config, "intermediate_size") + expert_inter = _get(config, "moe_intermediate_size") + experts = _get(config, "num_experts") + selected = tuple(range(schedule.num_layers)) if layers is None else tuple(layers) + for layer in selected: + if type(layer) is not int or not 0 <= layer < schedule.num_layers: + raise ValueError(f"layer {layer!r} is not a logical layer of the checkpoint") + if len(set(selected)) != len(selected): + raise ValueError(f"layers must be distinct, got {selected}") + + specs: "OrderedDict[str, TensorSpec]" = OrderedDict() + + def add(name, shape, dtype, family, layer=None, expert=None, split_dim=None): + specs[name] = TensorSpec(name, tuple(shape), dtype, family, layer, expert, split_dim) + + add(EMBED, (vocab, hidden), BF16, "embed_tokens", split_dim=0) + for layer in sorted(selected): + p = f"model.layers.{layer}." + for norm in ATTENTION_NORMS: + add(p + norm + ".weight", (hidden,), BF16, "block_norm", layer) + add( + p + "self_attn.q_proj.weight", + (heads * head_dim, hidden), + FP8, + "q_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.k_proj.weight", + (kv_heads * head_dim, hidden), + FP8, + "k_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.v_proj.weight", + (kv_heads * head_dim, hidden), + FP8, + "v_proj", + layer, + split_dim=0, + ) + add( + p + "self_attn.gate_proj.weight", + (heads * head_dim, hidden), + FP8, + "attn_gate", + layer, + split_dim=0, + ) + add( + p + "self_attn.o_proj.weight", + (hidden, heads * head_dim), + FP8, + "o_proj", + layer, + split_dim=1, + ) + add(p + "self_attn.q_norm.weight", (head_dim,), BF16, "qk_norm", layer) + add(p + "self_attn.k_norm.weight", (head_dim,), BF16, "qk_norm", layer) + add(p + "self_attn.gate_bias", (heads,), BF16, "gate_bias", layer, split_dim=0) + if schedule.is_sparse(layer): + add(p + "mlp.gate.weight", (experts, hidden), F32, "router_weight", layer) + add(p + "mlp.gate.e_score_correction_bias", (experts,), F32, "router_bias", layer) + for e in range(experts): + q = f"{p}mlp.experts.{e}." + add(q + "gate_proj.weight", (expert_inter, hidden), FP8, "expert_gate", layer, e, 0) + add(q + "up_proj.weight", (expert_inter, hidden), FP8, "expert_up", layer, e, 0) + add(q + "down_proj.weight", (hidden, expert_inter), FP8, "expert_down", layer, e, 1) + else: + add( + p + "mlp.gate_proj.weight", + (dense_inter, hidden), + FP8, + "dense_gate", + layer, + split_dim=0, + ) + add( + p + "mlp.up_proj.weight", (dense_inter, hidden), FP8, "dense_up", layer, split_dim=0 + ) + add( + p + "mlp.down_proj.weight", + (hidden, dense_inter), + FP8, + "dense_down", + layer, + split_dim=1, + ) + add(FINAL_NORM, (hidden,), BF16, "final_norm") + add(HEAD, (vocab, hidden), BF16, "lm_head", split_dim=0) + return specs + + +def check_partition(specs: Mapping[str, TensorSpec], world_size: int = WORLD_SIZE) -> None: + """Raise unless every split tensor's rank intervals tile its axis exactly, in rank order.""" + for spec in specs.values(): + if spec.split_dim is None: + continue + cursor = 0 + for rank in range(world_size): + dim, start, stop = spec.interval(rank, world_size) + if start != cursor or stop <= start: + raise ValueError( + f"{spec.name}: rank {rank} interval [{start}, {stop}) leaves a gap or overlap" + ) + cursor = stop + if cursor != spec.shape[spec.split_dim]: + raise ValueError( + f"{spec.name}: rank intervals cover {cursor} of {spec.shape[spec.split_dim]}" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/loader_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/loader_bf16.py new file mode 100644 index 000000000000..09fc75346ecf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_tp4/loader_bf16.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Metadata catalog and bounded streaming reads of a safetensors checkpoint directory. + +``SafetensorsCatalog`` reads ``model.safetensors.index.json`` and every shard's header (an +8-byte little-endian length and a JSON table); it never reads a payload. ``validate`` checks +the catalog against the expected logical tensors before anything is loaded: every expected +tensor and every FP8 matrix's scale present, nothing else present, stored dtypes and shapes +exact, each entry's byte span inside its file and matching its dtype and shape. + +``read_slice_into`` copies one rank-local slice of one tensor into its final device +storage through bounded byte-range tiles: a row-split or replicated tensor reads only its +rows; a column-split tensor reads whole rows, tile by tile, and keeps its columns on the +device. Tiles are at most ``TILE_BYTES`` and staged in at most two pinned host buffers, so +no tensor, shard or model is ever materialized on the host. An FP8 matrix's effective value +is ``bf16(float32(e4m3(q)) * float32(scale[row, 0]))``: the stored BF16 per-output-channel +scale multiplies its row, broadcast along the input columns, and the product is rounded to +BF16 once. BF16 and F32 tensors are copied byte for byte. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import struct +import threading +from collections import OrderedDict +from typing import Generator, Mapping, Optional, Sequence + +import torch + +from .layout_bf16 import BF16, F32, FP8, TensorSpec + +INDEX_FILE = "model.safetensors.index.json" +TILE_BYTES = 16 << 20 +_ITEMSIZE = {"U8": 1, FP8: 1, BF16: 2, F32: 4} +_TORCH_DTYPE = {FP8: torch.float8_e4m3fn, BF16: torch.bfloat16, F32: torch.float32} +_MAX_HEADER = 256 << 20 + + +class CheckpointFormatError(ValueError): + """The checkpoint's index, headers or tensors are not what the layout expects.""" + + +class Entry: + """One stored tensor: file, dtype, shape and absolute byte span.""" + + __slots__ = ("name", "path", "dtype", "shape", "begin", "end") + + def __init__( + self, name: str, path: str, dtype: str, shape: Sequence[int], begin: int, end: int + ) -> None: + self.name, self.path, self.dtype, self.shape = name, path, dtype, tuple(shape) + self.begin, self.end = begin, end + + @property + def row_bytes(self) -> int: + cols = 1 + for d in self.shape[1:]: + cols *= d + return cols * _ITEMSIZE[self.dtype] + + +def _read_header(path: str) -> tuple: + with open(path, "rb") as f: + raw = f.read(8) + if len(raw) != 8: + raise CheckpointFormatError(f"{path}: truncated safetensors header length") + (length,) = struct.unpack(" None: + self.checkpoint_dir = os.fspath(checkpoint_dir) + index_path = os.path.join(self.checkpoint_dir, INDEX_FILE) + try: + with open(index_path) as f: + index = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise CheckpointFormatError(f"cannot read {INDEX_FILE}: {exc}") from None + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise CheckpointFormatError(f"{INDEX_FILE} has no weight_map") + self.entries: "OrderedDict[str, Entry]" = OrderedDict() + for filename in sorted(set(weight_map.values())): + if not isinstance(filename, str) or os.path.basename(filename) != filename: + raise CheckpointFormatError(f"{INDEX_FILE} names an unsafe shard path {filename!r}") + path = os.path.join(self.checkpoint_dir, filename) + table, data_start = _read_header(path) + size = os.path.getsize(path) + spans = [] + for name, info in table.items(): + if name == "__metadata__": + continue + try: + dtype, shape = info["dtype"], info["shape"] + begin, end = info["data_offsets"] + except (KeyError, TypeError, ValueError): + raise CheckpointFormatError( + f"{filename}: malformed header entry {name!r}" + ) from None + if weight_map.get(name) != filename: + raise CheckpointFormatError( + f"{name} is stored in {filename} but the index maps it to {weight_map.get(name)!r}" + ) + if type(begin) is not int or type(end) is not int: + raise CheckpointFormatError(f"{filename}: {name} offsets must be integers") + if not 0 <= begin <= end or data_start + end > size: + raise CheckpointFormatError( + f"{filename}: {name} spans [{begin}, {end}) outside the file" + ) + spans.append((begin, end, name)) + self.entries[name] = Entry( + name, path, dtype, shape, data_start + begin, data_start + end + ) + spans.sort() + for (b0, e0, n0), (b1, _, n1) in zip(spans, spans[1:]): + if b1 < e0: + raise CheckpointFormatError(f"{filename}: {n0} and {n1} overlap") + missing = sorted(set(weight_map) - set(self.entries)) + if missing: + raise CheckpointFormatError( + f"index entries absent from their shard headers: {missing[:8]}" + ) + + def validate(self, specs: Mapping[str, TensorSpec], *, complete: bool = True) -> None: + """Raise CheckpointFormatError unless the stored tensors are exactly what ``specs`` needs. + + With ``complete`` (``specs`` describes the whole model) any stored tensor that is + neither an expected logical tensor nor its FP8 scale is rejected; otherwise only the + expected ones are checked. + """ + expected = {} + for spec in specs.values(): + expected[spec.name] = (spec.stored_dtype, tuple(spec.shape)) + if spec.scale_name is not None: + expected[spec.scale_name] = (BF16, (spec.shape[0], 1)) + self.validate_entries(expected, complete=complete) + + def validate_entries( + self, expected: Mapping[str, tuple[str, tuple[int, ...]]], *, complete: bool = True + ) -> None: + """Check stored names, dtypes, shapes and byte spans for any supported weight format.""" + missing = sorted(set(expected) - set(self.entries)) + if missing: + raise CheckpointFormatError( + f"{len(missing)} expected tensors are missing: {missing[:8]}" + ) + if complete: + extra = sorted(set(self.entries) - set(expected)) + if extra: + raise CheckpointFormatError(f"{len(extra)} unsupported tensors: {extra[:8]}") + for name, (dtype, shape) in expected.items(): + entry = self.entries[name] + if entry.dtype != dtype or entry.shape != shape: + raise CheckpointFormatError( + f"{name} is stored as {entry.dtype}{list(entry.shape)}; expected {dtype}{list(shape)}" + ) + if entry.end - entry.begin != math.prod(shape) * _ITEMSIZE[dtype]: + raise CheckpointFormatError( + f"{name}: byte span does not match {dtype}{list(shape)}" + ) + + +class TileReader: + """Bounded byte-range reads: at most two pinned staging buffers of ``tile_bytes``. + + ``tiles(entry, rows)`` yields ``(row_start, row_stop, host_view)`` for consecutive row + ranges; the next tile is read on a background thread while the caller consumes the + current one, and a buffer is refilled only after the caller's device copy of it + completed. Counts the bytes read and the peak staging footprint. + """ + + def __init__(self, tile_bytes: int = TILE_BYTES, pin: bool = True) -> None: + self.tile_bytes = tile_bytes + self.buffers = [ + torch.empty(tile_bytes, dtype=torch.uint8, pin_memory=pin) for _ in range(2) + ] + self.bytes_read = 0 + self.reads = 0 + self._fds: "OrderedDict[str, int]" = OrderedDict() + self._lock = threading.Lock() + + def close(self) -> None: + for fd in self._fds.values(): + os.close(fd) + self._fds.clear() + self.buffers = [] + + def _fd(self, path: str) -> int: + fd = self._fds.get(path) + if fd is None: + if len(self._fds) >= 4: + _, old = self._fds.popitem(last=False) + os.close(old) + fd = os.open(path, os.O_RDONLY) + self._fds[path] = fd + return fd + + def _read(self, buf: torch.Tensor, fd: int, offset: int, length: int) -> None: + view = memoryview(buf.numpy())[:length] + got = 0 + while got < length: + n = os.preadv(fd, [view[got:]], offset + got) + if n <= 0: + raise OSError(f"short read at offset {offset + got}") + got += n + with self._lock: + self.bytes_read += length + self.reads += 1 + + def tiles( + self, entry: Entry, row_start: int, row_stop: int + ) -> Generator[tuple[int, int, torch.Tensor], None, None]: + row_bytes = entry.row_bytes + if row_bytes > self.tile_bytes: + raise CheckpointFormatError(f"{entry.name}: one row ({row_bytes} B) exceeds the tile") + rows_per_tile = max(1, self.tile_bytes // row_bytes) + fd = self._fd(entry.path) + plan = [ + (r, min(r + rows_per_tile, row_stop)) for r in range(row_start, row_stop, rows_per_tile) + ] + if not plan: + return + pending = None + try: + for i, (r0, r1) in enumerate(plan): + buf = self.buffers[i % 2] + length = (r1 - r0) * row_bytes + if pending is None: + self._read(buf, fd, entry.begin + r0 * row_bytes, length) + else: + pending.join() + if pending.error is not None: + raise pending.error + pending = None + # The other buffer held tile i - 1, whose device copy the caller completed + # before asking for tile i: read tile i + 1 into it while tile i is consumed. + if i + 1 < len(plan): + n0, n1 = plan[i + 1] + pending = _Prefetch( + self, + self.buffers[(i + 1) % 2], + fd, + entry.begin + n0 * row_bytes, + (n1 - n0) * row_bytes, + ) + pending.start() + yield r0, r1, buf[:length] + finally: + if pending is not None: + pending.join() + + +class _Prefetch(threading.Thread): + def __init__( + self, reader: TileReader, buf: torch.Tensor, fd: int, offset: int, length: int + ) -> None: + super().__init__(daemon=True) + self.args = (buf, fd, offset, length) + self.reader = reader + self.error: BaseException | None = None + + def run(self) -> None: + try: + self.reader._read(*self.args) + except BaseException as exc: # Worker boundary: tiles() re-raises on the consumer. + self.error = exc + + +def _decode_tile(raw: torch.Tensor, entry: Entry, rows: int, device: torch.device) -> torch.Tensor: + """Host tile bytes -> device tensor ``[rows, ...]`` in the stored dtype.""" + dev = raw.to(device, non_blocking=True) + return dev.view(_TORCH_DTYPE[entry.dtype]).view(rows, *entry.shape[1:]) + + +@torch.no_grad() +def read_slice_into( + catalog: SafetensorsCatalog, + reader: TileReader, + spec: TensorSpec, + rank: int, + dst: torch.Tensor, + *, + digest: Optional["hashlib._Hash"] = None, +) -> None: + """Copy rank ``rank``'s effective slice of ``spec`` into ``dst`` (its final device storage). + + ``dst`` must have the rank-local shape and the effective dtype. ``digest`` (optional) is + updated with the source bytes actually read, in tile order, with each tile's row range. + """ + entry = catalog.entries[spec.name] + local = spec.local_shape(rank) + want = torch.bfloat16 if spec.stored_dtype in (FP8, BF16) else torch.float32 + if tuple(dst.shape) != local or dst.dtype != want: + raise ValueError( + f"{spec.name}: destination {dst.dtype}{list(dst.shape)} != {want}{list(local)}" + ) + interval = spec.interval(rank) + rows = spec.shape[0] + row_start, row_stop, col = 0, rows, None + if interval is not None: + dim, start, stop = interval + if dim == 0: + row_start, row_stop = start, stop + else: + col = (start, stop) + scale = None + if spec.stored_dtype == FP8: + scale_entry = catalog.entries[spec.scale_name] + scale = torch.empty(row_stop - row_start, 1, dtype=torch.bfloat16, device=dst.device) + for r0, r1, raw in reader.tiles(scale_entry, row_start, row_stop): + scale[r0 - row_start : r1 - row_start].copy_( + _decode_tile(raw, scale_entry, r1 - r0, dst.device) + ) + torch.cuda.current_stream(dst.device).synchronize() + if digest is not None: + digest.update(f"{scale_entry.name}:{r0}:{r1}|".encode()) + digest.update(memoryview(raw.numpy())) + for r0, r1, raw in reader.tiles(entry, row_start, row_stop): + tile = _decode_tile(raw, entry, r1 - r0, dst.device) + if col is not None: + tile = tile[:, col[0] : col[1]] + out = dst[r0 - row_start : r1 - row_start] + if spec.stored_dtype == FP8: + s = scale[r0 - row_start : r1 - row_start].float() + out.copy_((tile.float() * s).to(torch.bfloat16)) + else: + out.copy_(tile) + torch.cuda.current_stream(dst.device).synchronize() + if digest is not None: + digest.update(f"{entry.name}:{r0}:{r1}|".encode()) + digest.update(memoryview(raw.numpy())) + del scale diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_validation.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_validation.py new file mode 100644 index 000000000000..c1a7773d26f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/_validation.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host-side shape and device admission for directly constructed Beam components.""" + +import torch + + +def explicit_device(device: str | torch.device | None) -> torch.device: + """Resolve an unindexed CUDA device against the current device once at construction.""" + result = torch.device("cuda" if device is None else device) + if result.type == "cuda" and result.index is None: + result = torch.device("cuda", torch.cuda.current_device()) + return result + + +def check_positions(positions: torch.Tensor, shape: tuple[int, ...], device: torch.device) -> None: + """Require int32/int64 position metadata of the exact shape without reading device values.""" + if ( + not isinstance(positions, torch.Tensor) + or positions.dtype not in (torch.int32, torch.int64) + or tuple(positions.shape) != shape + or positions.device != device + ): + raise ValueError(f"position_ids must be an int32/int64 tensor of shape {shape} on {device}") + + +def metadata_values(values, count: int, name: str) -> list[int]: + """Read a complete prefix of a host metadata array; never silently truncate it.""" + if len(values) < count: + raise ValueError(f"{name} needs {count} entries, got {len(values)}") + prefix = values[:count] + return prefix.tolist() if hasattr(prefix, "tolist") else list(prefix) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dispatch_combine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dispatch_combine.py new file mode 100644 index 000000000000..7a1565db9089 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/dispatch_combine.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Beam MoE dispatch and weighted combination for TP1 components and TP4 models. + +A caller supplies hidden states, eight expert ids and mixing weights per token, and +an expert provider. This stateless helper sends token rows to their experts and +returns the weighted, combined and scaled outputs. The provider owns parameter +loading and expert execution. + +Ownership. This component owns dispatch and combination; the provider owns the expert +computation. ``expert_provider(expert_id, rows)`` is called once for every selected +expert, in ascending id order, with that expert's ``rows`` -- ``[n, 5120]`` +``torch.bfloat16``, ``n >= 1``, the routed tokens in ascending token order -- and returns +the expert's unweighted output for them: ``[n, 5120]`` ``torch.bfloat16`` on the same +device, row ``i`` for input row ``i``. It applies no mixing weight and no routed scaling +factor; both are applied here. A routed expert component per id adapts to it as +``lambda expert_id, rows: experts[expert_id](rows)``. The router owns selection: no expert +is re-selected, no bias is added and no weight is re-normalized here. No residual, +normalization or collective is applied. + +The combination, as the source model computes it (``ops/moe/beam_dispatch_combine.py``):: + + for each selected expert e, ascending: + out = expert_provider(e, hidden[tokens routed to e]) [n, 5120] bf16 + route = (out.float() * weights[:, None]).to(bfloat16) [n, 5120] each route rounded once + acc[tokens] += route.float() [T, 5120] fp32 + output = (acc * 2.5).to(bfloat16) [T, 5120] bf16 +""" + +from __future__ import annotations + +import functools +from typing import Callable + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.moe.beam_dispatch_combine import ( + beam_dispatch_combine, +) + +from ._constants import SM as _SM +from ._constants import SUPPORTED_CONFIG + +HIDDEN_SIZE = SUPPORTED_CONFIG["hidden_size"] +EXPERTS_PER_TOKEN = SUPPORTED_CONFIG["num_experts_per_tok"] +ROUTED_SCALING_FACTOR = SUPPORTED_CONFIG["routed_scaling_factor"] + +#: ``expert_provider(expert_id, rows) -> output``; see the module docstring. +ExpertProvider = Callable[[int, torch.Tensor], torch.Tensor] + + +def _check_tensor(name: str, tensor, dtype: torch.dtype) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}") + if tensor.layout != torch.strided: + raise ValueError(f"{name} must be a dense tensor, got {tensor.layout}") + + +def _check_inputs( + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: ExpertProvider, +) -> None: + """Raise on any input this component would otherwise compute wrongly. + + Reads tensor metadata only. The ids' values are not checked: each in ``[0, 256)`` and + the eight ids of a row distinct are the caller's preconditions -- the router's output + satisfies both -- because checking them would copy the routing table to the host on + every call. + """ + _check_tensor("hidden_states", hidden_states, torch.bfloat16) + _check_tensor("expert_ids", expert_ids, torch.int64) + _check_tensor("mixing_weights", mixing_weights, torch.float32) + if hidden_states.ndim != 2 or hidden_states.shape[1] != HIDDEN_SIZE: + raise ValueError( + f"hidden_states must be [tokens, {HIDDEN_SIZE}], got {list(hidden_states.shape)}" + ) + tokens = hidden_states.shape[0] + if tokens < 1: + raise ValueError("hidden_states must hold at least one token") + for name, tensor in (("expert_ids", expert_ids), ("mixing_weights", mixing_weights)): + if tuple(tensor.shape) != (tokens, EXPERTS_PER_TOKEN): + raise ValueError( + f"{name} must be [tokens, {EXPERTS_PER_TOKEN}] = [{tokens}, " + f"{EXPERTS_PER_TOKEN}] for {tokens} tokens, got {list(tensor.shape)}" + ) + device = hidden_states.device + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"hidden_states must be on an sm_{_SM[0]}{_SM[1]} CUDA device, got {device}" + ) + for name, tensor in (("expert_ids", expert_ids), ("mixing_weights", mixing_weights)): + if tensor.device != device: + raise ValueError( + f"{name} must be on hidden_states' device {device}, got {tensor.device}" + ) + if not callable(expert_provider): + raise TypeError(f"expert_provider must be callable, got {type(expert_provider).__name__}") + + +def _checked_expert_output( + expert_provider: ExpertProvider, expert_id: int, rows: torch.Tensor +) -> torch.Tensor: + """Call the provider once and return its output unchanged, or raise if it breaks the contract.""" + out = expert_provider(expert_id, rows) + if not isinstance(out, torch.Tensor): + raise TypeError( + f"expert_provider({expert_id}) must return a torch.Tensor, got {type(out).__name__}" + ) + if out.dtype != torch.bfloat16: + raise TypeError(f"expert_provider({expert_id}) must return torch.bfloat16, got {out.dtype}") + if out.shape != rows.shape: + raise ValueError( + f"expert_provider({expert_id}) must return {list(rows.shape)} for " + f"{rows.shape[0]} rows, got {list(out.shape)}" + ) + if out.device != rows.device or out.layout != torch.strided: + raise ValueError( + f"expert_provider({expert_id}) must return a dense tensor on {rows.device}, " + f"got {out.layout} on {out.device}" + ) + return out + + +@torch.no_grad() +def dispatch_and_combine( + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: ExpertProvider, +) -> torch.Tensor: + """Dispatch tokens to their experts and return the combined output. + + ``hidden_states`` ``[T, 5120]`` ``torch.bfloat16`` on an sm_103 CUDA device, T >= 1; + ``expert_ids`` ``[T, 8]`` ``torch.int64`` and ``mixing_weights`` ``[T, 8]`` + ``torch.float32`` on the same device, as the router returns them; ``expert_provider`` + as in the module docstring. Returns ``[T, 5120]`` ``torch.bfloat16`` on that device, + row ``t`` for token ``t``. Inputs outside that, and a provider output outside its + contract, are rejected with an error naming the argument; the three tensors are never + modified. + + Preconditions, not checked: every id in ``[0, 256)`` and the eight ids of each row + distinct. Mixing weights are used as given; they need not sum to one, and a zero + weight still calls its expert. + """ + _check_inputs(hidden_states, expert_ids, mixing_weights, expert_provider) + return beam_dispatch_combine( + hidden_states, + expert_ids, + mixing_weights, + functools.partial(_checked_expert_output, expert_provider), + routed_scaling_factor=ROUTED_SCALING_FACTOR, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/component.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/component.py new file mode 100644 index 000000000000..4c8660547915 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/component.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam NVFP4-expert checkpoint-layout component / sm_103 / tp4: packed experts and blocks. + +Owned by the routed NVFP4 target (this directory); the full model (``model``) builds on it. +Constructed on its own it is a component, not a target: a caller on each of four ranks +constructs this class with the checkpoint's configuration and its rank's ``Mapping`` (attention +TP4, MoE TP4/EP1, PP1), loads a caller-supplied ModelOpt NVFP4 checkpoint directory into the +rank's slices, and runs representative computations. It does not run a full-depth forward, +prefill/decode loop or sampler. + +This is the BF16 layout component (``component_bf16``) with NVFP4 expert +storage and execution. Everything outside the routed experts -- ownership, the embedding, +attention on TRTLLM attention and KVCacheManagerV2 metadata, the dense feed-forward, the +replicated FP32 router with its diagnostic-mode all-rank agreement check, dispatch/combine, the +norms, depth residuals and head -- uses that component's code on +the checkpoint's BF16/FP32 tensors copied byte for byte. The experts stay packed: + + storage per sparse layer and projection kind, banks indexed by global expert id: packed + E2M1 codes ``[256, rows, K/2]`` u8 and E4M3 block scales in the GEMM's 128x4 + swizzled order ``[256, rows * K/16]`` u8 (the rank's slice, see ``layout``), and + ``input_scale``, ``weight_scale_2``, ``g = 1/input_scale``, ``alpha = + input_scale * weight_scale_2`` ``[256]`` fp32 + compute ``expert.nvfp4_swiglu_expert``: per selected expert, activation quantization, + three CUTLASS NVFP4 GEMMs and the fp32 SwiGLU, as the expert provider of + ``beam_dispatch_combine``; the combination (bf16 routes, fp32 ascending-expert + sum, x2.5, bf16) is unchanged, and the one MoE all-reduce sums the four bf16 rank + partials in fp32 and rounds the sum to bf16 once + +Loading streams bounded tiles straight into that storage and is a transaction: metadata +(config.json, hf_quant_config.json, every index/header entry) is validated against the whole +model before any storage changes, and the component is usable only after every rank loaded +every selected tensor. Packed codes, block scales and derived scalars are all registered +parameters. +""" + +from __future__ import annotations + +import hashlib +import struct +from typing import Optional + +import torch +from torch import nn + +from .. import dispatch_combine +from .._tp4.component_bf16 import BeamCheckpointLayoutBf16Sm103Tp4, LoadError, _attr +from . import expert as nvfp4_expert +from .layout import ( + PROJECTIONS, + QuantSpec, + TensorSpec, + check_partition, + check_recipe, + read_quant_config, + stored_entries, + swizzled_numel, + tensor_specs, +) +from .loader import ( + CheckpointFormatError, + SafetensorsCatalog, + TileReader, + check_block_scales, + read_packed_into, + read_scalars, + validate, +) + +#: Bank fields of one projection kind of one sparse layer. +BANK_FIELDS = ("weight", "weight_scale", "input_scale", "weight_scale_2", "g", "alpha") +_KINDS = tuple(p.removesuffix("_proj") for p in PROJECTIONS) + + +def _f32_bits(value: float) -> int: + return struct.unpack(" None: + self.bank_names: list = [] + self.expert_names: list = [] + experts = self.config["num_experts"] + for spec in self.specs.values(): + if isinstance(spec, QuantSpec): + self.expert_names.append(spec.name) + if spec.bank in self.bank_names: + continue + rows, half = spec.local_packed_shape(self.rank) + _, scale_cols = spec.local_scale_shape(self.rank) + fields = { + "weight": ((experts, rows, half), torch.uint8), + "weight_scale": ((experts, swizzled_numel(rows, scale_cols)), torch.uint8), + "input_scale": ((experts,), torch.float32), + "weight_scale_2": ((experts,), torch.float32), + "g": ((experts,), torch.float32), + "alpha": ((experts,), torch.float32), + } + for field, (shape, dtype) in fields.items(): + self.register_parameter( + _attr(f"{spec.bank}.{field}"), + nn.Parameter( + torch.empty(shape, dtype=dtype, device=self.device), requires_grad=False + ), + ) + self.bank_names.append(spec.bank) + continue + dtype = torch.bfloat16 if spec.effective_dtype == "BF16" else torch.float32 + param = nn.Parameter( + torch.empty(spec.local_shape(self.rank), dtype=dtype, device=self.device), + requires_grad=False, + ) + self.register_parameter(_attr(spec.name), param) + #: Per sparse layer, per global expert: gate and up input scales bitwise equal. + self._shared_fc1: dict = {} + + # ------------------------------------------------------------------ storage + + def tensor(self, name: str) -> torch.Tensor: + """The rank-local storage of higher-precision tensor ``name``.""" + spec = self.specs[name] + if isinstance(spec, QuantSpec): + raise ValueError(f"{name} is an NVFP4 expert projection; see expert_storage()") + return getattr(self, _attr(name)) + + def bank(self, bank: str, field: str = "weight") -> torch.Tensor: + """The current ``[256, ...]`` storage of one field of one sparse layer's expert bank.""" + if field not in BANK_FIELDS: + raise ValueError(f"unknown bank field {field!r}; expected one of {BANK_FIELDS}") + return getattr(self, _attr(f"{bank}.{field}")) + + def expert_bank(self, layer: int) -> dict: + """``{"gate"|"up"|"down": {field: [256, ...] storage}}`` of sparse layer ``layer``.""" + self._check_usable(layer) + if not self.schedule.is_sparse(layer): + raise ValueError(f"layer {layer} has a dense feed-forward") + return { + kind: {f: self.bank(f"layers.{layer}.experts.{kind}", f) for f in BANK_FIELDS} + for kind in _KINDS + } + + def expert_storage(self, layer: int, expert: int, proj: str) -> dict: + """Views of one expert projection's rank-local storage (``proj`` in ``PROJECTIONS``).""" + spec = self.specs.get(f"model.layers.{layer}.mlp.experts.{expert}.{proj}") + if not isinstance(spec, QuantSpec): + raise ValueError(f"layer {layer} expert {expert} {proj!r} is not held") + return {f: self.bank(spec.bank, f)[expert] for f in BANK_FIELDS} + + # ------------------------------------------------------------------ loading + + def _load_expert(self, catalog, reader, spec: QuantSpec, digest) -> dict: + dim, p0, p1 = spec.packed_interval(self.rank) + _, s0, s1 = spec.scale_interval(self.rank) + if dim == 0: + rows, packed_cols, scale_cols = (p0, p1), None, None + else: + rows, packed_cols, scale_cols = (0, spec.shape[0]), (p0, p1), (s0, s1) + packed = self.bank(spec.bank, "weight")[spec.expert] + read_packed_into( + catalog, reader, spec.name + ".weight", rows, packed_cols, packed, digest=digest + ) + linear = torch.empty( + spec.local_scale_shape(self.rank), dtype=torch.uint8, device=self.device + ) + read_packed_into( + catalog, reader, spec.name + ".weight_scale", rows, scale_cols, linear, digest=digest + ) + check_block_scales(spec.name + ".weight_scale", linear) + swizzled = torch.ops.trtllm.block_scale_interleave(linear).reshape(-1) + dst = self.bank(spec.bank, "weight_scale")[spec.expert] + if swizzled.numel() != dst.numel(): + raise CheckpointFormatError( + f"{spec.name}: swizzled scales hold {swizzled.numel()} bytes, storage {dst.numel()}" + ) + dst.copy_(swizzled) + torch.cuda.current_stream(self.device).synchronize() + return { + "interval": list(spec.interval(self.rank)), + "packed_interval": [dim, p0, p1], + "scale_interval": [dim, s0, s1], + "local_packed_shape": list(packed.shape), + "local_scale_shape": list(linear.shape), + "bank": spec.bank, + "slot": spec.expert, + "source_sha256": digest.hexdigest() if digest is not None else None, + } + + def _load_scalars(self, catalog, reader, layer: int) -> dict: + experts = [s for s in self.specs.values() if isinstance(s, QuantSpec) and s.layer == layer] + names = [s.name + suffix for s in experts for suffix in (".input_scale", ".weight_scale_2")] + values = read_scalars(catalog, reader, names) + bad = [n for n, v in values.items() if not (v > 0.0 and v != float("inf"))] + if bad: + raise CheckpointFormatError(f"non-finite or non-positive scalars: {bad[:4]}") + n_experts = self.config["num_experts"] + shared = [False] * n_experts + for kind in _KINDS: + s = torch.empty(n_experts, dtype=torch.float32) + w2 = torch.empty(n_experts, dtype=torch.float32) + for e in range(n_experts): + p = f"model.layers.{layer}.mlp.experts.{e}.{kind}_proj" + s[e], w2[e] = values[p + ".input_scale"], values[p + ".weight_scale_2"] + g = torch.ones(n_experts, dtype=torch.float32) / s + alpha = s * w2 + bank = f"layers.{layer}.experts.{kind}" + for field, value in ( + ("input_scale", s), + ("weight_scale_2", w2), + ("g", g), + ("alpha", alpha), + ): + self.bank(bank, field).copy_(value) + for e in range(n_experts): + p = f"model.layers.{layer}.mlp.experts.{e}." + shared[e] = _f32_bits(values[p + "gate_proj.input_scale"]) == _f32_bits( + values[p + "up_proj.input_scale"] + ) + self._shared_fc1[layer] = shared + return {"scalars": len(values), "gate_up_input_scales_shared": sum(shared)} + + def _validate_catalog( + self, checkpoint_dir: str, catalog: SafetensorsCatalog, stored: dict + ) -> dict: + full = tensor_specs(self.config) + recipe = check_recipe(read_quant_config(checkpoint_dir, stored), full) + validate(catalog, stored_entries(full), complete=True) + return { + "recipe": recipe, + "contract": dict(nvfp4_expert.CONTRACT), + "expert_projections": len(self.expert_names), + } + + def _reset_loaded_state(self) -> None: + self._shared_fc1 = {} + + @staticmethod + def _entry_name(spec: TensorSpec | QuantSpec) -> str: + return spec.name + ".weight" if isinstance(spec, QuantSpec) else spec.name + + def _load_tensor( + self, + catalog: SafetensorsCatalog, + reader: TileReader, + spec: TensorSpec | QuantSpec, + digest: Optional["hashlib._Hash"], + ) -> dict: + if isinstance(spec, QuantSpec): + return self._load_expert(catalog, reader, spec, digest) + record = super()._load_tensor(catalog, reader, spec, digest) + # Packed expert ownership is recorded by _load_expert; dense tensors are unbanked. + record.pop("bank") + record.pop("slot") + return record + + def _load_extra(self, catalog: SafetensorsCatalog, reader: TileReader) -> dict: + scalars = { + layer: self._load_scalars(catalog, reader, layer) + for layer in self.layers + if self.schedule.is_sparse(layer) + } + return {"scalars": scalars} + + # ------------------------------------------------------------------ computations + + @torch.no_grad() + def moe(self, layer: int, h: torch.Tensor) -> torch.Tensor: + """Normalized input ``[T, hidden]`` -> MoE output after the one all-reduce. + + The replicated FP32 router, the diagnostic-mode all-rank agreement on its input, ids + and weights and the combination are the BF16 component's; each selected expert runs + ``expert.nvfp4_swiglu_expert`` on this rank's packed slice. The rank partials are + reduced once, in FP32, then rounded to BF16. + """ + routing = self.route(layer, h) + ids, weights = routing["ids"], routing["weights"] + if self.diagnostics: + self._digests_agree((h, ids, weights)) + bank = self.expert_bank(layer) + shared = self._shared_fc1[layer] + + def provider(expert_id: int, rows: torch.Tensor) -> torch.Tensor: + return nvfp4_expert.nvfp4_swiglu_expert(rows, bank, expert_id, shared[expert_id]) + + local = dispatch_combine.dispatch_and_combine(h, ids, weights, provider) + # The four BF16 rank partials are summed once in FP32 and the sum is rounded once. + out = self.all_reduce(local.float()).to(torch.bfloat16) + return out + + +__all__ = [ + "BANK_FIELDS", + "BeamCheckpointLayoutNvfp4Sm103Tp4", + "LoadError", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/expert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/expert.py new file mode 100644 index 000000000000..238557f10644 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/expert.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NVFP4 W4A4 SwiGLU expert on one rank's intermediate slice: the MoE expert provider. + +``nvfp4_swiglu_expert`` is the provider ``beam_dispatch_combine`` calls once per selected expert, +in ascending expert order, with that expert's routed rows ``x`` ``[n, 5120]`` bf16. With the +rank's packed codes ``W``, 128x4-swizzled block scales ``S`` and FP32 scalars of expert ``e``:: + + fc1, sf1 = fp4_quantize(x, g_gate[e]) [n, 2560] u8, swizzled scales + gate = nvfp4_gemm(fc1, W_gate[e], sf1, S_gate[e], alpha_gate[e]) [n, 640] bf16 + up = nvfp4_gemm(fc1, W_up[e], sf1, S_up[e], alpha_up[e]) [n, 640] bf16 + act = bf16(silu(gate.float()) * up.float()) swiglu_fp32 + fc2, sf2 = fp4_quantize(act, g_down[e]) [n, 320] u8, swizzled scales + y = nvfp4_gemm(fc2, W_down[e], sf2, S_down[e], alpha_down[e]) [n, 5120] bf16 + +``g = 1 / input_scale`` and ``alpha = input_scale * weight_scale_2`` per projection. Gate and up +share one activation quantization only when their stored input scales are bitwise equal; +otherwise up quantizes ``x`` with its own ``g``. Every GEMM is the CUTLASS NVFP4 GEMM +(``allowed_backends="cutlass"``) with BF16 output; there is no BF16 expert path. The output is +the expert's unweighted rank partial: mixing weights, the routed scaling factor and the +tensor-parallel sum are applied by the caller. + +Per expert this is two quantizer calls (three when the gate/up input scales differ), three +GEMMs and one elementwise SwiGLU stage. Expert enumeration happens on the host in +``beam_dispatch_combine``, so the path synchronizes with the device and is not CUDA-graph +capturable; a batched pre-routed backend can replace this provider without changing routing, +combination or the reduction. +""" + +from __future__ import annotations + +from typing import Mapping + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.swiglu_fp32 import ( + swiglu_fp32, +) + +from .layout import GROUP + +#: The only NVFP4 GEMM backend this provider runs. +GEMM_BACKENDS = "cutlass" + +#: The kernel contract, as recorded in load and run records. +CONTRACT = { + "activation_quantizer": "torch.ops.trtllm.fp4_quantize(x, g, sf_vec_size=16, " + "sf_use_ue8m0=False, is_sf_swizzled_layout=True)", + "gemm": "torch.ops.trtllm.nvfp4_gemm(..., output_dtype=torch.bfloat16, " + f"allowed_backends={GEMM_BACKENDS!r})", + "activation": "swiglu_fp32: bf16(silu(gate.float()) * up.float())", + "g": "float32(1 / input_scale)", + "alpha": "float32(input_scale * weight_scale_2)", + "weight_scales": "E4M3 bytes, rank slice first, then torch.ops.trtllm.block_scale_interleave", +} + + +def _gemm(a: torch.Tensor, a_sf: torch.Tensor, proj: Mapping[str, torch.Tensor], e: int): + return nvfp4_gemm( + a, + proj["weight"][e], + a_sf, + proj["weight_scale"][e], + proj["alpha"][e : e + 1], + torch.bfloat16, + 0, + GEMM_BACKENDS, + ) + + +def _quantize(x: torch.Tensor, proj: Mapping[str, torch.Tensor], e: int): + return fp4_quantize(x, proj["g"][e : e + 1], GROUP, False, True) + + +@torch.no_grad() +def nvfp4_swiglu_expert( + rows: torch.Tensor, + bank: Mapping[str, Mapping[str, torch.Tensor]], + expert: int, + shared_fc1: bool, +) -> torch.Tensor: + """Expert ``expert``'s rank partial for ``rows`` ``[n, 5120]`` bf16 -> ``[n, 5120]`` bf16. + + ``bank`` maps ``"gate"``, ``"up"``, ``"down"`` to the layer's banks (``weight``, + ``weight_scale``, ``g``, ``alpha``, each indexed by global expert id); ``shared_fc1`` says + the gate and up input scales of this expert are bitwise equal. + """ + gate_p, up_p, down_p = bank["gate"], bank["up"], bank["down"] + fc1, sf1 = _quantize(rows, gate_p, expert) + fc1_up, sf1_up = (fc1, sf1) if shared_fc1 else _quantize(rows, up_p, expert) + gate = _gemm(fc1, sf1, gate_p, expert) + up = _gemm(fc1_up, sf1_up, up_p, expert) + act = swiglu_fp32(gate, up) + fc2, sf2 = _quantize(act["act"], down_p, expert) + y = _gemm(fc2, sf2, down_p, expert) + return y diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/layout.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/layout.py new file mode 100644 index 000000000000..4e9c8aeaf320 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/layout.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Recipe, stored geometry and four-rank ownership of the Beam NVFP4-expert checkpoint layout. + +The checkpoint is a ModelOpt NVFP4 export (``hf_quant_config.json``: ``quant_algo`` NVFP4, +group size 16, static activation scales, no KV-cache quantization). Only the routed experts' +gate/up/down projections are quantized; every other tensor keeps its stored precision: the +attention, dense feed-forward, norm, embedding and head tensors are BF16, the router's weight and +correction bias FP32. The four-rank configuration and every logical tensor's interval are those +of the BF16 layout (``layout_bf16``); this module changes only +what is stored. + +An expert projection of logical shape ``[N, K]`` is stored as four tensors:: + + .weight U8 [N, K/2] two E2M1 codes per byte; the low nibble holds + the even K index, the high nibble the next one + .weight_scale F8_E4M3 [N, K/16] one block scale per 16 consecutive K elements + .weight_scale_2 F32 [] the projection's global weight scale + .input_scale F32 [] the static activation scale s + +Its logical weight is ``W[n, k] = e2m1(code[n, k]) * e4m3(weight_scale[n, k // 16]) * +weight_scale_2``. An activation row ``x`` is quantized with ``g = 1 / s`` (per 16-element block +``b``: ``scale_b = e4m3(g * max|x_b| / 6)``, ``code = e2m1(x * g / scale_b)``), and a projection +is ``alpha * (codes x block scales) @ (codes x block scales)^T`` with ``alpha = s * +weight_scale_2``. + +Rank r owns intermediate ``[640 r, 640 r + 640)`` of every expert, on whole 16-element blocks:: + + gate_proj, up_proj [2560, 5120] rows [640 r, 640 r + 640) packed [640, 2560], scales [640, 320] + down_proj [5120, 2560] cols [640 r, 640 r + 640) packed [5120, 320] = byte columns + [320 r, 320 r + 320); scales + [5120, 40] = scale columns + [40 r, 40 r + 40) + +Every rank holds every projection's two FP32 scalars unchanged. Each sparse layer keeps each +projection kind of its 256 experts in banks (``layers..experts.``): packed codes +``[256, rows, K/2]``, block scales in the GEMM's 128x4 swizzled order ``[256, rows * K/16]``, and +the four scalars ``input_scale``, ``weight_scale_2``, ``g`` and ``alpha`` ``[256]`` FP32. +""" + +from __future__ import annotations + +import dataclasses +import json +import os +from collections import OrderedDict +from typing import Any, Mapping, Optional + +from tensorrt_llm.quantization.modelopt_config import read_modelopt_quant_config + +from .._tp4 import layout_bf16 as bf16_layout +from .._tp4.layout_bf16 import ( + BF16, + EMBED, + F32, + FINAL_NORM, + HEAD, + WORLD_SIZE, + TensorSpec, + check_config, + check_rank, +) + +CONFIG_FILE = "config.json" +QUANT_CONFIG_FILE = "hf_quant_config.json" +U8 = "U8" +E4M3 = "F8_E4M3" +#: Elements per NVFP4 block scale. +GROUP = 16 +PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +QUANT_ALGO = "NVFP4" +#: The quantizer's 128-row x 4-column scale tile. +SWIZZLE_ROWS, SWIZZLE_COLS = 128, 4 + +__all__ = [ + "BF16", + "E4M3", + "EMBED", + "F32", + "FINAL_NORM", + "GROUP", + "HEAD", + "PROJECTIONS", + "U8", + "WORLD_SIZE", + "QuantSpec", + "TensorSpec", + "check_config", + "check_partition", + "check_rank", + "check_recipe", + "stored_entries", + "swizzled_numel", + "tensor_specs", +] + + +def swizzled_numel(rows: int, cols: int) -> int: + """Bytes of a ``[rows, cols]`` block-scale matrix in the 128x4 swizzled layout (padded).""" + pad = lambda v, m: -(-v // m) * m # noqa: E731 + return pad(rows, SWIZZLE_ROWS) * pad(cols, SWIZZLE_COLS) + + +@dataclasses.dataclass(frozen=True) +class QuantSpec: + """One NVFP4 expert projection: logical geometry, four stored tensors, four-rank rule.""" + + name: str # "model.layers..mlp.experts.." + shape: tuple # logical [N, K] + layer: int + expert: int + proj: str # one of PROJECTIONS + split_dim: int # 0: output rows (gate/up), 1: input columns (down) + + @property + def family(self) -> str: + return "expert_" + self.proj.removesuffix("_proj") + + @property + def bank(self) -> str: + return f"layers.{self.layer}.experts.{self.proj.removesuffix('_proj')}" + + @property + def packed_shape(self) -> tuple: + return (self.shape[0], self.shape[1] // 2) + + @property + def scale_shape(self) -> tuple: + return (self.shape[0], self.shape[1] // GROUP) + + def stored(self) -> "OrderedDict[str, tuple]": + """The four stored tensors: name -> (dtype, shape).""" + return OrderedDict( + ( + (self.name + ".weight", (U8, self.packed_shape)), + (self.name + ".weight_scale", (E4M3, self.scale_shape)), + (self.name + ".weight_scale_2", (F32, ())), + (self.name + ".input_scale", (F32, ())), + ) + ) + + def interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of rank ``rank``'s logical slice, on whole quantization blocks.""" + check_rank(rank, world_size) + extent = self.shape[self.split_dim] + if extent % (world_size * GROUP): + raise ValueError( + f"{self.name}: dim {self.split_dim} of {extent} does not split into " + f"{world_size} whole {GROUP}-element blocks" + ) + size = extent // world_size + return self.split_dim, rank * size, (rank + 1) * size + + def packed_interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of the rank's slice of the packed U8 tensor (byte columns).""" + dim, start, stop = self.interval(rank, world_size) + return (dim, start, stop) if dim == 0 else (dim, start // 2, stop // 2) + + def scale_interval(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + """``(dim, start, stop)`` of the rank's slice of the block-scale tensor (scale columns).""" + dim, start, stop = self.interval(rank, world_size) + return (dim, start, stop) if dim == 0 else (dim, start // GROUP, stop // GROUP) + + def local_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + dim, start, stop = self.interval(rank, world_size) + shape = list(self.shape) + shape[dim] = stop - start + return tuple(shape) + + def local_packed_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + rows, cols = self.local_shape(rank, world_size) + return (rows, cols // 2) + + def local_scale_shape(self, rank: int, world_size: int = WORLD_SIZE) -> tuple: + rows, cols = self.local_shape(rank, world_size) + return (rows, cols // GROUP) + + +def tensor_specs(config: Any, layers: Optional[tuple] = None) -> "OrderedDict[str, Any]": + """Every logical tensor of the selected layers and the outer tensors. + + Keys are checkpoint names for higher-precision tensors (``TensorSpec``, stored BF16 or F32) + and projection prefixes for NVFP4 expert projections (``QuantSpec``). Geometry and + intervals are the BF16 layout's; the stored dtype of every non-expert matrix is BF16. + """ + specs: "OrderedDict[str, Any]" = OrderedDict() + for name, spec in bf16_layout.tensor_specs(config, layers).items(): + if spec.expert is not None: + prefix = name.removesuffix(".weight") + proj = prefix.rsplit(".", 1)[1] + specs[prefix] = QuantSpec( + prefix, tuple(spec.shape), spec.layer, spec.expert, proj, spec.split_dim + ) + elif spec.stored_dtype == bf16_layout.FP8: + specs[name] = dataclasses.replace(spec, stored_dtype=BF16) + else: + specs[name] = spec + return specs + + +def stored_entries(specs: Mapping[str, Any]) -> "OrderedDict[str, tuple]": + """Every stored tensor ``specs`` needs: name -> (dtype, shape).""" + out: "OrderedDict[str, tuple]" = OrderedDict() + for spec in specs.values(): + if isinstance(spec, QuantSpec): + out.update(spec.stored()) + else: + out[spec.name] = (spec.stored_dtype, tuple(spec.shape)) + return out + + +def check_partition(specs: Mapping[str, Any], world_size: int = WORLD_SIZE) -> None: + """Raise unless every split tensor's rank intervals tile its axis exactly, in rank order. + + For an expert projection the logical, packed-byte and block-scale intervals must each tile + their axis, and every logical boundary must fall on a quantization block boundary. + """ + bf16_layout.check_partition( + {n: s for n, s in specs.items() if not isinstance(s, QuantSpec)}, world_size + ) + for spec in specs.values(): + if not isinstance(spec, QuantSpec): + continue + for kind, extent, interval in ( + ("logical", spec.shape[spec.split_dim], spec.interval), + ("packed", spec.packed_shape[spec.split_dim], spec.packed_interval), + ("scale", spec.scale_shape[spec.split_dim], spec.scale_interval), + ): + cursor = 0 + for rank in range(world_size): + _, start, stop = interval(rank, world_size) + if start != cursor or stop <= start: + raise ValueError( + f"{spec.name}: rank {rank} {kind} interval [{start}, {stop}) leaves a gap " + "or overlap" + ) + cursor = stop + if cursor != extent: + raise ValueError(f"{spec.name}: {kind} intervals cover {cursor} of {extent}") + + +def _module_names(specs: Mapping[str, Any]) -> tuple: + """(quantized expert modules, higher-precision Linear modules) of ``specs``.""" + quantized, kept = [], [] + for name, spec in specs.items(): + if isinstance(spec, QuantSpec): + quantized.append(name) + elif name.endswith(".weight") and len(spec.shape) == 2 and name != EMBED: + kept.append(name.removesuffix(".weight")) + return quantized, kept + + +def check_recipe(quant_config: Any, specs: Mapping[str, Any]) -> dict: + """Raise ValueError unless ``hf_quant_config.json`` is the recipe this layout computes. + + ModelOpt NVFP4, group size 16, no KV-cache quantization, and the shared QuantConfig + exclusion semantics for every higher-precision Linear module of ``specs`` + (attention, router, dense feed-forward, head) and no expert projection. Returns a summary. + """ + if not isinstance(quant_config, Mapping): + raise ValueError("hf_quant_config.json must be a JSON object") + if quant_config.get("kv_cache_scheme") is not None: + raise ValueError("kv_cache_scheme is not supported; this target requires BF16 KV cache") + q = read_modelopt_quant_config(dict(quant_config)) + if q.get("quant_algo") != QUANT_ALGO: + raise ValueError(f"unsupported quant_algo {q.get('quant_algo')!r}; expected {QUANT_ALGO}") + if q.get("group_size") != GROUP or type(q.get("group_size")) is not int: + raise ValueError(f"unsupported group_size {q.get('group_size')!r}; expected {GROUP}") + if q.get("kv_cache_quant_algo") is not None: + raise ValueError( + f"unsupported kv_cache_quant_algo {q.get('kv_cache_quant_algo')!r}; expected none" + ) + patterns = q.get("exclude_modules") + if not isinstance(patterns, list) or not all(isinstance(p, str) for p in patterns): + raise ValueError("exclude_modules must be a list of strings") + from tensorrt_llm.models.modeling_utils import QuantConfig + + excluded = QuantConfig(exclude_modules=patterns).is_module_excluded_from_quantization + quantized, kept = _module_names(specs) + wrongly_excluded = [m for m in quantized if excluded(m)] + if wrongly_excluded: + raise ValueError(f"exclude_modules excludes quantized experts: {wrongly_excluded[:4]}") + not_excluded = [m for m in kept if not excluded(m)] + if not_excluded: + raise ValueError(f"higher-precision modules are not excluded: {not_excluded[:4]}") + return { + "quant_algo": QUANT_ALGO, + "group_size": GROUP, + "kv_cache_quant_algo": None, + "exclude_patterns": len(patterns), + "quantized_modules": len(quantized), + "excluded_linear_modules": len(kept), + } + + +def checkpoint_dir_of(pretrained_config: Any) -> Optional[str]: + """The checkpoint directory a loaded configuration came from, or None.""" + return getattr(pretrained_config, "name_or_path", None) or getattr( + pretrained_config, "_name_or_path", None + ) + + +def read_config(directory: Any, name: str = CONFIG_FILE) -> dict: + """The JSON object stored as ``name`` in checkpoint ``directory``; ValueError otherwise.""" + if not directory: + raise ValueError(f"no checkpoint directory names {name} ({directory!r})") + path = os.path.join(os.fspath(directory), name) + try: + with open(path) as f: + data = json.load(f) + except (OSError, ValueError) as exc: + raise ValueError(f"{path} is not a readable JSON file: {exc}") from exc + if not isinstance(data, dict): + raise ValueError(f"{path} must hold a JSON object") + return data + + +def read_quant_config(checkpoint_dir: Any, config: Mapping[str, Any]) -> dict: + """Use the separate ModelOpt file when present, otherwise the inline recipe.""" + if os.path.exists(os.path.join(os.fspath(checkpoint_dir), QUANT_CONFIG_FILE)): + return read_config(checkpoint_dir, QUANT_CONFIG_FILE) + inline = config.get("quantization_config") + if not isinstance(inline, dict): + raise ValueError( + "checkpoint has neither hf_quant_config.json nor an inline quantization_config" + ) + return inline + + +def check_checkpoint(checkpoint_dir: Any) -> dict: + """Raise ValueError unless this target computes the checkpoint; return its config.json. + + The one validator of what the target serves, used by routing (a refusal there is the + ``NoTarget`` reason) and by the target's construction, so the two cannot disagree: the + checkpoint's ``config.json`` is what the layout computes (``check_config``) and its + separate or inline quantization config the NVFP4 recipe of the whole model + (``check_recipe``). Reads configuration files; allocates nothing. + """ + config = read_config(checkpoint_dir) + check_config(config) + check_recipe(read_quant_config(checkpoint_dir, config), tensor_specs(config)) + return config diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/loader.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/loader.py new file mode 100644 index 000000000000..8a42e0f5510f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/loader.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Validation and bounded, byte-preserving reads of an NVFP4-expert safetensors checkpoint. + +The index/header catalog and the two-buffer tile reader are the BF16 layout's +(``SafetensorsCatalog``, ``TileReader``): metadata only at construction, at most two pinned +staging buffers of ``TILE_BYTES``, never a whole tensor, shard or model on the host. + +``validate`` checks the catalog against every stored tensor the layout needs before anything +is loaded: present, nothing else present, exact dtype and shape, and a byte span matching them. + +``read_packed_into`` copies the rank's rows (and, for a column split, its byte columns) of a +stored 2-D byte tensor -- packed U8 codes or E4M3 block-scale bytes -- into device storage +unchanged. A column split reads whole rows tile by tile and keeps its columns on the device. +Values are never decoded, rescaled or cast here. ``read_scalars`` reads FP32 scalars with one +bounded read per contiguous run. +""" + +from __future__ import annotations + +import hashlib +import struct +from typing import Mapping, Optional + +import torch + +from .._tp4.loader_bf16 import ( + TILE_BYTES, + CheckpointFormatError, + SafetensorsCatalog, + TileReader, + read_slice_into, +) +from .layout import E4M3, F32, U8 + +__all__ = [ + "TILE_BYTES", + "CheckpointFormatError", + "SafetensorsCatalog", + "TileReader", + "check_block_scales", + "read_packed_into", + "read_scalars", + "read_slice_into", + "validate", +] + +#: Scalar runs read in one request when they are at most this many bytes apart end to end. +_SCALAR_RUN_BYTES = 1 << 20 + + +class _Rows: + """A stored tensor seen as ``rows`` rows of ``row_bytes`` bytes, for ``TileReader.tiles``.""" + + __slots__ = ("name", "path", "begin", "end", "row_bytes") + + def __init__(self, name: str, path: str, begin: int, end: int, row_bytes: int): + self.name, self.path, self.begin, self.end, self.row_bytes = ( + name, + path, + begin, + end, + row_bytes, + ) + + +def validate(catalog: SafetensorsCatalog, expected: Mapping[str, tuple], *, complete=True) -> None: + """Raise CheckpointFormatError unless the stored tensors are exactly ``expected``. + + ``expected`` maps each stored name to ``(dtype, shape)`` (``layout.stored_entries``). With + ``complete`` any other stored tensor is rejected. + """ + catalog.validate_entries(expected, complete=complete) + + +def _sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() + + +@torch.no_grad() +def read_packed_into( + catalog: SafetensorsCatalog, + reader: TileReader, + name: str, + rows: tuple, + cols: Optional[tuple], + dst: torch.Tensor, + *, + digest: Optional["hashlib._Hash"] = None, +) -> None: + """Copy rows ``[rows[0], rows[1])`` (byte columns ``cols``, or all) of 2-D byte tensor ``name``. + + ``dst`` is uint8 ``[rows, cols]`` storage (any device), filled with the stored bytes + unchanged. ``digest`` (optional) is updated with the source bytes read, in tile order. + """ + entry = catalog.entries[name] + if entry.dtype not in (U8, E4M3) or len(entry.shape) != 2: + raise CheckpointFormatError(f"{name}: expected a 2-D U8/F8_E4M3 tensor, got {entry.dtype}") + row_bytes = entry.shape[1] + r0, r1 = rows + c0, c1 = (0, row_bytes) if cols is None else cols + if not (0 <= r0 < r1 <= entry.shape[0] and 0 <= c0 < c1 <= row_bytes): + raise ValueError(f"{name}: slice rows {rows} cols {cols} outside {list(entry.shape)}") + if dst.dtype != torch.uint8 or tuple(dst.shape) != (r1 - r0, c1 - c0): + raise ValueError( + f"{name}: destination {dst.dtype}{list(dst.shape)} != uint8{[r1 - r0, c1 - c0]}" + ) + view = _Rows(entry.name, entry.path, entry.begin, entry.end, row_bytes) + for t0, t1, raw in reader.tiles(view, r0, r1): + tile = raw.to(dst.device, non_blocking=True).view(t1 - t0, row_bytes) + dst[t0 - r0 : t1 - r0].copy_(tile[:, c0:c1]) + _sync(dst.device) + if digest is not None: + digest.update(f"{entry.name}:{t0}:{t1}|".encode()) + digest.update(memoryview(raw.numpy())) + + +def read_scalars(catalog: SafetensorsCatalog, reader: TileReader, names) -> dict: + """FP32 scalar tensors ``names`` -> ``{name: float}``, one bounded read per contiguous run.""" + entries = sorted((catalog.entries[n] for n in names), key=lambda e: (e.path, e.begin)) + for e in entries: + if e.dtype != F32 or tuple(e.shape) != () or e.end - e.begin != 4: + raise CheckpointFormatError(f"{e.name}: expected an F32 scalar, got {e.dtype}{e.shape}") + out = {} + limit = min(_SCALAR_RUN_BYTES, reader.tile_bytes) + i = 0 + while i < len(entries): + j = i + 1 + while ( + j < len(entries) + and entries[j].path == entries[i].path + and entries[j].end - entries[i].begin <= limit + ): + j += 1 + run = entries[i:j] + begin, end = run[0].begin, max(e.end for e in run) + view = _Rows(run[0].name, run[0].path, begin, end, end - begin) + for _, _, raw in reader.tiles(view, 0, 1): + data = bytes(raw.numpy()) + for e in run: + off = e.begin - begin + out[e.name] = struct.unpack(" None: + """Raise CheckpointFormatError on a NaN or negative E4M3 block scale (zero is allowed).""" + magnitude = scale_bytes & 0x7F + bad_nan = magnitude == 0x7F + bad_negative = ((scale_bytes & 0x80) != 0) & (magnitude != 0) + if bool((bad_nan | bad_negative).any()): + raise CheckpointFormatError(f"{name}: block scales hold NaN or negative E4M3 values") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py new file mode 100644 index 000000000000..85df467851bd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/full_model_nvfp4__sm_103__tp4/model.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam full model, NVFP4 experts / sm_103 / tp4: token ids to logits through all 52 layers. + +The packed NVFP4 checkpoint-layout component (``component``) with a +full-depth forward. Every operation is the component's own: the masked embedding with its +all-reduce, each logical layer's ``block`` (attention on TRTLLM attention and the call's +``KVCacheManagerV2`` slot for that layer, the dense SwiGLU on layers 0-2, the replicated FP32 +router and packed W4A4 experts on layers 3-51, the four sandwich norms and the depth residuals), +the final norm, the untied head with its vocabulary all-gather and the FP32 softcap. This class +adds the order and the contract of one call:: + + x = embed(input_ids) replicated [T, 5120] bf16 + for layer in 0 .. 51: x = block(layer, x) in logical order, each exactly once + logits = head(final_norm(x[rows])) fp32 [len(rows), 151936], softcapped once + +One call is one engine step of one or more requests (``_attention.check_batch_metadata``): +the step's context requests first, then its generation requests, their new tokens +concatenated into the call's T rows. A context request appends its ``n`` tokens after the +``past`` tokens the cache already holds for it (none, a reused prefix or earlier chunks), a +generation request appends one token. Each request's ``position_ids`` must be exactly +``past .. past + n - 1``. ``logit_rows`` selects the hidden rows that reach the final norm and +the head before either runs: ``None`` (default) is the last row of each request, in request +order -- the next-token logits a sampler needs -- ``"all"`` every row, or explicit row +indices. + +The cache must be a BF16 ``KVCacheManagerV2`` holding all 52 logical layers in logical order +with 2 kv heads of 128 per rank, a 512-token window on every sliding layer and the whole +sequence (``max_seq_len``) on the full-attention layers 5, 11, 17, 23, 29, 35, 41, 47 and 51 +(``attention_windows``). Input geometry, metadata and cache are validated before execution; absolute position +values are compared with the request spans only in diagnostic mode. A call with T above +``max_tokens`` is rejected (the default ``MAX_NUM_TOKENS`` matches +the supported deployment's native ``max_num_tokens``). + +Host synchronization: diagnostic mode checks position values once per call and router +agreement per sparse layer. Expert dispatch synchronizes per sparse layer in all modes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, Mapping, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.component import ( # noqa: E501 + BeamCheckpointLayoutNvfp4Sm103Tp4, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm.bindings import DataType +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from .. import _attention +from .._constants import SUPPORTED_CONFIG + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + +#: Logical decoder layers, all executed by every call. +NUM_LAYERS = SUPPORTED_CONFIG["num_hidden_layers"] +#: Tokens per call of the supported deployment (its native ``max_num_tokens``). +MAX_NUM_TOKENS = 8192 + +LogitRows = Literal["all"] | Sequence[int] | None + + +class BeamFullModelNvfp4Sm103Tp4(BeamCheckpointLayoutNvfp4Sm103Tp4): + """All 52 layers of the Beam model with packed NVFP4 experts on one of four TP ranks. + + ``config`` is the checkpoint's config.json (a mapping) and ``mapping`` this rank's + ``tensorrt_llm.mapping.Mapping`` (world size 4, TP4, PP1, MoE TP4/EP1, no attention DP); + every logical layer is held. ``load(checkpoint_dir)`` (collective) fills the weights; the + forward is usable only after every rank completed it. + """ + + def __init__( + self, + config: Mapping[str, Any], + mapping: ParallelMapping, + device: torch.device | None = None, + *, + max_tokens: int = MAX_NUM_TOKENS, + ) -> None: + super().__init__(config, mapping, layers=None, device=device, max_tokens=max_tokens) + if self.layers != tuple(range(NUM_LAYERS)): + raise ValueError(f"the full model holds logical layers 0..{NUM_LAYERS - 1}") + self.vocab_size = config["vocab_size"] + self._checked_manager = None + + def attention_windows(self, max_seq_len: int) -> list[int]: + """Per logical layer: the attention window its cache must keep (512 or ``max_seq_len``).""" + return [ + self.window if self.schedule.is_sliding(layer) else max_seq_len + for layer in range(NUM_LAYERS) + ] + + def check_cache_manager(self, manager: KVCacheManagerV2) -> None: + """Raise ValueError unless ``manager`` holds every logical layer as this model needs. + + Required: BF16, ``head_dim`` 128, 2 kv heads per rank on every layer, logical layers + 0..51 in local slots 0..51, sliding or full-context retention on sliding layers, + and the whole sequence on full-attention layers. Reads host-side manager state only. + """ + offsets = getattr(manager, "layer_offsets", None) + if not isinstance(offsets, Mapping) or dict(offsets) != { + layer: layer for layer in range(NUM_LAYERS) + }: + held = sorted(offsets) if isinstance(offsets, Mapping) else None + raise ValueError( + f"the KV cache manager must hold logical layers 0..{NUM_LAYERS - 1} in slots " + f"0..{NUM_LAYERS - 1}; it holds {held}" + ) + heads = list(getattr(manager, "num_kv_heads_per_layer", None) or []) + if ( + manager.dtype != DataType.BF16 + or manager.head_dim != self.head_dim + or heads != [self.local_kv_heads] * NUM_LAYERS + ): + raise ValueError( + f"the KV cache must be BF16 with {self.local_kv_heads} kv heads of " + f"{self.head_dim} per layer; got {manager.dtype}, head_dim {manager.head_dim}" + ) + windows = list(getattr(manager, "max_attention_window_vec", None) or []) + if len(windows) != NUM_LAYERS: + raise ValueError(f"the KV cache manager states {len(windows)} attention windows") + for layer, window in enumerate(windows): + if self.schedule.is_sliding(layer): + # The estimator uses a full-context pool, and short sequences + # normalize the sliding window to full context. The attention + # operator still applies self.window independently of retention. + ok = window is None or int(window) in (self.window, manager.max_seq_len) + else: + ok = window is None or int(window) >= manager.max_seq_len + if not ok: + kind = "sliding" if self.schedule.is_sliding(layer) else "full-attention" + raise ValueError( + f"logical layer {layer} is a {kind} layer; the cache keeps a window of " + f"{window} (max_seq_len {manager.max_seq_len})" + ) + + def _select_rows(self, logit_rows: LogitRows, spans: Sequence[tuple[int, int]]) -> list[int]: + num_tokens = sum(n for _, n in spans) + if logit_rows is None: + ends, end = [], 0 + for _, n in spans: + end += n + ends.append(end - 1) + return ends + if isinstance(logit_rows, str): + if logit_rows != "all": + raise ValueError( + f"logit_rows must be None, 'all' or row indices, got {logit_rows!r}" + ) + return list(range(num_tokens)) + rows = list(logit_rows) + if ( + not rows + or any(type(r) is not int or not 0 <= r < num_tokens for r in rows) + or len(set(rows)) != len(rows) + ): + raise ValueError( + f"logit_rows must be distinct row indices in [0, {num_tokens}), got {logit_rows!r}" + ) + return rows + + def check_call( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + attn_metadata: TrtllmAttentionMetadata, + ) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]]]: + """Validate one call; return the flat ids, the flat positions and the request spans. + + ``input_ids``: ``[T]`` or ``[1, T]`` torch.int64 ids in ``[0, vocab)`` on the model's + device, 1 <= T <= ``max_tokens``; ``attn_metadata``: one prepared step of its requests + on a manager that passes ``check_cache_manager``; ``position_ids``: an integer tensor + of T positions, each request's ``past .. past + n - 1`` in request order. Returns the + spans ``[(past, n), ...]`` of ``_attention.check_batch_metadata``. Position values are + compared with those spans only in diagnostic mode; callers otherwise supply them. + """ + if not isinstance(input_ids, torch.Tensor) or input_ids.dtype != torch.int64: + raise ValueError("input_ids must be a torch.int64 tensor") + if input_ids.dim() == 2 and input_ids.shape[0] == 1: + input_ids = input_ids[0] + if input_ids.dim() != 1 or not 1 <= input_ids.shape[0] <= self.max_tokens: + raise ValueError( + f"input_ids must be [T] or [1, T] with 1 <= T <= {self.max_tokens}, " + f"got {list(input_ids.shape)}" + ) + if input_ids.device != self.device: + raise ValueError("input_ids must be on the model's device") + num_tokens = input_ids.shape[0] + if ( + not isinstance(position_ids, torch.Tensor) + or position_ids.dtype not in (torch.int32, torch.int64) + or position_ids.numel() != num_tokens + or position_ids.device != self.device + ): + raise ValueError( + f"position_ids must be an integer tensor of {num_tokens} positions on the " + "model's device" + ) + if not isinstance(attn_metadata, TrtllmAttentionMetadata): + raise ValueError( + f"attn_metadata must be TrtllmAttentionMetadata, got {type(attn_metadata).__name__}" + ) + if attn_metadata.kv_cache_manager is None: + raise ValueError("attn_metadata needs a KV cache manager") + if attn_metadata.kv_cache_manager is not self._checked_manager: + self.check_cache_manager(attn_metadata.kv_cache_manager) + self._checked_manager = attn_metadata.kv_cache_manager + spans = _attention.check_batch_metadata( + attn_metadata, + num_tokens, + layer_idx=0, + num_kv_heads=self.local_kv_heads, + head_dim=self.head_dim, + ) + positions = position_ids.reshape(-1).to(torch.int64) + if not self.diagnostics: + return input_ids, positions, spans + expected = torch.cat([torch.arange(past, past + n, dtype=torch.int64) for past, n in spans]) + if not torch.equal(positions, expected.to(self.device, non_blocking=True)): + got, start = positions.cpu(), 0 + for i, (past, n) in enumerate(spans): + if not torch.equal(got[start : start + n], expected[start : start + n]): + break + start += n + raise ValueError( + f"request {i} of the call appends tokens {past}..{past + n - 1} after {past} " + "cached tokens; its positions must be exactly those" + ) + return input_ids, positions, spans + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + attn_metadata: TrtllmAttentionMetadata, + *, + logit_rows: LogitRows = None, + trace: list[dict] | None = None, + ) -> torch.Tensor: + """One engine step through all 52 layers; FP32 softcapped logits ``[R, vocab]``. + + ``trace``, when a list, receives one record per + executed stage in execution order (embedding with the step's request spans, each + logical layer with its kind and window, final norm, head). + """ + self._check_usable() + ids, positions, spans = self.check_call(input_ids, position_ids, attn_metadata) + num_tokens = ids.shape[0] + rows = self._select_rows(logit_rows, spans) + x = self.embed(ids) + if trace is not None: + trace.append({"stage": "embed", "tokens": num_tokens, "requests": spans}) + for layer in range(NUM_LAYERS): + x = self.block(layer, x, positions, attn_metadata, call_checked=True) + if trace is not None: + sliding = self.schedule.is_sliding(layer) + trace.append( + { + "stage": "layer", + "layer": layer, + "kv_slot": _attention.local_layer_index(attn_metadata, layer), + "mlp": "sparse" if self.schedule.is_sparse(layer) else "dense", + "attention": "sliding" if sliding else "full", + "window": self.window if sliding else attn_metadata.max_seq_len, + } + ) + selected = x if len(rows) == num_tokens and rows == list(range(num_tokens)) else x[rows] + normed = self.final_norm(selected.contiguous()) + logits = self.head(normed) + if trace is not None: + trace.append({"stage": "final_norm", "rows": rows}) + trace.append({"stage": "head", "rows": len(rows), "vocab": int(logits.shape[-1])}) + return logits + + +__all__ = [ + "MAX_NUM_TOKENS", + "NUM_LAYERS", + "BeamFullModelNvfp4Sm103Tp4", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/beam_logit_softcap.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/beam_logit_softcap.py new file mode 100644 index 000000000000..2d0f530dac5d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/beam_logit_softcap.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Final-logit tanh softcap in fp32: `tanh(logits.float() / cap) * cap`.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.logits_processor import scale_and_soft_cap_logits + +#: The head's output: every token's 151936 vocabulary logits, bf16, cap 30. +_HEAD = dict(vocab=151936, cap=30.0, dtype=torch.bfloat16) + + +class BeamLogitSoftcap(OpWrapper): + """Softcap the head projection `[T, V]`: `(logits.float() / cap).tanh() * cap`, fp32. + + The head's activation-dtype output is widened to fp32, divided by `cap`, passed + through tanh and multiplied by `cap`, each step in fp32; the fp32 result is + returned, in `(-cap, cap)` for finite inputs and exactly `+-cap` once tanh + saturates. Applied once, after the head projection and nothing else. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="head, 1 token", spec=dict(_HEAD, tokens=1)), + Cell(why="head, 17 tokens", spec=dict(_HEAD, tokens=17)), + Cell(why="head, 129 tokens: the longest context", spec=dict(_HEAD, tokens=129)), + Cell( + why="signed values from 0 to beyond saturation, -0 included", + spec=dict(_HEAD, tokens=1, saturating=True), + ), + ) + + note = """ + Not idempotent: a second application moves every nonzero logit toward zero + (`30 * tanh(30 * tanh(x / 30) / 30)` < `30 * tanh(x / 30)` for x > 0), and + omitting it changes every logit whose magnitude exceeds about 0.05. Both are + therefore visible on ordinary logits by comparing the returned values with this + expression of the captured head output, without any saturating input. + """ + + def __call__(self, logits: torch.Tensor, cap: float) -> torch.Tensor: + owned = logits.to(dtype=torch.float32, copy=True) + return scale_and_soft_cap_logits(owned, multiplier=1.0, soft_cap=cap) + + def reference(self, logits: torch.Tensor, cap: float) -> torch.Tensor: + """The same three fp32 steps, written out one by one.""" + widened = logits.to(torch.float32) + scaled = torch.div(widened, cap) + return torch.mul(torch.tanh(scaled), cap) + + def is_valid(self, logits: torch.Tensor, cap: float) -> None: + # A tensor cap broadcasts and a non-positive cap flips or removes the cap; an + # fp32 input that was rounded upstream is indistinguishable here. None raises. + assert isinstance(cap, float) and cap > 0, f"cap must be a positive float, got {cap!r}" + assert logits.dim() == 2, "logits must be [T, vocab]" + assert logits.dtype == torch.bfloat16, "logits must be the bf16 head output" + + +beam_logit_softcap = BeamLogitSoftcap() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/swiglu_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/swiglu_fp32.py new file mode 100644 index 000000000000..b5c482e08a1c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/activation/swiglu_fp32.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SwiGLU with the SiLU and the product in fp32, rounded once to the projections' dtype.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 3's experts: the gate and up projections are [T, 2560] bf16. +_EXPERT = dict(intermediate=2560, dtype=torch.bfloat16) + + +class SwigluFp32(OpWrapper): + """`silu(gate.float()) * up.float()`, rounded once to `gate.dtype`, every step returned. + + `gate` and `up` are an expert's two projections, `[T, I]` bf16 each. The steps, as the + source model computes them:: + + silu = silu(gate.float()) [T, I] fp32 + product = silu * up.float() [T, I] fp32 + act = product.to(gate.dtype) [T, I] bf16, the down projection's input + + Returns all three in a dict under those names, so a caller that needs only `act` and a + test that checks each step observe the same call. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3 expert, 1 token", spec=dict(_EXPERT, tokens=1)), + Cell(why="block 3 expert, 17 tokens", spec=dict(_EXPERT, tokens=17)), + Cell(why="block 3 expert, 129 tokens: the longest context", spec=dict(_EXPERT, tokens=129)), + ) + + note = """ + Not the shared `flashinfer_silu_and_mul` entry. That kernel also computes in fp32 and + rounds once, but it returns only the rounded activation: the fp32 SiLU and product it + computed cannot be observed, and here they are outputs of the op. + + Rounding the SiLU or the product to bf16 before the final cast is a different answer: + it rounds twice, and on inputs away from any bf16 rounding boundary the activation + moves by a bf16 unit wherever the two roundings disagree. + """ + + def __call__(self, gate: torch.Tensor, up: torch.Tensor) -> dict[str, torch.Tensor]: + silu = F.silu(gate.float()) + product = silu * up.float() + return dict(silu=silu, product=product, act=product.to(gate.dtype)) + + def reference(self, gate: torch.Tensor, up: torch.Tensor) -> dict[str, torch.Tensor]: + """Each step in fp64, rounded to fp32 where the source holds fp32, then to `gate.dtype`. + + The product of an fp32 SiLU and a bf16 `up` is exact in fp64, so `product` is the + correctly rounded fp32 product of the reference's own SiLU. + """ + silu = F.silu(gate.double()).float() + product = (silu.double() * up.double()).float() + return dict(silu=silu, product=product, act=product.to(gate.dtype)) + + def is_valid(self, gate: torch.Tensor, up: torch.Tensor) -> None: + # Mismatched shapes broadcast, and a gate that is not bf16 removes (fp32) or moves + # (fp16) the one rounding the down projection's input carries. Neither raises. + assert gate.shape == up.shape, ( + f"gate {tuple(gate.shape)} and up {tuple(up.shape)} must have one shape" + ) + assert gate.dtype == torch.bfloat16 and up.dtype == torch.bfloat16, ( + f"gate and up must be torch.bfloat16, got {gate.dtype} and {up.dtype}" + ) + + +swiglu_fp32 = SwigluFp32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_neox_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_neox_rope.py new file mode 100644 index 000000000000..4b5b952f3f71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_neox_rope.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""NeoX half-rotation RoPE on per-head q and k, in fp32 with one rounding per tensor.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 0's sliding-attention layer: 48 q heads and 8 kv heads of 128, theta 10000, factor 1. +#: Block 3's sliding-attention layer has the same geometry and RoPE parameters. +_BLOCK0 = dict( + num_heads=48, num_kv_heads=8, head_dim=128, theta=10000.0, factor=1.0, dtype=torch.bfloat16 +) + + +def _cos_sin( + positions: torch.Tensor, head_dim: int, theta: float, factor: float +) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 `[T, 1, head_dim]` cos and sin, both halves carrying the same angles.""" + exponent = torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device) + inv_freq = 1.0 / (theta ** (exponent / head_dim)) / factor + freqs = positions.float().unsqueeze(-1) * inv_freq + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().unsqueeze(-2), emb.sin().unsqueeze(-2) + + +def _rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + x32 = x.float() + first, second = x32.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return (x32 * cos + rotated * sin).to(x.dtype) + + +class BeamNeoxRope(OpWrapper): + """Rotate `q` `[T, Hq, D]` and `k` `[T, Hkv, D]` by their token's position. + + Frequencies are `1 / theta ** (2i / D) / factor` for `i < D / 2`; angles are + `position * frequency` in fp32, and the two halves of the head dimension rotate + together -- `x * cos + rotate_half(x) * sin`, NeoX order, not interleaved pairs. + Each output is rounded once to its input's dtype. `q` and `k` arrive already + normalized and rounded by their per-head RMSNorm; this op does not normalize, and + `v` never passes through it. Returns two new tensors. + """ + + ARCHS = frozenset({Arch.SM_103}) + + # The three fresh-context lengths the dense block runs: a single token (position 0, + # the identity rotation), a short and a long context. + CELLS: tuple[Cell, ...] = ( + Cell( + why="block 0, 1 token: position 0 is the identity rotation", + spec=dict(_BLOCK0, tokens=1), + ), + Cell(why="block 0, 17 tokens: positions 0..16", spec=dict(_BLOCK0, tokens=17)), + Cell( + why="block 0, 129 tokens: the longest context the block accepts", + spec=dict(_BLOCK0, tokens=129), + ), + ) + + note = """ + Native torch, not a fused kernel: the source model rounds the per-head norm + output to the activation dtype before promoting it again for the rotation, and + no fused norm-plus-RoPE op in this tree has that intermediate rounding (the + catalog's fused_qk_norm_rope keeps the normalized values in fp32 through the + rotation). `q` and `k` share one cos/sin evaluation. + + `positions` is per token, not per sequence: a fresh context passes + `arange(T)`. No scaling other than the linear `factor` is modelled. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + theta: float, + factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + cos, sin = _cos_sin(positions, q.shape[-1], theta, factor) + return _rotate(q, cos, sin), _rotate(k, cos, sin) + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + theta: float, + factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + """The rotation written per half: `(x1 c - x2 s, x2 c + x1 s)`, fp32, one cast.""" + half = q.shape[-1] // 2 + index = torch.arange(half, dtype=torch.float32, device=positions.device) + inv_freq = 1.0 / (theta ** (2.0 * index / q.shape[-1])) / factor + angle = positions.float()[:, None] * inv_freq[None, :] + c, s = angle.cos()[:, None, :], angle.sin()[:, None, :] + + def rotate(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.float()[..., :half], x.float()[..., half:] + return torch.cat((x1 * c - x2 * s, x2 * c + x1 * s), dim=-1).to(x.dtype) + + return rotate(q), rotate(k) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + theta: float, + factor: float, + ) -> None: + # A [1] or [B, 1] positions tensor broadcasts across every token and rotates + # them all by one angle; a float tensor is accepted and rotates by fractional + # positions. Neither raises. + assert q.dim() == 3 and k.dim() == 3, "q and k must be [T, heads, head_dim]" + assert positions.shape == (q.shape[0],) and k.shape[0] == q.shape[0], ( + f"positions must be [T] with T = {q.shape[0]} tokens; got {tuple(positions.shape)} " + f"for k of {k.shape[0]} tokens" + ) + assert not positions.is_floating_point(), "positions must be integer token positions" + assert q.shape[-1] == k.shape[-1] and q.shape[-1] % 2 == 0, ( + "q and k must share an even head_dim" + ) + + +beam_neox_rope = BeamNeoxRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_nope_query_scale.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_nope_query_scale.py new file mode 100644 index 000000000000..033e47f15aa8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_nope_query_scale.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Position-dependent query scale of a full-attention layer without position embedding.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 5's full-attention layer: 48 q heads of 128, bf16. +_BLOCK5 = dict(num_heads=48, head_dim=128, dtype=torch.bfloat16) + + +class BeamNopeQueryScale(OpWrapper): + """Scale `q` `[T, H, D]` per token: `q * exp(0.1 * log1p((position + 1) / 2048))`. + + The scale is computed in fp32 from the token's position, the product is the + activation-dtype `q` promoted against that fp32 scale, and the result is rounded + once to `q`'s dtype. It replaces the rotation of a full-attention layer that has no + position embedding: it applies to the normalized `q` only, never to `k` or `v`, + and it is not the attention's `1 / sqrt(D)` softmax scale. Returns a new tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 5, 1 token: position 0 alone", spec=dict(_BLOCK5, tokens=1)), + Cell(why="block 5, 17 tokens: positions 0..16", spec=dict(_BLOCK5, tokens=17)), + Cell( + why="block 5, 129 tokens: the longest context the block accepts", + spec=dict(_BLOCK5, tokens=129), + ), + ) + + note = """ + The scale grows from 1.00005 at position 0 to 1.0063 at position 128, so most + bf16 query values are unchanged by it at these lengths: only values close to a + rounding boundary move, and position 0 moves almost none. A test that has to see + the scale therefore compares the rounded result exactly at nonzero positions. + """ + + def __call__(self, q: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + scale = torch.exp(0.1 * torch.log1p((positions.float() + 1) / 2048)) + return (q * scale[:, None, None]).to(q.dtype) + + def reference(self, q: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Each step in fp32, the product of the widened `q` and the scale, one cast.""" + position = positions.to(torch.float32) + ratio = torch.div(position + 1.0, 2048.0) + scale = torch.exp(torch.mul(torch.log1p(ratio), 0.1)) + return (q.to(torch.float32) * scale.view(-1, 1, 1)).to(q.dtype) + + def is_valid(self, q: torch.Tensor, positions: torch.Tensor) -> None: + # A [1] positions tensor broadcasts one scale over every token, and a float + # tensor scales by fractional positions. Neither raises. + assert q.dim() == 3, "q must be [T, heads, head_dim]" + assert positions.shape == (q.shape[0],), ( + f"positions must be [T] with T = {q.shape[0]} tokens, got {tuple(positions.shape)}" + ) + assert not positions.is_floating_point(), "positions must be integer token positions" + + +beam_nope_query_scale = BeamNopeQueryScale() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_output_gate.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_output_gate.py new file mode 100644 index 000000000000..5b87f8713ee1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/attention/beam_output_gate.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-head sigmoid output gate with a per-head bias, in the activation dtype.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 0's gate: 48 heads of 128, bf16. Block 3's gate has the same geometry. +_BLOCK0 = dict(num_heads=48, head_dim=128, dtype=torch.bfloat16) + + +class BeamOutputGate(OpWrapper): + """`attention_output * sigmoid(gate + gate_bias)`, all `[T, H, D]` but `gate_bias` `[H]`. + + `gate_bias` is broadcast over tokens and the head dimension. The add, the sigmoid + and the product are separate activation-dtype operations, each rounded to that + dtype, which is what the source model computes; an fp32 chain rounded once is a + different result. Returns a new tensor. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 0, 1 token", spec=dict(_BLOCK0, tokens=1)), + Cell(why="block 0, 17 tokens", spec=dict(_BLOCK0, tokens=17)), + Cell( + why="block 0, 129 tokens: the longest context the block accepts", + spec=dict(_BLOCK0, tokens=129), + ), + ) + + note = """ + The gate is the second q-width projection of the attention input, not a slice of + q, and it is applied to the attention core's output before the output projection. + Omitting the bias, or using one value for every head, still returns a plausible + gated output; only a head-varying bias makes either mistake visible, which is why + the block's tests use one. + """ + + def __call__( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> torch.Tensor: + return attention_output * torch.sigmoid(gate + gate_bias.unsqueeze(-1)) + + def reference( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> torch.Tensor: + """Each step in fp32, rounded to the activation dtype before the next.""" + dtype = attention_output.dtype + biased = (gate.float() + gate_bias.float().unsqueeze(-1)).to(dtype) + weight = torch.sigmoid(biased.float()).to(dtype) + return (attention_output.float() * weight.float()).to(dtype) + + def is_valid( + self, + attention_output: torch.Tensor, + gate: torch.Tensor, + gate_bias: torch.Tensor, + ) -> None: + # Broadcasting accepts a bias of the wrong length-1 shape and a gate of fewer + # dims; mixed dtypes promote and change every rounding. None of these raise. + assert attention_output.dim() == 3 and gate.shape == attention_output.shape, ( + "attention_output and gate must both be [T, heads, head_dim]" + ) + assert gate_bias.shape == (attention_output.shape[1],), ( + f"gate_bias must be [heads] = [{attention_output.shape[1]}], " + f"got {tuple(gate_bias.shape)}" + ) + assert attention_output.dtype == gate.dtype == gate_bias.dtype, ( + "attention_output, gate and gate_bias must share one dtype" + ) + + +beam_output_gate = BeamOutputGate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_dispatch_combine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_dispatch_combine.py new file mode 100644 index 000000000000..760333808b08 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_dispatch_combine.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Dispatch tokens to their selected experts and combine the weighted outputs in fp32.""" + +from __future__ import annotations + +from typing import Callable + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Block 3's MoE: 5120 hidden, 256 experts, 8 selected per token, routed scaling factor 2.5. +_BLOCK3 = dict(hidden_size=5120, num_experts=256, top_k=8, routed_scaling_factor=2.5) + + +class BeamDispatchCombine(OpWrapper): + """Run every selected expert on its tokens and sum the weighted outputs per token. + + `hidden_states [T, H]` bf16, `expert_ids [T, K]` int64 and `mixing_weights [T, K]` fp32 + are one routing call's tokens and result: slot `j` of row `t` sends token `t` to expert + `expert_ids[t, j]` with weight `mixing_weights[t, j]`. `expert_provider(e, rows)` is the + expert computation; it returns expert `e`'s unweighted bf16 output, one row per row of + `rows`, in their order. The steps, as the source model computes them:: + + assignments = the T*K slots in token-major order, each with its token and weight + for e in the selected experts, ascending: + tokens = tokens of the assignments to e, ascending [n] int64 + out = expert_provider(e, hidden_states[tokens]) [n, H] bf16 + route = (out.float() * weights[:, None]).to(bfloat16) [n, H] each route rounded once + acc[tokens] += route.float() [T, H] fp32 + output = (acc * routed_scaling_factor).to(bfloat16) [T, H] bf16 + + The provider is called once per selected expert, with at least one row; an expert no + token selected is never called. Returns `output`. `expert_ids`, `mixing_weights` and + `hidden_states` are read, never written; the weights are used as given, without + re-normalization or bias. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3, 1 token: eight experts of one row each", spec=dict(_BLOCK3, tokens=1)), + Cell( + why="block 3, 17 tokens: experts shared by several tokens", + spec=dict(_BLOCK3, tokens=17), + ), + Cell( + why="block 3, 129 tokens: almost every expert selected, the longest context", + spec=dict(_BLOCK3, tokens=129), + ), + ) + + note = """ + Every weighted route is rounded to bf16 before it is added, and the sum is fp32. On + finite values that is a different answer from an fp32 route, a bf16 product, a bf16 + accumulator or a sum rounded only at the end, and from folding the scaling factor into + the weights: each moves outputs by bf16 units. A token's eight routes are added in + ascending expert order, one expert at a time; within an expert a token appears once, + so no addition order is left to the hardware. Adding them in slot order, or all at + once with one scattered add, changes the fp32 sum on elements where routes cancel. + + The selected experts and their tokens are host metadata: `unique()` and `tolist()` + synchronize with the device twice for the expert list, and the two boolean selections + twice per selected expert -- 512 synchronizations when 255 experts are selected. Token + rows, weights, products, roundings, sums and the scaling run on the device of + `hidden_states`. + + Not the fused MoE runners of the catalog: those own the expert weights and activation + (and their quantization) as well as the combination, so they cannot call an expert + given as a function, and they combine without the per-route bf16 rounding. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: Callable[[int, torch.Tensor], torch.Tensor], + routed_scaling_factor: float, + ) -> torch.Tensor: + tokens, top_k = expert_ids.shape + flat_ids = expert_ids.reshape(-1) + flat_weights = mixing_weights.reshape(-1) + token_ids = torch.arange(tokens, device=hidden_states.device).repeat_interleave(top_k) + acc = torch.zeros( + tokens, hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device + ) + for expert_id in flat_ids.unique().tolist(): + assigned = flat_ids == expert_id + expert_tokens = token_ids[assigned] + out = expert_provider(expert_id, hidden_states[expert_tokens]) + route = (out.float() * flat_weights[assigned].unsqueeze(-1)).to(torch.bfloat16) + acc.index_add_(0, expert_tokens, route.float()) + return (acc * routed_scaling_factor).to(torch.bfloat16) + + def reference( + self, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + mixing_weights: torch.Tensor, + expert_provider: Callable[[int, torch.Tensor], torch.Tensor], + routed_scaling_factor: float, + ) -> torch.Tensor: + """The combination in plain torch, expert by expert over every id up to the largest. + + Each expert's (token, slot) pairs are looked up in the routing table directly, so + nothing is shared with `__call__`'s flattening and grouping. The product of a bf16 + output and an fp32 weight is exact in fp64; rounding it to fp32 is the fp32 product, + and that is rounded to bf16 as the source does. The bf16 routes are then added in + fp32 in ascending expert order, row by row, and the sum is scaled and rounded once. + Rounding each route is the source's semantics, not a loss of reference accuracy. + """ + acc = torch.zeros( + hidden_states.shape[0], + hidden_states.shape[1], + dtype=torch.float32, + device=hidden_states.device, + ) + for expert_id in range(int(expert_ids.max()) + 1): + token_slot = (expert_ids == expert_id).nonzero() + if token_slot.shape[0] == 0: + continue + token, slot = token_slot[:, 0], token_slot[:, 1] + out = expert_provider(expert_id, hidden_states[token]) + weight = mixing_weights[token, slot].to(torch.float64) + route = (out.to(torch.float64) * weight[:, None]).to(torch.float32).to(torch.bfloat16) + acc[token] = acc[token] + route.to(torch.float32) + return (acc * routed_scaling_factor).to(torch.bfloat16) + + +beam_dispatch_combine = BeamDispatchCombine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_grouped_topk_router.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_grouped_topk_router.py new file mode 100644 index 000000000000..5ef5bb8cffc4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/moe/beam_grouped_topk_router.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Grouped sigmoid top-k routing with a selection-only correction bias, all in fp32.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm + +#: Block 3's router: 256 experts in 8 groups of 32, 4 groups kept, 8 experts per token. +_BLOCK3 = dict(hidden_size=5120, num_experts=256, n_group=8, topk_group=4, top_k=8, eps=1e-20) + + +class BeamGroupedTopkRouter(OpWrapper): + """Route `hidden_states [T, hidden]` with `weight [E, hidden]` and `correction_bias [E]`. + + Every step is fp32. The projection uses explicit FP32 cuBLAS computation; + autocast is disabled without changing the caller's global matmul policy:: + + logits = hidden_states @ weight^T [T, E] + scores = sigmoid(logits) unbiased + selection_scores = scores + correction_bias used only to choose + group_scores = sum of the two largest selection_scores in each of n_group + contiguous groups [T, n_group] + groups = top-topk_group group_scores, descending [T, topk_group] + masked_selection_scores = selection_scores, -inf outside those groups + ids = top-top_k masked_selection_scores, descending [T, top_k] int64 + weights = scores[ids] / (sum(scores[ids]) + eps) [T, top_k] + + Returns every tensor above in a dict under those names, so a caller that needs only + `ids` and `weights` and a test that checks each step observe the same call. The + correction bias decides the selection and never reaches `weights`; no routed scaling + factor is applied. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block 3, 1 token", spec=dict(_BLOCK3, tokens=1)), + Cell(why="block 3, 17 tokens", spec=dict(_BLOCK3, tokens=17)), + Cell( + why="block 3, 129 tokens: several hundred decision margins", + spec=dict(_BLOCK3, tokens=129), + ), + ) + + note = """ + Ties follow `torch.topk(..., sorted=True)` and nothing else: there is no secondary + index rule. Measured on sm_103 at these shapes, a tie straddling the 4th/5th group or + the 8th/9th expert boundary keeps the lower index, and tied values inside the + selection are not in index order -- a row of 256 equal scores selects groups + [1, 0, 2, 3] and experts [7, 6, 4, 5, 1, 0, 2, 3]. That is the source model's order, + which is why the selection is not delegated to a kernel with its own tie rule. + + The fused grouped top-k op (`noaux_tc_op`) is not this computation: it takes logits, + uses a tanh-form sigmoid that saturates where this one does not, renormalizes in + fp64, applies the routed scaling factor itself and breaks ties by lower expert index. + `reference` keeps the source's `topk` selections because no other formulation + reproduces their order on ties. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> dict[str, torch.Tensor]: + with torch.autocast("cuda", enabled=False): + logits = cublas_mm( + hidden_states.float().contiguous(), weight.contiguous().t(), out_dtype=torch.float32 + ) + scores = logits.sigmoid() + selection_scores = scores + correction_bias + tokens, num_experts = selection_scores.shape + per_group = num_experts // n_group + grouped = selection_scores.view(tokens, n_group, per_group) + group_scores = grouped.topk(2, dim=-1).values.sum(dim=-1) + groups = group_scores.topk(topk_group, dim=-1, sorted=True).indices + group_mask = torch.zeros_like(group_scores).scatter_(1, groups, 1) + expert_mask = ( + group_mask.unsqueeze(-1) + .expand(tokens, n_group, per_group) + .reshape(tokens, num_experts) + ) + masked_selection_scores = selection_scores.masked_fill( + ~expert_mask.bool(), float("-inf") + ) + ids = masked_selection_scores.topk(top_k, dim=-1, sorted=True).indices + selected = scores.gather(-1, ids) + weights = selected / (selected.sum(dim=-1, keepdim=True) + eps) + return dict( + logits=logits, + scores=scores, + selection_scores=selection_scores, + group_scores=group_scores, + groups=groups, + masked_selection_scores=masked_selection_scores, + ids=ids, + weights=weights, + ) + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> dict[str, torch.Tensor]: + """The routing in plain torch, with FP32 reference precision set by the test caller. + + Exclusion is written as group membership and each group's score as the first two + entries of a descending sort; both are exact, so the result is the source routing's + bit for bit. The two selections stay `topk(..., sorted=True)`, the source's own + call, because the order of tied scores is whatever that call returns. + """ + with torch.autocast("cuda", enabled=False): + logits = F.linear(hidden_states.to(torch.float32), weight.to(torch.float32)) + scores = torch.sigmoid(logits) + selection_scores = scores + correction_bias.to(torch.float32) + tokens, num_experts = selection_scores.shape + per_group = num_experts // n_group + ranked = selection_scores.reshape(tokens, n_group, per_group).sort( + dim=-1, descending=True + ) + group_scores = ranked.values[..., :2].sum(dim=-1) + groups = torch.topk(group_scores, topk_group, dim=-1, sorted=True).indices + expert_group = torch.arange(num_experts, device=logits.device) // per_group + eligible = (expert_group.view(1, num_experts, 1) == groups.unsqueeze(1)).any(dim=-1) + masked_selection_scores = torch.where( + eligible, selection_scores, torch.full_like(selection_scores, float("-inf")) + ) + ids = torch.topk(masked_selection_scores, top_k, dim=-1, sorted=True).indices + selected = torch.take_along_dim(scores, ids, dim=-1) + weights = selected / (selected.sum(dim=-1, keepdim=True) + eps) + return dict( + logits=logits, + scores=scores, + selection_scores=selection_scores, + group_scores=group_scores, + groups=groups, + masked_selection_scores=masked_selection_scores, + ids=ids, + weights=weights, + ) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + correction_bias: torch.Tensor, + n_group: int, + topk_group: int, + top_k: int, + eps: float, + ) -> None: + # A [1, T, hidden] batch projects fine, and the group view then folds its leading + # dimension away -- silently when T == 1. + assert hidden_states.dim() == 2, ( + f"hidden_states must be [tokens, hidden], got {tuple(hidden_states.shape)}" + ) + # A length-1 bias broadcasts over every expert, and a wider one promotes every + # selection score out of fp32; neither raises. + assert correction_bias.shape == (weight.shape[0],), ( + f"correction_bias must be [{weight.shape[0]}], got {tuple(correction_bias.shape)}" + ) + assert correction_bias.dtype == torch.float32, ( + f"correction_bias must be torch.float32, got {correction_bias.dtype}" + ) + num_experts = weight.shape[0] + assert num_experts % n_group == 0 and num_experts // n_group >= 2, ( + f"{num_experts} experts do not form {n_group} groups of at least two" + ) + # With fewer eligible experts than top_k, the expert top-k returns excluded ones. + assert 0 < topk_group <= n_group and 0 < top_k <= topk_group * (num_experts // n_group), ( + f"top_k={top_k} cannot be chosen from {topk_group} of {n_group} groups" + ) + + +beam_grouped_topk_router = BeamGroupedTopkRouter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/beam_rmsnorm_fp32.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/beam_rmsnorm_fp32.py new file mode 100644 index 000000000000..6c1a635fc0db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/norm/beam_rmsnorm_fp32.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RMS normalization over the last dim in the source model's fp32 op order, rounded once.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The Beam decoder block's input norm: one 5120-wide row per token, bf16, eps 1e-6. +_BLOCK_INPUT = dict(width=5120, dtype=torch.bfloat16, eps=1e-6) + + +class BeamRmsnormFp32(OpWrapper): + """`x / sqrt(mean(x^2, -1) + eps) * weight`, one eager torch op per source step. + + `x` is `[T, D]` bf16 and `weight` `[D]` bf16. The steps, as the source model computes + them:: + + x32 = x.float() [T, D] fp32 + rms = x32.pow(2).mean(-1, keepdim=True) [T, 1] fp32 + .add(eps).rsqrt() + normed = x32 * rms [T, D] fp32 + out = (normed * weight).to(x.dtype) [T, D] bf16, weight promoted + + Returns `out`, a new tensor; `x` is not mutated. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="block input norm, 1 token", spec=dict(_BLOCK_INPUT, tokens=1)), + Cell(why="block input norm, 17 tokens", spec=dict(_BLOCK_INPUT, tokens=17)), + Cell( + why="block input norm, 129 tokens: the longest context", + spec=dict(_BLOCK_INPUT, tokens=129), + ), + ) + + note = """ + Not the catalog `flashinfer_rmsnorm` entry. That kernel also reduces in fp32 and rounds + once, but it sums the squares in its own order and computes its own reciprocal square + root. On identical bf16 inputs its output differs from the source model's by one bf16 + unit wherever the fp32 product lands near a rounding boundary. The difference is a few + elements per 129 tokens of real hidden states. Every projection that reads the normalized + row then carries it. + + This entry issues the source's torch ops in the source's order. Its result equals the + source's when both run the same torch reduction kernel for the same row width. The + entry does not guard against a different torch build that changes that kernel's + summation order. + """ + + def __call__(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(eps).rsqrt() + normed = x32 * rms + return (normed * weight).to(x.dtype) + + def reference(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Statistics and products in fp64, rounded once to `x.dtype`.""" + x64 = x.double() + rms = torch.rsqrt(x64.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x64 * rms * weight.double()).to(x.dtype) + + def is_valid(self, x: torch.Tensor, weight: torch.Tensor, eps: float) -> None: + # A weight of another dtype or width broadcasts or promotes without raising, and an + # fp32 input removes the one rounding the projections' input carries. + assert x.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16, ( + f"x and weight must be torch.bfloat16, got {x.dtype} and {weight.dtype}" + ) + assert weight.dim() == 1 and weight.shape[0] == x.shape[-1], ( + f"weight must be [{x.shape[-1]}], got {tuple(weight.shape)}" + ) + + +beam_rmsnorm_fp32 = BeamRmsnormFp32() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/beam_depth_residual.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/beam_depth_residual.py new file mode 100644 index 000000000000..46ab9d40dc7f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/beam/ops/residual/beam_depth_residual.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Depth-scaled residual add, in the activation dtype.""" + +from __future__ import annotations + +import math + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: A block's two adds as (branch, residual) scales: (1, sqrt(n - 1)) / sqrt(n) for the n-th +#: add of the residual stream. Block i's attention add is n = 2(i + 1) and its feed-forward +#: add the next one: n = 2 and 3 for block 0, n = 8 and 9 for block 3. +_ATTENTION_ADD = (1.0 / math.sqrt(2.0), math.sqrt(1.0) / math.sqrt(2.0)) +_FEED_FORWARD_ADD = (1.0 / math.sqrt(3.0), math.sqrt(2.0) / math.sqrt(3.0)) +_BLOCK3_ATTENTION_ADD = (1.0 / math.sqrt(8.0), math.sqrt(7.0) / math.sqrt(8.0)) +_BLOCK3_FEED_FORWARD_ADD = (1.0 / math.sqrt(9.0), math.sqrt(8.0) / math.sqrt(9.0)) + + +def depth_scales(layer_idx: int, is_attention: bool) -> tuple[float, float]: + """(branch, residual) scales of a block's attention or feed-forward residual add. + + The n-th add of the residual stream scales the branch by 1 / sqrt(n) and the stream + by sqrt(n - 1) / sqrt(n); block i's attention add is the 2(i + 1)-th, its + feed-forward add the next one. + """ + num_adds = 2 * (layer_idx + 1) + (0 if is_attention else 1) + denom = math.sqrt(num_adds) + return 1.0 / denom, math.sqrt(num_adds - 1) / denom + + +_HIDDEN = dict(hidden=5120, dtype=torch.bfloat16) + + +class BeamDepthResidual(OpWrapper): + """`branch * branch_scale + residual * residual_scale`, a new tensor. + + Three activation-dtype operations -- two scalings and the add -- each rounded to + that dtype, as the source model computes them. The scales are Python floats and + reach the kernel as fp32. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="block 0 attention add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_ATTENTION_ADD), + ), + Cell( + why="block 0 feed-forward add, 1 token", + spec=dict(_HIDDEN, tokens=1, scales=_FEED_FORWARD_ADD), + ), + Cell( + why="block 0 feed-forward add, 129 tokens: the longest context", + spec=dict(_HIDDEN, tokens=129, scales=_FEED_FORWARD_ADD), + ), + Cell( + why="block 3 attention add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_BLOCK3_ATTENTION_ADD), + ), + Cell( + why="block 3 feed-forward add, 17 tokens", + spec=dict(_HIDDEN, tokens=17, scales=_BLOCK3_FEED_FORWARD_ADD), + ), + ) + + note = """ + Not fused with the norm before it: the source rounds the norm output, then each + scaled term, then the sum. A fused fp32 add-and-scale rounds once and is a + different result. + """ + + def __call__( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> torch.Tensor: + return branch * branch_scale + residual * residual_scale + + def reference( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> torch.Tensor: + """Each product and the sum in fp32, rounded to the activation dtype in turn.""" + dtype = branch.dtype + scaled_branch = (branch.float() * branch_scale).to(dtype) + scaled_residual = (residual.float() * residual_scale).to(dtype) + return (scaled_branch.float() + scaled_residual.float()).to(dtype) + + def is_valid( + self, + branch: torch.Tensor, + residual: torch.Tensor, + branch_scale: float, + residual_scale: float, + ) -> None: + # Broadcasting accepts mismatched shapes, mixed dtypes promote, and a tensor + # scale in the activation dtype rounds the coefficient itself. None raise. + assert branch.shape == residual.shape, ( + f"branch {tuple(branch.shape)} and residual {tuple(residual.shape)} must match" + ) + assert branch.dtype == residual.dtype, "branch and residual must share one dtype" + assert isinstance(branch_scale, float) and isinstance(residual_scale, float), ( + "scales must be Python floats" + ) + + +beam_depth_residual = BeamDepthResidual() diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 22d315ee0f0e..d07a2e30f9e2 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -148,6 +148,10 @@ l0_b300: - unittest/_torch/modeling_v2/test_modeling_v2_mnnvl_workspace.py TIMEOUT (30) - unittest/_torch/modeling_v2/test_modeling_v2_no_stale_claims.py TIMEOUT (30) - unittest/_torch/modeling_v2/activation TIMEOUT (30) + - unittest/_torch/modeling_v2/models/beam/ops TIMEOUT (30) + - unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_cublas_mm.py TIMEOUT (30) + - unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_nvfp4_gemm.py TIMEOUT (30) + - unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_fp4_quantize.py TIMEOUT (30) - unittest/_torch/modeling_v2/attention TIMEOUT (30) - unittest/_torch/modeling_v2/gemm TIMEOUT (30) - unittest/_torch/modeling_v2/graph TIMEOUT (30) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 71abe4b3583e..3320df31638a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -82,6 +82,7 @@ l0_cpu: # compute capability, no device. The default targets' shard tests: small CPU # tensors and the CPU block_scale_interleave. The rest of unittest/_torch/modeling_v2 # needs a GPU. + - unittest/_torch/modeling_v2/models/beam/test_component_contracts.py - unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py - unittest/_torch/modeling_v2/test_modeling_v2_contracts.py - unittest/_torch/modeling_v2/test_modeling_v2_native_loader.py diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index b534b1a08b47..96e33a25fe78 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -5668,6 +5668,7 @@ def _op_reference( causal: bool = True, e4m3: bool = False, layer: Optional[int] = None, + window: Optional[int] = None, nvfp4: bool = False, ) -> torch.Tensor: """`op.reference` over the history the env mirrored for these requests. @@ -5699,6 +5700,7 @@ def _op_reference( v_head_dim=d, softmax_scale=1.0 / math.sqrt(d), causal=causal, + window=window, e4m3_inputs=e4m3, e4m3_probabilities=e4m3, ) @@ -7242,10 +7244,11 @@ def _cell_paged(cell, arm) -> None: tokens_per_block=spec["tokens_per_block"], num_blocks=spec["num_blocks"], max_batch=spec["max_num_requests"], - max_blocks_per_seq=16, + max_blocks_per_seq=spec.get("max_blocks_per_seq", 16), pool_dtype=torch.float8_e4m3fn if fp8 else torch.bfloat16, quant_mode=spec["quant_mode"], kv_scaling_factor=spec["kv_scale"] if fp8 or nvfp4 else None, + attention_window_size=spec["window"], ) cached, new, nc = spec["cached"], spec["new"], spec["num_contexts"] p, paged = spec["predicted_tokens_per_seq"], spec["use_paged_context_fmha"] @@ -7283,7 +7286,10 @@ def _cell_paged(cell, arm) -> None: assert not (nvfp4 and not paged), "an NVFP4 cell runs the paged-context path" e4m3 = (fp8 or nvfp4) and (paged or nc == 0) causal = spec["mask_type"] == MASK_CAUSAL - _compare(out, _op_reference(env, qkv, list(new), rids, causal=causal, e4m3=e4m3, nvfp4=nvfp4)) + ref = _op_reference( + env, qkv, list(new), rids, causal=causal, e4m3=e4m3, window=spec["window"], nvfp4=nvfp4 + ) + _compare(out, ref) for rid in rids: env.check_cache(rid) env.check_unwritten_pool_zero() diff --git a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py index 307f6e3e700b..9e8d651f0787 100644 --- a/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py +++ b/tests/unittest/_torch/modeling_v2/gemm/test_modeling_v2_nvfp4_gemm.py @@ -80,7 +80,15 @@ def test_certified_cells(cell) -> None: alpha = (1.0 / (act_gs * weight_gs)).float().reshape(1) with validating(op): - out = op(act, weight, act_sf, weight_sf, alpha, torch.bfloat16) + out = op( + act, + weight, + act_sf, + weight_sf, + alpha, + torch.bfloat16, + allowed_backends=spec.get("backend", "cutlass,cublaslt,cuda_core"), + ) assert out.shape == (spec["m"], spec["n"]) _compare(out, op.reference(act, weight, act_sf, weight_sf, alpha, torch.bfloat16)) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/README.md b/tests/unittest/_torch/modeling_v2/models/beam/_components/README.md new file mode 100644 index 000000000000..a61700b34aaf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/README.md @@ -0,0 +1,16 @@ + + +# Local bring-up components + +These TP1 components are checkout-only test helpers, not supported serving APIs or a +reference for TP4 bitwise results. Their names describe the isolated component; device and geometry contracts are +documented in each helper. They do not register ModelingV2 targets. The six-layer prefill helper also +supports the local cached-generation tests. Use `LLM()` or `trtllm-serve` for serving. + +Run tests through pytest with the repository's test import plugin. Existing L0 selectors +retain the small single-GPU tests; the large component suites remain manual-only. + +TP1 normalization and activation use their FlashInfer contracts; the TP4 input norm +and dense MLP use explicit FP32 Beam operators, so BF16 rounding can differ. TP1 also +accepts a narrower RoPE schema. Compare each helper against its declared contract, +not against the TP4 model as a bitwise reference. diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/__init__.py new file mode 100644 index 000000000000..ccaaecb91ed8 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Checkout-only Beam bring-up components; never installed with the runtime.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/__init__.py new file mode 100644 index 000000000000..10b0e507d6c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder block 0 (dense, sliding attention) / sm_103 / tp1.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/block.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/block.py new file mode 100644 index 000000000000..98722d6c6d6a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/dense_block/block.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder blocks 0-2 / sm_103 / tp1: one dense, sliding-attention block, constructed directly. + +This TP1 component uses FlashInfer normalization/activation kernels; see ``../README.md`` +for its numerical differences from the TP4 serving model. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration for one of the dense blocks +0, 1 or 2, loads the block's fifteen tensors and runs one step of one BF16 request through +it -- a context step after any cached prefix, or a one-token generation step -- with the +attention batch state prepared on a ``KVCacheManagerV2``. The forward is +assembled from catalog entries, with the attention sublayer shared with the sparse block +component in the production family ``_attention.py``, and it does not call the built-in Beam modeling. + +The block, with block 0's residual coefficients (block i's are in ``depth_scales``):: + + h = rmsnorm(x, input_layernorm) + q, k, v, g = h @ q_proj^T, h @ k_proj^T, h @ v_proj^T, h @ gate_proj^T + q, k = rmsnorm(q, q_norm), rmsnorm(k, k_norm) per head, rounded to bf16 + q, k = rope(q, k) fp32 NeoX rotation, rounded + a = causal attention(q, k, v), window 512, scale 1/sqrt(128), 6 q heads per kv head + a = a * sigmoid(g + gate_bias) per head, bf16 steps + x = rmsnorm(a @ o_proj^T, post_attention_layernorm) / sqrt(2) + x / sqrt(2) + h = rmsnorm(x, pre_feedforward_layernorm) + m = (silu(h @ gate^T) * (h @ up^T)) @ down^T silu and product in fp32 + out = rmsnorm(m, post_feedforward_layernorm) / sqrt(3) + x * sqrt(2) / sqrt(3) + +RMSNorm and the dense activation follow their FlashInfer operator contracts. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( # noqa: E501 + flashinfer_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.residual.beam_depth_residual import ( + beam_depth_residual, + depth_scales, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "intermediate_size": 20480, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "attention_bias": False, + "attention_dropout": 0.0, + "sliding_window": 512, +} +DEFAULT_LAYER_IDX = 0 +#: The dense blocks of the checkpoint's schedule; ``layer_idx`` fixes the residual coefficients. +SUPPORTED_LAYER_INDICES = (0, 1, 2) +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The most tokens one call may append: the six-layer model's sequence capacity. Longer +#: calls are rejected; the request's whole KV length is bounded by its cache's capacity. +MAX_TOKENS = 1024 + + +def weight_shapes(config) -> "OrderedDict[str, tuple[int, ...]]": + """The block's fifteen tensors by checkpoint name, in checkpoint order, [out, in] for matrices.""" + hidden, inter = config.hidden_size, config.intermediate_size + q_width = config.num_attention_heads * config.head_dim + kv_width = config.num_key_value_heads * config.head_dim + return OrderedDict( + [ + ("input_layernorm.weight", (hidden,)), + ("self_attn.gate_bias", (config.num_attention_heads,)), + ("self_attn.q_proj.weight", (q_width, hidden)), + ("self_attn.k_proj.weight", (kv_width, hidden)), + ("self_attn.v_proj.weight", (kv_width, hidden)), + ("self_attn.gate_proj.weight", (q_width, hidden)), + ("self_attn.o_proj.weight", (hidden, q_width)), + ("self_attn.q_norm.weight", (config.head_dim,)), + ("self_attn.k_norm.weight", (config.head_dim,)), + ("pre_feedforward_layernorm.weight", (hidden,)), + ("mlp.gate_proj.weight", (inter, hidden)), + ("mlp.up_proj.weight", (inter, hidden)), + ("mlp.down_proj.weight", (hidden, inter)), + ("post_attention_layernorm.weight", (hidden,)), + ("post_feedforward_layernorm.weight", (hidden,)), + ] + ) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight").replace(".", "_") + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only dense blocks {SUPPORTED_LAYER_INDICES} are supported, got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this block computes {expected!r}") + layer_types = getattr(config, "layer_types", None) or [] + if len(layer_types) <= layer_idx or layer_types[layer_idx] != "sliding_attention": + raise ValueError( + f"block {layer_idx} must be a sliding_attention layer: " + f"{layer_types[layer_idx] if len(layer_types) > layer_idx else None!r}" + ) + mlp_types = getattr(config, "mlp_layer_types", None) or [] + if len(mlp_types) <= layer_idx or mlp_types[layer_idx] != "dense": + raise ValueError( + f"block {layer_idx} must have a dense MLP: " + f"{mlp_types[layer_idx] if len(mlp_types) > layer_idx else None!r}" + ) + rope = (getattr(config, "rope_parameters", None) or {}).get("sliding_attention") + if rope != SUPPORTED_ROPE: + raise ValueError( + f"unsupported sliding_attention rope parameters {rope!r}; need {SUPPORTED_ROPE!r}" + ) + + +class BeamDenseBlockBf16Sm103Tp1(nn.Module): + """Beam dense decoder block 0, 1 or 2: BF16 weights and activations, one GPU, one request step. + + ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` + returns ``[1, T, 5120]`` bf16 for 1 <= T <= ``MAX_TOKENS``, a prepared + ``TrtllmAttentionMetadata`` holding one step of one request on a BF16 + ``KVCacheManagerV2`` (8 kv heads of 128) that holds the block's logical layer + ``layer_idx`` -- a context step of T tokens after ``past`` cached ones (paged-context + FMHA when ``past > 0``) or a one-token generation step -- and positions + ``arange(past, past + T)``. Inputs outside that are rejected before any kernel runs; + batches of several requests and other blocks are not computed here. + """ + + def __init__( + self, config, layer_idx: int = DEFAULT_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this block is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.eps = config.rms_norm_eps + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.window = config.sliding_window + self.rope_theta = SUPPORTED_ROPE["rope_theta"] + self.attention_scales = depth_scales(layer_idx, is_attention=True) + self.feed_forward_scales = depth_scales(layer_idx, is_attention=False) + self._shapes = weight_shapes(config) + for name, shape in self._shapes.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The block's tensors by checkpoint name, in checkpoint order.""" + return OrderedDict((name, getattr(self, _attribute(name))) for name in self._shapes) + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the block's fifteen tensors, by checkpoint name, BF16, [out, in].""" + if not isinstance(weights, Mapping): + raise TypeError("weights must be a mapping") + missing = [n for n in self._shapes if n not in weights] + unexpected = sorted(n for n in weights if n not in self._shapes) + if missing or unexpected: + raise ValueError(f"block weights mismatch: missing {missing}, unexpected {unexpected}") + for name, shape in self._shapes.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tuple(tensor.shape) != shape or tensor.dtype != torch.bfloat16: + raise ValueError( + f"{name}: got {tuple(tensor.shape)} {tensor.dtype}, need {shape} torch.bfloat16" + ) + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError(f"{name} must be a dense host or CUDA tensor") + self._loaded = False + for name in self._shapes: + getattr(self, _attribute(name)).copy_(weights[name]) + self._loaded = True + + def _check_step( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> None: + """Raise ValueError on anything this block would otherwise compute wrongly.""" + _attention.check_step( + hidden_states, + position_ids, + md, + layer_idx=self.layer_idx, + hidden_size=self.hidden_size, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + device=self.input_layernorm.device, + max_tokens=MAX_TOKENS, + ) + if not self._loaded: + raise ValueError("load_weights must complete before running the block") + + @property + def hidden_size(self) -> int: + return self._shapes["input_layernorm.weight"][0] + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state.""" + return _attention.step_args(md, self.layer_idx) + + def attention(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm.""" + return _attention.attention( + h, + positions, + step, + q_proj=self.self_attn_q_proj, + k_proj=self.self_attn_k_proj, + v_proj=self.self_attn_v_proj, + gate_proj=self.self_attn_gate_proj, + o_proj=self.self_attn_o_proj, + q_norm=self.self_attn_q_norm, + k_norm=self.self_attn_k_norm, + gate_bias=self.self_attn_gate_bias, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.eps, + rope_theta=self.rope_theta, + window=self.window, + ) + + def feed_forward(self, h: torch.Tensor) -> torch.Tensor: + """Normalized input [T, hidden] -> dense SwiGLU output [T, hidden], before its post-norm.""" + gate = cublas_mm(h, self.mlp_gate_proj.t()) + up = cublas_mm(h, self.mlp_up_proj.t()) + act = flashinfer_silu_and_mul(torch.cat((gate, up), dim=-1)) + return cublas_mm(act, self.mlp_down_proj.t()) + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + self._check_step(hidden_states, position_ids, attn_metadata) + step = self.step_args(attn_metadata) + num_tokens = hidden_states.shape[1] + x = hidden_states.view(num_tokens, self.hidden_size) + positions = position_ids.view(num_tokens) + + h = flashinfer_rmsnorm(x, self.input_layernorm, self.eps) + h = self.attention(h, positions, step) + h = flashinfer_rmsnorm(h, self.post_attention_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.attention_scales) + + h = flashinfer_rmsnorm(x, self.pre_feedforward_layernorm, self.eps) + h = self.feed_forward(h) + h = flashinfer_rmsnorm(h, self.post_feedforward_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.feed_forward_scales) + return x.view(1, num_tokens, self.hidden_size) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/__init__.py new file mode 100644 index 000000000000..e6278158bae4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE feed-forward of decoder block 3 / sm_103 / tp1, as a directly constructed component.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/moe.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/moe.py new file mode 100644 index 000000000000..2e3ee6e1612d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe/moe.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE feed-forward of the sparse decoder blocks 3-5 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration, loads the MoE's 770 tensors +and passes it the already-normalized hidden states that enter its sparse block's +feed-forward sublayer; ``layer_idx`` (3, 4 or 5, as the router accepts) names that block, +and each block owns its own router and experts. It composes three accepted components and +adds no arithmetic of its own:: + + weights, ids = router(hidden) ../moe_router + out = dispatch_combine(hidden, ids, weights, ../dispatch_combine.py + lambda e, rows: experts[e](rows)) + ../moe_expert, 256 of them + +That is, as the source model computes it: fp32 grouped sigmoid top-k routing with a +selection-only correction bias (8 of 256 experts per token, from 4 of 8 groups) and +unbiased mixing weights; each selected expert, in ascending id order, runs its own bf16 +SwiGLU on its tokens' rows; each expert output times its fp32 weight is rounded to bf16, +the routes are summed in fp32, scaled by the routed scaling factor 2.5 once and rounded +to bf16. No residual, normalization or collective is applied here. + +Host synchronization: routing runs on the GPU without one. Dispatch reads the selected +experts and each expert's tokens back to the host -- two synchronizations for the expert +list and two per selected expert -- so a call is not capturable in a CUDA graph. Token +rows, products, roundings, sums and the scaling run on the GPU. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import dispatch_combine +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM + +from ..moe_expert.expert import WEIGHT_SHAPES as EXPERT_WEIGHT_SHAPES +from ..moe_expert.expert import BeamMoeExpertBf16Sm103Tp1 +from ..moe_router.router import BeamMoeRouterBf16Sm103Tp1 +from ..moe_router.router import check_config as check_router_config + +#: Fields of the checkpoint's configuration this component reads beyond the router's. +SUPPORTED_CONFIG = { + "moe_intermediate_size": 2560, + "hidden_act": "silu", + "routed_scaling_factor": 2.5, +} +DEFAULT_LAYER_IDX = 3 + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + check_router_config(config, layer_idx) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this MoE computes {expected!r}") + + +def weight_shapes(config) -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """The MoE's tensors by checkpoint name (relative to ``mlp.``), in order, with dtypes.""" + shapes = OrderedDict( + [ + ("gate.weight", ((config.num_experts, config.hidden_size), torch.float32)), + ("gate.e_score_correction_bias", ((config.num_experts,), torch.float32)), + ] + ) + for expert_id in range(config.num_experts): + for name, shape in EXPERT_WEIGHT_SHAPES.items(): + shapes[f"experts.{expert_id}.{name}"] = (shape, torch.bfloat16) + return shapes + + +def check_weights(expected: Mapping, weights: Mapping[str, torch.Tensor], what: str) -> None: + """Raise unless ``weights`` holds exactly ``expected``'s names, dtypes and shapes. + + Reads tensor metadata only, never values, and copies nothing. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + missing = [n for n in expected if n not in weights] + unexpected = sorted(str(n) for n in weights if n not in expected) + if missing or unexpected: + raise ValueError( + f"{what} weights mismatch: missing {missing[:8]}{'...' if len(missing) > 8 else ''}, " + f"unexpected {unexpected[:8]}{'...' if len(unexpected) > 8 else ''}" + ) + for name, (shape, dtype) in expected.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError( + f"{name} must be a dense host or CUDA tensor, got {tensor.layout} on {tensor.device}" + ) + + +class BeamMoeBf16Sm103Tp1(nn.Module): + """A sparse block's MoE: normalized BF16 hidden states in, the combined BF16 MoE output out. + + State: ``gate``, the router (FP32 ``weight [256, 5120]`` and selection-only + ``e_score_correction_bias [256]``), and ``experts``, 256 routed experts, each with its + own bias-free BF16 ``gate_proj``/``up_proj [2560, 5120]`` and ``down_proj [5120, 2560]``. + All are loaded together with ``load_weights``. + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the MoE's device, + T >= 1 tokens already normalized by the caller, and returns ``[T, 5120]`` + ``torch.bfloat16`` on that device, row ``t`` for token ``t``. Inputs outside that are + rejected before any kernel runs, and the input is never modified. Measured at T = 1, 4, + 17 and 129. + """ + + def __init__( + self, config, layer_idx: int = DEFAULT_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this MoE is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self._shapes = weight_shapes(config) + self.gate = BeamMoeRouterBf16Sm103Tp1(config, layer_idx, device=device) + self.experts = nn.ModuleList( + BeamMoeExpertBf16Sm103Tp1(device=device) for _ in range(config.num_experts) + ) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The MoE's tensors by checkpoint name (relative to ``mlp.``), in checkpoint order.""" + tensors = OrderedDict( + [ + ("gate.weight", self.gate.weight), + ("gate.e_score_correction_bias", self.gate.e_score_correction_bias), + ] + ) + for expert_id, expert in enumerate(self.experts): + for name, tensor in expert.checkpoint_tensors().items(): + tensors[f"experts.{expert_id}.{name}"] = tensor + return tensors + + def check_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Raise unless ``weights`` is exactly the MoE's 770 tensors by name, dtype and shape.""" + check_weights(self._shapes, weights, "MoE") + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the MoE's 770 tensors, by checkpoint name relative to ``mlp.``. + + ``gate.weight`` and ``gate.e_score_correction_bias`` are FP32; each + ``experts..{gate,up,down}_proj.weight`` is BF16 and loads into expert ``e``. + Values are copied as given, [out, in], without transposing or converting; host + tensors are copied to the MoE's device. Every tensor is checked before any + parameter changes, so a rejected load leaves the MoE as it was. The MoE runs only + after a load completed: a failure while copying leaves it unloaded. + """ + self.check_weights(weights) + self._loaded = False + self.gate.load_weights(weights["gate.weight"], weights["gate.e_score_correction_bias"]) + for expert_id, expert in enumerate(self.experts): + prefix = f"experts.{expert_id}." + expert.load_weights({name: weights[prefix + name] for name in EXPERT_WEIGHT_SHAPES}) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this MoE would otherwise compute wrongly.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before the MoE runs") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.layout != torch.strided: + raise ValueError(f"hidden_states must be a dense tensor, got {hidden_states.layout}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"hidden_states must be [tokens, {self.hidden_size}], " + f"got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.gate.weight.device: + raise ValueError( + f"hidden_states must be on the MoE's device {self.gate.weight.device}, " + f"got {hidden_states.device}" + ) + + def _expert(self, expert_id: int, rows: torch.Tensor) -> torch.Tensor: + """The dispatch provider: expert ``expert_id``'s own unweighted output for ``rows``.""" + return self.experts[expert_id](rows) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self._check_input(hidden_states) + mixing_weights, expert_ids = self.gate(hidden_states) + return dispatch_combine.dispatch_and_combine( + hidden_states, expert_ids, mixing_weights, self._expert + ) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/__init__.py new file mode 100644 index 000000000000..434f263710d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Beam MoE routed expert of decoder block 3 (SwiGLU, H5120 / I2560) / sm_103 / tp1.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/_expert_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/_expert_cases.py new file mode 100644 index 000000000000..76d0d76333d9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/_expert_cases.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated inputs, the plain-torch golden and the comparison rules of the expert tests. + +The expert is one SwiGLU expert of Beam decoder block 3: hidden 5120, intermediate 2560, +bias-free BF16 weights [out, in], BF16 projections, fp32 SiLU and product, one BF16 +rounding before the down projection. Every tensor here is generated on the host from a +seed or written down exactly, so the tests need no files. +""" + +from __future__ import annotations + +import contextlib +import math +import os + +import torch +import torch.nn.functional as F + +HIDDEN, INTER = 5120, 2560 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +STAGES = ("gate", "up", "silu", "product", "act", "y") +STAGE_DTYPES = { + "gate": torch.bfloat16, + "up": torch.bfloat16, + "silu": torch.float32, + "product": torch.float32, + "act": torch.bfloat16, + "y": torch.bfloat16, +} +BF16_UNIT = 2.0**-7 +FP32_FORMULA = 2.0**-16 + +#: (atol, rtol) per stage for the generated matrix, elementwise +#: ``|actual - expected| <= atol + rtol * |expected|``. rtol is one bf16 unit at any +#: magnitude; each atol is four times, rounded up, the largest excess over that measured +#: between valid implementations of this expert at these shapes (cuBLAS kernels, +#: fp32 and fp64 GEMMs, split-K orders, SiLU and product formulas). +MATRIX_RULES = { + "gate": (4e-5, BF16_UNIT), + "up": (4e-5, BF16_UNIT), + "silu": (0.02, BF16_UNIT), + "product": (0.009, BF16_UNIT), + "act": (0.04, BF16_UNIT), + "y": (0.008, BF16_UNIT), +} +#: Fixtures whose every projection output is one exact product, with values away from +#: any bf16 rounding boundary: projections, activation and output are exact, and SiLU and +#: product are fp32 formulas of exact inputs, within 2^-16 relative. +EXACT_RULES = { + "gate": (0.0, 0.0), + "up": (0.0, 0.0), + "silu": (0.0, FP32_FORMULA), + "product": (0.0, FP32_FORMULA), + "act": (0.0, 0.0), + "y": (0.0, 0.0), +} + +_TF32_OVERRIDES = ("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "NVIDIA_TF32_OVERRIDE") + + +# ----------------------------------------------------------------------------- numerics + + +def numerical_state() -> dict: + m = torch.backends.cuda.matmul + return { + "float32_matmul_precision": torch.get_float32_matmul_precision(), + "cuda_matmul_fp32_precision": m.fp32_precision, + "cudnn_fp32_precision": torch.backends.cudnn.fp32_precision, + "cuda_matmul_allow_tf32": m.allow_tf32, + "cudnn_allow_tf32": torch.backends.cudnn.allow_tf32, + "bf16_reduced_precision_reduction": m.allow_bf16_reduced_precision_reduction, + "autocast_cuda": torch.is_autocast_enabled("cuda"), + } + + +_CONTROLLED = { + "float32_matmul_precision": "highest", + "cuda_matmul_fp32_precision": "ieee", + "cudnn_fp32_precision": "ieee", + "cuda_matmul_allow_tf32": False, + "cudnn_allow_tf32": False, + "bf16_reduced_precision_reduction": False, + "autocast_cuda": False, +} + + +def assert_numerical_controls() -> None: + for var in _TF32_OVERRIDES: + assert var not in os.environ, ( + f"{var} substitutes TF32 for fp32 GEMMs; run with env -u {var}" + ) + state = numerical_state() + assert state == _CONTROLLED, f"numerical controls not in force: {state}" + + +@contextlib.contextmanager +def numerical_controls(): + """IEEE fp32 matmul, bf16 GEMMs without reduced-precision reduction, autocast off. + + The golden's fp32 steps and its bf16 GEMMs follow these switches; the comparison rules + assume them. They are asserted on entry and exit and restored afterwards. + """ + m = torch.backends.cuda.matmul + saved = ( + torch.get_float32_matmul_precision(), + torch.backends.cudnn.allow_tf32, + m.fp32_precision, + torch.backends.cudnn.fp32_precision, + m.allow_bf16_reduced_precision_reduction, + ) + torch.set_float32_matmul_precision("highest") + m.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + m.fp32_precision = "ieee" + torch.backends.cudnn.fp32_precision = "ieee" + m.allow_bf16_reduced_precision_reduction = False + try: + with torch.autocast("cuda", enabled=False): + assert_numerical_controls() + yield + assert_numerical_controls() + finally: + # The legacy switches first: their setters also move the newer ones restored after. + torch.set_float32_matmul_precision(saved[0]) + torch.backends.cudnn.allow_tf32 = saved[1] + m.fp32_precision = saved[2] + torch.backends.cudnn.fp32_precision = saved[3] + m.allow_bf16_reduced_precision_reduction = saved[4] + + +def golden(x: torch.Tensor, gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor): + """The source expert in plain torch, every step, on the inputs' device.""" + gate = F.linear(x, gate_w) + up = F.linear(x, up_w) + silu = F.silu(gate.float()) + product = silu * up.float() + act = product.to(torch.bfloat16) + return dict(gate=gate, up=up, silu=silu, product=product, act=act, y=F.linear(act, down_w)) + + +def violations(actual: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float) -> dict: + """Elementwise rule check, both sides widened exactly to fp64; a nonfinite value fails.""" + a = actual.detach().double() + e = expected.detach().to(actual.device).double() + diff = (a - e).abs() + ok = diff <= atol + rtol * e.abs() # NaN compares False: a nonfinite element fails + bad = ~ok + worst = int(torch.argmax(torch.where(ok, torch.zeros_like(diff), diff.nan_to_num(math.inf)))) + return { + "failing": int(bad.sum()), + "nonfinite": int((~torch.isfinite(a)).sum()), + "max_abs": float(diff.nan_to_num(math.inf).max()) if diff.numel() else 0.0, + "worst": (worst, float(a.reshape(-1)[worst]), float(e.reshape(-1)[worst])) + if diff.numel() + else None, + } + + +def assert_stages(actual: dict, expected: dict, rules: dict, what: str) -> None: + """Every stage: its dtype and shape, and each element within that stage's rule.""" + problems = [] + for stage in STAGES: + a, e = actual[stage], expected[stage] + if a.dtype != STAGE_DTYPES[stage]: + problems.append(f"{stage}: dtype {a.dtype}, expected {STAGE_DTYPES[stage]}") + if a.shape != e.shape: + problems.append(f"{stage}: shape {tuple(a.shape)}, expected {tuple(e.shape)}") + continue + found = violations(a, e, *rules[stage]) + if found["failing"] or found["nonfinite"]: + problems.append(f"{stage}: {found} under (atol, rtol) {rules[stage]}") + assert not problems, f"{what}: " + "; ".join(problems) + + +# ----------------------------------------------------------------------------- generated cases + + +def matrix_case(seed: int) -> tuple[dict, torch.Tensor]: + """Seed -> (weights by checkpoint name, the [129, 5120] input; case T uses its first T rows). + + Draw order gate, up, down, input, standard normal fp32 from a host generator; gate and + up scaled by 1/sqrt(5120), down by 1/sqrt(2560), each rounded once to bf16. + """ + g = torch.Generator(device="cpu").manual_seed(seed) + + def draw(rows: int, cols: int, scale: float) -> torch.Tensor: + return (torch.randn(rows, cols, generator=g, dtype=torch.float32) * scale).to( + torch.bfloat16 + ) + + weights = { + "gate_proj.weight": draw(INTER, HIDDEN, 1.0 / math.sqrt(HIDDEN)), + "up_proj.weight": draw(INTER, HIDDEN, 1.0 / math.sqrt(HIDDEN)), + "down_proj.weight": draw(HIDDEN, INTER, 1.0 / math.sqrt(INTER)), + } + return weights, draw(max(TOKENS), HIDDEN, 1.0) + + +def as_weights(gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor) -> dict: + return {"gate_proj.weight": gate, "up_proj.weight": up, "down_proj.weight": down} + + +# ----------------------------------------------------------------------------- exact fixtures + + +def silu64(t: torch.Tensor) -> torch.Tensor: + return F.silu(t.double()) + + +def bf16_margin(v: torch.Tensor) -> torch.Tensor: + """Relative distance of each fp64 value from its nearest bf16 rounding boundary.""" + v = v.double() + nearest = v.to(torch.bfloat16) + bits = nearest.view(torch.int16) + above = (bits + 1).view(torch.bfloat16).double() # one unit further from zero + below = (bits - 1).view(torch.bfloat16).double() # one unit nearer to zero + n = nearest.double() + dist = torch.minimum((v - (n + above) / 2).abs(), (v - (n + below) / 2).abs()) + return torch.where(v == 0, torch.full_like(v, math.inf), dist / v.abs()) + + +def bf16_grid(lo_exp: int, hi_exp: int, powers_of_two: bool = True) -> torch.Tensor: + """Every positive bf16 value in [2^lo_exp, 2^(hi_exp + 1)), as fp64.""" + mantissas = torch.arange(0 if powers_of_two else 1, 128, dtype=torch.float64) + return torch.cat([(1 + mantissas / 128) * 2.0**e for e in range(lo_exp, hi_exp + 1)]) + + +def identity_block_weights() -> dict: + """gate[i] = x[i], up[i] = x[2560 + i], y[5119 - 2i] = act[i]: one exact term each.""" + i = torch.arange(INTER) + gate = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + up = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + down = torch.zeros(HIDDEN, INTER, dtype=torch.bfloat16) + gate[i, i] = 1 + up[i, i + INTER] = 1 + down[HIDDEN - 1 - 2 * i, i] = 1 + return as_weights(gate, up, down) + + +def identity_block_expected(gate: torch.Tensor, up: torch.Tensor) -> dict: + """What every stage must be, exactly, for exact bf16 `gate`/`up` through identity blocks.""" + s = silu64(gate) + act = (s.float().double() * up).float().to(torch.bfloat16) + y = torch.zeros(gate.shape[0], HIDDEN, dtype=torch.bfloat16) + y[:, HIDDEN - 1 - 2 * torch.arange(INTER)] = act + return dict( + gate=gate.to(torch.bfloat16), + up=up.to(torch.bfloat16), + silu=s.float(), + product=(s.float().double() * up).float(), + act=act, + y=y, + ) + + +def signed_gate_fixture() -> tuple[dict, torch.Tensor, dict]: + """Negative, zero and positive gates over four decades, power-of-two signed ups. + + Gates whose SiLU lies within 2^-16 of a bf16 rounding boundary are dropped, so every + stage is exact; with a power-of-two up the product only rescales the SiLU. + """ + magnitudes = [1000.0, 60.0, 24.0, 10.0, 6.0, 4.5, 3.0, 2.25, 1.75, 1.28125, 1.0, 0.875] + magnitudes += [0.625, 0.375, 0.25, 0.125, 2.0**-4, 2.0**-6, 2.0**-10, 2.0**-20, 2.0**-60] + candidates = torch.tensor( + [-m for m in magnitudes] + [0.0, -0.0] + magnitudes, dtype=torch.float64 + ) + gates = candidates[bf16_margin(silu64(candidates)) > FP32_FORMULA] + ups = torch.tensor( + [1.0, -1.0, 2.0, -2.0, 0.5, -0.5, 4.0, -0.25, 0.125, -8.0], dtype=torch.float64 + ) + i = torch.arange(INTER) + gate = torch.stack([gates[(i + 37 * t) % len(gates)] for t in range(4)]) + up = torch.stack([ups[(5 * i + 3 * t) % len(ups)] for t in range(4)]) + x = torch.cat([gate, up], dim=1).to(torch.bfloat16) + return identity_block_weights(), x, identity_block_expected(gate, up) + + +def precision_fixture() -> tuple[dict, torch.Tensor, dict, torch.Tensor]: + """Gate/up pairs where the fp32 SiLU and product, rounded once, differ from bf16 steps. + + Over every bf16 gate with |g| in [2^-8, 8) and every non-power-of-two bf16 up with |u| + in [0.5, 4), keep the pairs whose correct activation and whose prematurely rounded one + (bf16 SiLU, bf16 product) are both 2^-16 or more from a bf16 rounding boundary and + differ, then draw 4 x 2560 of them with a fixed seed. Returns the weights, the input, + the exact stages and the premature activation every element must differ from. + """ + g_pos, u_pos = bf16_grid(-8, 2), bf16_grid(-1, 1, powers_of_two=False) + g = torch.cat([-g_pos.flip(0), g_pos]) + u = torch.cat([-u_pos.flip(0), u_pos]) + s = silu64(g) + product = s.float().double()[:, None] * u[None, :] + premature = s.float().to(torch.bfloat16).double()[:, None] * u[None, :] + robust = (bf16_margin(product) > FP32_FORMULA) & (bf16_margin(s) > FP32_FORMULA)[:, None] + robust &= bf16_margin(premature) > FP32_FORMULA + differs = product.float().to(torch.bfloat16) != premature.to(torch.bfloat16) + rows, cols = torch.nonzero(robust & differs, as_tuple=True) + pick = torch.randperm(rows.numel(), generator=torch.Generator().manual_seed(20261020)) + pick = pick[: 4 * INTER] + gate, up = g[rows[pick]].reshape(4, INTER), u[cols[pick]].reshape(4, INTER) + x = torch.cat([gate, up], dim=1).to(torch.bfloat16) + premature_act = (s.float().to(torch.bfloat16).double()[rows[pick]] * u[cols[pick]]).to( + torch.bfloat16 + ) + return ( + identity_block_weights(), + x, + identity_block_expected(gate, up), + premature_act.reshape(4, INTER), + ) + + +def sparse_mapping_fixture() -> tuple[dict, torch.Tensor, dict]: + """Asymmetric one-term maps with power-of-two markers: + gate[i] = m_g(i) x[a(i)], up[i] = m_u(i) x[b(i)], y[d(i)] = m_d(i) act[i]. + + a and b together cover all 5120 input features, d is a permutation into the 5120 + outputs, and gate and up read different inputs with different markers, so a transposed + weight, swapped gate/up roles or a wrong down mapping all move some output. + """ + i = torch.arange(INTER) + a, b, d = (3 * i + 7) % HIDDEN, (3 * i + 2567) % HIDDEN, (HIDDEN - 1 - 3 * i) % HIDDEN + m_g = 2.0 ** ((i % 5) - 2).double() + m_u = -(2.0 ** ((i % 3) - 1).double()) + m_d = 2.0 ** ((i % 4) - 1).double() + candidates = torch.cat([-bf16_grid(-3, 0).flip(0), bf16_grid(-3, 0)]) + robust = torch.ones_like(candidates, dtype=torch.bool) + for marker in (0.25, 0.5, 1.0, 2.0, 4.0): + robust &= bf16_margin(silu64(candidates * marker)) > FP32_FORMULA + palette = candidates[robust] + palette = palette[torch.randperm(palette.numel(), generator=torch.Generator().manual_seed(7))] + x = torch.zeros(3, HIDDEN, dtype=torch.float64) + for t in range(3): + x[t, a] = palette[(7 * i + 13 * t) % 64] + sign = torch.where((i + 2 * t) % 3 == 0, -1.0, 1.0).double() + x[t, b] = sign * 2.0 ** (((i + t) % 5) - 2).double() + gate_w = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + up_w = torch.zeros(INTER, HIDDEN, dtype=torch.bfloat16) + down_w = torch.zeros(HIDDEN, INTER, dtype=torch.bfloat16) + gate_w[i, a] = m_g.to(torch.bfloat16) + up_w[i, b] = m_u.to(torch.bfloat16) + down_w[d, i] = m_d.to(torch.bfloat16) + gate, up = x[:, a] * m_g, x[:, b] * m_u + expected = identity_block_expected(gate, up) + y = torch.zeros(3, HIDDEN, dtype=torch.bfloat16) + y[:, d] = (expected["act"].double() * m_d).to(torch.bfloat16) + expected["y"] = y + return as_weights(gate_w, up_w, down_w), x.to(torch.bfloat16), expected diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/expert.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/expert.py new file mode 100644 index 000000000000..f8716cd7fe8c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_expert/expert.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One Beam MoE routed expert of decoder block 3 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it, loads one expert's three tensors and passes it the +already-normalized hidden states of the tokens routed to that expert. Which expert it is +decides only which tensors are loaded. It routes no token and applies no mixing weight, +routed scaling factor, residual or normalization, and it makes no collective: weighting +and combining expert outputs, and any reduction across ranks, belong to the caller. It +does not call the built-in Beam modeling. + +The expert, bias-free, with the projections in bf16 and the activation in fp32:: + + gate = hidden @ gate_proj^T [T, 2560] bf16 catalog cublas_mm + up = hidden @ up_proj^T [T, 2560] bf16 catalog cublas_mm + silu = silu(gate.float()) [T, 2560] fp32 ops/activation/swiglu_fp32.py + product = silu * up.float() [T, 2560] fp32 + act = product.to(bfloat16) [T, 2560] bf16 the one rounding before down + out = act @ down_proj^T [T, 5120] bf16 catalog cublas_mm +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.swiglu_fp32 import ( + swiglu_fp32, +) + +HIDDEN_SIZE = 5120 +INTERMEDIATE_SIZE = 2560 + +#: The expert's three tensors by checkpoint name (relative to ``mlp.experts..``), [out, in]. +WEIGHT_SHAPES = OrderedDict( + [ + ("gate_proj.weight", (INTERMEDIATE_SIZE, HIDDEN_SIZE)), + ("up_proj.weight", (INTERMEDIATE_SIZE, HIDDEN_SIZE)), + ("down_proj.weight", (HIDDEN_SIZE, INTERMEDIATE_SIZE)), + ] +) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight") + + +class BeamMoeExpertBf16Sm103Tp1(nn.Module): + """One routed expert: BF16 hidden states in, the expert's unweighted BF16 output out. + + State is exactly three bias-free BF16 matrices, loaded with ``load_weights``: + ``gate_proj`` and ``up_proj`` ``[2560, 5120]`` and ``down_proj`` ``[5120, 2560]``, stored + as given, [out, in]. + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the expert's + device, T >= 1 tokens already normalized by the caller, and returns ``[T, 5120]`` + ``torch.bfloat16`` on that device, one row per input row. Inputs outside that are + rejected before any kernel runs. Measured at T = 1, 17 and 129; zero tokens (an expert + no token was routed to) is rejected, not computed. + + Op ownership: the three projections are the catalog's ``cublas_mm``, each with the + stored weight's ``.t()`` as its column-major operand; the fp32 SiLU, the fp32 product + and the bf16 rounding are this component's ``swiglu_fp32``. + """ + + def __init__(self, device: str | torch.device = "cuda"): + super().__init__() + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this expert is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + for name, shape in WEIGHT_SHAPES.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self._loaded = False + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The expert's tensors by checkpoint name, in checkpoint order.""" + return OrderedDict((name, getattr(self, _attribute(name))) for name in WEIGHT_SHAPES) + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the expert's three tensors, by checkpoint name: BF16, [out, in]. + + Values are copied as given, without transposing or converting; host tensors are + copied to the expert's device. Every tensor is checked, and copied to the device, + before any parameter changes, so a rejected load leaves the expert as it was. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + missing = [n for n in WEIGHT_SHAPES if n not in weights] + unexpected = sorted(str(n) for n in weights if n not in WEIGHT_SHAPES) + if missing or unexpected: + raise ValueError(f"expert weights mismatch: missing {missing}, unexpected {unexpected}") + for name, shape in WEIGHT_SHAPES.items(): + tensor = weights[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must be torch.bfloat16, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + if tensor.layout != torch.strided or tensor.device.type not in ("cpu", "cuda"): + raise ValueError( + f"{name} must be a dense host or CUDA tensor, got {tensor.layout} " + f"on {tensor.device}" + ) + device = self.gate_proj.device + staged = {name: weights[name].to(device=device, copy=True) for name in WEIGHT_SHAPES} + for name, tensor in staged.items(): + getattr(self, _attribute(name)).copy_(tensor) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this expert would otherwise compute wrongly.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before the expert runs") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != HIDDEN_SIZE: + raise ValueError( + f"hidden_states must be [tokens, {HIDDEN_SIZE}], got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.gate_proj.device: + raise ValueError( + f"hidden_states must be on the expert's device {self.gate_proj.device}, " + f"got {hidden_states.device}" + ) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self._check_input(hidden_states) + x = hidden_states.contiguous() + gate = cublas_mm(x, self.gate_proj.t()) + up = cublas_mm(x, self.up_proj.t()) + act = swiglu_fp32(gate, up)["act"] + return cublas_mm(act, self.down_proj.t()) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/__init__.py new file mode 100644 index 000000000000..539171ed4825 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE router of decoder block 3 (grouped sigmoid top-k) / sm_103 / tp1.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/_router_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/_router_cases.py new file mode 100644 index 000000000000..6ef460f0d26e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/_router_cases.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs and comparison rules shared by the Beam MoE router's tests. + +Everything here is generated on the CPU from fixed seeds, so the bytes a test routes are +the same on every machine and can be checked against routes captured elsewhere. + +Matrix cases, per seed: a CPU generator seeded with it draws the FP32 router weight +``randn(256, 5120) / sqrt(5120)``, then the FP32 correction bias ``0.05 * randn(256)``, +then FP32 token rows ``randn(T, 5120)``, scaled to unit RMS and rounded to BF16 -- the +already-normalized input the router receives. Weight and bias are drawn first, so they +are the same for every token count of a seed. + +Controlled fixtures fix the logits exactly instead: token row ``r`` is one-hot at column +``r`` and the seed-17 weight's column ``r`` is replaced by the logits that row should +produce, so ``logits[r] = weight[:, r]`` whatever the GEMM's reduction order. Each fixture +is a decision the router has to get right and says which. +""" + +from __future__ import annotations + +import math + +import torch + +HIDDEN, EXPERTS, GROUPS, PER_GROUP, TOPK_GROUP, TOP_K = 5120, 256, 8, 32, 4, 8 +NORM_EPS = 1e-20 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +CASES = tuple((seed, tokens) for seed in SEEDS for tokens in TOKENS) +CASE_IDS = [f"seed{seed}-T{tokens}" for seed, tokens in CASES] + +#: Elementwise |actual - reference| <= atol (rtol 0) per compared tensor, both widened to +#: fp64. Each atol is four times the largest difference measured between valid fp32 +#: evaluations of this routing at these shapes and inputs -- GEMM reduction orders and +#: layouts, sigmoid forms, normalization forms -- rounded up to one significant digit. +#: A route is never excused by a tolerance: groups and ids must match exactly, in order. +RULES = { + "logits": 8e-6, + "scores": 3e-6, + "selection_scores": 3e-6, + "group_scores": 2e-6, + "weights": 8e-7, +} + +#: |sum(weights) - 1| bound for every row whose selected scores are ordinary: eight fp32 +#: terms summed, divided and summed again err by at most 15 units of 2^-24 (8.9e-7). +NORM_ATOL = 1e-6 + + +# ----------------------------------------------------------------------------- matrix + + +def matrix_case(seed: int, tokens: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(weight [256, 5120] fp32, correction bias [256] fp32, hidden [T, 5120] bf16) on CPU.""" + g = torch.Generator(device="cpu").manual_seed(seed) + weight = torch.randn(EXPERTS, HIDDEN, generator=g, dtype=torch.float32) * ( + 1.0 / math.sqrt(HIDDEN) + ) + bias = torch.randn(EXPERTS, generator=g, dtype=torch.float32) * 0.05 + x = torch.randn(tokens, HIDDEN, generator=g, dtype=torch.float32) + x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True)) + return weight, bias, x.to(torch.bfloat16) + + +# ----------------------------------------------------------------------------- controlled fixtures + + +def _logit(p: float) -> float: + return math.log(p / (1.0 - p)) + + +def _row(default: float) -> list: + return [default] * EXPERTS + + +def _set_group(row: list, group: int, values) -> None: + for j, v in enumerate(values): + row[group * PER_GROUP + j] = v + + +def _logits_for_scores(rows) -> torch.Tensor: + """FP32 logits whose sigmoid is each row's intended score.""" + return torch.tensor([[_logit(p) for p in r] for r in rows], dtype=torch.float64).float() + + +def _bias_fixtures() -> list: + # Each group has two lead experts; the rest sit at distinct values just above 0.30. + p = [0.30 + 0.0001 * ((e * 37) % EXPERTS) for e in range(EXPERTS)] + leads = { + 0: (0.70, 0.69), + 1: (0.68, 0.67), + 2: (0.66, 0.65), + 3: (0.64, 0.63), + 4: (0.62, 0.61), + 5: (0.60, 0.59), + 6: (0.50, 0.49), + 7: (0.48, 0.47), + } + for group, (a, b) in leads.items(): + p[group * PER_GROUP], p[group * PER_GROUP + 1] = a, b + logits = _logits_for_scores([p]) + shifted = torch.zeros(EXPERTS, dtype=torch.float32) + shifted[6 * PER_GROUP] = 0.25 + shifted[6 * PER_GROUP + 1] = 0.27 + return [ + dict( + name="bias_zero", + rows=("lead_experts",), + logits=logits, + bias=torch.zeros(EXPERTS), + intent="no bias: groups [0, 1, 2, 3], experts 0, 1, 32, 33, 64, 65, 96, 97", + ), + dict( + name="bias_shift", + rows=("lead_experts",), + logits=logits, + bias=shifted, + intent="+0.25/+0.27 on experts 192/193 lifts group 6 to first: groups [6, 0, 1, 2]; " + "the weights stay the normalized unbiased scores of the selected experts", + ), + ] + + +def _group_boundary() -> dict: + rows = [] + for second in (0.62, 0.67): + r = _row(0.20) + _set_group(r, 0, [0.90, 0.89]) + _set_group(r, 1, [0.88, 0.87]) + _set_group(r, 2, [0.86, 0.85]) + _set_group(r, 3, [0.655, 0.645]) # incumbent: top-2 sum 1.300 + _set_group(r, 4, [0.66, second]) # challenger: 1.28 below, 1.33 above + _set_group(r, 5, [0.95] + [0.05] * (PER_GROUP - 1)) # best single expert, sum 1.00 + _set_group(r, 6, [0.30, 0.29]) + _set_group(r, 7, [0.28, 0.27]) + rows.append(r) + return dict( + name="group_boundary", + rows=("challenger_below", "challenger_above"), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="only group 4's second expert moves (0.62 -> 0.67): groups [0, 1, 2, 3] -> " + "[0, 1, 2, 4]; scoring groups by their best expert alone would pick group 5", + ) + + +def _expert_boundary() -> dict: + rows = [] + for mover in (0.79, 0.82): + r = _row(0.30) + _set_group(r, 0, [0.95, 0.93, mover]) # third in its group: group scores do not move + _set_group(r, 1, [0.91, 0.89]) + _set_group(r, 2, [0.87, 0.85]) + _set_group(r, 3, [0.83, 0.81]) + _set_group(r, 4, [0.60, 0.58]) + _set_group(r, 5, [0.59, 0.57]) + _set_group(r, 6, [0.56, 0.55]) + _set_group(r, 7, [0.54, 0.53]) + rows.append(r) + return dict( + name="expert_boundary", + rows=("mover_below", "mover_above"), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="expert 2 moves 0.79 -> 0.82 across the 8th/9th boundary and expert 97 (0.81) " + "drops out; the selected groups do not change", + ) + + +EXCLUSION_SPIKE = 7 * PER_GROUP + 5 + + +def _exclusion() -> dict: + logits = [0.0005 * ((e * 37) % EXPERTS) for e in range(EXPERTS)] + logits[EXCLUSION_SPIKE] = 6.0 + bias = [] + for e in range(EXPERTS): + group = e // PER_GROUP + if group <= 3: + bias.append(-2.0) + elif group <= 6: + bias.append(-5.0) + else: + bias.append(0.0 if e == EXCLUSION_SPIKE else -10.0) + return dict( + name="exclusion", + rows=("spike_in_excluded_group",), + logits=torch.tensor([logits], dtype=torch.float32), + bias=torch.tensor(bias, dtype=torch.float32), + intent=f"expert {EXCLUSION_SPIKE} has the highest selection score but its group loses; " + "every retained selection score is negative, so skipping the exclusion or masking with " + "zero instead of -inf both select an excluded expert", + ) + + +def _ties() -> dict: + def top7() -> list: + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89, 0.60]) + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + return r + + rows, names = [], [] + r = top7() + r[1 * PER_GROUP + 7] = 0.70 + r[1 * PER_GROUP + 25] = 0.70 + rows.append(r) + names.append("expert_boundary_tie_same_group_39_57") + r = top7() + r[0 * PER_GROUP + 3] = 0.70 + r[3 * PER_GROUP + 30] = 0.70 + rows.append(r) + names.append("expert_boundary_tie_distant_3_126") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + r[1 * PER_GROUP + 1] = 0.92 + r[1 * PER_GROUP + 30] = 0.92 + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89, 0.88]) + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + rows.append(r) + names.append("expert_order_tie_33_62") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.80, 0.78]) + _set_group(r, 6, [0.80, 0.78]) + _set_group(r, 4, [0.50, 0.49]) + _set_group(r, 5, [0.48, 0.47]) + _set_group(r, 7, [0.46, 0.45]) + rows.append(r) + names.append("group_boundary_tie_3_6") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.90, 0.89]) + _set_group(r, 4, [0.90, 0.89]) + _set_group(r, 2, [0.85, 0.84]) + _set_group(r, 3, [0.50, 0.49]) + _set_group(r, 5, [0.48, 0.47]) + _set_group(r, 6, [0.46, 0.45]) + _set_group(r, 7, [0.44, 0.43]) + rows.append(r) + names.append("group_order_tie_1_4") + rows.append(_row(0.5)) + names.append("all_equal") + r = _row(0.30) + _set_group(r, 0, [0.95, 0.94]) + _set_group(r, 1, [0.93, 0.92]) + _set_group(r, 2, [0.91, 0.90]) + _set_group(r, 3, [0.89]) + for g, js in ((0, (5, 17, 31)), (1, (2, 28)), (2, (9, 10, 20)), (3, (4, 30))): + for j in js: + r[g * PER_GROUP + j] = 0.70 + for g in range(4, 8): + _set_group(r, g, [0.50 - 0.01 * g, 0.49 - 0.01 * g]) + rows.append(r) + names.append("expert_boundary_tie_10way") + return dict( + name="ties", + rows=tuple(names), + logits=_logits_for_scores(rows), + bias=torch.zeros(EXPERTS), + intent="bitwise-equal scores at the 8th/9th expert boundary, inside the top-8 order, at " + "the 4th/5th group boundary, inside the top-4 group order, everywhere, and a 10-way " + "boundary tie", + ) + + +#: The order `torch.topk(..., sorted=True)` returns for each tie row -- the source router's +#: order, measured on sm_103. A tie at a selection boundary keeps the lower index; tied +#: values inside the selection are not in index order. +TIE_ROUTES = { + "expert_boundary_tie_same_group_39_57": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 39]), + "expert_boundary_tie_distant_3_126": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 3]), + "expert_order_tie_33_62": ([0, 1, 2, 3], [0, 1, 33, 62, 64, 65, 96, 97]), + "group_boundary_tie_3_6": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 97]), + "group_order_tie_1_4": ([0, 4, 1, 2], [0, 1, 128, 32, 33, 129, 64, 65]), + "all_equal": ([1, 0, 2, 3], [7, 6, 4, 5, 1, 0, 2, 3]), + "expert_boundary_tie_10way": ([0, 1, 2, 3], [0, 1, 32, 33, 64, 65, 96, 5]), +} + +EPSILON_IDS = [0, 1, 32, 33, 64, 65, 96, 97] + + +def _epsilon() -> dict: + def bias_of(e: int) -> float: + group, j = divmod(e, PER_GROUP) + return 0.9 - 0.1 * group - 0.001 * j if j < 2 else -0.5 - 0.01 * group - 0.0001 * j + + zero_row = [-120.0] * EXPERTS + near_row = [-60.0] * EXPERTS + for k, e in enumerate(EPSILON_IDS): + near_row[e] = -47.5 - 0.15 * k + return dict( + name="epsilon", + rows=("zero_sum", "near_epsilon_sum"), + logits=torch.tensor([zero_row, near_row], dtype=torch.float32), + bias=torch.tensor([bias_of(e) for e in range(EXPERTS)], dtype=torch.float32), + intent="the bias fixes the selection to experts 0, 1, 32, 33, 64, 65, 96, 97; row 0's " + "scores underflow to 0 (weights 0, not NaN), row 1's selected scores sum to ~1e-20, so " + "its weights sum visibly below 1 with exactly the 1e-20 denominator term", + ) + + +def fixtures() -> dict: + """Every controlled fixture by name, in a fixed order.""" + specs = [*_bias_fixtures(), _group_boundary(), _expert_boundary(), _exclusion(), _ties()] + specs.append(_epsilon()) + return {spec["name"]: spec for spec in specs} + + +def fixture_tensors(spec: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(weight, correction bias, hidden) on CPU that make the router see `spec`'s logits.""" + logits = spec["logits"].float() + rows = logits.shape[0] + weight = matrix_case(17, 1)[0].clone() + weight[:, :rows] = logits.t() + hidden = torch.zeros(rows, HIDDEN, dtype=torch.bfloat16) + hidden[torch.arange(rows), torch.arange(rows)] = 1.0 + return weight.contiguous(), spec["bias"].float().contiguous(), hidden + + +# ----------------------------------------------------------------------------- comparison + + +def compare(name: str, actual: torch.Tensor, expected: torch.Tensor, atol: float) -> dict: + """Score one fp32 tensor against `atol` (rtol 0); non-finite values must match exactly.""" + result = {"tensor": name, "atol": atol, "rtol": 0.0, "shape": list(expected.shape)} + problems = [] + if tuple(actual.shape) != tuple(expected.shape): + problems.append(f"shape {tuple(actual.shape)} != {tuple(expected.shape)}") + if actual.dtype != torch.float32 or expected.dtype != torch.float32: + problems.append(f"dtypes {actual.dtype}/{expected.dtype}, need torch.float32") + if not problems: + a, e = actual.detach().double(), expected.detach().double() + finite_a, finite_e = torch.isfinite(a), torch.isfinite(e) + nonfinite_mismatch = int(((finite_a != finite_e) | (~finite_e & (a != e))).sum()) + both = finite_a & finite_e + diff = (a[both] - e[both]).abs() + violations = int((diff > atol).sum()) + result.update( + max_abs=float(diff.max()) if diff.numel() else 0.0, + unequal=int((diff > 0).sum()), + violations=violations, + nonfinite_mismatch=nonfinite_mismatch, + numel=int(e.numel()), + ) + if violations: + problems.append(f"{violations} elements exceed atol {atol:g} (max {diff.max():.3e})") + if nonfinite_mismatch: + problems.append(f"{nonfinite_mismatch} non-finite positions differ") + result["problems"] = problems + result["passed"] = not problems + return result + + +def compare_exact(name: str, actual: torch.Tensor, expected: torch.Tensor) -> dict: + """Ordered integer ids: equal element for element, never sorted first.""" + same = tuple(actual.shape) == tuple(expected.shape) and actual.dtype == expected.dtype + rows = ( + (actual == expected).all(dim=-1) + if same + else torch.zeros(expected.shape[0], dtype=torch.bool, device=expected.device) + ) + mismatched = [int(i) for i in torch.nonzero(~rows).flatten().cpu()] + return { + "tensor": name, + "rule": "exact ordered equality", + "shape": list(expected.shape), + "dtype": str(actual.dtype), + "mismatched_rows": mismatched, + "passed": same and not mismatched, + "problems": [] if same and not mismatched else [f"rows {mismatched[:8]} differ"], + } + + +def compare_routes(actual: dict, expected: dict, label: str) -> list: + """Every step of a route against another route of the same input.""" + results = [ + compare_exact("groups", actual["groups"], expected["groups"]), + compare_exact("ids", actual["ids"], expected["ids"]), + ] + for name, atol in RULES.items(): + if name == "weights": + continue + results.append(compare(name, actual[name], expected[name], atol)) + results.append( + compare( + "masked_selection_scores", + actual["masked_selection_scores"], + expected["masked_selection_scores"], + RULES["selection_scores"], + ) + ) + # Weights belong to the experts they weigh: compared on rows routed identically. + same_ids = (actual["ids"] == expected["ids"]).all(dim=-1) + results.append( + compare( + "weights", actual["weights"][same_ids], expected["weights"][same_ids], RULES["weights"] + ) + ) + for r in results: + r["against"] = label + return results + + +def assert_all_pass(results: list) -> None: + failed = [ + f"{r['against']}:{r['tensor']}: {'; '.join(r['problems'])}" + for r in results + if not r["passed"] + ] + assert not failed, "\n".join(failed) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/router.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/router.py new file mode 100644 index 000000000000..7e9f84c0e97e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/moe_router/router.py @@ -0,0 +1,178 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam MoE router of the sparse decoder blocks 3-5 / sm_103 / tp1, constructed directly. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration, loads the router's two +tensors and passes it the already-normalized hidden states that enter its sparse +block's feed-forward sublayer; each sparse block owns its own router. It returns the +eight experts each token is routed to and their mixing weights. It runs no expert, +dispatches no token and combines no expert output; the routed scaling factor belongs to +that combine and is not applied here. It does not call the built-in Beam modeling. + +The routing, every step fp32 (``ops/moe/beam_grouped_topk_router.py``):: + + logits = hidden.float() @ weight^T [T, 256] + scores = sigmoid(logits) unbiased: the weights use these + choice = scores + e_score_correction_bias biased: only selection uses these + groups = top-4 of (sum of the 2 largest choice in each of 8 groups of 32), descending + ids = top-8 choice among the experts of those groups, descending [T, 8] int64 + weights = scores[ids] / (sum(scores[ids]) + 1e-20) [T, 8] fp32 +""" + +from __future__ import annotations + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import ( + ROUTER_EPS as NORM_EPS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.moe.beam_grouped_topk_router import ( + beam_grouped_topk_router, +) + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "num_experts": 256, + "num_experts_per_tok": 8, + "n_group": 8, + "topk_group": 4, + "norm_topk_prob": True, +} +DEFAULT_LAYER_IDX = 3 +#: The sparse blocks this router has been built for: the checkpoint's first three. +SUPPORTED_LAYER_INDICES = (3, 4, 5) + +#: Added to the selected scores' sum before dividing, as the source model does. + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only the routers of sparse blocks {SUPPORTED_LAYER_INDICES} are supported, " + f"got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this router computes {expected!r}") + mlp_types = getattr(config, "mlp_layer_types", None) or [] + if len(mlp_types) <= layer_idx or mlp_types[layer_idx] != "sparse": + raise ValueError( + f"block {layer_idx} must have a sparse (MoE) feed-forward: " + f"mlp_layer_types[{layer_idx}]=" + f"{mlp_types[layer_idx] if len(mlp_types) > layer_idx else None!r}" + ) + + +class BeamMoeRouterBf16Sm103Tp1(nn.Module): + """A sparse block's router: BF16 hidden states in, expert ids and FP32 mixing weights out. + + State is exactly two FP32 tensors, loaded with ``load_weights``: ``weight`` + ``[256, 5120]`` (checkpoint ``mlp.gate.weight``) and ``e_score_correction_bias`` + ``[256]`` (``mlp.gate.e_score_correction_bias``). + + ``forward(hidden_states)`` takes ``[T, 5120]`` ``torch.bfloat16`` on the router's + device, T >= 1 tokens already normalized by the caller, and returns + ``(mixing_weights, expert_ids)`` -- the source model's order -- both ``[T, 8]`` on that + device: ``mixing_weights`` ``torch.float32``, ``expert_ids`` ``torch.int64`` in + ``[0, 256)``, in descending order of the biased selection score. Inputs outside that + are rejected before any kernel runs. Routing runs entirely on the GPU and makes no + host synchronization. + """ + + def __init__( + self, config, layer_idx: int = DEFAULT_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this router is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.n_group = config.n_group + self.topk_group = config.topk_group + self.weight = nn.Parameter( + torch.empty(self.num_experts, self.hidden_size, dtype=torch.float32, device=device), + requires_grad=False, + ) + self.register_buffer( + "e_score_correction_bias", + torch.empty(self.num_experts, dtype=torch.float32, device=device), + ) + self._loaded = False + + @torch.no_grad() + def load_weights(self, weight: torch.Tensor, correction_bias: torch.Tensor) -> None: + """Copy the router's two tensors: ``weight [256, 5120]`` and ``correction_bias [256]``. + + Both must already be ``torch.float32``: the router computes in fp32 and does not + accept parameters that were rounded below it. Both are checked before either is + copied, so a rejected load leaves the router as it was. + """ + expected = ( + ("weight", weight, (self.num_experts, self.hidden_size)), + ("correction_bias", correction_bias, (self.num_experts,)), + ) + for name, tensor, shape in expected: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must be torch.float32, got {tensor.dtype}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must be {list(shape)}, got {list(tensor.shape)}") + self.weight.copy_(weight) + self.e_score_correction_bias.copy_(correction_bias) + self._loaded = True + + def _check_input(self, hidden_states: torch.Tensor) -> None: + """Raise on anything this router would otherwise compute wrongly or move to the host.""" + if not self._loaded: + raise RuntimeError("load_weights must be called before routing") + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"hidden_states must be a torch.Tensor, got {type(hidden_states).__name__}" + ) + if hidden_states.dtype != torch.bfloat16: + raise TypeError(f"hidden_states must be torch.bfloat16, got {hidden_states.dtype}") + if hidden_states.dim() != 2 or hidden_states.shape[1] != self.hidden_size: + raise ValueError( + f"hidden_states must be [tokens, {self.hidden_size}], " + f"got {list(hidden_states.shape)}" + ) + if hidden_states.shape[0] < 1: + raise ValueError("hidden_states must hold at least one token") + if hidden_states.device != self.weight.device: + raise ValueError( + f"hidden_states must be on the router's device {self.weight.device}, " + f"got {hidden_states.device}" + ) + + def route(self, hidden_states: torch.Tensor) -> dict[str, torch.Tensor]: + """Every intermediate of one routing call, by the op's names; see the op's docstring.""" + self._check_input(hidden_states) + return beam_grouped_topk_router( + hidden_states, + self.weight, + self.e_score_correction_bias, + n_group=self.n_group, + topk_group=self.topk_group, + top_k=self.top_k, + eps=NORM_EPS, + ) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + routing = self.route(hidden_states) + return routing["weights"], routing["ids"] diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/__init__.py new file mode 100644 index 000000000000..77397f5b34ef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam six-block small model, token ids to prefill logits (H5120, vocab 151936) / sm_103 / tp1.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/generation.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/generation.py new file mode 100644 index 000000000000..2de6b0ffcb8a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/generation.py @@ -0,0 +1,693 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Stateful greedy generation for the Beam six-block model / sm_103 / tp1. + +A thin serial controller that drives the stateless six-block model +(``BeamSixLayerBf16Sm103Tp1``) through TensorRT-LLM's own request, cache and +sampler objects: one ``LlmRequest`` per request, one long-lived ``KVCacheManagerV2`` with +block reuse for all of them, a ``SeqSlotManager`` for sampler slots and the existing +``TorchSampler``. It owns no cache allocator, prefix table or sampler of its own; it calls +those objects in the order the non-overlap executor does -- forward, ``sample_async``, +context-cursor advance, ``update_requests``, ``update_context_resources``, finish +handling, ``update_resources``. It is a component, not a serving engine: one request +executes at a time, at most ``GENERATOR_CACHE["max_batch_size"]`` (two) request states +are live, and there is no scheduler, batching, streaming or LLM API registration. + +Lifecycle:: + + with BeamSmallModelGenerator(model) as gen: # binds the loaded model + state = gen.create_request(prompt_ids, max_new_tokens=8, end_id=None) + first = gen.prefill(state) # prompt tokens not already cached; first token + while state.lifecycle is not Lifecycle.FINISHED: + step = gen.decode(state) # consumes the pending sampled token + tokens, reason = state.output_token_ids, state.finish_reason + gen.release(state) + +``generate(state)`` runs the prefill and decode steps of a new sampling request to its +end. A compute-only request (``sample=False``) runs the same steps without the sampler: +``prefill`` returns the last prompt token's logits and ``decode(state, token_id)`` +consumes an explicit (teacher-forced) token; its token budget is ``max_new_tokens`` +explicit tokens. + +Semantics. + +* Positions are absolute. A prefill computes prompt tokens ``reused .. P - 1`` at those + positions after ``reused`` cached ones; a decode consumes one token at position + ``consumed`` (the number of tokens whose keys and values the cache holds) and leaves + ``consumed + 1``. Every model call is checked by the model against the manager before + it runs (``model.check_request``). +* Prefix reuse is the manager's. A prefill first asks the manager for the request's + cache (``prepare_context``): the manager matches the prompt against what earlier + requests committed and claims the longest reusable full-page prefix -- at most + ``P - 1`` tokens, since the last prompt token is always recomputed for its logits; + with partial reuse off, only whole 32-token pages. The prefill computes the rest and the + manager commits it (``update_context_resources``), so later requests can reuse it, + also after this request is released. ``reuse=False`` claims nothing + (``prepare_context_cache(reuse_limit=0)``) and computes the whole prompt: the + reuse-disabled control. Under the default ``all_reusable`` policy, + ``update_resources`` also commits computed generated tokens for reuse. +* Sampling is greedy: ``temperature=0``, ``top_k=1``, beam width 1, no penalties, stop + words, logits processors or draft tokens. ``TorchSampler`` chooses each token from the + step's FP32 logits (the lowest vocabulary index among exact maxima). +* Stopping follows the sampler: ``FinishReason.END_ID`` when the sampled token equals + ``end_id``, ``FinishReason.LENGTH`` after ``max_new_tokens`` sampled tokens; END_ID wins + when both hold. The last sampled token is returned and never consumed: after P prompt + tokens and G sampled tokens the cache holds P + G - 1 tokens, and no model call follows + the finishing step. An explicit token never counts as sampled. +* Isolation: each request owns its ``LlmRequest``, cache, sampler slot and ledger; the + model holds no request state. Immutable committed prefix pages may be shared between + requests at the same layer; every page a request writes is its own. +* Identity: a generator binds one model and its current weights -- each of the model's + checkpoint tensors as the very tensor object it is, with its version counter, storage + address, shape and strides. While it is open the model refuses ``load_weights``, and + every step first verifies that the model's checkpoint tensors are still exactly those + bound ones, unmodified: a tensor changed in place, replaced by another tensor (for + example ``load_state_dict(..., assign=True)`` or attribute assignment) or given other + storage is refused before the step looks up a prefix, allocates or runs a kernel, for as + long as the change persists. ``close`` releases every request and shuts the manager + down, after which a new generator starts from an empty cache and binds the weights the + model then holds. + +Limits: requests of 1 .. ``max_seq_len - max_new_tokens`` prompt tokens, ``P + +max_new_tokens <= 1024``; at most two live states; one GB300 (sm_103), eager execution only +(no CUDA graph, no overlap scheduler): every step synchronizes the host several times. +Requests are validated before any kernel runs; a request whose step failed after it began +is FAILED and can only be released. +""" + +from __future__ import annotations + +import enum +import itertools +import weakref +from dataclasses import dataclass +from typing import Optional, Sequence + +import numpy as np +import torch + +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm._torch.pyexecutor.sampler import TorchSampler +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.bindings.executor import SamplingConfig as ExecutorSamplingConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping as ParallelMapping + +from .model import ATTENTION_WINDOWS, MAX_TOKENS, NUM_LAYERS, BeamSixLayerBf16Sm103Tp1 + +#: The generator's cache: bf16, 32-token pages, a 1024-token sequence capacity, two live +#: requests, block reuse of whole pages (partial reuse off), no host tier and no +#: sliding-window scratch pages. ``max_tokens`` sizes the device quota so that two live +#: requests at full capacity fit beside the reusable pages of released ones. +GENERATOR_CACHE = { + "tokens_per_block": 32, + "max_seq_len": 1024, + "max_batch_size": 2, + "max_tokens": 4096, +} + +#: The sampler: greedy, beam width 1, no draft tokens, no overlap scheduler or async worker. +SAMPLER_ARGS = TorchSampler.Args( + max_seq_len=GENERATOR_CACHE["max_seq_len"], + max_draft_len=0, + max_num_sequences=GENERATOR_CACHE["max_batch_size"], + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + enable_async_worker=False, +) + + +def build_cache_manager(vocab_size: int = 151936) -> KVCacheManagerV2: + """The generator's six-layer ``KVCacheManagerV2`` (``GENERATOR_CACHE``), block reuse on. + + Logical layers 0-5 in slots 0-5, 8 bf16 kv heads of 128, per-layer windows + ``ATTENTION_WINDOWS`` (V2 keeps the five sliding layers and layer 5 in separate pools), + whole-page reuse under the ``all_reusable`` policy, no host tier, no scratch pages, + no overlap scheduler. The caller shuts it down. + """ + cache = GENERATOR_CACHE + return KVCacheManagerV2( + KvCacheConfig( + max_tokens=cache["max_tokens"], + max_attention_window=list(ATTENTION_WINDOWS), + enable_block_reuse=True, + enable_partial_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + ), + CacheTypeCpp.SELF, + num_layers=NUM_LAYERS, + num_kv_heads=8, + head_dim=128, + tokens_per_block=cache["tokens_per_block"], + max_seq_len=cache["max_seq_len"], + max_batch_size=cache["max_batch_size"], + mapping=ParallelMapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + vocab_size=vocab_size, + disable_overlap_scheduler=True, + ) + + +def _tensor_state(tensor: torch.Tensor) -> tuple: + """What changes when a tensor's contents or storage change: version, address, shape, strides.""" + return tensor._version, tensor.data_ptr(), tuple(tensor.shape), tensor.stride() + + +def greedy_sampling_config() -> ExecutorSamplingConfig: + """The request sampling config: temperature 0 and top_k 1 at beam width 1, nothing else set. + + ``TorchSampler`` resolves it to ``SamplingMode.GREEDY``. + """ + return ExecutorSamplingConfig(beam_width=1, top_k=1, temperature=0.0) + + +def sample_step( + sampler: TorchSampler, request: LlmRequest, logits: torch.Tensor, *, is_context: bool +) -> None: + """Choose ``request``'s next token from ``logits`` (``[vocab]`` FP32) with ``sampler``. + + Runs ``sample_async`` on a batch holding only ``request`` -- as its last context chunk + or as its one generation row -- then advances a context request's cursor to the end of + its prompt as the executor does, then ``update_requests``, which appends the chosen + token to ``request`` and finishes it when the sampler's stopping rules hold. The + request must already hold a sampler slot. + """ + batch = ScheduledRequests() + if is_context: + batch.context_requests_last_chunk = [request] + prefix_sum = [0, 1] + else: + batch.generation_requests = [request] + prefix_sum = [0] + sample_state = sampler.sample_async(batch, {"logits": logits.view(1, -1)}, prefix_sum) + if is_context: + request.py_last_context_chunk = ( + request.context_current_position, + request.context_current_position + request.context_chunk_size, + ) + request.move_to_next_context_chunk() + if request.context_remaining_length == 0: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + sampler.update_requests(sample_state) + + +class _Request(LlmRequest): + """An ``LlmRequest`` that keeps the finish reason the sampler records on it.""" + + py_finish_reason: Optional[FinishReason] = None + + def finish_by(self, reason: FinishReason, beam: int) -> None: + super().finish_by(reason, beam) + self.py_finish_reason = reason + + +class Lifecycle(enum.Enum): + NEW = "new" + PREFILLED = "prefilled" + DECODING = "decoding" + FINISHED = "finished" + RELEASED = "released" + FAILED = "failed" + + +@dataclass(frozen=True) +class StepOutput: + """One model step of one request. + + ``logits`` are the FP32 ``[151936]`` softcapped logits of the step's last token, at + absolute ``position``. ``token_id`` is the token the sampler chose from them (None for + a compute-only request) and ``finish_reason`` is set on the step that finished it. + """ + + logits: torch.Tensor + position: int + num_computed: int + token_id: Optional[int] + finish_reason: Optional[FinishReason] + + +@dataclass(frozen=True) +class GenerationOutput: + token_ids: tuple + finish_reason: FinishReason + steps: tuple + + +class RequestState: + """A handle to one request of one generator; read-only to callers. + + ``consumed``: tokens whose keys and values the cache holds. ``reused``: prompt tokens + the manager claimed from earlier requests. ``computed_prompt_tokens``: prompt tokens + the prefill computed. ``output_token_ids``: sampled tokens, in order. + ``ledger``: one record per step (positions, token source, manager claim, capacity, + history, committed tokens and page indices after the step). + """ + + def __init__( + self, + generator, + request: LlmRequest, + prompt: tuple, + *, + sample: bool, + reuse: bool, + max_new_tokens: int, + end_id: Optional[int], + ): + self._generator = weakref.ref(generator) + self._request = request + self.request_id = request.py_request_id + self.prompt = prompt + self.sample = sample + self.reuse = reuse + self.max_new_tokens = max_new_tokens + self.end_id = end_id + self.lifecycle = Lifecycle.NEW + self.consumed = 0 + self.reused = 0 + self.computed_prompt_tokens = 0 + self.explicit_tokens = 0 + self.output_token_ids: list = [] + self.pending_token: Optional[int] = None + self.finish_reason: Optional[FinishReason] = None + self.ledger: list = [] + + @property + def prompt_len(self) -> int: + return len(self.prompt) + + def __repr__(self) -> str: + return ( + f"RequestState(id={self.request_id}, {self.lifecycle.value}, consumed={self.consumed}, " + f"reused={self.reused}, output={self.output_token_ids})" + ) + + +class BeamSmallModelGenerator: + """Greedy cached generation and cross-request prefix reuse for one loaded six-block model. + + See the module docstring for the lifecycle, semantics and limits. ``model`` must be a + loaded ``BeamSixLayerBf16Sm103Tp1``; the generator builds its own + ``KVCacheManagerV2`` (``GENERATOR_CACHE``), ``SeqSlotManager`` and ``TorchSampler`` + (``SAMPLER_ARGS``) and releases them in ``close``. + """ + + _request_ids = itertools.count(1) + + def __init__(self, model: BeamSixLayerBf16Sm103Tp1): + if not isinstance(model, BeamSixLayerBf16Sm103Tp1): + raise TypeError(f"model must be BeamSixLayerBf16Sm103Tp1, got {type(model).__name__}") + if not model._loaded: + raise ValueError("load_weights must complete before a generator binds the model") + self.model = model + self.closed = False + self._weights = { + name: (tensor, _tensor_state(tensor)) + for name, tensor in model.checkpoint_tensors().items() + } + self.manager = build_cache_manager(model.vocab_size) + try: + self.seq_slots = SeqSlotManager(GENERATOR_CACHE["max_batch_size"]) + self.sampler = TorchSampler(SAMPLER_ARGS) + except BaseException: + self.manager.shutdown() + raise + self._live: dict = {} + model.register_cache_holder(self) + + # ------------------------------------------------------------------ lifecycle + + def __enter__(self) -> "BeamSmallModelGenerator": + return self + + def __exit__(self, *exc) -> None: + self.close() + + @property + def live_requests(self) -> tuple: + return tuple(self._live.values()) + + def create_request( + self, + prompt_ids, + *, + max_new_tokens: int, + end_id: Optional[int] = None, + sample: bool = True, + reuse: bool = True, + ) -> RequestState: + """Admit a new request; nothing runs on the GPU yet. + + ``prompt_ids``: a non-empty sequence (or 1-D integer tensor) of token ids in + ``[0, vocab)``, copied. ``max_new_tokens >= 1``: tokens to sample (a sampling + request) or explicit tokens to consume (``sample=False``), with + ``P + max_new_tokens <= max_seq_len``. ``end_id``: a token id that stops sampling, + or None. ``reuse``: whether the prefill may claim a cached prefix. Raises + ValueError, without side effects, on anything else or when two states are live. + """ + self._check_open() + if isinstance(prompt_ids, torch.Tensor): + if ( + prompt_ids.dim() != 1 + or prompt_ids.is_floating_point() + or prompt_ids.dtype == torch.bool + ): + raise ValueError("prompt_ids must be a 1-D integer tensor") + prompt = tuple(int(t) for t in prompt_ids.tolist()) + else: + prompt = tuple(prompt_ids) + if not all( + isinstance(t, (int, np.integer)) and not isinstance(t, bool) for t in prompt + ): + raise ValueError("prompt_ids must hold integer token ids") + prompt = tuple(int(t) for t in prompt) + vocab = self.model.vocab_size + capacity = self.manager.max_seq_len + if not prompt: + raise ValueError("the prompt must hold at least one token") + if any(not 0 <= t < vocab for t in prompt): + raise ValueError(f"prompt token ids must be in [0, {vocab})") + if type(max_new_tokens) is not int or max_new_tokens < 1: + raise ValueError(f"max_new_tokens must be a positive int, got {max_new_tokens!r}") + if len(prompt) + max_new_tokens > capacity: + raise ValueError( + f"{len(prompt)} prompt tokens + {max_new_tokens} new tokens exceed the " + f"{capacity}-token sequence capacity" + ) + if end_id is not None and (type(end_id) is not int or not 0 <= end_id < vocab): + raise ValueError(f"end_id must be None or a token id in [0, {vocab}), got {end_id!r}") + if end_id is not None and not sample: + raise ValueError("end_id stops sampling; a compute-only request samples nothing") + if len(self._live) >= GENERATOR_CACHE["max_batch_size"]: + raise ValueError( + f"{len(self._live)} requests are live; release one before creating another" + ) + request = _Request( + request_id=next(self._request_ids), + max_new_tokens=max_new_tokens, + input_tokens=list(prompt), + sampling_config=greedy_sampling_config(), + is_streaming=False, + end_id=end_id, + exclude_input_from_output=True, + ) + state = RequestState( + self, + request, + prompt, + sample=sample, + reuse=reuse, + max_new_tokens=max_new_tokens, + end_id=end_id, + ) + self._live[state.request_id] = state + return state + + def release(self, state: RequestState) -> None: + """Free the request's cache and sampler slot; its handle can no longer be used. + + Waits for issued work first. Committed pages stay reusable by later requests until + the manager evicts them. Releasing a released or foreign handle raises ValueError. + """ + self._check_open() + self._check_state(state, tuple(Lifecycle)) + torch.cuda.current_stream().synchronize() + request = state._request + try: + self.manager.free_resources(request) + if request.py_seq_slot is not None: + self.seq_slots.free_resources(request) + finally: + state.lifecycle = Lifecycle.RELEASED + del self._live[state.request_id] + + def close(self) -> None: + """Release every live request and shut the manager down. Idempotent.""" + if self.closed: + return + try: + torch.cuda.current_stream().synchronize() + for state in list(self._live.values()): + try: + self.manager.free_resources(state._request) + finally: + state.lifecycle = Lifecycle.RELEASED + self._live.clear() + finally: + self.closed = True + self.manager.shutdown() + + # ------------------------------------------------------------------ steps + + def prefill(self, state: RequestState) -> StepOutput: + """Compute the prompt tokens the cache does not hold yet, after any reused prefix. + + Claims the reusable prefix through the manager (none with ``reuse=False``), sizes + the cache for the prompt, runs the model on tokens ``reused .. P - 1`` and commits + them for reuse. A sampling request then samples its first token, which may finish + it. Returns the last prompt token's logits. + """ + self._check_open() + self._check_state(state, (Lifecycle.NEW,)) + self._check_weights() + request = state._request + try: + batch = ScheduledRequests() + batch.context_requests_last_chunk = [request] + self.seq_slots.prepare_resources(batch) + if not state.reuse: + claim = self.manager.prepare_context_cache(request, reuse_limit=0) + cache = self.manager.kv_cache_map.get(state.request_id) + if claim != 0 or cache is None or cache.num_committed_tokens != 0: + raise RuntimeError(f"the reuse-disabled control claimed {claim} tokens") + if not self.manager.prepare_context(request): + raise RuntimeError(f"the manager could not admit request {state.request_id}") + reused = request.context_current_position + num_tokens = request.context_remaining_length + if reused + num_tokens != state.prompt_len or num_tokens < 1: + raise RuntimeError( + f"the manager claimed {reused} of {state.prompt_len} prompt tokens and left " + f"{num_tokens} to compute" + ) + if not self.manager.resize_context(request, num_tokens): + raise RuntimeError(f"no cache capacity for request {state.request_id}'s prompt") + self.manager.prepare_resources(batch) + state.reused = reused + state.computed_prompt_tokens = num_tokens + logits = self._forward(state, state.prompt[reused:], reused, is_context=True) + row = logits[0, -1] + out_logits = row.clone() + token = reason = None + if state.sample: + sample_step(self.sampler, request, row, is_context=True) + else: + request.py_last_context_chunk = (reused, reused + num_tokens) + request.move_to_next_context_chunk() + request.state = LlmRequestState.GENERATION_IN_PROGRESS + self.manager.update_context_resources(batch) + state.consumed = state.prompt_len + if state.sample: + token, reason = self._after_sampling(state) + state.lifecycle = Lifecycle.FINISHED if reason is not None else Lifecycle.PREFILLED + self._record( + state, "prefill", reused, num_tokens, "prompt", token, reason, claimed=reused + ) + return StepOutput(out_logits, state.prompt_len - 1, num_tokens, token, reason) + except Exception: + state.lifecycle = Lifecycle.FAILED + raise + + def decode(self, state: RequestState, token_id: Optional[int] = None) -> StepOutput: + """Consume one token at position ``consumed`` and return the next token's logits. + + A sampling request consumes its pending sampled token (``token_id`` must be None or + equal to it) and samples the next one; a compute-only request consumes the explicit + ``token_id``. The cache grows by one token, which the step writes in every layer. + """ + self._check_open() + self._check_state(state, (Lifecycle.PREFILLED, Lifecycle.DECODING)) + request = state._request + if state.sample: + if state.pending_token is None: + raise ValueError("a sampling request has no pending token to consume") + if token_id is not None and token_id != state.pending_token: + raise ValueError( + f"a sampling request consumes its sampled token {state.pending_token}, " + f"not {token_id}" + ) + token, source = state.pending_token, "sampled" + else: + if ( + token_id is None + or type(token_id) is not int + or not 0 <= token_id < self.model.vocab_size + ): + raise ValueError( + f"a compute-only request needs an explicit token id, got {token_id!r}" + ) + if state.explicit_tokens >= state.max_new_tokens: + raise ValueError(f"the request's budget of {state.max_new_tokens} tokens is spent") + token, source = token_id, "teacher_forced" + if state.consumed + 1 > self.manager.max_seq_len: + raise ValueError(f"the request already holds {state.consumed} tokens, the capacity") + self._check_weights() + try: + if not state.sample: + request.add_new_token(token, 0) + state.explicit_tokens += 1 + past = state.consumed + if request.max_beam_num_tokens - 1 != past or request.get_last_tokens(0) != token: + raise RuntimeError( + f"request {state.request_id} holds {request.max_beam_num_tokens} tokens " + f"ending in {request.get_last_tokens(0)}; the step consumes {token} at {past}" + ) + if not self.manager.try_allocate_generation(request): + raise RuntimeError(f"no cache capacity for request {state.request_id}'s next token") + batch = ScheduledRequests() + batch.generation_requests = [request] + self.manager.prepare_resources(batch) + logits = self._forward(state, (token,), past, is_context=False) + row = logits[0, -1] + out_logits = row.clone() + state.consumed = past + 1 + state.pending_token = None + reason = sampled = None + if state.sample: + sample_step(self.sampler, request, row, is_context=False) + sampled, reason = self._after_sampling(state) + self.manager.update_resources(batch) + else: + self.manager.update_resources(batch) + cache = self.manager.kv_cache_map[state.request_id] + if not cache.resize(None, state.consumed) or cache.history_length != state.consumed: + raise RuntimeError(f"the manager did not record {state.consumed} cached tokens") + state.lifecycle = Lifecycle.FINISHED if reason is not None else Lifecycle.DECODING + self._record(state, "decode", past, 1, source, sampled, reason) + return StepOutput(out_logits, past, 1, sampled, reason) + except Exception: + state.lifecycle = Lifecycle.FAILED + raise + + def generate(self, state: RequestState) -> GenerationOutput: + """Prefill a new sampling request and decode until the sampler finishes it.""" + self._check_open() + self._check_state(state, (Lifecycle.NEW,)) + if not state.sample: + raise ValueError("generate needs a sampling request") + steps = [self.prefill(state)] + while state.lifecycle is not Lifecycle.FINISHED: + steps.append(self.decode(state)) + return GenerationOutput(tuple(state.output_token_ids), state.finish_reason, tuple(steps)) + + # ------------------------------------------------------------------ internals + + def _check_open(self) -> None: + if self.closed: + raise ValueError("the generator is closed") + + def _check_state(self, state, allowed) -> None: + if not isinstance(state, RequestState) or state._generator() is not self: + raise ValueError("the request belongs to another generator") + if self._live.get(state.request_id) is not state: + raise ValueError(f"request {state.request_id} was released") + if state.lifecycle not in allowed: + raise ValueError( + f"request {state.request_id} is {state.lifecycle.value}; this call needs " + f"{', '.join(s.value for s in allowed)}" + ) + + def _check_weights(self) -> None: + """Raise ValueError unless the model's checkpoint tensors are the bound, unmodified ones.""" + current = self.model.checkpoint_tensors() + if current.keys() != self._weights.keys(): + raise ValueError( + "the model's tensors changed after the generator bound it; its cache no longer " + "matches the weights. Close the generator and build a new one" + ) + for name, tensor in current.items(): + bound, state = self._weights[name] + if tensor is not bound or _tensor_state(tensor) != state: + raise ValueError( + f"{name} changed after the generator bound the model; its cache no longer " + "matches the weights. Close the generator and build a new one" + ) + + def _metadata(self, state: RequestState, num_tokens: int, past: int, is_context: bool): + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=MAX_TOKENS, + kv_cache_manager=self.manager, + num_contexts=1 if is_context else 0, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[past]), + seq_lens=torch.tensor([num_tokens], dtype=torch.int), + request_ids=[state.request_id], + prompt_lens=[num_tokens if is_context else state.prompt_len], + runtime_features=AttentionRuntimeFeatures(cache_reuse=True), + ) + md.prepare() + return md + + def _forward(self, state: RequestState, tokens: Sequence[int], past: int, is_context: bool): + num_tokens = len(tokens) + device = self.model.device + md = self._metadata(state, num_tokens, past, is_context) + ids = torch.tensor([list(tokens)], dtype=torch.int64, device=device) + positions = torch.arange(past, past + num_tokens, dtype=torch.int64, device=device).view( + 1, -1 + ) + return self.model(ids, positions, md) + + def _after_sampling(self, state: RequestState): + """Read the sampler's choice and stopping decision off the request.""" + request = state._request + token = int(request.get_last_tokens(0)) + state.output_token_ids.append(token) + reason = None + if request.is_finished: + reason = request.py_finish_reason + if reason not in (FinishReason.END_ID, FinishReason.LENGTH): + raise RuntimeError(f"request {state.request_id} finished with reason {reason}") + state.finish_reason = reason + else: + state.pending_token = token + return token, reason + + def _record(self, state, step, past, num_tokens, source, token, reason, claimed=None) -> None: + cache = self.manager.kv_cache_map.get(state.request_id) + record = { + "step": step, + "past": past, + "tokens": num_tokens, + "positions": [past, past + num_tokens - 1], + "input": source, + "consumed": state.consumed, + "sampled": token, + "finish_reason": None if reason is None else reason.name, + "seq_slot": state._request.py_seq_slot, + } + if claimed is not None: + record["claimed"] = claimed + if cache is not None: + num_blocks = int(cache.num_blocks) + record.update( + capacity=int(cache.capacity), + history=int(cache.history_length), + committed=int(cache.num_committed_tokens), + pages=[ + np.frombuffer( + cache.get_base_page_indices(pool), dtype=np.int32, count=num_blocks + ).tolist() + for pool in range(len(self.manager.impl.layer_grouping)) + ], + ) + state.ledger.append(record) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/model.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/model.py new file mode 100644 index 000000000000..0ac7c2303bc7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/six_layer_model/model.py @@ -0,0 +1,685 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam six-block small model / sm_103 / tp1: token ids to logits, one request step per call. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no registered architecture: nothing in the engine builds this class, and it serves no +request on its own. A caller constructs it from a six-block configuration of the +checkpoint -- the checkpoint's own fields with ``num_hidden_layers = 6`` and the first six +entries of its two layer schedules -- loads its 2394 tensors and runs one step of one BF16 +request through it, with the attention batch state prepared on a six-layer +``KVCacheManagerV2``: ``prefill_context`` builds a fresh one for a single context, and the +stateful generator in ``generation.py`` keeps one across requests for cached decoding and +prefix reuse. It does not call the built-in Beam modeling. + +The model, all blocks with their own parameters and logical index:: + + x = embed_tokens[input_ids] [1, T, 5120] bf16, no scale or norm + x = block_0(x); x = block_1(x); x = block_2(x) dense, sliding attention + RoPE + x = block_3(x); x = block_4(x) sparse, sliding attention + RoPE + x = block_5(x) sparse, full attention, no RoPE + h = rmsnorm(x, norm) fp32 statistics and affine, one rounding + head = h @ lm_head^T [T, 151936] bf16, untied, no bias + logits = tanh(head.float() / 30) * 30 [1, T, 151936] fp32 + +The blocks are the dense and sparse block components, built for logical layers 0-2 and +3-5; each block's residual coefficients, attention kind and router are its own layer's. + +Steps. A call is one step of one request: a context step computes the T tokens of the +request's prompt that are not yet in its cache, after the ``past`` tokens that are (none +for a fresh context; a prefix reused from another request or computed by an earlier step +otherwise), and a generation step appends one token. The tokens sit at absolute positions +``past .. past + T - 1``; the returned rows are those tokens' logits. + +Cache. One call appends every layer's keys and values to one ``KVCacheManagerV2`` that +holds exactly logical layers 0-5, each in its own local slot: ``layer_offsets`` maps +logical layer i to slot i, and slot i's rows in the pool tables belong to it alone. The +five sliding layers attend through a 512-token window and layer 5 to the whole history; +the manager states this as its per-layer window vector ``[512] * 5 + [None]`` (``None`` is +the full ``max_seq_len``), and the attention op of each layer is called with that layer's +window. Before the first layer runs, ``check_cache`` verifies the manager, its slots and +the host pool tables the attention op reads -- that they address each slot's own key and +value pages as the manager allocated them -- and ``check_request`` verifies the request: +that the manager holds an active cache for it whose history ends exactly where the step +begins (a step never rewrites a cached token) and with the capacity the step needs, that +every page the step reads or writes is a valid page of its pool, and that the block-offset +table the op reads holds exactly the manager's addresses of those pages. Earlier tokens' keys +and values are read from the pool on a cached context step (paged-context FMHA) and on +every generation step. + +Host synchronization per call: one for the model's token-id range, position and +block-offset check, one for each block's position check, and each sparse block's MoE +dispatch (two, plus two per selected expert). A call is therefore not capturable in a CUDA +graph. +""" + +from __future__ import annotations + +import contextlib +import weakref +from collections import OrderedDict +from typing import Iterator, Mapping + +import numpy as np +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.beam_logit_softcap import ( + beam_logit_softcap, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp +from tensorrt_llm.bindings import DataType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping as ParallelMapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, LayerId, PageIndexMode + +from ..dense_block import block as dense_block +from ..moe.moe import check_weights +from ..moe.moe import weight_shapes as moe_weight_shapes +from ..sparse_block import block as sparse_block + +NUM_LAYERS = 6 +DENSE_LAYERS = (0, 1, 2) +SPARSE_LAYERS = (3, 4, 5) +#: The six-block schedules: the first six entries of the checkpoint's. +LAYER_TYPES = ("sliding_attention",) * 5 + ("full_attention",) +MLP_LAYER_TYPES = ("dense",) * 3 + ("sparse",) * 3 + +#: The configuration this component computes, beyond what its blocks check themselves. +#: Anything else is rejected at construction rather than computed approximately. +SUPPORTED_CONFIG = { + "num_hidden_layers": NUM_LAYERS, + "hidden_size": 5120, + "vocab_size": 151936, + "rms_norm_eps": 1e-6, + "tie_word_embeddings": False, + "final_logit_softcapping": 30.0, +} + +#: The KV cache ``prefill_context`` builds: bf16, 32-token pages, a 1024-token sequence +#: capacity and one request. A window equal to ``max_seq_len`` is the whole context. +KV_CACHE = {"tokens_per_block": 32, "max_seq_len": 1024, "max_tokens": 1024} +ATTENTION_WINDOWS = (512,) * 5 + (KV_CACHE["max_seq_len"],) + +#: The most tokens one call may compute: the sequence capacity. A request's whole KV +#: length is bounded by its cache's ``max_seq_len``. +MAX_TOKENS = KV_CACHE["max_seq_len"] + + +#: Checkpoint names of the model's own tensors, beside its blocks'. +EMBED = "model.embed_tokens.weight" +NORM = "model.norm.weight" +HEAD = "lm_head.weight" + + +def block_prefix(layer_idx: int) -> str: + return f"model.layers.{layer_idx}." + + +def check_config(config) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this model computes {expected!r}") + for field, expected in (("layer_types", LAYER_TYPES), ("mlp_layer_types", MLP_LAYER_TYPES)): + actual = getattr(config, field, None) + if actual is None or tuple(actual) != expected: + raise ValueError( + f"unsupported {field}={actual!r}; this model computes {list(expected)}" + ) + for layer_idx in DENSE_LAYERS: + dense_block.check_config(config, layer_idx) + for layer_idx in SPARSE_LAYERS: + sparse_block.check_config(config, layer_idx) + + +def weight_schema(config) -> "OrderedDict[str, tuple[tuple[int, ...], torch.dtype]]": + """Every tensor by checkpoint name, in checkpoint order, with its shape and dtype. + + 2394 tensors: the embedding, each block's (15 per dense block, 782 per sparse block: + 12 attention and norm tensors and the MoE's FP32 router and 256 BF16 experts), the + final norm and the head. Matrices are [out, in]; the embedding is [vocab, hidden]. + """ + vocab, hidden = config.vocab_size, config.hidden_size + schema = OrderedDict([(EMBED, ((vocab, hidden), torch.bfloat16))]) + for layer_idx in range(NUM_LAYERS): + prefix = block_prefix(layer_idx) + if layer_idx in DENSE_LAYERS: + names = dense_block.weight_shapes(config).items() + else: + names = sparse_block.weight_shapes(config).items() + for name, shape in names: + schema[prefix + name] = (shape, torch.bfloat16) + if layer_idx in SPARSE_LAYERS: + for name, shape_dtype in moe_weight_shapes(config).items(): + schema[prefix + sparse_block.MLP_PREFIX + name] = shape_dtype + schema[NORM] = ((hidden,), torch.bfloat16) + schema[HEAD] = ((vocab, hidden), torch.bfloat16) + return schema + + +def check_cache(md: TrtllmAttentionMetadata, kv_len: int) -> None: + """Raise ValueError unless the call's cache holds the six layers in distinct valid slots. + + Required, before any layer writes: a ``KVCacheManagerV2`` holding exactly logical + layers 0-5, mapped one-to-one onto its six local slots; 8 bf16 kv heads of 128 in + every slot; a 512-token window on the sliding layers and the whole context on layer + 5, the full window holding the request's ``kv_len`` tokens; keys and values in one + shared page table per pool; host pool tables that address, for every slot, exactly + the key and value pages the manager allocated to that slot, so that no two slots + share storage (``_check_pool_tables``); and an int32 CUDA block-offset table of the + shape the attention op indexes. Reads host-side manager state and tensor metadata only. + """ + mgr = getattr(md, "kv_cache_manager", None) + if not isinstance(mgr, KVCacheManagerV2): + raise ValueError(f"the KV cache must be a KVCacheManagerV2, got {type(mgr).__name__}") + offsets = dict(getattr(mgr, "layer_offsets", None) or {}) + if set(offsets) != set(range(NUM_LAYERS)) or mgr.num_local_layers != NUM_LAYERS: + raise ValueError( + f"the KV cache must hold exactly logical layers 0-{NUM_LAYERS - 1}; it holds " + f"{sorted(offsets)} in {mgr.num_local_layers} local slots" + ) + slots = [offsets[layer_idx] for layer_idx in range(NUM_LAYERS)] + if any(type(slot) is not int for slot in slots) or sorted(slots) != list(range(NUM_LAYERS)): + raise ValueError(f"the six layers must map one-to-one onto slots 0-5, got {slots}") + heads = list(getattr(mgr, "num_kv_heads_per_layer", None) or []) + head_dims = list(getattr(mgr, "head_dim_per_layer", None) or []) + if ( + mgr.dtype != DataType.BF16 + or heads != [8] * NUM_LAYERS + or (head_dims and head_dims != [128] * NUM_LAYERS) + or mgr.head_dim != 128 + ): + raise ValueError( + f"every slot must hold 8 bf16 kv heads of 128; got {mgr.dtype}, {heads}, " + f"head_dim {mgr.head_dim}" + ) + windows = list(mgr.max_attention_window_vec) + for layer_idx, slot in enumerate(slots): + window = windows[slot] + if LAYER_TYPES[layer_idx] == "sliding_attention": + retained = mgr.max_seq_len if window is None else window + if retained != ATTENTION_WINDOWS[layer_idx]: + raise ValueError( + f"sliding layer {layer_idx} must keep a {ATTENTION_WINDOWS[layer_idx]}-token " + f"window in slot {slot}; the manager keeps {retained}" + ) + elif window is not None: + raise ValueError( + f"full-attention layer {layer_idx} must keep the whole context in slot {slot}; " + f"the manager keeps a {window}-token window" + ) + if mgr.max_seq_len < kv_len or md.max_seq_len != mgr.max_seq_len: + raise ValueError( + f"the cache holds {mgr.max_seq_len} tokens per sequence and the metadata states " + f"{md.max_seq_len}; the request reaches {kv_len}" + ) + _check_pool_tables(md, mgr, slots) + offsets = md.kv_cache_block_offsets + if ( + not isinstance(offsets, torch.Tensor) + or offsets.device.type != "cuda" + or offsets.dtype != torch.int32 + or offsets.dim() != 4 + or tuple(offsets.shape[i] for i in (0, 2, 3)) + != (mgr.num_attention_op_pools, 2, mgr.max_blocks_per_seq) + or offsets.shape[1] < 1 + ): + raise ValueError( + "the block-offset table must be a CUDA torch.int32 [pools, sequences, 2, " + f"{mgr.max_blocks_per_seq}] tensor with {mgr.num_attention_op_pools} pools" + ) + + +def _pool_windows(mgr: KVCacheManagerV2) -> dict: + """Each pool's attention window, from the windows of the slots it holds (``None``: full).""" + windows: dict = {} + for slot in range(NUM_LAYERS): + pool = int(mgr.impl.get_layer_group_id(LayerId(slot))) + window = mgr.max_attention_window_vec[slot] + if windows.setdefault(pool, window) != window: + raise ValueError(f"pool {pool} holds slots with windows {windows[pool]} and {window}") + return windows + + +def check_request(md: TrtllmAttentionMetadata, num_tokens: int, kv_len: int): + """Raise ValueError unless the call continues its request's cache on valid pages. + + The call reads keys and values of earlier tokens and writes its own ``num_tokens`` + after ``kv_len - num_tokens`` cached ones. Required, from the manager's own state: one + request id, whose cache in the manager is active, holds exactly those ``kv_len - + num_tokens`` tokens (its ``history_length``: the call appends right after them, so it + never rewrites a token the cache holds -- committed or reused prefix pages included, + which the manager keeps within the history -- and leaves no unwritten gap) and has + capacity for ``kv_len`` tokens; and for every pool, a valid page of that pool -- not + ``BAD_PAGE_INDEX``, inside the pool -- at every block ordinal a token of the call can + read or write: on the full layer's pool every ordinal below ``kv_len``, on the sliding + layers' pool those from the oldest key the call's first query attends to. Ordinals + older than that may carry the manager's out-of-window sentinel. + + Returns ``(expected, live)``: host int32 ``[pools, 2, n]`` block-offset entries the op + must read for those pages (the manager's page-to-offset conversion of each page) and + a bool ``[pools, n]`` mask of the ordinals checked, ``n = ceil(kv_len / tokens_per_block)``. + The caller compares them with the CUDA table. Reads host-side manager state only. + """ + mgr = md.kv_cache_manager + request_ids = list(md.request_ids or ()) + if len(request_ids) != 1: + raise ValueError(f"one request per call is supported; metadata names {request_ids}") + (request_id,) = request_ids + kv_cache = mgr.kv_cache_map.get(request_id) + if kv_cache is None or not kv_cache.is_active: + raise ValueError(f"request {request_id} holds no active cache in this manager") + history = int(kv_cache.history_length) + if history != kv_len - num_tokens: + raise ValueError( + f"request {request_id}'s cache holds {history} tokens; the call writes tokens " + f"{kv_len - num_tokens}..{kv_len - 1}, but a step must append right after the tokens " + "the cache holds" + ) + if kv_cache.capacity < kv_len: + raise ValueError( + f"request {request_id} has capacity for {kv_cache.capacity} tokens; the call reaches " + f"{kv_len}" + ) + tokens_per_block = mgr.tokens_per_block + num_blocks = -(-kv_len // tokens_per_block) + held = int(kv_cache.num_blocks) + if held < num_blocks: + raise ValueError(f"request {request_id} holds {held} blocks; the call needs {num_blocks}") + first_query = kv_len - num_tokens + pools = _pool_windows(mgr) + num_pools = len(mgr.impl.layer_grouping) + if sorted(pools) != list(range(num_pools)): + raise ValueError(f"the six slots occupy pools {sorted(pools)} of {num_pools}") + expected = torch.zeros((num_pools, 2, num_blocks), dtype=torch.int32) + live = torch.zeros((num_pools, num_blocks), dtype=torch.bool) + for pool, window in sorted(pools.items()): + oldest = 0 if window is None else max(0, first_query - window + 1) + pages = np.frombuffer(kv_cache.get_base_page_indices(pool), dtype=np.int32, count=held) + layer = next( + LayerId(slot) + for slot in range(NUM_LAYERS) + if int(mgr.impl.get_layer_group_id(LayerId(slot))) == pool + ) + scale, value_offset = int(mgr.index_scales[pool]), int(mgr.kv_offset[pool]) + bound = int(mgr.impl.get_page_index_upper_bound(layer, Role.KEY)) + for ordinal in range(oldest // tokens_per_block, num_blocks): + page = int(pages[ordinal]) + if page == BAD_PAGE_INDEX or page < 0 or scale * page + value_offset >= bound: + raise ValueError( + f"request {request_id} has page {page} at block {ordinal} of pool {pool}; " + f"the call reads or writes that block and the pool holds pages below " + f"{bound // scale}" + ) + expected[pool, 0, ordinal] = scale * page + expected[pool, 1, ordinal] = scale * page + value_offset + live[pool, ordinal] = True + return expected, live + + +def _overlap(spans) -> bool: + spans = sorted(spans) + return any(end > start for (_, end), (start, _) in zip(spans, spans[1:])) + + +def _check_pool_tables(md: TrtllmAttentionMetadata, mgr: KVCacheManagerV2, slots) -> None: + """Raise ValueError unless the host pool tables address each slot's own cache storage. + + The attention op reads slot s's row of ``host_kv_cache_pool_mapping`` as int32 + ``(pool, offset)`` and finds the slot's keys at ``host_kv_cache_pool_pointers[pool, 0] + + offset * 2 * page``, a page holding one slot's keys of ``tokens_per_block`` tokens and + the next page their values, with the pool's page indices from its block offsets. The + tables are host tensors a caller can replace, so each row is checked against the + manager's allocation itself: the pool is the slot's layer group, the offset keeps the + slot's key and value pages inside one slot of that pool, the address is where the + manager keeps the slot's keys, its values follow one page later, and no two slots' + pages, and no two pools, overlap. The five sliding layers and layer 5 are layer groups + of different slot sizes, which the manager keeps in separate pools; a cache whose layer + groups shared one pool is rejected rather than reasoned about. + """ + impl = mgr.impl + num_pools = len(impl.layer_grouping) + mapping = md.host_kv_cache_pool_mapping + pointers = md.host_kv_cache_pool_pointers + if ( + not isinstance(mapping, torch.Tensor) + or mapping.device.type != "cpu" + or mapping.dtype != torch.int32 + or tuple(mapping.shape) != (NUM_LAYERS, 2) + ): + raise ValueError(f"the pool mapping must be a host torch.int32 [{NUM_LAYERS}, 2] table") + if ( + not isinstance(pointers, torch.Tensor) + or pointers.device.type != "cpu" + or pointers.dtype != torch.int64 + or tuple(pointers.shape) != (num_pools, 2) + ): + raise ValueError( + f"the pool pointers must be a host torch.int64 [{num_pools}, 2] table, one row per pool" + ) + if ( + mgr.kv_cache_type != CacheTypeCpp.SELF + or mgr.enable_swa_scratch_reuse + or mgr.num_attention_op_pools != num_pools + ): + raise ValueError( + "the cache must hold keys and values with one shared page table per pool, " + "without sliding-window scratch reuse" + ) + page = mgr.tokens_per_block * 8 * 128 * torch.bfloat16.itemsize + pairs, pools = [], {} + for layer_idx, slot in enumerate(slots): + layer = LayerId(slot) + pool, offset = mapping[slot].tolist() + group = int(impl.get_layer_group_id(layer)) + scale = impl.get_page_index_scale(layer, Role.KEY) + if pool != group or not 0 <= offset < scale // 2: + raise ValueError( + f"layer {layer_idx} has pool row ({pool}, {offset}); the manager keeps slot {slot} " + f"in pool {group}, whose slots hold {scale // 2} key and value page pairs" + ) + start = impl.get_mem_pool_base_address(layer, Role.KEY, PageIndexMode.PER_LAYER) + key = impl.get_mem_pool_base_address(layer, Role.KEY, PageIndexMode.SHARED) + value = impl.get_mem_pool_base_address(layer, Role.VALUE, PageIndexMode.SHARED) + addressed = int(pointers[pool, 0]) + offset * 2 * page + if ( + addressed != key + or int(pointers[pool, 1]) != 0 + or value != key + page + or impl.get_page_stride(layer, Role.KEY) != page + or impl.get_page_stride(layer, Role.VALUE) != page + or not start <= key <= key + 2 * page <= start + scale * page + ): + raise ValueError( + f"layer {layer_idx}'s pool row ({pool}, {offset}) with pool pointers " + f"{pointers[pool].tolist()} addresses {addressed:#x}; the manager keeps slot " + f"{slot}'s keys at {key:#x} and values at {value:#x} in pool {pool} from " + f"{start:#x}, in {page}-byte pages" + ) + extent = (start, key + impl.get_page_index_upper_bound(layer, Role.KEY) * page) + if pools.setdefault(pool, extent) != extent: + raise ValueError( + f"pool {pool} spans {pools[pool]} for one slot and {extent} for another" + ) + pairs.append((key, key + 2 * page)) + if _overlap(pairs) or _overlap(pools.values()): + raise ValueError( + f"cache slots share storage: key and value pages {pairs}, pools {sorted(pools.values())}" + ) + + +@contextlib.contextmanager +def prefill_context(num_tokens: int, request_id: int = 0) -> Iterator[TrtllmAttentionMetadata]: + """A fresh six-layer ``KVCacheManagerV2`` and prepared metadata for one new request. + + The manager holds logical layers 0-5 in slots 0-5 (bf16, 8 kv heads of 128, + ``KV_CACHE`` pages and capacity, windows ``ATTENTION_WINDOWS``) with block reuse, + partial reuse, sliding-window scratch reuse and the overlap scheduler off; the + metadata holds one context request of ``num_tokens`` tokens starting at an empty + cache. The manager is shut down on exit, also after an error. + """ + mgr = KVCacheManagerV2( + KvCacheConfig( + max_tokens=KV_CACHE["max_tokens"], + max_attention_window=list(ATTENTION_WINDOWS), + enable_block_reuse=False, + enable_partial_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheTypeCpp.SELF, + num_layers=NUM_LAYERS, + num_kv_heads=8, + head_dim=128, + tokens_per_block=KV_CACHE["tokens_per_block"], + max_seq_len=KV_CACHE["max_seq_len"], + max_batch_size=1, + mapping=ParallelMapping(world_size=1, tp_size=1, rank=0), + dtype=DataType.BF16, + disable_overlap_scheduler=True, + ) + try: + mgr.add_dummy_requests([request_id], token_nums=[num_tokens]) + md = TrtllmAttentionMetadata( + max_num_requests=1, + max_num_tokens=MAX_TOKENS, + kv_cache_manager=mgr, + num_contexts=1, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]), + seq_lens=torch.tensor([num_tokens], dtype=torch.int), + request_ids=[request_id], + prompt_lens=[num_tokens], + ) + md.prepare() + yield md + finally: + mgr.shutdown() + + +class BeamSixLayerBf16Sm103Tp1(nn.Module): + """Beam six-block model: token ids in, every token's FP32 vocabulary logits out. + + State: ``embed_tokens`` ``[151936, 5120]`` and ``lm_head`` ``[151936, 5120]``, two + separate BF16 tensors; ``layers``, dense blocks 0-2 and sparse blocks 3-5, each with + its own tensors, router and 256 experts; ``norm`` ``[5120]`` BF16. ``load_weights`` + loads all 2394 tensors by checkpoint name; the model runs only after that succeeded. + + ``forward(input_ids [1, T] int32/int64, position_ids [1, T], attn_metadata)`` returns + ``[1, T, 151936]`` ``torch.float32`` on the model's device: row t, entry v is the + softcapped logit for vocabulary id v of the t-th token of the call. Accepted: + 1 <= T <= ``MAX_TOKENS`` token ids in ``[0, 151936)``, contiguous on the model's device; + the prepared ``TrtllmAttentionMetadata`` of one step of one request -- a context step + of T tokens after ``past`` cached ones (paged-context FMHA when ``past > 0``) or a + one-token generation step -- on the six-layer cache ``check_cache`` describes, its + request's pages valid as ``check_request`` describes and its KV length at most the + cache's ``max_seq_len``; and integer positions equal to ``arange(past, past + T)``. + ``prefill_context(T)`` yields such metadata for a fresh context; the generator in + ``generation.py`` for every step of its requests. Anything else is rejected with + ValueError before the embedding lookup or any other kernel runs, including an id + outside the vocabulary. Inputs and parameters are never modified. Batches of several + requests and padding are not computed here. + + Weights cannot change under a cache built from them: ``load_weights`` raises while a + generator holding this model's cache is open (``close`` it first). + """ + + def __init__(self, config, device: str | torch.device = "cuda"): + super().__init__() + check_config(config) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this model is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.eps = config.rms_norm_eps + self.logit_softcap = config.final_logit_softcapping + self._schema = weight_schema(config) + + def empty(name: str) -> nn.Parameter: + shape, dtype = self._schema[name] + return nn.Parameter(torch.empty(shape, dtype=dtype, device=device), requires_grad=False) + + self.embed_tokens = empty(EMBED) + self.layers = nn.ModuleList( + dense_block.BeamDenseBlockBf16Sm103Tp1(config, layer_idx, device=device) + if layer_idx in DENSE_LAYERS + else sparse_block.BeamSparseBlockBf16Sm103Tp1(config, layer_idx, device=device) + for layer_idx in range(NUM_LAYERS) + ) + self.norm = empty(NORM) + self.lm_head = empty(HEAD) + self._loaded = False + #: Generators whose KV cache holds keys and values computed with these weights. + self._cache_holders: "weakref.WeakSet" = weakref.WeakSet() + + @property + def device(self) -> torch.device: + return self.embed_tokens.device + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """All 2394 tensors by checkpoint name, in checkpoint order.""" + tensors = OrderedDict([(EMBED, self.embed_tokens)]) + for layer_idx, block in enumerate(self.layers): + for name, tensor in block.checkpoint_tensors().items(): + tensors[block_prefix(layer_idx) + name] = tensor + tensors[NORM] = self.norm + tensors[HEAD] = self.lm_head + return tensors + + def check_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Raise unless ``weights`` holds exactly the 2394 names with their dtypes and shapes. + + Reads tensor metadata only, never values, and copies nothing. + """ + check_weights(self._schema, weights, "model") + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the model's 2394 tensors, by checkpoint name, [out, in]. + + Names are ``model.embed_tokens.weight``, ``model.layers..``, + ``model.norm.weight`` and ``lm_head.weight``; dtypes are BF16 except each sparse + block's FP32 ``mlp.gate.weight`` and ``mlp.gate.e_score_correction_bias``; host or + CUDA tensors are copied as given, without transposing or converting. Every tensor + is checked before any parameter changes, so a rejected load leaves the model as it + was. The model runs only after a load completed: a failure while copying leaves it + unloaded until a later load succeeds. Copies proceed one tensor (one expert, one + block's attention tensors) at a time; no copy of a whole bank is made. + + Raises RuntimeError, before checking or copying anything, while an open generator + holds a KV cache computed with the current weights: its cached and reusable keys + and values would no longer match them. + """ + holders = [holder for holder in self._cache_holders if not holder.closed] + if holders: + raise RuntimeError( + f"{len(holders)} open generator(s) hold KV caches computed with these weights; " + "close them before loading new weights" + ) + self.check_weights(weights) + self._loaded = False + for name, param in ((EMBED, self.embed_tokens), (NORM, self.norm), (HEAD, self.lm_head)): + param.copy_(weights[name]) + for layer_idx, block in enumerate(self.layers): + prefix = block_prefix(layer_idx) + block.load_weights( + {name[len(prefix) :]: w for name, w in weights.items() if name.startswith(prefix)} + ) + self._loaded = True + + def register_cache_holder(self, holder) -> None: + """Record a generator whose KV cache holds keys and values computed with these weights. + + ``holder`` must expose a ``closed`` attribute; while it is False ``load_weights`` + raises. Holders are referenced weakly. + """ + self._cache_holders.add(holder) + + def _check_call( + self, input_ids: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> int: + """Raise ValueError on anything the model would otherwise compute wrongly; return T.""" + if not self._loaded: + raise ValueError("load_weights must complete before the model runs") + if ( + not isinstance(input_ids, torch.Tensor) + or input_ids.dtype not in (torch.int32, torch.int64) + or input_ids.layout != torch.strided + ): + raise ValueError("input_ids must be a dense torch.int32 or torch.int64 tensor") + if input_ids.dim() != 2 or input_ids.shape[0] != 1: + raise ValueError(f"input_ids must be [1, T], got {tuple(input_ids.shape)}") + num_tokens = input_ids.shape[1] + if not 1 <= num_tokens <= MAX_TOKENS: + raise ValueError( + f"supported token counts per call are 1..{MAX_TOKENS}, got {num_tokens}" + ) + if input_ids.device != self.device or not input_ids.is_contiguous(): + raise ValueError(f"input_ids must be contiguous on the model's device {self.device}") + if ( + not isinstance(position_ids, torch.Tensor) + or position_ids.dtype not in (torch.int32, torch.int64) + or tuple(position_ids.shape) != (1, num_tokens) + or position_ids.device != self.device + or not position_ids.is_contiguous() + ): + raise ValueError( + f"position_ids must be a contiguous integer [1, {num_tokens}] tensor on {self.device}" + ) + if not isinstance(md, TrtllmAttentionMetadata): + raise ValueError( + f"attn_metadata must be TrtllmAttentionMetadata, got {type(md).__name__}" + ) + mgr = getattr(md, "kv_cache_manager", None) + if not isinstance(mgr, KVCacheManagerV2): + raise ValueError(f"the KV cache must be a KVCacheManagerV2, got {type(mgr).__name__}") + pasts = { + _attention.check_metadata( + md, num_tokens, layer_idx=layer_idx, num_kv_heads=8, head_dim=128 + ) + for layer_idx in range(NUM_LAYERS) + } + (past,) = pasts + kv_len = past + num_tokens + check_cache(md, kv_len) + expected_table, live = check_request(md, num_tokens, kv_len) + table = md.kv_cache_block_offsets[:, 0, :, : expected_table.shape[-1]] + table_ok = ( + (table == expected_table.to(table.device, non_blocking=True)) + | ~live.to(table.device, non_blocking=True)[:, None, :] + ).all() + expected = torch.arange( + past, past + num_tokens, dtype=position_ids.dtype, device=self.device + ) + low, high = torch.aminmax(input_ids) + positions_ok, ids_ok, table_ok = torch.stack( + ( + (position_ids[0] == expected).all(), + (low >= 0) & (high < self.vocab_size), + table_ok, + ) + ).tolist() + if not positions_ok: + raise ValueError( + f"the call appends tokens {past}..{kv_len - 1} after {past} cached tokens; its " + "positions must be exactly those" + ) + if not table_ok: + raise ValueError( + "the block-offset table does not address the pages the manager allocated to " + "the request" + ) + if not ids_ok: + raise ValueError( + f"token ids must be in [0, {self.vocab_size}); got values in " + f"[{int(low)}, {int(high)}]" + ) + return num_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + num_tokens = self._check_call(input_ids, position_ids, attn_metadata) + x = nn.functional.embedding(input_ids, self.embed_tokens) + for block in self.layers: + x = block(x, position_ids, attn_metadata) + h = flashinfer_rmsnorm(x.view(num_tokens, self.hidden_size), self.norm, self.eps) + head = cublas_mm(h, self.lm_head.t()) + logits = beam_logit_softcap(head, self.logit_softcap) + return logits.view(1, num_tokens, self.vocab_size) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/__init__.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/__init__.py new file mode 100644 index 000000000000..c10dd6cb1a2e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder block 3 (sparse MoE, sliding attention) / sm_103 / tp1.""" diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/block.py b/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/block.py new file mode 100644 index 000000000000..bb82eb5bc4ec --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_components/sparse_block/block.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam decoder blocks 3-5 / sm_103 / tp1: one sparse (MoE) block, constructed directly. + +This TP1 component uses FlashInfer normalization/activation kernels; see ``../README.md`` +for its numerical differences from the TP4 serving model. + +A component, not a routed target. There is no ``routing.py``, no router-index entry and +no ``ForCausalLM``: nothing in the engine builds this class, and it serves no request. +A caller constructs it from the checkpoint's configuration for one of the sparse blocks +3, 4 or 5, loads the block's 782 tensors and runs one step of one BF16 request through it +-- a context step after any cached prefix, or a one-token generation step -- with the +attention batch state prepared on a ``KVCacheManagerV2`` that holds the block's logical +layer. The forward is assembled from catalog entries, the attention sublayer +shared with the dense block component (``../_attention.py``) and the MoE component +(``../moe``); it does not call the built-in Beam modeling and owns no +dense MLP. + +The block, with block 3's residual coefficients (block i's are in ``depth_scales``):: + + h = rmsnorm(x, input_layernorm) + a = attention(h) see ../_attention.py + x = rmsnorm(a, post_attention_layernorm) / sqrt(8) + x * sqrt(7) / sqrt(8) + h = rmsnorm(x, pre_feedforward_layernorm) + m = moe(h) routed scaling 2.5 included + out = rmsnorm(m, post_feedforward_layernorm) / sqrt(9) + x * sqrt(8) / sqrt(9) + +Blocks 3 and 4 attend through the 512-token sliding window with NeoX RoPE; block 5 is a +full-attention layer without position embedding: no rotation of q or k, the per-token +query scale on q, and the whole context as its window. + +Every RMSNorm computes its statistics, normalization and affine in fp32 and rounds once; +each residual add rounds both scaled terms to bf16, then their sum. The residual +coefficients are the logical block's: block i's attention add is the 2(i + 1)-th add of +the residual stream, its feed-forward add the next one. + +Logical layer and cache slot: ``layer_idx`` fixes the block's semantics; the cache slot is +the call's manager's local index for that layer -- 0 for a manager whose ``layer_mask`` +selects only it, its own position in a manager holding several layers. A manager that +does not hold the layer -- a default one-layer manager holds layer 0 -- is rejected +before any kernel runs. + +Host synchronization per call: one for the position check, and the MoE's dispatch +synchronizations (two, plus two per selected expert). Positions are absolute: a sliding +layer attends to the newest 512 keys of the request's history up to each query, the +full-attention layer to all of them. +""" + +from __future__ import annotations + +from collections import OrderedDict +from typing import Mapping + +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.beam import _attention +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._constants import SM as _SM +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.residual.beam_depth_residual import ( + beam_depth_residual, + depth_scales, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +from ..moe.moe import BeamMoeBf16Sm103Tp1, check_weights +from ..moe.moe import check_config as check_moe_config + +#: The configuration this component computes, as the checkpoint's config.json states it. +#: Anything else is rejected at construction rather than computed approximately; the +#: MoE component checks its own fields. +SUPPORTED_CONFIG = { + "hidden_size": 5120, + "num_attention_heads": 48, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-6, + "attention_bias": False, + "attention_dropout": 0.0, + "sliding_window": 512, +} +DEFAULT_LAYER_IDX = 3 +#: The sparse blocks of the checkpoint's schedule and the attention each one runs. +ATTENTION_TYPES = {3: "sliding_attention", 4: "sliding_attention", 5: "full_attention"} +SUPPORTED_LAYER_INDICES = tuple(ATTENTION_TYPES) +#: RoPE of the sliding layers; the full-attention layers have none. +SUPPORTED_ROPE = {"rope_type": "default", "rope_theta": 10000.0} + +#: The most tokens one call may append: the six-layer model's sequence capacity. Longer +#: calls are rejected; the request's whole KV length is bounded by its cache's capacity. +MAX_TOKENS = 1024 + + +#: The MoE's tensors carry this prefix in the block's checkpoint names. +MLP_PREFIX = "mlp." + + +def weight_shapes(config) -> "OrderedDict[str, tuple[int, ...]]": + """The block's twelve attention and norm tensors by checkpoint name, [out, in], all BF16.""" + hidden = config.hidden_size + q_width = config.num_attention_heads * config.head_dim + kv_width = config.num_key_value_heads * config.head_dim + return OrderedDict( + [ + ("input_layernorm.weight", (hidden,)), + ("self_attn.gate_bias", (config.num_attention_heads,)), + ("self_attn.q_proj.weight", (q_width, hidden)), + ("self_attn.k_proj.weight", (kv_width, hidden)), + ("self_attn.v_proj.weight", (kv_width, hidden)), + ("self_attn.gate_proj.weight", (q_width, hidden)), + ("self_attn.o_proj.weight", (hidden, q_width)), + ("self_attn.q_norm.weight", (config.head_dim,)), + ("self_attn.k_norm.weight", (config.head_dim,)), + ("pre_feedforward_layernorm.weight", (hidden,)), + ("post_attention_layernorm.weight", (hidden,)), + ("post_feedforward_layernorm.weight", (hidden,)), + ] + ) + + +def _attribute(name: str) -> str: + return name.removesuffix(".weight").replace(".", "_") + + +def check_config(config, layer_idx: int) -> None: + """Raise ValueError naming the first field outside what this component computes.""" + if type(layer_idx) is not int or layer_idx not in SUPPORTED_LAYER_INDICES: + raise ValueError( + f"only sparse blocks {SUPPORTED_LAYER_INDICES} are supported, got layer_idx={layer_idx!r}" + ) + for field, expected in SUPPORTED_CONFIG.items(): + actual = getattr(config, field, None) + if actual != expected or type(actual) is not type(expected): + raise ValueError(f"unsupported {field}={actual!r}; this block computes {expected!r}") + attention_type = ATTENTION_TYPES[layer_idx] + layer_types = getattr(config, "layer_types", None) or [] + if len(layer_types) <= layer_idx or layer_types[layer_idx] != attention_type: + raise ValueError( + f"block {layer_idx} must be a {attention_type} layer: " + f"{layer_types[layer_idx] if len(layer_types) > layer_idx else None!r}" + ) + rope_parameters = getattr(config, "rope_parameters", None) or {} + rope = rope_parameters.get(attention_type) + expected_rope = SUPPORTED_ROPE if attention_type == "sliding_attention" else None + if rope != expected_rope or attention_type not in rope_parameters: + raise ValueError( + f"unsupported {attention_type} rope parameters {rope!r}; need {expected_rope!r}" + ) + check_moe_config(config, layer_idx) + + +class BeamSparseBlockBf16Sm103Tp1(nn.Module): + """Beam sparse decoder block 3, 4 or 5: BF16 weights and activations, an FP32 router, one GPU, + one request step. + + State: the twelve BF16 attention and norm tensors of ``weight_shapes`` and ``mlp``, the + MoE component with its router and 256 experts; ``load_weights`` loads all 782. + + ``forward(hidden_states [1, T, 5120] bf16, position_ids [1, T], attn_metadata)`` + returns ``[1, T, 5120]`` bf16 for 1 <= T <= ``MAX_TOKENS``, a prepared + ``TrtllmAttentionMetadata`` holding one step of one request on a BF16 + ``KVCacheManagerV2`` (8 kv heads of 128) that holds the block's logical layer -- a + context step of T tokens after ``past`` cached ones (paged-context FMHA when + ``past > 0``) or a one-token generation step -- and positions ``arange(past, past + T)``. + Inputs outside that are rejected before any kernel runs; batches of several requests + and other blocks are not computed here. Inputs and parameters are never modified. + """ + + def __init__( + self, config, layer_idx: int = DEFAULT_LAYER_IDX, device: str | torch.device = "cuda" + ): + super().__init__() + check_config(config, layer_idx) + device = torch.device(device) + if device.type != "cuda" or torch.cuda.get_device_capability(device) != _SM: + raise ValueError( + f"this block is built for sm_{_SM[0]}{_SM[1]} CUDA devices, got {device}" + ) + self.layer_idx = layer_idx + self.eps = config.rms_norm_eps + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.is_sliding = ATTENTION_TYPES[layer_idx] == "sliding_attention" + self.window = config.sliding_window if self.is_sliding else None + self.rope_theta = SUPPORTED_ROPE["rope_theta"] if self.is_sliding else None + self.attention_scales = depth_scales(layer_idx, is_attention=True) + self.feed_forward_scales = depth_scales(layer_idx, is_attention=False) + self._shapes = weight_shapes(config) + for name, shape in self._shapes.items(): + tensor = torch.empty(shape, dtype=torch.bfloat16, device=device) + self.register_parameter(_attribute(name), nn.Parameter(tensor, requires_grad=False)) + self.mlp = BeamMoeBf16Sm103Tp1(config, layer_idx, device=device) + self._loaded = False + + @property + def hidden_size(self) -> int: + return self._shapes["input_layernorm.weight"][0] + + def checkpoint_tensors(self) -> "OrderedDict[str, torch.Tensor]": + """The block's 782 tensors by checkpoint name: its own twelve, then the MoE's.""" + tensors = OrderedDict((name, getattr(self, _attribute(name))) for name in self._shapes) + for name, tensor in self.mlp.checkpoint_tensors().items(): + tensors[MLP_PREFIX + name] = tensor + return tensors + + @torch.no_grad() + def load_weights(self, weights: Mapping[str, torch.Tensor]) -> None: + """Copy exactly the block's 782 tensors, by checkpoint name, [out, in]. + + The twelve attention and norm tensors are BF16; the ``mlp.``-prefixed MoE tensors + are as ``BeamMoeBf16Sm103Tp1.load_weights`` describes (FP32 router, BF16 experts). + Every tensor is checked before any parameter changes, so a rejected load leaves + the block as it was. The block runs only after a load completed: a failure while + copying leaves it unloaded. + """ + if not isinstance(weights, Mapping): + raise TypeError(f"weights must be a mapping, got {type(weights).__name__}") + own = {n: w for n, w in weights.items() if not str(n).startswith(MLP_PREFIX)} + moe = {n[len(MLP_PREFIX) :]: w for n, w in weights.items() if str(n).startswith(MLP_PREFIX)} + check_weights( + OrderedDict((n, (shape, torch.bfloat16)) for n, shape in self._shapes.items()), + own, + "block", + ) + self.mlp.check_weights(moe) + self._loaded = False + device = self.input_layernorm.device + staged = {name: own[name].to(device=device, copy=True) for name in self._shapes} + for name, tensor in staged.items(): + getattr(self, _attribute(name)).copy_(tensor) + self.mlp.load_weights(moe) + self._loaded = True + + def _check_step( + self, hidden_states: torch.Tensor, position_ids: torch.Tensor, md: TrtllmAttentionMetadata + ) -> None: + """Raise ValueError on anything this block would otherwise compute wrongly.""" + if not self._loaded: + raise ValueError("load_weights must be called before the block runs") + _attention.check_step( + hidden_states, + position_ids, + md, + layer_idx=self.layer_idx, + hidden_size=self.hidden_size, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + device=self.input_layernorm.device, + max_tokens=MAX_TOKENS, + ) + + def step_args(self, md: TrtllmAttentionMetadata) -> dict: + """Project the prepared metadata onto thop_attention's explicit batch state.""" + return _attention.step_args(md, self.layer_idx) + + def attention(self, h: torch.Tensor, positions: torch.Tensor, step: dict) -> torch.Tensor: + """Normalized input [T, hidden] -> attention output [T, hidden], before post_attention_layernorm.""" + return _attention.attention( + h, + positions, + step, + q_proj=self.self_attn_q_proj, + k_proj=self.self_attn_k_proj, + v_proj=self.self_attn_v_proj, + gate_proj=self.self_attn_gate_proj, + o_proj=self.self_attn_o_proj, + q_norm=self.self_attn_q_norm, + k_norm=self.self_attn_k_norm, + gate_bias=self.self_attn_gate_bias, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + eps=self.eps, + rope_theta=self.rope_theta, + window=self.window if self.is_sliding else step["max_seq_len"], + ) + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + attn_metadata: TrtllmAttentionMetadata, + ) -> torch.Tensor: + self._check_step(hidden_states, position_ids, attn_metadata) + step = self.step_args(attn_metadata) + num_tokens = hidden_states.shape[1] + x = hidden_states.view(num_tokens, self.hidden_size) + positions = position_ids.view(num_tokens) + + h = flashinfer_rmsnorm(x, self.input_layernorm, self.eps) + h = self.attention(h, positions, step) + h = flashinfer_rmsnorm(h, self.post_attention_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.attention_scales) + + h = flashinfer_rmsnorm(x, self.pre_feedforward_layernorm, self.eps) + h = self.mlp(h) + h = flashinfer_rmsnorm(h, self.post_feedforward_layernorm, self.eps) + x = beam_depth_residual(h, x, *self.feed_forward_scales) + return x.view(1, num_tokens, self.hidden_size) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/_dispatch_combine_cases.py b/tests/unittest/_torch/modeling_v2/models/beam/_dispatch_combine_cases.py new file mode 100644 index 000000000000..f2328f547130 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/_dispatch_combine_cases.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generated inputs, controlled experts and the plain-torch golden of the dispatch/combine tests. + +Block 3's MoE combination: hidden 5120, 256 experts, eight per token, routed scaling factor +2.5. The experts here are controlled functions rather than matrices: expert ``e`` maps a +BF16 row ``x`` to ``bf16(fp32(x) * A[e] + B[e])`` with fixed FP32 tables ``A`` and ``B`` +``[256, 5120]``, so every output depends on the expert's identity and on the gathered +token's values, and a row sent to the wrong expert or taken from the wrong token changes +it. Every tensor is generated on the host from a seed, so the tests need no files. +""" + +from __future__ import annotations + +import torch + +HIDDEN, EXPERTS, TOP_K, SCALE = 5120, 256, 8, 2.5 +SEEDS = (17, 41) +TOKENS = (1, 17, 129) +BF16, FP32, I64 = torch.bfloat16, torch.float32, torch.int64 +PROVIDER_SEED = 20261019 +PERM_SEED = 20261020 +FIXTURE_SEED = 20261021 + + +class AffineExperts: + """Expert ``e``: ``rows -> bf16(fp32(rows) * A[e] + B[e])``, row-wise; every call logged. + + ``calls`` holds ``(expert_id, rows, output)`` per call, in call order. An empty batch is + refused, so an expert called for no token fails the test that made the call. + """ + + def __init__(self, a: torch.Tensor, b: torch.Tensor): + self.a, self.b = a, b + self.calls: list[tuple[int, torch.Tensor, torch.Tensor]] = [] + + def __call__(self, expert_id: int, rows: torch.Tensor) -> torch.Tensor: + assert type(expert_id) is int, f"expert id must be a Python int, got {type(expert_id)}" + assert rows.shape[0] > 0, f"expert {expert_id} called with an empty batch" + out = (rows.float() * self.a[expert_id] + self.b[expert_id]).to(BF16) + self.calls.append((expert_id, rows, out)) + return out + + def fresh(self) -> "AffineExperts": + """The same experts with an empty call log.""" + return AffineExperts(self.a, self.b) + + +def affine_experts(device: str = "cuda", zero: bool = False) -> AffineExperts: + """Signed coefficients |A| in [0.5, 1.5) and offsets B ~ 0.25 N(0, 1); all zero if asked.""" + if zero: + return AffineExperts(*(torch.zeros(EXPERTS, HIDDEN, dtype=FP32, device=device),) * 2) + g = torch.Generator(device="cpu").manual_seed(PROVIDER_SEED) + magnitude = 0.5 + torch.rand(EXPERTS, HIDDEN, generator=g, dtype=FP32) + negative = torch.rand(EXPERTS, HIDDEN, generator=g, dtype=FP32) < 0.5 + a = torch.where(negative, -magnitude, magnitude) + b = 0.25 * torch.randn(EXPERTS, HIDDEN, generator=g, dtype=FP32) + return AffineExperts(a.to(device), b.to(device)) + + +def matrix_case(seed: int, tokens: int): + """``(hidden, ids, weights)`` on the host: 129 generated tokens, the first ``tokens`` used. + + Hidden states are unit-variance BF16. Each row's ids are eight distinct experts in + random (unsorted) order and its weights are normalized as the router normalizes them. + """ + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randn(max(TOKENS), HIDDEN, generator=g, dtype=FP32).to(BF16) + ids = torch.stack([torch.randperm(EXPERTS, generator=g)[:TOP_K] for _ in range(max(TOKENS))]) + scores = 0.05 + torch.rand(max(TOKENS), TOP_K, generator=g, dtype=FP32) + w = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) + return x[:tokens].clone(), ids[:tokens].to(I64).clone(), w[:tokens].clone() + + +def permutation(seed: int, tokens: int) -> torch.Tensor: + g = torch.Generator(device="cpu").manual_seed(PERM_SEED + 1000 * seed + tokens) + return torch.randperm(tokens, generator=g) + + +SHARED_IDS = [ + [255, 3, 0, 128, 254, 1, 77, 200], + [0, 255, 9, 128, 31, 254, 2, 100], + [7, 0, 255, 1, 64, 129, 250, 33], + [254, 128, 0, 255, 5, 6, 11, 99], + [1, 2, 3, 4, 255, 0, 253, 252], + [128, 129, 130, 0, 131, 132, 255, 133], +] +WITNESS_IDS = [ + [12, 200, 57, 255, 0, 91, 143, 33], + [200, 12, 255, 64, 18, 91, 222, 5], + [0, 57, 143, 129, 200, 77, 255, 111], + [91, 33, 0, 12, 250, 129, 64, 7], +] +WITNESS_BASE_WEIGHTS = [0.34, 0.21, 0.14, 0.11, 0.08, 0.06, 0.04, 0.02] +ZERO_IDS = [ + [0, 17, 255, 42, 128, 3, 200, 91], + [255, 128, 7, 0, 64, 33, 250, 1], + [42, 0, 91, 254, 17, 255, 5, 129], +] +#: Tokens 0, 1 and 2 give experts 255, 33 and 129 a weight of exactly zero. +ZERO_ROUTE_SLOT = [2, 5, 7] + + +def fixtures() -> dict: + """The focused fixtures on the host, by name: ``(hidden, ids, weights, zero_experts)``. + + shared_extremes six distinct tokens all sharing experts 0 and 255 (and 128 / 254 / + 1 several times), unsorted slots, unequal weights that do not sum + to one, so a re-normalization would show + cancellation_witness signed routes of unequal size (weights 0.02 to 0.34, none a BF16 + value), so each rounding and the accumulation order are visible + zero_weight_route one route per token with weight exactly 0 + all_zero_weights every weight 0 + all_zero_outputs ordinary weights; the zero experts return exact zeros + """ + g = torch.Generator(device="cpu").manual_seed(FIXTURE_SEED) + fx = {} + x = torch.randn(6, HIDDEN, generator=g, dtype=FP32).to(BF16) + w = 0.05 + 0.35 * torch.rand(6, TOP_K, generator=g, dtype=FP32) + fx["shared_extremes"] = (x, torch.tensor(SHARED_IDS, dtype=I64), w, False) + x = torch.randn(4, HIDDEN, generator=g, dtype=FP32).to(BF16) + base = torch.tensor(WITNESS_BASE_WEIGHTS, dtype=FP32) + rows = [] + for _ in range(4): + perm = torch.randperm(TOP_K, generator=g) + jitter = 1.0 + 0.01 * torch.randn(TOP_K, generator=g, dtype=FP32) + rows.append(base[perm] * jitter) + raw = torch.stack(rows) + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + fx["cancellation_witness"] = (x, torch.tensor(WITNESS_IDS, dtype=I64), w, False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + raw = 0.05 + torch.rand(3, TOP_K, generator=g, dtype=FP32) + for t, j in enumerate(ZERO_ROUTE_SLOT): + raw[t, j] = 0.0 + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + for t, j in enumerate(ZERO_ROUTE_SLOT): + w[t, j] = 0.0 + fx["zero_weight_route"] = (x, torch.tensor(ZERO_IDS, dtype=I64), w, False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + fx["all_zero_weights"] = (x, torch.tensor(ZERO_IDS, dtype=I64), torch.zeros(3, TOP_K), False) + x = torch.randn(3, HIDDEN, generator=g, dtype=FP32).to(BF16) + raw = 0.05 + torch.rand(3, TOP_K, generator=g, dtype=FP32) + w = raw / (raw.sum(dim=-1, keepdim=True) + 1e-20) + fx["all_zero_outputs"] = (x, torch.tensor(ZERO_IDS, dtype=I64), w, True) + return fx + + +#: Wrong answers the golden can produce on request, one step changed each. +FAULTS = ( + "wrong_scatter_mapping", # contributions added to token (t + 1) mod T + "no_route_cast", # weighted route kept fp32 + "bf16_product", # weight rounded to bf16, product in bf16 + "cast_after_accumulate", # fp32 routes summed, the sum rounded to bf16 before scaling + "bf16_accumulator", # accumulator held in bf16 + "overwrite", # a contribution replaces the accumulator row instead of adding + "weights_slot_rotated", # each route takes the next slot's weight + "renormalized_weights", # weights re-normalized per token + "scale_omitted", + "scale_duplicated", +) + + +def golden(x, ids, w, experts, fault: str | None = None, skip_slots=None) -> dict: + """The source combination in plain torch, every step, on the inputs' device. + + Written independently of the component: every expert id 0..255 is enumerated, each + expert's (token, slot) pairs come from the routing table by ``nonzero``, and each + expert's rounded routes are added to their tokens' rows of an fp32 accumulator + (``acc[t] += route``; a token appears at most once per expert). Returns the stages + ``product`` / ``route`` (assignment rows in call order), ``acc`` (before scaling), + ``scaled`` and ``output``, plus ``calls`` = ``[(expert, tokens)]``. ``skip_slots[t]`` + drops slot ``skip_slots[t]`` of token ``t`` from the table. + """ + tokens = x.shape[0] + wk = w / (w.sum(dim=-1, keepdim=True) + 1e-20) if fault == "renormalized_weights" else w + acc = torch.zeros( + tokens, HIDDEN, dtype=BF16 if fault == "bf16_accumulator" else FP32, device=x.device + ) + stages = {"product": [], "route": [], "calls": []} + for e in range(EXPERTS): + token_slot = (ids == e).nonzero() + if skip_slots is not None: + keep = [slot != skip_slots[token] for token, slot in token_slot.tolist()] + token_slot = token_slot[torch.tensor(keep, dtype=torch.bool, device=ids.device)] + if token_slot.shape[0] == 0: + continue + t, j = token_slot[:, 0], token_slot[:, 1] + out = experts(e, x[t]) + jw = (j + 1) % TOP_K if fault == "weights_slot_rotated" else j + weight = wk[t, jw] + if fault == "bf16_product": + product = out * weight.to(BF16).unsqueeze(1) + route = product + else: + product = out.float() * weight.unsqueeze(1) + route = ( + product if fault in ("no_route_cast", "cast_after_accumulate") else product.to(BF16) + ) + dst = (t + 1) % tokens if fault == "wrong_scatter_mapping" else t + if fault == "overwrite": + acc[dst] = route.float() + elif fault == "bf16_accumulator": + acc[dst] = (acc[dst].float() + route.float()).to(BF16) + else: + acc[dst] += route.float() + stages["product"].append(product) + stages["route"].append(route) + stages["calls"].append((e, t)) + acc = acc.float() + if fault == "cast_after_accumulate": + acc = acc.to(BF16).float() + scaled = {"scale_omitted": acc, "scale_duplicated": acc * SCALE * SCALE}.get(fault, acc * SCALE) + stages.update(acc=acc, scaled=scaled, output=scaled.to(BF16)) + stages["product"] = torch.cat(stages["product"]) + stages["route"] = torch.cat(stages["route"]) + return stages + + +def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + """Same dtype, shape and bits (so -0.0 differs from +0.0).""" + if a.dtype != b.dtype or a.shape != b.shape: + return False + view = {BF16: torch.int16, FP32: torch.int32}.get(a.dtype) + if view is None: + return bool(torch.equal(a, b.to(a.device))) + return bool(torch.equal(a.contiguous().view(view), b.to(a.device).contiguous().view(view))) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_flashinfer_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_flashinfer_silu_and_mul.py new file mode 100644 index 000000000000..a3fdb9feef35 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_flashinfer_silu_and_mul.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for flashinfer_silu_and_mul, using the shared operator comparison.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.activation.flashinfer_silu_and_mul import ( + flashinfer_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../../activation", "../../../.."] +import test_modeling_v2_flashinfer_silu_and_mul as shared # noqa: E402 + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/dense_block": tuple( + ( + Cell( + why=f"beam dense block, {tokens} tokens: 2 x 20480", + spec=dict(shape=(tokens, 40960), dtype=torch.bfloat16), + ) + for tokens in (1, 17, 129) + ) + ), +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell) -> None: + assert shared.op is op + shared.test_certified_cells(cell) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_logit_softcap.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_logit_softcap.py new file mode 100644 index 000000000000..1f1cbfe5a6fc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_logit_softcap.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the six-block model's beam_logit_softcap entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.beam_logit_softcap import ( # noqa: E501 + beam_logit_softcap as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_logit_softcap requires a CUDA device" + +#: Signed pre-softcap values from zero to far beyond saturation, -0 included. +SATURATION = (0.0, -0.0, 2.0**-10, -(2.0**-10), 0.5, -0.5, 1.0, -1.0, 5.0, -5.0, 10.0, -10.0) +SATURATION += (29.0, -29.0, 30.0, -30.0, 60.0, -60.0, 150.0, -150.0, 300.0, -300.0, 1024.0, -1024.0) +SATURATION += (32768.0, -32768.0, 3.0e38, -3.0e38) + + +def _build(spec, seed): + if spec.get("saturating"): + return torch.tensor(SATURATION, dtype=torch.float32, device="cuda").to(spec["dtype"])[None] + g = torch.Generator(device="cuda").manual_seed(seed) + logits = 2.0 * torch.randn(spec["tokens"], spec["vocab"], generator=g, device="cuda") + return logits.to(spec["dtype"]) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the same three fp32 steps; finite, within the cap, fp32.""" + certifying(op) + logits = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(logits, cell.spec["cap"]) + assert out.shape == logits.shape and out.dtype == torch.float32 + torch.testing.assert_close(out, op.reference(logits, cell.spec["cap"]), rtol=0.0, atol=0.0) + assert bool(torch.isfinite(out).all()) and float(out.abs().max()) <= cell.spec["cap"] + + +def test_saturation_and_signs() -> None: + """Large magnitudes map to exactly +-30, signs are kept (-0 stays -0), and the result + tracks the float64 expression within a few fp32 units of the cap.""" + logits = _build(dict(saturating=True, dtype=torch.bfloat16), 0) + out = op(logits, 30.0)[0] + ref = torch.tanh(logits[0].double() / 30.0) * 30.0 + assert float((out.double() - ref).abs().max()) <= 8 * 30.0 * 2.0**-24 + big = logits[0].float().abs() >= 300.0 + assert torch.equal(out[big].abs(), torch.full_like(out[big], 30.0)) + assert torch.equal(torch.signbit(out), torch.signbit(logits[0].float())) + + +def test_once_is_not_twice_or_none() -> None: + """Applying the cap twice, or not at all, changes ordinary unsaturated logits.""" + logits = _build(dict(tokens=17, vocab=151936, dtype=torch.bfloat16), 5) + once = op(logits, 30.0) + assert not torch.equal(once, logits.float()) + assert not torch.equal(once, op(once.to(torch.bfloat16), 30.0)) + + +@pytest.mark.parametrize( + "case,match", [("int_cap", "positive float"), ("fp32_logits", "bf16 head output")] +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + logits = _build(dict(tokens=1, vocab=151936, dtype=torch.bfloat16), 9) + cap = 30.0 + if case == "int_cap": + cap = 30 + else: + logits = logits.float() + with pytest.raises(AssertionError, match=match): + with validating(op): + op(logits, cap) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_swiglu_fp32.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_swiglu_fp32.py new file mode 100644 index 000000000000..7e64b9da7961 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/activation/test_modeling_v2_beam_swiglu_fp32.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the swiglu_fp32 entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.activation.swiglu_fp32 import ( # noqa: E501 + swiglu_fp32 as op, +) + +__extra_import_path__ = ["../../_components/moe_expert", "../../../.."] +import _expert_cases as cases # noqa: E402 +from _cells import certified_cells # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "swiglu_fp32 requires a CUDA device" + + +TARGET_CELLS = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 SwiGLU {width} wide, {tokens} tokens", + spec=dict(tokens=tokens, intermediate=width, dtype=torch.bfloat16), + ) + for width in (640, 5120) + for tokens in (1, 17, 129) + ), +} +CELLS = certified_cells(op, TARGET_CELLS) + + +def _build(spec, seed): + """Gate and up as the expert's projections produce them: unit-variance bf16.""" + g = torch.Generator(device="cpu").manual_seed(seed) + shape = (spec["tokens"], spec["intermediate"]) + gate = torch.randn(shape, generator=g).to(spec["dtype"]).cuda() + up = torch.randn(shape, generator=g).to(spec["dtype"]).cuda() + return gate, up + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:26] for c in CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """SiLU and product within 2^-16 relative of the fp64 steps (fp32 formula spread on the + same inputs), the activation within one bf16 unit (a product that close to a rounding + boundary may round either way).""" + certifying(op) + gate, up = _build(cell.spec, seed) + out = op(gate, up) + ref = op.reference(gate, up) + assert {k: v.dtype for k, v in out.items()} == { + "silu": torch.float32, + "product": torch.float32, + "act": torch.bfloat16, + } + for name, rtol in (("silu", cases.FP32_FORMULA), ("product", cases.FP32_FORMULA)): + torch.testing.assert_close(out[name], ref[name], rtol=rtol, atol=0.0) + torch.testing.assert_close(out["act"], ref["act"], rtol=cases.BF16_UNIT, atol=0.0) + + +def test_bf16_steps_are_a_different_answer() -> None: + """On the precision fixture's pairs every activation is exact here, and every one + differs from the activation of a bf16 SiLU times up computed in bf16.""" + _, x, expected, premature_act = cases.precision_fixture() + gate, up = x[:, : cases.INTER].cuda(), x[:, cases.INTER :].cuda() + out = op(gate, up) + assert torch.equal(out["act"].cpu(), expected["act"]) + bf16_steps = torch.nn.functional.silu(gate) * up + assert torch.equal(bf16_steps.cpu(), premature_act) + assert bool((out["act"] != bf16_steps).all()) + + +@pytest.mark.parametrize( + "case,match", + [ + ("broadcast_up", "must have one shape"), + ("fp32_gate", "must be torch.bfloat16"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + gate, up = _build(op.CELLS[1].spec, seed=17) + if case == "broadcast_up": + up = up[:1] + else: + gate = gate.float() + with pytest.raises(AssertionError, match=match): + with validating(op): + op(gate, up) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_neox_rope.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_neox_rope.py new file mode 100644 index 000000000000..6c2bdfea6ade --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_neox_rope.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_neox_rope entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_neox_rope import ( + beam_neox_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _cells import certified_cells # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_neox_rope requires a CUDA device" + + +TARGET_CELLS = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell(why=f"beam TP4: {cell.why}", spec=dict(cell.spec, num_heads=12, num_kv_heads=2)) + for cell in op.CELLS + ), +} +CELLS = certified_cells(op, TARGET_CELLS) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, d, dtype = spec["tokens"], spec["head_dim"], spec["dtype"] + q = torch.randn(t, spec["num_heads"], d, generator=g, device="cuda").to(dtype) + k = torch.randn(t, spec["num_kv_heads"], d, generator=g, device="cuda").to(dtype) + return q, k, torch.arange(t, device="cuda") + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the op and its reference make the same fp32 operations in the same order + and round once, so any difference is a changed formula, not accumulation noise.""" + certifying(op) + spec = cell.spec + q, k, positions = _build(spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + q_out, k_out = op(q, k, positions, spec["theta"], spec["factor"]) + q_ref, k_ref = op.reference(q, k, positions, spec["theta"], spec["factor"]) + for out, ref, src in ((q_out, q_ref, q), (k_out, k_ref, k)): + assert out.shape == src.shape and out.dtype == src.dtype + assert out.data_ptr() != src.data_ptr(), "the op must return a new tensor" + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + if spec["tokens"] == 1: + assert torch.equal(q_out, q) and torch.equal(k_out, k), "position 0 must rotate by 0" + else: + assert not torch.equal(q_out[1:], q[1:]), "positions past 0 must rotate" + + +def test_halves_rotate_together_not_interleaved_pairs() -> None: + """A unit vector in dimension 0 at position p goes to (cos p, ..., sin p at D/2, ...): + NeoX pairs dimension i with i + D/2, not with i + 1.""" + d = 128 + q = torch.zeros(2, 1, d, dtype=torch.float32, device="cuda") + q[:, 0, 0] = 1.0 + k = q.clone() + positions = torch.arange(2, device="cuda") + q_out, _ = op(q, k, positions, 10000.0, 1.0) + assert q_out[1, 0, 0].item() == pytest.approx(torch.cos(torch.tensor(1.0)).item(), abs=1e-6) + assert q_out[1, 0, d // 2].item() == pytest.approx( + torch.sin(torch.tensor(1.0)).item(), abs=1e-6 + ) + assert q_out[1, 0, 1].item() == 0.0 + + +@pytest.mark.parametrize( + "case,match", + [ + ("positions_broadcast", "positions must be"), + ("float_positions", "integer token positions"), + ("odd_head_dim", "even head_dim"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + q = torch.randn(17, 48, 128, dtype=torch.bfloat16, device="cuda") + k = torch.randn(17, 8, 128, dtype=torch.bfloat16, device="cuda") + positions = torch.arange(17, device="cuda") + if case == "positions_broadcast": + positions = positions[:1] + elif case == "float_positions": + positions = positions.float() + else: + q, k = q[..., :127], k[..., :127] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(q, k, positions, 10000.0, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_nope_query_scale.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_nope_query_scale.py new file mode 100644 index 000000000000..7f0b529cb771 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_nope_query_scale.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_nope_query_scale entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_nope_query_scale import ( + beam_nope_query_scale as op, +) + +__extra_import_path__ = ["../../../.."] +from _cells import certified_cells # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_nope_query_scale requires a CUDA device" + + +TARGET_CELLS = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell(why=f"beam TP4: {cell.why}", spec=dict(cell.spec, num_heads=12)) for cell in op.CELLS + ), +} +CELLS = certified_cells(op, TARGET_CELLS) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, h, d, dtype = spec["tokens"], spec["num_heads"], spec["head_dim"], spec["dtype"] + q = torch.randn(t, h, d, generator=g, device="cuda").to(dtype) + positions = torch.arange(t, device="cuda") + return q, positions + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: the fp32 scale and the single rounding of the product are what the + reference spells out step by step.""" + certifying(op) + q, positions = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + out = op(q, positions) + assert out.shape == q.shape and out.dtype == q.dtype + torch.testing.assert_close(out, op.reference(q, positions), rtol=0.0, atol=0.0) + + +def test_scale_against_float64() -> None: + """The fp32 scale is within four fp32 units of the float64 expression at every position, so + the bf16 result differs from a float64-scaled one only where the product sits within that + distance of a bf16 rounding boundary.""" + positions = torch.arange(129, device="cuda") + q = torch.ones(129, 1, 1, dtype=torch.float32, device="cuda") + scale32 = op(q, positions).view(-1) + scale64 = torch.exp(0.1 * torch.log1p((positions.double() + 1) / 2048)) + assert float(((scale32.double() - scale64).abs() / scale64).max()) <= 4 * 2.0**-23 + + +def test_scale_is_visible_at_nonzero_positions() -> None: + """Values one bf16 unit below a power of two round up under the scale once it exceeds half + a unit -- at the later positions of a 129-token context -- while position 0's scale + leaves them unchanged: an omitted scale is detectable only at nonzero positions.""" + positions = torch.arange(129, device="cuda") + q = torch.full((129, 1, 1), 1.9921875, dtype=torch.bfloat16, device="cuda") + out = op(q, positions).view(-1) + assert float(out[0]) == 1.9921875 + assert float(out[-1]) == 2.0 + + +def test_omitting_the_scale_or_rotating_differs() -> None: + """A full-attention q that is left unscaled is a different tensor on ordinary inputs.""" + q, positions = _build(dict(tokens=129, num_heads=48, head_dim=128, dtype=torch.bfloat16), 3) + assert not torch.equal(op(q, positions), q) + + +@pytest.mark.parametrize( + "case,match", + [ + ("positions_length_one", "positions must be"), + ("float_positions", "integer token positions"), + ("q_rank", "q must be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + q, positions = _build(dict(tokens=17, num_heads=48, head_dim=128, dtype=torch.bfloat16), 11) + if case == "positions_length_one": + positions = positions[:1] + elif case == "float_positions": + positions = positions.float() + else: + q = q.view(17, 48 * 128) + with pytest.raises(AssertionError, match=match): + with validating(op): + op(q, positions) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_output_gate.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_output_gate.py new file mode 100644 index 000000000000..c3801f0d0313 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_output_gate.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_output_gate entry.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.attention.beam_output_gate import ( + beam_output_gate as op, +) + +__extra_import_path__ = ["../../../.."] +from _cells import certified_cells # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_output_gate requires a CUDA device" + + +TARGET_CELLS = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell(why=f"beam TP4: {cell.why}", spec=dict(cell.spec, num_heads=12)) for cell in op.CELLS + ), +} +CELLS = certified_cells(op, TARGET_CELLS) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + t, h, d, dtype = spec["tokens"], spec["num_heads"], spec["head_dim"], spec["dtype"] + attention_output = torch.randn(t, h, d, generator=g, device="cuda").to(dtype) + gate = (2.0 * torch.randn(t, h, d, generator=g, device="cuda")).to(dtype) + gate_bias = (0.75 * torch.randn(h, generator=g, device="cuda")).to(dtype) + return attention_output, gate, gate_bias + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why[:44] for c in CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: each activation-dtype step of the op is an fp32 computation rounded once, + which is what the reference spells out.""" + certifying(op) + attention_output, gate, gate_bias = _build( + cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + ) + out = op(attention_output, gate, gate_bias) + assert out.shape == attention_output.shape and out.dtype == attention_output.dtype + torch.testing.assert_close( + out, op.reference(attention_output, gate, gate_bias), rtol=0.0, atol=0.0 + ) + + +def test_rounding_between_steps_is_observable() -> None: + """One fp32 chain rounded once is a different answer on these inputs, so the cell + gate above distinguishes the source's per-step rounding from a fused gate.""" + attention_output, gate, gate_bias = _build( + dict(tokens=129, num_heads=48, head_dim=128, dtype=torch.bfloat16), seed=7 + ) + fused = ( + attention_output.float() * torch.sigmoid(gate.float() + gate_bias.float().unsqueeze(-1)) + ).to(torch.bfloat16) + assert not torch.equal(op(attention_output, gate, gate_bias), fused) + + +@pytest.mark.parametrize( + "case,match", + [ + ("bias_length_one", "gate_bias must be"), + ("fp32_bias", "share one dtype"), + ("gate_shape", "must both be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + attention_output, gate, gate_bias = _build( + dict(tokens=17, num_heads=48, head_dim=128, dtype=torch.bfloat16), seed=11 + ) + if case == "bias_length_one": + gate_bias = gate_bias[:1] + elif case == "fp32_bias": + gate_bias = gate_bias.float() + else: + gate = gate[:, :1] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(attention_output, gate, gate_bias) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_thop_attention.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_thop_attention.py new file mode 100644 index 000000000000..3252d1bb6609 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/attention/test_modeling_v2_beam_thop_attention.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for thop_attention, using the shared operator comparison.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import STANDARD +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.thop_attention import ( + thop_attention as op, +) + +__extra_import_path__ = ["../../../../attention", "../../../.."] +import test_modeling_v2_thop_attention as shared # noqa: E402 +from _validating import certifying # noqa: E402,F401 + +_BEAM_DENSE = dict( + input_type=STANDARD, + num_heads=48, + num_kv_heads=8, + head_size=128, + v_head_dim=128, + sinks=False, + window=512, + quant_mode=0, + mask_type=1, + tokens_per_block=32, + num_blocks=16, + max_num_requests=1, + use_paged_context_fmha=False, + num_contexts=1, + predicted_tokens_per_seq=1, +) + +_BEAM_SLIDING = dict( + _BEAM_DENSE, + num_blocks=64, + max_num_requests=3, + max_blocks_per_seq=20, + use_paged_context_fmha=True, +) + +_BEAM_FULL = dict(_BEAM_SLIDING, window=None) + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 {name}, {'full' if window is None else 'sliding'} attention", + spec=dict( + _BEAM_SLIDING, + num_heads=12, + num_kv_heads=2, + window=window, + cached=cached, + new=new, + num_contexts=contexts, + predicted_tokens_per_seq=1, + ), + ) + for window in (512, None) + for name, cached, new, contexts in ( + ("fresh decode-sized prefill", (0,), (1,), 1), + ("fresh prefill", (0,), (129,), 1), + ("prefix reuse", (64,), (17,), 1), + ("prefill beyond window", (0,), (513,), 1), + ("prefix at window", (512,), (1,), 1), + ("decode across window", (511, 512, 513), (1, 1, 1), 0), + ) + ), + "beam/six_layer_model": ( + *( + Cell( + why=f"beam dense block: GQA 48q/8kv d128, fresh {tokens}-row context", + spec=dict(_BEAM_DENSE, cached=(0,), new=(tokens,)), + ) + for tokens in (1, 17, 129) + ), + Cell( + why="beam sliding, paged context: 17 rows over a 64-token reused prefix", + spec=dict( + _BEAM_SLIDING, cached=(64,), new=(17,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding, paged context: a fresh 513-row context, key 0 leaves the " + "window at row 512", + spec=dict( + _BEAM_SLIDING, cached=(0,), new=(513,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding, paged context: one row at position 512 after a 512-token " + "reused prefix", + spec=dict( + _BEAM_SLIDING, cached=(512,), new=(1,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + Cell( + why="beam sliding decode across the window: histories 511, 512 and 513", + spec=dict( + _BEAM_SLIDING, + cached=(511, 512, 513), + new=(1, 1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="beam full layer decode past the sliding window: histories 511, 512 and 513", + spec=dict( + _BEAM_FULL, + cached=(511, 512, 513), + new=(1, 1, 1), + num_contexts=0, + predicted_tokens_per_seq=1, + ), + ), + Cell( + why="beam full layer, paged context: 17 rows over a 64-token reused prefix", + spec=dict( + _BEAM_FULL, cached=(64,), new=(17,), num_contexts=1, predicted_tokens_per_seq=1 + ), + ), + ), +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell, certifying) -> None: # noqa: F811 + assert shared.op is op + shared.test_certified_cells(cell, certifying) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_dispatch_combine.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_dispatch_combine.py new file mode 100644 index 000000000000..c3e72cdd620e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_dispatch_combine.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_dispatch_combine entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.moe.beam_dispatch_combine import ( # noqa: E501 + beam_dispatch_combine as op, +) + +__extra_import_path__ = ["../..", "../../../.."] +import _dispatch_combine_cases as cases # noqa: E402 +from _validating import certifying # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_dispatch_combine requires a CUDA device" + + +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:26] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit for bit: both sides form the same correctly rounded fp32 products, round each + route to bf16 once and add the routes in fp32 in the same expert order, so a single + differing bit is a different computation, not a rounding allowance.""" + certifying(op) + spec = cell.spec + x, ids, w = (t.cuda() for t in cases.matrix_case(seed, spec["tokens"])) + experts = cases.affine_experts() + assert (spec["hidden_size"], spec["num_experts"], spec["top_k"]) == ( + cases.HIDDEN, + cases.EXPERTS, + cases.TOP_K, + ) + called, referenced = experts.fresh(), experts.fresh() + out = op(x, ids, w, called, spec["routed_scaling_factor"]) + ref = op.reference(x, ids, w, referenced, spec["routed_scaling_factor"]) + assert out.dtype == torch.bfloat16 and out.shape == (spec["tokens"], cases.HIDDEN) + assert cases.bitwise_equal(out, ref) + # One call per selected expert, ascending, on both sides. + assert ( + [e for e, _, _ in called.calls] + == [e for e, _, _ in referenced.calls] + == sorted(set(ids.reshape(-1).tolist())) + ) + + +def test_reference_keeps_the_rounded_route() -> None: + """On the cancellation witness the reference differs from the same combination with + fp32 routes: rounding each weighted route to bf16 is the source's semantics.""" + x, ids, w, _ = ( + t.cuda() if isinstance(t, torch.Tensor) else t + for t in cases.fixtures()["cancellation_witness"] + ) + experts = cases.affine_experts() + ref = op.reference(x, ids, w, experts.fresh(), cases.SCALE) + gold = cases.golden(x, ids, w, experts.fresh()) + unrounded = cases.golden(x, ids, w, experts.fresh(), fault="no_route_cast") + assert cases.bitwise_equal(ref, gold["output"]) + assert not torch.equal(ref, unrounded["output"]) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_grouped_topk_router.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_grouped_topk_router.py new file mode 100644 index 000000000000..9ba3eeb986fd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/moe/test_modeling_v2_beam_grouped_topk_router.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_grouped_topk_router entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.moe.beam_grouped_topk_router import ( # noqa: E501 + beam_grouped_topk_router as op, +) + +__extra_import_path__ = ["../../_components/moe_router", "../../../.."] +import _router_cases as cases # noqa: E402 +from _precision import true_fp32_matmul # noqa: E402 +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_grouped_topk_router requires a CUDA device" + + +def _build(spec, seed): + weight, bias, hidden = cases.matrix_case(seed, spec["tokens"]) + kwargs = {k: spec[k] for k in ("n_group", "topk_group", "top_k", "eps")} + return (hidden.cuda(), weight.cuda(), bias.cuda()), kwargs + + +@pytest.mark.parametrize("precision", ["ieee", "tf32"]) +@pytest.mark.parametrize("seed", cases.SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:24] for c in op.CELLS]) +def test_certified_cells(cell, seed, precision, monkeypatch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Every step within the rules in `_router_cases.RULES`; groups and ids exactly, in order.""" + certifying(op) + args, kwargs = _build(cell.spec, seed) + with true_fp32_matmul(): + ref = op.reference(*args, **kwargs) + matmul = torch.backends.cuda.matmul + previous = matmul.fp32_precision + try: + matmul.fp32_precision = precision + + # The production call must not even read either global precision API. + class NoGlobalPrecisionAccess: + def __getattr__(self, name): + raise AssertionError(f"router read global precision: {name}") + + def __setattr__(self, name, value): + raise AssertionError(f"router changed global precision: {name}") + + with monkeypatch.context() as patch: + patch.setattr(torch.backends.cuda, "matmul", NoGlobalPrecisionAccess()) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = op(*args, **kwargs) + assert matmul.fp32_precision == precision + finally: + matmul.fp32_precision = previous + assert set(out) == set(ref) + assert out["ids"].shape == (cell.spec["tokens"], cell.spec["top_k"]) + cases.assert_all_pass(cases.compare_routes(out, ref, "reference")) + + +@pytest.mark.parametrize("fixture", cases.fixtures().values(), ids=cases.fixtures().keys()) +def test_controlled_routing_boundaries(fixture, certifying) -> None: # noqa: F811 + """Projection changes must preserve bias, exclusion, tie and epsilon behavior.""" + certifying(op) + weight, bias, hidden = cases.fixture_tensors(fixture) + args = (hidden.cuda(), weight.cuda(), bias.cuda()) + kwargs = {k: op.CELLS[0].spec[k] for k in ("n_group", "topk_group", "top_k", "eps")} + out = op(*args, **kwargs) + with true_fp32_matmul(): + ref = op.reference(*args, **kwargs) + cases.assert_all_pass(cases.compare_routes(out, ref, "reference")) + + +@pytest.mark.parametrize( + "case,match", + [ + ("batched_3d", "hidden_states must be"), + ("bias_length_one", "correction_bias must be"), + ("fp64_bias", "correction_bias must be torch.float32"), + ("top_k_beyond_kept_groups", "cannot be chosen"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + (hidden, weight, bias), kwargs = _build(op.CELLS[0].spec, seed=17) + if case == "batched_3d": + hidden = hidden.unsqueeze(0) + elif case == "bias_length_one": + bias = bias[:1] + elif case == "fp64_bias": + bias = bias.double() + else: + kwargs["top_k"] = kwargs["topk_group"] * 32 + 1 + with pytest.raises(AssertionError, match=match): + with validating(op): + op(hidden, weight, bias, **kwargs) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_flashinfer_rmsnorm.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_flashinfer_rmsnorm.py new file mode 100644 index 000000000000..71b64baa0385 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_flashinfer_rmsnorm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for flashinfer_rmsnorm, using the shared operator comparison.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm as op, +) + +__extra_import_path__ = ["../../../../norm", "../../../.."] +import test_modeling_v2_flashinfer_rmsnorm as shared # noqa: E402 +from _validating import certifying # noqa: E402,F401 + +_BEAM_DENSE_BLOCK = tuple( + Cell( + why=f"beam dense {name}, {tokens} tokens: {rows} rows of {width}", + spec=dict(shape=(rows, width), dtype=torch.bfloat16, eps=1e-6), + ) + for tokens in (1, 17, 129) + for name, rows, width in ( + ("sandwich norm", tokens, 5120), + ("q_norm", tokens * 48, 128), + ("k_norm", tokens * 8, 128), + ) +) + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 {name}, {tokens} tokens: {rows} rows of {width}", + spec=dict(shape=(rows, width), dtype=torch.bfloat16, eps=1e-6), + ) + for tokens in (1, 17, 129) + for name, rows, width in ( + ("sandwich/final norm", tokens, 5120), + ("q_norm", tokens * 12, 128), + ("k_norm", tokens * 2, 128), + ) + ), + "beam/dense_block": _BEAM_DENSE_BLOCK, +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell, certifying) -> None: # noqa: F811 + assert shared.op is op + shared.test_certified_cells(cell, certifying) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_rmsnorm_fp32.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_rmsnorm_fp32.py new file mode 100644 index 000000000000..4e6e8563e23a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/norm/test_modeling_v2_beam_rmsnorm_fp32.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_rmsnorm_fp32 entry.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.norm.beam_rmsnorm_fp32 import ( # noqa: E501 + beam_rmsnorm_fp32 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_rmsnorm_fp32 requires a CUDA device" + +SEEDS = (17, 41) + + +def _build(spec, seed): + """Hidden rows with a residual stream's spread (a few large channels), a weight near 1.""" + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randn(spec["tokens"], spec["width"], generator=g) * 4.0 + x[:, :: spec["width"] // 8] *= 64.0 + w = 1.0 + 0.25 * torch.randn(spec["width"], generator=g) + return x.to(spec["dtype"]).cuda(), w.to(spec["dtype"]).cuda() + + +def _bf16_units_apart(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Distance between same-signed bf16 values in units of the last place.""" + assert a.dtype == b.dtype == torch.bfloat16 + assert bool(((a.float() >= 0) == (b.float() >= 0)).all()), "signs differ" + return (a.view(torch.int16).int() - b.view(torch.int16).int()).abs() + + +@pytest.mark.parametrize("seed", SEEDS) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:26] for c in op.CELLS]) +def test_certified_cells(cell, seed, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 unit of the fp64 reference: the fp32 steps carry a few fp32 units of + error, so the single rounding at the end may land on either neighbour of a boundary.""" + certifying(op) + x, w = _build(cell.spec, seed) + before = x.clone() + out = op(x, w, cell.spec["eps"]) + ref = op.reference(x, w, cell.spec["eps"]) + assert out.dtype == torch.bfloat16 and out.shape == x.shape + assert out.data_ptr() != x.data_ptr() and torch.equal(x, before) + assert int(_bf16_units_apart(out, ref).max()) <= 1 + + +def test_rounds_once_at_the_end() -> None: + """Rounding the normalized row to bf16 before the weight is a different answer: it + rounds twice, and on these rows some outputs move by a bf16 unit.""" + spec = op.CELLS[2].spec + x, w = _build(spec, seed=17) + out = op(x, w, spec["eps"]) + x32 = x.float() + rms = x32.pow(2).mean(dim=-1, keepdim=True).add(spec["eps"]).rsqrt() + rounded_twice = (x32 * rms).to(torch.bfloat16) * w + assert torch.equal(out, (x32 * rms * w).to(torch.bfloat16)) + assert bool((out != rounded_twice).any()) + + +@pytest.mark.parametrize( + "case,match", + [ + ("fp32_weight", "must be torch.bfloat16"), + ("narrow_weight", "weight must be"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + x, w = _build(op.CELLS[1].spec, seed=17) + w = w.float() if case == "fp32_weight" else w[:1] + with pytest.raises(AssertionError, match=match): + with validating(op): + op(x, w, 1e-6) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/ops/residual/test_modeling_v2_beam_depth_residual.py b/tests/unittest/_torch/modeling_v2/models/beam/ops/residual/test_modeling_v2_beam_depth_residual.py new file mode 100644 index 000000000000..a7c13c80f5da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/ops/residual/test_modeling_v2_beam_depth_residual.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the beam_depth_residual entry.""" + +import math +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.ops.residual.beam_depth_residual import ( + beam_depth_residual as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402,F401 — path above; fixture by name + +assert torch.cuda.is_available(), "beam_depth_residual requires a CUDA device" + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + shape = (spec["tokens"], spec["hidden"]) + branch = torch.randn(shape, generator=g, device="cuda").to(spec["dtype"]) + residual = (2.0 * torch.randn(shape, generator=g, device="cuda")).to(spec["dtype"]) + return branch, residual + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bitwise: two scalings and an add, each an fp32 operation rounded once.""" + certifying(op) + branch, residual = _build(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + branch_scale, residual_scale = cell.spec["scales"] + out = op(branch, residual, branch_scale, residual_scale) + assert out.shape == branch.shape and out.dtype == branch.dtype + ref = op.reference(branch, residual, branch_scale, residual_scale) + torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0) + + +def test_scales_are_the_source_depth_coefficients() -> None: + """The cells carry blocks 0 and 3's coefficients: the n-th residual add scales by + (1, sqrt(n - 1)) / sqrt(n); block i's attention add is n = 2(i + 1), its feed-forward + add n = 2(i + 1) + 1.""" + scales = {cell.why.split(",")[0]: cell.spec["scales"] for cell in op.CELLS} + assert scales["block 0 attention add"] == (1 / math.sqrt(2), 1 / math.sqrt(2)) + assert scales["block 0 feed-forward add"] == (1 / math.sqrt(3), math.sqrt(2) / math.sqrt(3)) + assert scales["block 3 attention add"] == (1 / math.sqrt(8), math.sqrt(7) / math.sqrt(8)) + assert scales["block 3 feed-forward add"] == (1 / 3, math.sqrt(8) / 3) + for block, (attention, feed_forward) in ((0, (2, 3)), (3, (8, 9))): + for kind, n in (("attention", attention), ("feed-forward", feed_forward)): + branch, residual = scales[f"block {block} {kind} add"] + assert (branch, residual) == (1 / math.sqrt(n), math.sqrt(n - 1) / math.sqrt(n)) + + +def test_fused_fp32_add_is_a_different_answer() -> None: + branch, residual = _build(dict(tokens=129, hidden=5120, dtype=torch.bfloat16), seed=5) + a, b = 1 / math.sqrt(3), math.sqrt(2) / math.sqrt(3) + fused = (branch.float() * a + residual.float() * b).to(torch.bfloat16) + assert not torch.equal(op(branch, residual, a, b), fused) + + +@pytest.mark.parametrize( + "case,match", + [ + ("broadcast_residual", "must match"), + ("fp32_residual", "share one dtype"), + ("tensor_scale", "Python floats"), + ], +) +def test_guard_refuses_inputs_the_op_answers_wrongly(case, match) -> None: + branch, residual = _build(dict(tokens=17, hidden=5120, dtype=torch.bfloat16), seed=3) + branch_scale, residual_scale = 1 / math.sqrt(2), 1 / math.sqrt(2) + if case == "broadcast_residual": + residual = residual[:1] + elif case == "fp32_residual": + residual = residual.float() + else: + branch_scale = torch.tensor(branch_scale, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match=match): + with validating(op): + op(branch, residual, branch_scale, residual_scale) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/test_component_contracts.py b/tests/unittest/_torch/modeling_v2/models/beam/test_component_contracts.py new file mode 100644 index 000000000000..3c8cf45ce12d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/test_component_contracts.py @@ -0,0 +1,537 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU component, cache geometry and checkpoint contracts; no engine registration.""" + +import ast +from importlib.util import resolve_name +from pathlib import Path +from types import SimpleNamespace as NS + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models import beam +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._tp4 import layout_bf16 as layout +from tensorrt_llm._torch._experimental.modeling_v2.models.beam._tp4 import loader_bf16 as loader +from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.model import ( + BeamFullModelNvfp4Sm103Tp4 as Component, +) +from tensorrt_llm.bindings import DataType + +pytestmark = pytest.mark.cpu_only + + +def config(): + return dict( + layout.SUPPORTED_CONFIG, + layer_types=["sliding_attention"] * 51 + ["full_attention"], + mlp_layer_types=["dense"] * 3 + ["sparse"] * 49, + rope_parameters={ + "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0}, + "full_attention": None, + }, + ) + + +@pytest.mark.parametrize("length", [128, 512, 8192]) +@pytest.mark.parametrize("retention", ["sliding", "full", "normalized"]) +def test_cache_allows_estimator_and_short_sequence_geometry(length, retention): + component = NS( + head_dim=128, local_kv_heads=2, window=512, schedule=layout.check_config(config()) + ) + window = {"sliding": 512, "full": length, "normalized": None}[retention] + manager = NS( + layer_offsets=dict(enumerate(range(52))), + dtype=DataType.BF16, + head_dim=128, + num_kv_heads_per_layer=[2] * 52, + max_attention_window_vec=[window] * 51 + [None], + max_seq_len=length, + ) + Component.check_cache_manager(component, manager) + # Cache retention does not change the mathematical attention window. + assert Component.attention_windows(component, length) == [512] * 51 + [length] + + +@pytest.mark.parametrize("layer,window", [(0, 256), (0, 1024), (51, 512)]) +def test_cache_rejects_incompatible_retention(layer, window): + component = NS( + head_dim=128, local_kv_heads=2, window=512, schedule=layout.check_config(config()) + ) + manager = NS( + layer_offsets=dict(enumerate(range(52))), + dtype=DataType.BF16, + head_dim=128, + num_kv_heads_per_layer=[2] * 52, + max_attention_window_vec=[512] * 51 + [None], + max_seq_len=8192, + ) + manager.max_attention_window_vec[layer] = window + with pytest.raises(ValueError, match="cache keeps a window"): + Component.check_cache_manager(component, manager) + + +@pytest.mark.parametrize("error_type", [OSError, RuntimeError, ValueError, TypeError]) +def test_prefetch_failure_is_raised_in_consumer(tmp_path, monkeypatch, error_type): + path = tmp_path / "weights.bin" + path.write_bytes(bytes(range(8))) + entry = loader.Entry("weight", str(path), layout.F32, (2, 1), 0, 8) + reader = loader.TileReader(tile_bytes=4, pin=False) + original = reader._read + failure = error_type("injected background read failure") + + def read(buf, fd, offset, length): + if offset == 4: + raise failure + original(buf, fd, offset, length) + + monkeypatch.setattr(reader, "_read", read) + tiles = reader.tiles(entry, 0, 2) + try: + assert next(tiles)[:2] == (0, 1) + with pytest.raises(error_type) as caught: + next(tiles) + assert caught.value is failure + finally: + tiles.close() + reader.close() + + +@pytest.mark.parametrize( + "field,value", [("partial_rotary_factor", 0.5), ("use_sliding_window", False)] +) +def test_math_changing_config_is_rejected(field, value): + cfg = config() + cfg[field] = value + with pytest.raises(ValueError, match=field): + layout.check_config(cfg) + + +def test_nested_partial_rope_is_rejected(): + cfg = config() + cfg["rope_parameters"]["sliding_attention"]["partial_rotary_factor"] = 0.5 + with pytest.raises(ValueError, match="rope_parameters"): + layout.check_config(cfg) + + +@pytest.mark.parametrize( + "field", ["factor", "attention_factor", "original_max_position_embeddings", "unknown"] +) +def test_unknown_sliding_rope_options_are_rejected(field: str) -> None: + cfg = config() + cfg["rope_parameters"]["sliding_attention"][field] = 1.0 + with pytest.raises(ValueError, match="rope_parameters"): + layout.check_config(cfg) + + +@pytest.mark.parametrize("field", ["factor", "global_attention", "unknown"]) +def test_unknown_rope_groups_are_rejected(field: str) -> None: + cfg = config() + cfg["rope_parameters"][field] = None + with pytest.raises(ValueError, match="rope_parameters"): + layout.check_config(cfg) + + +def test_explicit_full_head_rope_is_supported() -> None: + cfg = config() + cfg["rope_parameters"]["sliding_attention"]["partial_rotary_factor"] = 1.0 + assert layout.check_config(cfg).num_layers == 52 + + +def test_benign_config_metadata_and_supported_optional_values(): + cfg = dict( + config(), + partial_rotary_factor=1.0, + use_sliding_window=True, + transformers_version="test", + _name_or_path="checkpoint", + ) + assert layout.check_config(cfg).num_layers == 52 + + +@pytest.fixture +def recipe_specs(): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + expert = "model.layers.3.mlp.experts.0.gate_proj" + return { + expert: nvfp4_layout.QuantSpec(expert, (2560, 5120), 3, 0, "gate_proj", 0), + "model.layers.3.self_attn.q_proj.weight": NS(shape=(6144, 5120)), + "lm_head.weight": NS(shape=(151936, 5120)), + } + + +@pytest.mark.parametrize("flat", [False, True]) +@pytest.mark.parametrize( + "attention_pattern", + ["model.layers.*.self_attn.*", "model.layers.3.self_attn", "re:.*self_attn"], +) +def test_shared_quantization_shapes_and_exclusions(recipe_specs, flat, attention_pattern): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + values = dict(quant_algo="NVFP4", group_size=16, exclude_modules=[attention_pattern, "lm_head"]) + recipe = {"quantization": values} + if flat: + recipe = dict(values, quant_method="modelopt", ignore=values["exclude_modules"]) + del recipe["exclude_modules"] + summary = nvfp4_layout.check_recipe(recipe, recipe_specs) + assert summary["quantized_modules"] == 1 + assert summary["excluded_linear_modules"] == 2 + + +def test_parent_exclusion_cannot_hide_quantized_experts(recipe_specs): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + recipe = { + "quantization": dict( + quant_algo="NVFP4", group_size=16, exclude_modules=["model.layers.3", "lm_head"] + ) + } + with pytest.raises(ValueError, match="excludes quantized experts"): + nvfp4_layout.check_recipe(recipe, recipe_specs) + + +def test_quantization_file_precedes_inline_recipe(tmp_path): + import json + + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + inline = {"quant_method": "modelopt", "quant_algo": "NVFP4"} + config = {"quantization_config": inline} + assert nvfp4_layout.read_quant_config(tmp_path, config) == inline + file_recipe = {"quantization": {"quant_algo": "FP8"}} + (tmp_path / "hf_quant_config.json").write_text(json.dumps(file_recipe)) + assert nvfp4_layout.read_quant_config(tmp_path, config) == file_recipe + + +@pytest.mark.parametrize("scheme", ["FP8", "NVFP4", "unknown", {"type": "float", "num_bits": 8}]) +def test_flat_recipe_never_silently_disables_requested_kv_quantization(recipe_specs, scheme): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + with pytest.raises(ValueError, match="BF16 KV cache"): + nvfp4_layout.check_recipe( + dict( + quant_method="modelopt", quant_algo="NVFP4", group_size=16, kv_cache_scheme=scheme + ), + recipe_specs, + ) + + +@pytest.mark.parametrize("in_file", [False, True]) +def test_checkpoint_validation_uses_shared_recipe_reader( + tmp_path, monkeypatch, recipe_specs, in_file +): + import json + + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + layout as nvfp4_layout, + ) + + recipe = dict( + quant_method="modelopt", + quant_algo="NVFP4", + group_size=16, + ignore=["model.layers.3.self_attn", "lm_head"], + ) + cfg = config() + if in_file: + (tmp_path / "hf_quant_config.json").write_text(json.dumps(recipe)) + else: + cfg["quantization_config"] = recipe + (tmp_path / "config.json").write_text(json.dumps(cfg)) + monkeypatch.setattr(nvfp4_layout, "tensor_specs", lambda config: recipe_specs) + assert nvfp4_layout.check_checkpoint(tmp_path) == cfg + + +def test_components_do_not_import_full_model_internals(): + family = Path(next(iter(beam.__path__))) + target = family / "full_model_nvfp4__sm_103__tp4" + target_module = f"{beam.__name__}.{target.name}" + for source in family.rglob("*.py"): + if source.is_relative_to(target): + continue + package = ".".join((beam.__name__, *source.parent.relative_to(family).parts)) + for node in ast.walk(ast.parse(source.read_text())): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + module = resolve_name("." * node.level + (node.module or ""), package) + modules = [f"{module}.{alias.name}" for alias in node.names] + else: + continue + for module in modules: + assert module != target_module and not module.startswith(target_module + "."), ( + f"{source.relative_to(family)} imports target-private code: {module}" + ) + + +@pytest.mark.parametrize("local_rank", range(4)) +def test_load_failures_identify_remote_rank_on_every_rank(monkeypatch, local_rank): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam._load import load_failure_details + from tensorrt_llm._torch.distributed.communicator import Distributed + + outcomes = [ + (True, ""), + (False, "ValueError('bad block scale')"), + (True, ""), + (False, "short read"), + ] + monkeypatch.setattr(Distributed, "get", lambda mapping: NS(allgather=lambda outcome: outcomes)) + error = ValueError("bad block scale") if local_rank == 1 else None + detail = load_failure_details( + NS(world_size=4, rank=local_rank), local_rank not in (1, 3), error + ) + assert detail == "rank 1: ValueError('bad block scale'); rank 3: short read" + + +def test_local_load_success_and_incomplete_coverage(): + from tensorrt_llm._torch._experimental.modeling_v2.models.beam._load import load_failure_details + + mapping = NS(world_size=1) + assert load_failure_details(mapping, True, None) is None + assert load_failure_details(mapping, False, None) == "rank 0: incomplete tensor coverage" + + +@pytest.mark.parametrize("field", ["layer_types", "mlp_layer_types"]) +@pytest.mark.parametrize("value", [1, "dense", None, [[]]]) +def test_malformed_layer_schedule_is_value_error(field, value): + cfg = config() + cfg[field] = value + with pytest.raises(ValueError, match=field): + layout.check_config(cfg) + + +def test_direct_call_device_and_metadata_admission(monkeypatch): + import torch + + from tensorrt_llm._torch._experimental.modeling_v2.models.beam._validation import ( + check_positions, + explicit_device, + metadata_values, + ) + + monkeypatch.setattr(torch.cuda, "current_device", lambda: 2) + assert explicit_device("cuda") == torch.device("cuda:2") + assert explicit_device(None) == torch.device("cuda:2") + assert explicit_device("cuda:1") == torch.device("cuda:1") + for dtype in (torch.int32, torch.int64): + check_positions(torch.arange(3, dtype=dtype), (3,), torch.device("cpu")) + for dtype in (torch.uint8, torch.int16, torch.bool, torch.float32): + with pytest.raises(ValueError, match="int32/int64"): + check_positions(torch.ones(3, dtype=dtype), (3,), torch.device("cpu")) + with pytest.raises(ValueError, match="shape"): + check_positions(torch.ones(1, dtype=torch.int64), (3,), torch.device("cpu")) + for values in ([2, 3], torch.tensor([2, 3])): + assert metadata_values(values, 2, "seq_lens") == [2, 3] + with pytest.raises(ValueError, match="seq_lens needs 3"): + metadata_values(values, 3, "seq_lens") + + +@pytest.mark.parametrize("offset", [0.0, False, "0"]) +def test_safetensors_offsets_must_be_integers(tmp_path, monkeypatch, offset): + import json + + (tmp_path / loader.INDEX_FILE).write_text(json.dumps({"weight_map": {"w": "shard"}})) + (tmp_path / "shard").write_bytes(bytes(8)) + monkeypatch.setattr( + loader, + "_read_header", + lambda path: ( + {"w": {"dtype": "F32", "shape": [1], "data_offsets": [offset, 4]}}, + 0, + ), + ) + with pytest.raises(loader.CheckpointFormatError, match="offsets must be integers"): + loader.SafetensorsCatalog(str(tmp_path)) + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("failure", [None, "metadata", "copy", "close", "remote"]) +def test_checkpoint_load_transaction( + tmp_path, monkeypatch, packed: bool, failure: str | None +) -> None: + """Real CPU tensor reads preserve load atomicity and propagate every failure phase.""" + import json + + import torch + from safetensors.torch import save_file + + from tensorrt_llm._torch._experimental.modeling_v2.models.beam._tp4 import ( + component_bf16 as bf16, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + component as nvfp4, + ) + + cls = ( + nvfp4.BeamCheckpointLayoutNvfp4Sm103Tp4 if packed else bf16.BeamCheckpointLayoutBf16Sm103Tp4 + ) + obj = cls.__new__(cls) + torch.nn.Module.__init__(obj) + # CPU copies are synchronous; omit only the CUDA stream synchronization. + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: NS(synchronize=lambda: None)) + source = torch.arange(8, dtype=torch.float32).reshape(4, 2) + spec = layout.TensorSpec("weight", (4, 2), layout.F32, "dense", None, None, None) + obj.config = {"test": True} + obj.specs = {spec.name: spec} + obj.mapping = NS(world_size=1, rank=0) + obj.rank, obj.layers, obj.device = 0, (), torch.device("cpu") + obj.schedule = NS(is_sparse=lambda layer: False) + obj.expert_names, obj._shared_fc1 = [], {3: [True]} + obj.usable, obj.load_record = True, {"previous": True} + obj.register_parameter( + bf16._attr(spec.name), torch.nn.Parameter(torch.zeros_like(source), requires_grad=False) + ) + obj._check_diagnostics_agree = lambda: None + monkeypatch.setattr(cls, "_tensor_specs", staticmethod(lambda *args: obj.specs)) + if packed: + monkeypatch.setattr(nvfp4, "tensor_specs", lambda *args: obj.specs) + monkeypatch.setattr(nvfp4, "read_quant_config", lambda *args: {}) + monkeypatch.setattr(nvfp4, "check_recipe", lambda *args: {"verified": True}) + (tmp_path / "config.json").write_text(json.dumps(obj.config)) + save_file({"weight": source}, str(tmp_path / "weights.safetensors")) + (tmp_path / loader.INDEX_FILE).write_text( + json.dumps({"weight_map": {"weight": "weights.safetensors"}}) + ) + if failure == "metadata": + (tmp_path / "config.json").write_text("{}") + reader = loader.TileReader(tile_bytes=16, pin=False) + monkeypatch.setattr(bf16, "TileReader", lambda: reader) + closed = [] + real_close = reader.close + + def close() -> None: + real_close() + closed.append(True) + if failure == "close": + raise OSError("close failed") + + monkeypatch.setattr(reader, "close", close) + if failure == "copy": + + def copy_failed(*args, **kwargs) -> None: + raise OSError("copy failed") + + monkeypatch.setattr(bf16, "read_slice_into", copy_failed) + phases = [] + + def agree(mapping, ok, error): + phases.append(ok) + if failure == "remote" and len(phases) == 2: + return "rank 2: remote copy failed" + return None if ok else str(error) + + monkeypatch.setattr(bf16, "load_failure_details", agree) + try: + if failure is not None: + with pytest.raises( + bf16.LoadError, + match={ + "metadata": "checkpoint rejected before loading", + "copy": "copy failed", + "close": "close failed", + "remote": "rank 2", + }[failure], + ): + obj.load(str(tmp_path), record_digests=False) + if failure == "metadata": + assert obj.usable and obj.load_record == {"previous": True} + assert torch.count_nonzero(obj.tensor("weight")) == 0 + assert phases == [False] and not closed + else: + assert not obj.usable and obj.load_record is None + assert len(phases) == 2 and closed == [True] + if packed: + assert obj._shared_fc1 == {} + else: + record = obj.load(str(tmp_path), record_digests=False) + assert torch.equal(obj.tensor("weight"), source) + assert obj.usable and record is obj.load_record + assert record["tensors"] == 1 and record["bytes_read"] == 32 + assert record["coverage"]["weight"]["source_sha256"] is None + assert phases == [True, True] and closed == [True] + if packed: + assert record["recipe"] == {"verified": True} and record["scalars"] == {} + finally: + real_close() + + +@pytest.mark.parametrize("dtype,itemsize", [("BF16", 2), ("F32", 4), ("F8_E4M3", 1), ("U8", 1)]) +def test_stored_entry_validation_is_shared(dtype: str, itemsize: int) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4.loader import ( + validate, + ) + + catalog = loader.SafetensorsCatalog.__new__(loader.SafetensorsCatalog) + catalog.entries = {"weight": loader.Entry("weight", "unused", dtype, (2, 4), 0, 8 * itemsize)} + validate(catalog, {"weight": (dtype, (2, 4))}) + catalog.entries["weight"].end -= 1 + with pytest.raises(loader.CheckpointFormatError, match="byte span"): + validate(catalog, {"weight": (dtype, (2, 4))}) + + +@pytest.mark.parametrize("packed", [False, True]) +def test_component_constructor_uses_format_storage(monkeypatch, packed: bool) -> None: + import torch + + from tensorrt_llm._torch._experimental.modeling_v2.models.beam._tp4 import ( + component_bf16 as bf16, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.beam.full_model_nvfp4__sm_103__tp4 import ( + component as nvfp4, + ) + + cls = ( + nvfp4.BeamCheckpointLayoutNvfp4Sm103Tp4 if packed else bf16.BeamCheckpointLayoutBf16Sm103Tp4 + ) + spec = layout.TensorSpec(layout.EMBED, (8, 2), layout.BF16, "embedding", None, None, 0) + monkeypatch.setattr(cls, "_tensor_specs", staticmethod(lambda *args: {spec.name: spec})) + monkeypatch.setattr(bf16, "AllReduce", lambda **kwargs: NS()) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (10, 3)) + real_empty = torch.empty + + def empty_on_cpu(*args, **kwargs): + kwargs["device"] = "cpu" + return real_empty(*args, **kwargs) + + monkeypatch.setattr(torch, "empty", empty_on_cpu) + mapping = NS( + rank=0, + world_size=4, + tp_size=4, + pp_size=1, + moe_tp_size=4, + moe_ep_size=1, + cp_size=1, + enable_attention_dp=False, + tp_group=list(range(4)), + ) + obj = cls(config(), mapping, layers=[], device=torch.device("cuda:0"), max_tokens=17) + assert obj.tensor(layout.EMBED).shape == (2, 2) + assert obj.vocab_interval == (0, 2) and obj.max_tokens == 17 + assert not obj.usable and obj.load_record is None + assert obj.local_heads == 12 and obj.local_kv_heads == 2 + if packed: + assert obj.expert_names == [] and obj._shared_fc1 == {} + + +def test_bringup_helpers_and_tensor_captures_are_not_runtime_apis(): + family = Path(next(iter(beam.__path__))) + assert not list(family.glob("*_bf16__sm_103__tp1")) + for source in family.rglob("*.py"): + tree = ast.parse(source.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef): + names = {a.arg for a in (*node.args.args, *node.args.kwonlyargs)} + assert not names.intersection({"capture", "layer_capture"}), source + elif isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith("_components"), source diff --git a/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_cublas_mm.py b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_cublas_mm.py new file mode 100644 index 000000000000..5e9af3f5d4e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_cublas_mm.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for cublas_mm, using the shared operator comparison.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm as op + +__extra_import_path__ = ["../../gemm", "../.."] +import test_modeling_v2_cublas_mm as shared # noqa: E402 + +_BEAM_DENSE_BLOCK = tuple( + Cell( + why=f"beam dense {name} {k}->{n}, {m} tokens", + spec=dict(m=m, k=k, n=n, dtype=torch.bfloat16, bias=False, out_dtype=None), + ) + for name, k, n in ( + ("q/gate", 5120, 6144), + ("k/v", 5120, 1024), + ("o", 6144, 5120), + ("gate/up", 5120, 20480), + ("down", 20480, 5120), + ) + for m in (1, 17, 129) +) + +_BEAM_HEAD = tuple( + Cell( + why=f"beam head 5120->151936, {m} tokens", + spec=dict(m=m, k=5120, n=151936, dtype=torch.bfloat16, bias=False, out_dtype=None), + ) + for m in (1, 17, 129) +) + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 {name} {k}->{n}, {m} tokens", + spec=dict(m=m, k=k, n=n, dtype=torch.bfloat16, bias=False, out_dtype=None), + ) + for name, k, n in ( + ("q/gate", 5120, 1536), + ("k/v", 5120, 256), + ("o", 1536, 5120), + ("dense gate/up", 5120, 5120), + ("dense down", 5120, 5120), + ("head", 5120, 37984), + ) + for m in (1, 17, 129) + ), + "beam/dense_block": _BEAM_DENSE_BLOCK, + "beam/six_layer_model": _BEAM_HEAD, +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell) -> None: + assert shared.op is op + shared.test_certified_cells(cell) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_fp4_quantize.py b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_fp4_quantize.py new file mode 100644 index 000000000000..12448c151c49 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_fp4_quantize.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for fp4_quantize, using the shared operator comparison.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize as op, +) + +__extra_import_path__ = ["../../quantization", "../.."] +import test_modeling_v2_fp4_quantize as shared # noqa: E402 + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 expert activation, {rows} rows at width {width}", + spec=dict(shape=(rows, width), dtype=torch.bfloat16, vec=16, swizzled=True), + ) + for width in (5120, 640) + for rows in (1, 17, 129) + ), +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell) -> None: + assert shared.op is op + shared.test_certified_cells(cell) diff --git a/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_nvfp4_gemm.py b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_nvfp4_gemm.py new file mode 100644 index 000000000000..f7366a571f9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/beam/test_modeling_v2_beam_nvfp4_gemm.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Beam component shapes for nvfp4_gemm, using the shared operator comparison.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.nvfp4_gemm import nvfp4_gemm as op + +__extra_import_path__ = ["../../gemm", "../.."] +import test_modeling_v2_nvfp4_gemm as shared # noqa: E402 + +TARGET_CELLS: dict[str, tuple[Cell, ...]] = { + "beam/full_model_nvfp4__sm_103__tp4": tuple( + Cell( + why=f"beam TP4 expert {name} {k}->{n}, {m} rows", + spec=dict(m=m, k=k, n=n, bias=False, backend="cutlass"), + ) + for name, k, n in (("gate/up", 5120, 640), ("down", 640, 5120)) + for m in (1, 17, 129) + ), +} + +# A shared shape used by multiple components needs only one kernel comparison. +CELLS = tuple( + { + repr(sorted(cell.spec.items())): cell for cells in TARGET_CELLS.values() for cell in cells + }.values() +) + + +@pytest.mark.parametrize("cell", CELLS, ids=[c.why for c in CELLS]) +def test_certified_cells(cell: Cell) -> None: + assert shared.op is op + shared.test_certified_cells(cell) diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py index 90499069c25f..6efd2e5703a7 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py @@ -240,6 +240,11 @@ def test_accuracy_directories_are_targets(repo: Path = _REPO): assert not problems, "\n".join(problems) +# Families that ship directly constructed components without a registered model. +# An exemption is invalid as soon as routing or registry imports are introduced. +_COMPONENT_ONLY_FAMILIES: set[str] = {"beam"} + + def test_every_routed_target_has_accuracy_tests(repo: Path = _REPO): """Every class a family's ``route`` returns lives in a target with accuracy tests. @@ -251,6 +256,15 @@ def test_every_routed_target_has_accuracy_tests(repo: Path = _REPO): for family_dir in _families(repo): routing = family_dir / "routing.py" where = routing.relative_to(repo) + if family_dir.name in _COMPONENT_ONLY_FAMILIES: + registry = repo / _PY / "models/__init__.py" + prefix = f"{_MODELS_PACKAGE}.{family_dir.name}" + imported = any( + name == prefix or name.startswith(prefix + ".") for name in _imports(registry, repo) + ) + if routing.exists() or imported: + problems.append(f"{where}: component-only exemption has routing or registration") + continue if not routing.is_file(): problems.append(f"{where}: missing") continue diff --git a/tests/unittest/scripts/cbts/test_modeling_v2.py b/tests/unittest/scripts/cbts/test_modeling_v2.py index cc325b13baff..8c3666b27ef4 100644 --- a/tests/unittest/scripts/cbts/test_modeling_v2.py +++ b/tests/unittest/scripts/cbts/test_modeling_v2.py @@ -447,6 +447,48 @@ def test_ci_selection_check_reads_imports_with_the_rule_parser() -> None: assert modeling_v2.routed_targets is modeling_v2_source.routed_targets +@pytest.mark.parametrize( + ("component_only", "routing", "imported", "accuracy", "error"), + [ + (True, False, False, False, None), + (False, False, False, False, "missing"), + (True, True, False, False, "component-only exemption"), + (True, False, True, False, "component-only exemption"), + (False, True, True, False, "has no"), + (False, True, True, True, None), + ], +) +def test_component_only_exemption_preserves_registered_accuracy_gate( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + component_only: bool, + routing: bool, + imported: bool, + accuracy: bool, + error: str | None, +) -> None: + """Only an explicitly unregistered component family can omit its accuracy gate.""" + repo = Path(__file__).resolve().parents[4] + path = repo / "tests/unittest/_torch/modeling_v2/test_modeling_v2_ci_selection.py" + spec = importlib.util.spec_from_file_location("modeling_v2_ci_selection", path) + check = importlib.util.module_from_spec(spec) + spec.loader.exec_module(check) + monkeypatch.setattr(check, "_COMPONENT_ONLY_FAMILIES", {"fam"} if component_only else set()) + family = tmp_path / _SRC / "models/fam" + family.mkdir(parents=True) + (family / "__init__.py").write_text("") + (family.parent / "__init__.py").write_text("from .fam import routing\n" if imported else "") + if routing: + (family / "routing.py").write_text(_routing("a")) + if accuracy: + (tmp_path / "tests/integration/defs" / _ACC / "fam/a").mkdir(parents=True) + if error: + with pytest.raises(AssertionError, match=error): + check.test_every_routed_target_has_accuracy_tests(tmp_path) + else: + check.test_every_routed_target_has_accuracy_tests(tmp_path) + + def test_docs_are_not_claimed(tmp_path: Path) -> None: rule = _rule(tmp_path) assert rule.apply(PRInputs(changed_files=[f"{_SRC}/models/fam/a/README.md"], diffs={})) is None @@ -492,7 +534,7 @@ def repository_cpu_rule() -> ModelingV2Rule: @pytest.mark.parametrize( - "family", ["deepseek_v3", "gpt_oss", "glm_moe_dsa", "kimi_k3_vl", "qwen4_exp"] + "family", ["deepseek_v3", "gpt_oss", "glm_moe_dsa", "kimi_k3_vl", "qwen4_exp", "beam"] ) def test_repository_cpu_selection_keeps_model_tests_in_their_family( repository_cpu_rule: ModelingV2Rule, family: str @@ -538,3 +580,26 @@ def test_repository_catalog_change_runs_unit_tests_importing_it( f"{_UNIT}/quantization", f"{_UNIT}/gemm/test_modeling_v2_nvfp4_gemm.py", } <= selected + + +@pytest.mark.parametrize( + ("category", "op", "test_path"), + [ + ("gemm", "cublas_mm", "test_modeling_v2_beam_cublas_mm.py"), + ("gemm", "nvfp4_gemm", "test_modeling_v2_beam_nvfp4_gemm.py"), + ("quantization", "fp4_quantize", "test_modeling_v2_beam_fp4_quantize.py"), + ("norm", "flashinfer_rmsnorm", "ops/norm/test_modeling_v2_beam_flashinfer_rmsnorm.py"), + ( + "activation", + "flashinfer_silu_and_mul", + "ops/activation/test_modeling_v2_beam_flashinfer_silu_and_mul.py", + ), + ("attention", "thop_attention", "ops/attention/test_modeling_v2_beam_thop_attention.py"), + ], +) +def test_shared_kernel_change_selects_beam_workload_shapes( + repository_cpu_rule: ModelingV2Rule, category: str, op: str, test_path: str +) -> None: + selected = repository_cpu_rule._prefixes_for(f"{_SRC}/catalog/{category}/{op}.py") + path = f"{_UNIT}/models/beam/{test_path}" + assert any(path == prefix or path.startswith(prefix + "/") for prefix in selected)