diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 880f4b536d8e..1f94dbc3a267 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -19,6 +19,28 @@ Pass the Hub ID or local path via `--model`: - [`nvidia/Cosmos3-Super-Image2Video-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step) — DMD2-distilled image-to-video: same fixed 4-step, guidance-baked-in contract. The default omni video shape (720p × 189 frames) is the deployed shape, so no dedicated config is needed. This checkpoint declares `default_use_system_prompt: true` in its `model_index.json`, which the pipeline applies automatically (override with `--use_system_prompt` / `--no-use_system_prompt`). - [`nvidia/Cosmos3-Edge`](https://huggingface.co/nvidia/Cosmos3-Edge) — 4B Nemotron-dense backbone supporting **T2I / T2V / I2V only**: no audio tower, and the checkpoint's action weights are not supported by this pipeline yet. 480p-native defaults (832×480 × 121 frames, 50 UniPC steps on the checkpoint-declared native flow schedule with shift 3.0, guidance 5.0; T2I defaults to 640×640), so no dedicated config is needed. The model card validates 256p/480p, 50–150 frames, and 12–30 FPS; requests outside that envelope run with an advisory log. +### Static FP8 checkpoints + +Statically quantized (ModelOpt) FP8 builds of Nano and Super run on this path. +Quantization is detected from the checkpoint's own metadata — pass the directory +to `--model` exactly as you would a BF16 one, with no extra flag: + +```bash +python cosmos3.py --model /path/to/Cosmos3-Nano-FP8 \ + --prompt_file prompts/t2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml +``` + +There are no FP8 Hub IDs yet, so use a local path. T2V, T2I, I2V and V2V are +validated on a **single GPU**; every multi-GPU configuration is refused with an +explicit error, so use BF16 there. + +These checkpoints ship the audio tower (`sound_gen: true`), so T2AV/TI2AV run +rather than being refused — audio is quantized and generated like any other +supported task. It simply has not been exercised as thoroughly as the four +video/image tasks above, and no FP8 audio quality claim is made. FP8 output +quality in general has not been benchmarked against BF16. + ## Guardrails Guardrails are enabled by default (required by the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)). Install and authenticate as follows: diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 2c6fe00303e5..02712cdc734b 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -44,7 +44,8 @@ from flashinfer.fp4_quantization import nvfp4_quantize as _flashinfer_nvfp4_quantize from ..modules.multi_stream_utils import do_multi_stream -from ..modules.swiglu import silu_and_mul_kernel +from ..modules.swiglu import (get_silu_b200_tuning_params, + silu_and_mul_2in_kernel, silu_and_mul_kernel) from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets, fp4_scale_infer_shape, get_last_power_of_2_num_tokens_buckets, @@ -2087,6 +2088,77 @@ def _( return x.new_empty((b, d), dtype=o_dtype) +@torch.library.custom_op("trtllm::silu_and_mul_2in", mutates_args=()) +def silu_and_mul_2in(gate: torch.Tensor, + up: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None) -> torch.Tensor: + """silu_and_mul for gate and up held in separate tensors. + + Equivalent to silu_and_mul(cat([gate, up], -1), ...) with no concatenation. + """ + # Rank-agnostic: the kernel walks a flat run of numel() elements, so any + # matching contiguous shape works. Models carry rank-3 [batch, seq, hidden] + # activations, and requiring rank 2 here would force callers to reshape. + assert gate.shape == up.shape, ( + f"gate and up must have the same shape, got {tuple(gate.shape)} and " + f"{tuple(up.shape)}") + assert gate.dtype == up.dtype, ( + f"gate and up must have the same dtype, got {gate.dtype} and {up.dtype}" + ) + assert gate.device == up.device, ( + f"gate and up must be on the same device, got {gate.device} and " + f"{up.device}") + # The kernel indexes both operands as flat contiguous runs, so a strided + # view would read the wrong addresses and silently return garbage. Linear + # outputs are contiguous; reject anything else rather than copying. + assert gate.is_contiguous() and up.is_contiguous(), ( + f"gate and up must be contiguous, got strides {gate.stride()} and " + f"{up.stride()} for shape {tuple(gate.shape)}") + + o_dtype = dtype or gate.dtype + o = torch.empty(gate.shape, dtype=o_dtype, device=gate.device) + + # Validated and performance-tuned on B200 only. Other architectures are + # unvalidated and are not rejected; no performance guarantee is made. + block_elements, num_warps = get_silu_b200_tuning_params(o_dtype) + n_elements = gate.numel() + + silu_and_mul_2in_kernel[(triton.cdiv(n_elements, block_elements), )]( + o_ptr=o, + o_scale_ptr=scale, + gate_ptr=gate, + up_ptr=up, + n_elements=n_elements, + swiglu_limit=swiglu_limit or 0.0, + swiglu_alpha=swiglu_alpha if swiglu_alpha is not None else 1.0, + swiglu_beta=swiglu_beta if swiglu_beta is not None else 0.0, + BLOCK_SIZE=block_elements, + HAS_O_SCALE=scale is not None, + HAS_SWIGLU_LIMIT=swiglu_limit is not None and swiglu_limit > 0.0, + num_warps=num_warps, + ) + + return o + + +@silu_and_mul_2in.register_fake +def _( + gate: torch.Tensor, + up: torch.Tensor, + scale: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None, +) -> torch.Tensor: + o_dtype = dtype or gate.dtype + return gate.new_empty(gate.shape, dtype=o_dtype) + + class AllReduceRunner(TunableRunner): _prealloc_lock: ClassVar[threading.Lock] = threading.Lock() _prealloc_done: ClassVar[set] = set() diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index afd8796c3e10..76c030efa784 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -14,7 +14,7 @@ from ..utils import Fp4QuantizedTensor from .linear import (Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig, is_static_nvfp4_input_eligible) -from .swiglu import swiglu +from .swiglu import swiglu, swiglu_2in class GatedMLP(nn.Module): @@ -38,6 +38,7 @@ def __init__( swiglu_limit: Optional[float] = None, swiglu_alpha: Optional[float] = None, swiglu_beta: Optional[float] = None, + split_gate_up: bool = False, ): super().__init__() @@ -46,6 +47,20 @@ def __init__( self.intermediate_size = intermediate_size self.activation = activation self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm + # Keeps each projection's own calibrated scale, which fusing would + # discard. Off by default. + self.split_gate_up = split_gate_up + if split_gate_up and (config + or ModelConfig()).force_dynamic_quantization: + # The activation emits FP8 using down_proj's calibrated input_scale, + # which makes down_proj skip the dynamic quantization it was asked + # for. The fused path has the same behaviour, but rather than carry + # that ambiguity into a new topology, require the fused one -- the + # only consumer of split_gate_up is static per-tensor FP8. + raise ValueError( + "GatedMLP: split_gate_up is incompatible with " + "force_dynamic_quantization; dynamic quantization requires the " + "fused gate/up topology") self.swiglu_limit = float( swiglu_limit) if swiglu_limit is not None else None # SwiGLU-OAI shape parameters, left None for plain SwiGLU, where the @@ -103,15 +118,11 @@ def __init__( 'up': (local_intermediate_start, local_intermediate_end), } - self.gate_up_proj = Linear( - self.hidden_size, - self.intermediate_size * 2, + _common_proj_kwargs = dict( bias=bias, dtype=dtype, mapping=mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, - weights_loading_config=WeightsLoadingConfig( - weight_mode=WeightMode.FUSED_GATE_UP_LINEAR), quant_config=config.get_quant_config(), reduce_output=False, skip_create_weights_in_init=config.skip_create_weights_in_init, @@ -120,11 +131,30 @@ def __init__( use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, disable_deep_gemm=disable_deep_gemm, - fused_weight_shard_indices_mapping=gateup_shard_indices_mapping, use_custom_cublas_mm=use_custom_cublas_mm, - override_tp_sharding=override_tp_sharding, ) + if self.split_gate_up: + # Each Linear owns one checkpoint tensor and its own scale, so no + # fused weight mode, no shard-index mapping, and no gate/up-keyed + # override_tp_sharding -- Linear asserts that a dict tp_sharding + # only ever reaches a fused weight mode. + self.gate_proj = Linear(self.hidden_size, self.intermediate_size, + **_common_proj_kwargs) + self.up_proj = Linear(self.hidden_size, self.intermediate_size, + **_common_proj_kwargs) + self.gate_up_proj = None + else: + self.gate_up_proj = Linear( + self.hidden_size, + self.intermediate_size * 2, + weights_loading_config=WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_GATE_UP_LINEAR), + fused_weight_shard_indices_mapping=gateup_shard_indices_mapping, + override_tp_sharding=override_tp_sharding, + **_common_proj_kwargs, + ) + if is_shared_expert: down_type = LoraModuleType.SHARED_EXPERT_4H_TO_H h_to_4h_type = LoraModuleType.SHARED_EXPERT_H_TO_4H @@ -210,6 +240,102 @@ def _is_plain_swiglu(self): return ((self.swiglu_alpha is None or self.swiglu_alpha == 1.0) and (self.swiglu_beta is None or self.swiglu_beta == 0.0)) + def _apply_activation_2in(self, gate, up): + """Activation for the split path: gate and up arrive as two tensors. + + Mirrors _apply_activation, including emitting FP8 directly when down_proj + consumes FP8 -- concatenating the pair first would cost a full + intermediate-sized copy. + """ + if self.activation is not F.silu: + raise NotImplementedError( + f"split_gate_up requires SwiGLU activation, got {self.activation}" + ) + # As in _shares_gate_up_quantization: a method running this call in + # higher precision needs a 16-bit activation, so do not emit FP8 for it. + down_proj_is_fp8 = ( + self.down_proj.has_fp8_qdq + or self.down_proj.has_w4a8_nvfp4_fp8) and not getattr( + self.down_proj.quant_method, "high_precision", False) + if down_proj_is_fp8: + return swiglu_2in(gate, + up, + quant_scale=self.down_proj.input_scale, + quant_type=torch.float8_e4m3fn, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta) + return swiglu_2in(gate, + up, + swiglu_limit=self.swiglu_limit, + swiglu_alpha=self.swiglu_alpha, + swiglu_beta=self.swiglu_beta) + + def _shares_gate_up_quantization(self) -> bool: + """Config-only half of the condition, so post_load_weights() can reuse it. + + Dynamic quantization derives a scale per call, so differing calibrated + scales are legitimate there rather than an error. + """ + return ( + self.split_gate_up and self.gate_proj.has_fp8_qdq + and self.up_proj.has_fp8_qdq + and not self.gate_proj.force_dynamic_quantization + and not self.up_proj.force_dynamic_quantization + and self.gate_proj.input_scale is not None + and self.up_proj.input_scale is not None + # A quantization method may run a given call in higher + # precision than its checkpoint recipe -- Cosmos3 does this on + # the outer denoising steps -- by publishing ``high_precision``. + # Quantizing the shared activation here would hand such a call + # a tensor it must not receive, so leave it in its input dtype. + and + not getattr(self.gate_proj.quant_method, "high_precision", False)) + + def _can_share_gate_up_quantization(self, x) -> bool: + """Whether gate and up can consume one quantized activation. + + Reads no tensor values: that would sync the device every forward and + make the graph data-dependent. Scale equality is checked at load. + """ + return (self._shares_gate_up_quantization() + and not isinstance(x, Fp4QuantizedTensor) + and x.dtype != torch.float8_e4m3fn) + + def _split_gate_up_forward(self, x): + """Run the split projections, quantizing their shared input once. + + gate and up consume the same activation. With static per-tensor FP8 each + Linear would otherwise quantize it again, so quantize once here and hand + both the FP8 tensor -- FP8QDQLinearMethod.apply passes a pre-quantized + input straight through. + """ + if self._can_share_gate_up_quantization(x): + shape = x.shape + x2d = x.reshape(-1, shape[-1]) if x.dim() > 2 else x + qx, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( + x2d, self.gate_proj.input_scale) + x = qx.reshape(*shape[:-1], shape[-1]) if x.dim() > 2 else qx + return self._apply_activation_2in(self.gate_proj(x), self.up_proj(x)) + + def post_load_weights(self) -> None: + """Check the shared-activation invariant here, never in forward(). + + Reading scale tensors on the hot path would sync the device and break + fullgraph compilation. Only runs if the model's post-load walk includes + GatedMLP; several models restrict theirs to Linear. + """ + if not self._shares_gate_up_quantization(): + return + gate_scale, up_scale = (self.gate_proj.input_scale, + self.up_proj.input_scale) + if not torch.equal(gate_scale, up_scale): + raise ValueError( + "split gate/up share one quantized activation, so they must " + f"carry the same calibrated input_scale; got {gate_scale.item()} " + f"(gate) and {up_scale.item()} (up) at layer_idx=" + f"{self.layer_idx}") + def _can_fuse_gate_up_swiglu(self): """Check if fused GEMM + SwiGLU path is available. @@ -219,6 +345,8 @@ def _can_fuse_gate_up_swiglu(self): - gate_up_proj uses NVFP4 quantization - gate_up_proj has no bias (bias not supported in fused kernel) """ + if self.gate_up_proj is None: # split path has no fused projection + return False return (self.use_cute_dsl_blockscaling_mm and self.activation == F.silu and self._is_plain_swiglu() and self.gate_up_proj.has_nvfp4_activation_quantization @@ -311,9 +439,25 @@ def forward( **kwargs, ) -> torch.Tensor: if bool(lora_params): + if self.split_gate_up: + # forward_lora fuses gate/up LoRA into one projection, which the + # split path does not have. Not implemented rather than + # unsupported-by-accident: FP8 LoRA already falls back to BF16 + # activation (see _apply_activation), so a split FP8 + LoRA path + # would need the FP8 LoRA grouped GEMM first. + raise NotImplementedError( + "GatedMLP: LoRA is not supported with split_gate_up=True " + f"(layer_idx={self.layer_idx}); build the fused topology " + "if LoRA is required") return self.forward_lora(x, all_rank_num_tokens, final_all_reduce_params, lora_params) + if self.split_gate_up: + h2 = self._split_gate_up_forward(x) + return self.down_proj(h2, + all_reduce_params=final_all_reduce_params, + layer_idx=self.layer_idx) + if self._can_fuse_gate_up_swiglu_fp4out(): # During torch.compile the token dim is a SymInt, so `m >= MIN_M` # would create a SymBool guard that breaks piecewise CUDA graph diff --git a/tensorrt_llm/_torch/modules/swiglu.py b/tensorrt_llm/_torch/modules/swiglu.py index 6462705ec961..2633f11e8998 100644 --- a/tensorrt_llm/_torch/modules/swiglu.py +++ b/tensorrt_llm/_torch/modules/swiglu.py @@ -61,6 +61,68 @@ def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, tl.store(o_row_ptr + offsets, result, mask=mask) +@triton.jit +def silu_and_mul_2in_kernel(o_ptr, o_scale_ptr, gate_ptr, up_ptr, n_elements, + swiglu_limit: tl.constexpr, + swiglu_alpha: tl.constexpr, + swiglu_beta: tl.constexpr, BLOCK_SIZE: tl.constexpr, + HAS_O_SCALE: tl.constexpr, + HAS_SWIGLU_LIMIT: tl.constexpr) -> None: + """As silu_and_mul_kernel, but gate and up come from separate tensors. + + Used when a model keeps gate/up as distinct projections instead of one fused + Linear, so their outputs are never adjacent in memory. + + Indexes both operands as flat contiguous runs, so callers must pass + contiguous tensors -- the op enforces this. + """ + offsets = tl.program_id(axis=0).to(tl.int64) * BLOCK_SIZE + tl.arange( + 0, BLOCK_SIZE) + mask = offsets < n_elements + + a = tl.load(gate_ptr + offsets, mask=mask).to(tl.float32) + b = tl.load(up_ptr + offsets, mask=mask).to(tl.float32) + + if HAS_SWIGLU_LIMIT: + a = tl.minimum(a, swiglu_limit) + b = tl.clamp(b, -swiglu_limit, swiglu_limit) + + result = a * tl.sigmoid(swiglu_alpha * a) * (b + swiglu_beta) + + if HAS_O_SCALE: + o_scale = tl.load(o_scale_ptr) + result = scale_and_clamp(result, o_scale, o_ptr.dtype.element_ty) + + tl.store(o_ptr + offsets, result, mask=mask) + + +def get_silu_b200_tuning_params(out_dtype: torch.dtype) -> tuple[int, int]: + """Launch parameters for silu_and_mul_2in_kernel, measured on B200 (sm100). + + Returns ``(block_elements, num_warps)``. ``block_elements`` is how many + elements one Triton program handles -- Triton's ``BLOCK_SIZE`` -- not a CUDA + thread count; threads are ``num_warps * 32``, so 4 warps is 128 threads each + covering ``block_elements / 128`` elements. + + The output dtype matters because it sets the read/write balance: an FP8 + output writes one byte per element where BF16 writes two, leaving the FP8 + case more read-dominated and better served by more elements in flight. From + a sweep over ``block_elements`` {1024..8192} x ``num_warps`` {2, 4, 8} at the + Cosmos3 shapes, FP8 output was ~5% faster at 4096 than at 2048, while the + BF16 candidates tied within 0.6%. + + Validated and performance-tuned on B200 only. Other architectures are + unvalidated and are not rejected; no performance guarantee is made. + + When another architecture is tuned, add its own + ``get_silu__tuning_params`` and dispatch on device capability here + rather than widening this one. + """ + if out_dtype == torch.float8_e4m3fn: + return 4096, 4 + return 2048, 4 + + def swiglu(x, quant_scale: Optional[torch.Tensor] = None, quant_type=None, @@ -82,3 +144,34 @@ def swiglu(x, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta) + + +def swiglu_2in(gate, + up, + quant_scale: Optional[torch.Tensor] = None, + quant_type=None, + swiglu_limit: Optional[float] = None, + swiglu_alpha: Optional[float] = None, + swiglu_beta: Optional[float] = None): + """SiLU(gate) * up for separately projected gate and up tensors. + + Equivalent to ``swiglu(torch.cat([gate, up], dim=-1), ...)`` without + materializing the concatenation. + """ + if quant_scale is not None: + assert quant_type is not None + return torch.ops.trtllm.silu_and_mul_2in( + gate, + up, + scale=quant_scale, + dtype=quant_type, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + + return torch.ops.trtllm.silu_and_mul_2in(gate, + up, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index c7ff75392046..6ad96adfa007 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1625,6 +1625,14 @@ def forward_fn( Since Cosmos3 embeds text internally, we pass token IDs via extra_tensors rather than through encoder_hidden_states. """ + # Activation precision for this step. A pure function of step_index, + # so the conditional and unconditional CFG branches of one step + # always select the same path even though each calls this + # separately. No-op unless the checkpoint is static FP8. + self.transformer.set_denoising_step( + step_index=step_index, num_steps=len(self.scheduler.timesteps) + ) + current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None result = self.transformer( @@ -1673,26 +1681,33 @@ def post_step_fn(step_latents): if do_audio: extra_streams = {"audio": (audio_latents, self.audio_scheduler)} should_pin_condition_latents = condition_latents is not None and velocity_mask is not None - denoise_result = self.denoise( - latents=latents, - scheduler=self.scheduler, - prompt_embeds=cond_ids, # placeholder — actual conditioning via extra_cfg_tensors - neg_prompt_embeds=uncond_ids, - guidance_scale=guidance_scale, - forward_fn=forward_fn, - extra_cfg_tensors=extra_cfg_tensors, - extra_streams=extra_streams, - guidance_interval=guidance_interval, - # V2V pins the conditioning latents; distilled I2V re-anchors the - # conditioning frame. A request carries an image or a video, never - # both, so at most one of these applies. - post_step_fn=( - post_step_fn - if should_pin_condition_latents - else self._conditioning_anchor_post_step(image_latent) - ), - scheduler_step_kwargs=self.sampling.scheduler_step_kwargs(generator), - ) + try: + denoise_result = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=cond_ids, # placeholder — actual conditioning via extra_cfg_tensors + neg_prompt_embeds=uncond_ids, + guidance_scale=guidance_scale, + forward_fn=forward_fn, + extra_cfg_tensors=extra_cfg_tensors, + extra_streams=extra_streams, + guidance_interval=guidance_interval, + # V2V pins the conditioning latents; distilled I2V re-anchors the + # conditioning frame. A request carries an image or a video, never + # both, so at most one of these applies. + post_step_fn=( + post_step_fn + if should_pin_condition_latents + else self._conditioning_anchor_post_step(image_latent) + ), + scheduler_step_kwargs=self.sampling.scheduler_step_kwargs(generator), + ) + finally: + # In a finally because a failed request must not leave the selection + # latched. The transfer path runs the transformer without selecting + # a step, so it would inherit whatever the failed request left + # behind and run every call in 16-bit with nothing to indicate it. + self.transformer.reset_denoising_step() if extra_streams is not None: latents, extra_latents = denoise_result diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/step_precision.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/step_precision.py new file mode 100644 index 000000000000..ed30d17399e7 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/step_precision.py @@ -0,0 +1,293 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-denoising-step activation precision for static-FP8 Cosmos3. + +A ModelOpt-calibrated checkpoint carries one activation scale per projection, +taken as a max over the whole sampling trajectory. That single scale fits the +outer denoising steps worst, so this module lets those steps run the resident +FP8 weights through a 16-bit GEMM instead: the weight is dequantized with its +own ``weight_scale`` and ``input_scale`` goes unused. Middle steps keep the +checkpoint's fully quantized path. + +No extra weights are read -- the weights and scales are the ones already +loaded, and no second checkpoint or persistent dequantized copy is kept. What +the checkpoint does supply is the policy itself, under +``quantization_config.runtime.diffusion_step_policy``: which steps take the +16-bit path, and what the understanding tower does. Only the builds whose +calibration needs it carry one, so a checkpoint without a policy runs fully +quantized and there is no default to substitute. + +Precision is selected once per denoising step, before any transformer call for +that step, so the conditional and unconditional CFG branches of one step always +agree. +""" + +from dataclasses import dataclass +from typing import Any, Iterable, Mapping, Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.modules.linear import FP8QDQLinearMethod, Linear + +# The policy the checkpoint publishes under +# ``quantization_config.runtime.diffusion_step_policy``. Every field is +# required and every value is checked: a policy shape we do not implement must +# fail loudly rather than be half-honoured, since silently ignoring a field the +# producer set is indistinguishable from the feature not working. +_POLICY_FIELDS = frozenset( + { + "schema_version", + "type", + "index_space", + "scope", + "default_mode", + "first_steps", + "last_steps", + "overlap", + "reasoner", + } +) +_STEP_RANGE_FIELDS = frozenset({"count", "mode"}) +_SCOPE_COMPONENT = "transformer" + + +@dataclass(frozen=True) +class StepPrecisionPolicy: + """A validated ``diffusion_step_policy``. + + ``reasoner`` is not step-scoped. The understanding tower runs once per + request -- on the first transformer call, then cached -- so a step-indexed + rule would describe it only by accident of which step that call lands on. + The policy states its precision directly instead. + """ + + first_steps: int + last_steps: int + reasoner_high_precision: bool + + +def parse_diffusion_step_policy(quantization_config: Any) -> Optional[StepPrecisionPolicy]: + """Read the checkpoint's step policy, or None if it declares none. + + Absence is meaningful, not a default to fill in: the producer ships this + only for the checkpoints whose calibration needs it (the multi-step video + builds), and deliberately omits it for the image and distilled 4-step + builds, whose output does not show the artifact it targets. + """ + if not isinstance(quantization_config, Mapping): + return None + runtime = quantization_config.get("runtime") + if not isinstance(runtime, Mapping) or "diffusion_step_policy" not in runtime: + return None + policy = runtime["diffusion_step_policy"] + if not isinstance(policy, Mapping): + raise TypeError( + "quantization_config.runtime.diffusion_step_policy must be a mapping, " + f"got {type(policy).__name__}" + ) + + unknown = set(policy) - _POLICY_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy fields: {sorted(unknown)}") + missing = _POLICY_FIELDS - set(policy) + if missing: + raise ValueError(f"Missing diffusion_step_policy fields: {sorted(missing)}") + + schema_version = policy["schema_version"] + if ( + not isinstance(schema_version, int) + or isinstance(schema_version, bool) + or schema_version != 1 + ): + raise ValueError( + f"diffusion_step_policy.schema_version must be the integer 1, got {schema_version!r}" + ) + for field, expected in ( + ("type", "first_last_n"), + ("index_space", "denoising_loop_iteration"), + ("default_mode", "native"), + # Which mode wins where the first and last windows meet. "a16" is what + # the ``or`` below already yields, so this is checked rather than acted + # on; any other value would need a different predicate. + ("overlap", "a16"), + ): + if policy[field] != expected: + raise ValueError( + f"diffusion_step_policy.{field} must be {expected!r}, got {policy[field]!r}" + ) + + scope = policy["scope"] + if not isinstance(scope, list) or not scope or not all(isinstance(s, str) for s in scope): + raise TypeError("diffusion_step_policy.scope must be a non-empty list of strings") + if _SCOPE_COMPONENT not in scope: + return None + + reasoner = policy["reasoner"] + if reasoner not in ("native", "a16"): + raise ValueError( + f"diffusion_step_policy.reasoner must be 'native' or 'a16', got {reasoner!r}" + ) + + return StepPrecisionPolicy( + first_steps=_parse_step_range(policy["first_steps"], "first_steps"), + last_steps=_parse_step_range(policy["last_steps"], "last_steps"), + reasoner_high_precision=reasoner == "a16", + ) + + +def _parse_step_range(value: Any, name: str) -> int: + if not isinstance(value, Mapping): + raise TypeError(f"diffusion_step_policy.{name} must be a mapping") + unknown = set(value) - _STEP_RANGE_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy.{name} fields: {sorted(unknown)}") + missing = _STEP_RANGE_FIELDS - set(value) + if missing: + raise ValueError(f"Missing diffusion_step_policy.{name} fields: {sorted(missing)}") + if value["mode"] != "a16": + raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16', got {value['mode']!r}") + count = value["count"] + if not isinstance(count, int) or isinstance(count, bool) or count < 0: + raise TypeError(f"diffusion_step_policy.{name}.count must be a non-negative integer") + return count + + +class StepPrecisionController: + """Holds the activation precision selected for the current denoising step.""" + + def __init__(self, first_steps: int, last_steps: int) -> None: + if first_steps < 0 or last_steps < 0: + raise ValueError( + f"first_steps/last_steps must be non-negative, got {first_steps}/{last_steps}" + ) + self.first_steps = first_steps + self.last_steps = last_steps + self.high_precision = False + + def set_step(self, step_index: int, num_steps: int) -> None: + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + if step_index < 0 or step_index >= num_steps: + raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}") + # A single-step schedule is the warmup probe rather than a real + # request, and treating every step as an edge step would make warmup + # exercise a path the measured run never takes. + if num_steps == 1: + self.high_precision = False + return + self.high_precision = ( + step_index < self.first_steps or step_index >= num_steps - self.last_steps + ) + + def reset(self) -> None: + self.high_precision = False + + +def apply_fp8_w8a16_linear( + module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor] +) -> torch.Tensor: + """16-bit GEMM against the module's resident FP8 weight. + + ``FP8QDQLinearMethod.create_weights`` allocates ``weight`` as ``[out, in]`` + float8, which is the layout ``F.linear`` wants, so no transpose is needed. + ``weight_scale`` is the per-tensor scalar and broadcasts. ``input_scale`` is + deliberately unused -- leaving the activation unquantized is the point. + """ + if input.dtype == torch.float8_e4m3fn: + raise RuntimeError( + "step precision: a high-precision step received an already-quantized " + "activation. A caller that pre-quantizes a shared activation (fused " + "gate/up or shared q/k/v) must stand down while high_precision is set, " + "otherwise the step is not actually running in 16-bit." + ) + weight = module.weight.to(input.dtype) * module.weight_scale.to(input.dtype) + return F.linear(input, weight, bias) + + +class StepPrecisionFp8LinearMethod: + """Dispatches each call to the checkpoint's FP8 path or the 16-bit path. + + ``always_high`` serves the reasoner. The understanding tower builds its KV + cache on the first transformer call of a request and is cached after, so a + step-indexed decision would only match the policy while that call happens to + land inside a window -- true for the published 3/3 policy, false the moment + one ships ``first_steps: 0``. The policy states its precision directly. + """ + + def __init__( + self, + base_method: FP8QDQLinearMethod, + controller: StepPrecisionController, + always_high: bool = False, + ): + self.base_method = base_method + self.controller = controller + self.always_high = always_high + + @property + def high_precision(self) -> bool: + """Published so activation-sharing callers can stand down for this step.""" + return self.always_high or self.controller.high_precision + + def __getattr__(self, name: str): + # Everything not overridden here (create_weights, load_weights, + # process_weights_after_loading, the quantizes_* properties Linear + # queries) belongs to the wrapped method. + return getattr(self.base_method, name) + + def apply( + self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor] = None + ) -> torch.Tensor: + if self.high_precision: + return apply_fp8_w8a16_linear(module, input, bias) + return self.base_method.apply(module, input, bias) + + +def linear_runs_high_precision(module: Optional[Linear]) -> bool: + """Whether this Linear is currently on the 16-bit path. + + The contract a wrapped quantization method publishes: activation-sharing + callers quantize once above the Linear, which would defeat the 16-bit step, + so they consult this before doing so. + """ + if module is None: + return False + return bool(getattr(module.quant_method, "high_precision", False)) + + +def install_step_precision( + roots: Iterable[torch.nn.Module], + controller: StepPrecisionController, + always_high: bool = False, +) -> int: + """Wrap every static-FP8 Linear under *roots* for per-call dispatch. + + Must run after weight loading: the wrapper forwards the load-time hooks to + the base method, but wrapping earlier would put it in the path of the + loader's ``isinstance`` checks. Returns the number of wrapped modules; 0 + means this is not a static-FP8 model. + """ + wrapped = 0 + for root in roots: + for module in root.modules(): + if not isinstance(module, Linear): + continue + existing = module.quant_method + if isinstance(existing, StepPrecisionFp8LinearMethod): + # Installing twice must not leave wrappers pointing at a + # controller nobody drives: rebind them to the live one instead + # of skipping, or set_denoising_step would silently stop + # reaching the layers it is supposed to steer. + existing.controller = controller + existing.always_high = always_high + wrapped += 1 + continue + if not isinstance(existing, FP8QDQLinearMethod): + continue + module.quant_method = StepPrecisionFp8LinearMethod( + existing, controller, always_high=always_high + ) + wrapped += 1 + return wrapped diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index ebe83e605626..3585c4e468c4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -29,6 +29,11 @@ from tensorrt_llm._torch.modules.mlp import MLP from tensorrt_llm._torch.utils import relu2 from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.cosmos3.step_precision import ( + StepPrecisionController, + install_step_precision, + parse_diffusion_step_policy, +) from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader @@ -59,6 +64,27 @@ def apply_pretrained_config_compat_defaults( return pretrained_config +def uses_static_fp8(model_config: DiffusionModelConfig) -> bool: + """Whether this run consumes a statically quantized (calibrated) FP8 checkpoint. + + Such checkpoints store a separate calibrated scale per projection. Fusing + q/k/v or gate/up into one Linear forces a single scale on the group and + re-quantizes the other members onto it, discarding their calibration, so + those groups are kept as separate projections here instead. + + Dynamic quantization derives scales at load or per call, so it has no + calibration to preserve and keeps the fused topology. + """ + quant_config = model_config.quant_config + if quant_config is None or getattr(quant_config, "layer_quant_mode", None) is None: + return False + return ( + quant_config.layer_quant_mode.has_fp8_qdq() + and not model_config.force_dynamic_quantization + and not model_config.dynamic_weight_quant + ) + + COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" @@ -418,6 +444,7 @@ def __init__( layer_idx=layer_idx, module_name=module_name, enable_sequence_parallel=False, + share_qkv_input_quant=uses_static_fp8(model_config), ) # Attention Q/K norms run the fp32-weight-multiply flavor in both # recipes (this path has always been F.rms_norm); only the layernorms @@ -521,12 +548,14 @@ def __init__( key=(type(self).__name__, original_backend, "VANILLA"), ) + static_fp8 = uses_static_fp8(model_config) + super().__init__( hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_key_value_heads=num_key_value_heads, head_dim=head_dim, - qkv_mode=QKVMode.FUSE_QKV, + qkv_mode=QKVMode.SEPARATE_QKV if static_fp8 else QKVMode.FUSE_QKV, qk_norm=False, qk_norm_mode="per_head", bias=False, @@ -534,6 +563,7 @@ def __init__( layer_idx=layer_idx, module_name=module_name, enable_sequence_parallel=True, + share_qkv_input_quant=static_fp8, ) model_config.attention.backend = original_backend @@ -624,6 +654,7 @@ def _build_cosmos3_mlp( config=model_config, layer_idx=layer_idx, reduce_output=model_config.mapping.tp_size > 1, + split_gate_up=uses_static_fp8(model_config), ) return MLP( hidden_size=hidden_size, @@ -940,6 +971,8 @@ def __init__(self, model_config: DiffusionModelConfig): ) pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) self.recipe = resolve_arch_recipe(pretrained_config) + # Installed in post_load_weights when the checkpoint is static FP8. + self.step_precision_controller: Optional[StepPrecisionController] = None self.audio_gen = getattr(pretrained_config, "sound_gen", False) # Config fact only: the transformer never constructs action modules. self.has_action_weights = getattr(pretrained_config, "action_gen", False) @@ -1003,6 +1036,27 @@ def __init__(self, model_config: DiffusionModelConfig): "Ring parallelism is not supported for Cosmos3 cross-attention." ) + if uses_static_fp8(model_config): + # Static FP8 puts cross-attention on SEPARATE_QKV, which the parallel + # wrappers treat differently from a fused QKV: Attention2D silently + # falls back off Ulysses rather than failing. Reject the untested + # combinations outright instead of degrading quietly. + # cp_size unifies ring and Attention2D; ring is already rejected above. + unsupported = { + "tp_size": tp_size, + "ulysses_size": ulysses_size, + "cfg_size": vgm.cfg_size if vgm else 1, + "cp_size": vgm.cp_size if vgm else 1, + "parallel_vae_size": vgm.parallel_vae_size if vgm else 1, + } + engaged = {k: v for k, v in unsupported.items() if v > 1} + if engaged: + raise NotImplementedError( + "Static FP8 Cosmos3 is supported on one GPU only (each of " + f"{sorted(unsupported)} must be 1); got {engaged}. Use the " + "BF16 checkpoint for multi-GPU." + ) + self.language_model = Cosmos3LanguageModel(model_config, self.recipe) self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) @@ -1694,3 +1748,78 @@ def post_load_weights(self) -> None: for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): module.post_load_weights() + + # Second pass: GatedMLP and Attention validate invariants across their + # own projections, so they must run after those projections finalize. + # named_modules() yields parents before children, so folding this into + # the loop above would check the scales too early. + for _, module in self.named_modules(): + if isinstance(module, (GatedMLP, Attention)): + module.post_load_weights() + + self._maybe_install_step_precision() + + def _maybe_install_step_precision(self) -> None: + """Honor the checkpoint's ``diffusion_step_policy``, if it declares one. + + Runs last: the wrapper only ever needs to dispatch ``apply``, and the + projections must already be finalized. The policy is the checkpoint's + to state -- it ships only with the builds whose calibration needs it -- + so there is no default to apply when it is absent, and no knob here to + turn it on for a checkpoint that did not ask. + + Only static FP8 qualifies: under dynamic quantization the scale is + derived per call, so there is no calibration mismatch to avoid. + """ + policy = parse_diffusion_step_policy( + getattr(self.model_config.pretrained_config, "quantization_config", None) + ) + if policy is None: + return + if not uses_static_fp8(self.model_config): + logger.warning( + "Checkpoint declares a diffusion_step_policy but this run is not static " + "per-tensor FP8, so the policy does not apply and is ignored." + ) + return + + self.step_precision_controller = StepPrecisionController( + first_steps=policy.first_steps, + last_steps=policy.last_steps, + ) + # The generation tower follows the step windows. The reasoner's + # precision is stated outright: it runs once per request, on whichever + # transformer call builds its KV cache, so deriving it from a step index + # would match the policy only by coincidence. + wrapped = install_step_precision([self.gen_layers], self.step_precision_controller) + if policy.reasoner_high_precision: + wrapped += install_step_precision( + [self.language_model.layers], self.step_precision_controller, always_high=True + ) + if wrapped == 0: + logger.warning( + "Checkpoint declares a diffusion_step_policy, but no static-FP8 linears " + "were found to wrap; every step will run fully quantized." + ) + self.step_precision_controller = None + return + logger.info( + f"Cosmos3 diffusion_step_policy: {wrapped} FP8 linears wrapped; first " + f"{policy.first_steps} and last {policy.last_steps} denoising steps run with " + f"BF16 activations, reasoner " + f"{'always BF16' if policy.reasoner_high_precision else 'native'}." + ) + + def set_denoising_step(self, step_index: int, num_steps: int) -> None: + """Select this step's activation precision, before any transformer call. + + Called once per denoising step so a step's conditional and + unconditional CFG branches cannot disagree. + """ + if self.step_precision_controller is not None: + self.step_precision_controller.set_step(step_index, num_steps) + + def reset_denoising_step(self) -> None: + """Drop per-request precision state so it cannot leak into the next request.""" + if self.step_precision_controller is not None: + self.step_precision_controller.reset() diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 96256ea9a82a..d5f0dd02b993 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -58,6 +58,7 @@ def __init__( enable_sequence_parallel: bool = True, async_ulysses: bool = False, separate_qkv_is_self_attention: bool = False, + share_qkv_input_quant: bool = False, ): super().__init__() @@ -150,6 +151,27 @@ def __init__( and not self.force_dynamic_quantization ) + # Opt-in FP8 analog of the above, for the synchronous get_qkv() path. + # Kept opt-in rather than inferred from quant_config so that enabling it + # also commits the caller to running post_load_weights(), which checks + # the equal-input_scale invariant that makes the sharing sound. + self.share_qkv_input_quant = share_qkv_input_quant + if share_qkv_input_quant and self.qkv_mode != QKVMode.SEPARATE_QKV: + raise ValueError( + "Attention: share_qkv_input_quant requires " + f"qkv_mode=QKVMode.SEPARATE_QKV ('{QKVMode.SEPARATE_QKV.value}'), " + f"got QKVMode.{self.qkv_mode.name} ('{self.qkv_mode.value}'); a " + "fused QKV projection already quantizes its input once." + ) + if share_qkv_input_quant and self.force_dynamic_quantization: + # A dynamic scale is derived per Linear per call, so there is no + # shared calibrated scale to quantize against. + raise ValueError( + "Attention: share_qkv_input_quant is incompatible with " + "force_dynamic_quantization; sharing requires a static " + "calibrated input_scale." + ) + attention_metadata_state = getattr(config, "attention_metadata_state", None) if self.qk_norm: @@ -398,11 +420,83 @@ def get_qkv( kv_source = ( encoder_hidden_states if encoder_hidden_states is not None else hidden_states ) + if self._can_share_qkv_quantize(hidden_states, encoder_hidden_states): + hidden_states = self._static_quantize_fp8(hidden_states, self.to_q.input_scale) + kv_source = hidden_states q = self.to_q(hidden_states) k = self.to_k(kv_source) v = self.to_v(kv_source) return q, k, v + def _shares_qkv_input_quant(self) -> bool: + """Config-only half of the condition, so post_load_weights() can reuse it.""" + if not self.share_qkv_input_quant: + return False + projections = (self.to_q, self.to_k, self.to_v) + return all( + p.has_fp8_qdq + and p.input_scale is not None + and not p.force_dynamic_quantization + # A quantization method may run a given call in higher precision + # than its checkpoint recipe -- Cosmos3 does this on the outer + # denoising steps -- by publishing ``high_precision``. Quantizing + # the shared activation here would hand such a call a tensor it + # must not receive, so leave it in its input dtype. + and not getattr(p.quant_method, "high_precision", False) + for p in projections + ) + + def _can_share_qkv_quantize(self, hidden_states, encoder_hidden_states) -> bool: + """Whether q/k/v can consume one quantized activation. + + Reads no tensor values: that would sync the device every forward and make + the graph data-dependent. Scale equality is checked at load instead. + Cross-attention feeds k/v from a different tensor, so only self-attention + has a single activation to share. + """ + return ( + encoder_hidden_states is None + and self._shares_qkv_input_quant() + and not isinstance(hidden_states, Fp4QuantizedTensor) + and hidden_states.dtype != torch.float8_e4m3fn + ) + + @staticmethod + def _static_quantize_fp8(x: torch.Tensor, input_scale: torch.Tensor) -> torch.Tensor: + """Quantize once so each projection can skip its own quantize pass. + + FP8QDQLinearMethod.apply passes a pre-quantized input straight through. + The reshapes are views on contiguous activations, not copies. + """ + shape = x.shape + x2d = x.reshape(-1, shape[-1]) if x.dim() > 2 else x + qx, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x2d, input_scale) + return qx.reshape(shape) if x.dim() > 2 else qx + + def post_load_weights(self) -> None: + """Check the shared-activation invariant here, never in forward(). + + Each Linear applies its *own* input_scale in the GEMM epilogue, so + quantizing with to_q's scale is only correct when all three agree. + ModelOpt's self-attention calibration makes that an invariant (q/k/v see + the same input distribution), but a checkpoint that violates it would + otherwise produce silently wrong output. + """ + if not self._shares_qkv_input_quant(): + return + scales = {name: getattr(self, name).input_scale for name in ("to_q", "to_k", "to_v")} + mismatched = { + name: scale.item() + for name, scale in scales.items() + if not torch.equal(scale, scales["to_q"]) + } + if mismatched: + raise ValueError( + "q/k/v share one quantized activation, so they must carry the " + f"same calibrated input_scale; got to_q={scales['to_q'].item()} " + f"and mismatched {mismatched} at layer_idx={self.layer_idx}" + ) + def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: if self.qk_norm: q = self.norm_q(q) diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index 8c1c33a59adc..ed4c7e4c9e03 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/loader.py @@ -30,6 +30,11 @@ QuantAlgo.NVFP4: ("weight_scale", "weight_scale_2"), } +# Weight dtypes a module holds when it was *not* built for a quantized recipe. +# Quantized recipes allocate a narrow buffer instead (float8_e4m3fn for FP8, +# packed uint8 for NVFP4/AWQ). +_HIGH_PRECISION_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + class DynamicLinearWeightLoader: """ @@ -148,7 +153,11 @@ def _get_quant_algo_for_layer(self, name: str) -> Optional[QuantAlgo]: return None def _check_static_quant_scales( - self, weight_dict: Dict[str, torch.Tensor], quant_algo: Optional[QuantAlgo], name: str + self, + weight_dict: Dict[str, torch.Tensor], + quant_algo: Optional[QuantAlgo], + name: str, + module: Linear | None = None, ) -> None: """Refuse static quant recipes against checkpoints without scales. @@ -172,6 +181,21 @@ def _check_static_quant_scales( if self.quant_config.is_module_excluded_from_quantization(name): return + # Ask the destination buffer, not the module's name. ``quant_algo`` + # above falls back to the *global* recipe for any module that carries + # no ``quant_config`` of its own, which includes modules that cannot be + # quantized at all: ``Embedding`` reaches this loader because it + # subclasses ``LMHead`` -> ``Linear``, yet its ``__init__`` never + # exposes ``quant_config``, so it always keeps a high-precision buffer. + # A high-precision weight landing in a high-precision buffer is a + # correct load with nothing to corrupt, which is the condition this + # guard is about. Where the destination is unknown the check proceeds, + # keeping the fail-closed behaviour. + if module is not None: + destination = getattr(module, "weight", None) + if destination is not None and destination.dtype in _HIGH_PRECISION_DTYPES: + return + weight = weight_dict.get("weight") if weight is None or weight.dtype not in (torch.bfloat16, torch.float16, torch.float32): return @@ -337,7 +361,7 @@ def load_linear_weights( # Static (pre-quantized) recipes must not be loaded from checkpoints # that do not carry the expected scale tensors. for weight_dict in weight_dicts: - self._check_static_quant_scales(weight_dict, quant_algo, name) + self._check_static_quant_scales(weight_dict, quant_algo, name, module) # Special handling for fused NVFP4 dynamic quantization # Fused weights (Q,K,V or gate,up) must be quantized TOGETHER diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py index 290028266292..e36501ab6982 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py @@ -78,7 +78,6 @@ # golden/visual_gen_lpips/cosmos3_i2v_4step_lpips_golden_video.json. COSMOS3_I2V_4STEP_LPIPS_THRESHOLD = 0.10 - COSMOS3_FEATURE_LPIPS_THRESHOLD = 0.05 COSMOS3_QUANTIZATION_IGNORE = [ "language_model.*", @@ -130,13 +129,33 @@ def _build_cosmos3_accuracy_cases(): COSMOS3_ACCURACY_CASES = _build_cosmos3_accuracy_cases() -def _run_cosmos3_lpips_pipeline(num_frames, video=None): - """Run the Cosmos3-Nano pipeline (default setting, VANILLA attn, compile-off). +def _run_cosmos3_lpips_pipeline( + num_frames, + video=None, + image=None, + model_subpath=(COSMOS3_NANO_MODEL_SUBPATH,), + label="Cosmos3-Nano checkpoint", + height=COSMOS3_LPIPS_HEIGHT, + width=COSMOS3_LPIPS_WIDTH, + num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, + output_type="video", +): + """Run a Cosmos3 pipeline (default setting, VANILLA attn, compile-off). Returns the generated video tensor ``(B, T, H, W, C)`` (T == ``num_frames``), - or ``None`` if generation produced no video. ``num_frames=1`` yields the - single-frame text-to-image path; passing ``video`` (encoded MP4 bytes, - decoded on the worker's NVDEC) yields the video-to-video path. + or ``None`` if generation produced no video. Passing ``video`` (encoded MP4 + bytes, decoded on the worker's NVDEC) yields the video-to-video path; + passing ``image`` yields the image-to-video path. + + ``output_type="image"`` selects the real text-to-image path and returns + ``(B, H, W, C)`` instead. It is not the same as ``num_frames=1``: the + pipeline keys T2I off ``output_type``, which also swaps in + ``COSMOS3_T2I_PARAMS``, the T2I system prompt and the image resolution + template, so a one-frame video run exercises none of that. + + ``model_subpath`` selects the checkpoint under ``LLM_MODELS_ROOT`` and is a + tuple of path components, so nested checkpoints (the FP8 builds ship inside + a dated subdirectory) address the same way as top-level ones. """ # Cosmos3 re-reads the guardrail flag in __init__; set it before the pipeline loads. guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" @@ -151,8 +170,8 @@ def _run_cosmos3_lpips_pipeline(num_frames, video=None): VisualGenArgs, ) - model_path = _lpips_model_path(COSMOS3_NANO_MODEL_SUBPATH) - _skip_if_missing(model_path, "Cosmos3-Nano checkpoint", is_dir=True) + model_path = _lpips_model_path(*model_subpath) + _skip_if_missing(model_path, label, is_dir=True) _disable_inductor_compile_worker_quiesce() args = VisualGenArgs( model=model_path, @@ -171,18 +190,24 @@ def _run_cosmos3_lpips_pipeline(num_frames, video=None): # so pin it rather than inheriting the video-mode default. negative_prompt="", seed=COSMOS3_LPIPS_SEED, - height=COSMOS3_LPIPS_HEIGHT, - width=COSMOS3_LPIPS_WIDTH, + height=height, + width=width, num_frames=num_frames, - num_inference_steps=COSMOS3_LPIPS_NUM_INFERENCE_STEPS, + num_inference_steps=num_inference_steps, guidance_scale=COSMOS3_LPIPS_GUIDANCE_SCALE, frame_rate=COSMOS3_LPIPS_FRAME_RATE, use_guardrails=False, video=video, + image=image, + output_type=output_type, ) - if result is None or result.video is None: + if result is None: return None - return result.video.detach().cpu() + # T2I returns image (B, H, W, C) and leaves video unset. + produced = result.image if output_type == "image" else result.video + if produced is None: + return None + return produced.detach().cpu() finally: del pipeline _cleanup_cuda() diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6f0b5c320cc9..41073894481f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -125,6 +125,8 @@ l0_b200: - unittest/_torch/modules/test_triton_linear.py - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/test_rotary_embedding.py + - unittest/_torch/modules/test_swiglu_2in.py + - unittest/_torch/modules/test_gated_mlp_split.py - unittest/_torch/modules/test_low_m_gemm.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules @@ -267,8 +269,11 @@ l0_b200: - unittest/_torch/visual_gen/test_wan22_t2v_teacache.py - unittest/_torch/visual_gen/test_wan_transformer.py - unittest/_torch/visual_gen/test_cosmos3_transformer.py + - unittest/_torch/visual_gen/test_attention_qkv_share.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py - unittest/_torch/visual_gen/test_cosmos3_distilled.py + - unittest/_torch/visual_gen/test_cosmos3_fp8.py + - unittest/_torch/visual_gen/test_cosmos3_step_precision_component.py - unittest/_torch/visual_gen/test_cosmos3_transfer.py - unittest/_torch/visual_gen/test_control_kernels.py - unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c605b8e36d6e..ca172a5c6669 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -48,6 +48,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_flux_infer.py - unittest/_torch/visual_gen/test_ltx2_pipeline.py - unittest/_torch/visual_gen/test_ltx2_transformer.py + - unittest/_torch/visual_gen/test_cosmos3_step_precision.py - unittest/_torch/visual_gen/test_quant_static_guard.py - unittest/_torch/visual_gen/test_teacache.py - unittest/_torch/visual_gen/test_tensor_payload.py diff --git a/tests/unittest/_torch/modules/test_gated_mlp_split.py b/tests/unittest/_torch/modules/test_gated_mlp_split.py new file mode 100644 index 000000000000..93f82e8944d0 --- /dev/null +++ b/tests/unittest/_torch/modules/test_gated_mlp_split.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for GatedMLP's opt-in split gate/up topology. + +Statically quantized checkpoints calibrate a separate weight scale per +projection. Fusing gate and up into one Linear forces a single scale on the pair +and re-quantizes the other onto it, discarding calibration. `split_gate_up=True` +keeps them separate so each loads its checkpoint tensor and scale unchanged. + +The flag is opt-in and every existing caller relies on the fused default, so +these tests pin the default's topology and parameter names as much as they pin +the new path. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HIDDEN, INTERMEDIATE, TOKENS = 512, 1024, 256 +INPUT_SCALE, WEIGHT_SCALE = 1e-2, 1.1e-3 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _make(split, *, quant_algo=QuantAlgo.FP8, activation=F.silu, force_dynamic=False): + quant_config = QuantConfig(quant_algo=quant_algo) if quant_algo else QuantConfig() + config = ModelConfig(quant_config=quant_config) + config.force_dynamic_quantization = force_dynamic + return GatedMLP( + hidden_size=HIDDEN, + intermediate_size=INTERMEDIATE, + bias=False, + activation=activation, + dtype=torch.bfloat16, + config=config, + split_gate_up=split, + ).cuda() + + +def _capture(fn): + """Capture fn() after warming up on a side stream. + + Capture is sensitive to allocator state left by earlier tests, so the + documented warmup protocol is used rather than a single default-stream call. + """ + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + fn() + torch.cuda.current_stream().wait_stream(warmup) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = fn() + graph.replay() + torch.cuda.synchronize() + return out + + +def _fp8(rows, cols): + return (torch.randn(rows, cols, device="cuda", dtype=torch.bfloat16) * 0.05).to( + torch.float8_e4m3fn + ) + + +def _set_scales(linear, weight_scale=WEIGHT_SCALE, input_scale=INPUT_SCALE): + linear.weight_scale.data.fill_(weight_scale) + linear.input_scale.data.fill_(input_scale) + linear.inv_input_scale.data.fill_(1.0 / input_scale) + + +def _set_weights(linear): + """Give the projection finite weights. + + Linear allocates FP8 weights uninitialized, so a test that skips this reads + whatever the caching allocator hands back -- NaN often enough to matter. Any + torch.equal assertion then fails regardless of the behaviour under test, + because torch.equal is False for NaN even when both sides are bit-identical. + """ + linear.weight.data.copy_(_fp8(*linear.weight.shape)) + + +@requires_cuda +def test_default_topology_is_unchanged(): + """The default must stay byte-for-byte what existing callers already build.""" + mlp = _make(False) + + assert mlp.split_gate_up is False + assert isinstance(mlp.gate_up_proj, torch.nn.Module) + assert mlp.gate_up_proj.out_features == INTERMEDIATE * 2 + assert not hasattr(mlp, "gate_proj") + assert not hasattr(mlp, "up_proj") + + names = set(mlp.state_dict()) + assert any(n.startswith("gate_up_proj.") for n in names) + assert not any(n.startswith("gate_proj.") or n.startswith("up_proj.") for n in names) + + +@requires_cuda +def test_split_topology_builds_separate_projections(): + mlp = _make(True) + + assert mlp.gate_up_proj is None + assert mlp.gate_proj.out_features == INTERMEDIATE + assert mlp.up_proj.out_features == INTERMEDIATE + + names = set(mlp.state_dict()) + assert any(n.startswith("gate_proj.") for n in names) + assert any(n.startswith("up_proj.") for n in names) + assert not any(n.startswith("gate_up_proj.") for n in names) + + +@requires_cuda +def test_split_projections_load_vanilla_weights_directly(): + """Each split Linear owns one checkpoint tensor -- no fused shard mapping.""" + from tensorrt_llm._torch.modules.linear import WeightMode + + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj): + assert linear.weights_loading_config.weight_mode == WeightMode.VANILLA + assert getattr(linear, "fused_weight_shard_indices_mapping", None) is None + + weight = _fp8(INTERMEDIATE, HIDDEN) + mlp.gate_proj.load_weights( + [ + { + "weight": weight, + "weight_scale": torch.tensor(WEIGHT_SCALE), + "input_scale": torch.tensor(INPUT_SCALE), + } + ] + ) + # Loaded exactly: no requantization onto a shared scale. + assert torch.equal(mlp.gate_proj.weight.data, weight) + assert mlp.gate_proj.weight_scale.item() == pytest.approx(WEIGHT_SCALE) + + +@requires_cuda +def test_split_matches_fused_when_scales_agree(): + """With one scale for both, fusion requantizes nothing, so results must match. + + This isolates the topology change from the scale effect: any difference here + is a bug in the split path rather than the calibration it preserves. + """ + gate_w, up_w, down_w = ( + _fp8(INTERMEDIATE, HIDDEN), + _fp8(INTERMEDIATE, HIDDEN), + _fp8(HIDDEN, INTERMEDIATE), + ) + + fused, split = _make(False), _make(True) + fused.gate_up_proj.weight.data.copy_(torch.cat([gate_w, up_w], dim=0)) + _set_scales(fused.gate_up_proj) + split.gate_proj.weight.data.copy_(gate_w) + split.up_proj.weight.data.copy_(up_w) + _set_scales(split.gate_proj) + _set_scales(split.up_proj) + for mlp in (fused, split): + mlp.down_proj.weight.data.copy_(down_w) + _set_scales(mlp.down_proj) + split.post_load_weights() + + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + assert torch.equal(fused(x), split(x)) + + +@requires_cuda +def test_shared_input_is_quantized_once(): + """Both projections must receive the *same* already-quantized activation. + + Two Linears would otherwise each quantize the same values independently -- + which would also yield FP8 at both inputs, so dtype alone proves nothing. + Storage identity is the discriminating check. + """ + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + seen = {} + + def record(name): + def hook(_module, inputs): + seen[name] = (inputs[0].dtype, inputs[0].data_ptr()) + + return hook + + mlp.gate_proj.register_forward_pre_hook(record("gate")) + mlp.up_proj.register_forward_pre_hook(record("up")) + mlp(torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02) + + assert seen["gate"][0] == torch.float8_e4m3fn + assert seen["up"][0] == torch.float8_e4m3fn + # Same storage, not merely the same dtype: two independent quantizations of + # the same values would also both be FP8, so identity is what proves the + # activation was quantized once and shared. + assert seen["gate"][1] == seen["up"][1] + + +@requires_cuda +def test_fp8_output_feeds_down_proj(): + """SwiGLU emits FP8 for an FP8 down_proj, so down_proj skips its own quant.""" + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + seen = {} + mlp.down_proj.register_forward_pre_hook( + lambda _m, inputs: seen.__setitem__("dtype", inputs[0].dtype) + ) + mlp(torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02) + + assert seen["dtype"] == torch.float8_e4m3fn + + +@requires_cuda +def test_post_load_weights_rejects_mismatched_input_scales(): + """The shared quantization is only valid if both scales agree. + + Checked once here rather than in forward: reading the tensors on the hot path + would synchronize the device and make the graph data-dependent, which breaks + fullgraph compilation. + """ + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + mlp.up_proj.input_scale.data.fill_(INPUT_SCALE * 2) + with pytest.raises(ValueError, match="same calibrated input_scale"): + mlp.post_load_weights() + + +@requires_cuda +def test_post_load_weights_is_noop_for_fused(): + _make(False).post_load_weights() + + +@requires_cuda +def test_post_load_weights_ignores_scales_for_bf16(): + mlp = _make(True, quant_algo=None) + mlp.post_load_weights() # must not raise + + +@requires_cuda +def test_torch_compile_fullgraph(): + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + eager = mlp(x) + compiled = torch.compile(lambda t: mlp(t), fullgraph=True)(x) + + assert compiled.dtype == eager.dtype + assert torch.equal(compiled, eager) + + +@requires_cuda +def test_cuda_graph_capture(): + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + + captured = _capture(lambda: mlp(x)) + assert torch.equal(captured, mlp(x)) + + +@requires_cuda +def test_lora_is_rejected_on_split_path(): + """forward_lora fuses gate/up LoRA into a projection the split path lacks. + + Failing here beats dereferencing gate_up_proj=None with a bare TypeError. + """ + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + + with pytest.raises(NotImplementedError, match="LoRA is not supported"): + mlp(x, lora_params={"any": "value"}) + + +@requires_cuda +def test_non_swiglu_activation_is_rejected_on_split_path(): + mlp = _make(True, activation=F.gelu) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + + with pytest.raises(NotImplementedError, match="requires SwiGLU"): + mlp(x) + + +@requires_cuda +def test_split_rejects_forced_dynamic_quantization(): + """Dynamic quantization must use the fused topology. + + The activation emits FP8 with down_proj's calibrated scale, which would make + down_proj skip the dynamic quantization it was configured for. Rejected at + construction rather than silently downgraded. + """ + with pytest.raises(ValueError, match="force_dynamic_quantization"): + _make(True, force_dynamic=True) + + +@requires_cuda +def test_fused_topology_still_allows_forced_dynamic(): + """The rejection must be scoped to the split path only.""" + mlp = _make(False, force_dynamic=True) + assert mlp.gate_up_proj is not None + + +@requires_cuda +def test_bf16_split_runs_without_quantization(): + """split_gate_up must not assume a quantized checkpoint.""" + mlp = _make(True, quant_algo=None) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + + assert mlp._can_share_gate_up_quantization(x) is False + out = mlp(x) + assert out.shape == (TOKENS, HIDDEN) + assert out.dtype == torch.bfloat16 + + +@requires_cuda +@pytest.mark.parametrize( + "shape", [(2, 8, HIDDEN), (4, TOKENS, HIDDEN)], ids=["small_rank3", "model_rank3"] +) +def test_rank3_activations(shape): + """Models carry [batch, seq, hidden]; rank-2 tests alone would miss this. + + The op previously asserted rank 2 while its fake accepted rank 3, so tracing + succeeded and execution failed. + """ + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + x = torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.02 + out = mlp(x) + assert out.shape == (*shape[:-1], HIDDEN) + + # Same values through a rank-2 view must give the same result. + flat = mlp(x.reshape(-1, shape[-1])) + assert torch.equal(out.reshape(-1, HIDDEN), flat) + + +@requires_cuda +def test_rank3_compile_and_cuda_graph(): + mlp = _make(True) + for linear in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): + _set_weights(linear) + _set_scales(linear) + mlp.post_load_weights() + + x = torch.randn(2, 8, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + eager = mlp(x) + assert torch.equal(torch.compile(lambda t: mlp(t), fullgraph=True)(x), eager) + + captured = _capture(lambda: mlp(x)) + assert torch.equal(captured, mlp(x)) diff --git a/tests/unittest/_torch/modules/test_swiglu_2in.py b/tests/unittest/_torch/modules/test_swiglu_2in.py new file mode 100644 index 000000000000..63a2e3433214 --- /dev/null +++ b/tests/unittest/_torch/modules/test_swiglu_2in.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the two-input SwiGLU used by split gate/up projections. + +`silu_and_mul` requires gate and up adjacent in one tensor. Models that keep +gate and up as separate projections -- static-FP8 Cosmos3 does, to preserve each +projection's calibrated scale -- have no such tensor, and concatenating one is a +large GPU copy. `silu_and_mul_2in` consumes the two tensors directly. + +Every case asserts *bit-exact* equality with the fused op rather than a +tolerance: the two kernels perform the same arithmetic in the same order and +accumulate in fp32, so any difference is a defect rather than drift. +""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.swiglu import get_silu_b200_tuning_params, swiglu, swiglu_2in + +# (M, intermediate). The large case is the Cosmos3 Nano default T2V request: +# 720x1280 x 189 frames with CFG gives M = 88320. +SHAPES = [(8, 4), (13, 257), (1760, 12288), (88320, 12288)] +DTYPES = [torch.bfloat16, torch.float16] + + +def _reference(gate, up, **kwargs): + """The fused path, fed an explicitly concatenated tensor. + + silu_and_mul is rank-2 only, so higher-rank inputs are flattened for the + comparison and the result restored to the original shape. + """ + flat_gate = gate.reshape(-1, gate.shape[-1]) + flat_up = up.reshape(-1, up.shape[-1]) + out = swiglu(torch.cat([flat_gate, flat_up], dim=-1), **kwargs) + return out.reshape(gate.shape) + + +def _capture(fn): + """Capture fn() after warming up on a side stream. + + Capture is sensitive to allocator state left by earlier tests, so the + documented warmup protocol is used rather than a single default-stream call. + """ + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + fn() + torch.cuda.current_stream().wait_stream(warmup) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = fn() + graph.replay() + torch.cuda.synchronize() + return out + + +def _pair(m, n, dtype, seed=0): + torch.manual_seed(seed) + return ( + torch.randn(m, n, device="cuda", dtype=dtype), + torch.randn(m, n, device="cuda", dtype=dtype), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("m, n", SHAPES, ids=lambda v: str(v)) +@pytest.mark.parametrize("dtype", DTYPES, ids=["bf16", "fp16"]) +def test_matches_fused_silu_and_mul(m, n, dtype): + gate, up = _pair(m, n, dtype) + expected = _reference(gate, up) + actual = swiglu_2in(gate, up) + + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("swiglu_limit", [1.0, 7.0]) +def test_matches_fused_with_swiglu_limit(swiglu_limit): + """The limit clamps gate and up differently, so it must be plumbed through.""" + gate, up = _pair(256, 512, torch.bfloat16) + # Scale up so values actually exceed the limit and the clamp is exercised. + gate, up = gate * 10.0, up * 10.0 + + expected = _reference(gate, up, swiglu_limit=swiglu_limit) + actual = swiglu_2in(gate, up, swiglu_limit=swiglu_limit) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize( + "swiglu_alpha, swiglu_beta", + [(1.702, 1.0), (1.702, 0.0), (1.0, 1.0)], + ids=["swigluoai", "alpha_only", "beta_only"], +) +def test_matches_fused_with_alpha_beta(swiglu_alpha, swiglu_beta): + """alpha gains inside the sigmoid, beta offsets up; both must be plumbed. + + Defaulting them drops the kernel to the alpha=1, beta=0 special case, which + is numerically wrong rather than merely unsupported -- and silently so. + """ + gate, up = _pair(256, 512, torch.bfloat16) + + expected = _reference(gate, up, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta) + actual = swiglu_2in(gate, up, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_alpha_beta_defaults_match_plain_swiglu(): + """Omitting alpha/beta must stay bit-identical to plain silu_and_mul.""" + gate, up = _pair(256, 512, torch.bfloat16) + + assert torch.equal( + swiglu_2in(gate, up), swiglu_2in(gate, up, swiglu_alpha=1.0, swiglu_beta=0.0) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("m, n", [(256, 512), (1760, 12288)], ids=lambda v: str(v)) +def test_fp8_output_matches_fused(m, n): + """With an FP8 down_proj, GatedMLP asks SwiGLU to emit FP8 directly. + + The activation then carries down_proj's input quantization, so the two-input + form has to reproduce it exactly rather than returning BF16. + """ + gate, up = _pair(m, n, torch.bfloat16) + scale = torch.tensor(1e-2, device="cuda", dtype=torch.float32) + kwargs = dict(quant_scale=scale, quant_type=torch.float8_e4m3fn) + + expected = _reference(gate, up, **kwargs) + actual = swiglu_2in(gate, up, **kwargs) + + assert actual.dtype == torch.float8_e4m3fn + assert torch.equal(actual.float(), expected.float()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("quantized", [False, True], ids=["bf16_out", "fp8_out"]) +def test_torch_compile_fullgraph(quantized): + """Exercises register_fake: a wrong meta impl breaks tracing, not eager. + + The Cosmos3 E2E helper runs with compilation disabled, so nothing else in + this feature's test surface would catch it. + """ + gate, up = _pair(256, 512, torch.bfloat16) + scale = torch.tensor(1e-2, device="cuda", dtype=torch.float32) + + def fn(g, u): + if quantized: + return swiglu_2in(g, u, quant_scale=scale, quant_type=torch.float8_e4m3fn) + return swiglu_2in(g, u) + + eager = fn(gate, up) + compiled = torch.compile(fn, fullgraph=True)(gate, up) + + assert compiled.shape == eager.shape + assert compiled.dtype == eager.dtype + assert torch.equal(compiled.float(), eager.float()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_cuda_graph_capture(): + gate, up = _pair(256, 512, torch.bfloat16) + captured = _capture(lambda: swiglu_2in(gate, up)) + assert torch.equal(captured, swiglu_2in(gate, up)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_rejects_mismatched_inputs(): + gate = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + with pytest.raises(Exception, match="same shape"): + swiglu_2in(gate, torch.randn(8, 32, device="cuda", dtype=torch.bfloat16)) + with pytest.raises(Exception, match="same dtype"): + swiglu_2in(gate, torch.randn(8, 16, device="cuda", dtype=torch.float16)) + with pytest.raises(Exception, match="same device"): + swiglu_2in(gate, torch.randn(8, 16, dtype=torch.bfloat16)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_rejects_non_contiguous(): + """The kernel indexes flat runs, so a strided view must be rejected. + + Accepting one produced silently wrong output rather than an error: a (32, 2) + view with inner stride 2 differed from the reference by 2.7 in BF16 and 272 + with FP8 output. + """ + gate = torch.randn(32, 4, device="cuda", dtype=torch.bfloat16)[:, ::2] + up = torch.randn(32, 4, device="cuda", dtype=torch.bfloat16)[:, ::2] + assert not gate.is_contiguous() + + with pytest.raises(Exception, match="contiguous"): + swiglu_2in(gate, up) + # Contiguous in one operand only is still rejected. + with pytest.raises(Exception, match="contiguous"): + swiglu_2in(gate.contiguous(), up) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("quantized", [False, True], ids=["bf16_out", "fp8_out"]) +def test_opcheck(quantized): + """Full torch.library contract: schema, fake tensor, autograd registration.""" + gate, up = _pair(64, 128, torch.bfloat16) + kwargs = {} + if quantized: + kwargs = dict( + scale=torch.tensor(1e-2, device="cuda", dtype=torch.float32), dtype=torch.float8_e4m3fn + ) + torch.library.opcheck(torch.ops.trtllm.silu_and_mul_2in, (gate, up), kwargs) + + +def test_tuning_params_are_valid_launch_configs(): + """Guards the tuning table against edits that would not launch. + + Triton needs a power-of-two block, and threads (num_warps * 32) must stay + within the 1024-per-block limit -- easy to violate when hand-editing tuned + values. + """ + for out_dtype in (torch.bfloat16, torch.float16, torch.float8_e4m3fn): + block_elements, num_warps = get_silu_b200_tuning_params(out_dtype) + assert block_elements & (block_elements - 1) == 0, ( + f"{out_dtype}: block_elements {block_elements} is not a power of two" + ) + assert 1 <= num_warps <= 32 + assert num_warps * 32 <= 1024, ( + f"{out_dtype}: {num_warps} warps exceeds the per-block thread limit" + ) + assert block_elements % (num_warps * 32) == 0, ( + f"{out_dtype}: {block_elements} elements do not divide evenly across " + f"{num_warps * 32} threads" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize( + "shape", [(2, 8, 64), (4, 256, 512), (2, 3, 4, 16)], ids=["rank3_small", "rank3_model", "rank4"] +) +@pytest.mark.parametrize("quantized", [False, True], ids=["bf16_out", "fp8_out"]) +def test_rank_n_inputs(shape, quantized): + """The kernel walks a flat run, so any matching contiguous shape is valid. + + Model activations are rank-3 [batch, seq, hidden]. An earlier revision + asserted rank 2 in the op while its fake accepted rank 3, so compiled + execution traced cleanly and then failed. + """ + torch.manual_seed(0) + gate = torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + up = torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + kwargs = {} + if quantized: + kwargs = dict( + quant_scale=torch.tensor(1e-2, device="cuda", dtype=torch.float32), + quant_type=torch.float8_e4m3fn, + ) + + actual = swiglu_2in(gate, up, **kwargs) + assert actual.shape == gate.shape + assert torch.equal(actual.float(), _reference(gate, up, **kwargs).float()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_rank3_compile_and_cuda_graph(): + """Guards the fake/eager rank agreement that previously diverged.""" + gate, up = _pair(2 * 8, 64, torch.bfloat16) + gate, up = gate.reshape(2, 8, 64), up.reshape(2, 8, 64) + + eager = swiglu_2in(gate, up) + compiled = torch.compile(lambda g, u: swiglu_2in(g, u), fullgraph=True)(gate, up) + assert compiled.shape == eager.shape + assert torch.equal(compiled, eager) + + captured = _capture(lambda: swiglu_2in(gate, up)) + assert torch.equal(captured, eager) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_opcheck_rank3(): + gate, up = _pair(2 * 8, 64, torch.bfloat16) + torch.library.opcheck( + torch.ops.trtllm.silu_and_mul_2in, (gate.reshape(2, 8, 64), up.reshape(2, 8, 64)) + ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py b/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py new file mode 100644 index 000000000000..23810faeacbb --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_qkv_share.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for Attention's opt-in shared QKV input quantization. + +Static FP8 checkpoints keep q/k/v as separate calibrated projections, so a fused +QKV Linear would re-quantize two of them onto a third's scale. Splitting them +costs three quantizations of one identical activation instead of one, which +``share_qkv_input_quant=True`` recovers by quantizing once and handing the same +FP8 tensor to all three. + +The flag is opt-in: enabling it also commits the caller to running +``post_load_weights()``, which is where the equal-input_scale invariant that +makes the sharing sound is actually checked. These tests pin both halves. +""" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HIDDEN, HEADS, KV_HEADS, HEAD_DIM, TOKENS = 512, 8, 8, 64, 128 +INPUT_SCALE, WEIGHT_SCALE = 1e-2, 1.1e-3 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _make(*, share, qkv_mode=QKVMode.SEPARATE_QKV, quant_algo=QuantAlgo.FP8, force_dynamic=False): + config = DiffusionModelConfig( + quant_config=QuantConfig(quant_algo=quant_algo) if quant_algo else QuantConfig(), + force_dynamic_quantization=force_dynamic, + ) + return Attention( + hidden_size=HIDDEN, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + head_dim=HEAD_DIM, + qkv_mode=qkv_mode, + qk_norm=False, + bias=False, + config=config, + enable_sequence_parallel=False, + share_qkv_input_quant=share, + ).cuda() + + +def _set_scales(linear, input_scale=INPUT_SCALE): + linear.weight_scale.data.fill_(WEIGHT_SCALE) + linear.input_scale.data.fill_(input_scale) + linear.inv_input_scale.data.fill_(1.0 / input_scale) + + +def _calibrate(attn, input_scale=INPUT_SCALE): + for name in ("to_q", "to_k", "to_v"): + _set_scales(getattr(attn, name), input_scale) + + +@requires_cuda +def test_default_does_not_share(): + """Every existing caller relies on the unshared default.""" + attn = _make(share=False) + assert attn.share_qkv_input_quant is False + _calibrate(attn) + assert attn._shares_qkv_input_quant() is False + + +@requires_cuda +def test_fused_qkv_rejects_sharing(): + """A fused QKV projection already quantizes its input exactly once.""" + with pytest.raises(ValueError, match="SEPARATE_QKV"): + _make(share=True, qkv_mode=QKVMode.FUSE_QKV) + + +@requires_cuda +def test_forced_dynamic_rejects_sharing(): + with pytest.raises(ValueError, match="force_dynamic_quantization"): + _make(share=True, force_dynamic=True) + + +@requires_cuda +def test_unquantized_attention_does_not_share(): + """Without FP8 weights there is no static scale to quantize against.""" + attn = _make(share=True, quant_algo=None) + assert attn._shares_qkv_input_quant() is False + + +@requires_cuda +def test_shared_input_is_quantized_once(): + """All three projections must receive the *same* already-quantized tensor. + + Three independent quantizations of one activation would also yield FP8 at + every input, so dtype alone proves nothing -- storage identity is what + discriminates. + """ + attn = _make(share=True) + _calibrate(attn) + attn.post_load_weights() + + seen = {} + + def record(name): + def hook(_module, inputs): + seen[name] = (inputs[0].dtype, inputs[0].data_ptr()) + + return hook + + for name in ("to_q", "to_k", "to_v"): + getattr(attn, name).register_forward_pre_hook(record(name)) + + x = torch.randn(1, TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + attn.get_qkv(x) + + assert {dtype for dtype, _ in seen.values()} == {torch.float8_e4m3fn} + assert len({ptr for _, ptr in seen.values()}) == 1, ( + f"q/k/v should share one quantized activation, got {seen}" + ) + + +@requires_cuda +def test_sharing_matches_unshared_result(): + """Sharing must be a pure launch-count optimization, not a numerical change. + + Each Linear applies its own input_scale in the epilogue, and all three agree + here, so quantizing once must be bit-identical to quantizing three times. + """ + shared, unshared = _make(share=True), _make(share=False) + torch.manual_seed(0) + for name in ("to_q", "to_k", "to_v"): + out_features = getattr(shared, name).out_features + weight = (torch.randn(out_features, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.05).to( + torch.float8_e4m3fn + ) + for attn in (shared, unshared): + getattr(attn, name).weight.data.copy_(weight) + _set_scales(getattr(attn, name)) + shared.post_load_weights() + + x = torch.randn(1, TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + for actual, expected in zip(shared.get_qkv(x), unshared.get_qkv(x)): + assert torch.equal(actual, expected) + + +@requires_cuda +def test_cross_attention_source_is_not_shared(): + """k/v read a different tensor than q, so there is no single activation.""" + attn = _make(share=True) + _calibrate(attn) + attn.post_load_weights() + + x = torch.randn(1, TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + encoder = torch.randn(1, TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) * 0.02 + assert attn._can_share_qkv_quantize(x, encoder) is False + + seen = {} + attn.to_q.register_forward_pre_hook(lambda _m, inputs: seen.update(q=inputs[0].dtype)) + attn.get_qkv(x, encoder_hidden_states=encoder) + assert seen["q"] == torch.bfloat16 + + +@requires_cuda +def test_post_load_weights_rejects_mismatched_input_scales(): + """A checkpoint violating the shared-scale invariant must fail loudly. + + to_k's GEMM would apply its own scale to activations quantized with to_q's, + which silently corrupts the projection rather than erroring. + """ + attn = _make(share=True) + _calibrate(attn) + _set_scales(attn.to_k, input_scale=INPUT_SCALE * 2) + + with pytest.raises(ValueError, match="same calibrated input_scale"): + attn.post_load_weights() + + +@requires_cuda +def test_post_load_weights_is_a_noop_when_not_sharing(): + attn = _make(share=False) + _calibrate(attn) + _set_scales(attn.to_k, input_scale=INPUT_SCALE * 2) + attn.post_load_weights() + + +@requires_cuda +def test_prequantized_input_is_passed_through(): + """An upstream fused norm+quant already produced FP8; do not re-quantize.""" + attn = _make(share=True) + _calibrate(attn) + attn.post_load_weights() + + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + assert attn._can_share_qkv_quantize(x.to(torch.float8_e4m3fn), None) is False + + +@requires_cuda +@pytest.mark.parametrize("shape", [(TOKENS, HIDDEN), (2, TOKENS, HIDDEN)], ids=["rank2", "rank3"]) +def test_rank_is_preserved(shape): + """The quantize reshape is a view; projections must still see their rank.""" + attn = _make(share=True) + _calibrate(attn) + attn.post_load_weights() + + x = torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.02 + q, k, v = attn.get_qkv(x) + assert q.shape == (*shape[:-1], HEADS * HEAD_DIM) + assert k.shape == v.shape == (*shape[:-1], KV_HEADS * HEAD_DIM) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py b/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py new file mode 100644 index 000000000000..e49af21d1105 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py @@ -0,0 +1,580 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the statically quantized (ModelOpt FP8) Cosmos3 checkpoints. + +These checkpoints ship FP8 ``E4M3`` weights alongside calibrated per-tensor +weight *and* activation scales, so inference must use static activation +quantization. That is distinct from ``TestCosmos3FP8Load`` in +``test_cosmos3_pipeline.py``, which quantizes a BF16 checkpoint dynamically at +load time from a user-supplied quant config. + +The FP8 path is expected to work through the existing static-FP8 machinery +without Cosmos3-specific quantization code; these tests pin that contract so a +regression in config resolution, scale loading, or module exclusion is caught. + +Config tests need no checkpoint or GPU. Load tests require the checkpoints: + + DIFFUSION_MODEL_PATH_COSMOS3_NANO_FP8=/path/to/cosmos3-nano-fp8-14072026 \\ + DIFFUSION_MODEL_PATH_COSMOS3_SUPER_FP8=/path/to/cosmos3-super-fp8-14072026 \\ + pytest tests/unittest/_torch/visual_gen/test_cosmos3_fp8.py -v +""" + +import gc +import os +from pathlib import Path + +os.environ["TLLM_DISABLE_MPI"] = "1" + +import pytest +import torch + +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.config import DiffusionPipelineConfig +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + TorchCompileConfig, + VisualGenArgs, +) + +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] + + +@pytest.fixture(autouse=True, scope="module") +def _cleanup_mpi_env(): + """TLLM_DISABLE_MPI has to be set before the imports above, so it cannot be + a fixture -- but leaving it set makes any later module in the same process + inherit it. Drop it on the way out, as test_cosmos3_pipeline.py does.""" + yield + os.environ.pop("TLLM_DISABLE_MPI", None) + + +# Verbatim ``quantization_config`` shape exported by ModelOpt 0.44.0 into the +# published Cosmos3 FP8 checkpoints' ``transformer/config.json``. Only the keys +# TensorRT-LLM consumes are kept; ``dynamic: false`` on both weights and +# activations is what selects the static path. +MODELOPT_FP8_QUANT_CONFIG = { + "quant_method": "modelopt", + "quant_type": "FP8_FP8", + "quant_algo": "FP8", + "weight_only": False, + "config_groups": { + "group_0": { + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "targets": ["Linear"], + } + }, + "ignore": [ + "proj_in", + "proj_out", + "time_embedder*", + "audio_proj_in", + "audio_proj_out", + "action_proj_in", + "action_proj_out", + "lm_head", + "model.visual*", + "visual*", + ], + "producer": {"name": "modelopt", "version": "0.44.0"}, +} + + +def _llm_models_root() -> str: + """Resolve the checkpoint root, without asserting it exists. + + The path constants below call this at module scope, so raising here would + error the whole module during collection -- including the config tests this + module documents as needing neither a checkpoint nor a GPU. Returning a + non-existent path instead lets the per-test ``_skip_if_missing`` guards skip + only the tests that actually load weights. + """ + root = Path("/home/scratch.trt_llm_data_ci/llm-models/") + if "LLM_MODELS_ROOT" in os.environ: + root = Path(os.environ["LLM_MODELS_ROOT"]) + if not root.exists(): + root = Path("/scratch.trt_llm_data/llm-models/") + return str(root) + + +def _checkpoint(env_var: str, *default_parts: str) -> str: + return os.environ.get(env_var) or os.path.join(_llm_models_root(), *default_parts) + + +COSMOS3_NANO_FP8_PATH = _checkpoint( + "DIFFUSION_MODEL_PATH_COSMOS3_NANO_FP8", + "Cosmos3-Nano-FP8", + "cosmos3-nano-fp8-14072026", +) +COSMOS3_SUPER_FP8_PATH = _checkpoint( + "DIFFUSION_MODEL_PATH_COSMOS3_SUPER_FP8", + "Cosmos3-Super-FP8", + "cosmos3-super-fp8-14072026", +) +COSMOS3_NANO_BF16_PATH = _checkpoint("DIFFUSION_MODEL_PATH_COSMOS3", "Cosmos3-Nano") + +# Runtime TensorRT-LLM ``Linear`` counts per tower. Static FP8 keeps every +# projection separate, so these are exactly the checkpoint's own projection +# counts -- 7 per layer per tower (q, k, v, out, gate, up, down) over 36 Nano +# and 64 Super layers. Under the fused topology GEN QKV and both towers' +# gate/up pairs collapsed, giving 216/144 (Nano) and 384/256 (Super); the +# totals below are what a 1:1 checkpoint mapping looks like. Pinning the split +# catches a tower silently dropping out of quantization *and* catches the +# topology silently reverting to fused. +EXPECTED_FP8_LINEARS = { + "nano": {"UND": 252, "GEN": 252}, + "super": {"UND": 448, "GEN": 448}, +} + +# Boundary projections the checkpoint excludes from quantization. These are +# built as native ``nn.Linear`` (never TensorRT-LLM ``Linear``), so they are +# structurally incapable of being quantized -- assert that stays true. +EXPECTED_NATIVE_LINEARS = { + "vae2llm", + "llm2vae", + "audio2llm", + "llm2audio", + "time_embedder.mlp.linear_1", + "time_embedder.mlp.linear_2", +} + + +def _skip_if_missing(path: str, label: str) -> str: + if not path or not os.path.isdir(path): + pytest.skip(f"{label} not found: {path}") + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + return path + + +def _tower_of(module_name: str) -> str: + if module_name.startswith("language_model"): + return "UND" + if module_name.startswith("gen_layers"): + return "GEN" + return "other" + + +def _load_transformer(checkpoint_path: str): + args = VisualGenArgs( + model=checkpoint_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="VANILLA"), + ) + return PipelineLoader(args).load(skip_warmup=True) + + +@pytest.fixture +def _cleanup_gpu(): + yield + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class TestStaticFp8ConfigResolution: + """The ModelOpt recipe must resolve to *static* FP8 with no extra plumbing.""" + + def test_modelopt_recipe_resolves_to_static_fp8(self): + quant_config, layer_quant_config, dynamic_weight, dynamic_activation = ( + DiffusionPipelineConfig.load_diffusion_quant_config(MODELOPT_FP8_QUANT_CONFIG) + ) + + assert quant_config.quant_algo == QuantAlgo.FP8 + # ``dynamic: false`` in the checkpoint must not decay into runtime + # quantization: the calibrated weight/activation scales would be ignored. + assert dynamic_weight is False + assert dynamic_activation is False + assert layer_quant_config is None + + def test_checkpoint_ignore_list_becomes_exclude_modules(self): + quant_config, _, _, _ = DiffusionPipelineConfig.load_diffusion_quant_config( + MODELOPT_FP8_QUANT_CONFIG + ) + + assert quant_config.exclude_modules == MODELOPT_FP8_QUANT_CONFIG["ignore"] + for excluded in ("proj_in", "proj_out", "lm_head"): + assert quant_config.is_module_excluded_from_quantization(excluded) + + def test_absent_quantization_config_resolves_to_no_quantization(self): + quant_config, layer_quant_config, dynamic_weight, dynamic_activation = ( + DiffusionPipelineConfig.load_diffusion_quant_config({}) + ) + + assert quant_config.quant_algo is None + assert layer_quant_config is None + assert dynamic_weight is False + assert dynamic_activation is False + + +@pytest.mark.parametrize( + "checkpoint_path, label", + [ + (COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8"), + (COSMOS3_SUPER_FP8_PATH, "Cosmos3-Super-FP8"), + ], +) +def test_checkpoint_config_resolves_to_static_fp8(checkpoint_path, label): + """The real checkpoint on disk -- not just the recipe dict -- resolves to static FP8.""" + if not os.path.isdir(checkpoint_path): + pytest.skip(f"{label} not found: {checkpoint_path}") + + config = DiffusionPipelineConfig.from_pretrained( + checkpoint_path, args=VisualGenArgs(model=checkpoint_path) + ) + transformer_config = config.primary_model_config + + assert transformer_config.quant_config.quant_algo == QuantAlgo.FP8 + assert transformer_config.dynamic_weight_quant is False + assert transformer_config.force_dynamic_quantization is False + assert transformer_config.quant_config.exclude_modules is not None + + +def test_bf16_checkpoint_config_resolves_to_no_quantization(): + """Regression: adding FP8 support must not quantize the BF16 checkpoints.""" + if not os.path.isdir(COSMOS3_NANO_BF16_PATH): + pytest.skip(f"Cosmos3-Nano not found: {COSMOS3_NANO_BF16_PATH}") + + config = DiffusionPipelineConfig.from_pretrained( + COSMOS3_NANO_BF16_PATH, args=VisualGenArgs(model=COSMOS3_NANO_BF16_PATH) + ) + + assert config.primary_model_config.quant_config.quant_algo is None + + +def _build_two_layer_transformer(checkpoint_path): + """Build the transformer from a checkpoint's config, trimmed to two layers. + + Only the topology is under test, so the layer count is cut to keep the build + cheap. No weights are loaded. + """ + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + model_config = DiffusionPipelineConfig.from_pretrained( + checkpoint_path, args=VisualGenArgs(model=checkpoint_path) + ).primary_model_config + model_config.pretrained_config.num_hidden_layers = 2 + return Cosmos3VFMTransformer(model_config=model_config) + + +@pytest.mark.parametrize( + "checkpoint_path, label, static_fp8", + [ + (COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8", True), + (COSMOS3_NANO_BF16_PATH, "Cosmos3-Nano", False), + ], + ids=["fp8_splits", "bf16_stays_fused"], +) +def test_topology_follows_quantization(checkpoint_path, label, static_fp8): + """Only static FP8 unfuses; BF16 must keep the fused topology untouched. + + The split exists solely to preserve per-projection calibration, which BF16 + does not have. Pinning both directions here means a change to the predicate + cannot quietly alter the BF16 path -- the one every existing Cosmos3 user is + on. + """ + _skip_if_missing(checkpoint_path, label) + + transformer = _build_two_layer_transformer(checkpoint_path) + try: + names = set(dict(transformer.named_modules())) + gen_attn, gen_mlp = "gen_layers.0.cross_attention", "gen_layers.0.mlp" + und_mlp = "language_model.layers.0.mlp" + + if static_fp8: + for split in (f"{gen_attn}.to_q", f"{gen_attn}.to_k", f"{gen_attn}.to_v"): + assert split in names, f"{label}: expected split {split}" + assert f"{gen_attn}.qkv_proj" not in names, f"{label}: GEN QKV still fused" + for mlp in (gen_mlp, und_mlp): + assert f"{mlp}.gate_proj" in names and f"{mlp}.up_proj" in names + assert f"{mlp}.gate_up_proj" not in names, f"{label}: {mlp} still fused" + else: + assert f"{gen_attn}.qkv_proj" in names, f"{label}: GEN QKV unexpectedly split" + for mlp in (gen_mlp, und_mlp): + assert f"{mlp}.gate_up_proj" in names, f"{label}: {mlp} unexpectedly split" + assert f"{mlp}.gate_proj" not in names + + # The UND tower is SEPARATE_QKV in both configurations; only the shared + # activation quantization is conditional. + und_attn = transformer.language_model.layers[0].self_attn + assert und_attn.share_qkv_input_quant is static_fp8 + finally: + del transformer + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("dynamic_field", ["dynamic_weight_quant", "force_dynamic_quantization"]) +def test_dynamic_quantization_stays_fused(dynamic_field): + """Dynamic FP8 has no calibration to preserve, so it must keep fusing. + + Both dynamic flavors resolve to ``quant_algo == FP8``, so a predicate that + keyed on the algorithm alone would unfuse them too -- and the split path + would then quantize activations against a scale that does not exist yet. + """ + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import uses_static_fp8 + + _skip_if_missing(COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8") + + model_config = DiffusionPipelineConfig.from_pretrained( + COSMOS3_NANO_FP8_PATH, args=VisualGenArgs(model=COSMOS3_NANO_FP8_PATH) + ).primary_model_config + + assert uses_static_fp8(model_config) is True + setattr(model_config, dynamic_field, True) + assert uses_static_fp8(model_config) is False + + model_config.pretrained_config.num_hidden_layers = 2 + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + transformer = Cosmos3VFMTransformer(model_config=model_config) + try: + names = set(dict(transformer.named_modules())) + assert "gen_layers.0.cross_attention.qkv_proj" in names + assert "gen_layers.0.mlp.gate_up_proj" in names + assert transformer.language_model.layers[0].self_attn.share_qkv_input_quant is False + finally: + del transformer + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.integration +@pytest.mark.high_cuda_memory +@pytest.mark.parametrize( + "checkpoint_path, label, size", + [ + (COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8", "nano"), + (COSMOS3_SUPER_FP8_PATH, "Cosmos3-Super-FP8", "super"), + ], +) +def test_static_fp8_checkpoint_realizes_expected_module_layout( + checkpoint_path, label, size, _cleanup_gpu +): + """Load the real checkpoint and pin the realized dtype/scale layout. + + Super is exercised separately from Nano rather than inferred from it: it has + a different depth and width, and roughly twice the quantized linear count. + """ + _skip_if_missing(checkpoint_path, label) + + pipeline = _load_transformer(checkpoint_path) + try: + transformer = pipeline.transformer + + fp8_by_tower = {"UND": 0, "GEN": 0, "other": 0} + missing_scales = [] + native_linears = {} + + for name, module in transformer.named_modules(): + if isinstance(module, Linear): + weight = getattr(module, "weight", None) + if weight is not None and weight.dtype == torch.float8_e4m3fn: + fp8_by_tower[_tower_of(name)] += 1 + # Both scales must survive loading: ``weight_scale`` + # dequantizes the GEMM, ``input_scale`` is what makes the + # activation path static rather than dynamic. + if getattr(module, "weight_scale", None) is None: + missing_scales.append(f"{name}.weight_scale") + if getattr(module, "input_scale", None) is None: + missing_scales.append(f"{name}.input_scale") + elif isinstance(module, torch.nn.Linear): + native_linears[name] = module.weight.dtype + + assert not missing_scales, f"{label}: missing FP8 scales: {missing_scales[:10]}" + + expected = EXPECTED_FP8_LINEARS[size] + assert fp8_by_tower["UND"] == expected["UND"], ( + f"{label}: UND tower FP8 linears {fp8_by_tower['UND']} != {expected['UND']}" + ) + assert fp8_by_tower["GEN"] == expected["GEN"], ( + f"{label}: GEN tower FP8 linears {fp8_by_tower['GEN']} != {expected['GEN']}" + ) + + assert EXPECTED_NATIVE_LINEARS.issubset(set(native_linears)), ( + f"{label}: expected native boundary linears missing: " + f"{EXPECTED_NATIVE_LINEARS - set(native_linears)}" + ) + for boundary in ("vae2llm", "llm2vae", "audio2llm", "llm2audio"): + assert native_linears[boundary] == torch.bfloat16, ( + f"{label}: {boundary} should stay BF16, got {native_linears[boundary]}" + ) + + # ``post_load_weights`` deliberately promotes the timestep embedder to + # FP32 for precision; it is excluded from quantization in the checkpoint. + timestep_dtypes = {p.dtype for p in transformer.time_embedder.parameters()} + assert timestep_dtypes == {torch.float32}, ( + f"{label}: time_embedder should be FP32, got {timestep_dtypes}" + ) + finally: + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +# Groups TensorRT-LLM used to fuse, mapped checkpoint key -> runtime module. +# Fusion kept max(shard weight_scale) and requantized the other shards onto it, +# so these are precisely the projections whose calibration the split topology +# exists to preserve. Each entry is the worst shard-scale spread in its +# checkpoint (4.67x for Nano gen QKV, 6.10x for Super gen gate/up), per a full +# sweep of all fused groups -- the case with the most to lose. +PREVIOUSLY_FUSED_GROUPS = { + "nano": { + "layers.32.self_attn.add_q_proj": "gen_layers.32.cross_attention.to_q", + "layers.32.self_attn.add_k_proj": "gen_layers.32.cross_attention.to_k", + "layers.32.self_attn.add_v_proj": "gen_layers.32.cross_attention.to_v", + }, + "super": { + "layers.7.mlp_moe_gen.gate_proj": "gen_layers.7.mlp.gate_proj", + "layers.7.mlp_moe_gen.up_proj": "gen_layers.7.mlp.up_proj", + }, +} + + +def _load_checkpoint_tensors( + checkpoint_path, keys, suffixes=("weight", "weight_scale", "input_scale") +): + import json + + from safetensors.torch import load_file + + transformer_dir = os.path.join(checkpoint_path, "transformer") + with open( + os.path.join(transformer_dir, "diffusion_pytorch_model.safetensors.index.json") + ) as handle: + weight_map = json.load(handle)["weight_map"] + + shards, tensors = {}, {} + for key in keys: + for suffix in suffixes: + full_key = f"{key}.{suffix}" + if full_key not in weight_map: + continue + shard = weight_map[full_key] + if shard not in shards: + shards[shard] = load_file(os.path.join(transformer_dir, shard)) + tensors[full_key] = shards[shard][full_key] + return tensors + + +@pytest.mark.integration +@pytest.mark.high_cuda_memory +@pytest.mark.parametrize( + "checkpoint_path, label, size", + [ + (COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8", "nano"), + (COSMOS3_SUPER_FP8_PATH, "Cosmos3-Super-FP8", "super"), + ], +) +def test_previously_fused_groups_now_load_exactly(checkpoint_path, label, size, _cleanup_gpu): + """Every member of a formerly fused group must transcribe bit-for-bit. + + Fusion kept one weight scale per group and re-quantized the other members + onto it. Splitting the topology is only worth doing if each projection now + loads its own tensor and its own scale untouched, so this asserts exact + equality rather than a tolerance -- there is no arithmetic left to drift. + + The group's weight scales are asserted to actually differ first. Were they + equal, fusion would have been lossless and exactness here would hold + trivially, so the check would no longer discriminate between the two + topologies. + """ + _skip_if_missing(checkpoint_path, label) + + group = PREVIOUSLY_FUSED_GROUPS[size] + tensors = _load_checkpoint_tensors(checkpoint_path, list(group)) + first_key = next(iter(group)) + if f"{first_key}.weight" not in tensors: + pytest.skip(f"{label}: group {first_key} not present in checkpoint") + + weight_scales = {k: tensors[f"{k}.weight_scale"].float().item() for k in group} + input_scales = {k: tensors[f"{k}.input_scale"].float().item() for k in group} + + assert len(set(weight_scales.values())) > 1, ( + f"{label}: group {list(group)} has a single weight scale {weight_scales}, so it " + "cannot distinguish split loading from fused requantization -- pick a " + "group whose shard scales differ" + ) + + # q/k/v (and gate/up) see the same activation, so ModelOpt calibrates one + # shared input scale per group. The split path relies on that to quantize + # the activation once and hand the same tensor to each projection. + assert len(set(input_scales.values())) == 1, ( + f"{label}: group {list(group)} has differing input scales {input_scales}" + ) + + pipeline = _load_transformer(checkpoint_path) + try: + modules = dict(pipeline.transformer.named_modules()) + for checkpoint_key, runtime_name in group.items(): + parent = runtime_name.rsplit(".", 1)[0] + siblings = sorted(n for n in modules if n.startswith(parent))[:8] + assert runtime_name in modules, ( + f"{label}: expected split module {runtime_name}; the topology may " + f"have reverted to fused (present: {siblings})" + ) + module = modules[runtime_name] + + assert module.weight_scale.float().item() == pytest.approx( + weight_scales[checkpoint_key], rel=0, abs=0 + ), f"{label}: {runtime_name} weight_scale was rescaled" + assert module.input_scale.float().item() == pytest.approx( + input_scales[checkpoint_key], rel=0, abs=0 + ), f"{label}: {runtime_name} input_scale was rescaled" + + expected = tensors[f"{checkpoint_key}.weight"] + actual = module.weight.detach().cpu() + assert actual.dtype == expected.dtype == torch.float8_e4m3fn + # Compared as bits: FP8 has no exact torch.equal on all platforms and + # this must catch a single re-rounded element. + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), ( + f"{label}: {runtime_name} weight differs from the checkpoint; " + "it was re-quantized rather than loaded directly" + ) + finally: + del pipeline + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.integration +@pytest.mark.high_cuda_memory +def test_static_fp8_scales_match_checkpoint_calibration(_cleanup_gpu): + """Loaded scales must equal the checkpoint's calibrated values. + + ModelOpt stores ``.weight_scale``/``.input_scale`` next to + duplicate quantizer-internal tensors (``weight_quantizer._scale``, + ``*._amax``). Reading the wrong one -- or silently falling back to a + computed scale -- would still produce plausible images, so compare against + the raw checkpoint tensors. + """ + _skip_if_missing(COSMOS3_NANO_FP8_PATH, "Cosmos3-Nano-FP8") + + # An unfused UND projection: its scales must transcribe exactly, with none + # of the max-scale rescaling the fused groups undergo. + checkpoint_key = "layers.0.self_attn.to_q" + tensors = _load_checkpoint_tensors(COSMOS3_NANO_FP8_PATH, [checkpoint_key]) + expected_weight_scale = tensors[f"{checkpoint_key}.weight_scale"].float().item() + expected_input_scale = tensors[f"{checkpoint_key}.input_scale"].float().item() + + pipeline = _load_transformer(COSMOS3_NANO_FP8_PATH) + try: + module = dict(pipeline.transformer.named_modules())[ + "language_model.layers.0.self_attn.to_q" + ] + assert module.weight.dtype == torch.float8_e4m3fn + assert module.weight_scale.float().item() == pytest.approx(expected_weight_scale, rel=1e-6) + assert module.input_scale.float().item() == pytest.approx(expected_input_scale, rel=1e-6) + finally: + del pipeline + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision.py b/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision.py new file mode 100644 index 000000000000..6c61eee0176a --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-denoising-step activation precision for static-FP8 Cosmos3. + +The outer denoising steps run the resident FP8 weights through a 16-bit GEMM +instead of the checkpoint's quantized activation path. Two properties have to +hold or the feature is silently absent: the step policy must select the same +path for every call within a step, and the callers that pre-quantize a shared +activation must stand down while it is selected -- otherwise a "high precision" +step still receives FP8 activations and nothing changed. +""" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.models.cosmos3.step_precision import ( + StepPrecisionController, + StepPrecisionFp8LinearMethod, + apply_fp8_w8a16_linear, + parse_diffusion_step_policy, +) +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + + +def _policy(**overrides): + policy = { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": 3, "mode": "a16"}, + "last_steps": {"count": 3, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16", + } + policy.update(overrides) + return {"runtime": {"diffusion_step_policy": policy}} + + +# Pure-CPU test (stub modules, no device): runs in the CPU lane (l0_cpu.yml), +# which selects with `-m cpu_only`; GPU stages deselect it via `not cpu_only`. +pytestmark = pytest.mark.cpu_only + + +class _StubBaseMethod: + """Stands in for FP8QDQLinearMethod; records whether the FP8 path ran.""" + + def __init__(self): + self.calls = 0 + + def apply(self, module, input, bias=None): + self.calls += 1 + return torch.zeros(input.shape[0], module.weight.shape[0], dtype=input.dtype) + + +class _StubLinear(torch.nn.Module): + def __init__(self, out_features=4, in_features=8, scale=2.0): + super().__init__() + self.weight = torch.ones(out_features, in_features, dtype=torch.bfloat16) + self.weight_scale = torch.tensor(scale, dtype=torch.float32) + + +class TestStepPolicy: + @pytest.mark.parametrize( + "step_index,expected", + [ + (0, True), + (1, True), + (2, True), + (3, False), + (25, False), + (46, False), + (47, True), + (49, True), + ], + ) + def test_first_and_last_steps_are_high_precision(self, step_index, expected): + controller = StepPrecisionController(first_steps=3, last_steps=3) + controller.set_step(step_index, num_steps=50) + assert controller.high_precision is expected + + def test_selection_is_a_pure_function_of_the_step(self): + """CFG branches call separately; they must not disagree within a step.""" + controller = StepPrecisionController(first_steps=3, last_steps=3) + controller.set_step(1, num_steps=50) + first = controller.high_precision + controller.set_step(1, num_steps=50) + assert controller.high_precision is first is True + + def test_single_step_schedule_stays_on_the_quantized_path(self): + """A one-step schedule is the warmup probe, not an all-edge request.""" + controller = StepPrecisionController(first_steps=3, last_steps=3) + controller.set_step(0, num_steps=1) + assert controller.high_precision is False + + def test_zero_windows_disable_the_feature(self): + controller = StepPrecisionController(first_steps=0, last_steps=0) + for step in range(4): + controller.set_step(step, num_steps=4) + assert controller.high_precision is False + + def test_overlapping_windows_cover_every_step(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + for step in range(4): + controller.set_step(step, num_steps=4) + assert controller.high_precision is True + + def test_reset_clears_state(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + controller.set_step(0, num_steps=50) + controller.reset() + assert controller.high_precision is False + + @pytest.mark.parametrize("first,last", [(-1, 3), (3, -1)]) + def test_negative_windows_rejected(self, first, last): + with pytest.raises(ValueError, match="non-negative"): + StepPrecisionController(first_steps=first, last_steps=last) + + def test_out_of_range_step_rejected(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + with pytest.raises(IndexError): + controller.set_step(50, num_steps=50) + + def test_non_positive_num_steps_rejected(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + with pytest.raises(ValueError, match="num_steps must be positive"): + controller.set_step(0, num_steps=0) + + +class TestDispatch: + def test_middle_step_uses_the_checkpoint_path(self): + base = _StubBaseMethod() + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(base, controller) + controller.set_step(10, num_steps=50) + module = _StubLinear() + method.apply(module, torch.ones(2, 8, dtype=torch.bfloat16)) + assert base.calls == 1 + + def test_edge_step_bypasses_the_checkpoint_path(self): + base = _StubBaseMethod() + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(base, controller) + controller.set_step(0, num_steps=50) + module = _StubLinear() + out = method.apply(module, torch.ones(2, 8, dtype=torch.bfloat16)) + assert base.calls == 0 + # weight 1.0 * scale 2.0, summed over in_features=8 -> 16 per output. + assert torch.allclose(out, torch.full_like(out, 16.0)) + + def test_high_precision_is_published_for_sharing_callers(self): + """GatedMLP/Attention read this attribute to stand down. Contract test.""" + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(_StubBaseMethod(), controller) + controller.set_step(10, num_steps=50) + assert method.high_precision is False + controller.set_step(0, num_steps=50) + assert method.high_precision is True + + def test_wrapper_forwards_unknown_attributes(self): + base = _StubBaseMethod() + base.quantizes_nvfp4_activations = False + method = StepPrecisionFp8LinearMethod(base, StepPrecisionController(3, 3)) + assert method.quantizes_nvfp4_activations is False + + +class TestW8A16Apply: + def test_dequantized_weight_matches_scaled_reference(self): + module = _StubLinear(out_features=3, in_features=4, scale=0.5) + module.weight = torch.arange(12, dtype=torch.bfloat16).reshape(3, 4) + x = torch.ones(2, 4, dtype=torch.bfloat16) + out = apply_fp8_w8a16_linear(module, x, bias=None) + expected = torch.nn.functional.linear(x, module.weight.to(x.dtype) * 0.5) + assert torch.allclose(out, expected) + + def test_bias_is_applied(self): + module = _StubLinear(out_features=2, in_features=3, scale=1.0) + module.weight = torch.zeros(2, 3, dtype=torch.bfloat16) + bias = torch.tensor([1.0, -1.0], dtype=torch.bfloat16) + out = apply_fp8_w8a16_linear(module, torch.ones(1, 3, dtype=torch.bfloat16), bias) + assert torch.allclose(out, bias.unsqueeze(0)) + + def test_prequantized_activation_is_rejected(self): + """The failure this feature can have silently: an FP8 activation on a + 16-bit step means a sharing caller did not stand down, and the step is + not actually running in higher precision.""" + module = _StubLinear() + x = torch.ones(2, 8, dtype=torch.bfloat16).to(torch.float8_e4m3fn) + with pytest.raises(RuntimeError, match="must stand down"): + apply_fp8_w8a16_linear(module, x, bias=None) + + +class TestPolicyParsing: + """The checkpoint states the recipe; a shape we do not implement must fail. + + Half-honouring a policy is the dangerous outcome: a checkpoint that asked + for something we silently ignored looks exactly like the feature working. + """ + + def test_published_policy_is_accepted(self): + policy = parse_diffusion_step_policy(_policy()) + assert (policy.first_steps, policy.last_steps) == (3, 3) + assert policy.reasoner_high_precision is True + + @pytest.mark.parametrize( + "config", + [None, {}, {"runtime": {}}, {"runtime": {"other": 1}}, "not-a-mapping"], + ids=["none", "empty", "no-policy", "other-key", "not-mapping"], + ) + def test_absent_policy_returns_none(self, config): + assert parse_diffusion_step_policy(config) is None + + def test_scope_without_transformer_is_inert(self): + assert parse_diffusion_step_policy(_policy(scope=["vae"])) is None + + def test_reasoner_native_is_honoured(self): + policy = parse_diffusion_step_policy(_policy(reasoner="native")) + assert policy.reasoner_high_precision is False + + @pytest.mark.parametrize( + "overrides,match", + [ + ({"schema_version": 2}, "schema_version"), + ({"schema_version": True}, "schema_version"), + ({"type": "every_n"}, "type"), + ({"index_space": "sampler_step"}, "index_space"), + ({"default_mode": "a16"}, "default_mode"), + ({"overlap": "native"}, "overlap"), + ({"reasoner": "fp8"}, "reasoner"), + ], + ) + def test_unimplemented_policy_shapes_are_refused(self, overrides, match): + with pytest.raises(ValueError, match=match): + parse_diffusion_step_policy(_policy(**overrides)) + + def test_unknown_field_is_refused(self): + with pytest.raises(ValueError, match="Unknown diffusion_step_policy fields"): + parse_diffusion_step_policy(_policy(future_knob="x")) + + def test_missing_field_is_refused(self): + config = _policy() + del config["runtime"]["diffusion_step_policy"]["overlap"] + with pytest.raises(ValueError, match="Missing diffusion_step_policy fields"): + parse_diffusion_step_policy(config) + + @pytest.mark.parametrize( + "value", [{"count": -1, "mode": "a16"}, {"count": 3, "mode": "native"}, {"count": 3}] + ) + def test_bad_step_range_is_refused(self, value): + with pytest.raises((ValueError, TypeError)): + parse_diffusion_step_policy(_policy(first_steps=value)) + + def test_non_mapping_policy_is_refused(self): + with pytest.raises(TypeError, match="must be a mapping"): + parse_diffusion_step_policy({"runtime": {"diffusion_step_policy": []}}) + + +class TestReasonerPath: + """The reasoner runs once per request, on the first transformer call, so + its precision is stated by the policy rather than derived from a step + index -- which would match only while that call lands inside a window.""" + + def test_always_high_ignores_the_step(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(_StubBaseMethod(), controller, always_high=True) + for step in (0, 10, 25, 49): + controller.set_step(step, num_steps=50) + assert method.high_precision is True + + def test_generation_path_still_follows_the_step(self): + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(_StubBaseMethod(), controller, always_high=False) + controller.set_step(25, num_steps=50) + assert method.high_precision is False + + def test_always_high_takes_the_16bit_path_mid_schedule(self): + base = _StubBaseMethod() + controller = StepPrecisionController(first_steps=3, last_steps=3) + method = StepPrecisionFp8LinearMethod(base, controller, always_high=True) + controller.set_step(25, num_steps=50) + method.apply(_StubLinear(), torch.ones(2, 8, dtype=torch.bfloat16)) + assert base.calls == 0 + + +class TestTransformerWiring: + """Pins the wiring between checkpoint and towers. + + The component tests call install_step_precision themselves, so they cannot + see a mistake in how the transformer decides: reading the wrong config key, + or giving the unconditional path to the generation tower instead of the + reasoner. Both would leave every other test green. + """ + + @staticmethod + def _fp8_linear(): + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.modules.linear import Linear + + return Linear( + 8, + 8, + dtype=torch.bfloat16, + quant_config=ModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8) + ).get_quant_config(), + ) + + def _stub_transformer(self, quantization_config): + """A transformer-shaped object: the two towers and the model config.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + + model_config = DiffusionModelConfig( + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8), + pretrained_config=SimpleNamespace(quantization_config=quantization_config), + ) + return SimpleNamespace( + model_config=model_config, + gen_layers=torch.nn.ModuleList([self._fp8_linear()]), + language_model=SimpleNamespace(layers=torch.nn.ModuleList([self._fp8_linear()])), + step_precision_controller=None, + ) + + @staticmethod + def _install(stub): + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + Cosmos3VFMTransformer._maybe_install_step_precision(stub) + + def test_policy_wraps_generation_step_gated_and_reasoner_always(self): + stub = self._stub_transformer(_policy()) + self._install(stub) + + assert stub.step_precision_controller is not None + gen = stub.gen_layers[0].quant_method + reasoner = stub.language_model.layers[0].quant_method + assert isinstance(gen, StepPrecisionFp8LinearMethod) + assert isinstance(reasoner, StepPrecisionFp8LinearMethod) + # The distinguishing property: mid-schedule the towers disagree. + stub.step_precision_controller.set_step(25, num_steps=50) + assert gen.high_precision is False, "generation tower is not step-gated" + assert reasoner.high_precision is True, "reasoner tower is not unconditional" + + def test_reasoner_native_leaves_the_reasoner_alone(self): + stub = self._stub_transformer(_policy(reasoner="native")) + self._install(stub) + assert isinstance(stub.gen_layers[0].quant_method, StepPrecisionFp8LinearMethod) + assert not isinstance( + stub.language_model.layers[0].quant_method, StepPrecisionFp8LinearMethod + ) + + def test_checkpoint_without_a_policy_wraps_nothing(self): + stub = self._stub_transformer({}) + self._install(stub) + assert stub.step_precision_controller is None + for tower in (stub.gen_layers[0], stub.language_model.layers[0]): + assert not isinstance(tower.quant_method, StepPrecisionFp8LinearMethod) + + def test_policy_is_read_from_the_documented_config_key(self): + """A policy under any other key must not be picked up.""" + stub = self._stub_transformer({"diffusion_step_policy": _policy()["runtime"]}) + self._install(stub) + assert stub.step_precision_controller is None + + def test_installing_twice_keeps_wrappers_on_the_live_controller(self) -> None: + """post_load_weights running twice must not orphan the wrappers. + + The second install previously skipped already-wrapped modules and then + cleared the controller, leaving every wrapper bound to one nothing + drives: set_denoising_step would stop reaching the layers it steers, + silently, with the edge steps landing on the quantized path. + """ + stub = self._stub_transformer(_policy()) + self._install(stub) + first_controller = stub.step_precision_controller + + self._install(stub) + second_controller = stub.step_precision_controller + assert second_controller is not None, "controller was cleared by the second install" + + gen = stub.gen_layers[0].quant_method + reasoner = stub.language_model.layers[0].quant_method + assert gen.controller is second_controller + assert reasoner.controller is second_controller + # Not double-wrapped: the base method must still be the FP8 one. + assert not isinstance(gen.base_method, StepPrecisionFp8LinearMethod) + # And the live controller actually steers them. + second_controller.set_step(0, num_steps=50) + assert gen.high_precision is True + second_controller.set_step(25, num_steps=50) + assert gen.high_precision is False + assert reasoner.high_precision is True + assert first_controller is not None diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision_component.py b/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision_component.py new file mode 100644 index 000000000000..f663ce9274d5 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_step_precision_component.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Component tests for step precision: x . + +Level-1 policy tests live in test_cosmos3_step_precision.py and use stubs. +These build the real quantized components the feature acts on -- a split +GatedMLP -- drive it through the checkpoint's declared policy, and +compare feature-on against feature-off in the same job. There is no stored +reference: the comparisons are either exact arithmetic or an A/B against the +same module's other path. + +The property that actually matters is not "a flag flipped". These components +quantize the shared activation *above* the Linear, so if they do not stand +down while a high-precision step is selected, the step still runs on FP8 +activations and the feature is silently absent. That is what is pinned here. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.visual_gen.models.cosmos3.step_precision import ( + StepPrecisionController, + install_step_precision, + parse_diffusion_step_policy, +) +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HIDDEN, INTERMEDIATE, TOKENS = 512, 1024, 128 +# Chosen so the whole chain stays inside FP8's representable range. The +# structural tests next door use much smaller scales, which is fine when only +# topology is asserted, but here the intermediate activation would underflow +# FP8 to exactly zero and every numerical comparison would be vacuous. +INPUT_SCALE, WEIGHT_SCALE = 1e-2, 1.0 +NUM_STEPS = 50 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _make_mlp(split=True): + config = ModelConfig(quant_config=QuantConfig(quant_algo=QuantAlgo.FP8)) + mlp = GatedMLP( + hidden_size=HIDDEN, + intermediate_size=INTERMEDIATE, + bias=False, + activation=F.silu, + dtype=torch.bfloat16, + config=config, + split_gate_up=split, + ).cuda() + for linear in _linears(mlp): + linear.weight.data.copy_( + (torch.randn(*linear.weight.shape, device="cuda", dtype=torch.bfloat16) * 0.05).to( + torch.float8_e4m3fn + ) + ) + linear.weight_scale.data.fill_(WEIGHT_SCALE) + linear.input_scale.data.fill_(INPUT_SCALE) + linear.inv_input_scale.data.fill_(1.0 / INPUT_SCALE) + return mlp + + +def _linears(mlp): + names = ("gate_proj", "up_proj") if mlp.split_gate_up else ("gate_up_proj",) + return [getattr(mlp, n) for n in names] + [mlp.down_proj] + + +def _policy(first=3, last=3): + """The block the checkpoint publishes under quantization_config.runtime.""" + return { + "runtime": { + "diffusion_step_policy": { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": first, "mode": "a16"}, + "last_steps": {"count": last, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16", + } + } + } + + +def _install(mlp, quantization_config): + """Install exactly as the transformer does: parse the checkpoint, then wrap.""" + policy = parse_diffusion_step_policy(quantization_config) + if policy is None: + return None + controller = StepPrecisionController( + first_steps=policy.first_steps, last_steps=policy.last_steps + ) + assert install_step_precision([mlp], controller) == len(_linears(mlp)) + return controller + + +@requires_cuda +def test_checkpoint_without_a_policy_installs_nothing(): + """Absence is the signal, not a default to fill in. + + The image and distilled 4-step FP8 builds ship no policy and must run + fully quantized rather than inherit another checkpoint's windows. + """ + mlp = _make_mlp() + assert _install(mlp, {}) is None + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + # Without a wrapper nothing publishes high_precision, so the shared-input + # optimization stays engaged on every step. + assert mlp._can_share_gate_up_quantization(x) is True + + +@requires_cuda +def test_shared_quantization_stands_down_only_on_edge_steps(): + """The integration property the whole feature rests on. + + gate/up consume one activation, which the split path quantizes once above + the Linear. On a 16-bit step that must not happen, or the step receives FP8 + activations and is not actually running in higher precision. + """ + mlp = _make_mlp() + controller = _install(mlp, _policy()) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + + controller.set_step(0, NUM_STEPS) + assert mlp._can_share_gate_up_quantization(x) is False, "edge step still pre-quantizes" + + controller.set_step(NUM_STEPS // 2, NUM_STEPS) + assert mlp._can_share_gate_up_quantization(x) is True, "middle step lost the optimization" + + controller.set_step(NUM_STEPS - 1, NUM_STEPS) + assert mlp._can_share_gate_up_quantization(x) is False, "final step still pre-quantizes" + + +@requires_cuda +def test_edge_step_linear_matches_exact_dequantized_reference(): + """On a 16-bit step a projection is plain bf16 arithmetic, so it is exact. + + The reference is the same resident FP8 weight dequantized by its own + weight_scale -- nothing is read that the checkpoint did not already supply. + """ + mlp = _make_mlp() + controller = _install(mlp, _policy()) + controller.set_step(0, NUM_STEPS) + + linear = mlp.gate_proj + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + got = linear.quant_method.apply(linear, x, None) + expected = F.linear(x, linear.weight.to(x.dtype) * linear.weight_scale.to(x.dtype)) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +@requires_cuda +def test_edge_and_middle_steps_produce_different_output(): + """Feature-on vs feature-off, same module, same input, same job. + + If these matched, the 16-bit step would be doing nothing -- which is the + failure mode a flag-only test cannot see. + """ + mlp = _make_mlp() + controller = _install(mlp, _policy()) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + + controller.set_step(NUM_STEPS // 2, NUM_STEPS) + quantized = mlp(x).float() + controller.set_step(0, NUM_STEPS) + high_precision = mlp(x).float() + + assert torch.isfinite(high_precision).all() + assert high_precision.abs().mean().item() > 0, "output collapsed to zero" + # Scale-relative: an absolute tolerance here would pass trivially whenever + # the configured scales push the output near zero. + relative = ((high_precision - quantized).abs().mean() / high_precision.abs().mean()).item() + assert relative > 1e-3, ( + f"the 16-bit step produced the same output as the quantized step " + f"(relative difference {relative:.3g}), so it did not engage" + ) + + +@requires_cuda +def test_edge_step_is_closer_to_the_unquantized_reference(): + """The feature's actual claim: less activation-quantization error. + + Reference runs the same weights and activations with no activation + quantization anywhere, which is what a 16-bit step approximates. + """ + mlp = _make_mlp() + controller = _install(mlp, _policy()) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + + def dequant(linear): + return linear.weight.to(torch.float32) * linear.weight_scale.to(torch.float32) + + xf = x.float() + gate = F.linear(xf, dequant(mlp.gate_proj)) + up = F.linear(xf, dequant(mlp.up_proj)) + reference = F.linear(F.silu(gate) * up, dequant(mlp.down_proj)) + + controller.set_step(NUM_STEPS // 2, NUM_STEPS) + quantized_err = (mlp(x).float() - reference).abs().mean().item() + controller.set_step(0, NUM_STEPS) + high_precision_err = (mlp(x).float() - reference).abs().mean().item() + + assert high_precision_err < quantized_err, ( + f"16-bit step was not closer to the unquantized reference: " + f"{high_precision_err:.6g} vs {quantized_err:.6g}" + ) + + +@requires_cuda +def test_zero_windows_keep_every_step_quantized(): + """A policy may declare empty windows; every step then runs quantized.""" + mlp = _make_mlp() + controller = _install(mlp, _policy(first=0, last=0)) + x = torch.randn(TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + for step in (0, NUM_STEPS // 2, NUM_STEPS - 1): + controller.set_step(step, NUM_STEPS) + assert mlp._can_share_gate_up_quantization(x) is True diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py index af2a40e2bad9..c20a9a4f7cf5 100644 --- a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -101,6 +101,37 @@ def test_excluded_module_keeps_high_precision_weights(self): loader.load_linear_weights(module, "proj_out", [weight_dict]) assert module.loaded == [weight_dict] + def test_unquantizable_module_keeps_high_precision_weights(self) -> None: + """A module that was never built quantized must not be refused. + + ``quant_algo`` falls back to the global recipe for any module without + its own ``quant_config``, and that includes modules which cannot be + quantized: ``Embedding`` reaches this loader because it subclasses + ``LMHead`` -> ``Linear``, but its ``__init__`` never exposes + ``quant_config``, so its buffer stays high precision. ModelOpt does not + list it in ``ignore`` either, since only ``Linear`` targets were ever + candidates. Nothing was built quantized, so nothing can be corrupted. + """ + loader = _make_loader(QuantAlgo.FP8) + module = _StubLinear() + module.weight = torch.zeros(8, 16, dtype=torch.bfloat16) + weight_dict = _bf16_weights() + loader.load_linear_weights(module, "language_model.embed_tokens", [weight_dict]) + assert module.loaded == [weight_dict] + + def test_quantized_destination_still_refuses_unquantized_checkpoint(self) -> None: + """The destination check must not weaken the guard it sits in front of. + + A module built for FP8 holds a float8 buffer, so a high-precision + checkpoint weight is still the silent-corruption case and still raises. + """ + loader = _make_loader(QuantAlgo.FP8) + module = _StubLinear() + module.weight = torch.zeros(8, 16, dtype=torch.float8_e4m3fn) + with pytest.raises(ValueError, match="appears to be unquantized"): + loader.load_linear_weights(module, "blocks.0.attn1.to_q", [_bf16_weights()]) + assert module.loaded is None + def test_unquantized_recipe_is_unaffected(self): loader = _make_loader(None) module = _StubLinear()