[None][fix] handle rank-3 Fp4 and non-contiguous Ulysses shard in LTX-2 NVFP4 fused paths (#15718)

Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>

Author: luyiyun1021

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -216,6 +216,15 @@ def _fused_gate_up_swiglu(self, x, fp4_out=False):
         if not isinstance(x, (tuple, Fp4QuantizedTensor)) and x.dim() > 2:
             original_shape = x.shape
             x = x.reshape(-1, x.shape[-1])
+        elif isinstance(x, Fp4QuantizedTensor) and x.fp4_tensor.dim() > 2:
+            # Fused GEMM needs a 2D mat1: flatten a rank-3 fp4 activation's data
+            # to [M, D/2] (its SF is already flat-M, so reuse it); unflatten below.
+            original_shape = x.fp4_tensor.shape
+            x = Fp4QuantizedTensor(
+                fp4_tensor=x.fp4_tensor.reshape(-1, x.fp4_tensor.shape[-1]),
+                scaling_factor=x.scaling_factor,
+                is_sf_swizzled=x.is_sf_swizzled,
+            )
 
         # Get quantized inputs from Linear's NVFP4 pipeline
         act_fp4, act_sf, alpha = module.quant_method._input_prepare(module, x)
@@ -271,11 +280,10 @@ def forward(
             if torch.compiler.is_compiling():
                 fp4_out = True
             else:
-                if isinstance(x, (tuple, Fp4QuantizedTensor)):
-                    m = x[0].shape[0] if isinstance(x, tuple) else x.shape[0]
-                else:
-                    m = x.reshape(
-                        -1, x.shape[-1]).shape[0] if x.dim() > 2 else x.shape[0]
+                t = x.fp4_tensor if isinstance(x, Fp4QuantizedTensor) else (
+                    x[0] if isinstance(x, tuple) else x)
+                m = t.reshape(
+                    -1, t.shape[-1]).shape[0] if t.dim() > 2 else t.shape[0]
                 fp4_out = m >= GatedMLP._FP4OUT_MIN_M
             h2 = self._fused_gate_up_swiglu(x, fp4_out=fp4_out)
         elif self._can_fuse_gate_up_swiglu():

tensorrt_llm/_torch/modules/mlp.py

@@ -180,10 +180,12 @@ def _gelu_fusion_eligibility(self) -> Tuple[bool, bool]:
 
     @staticmethod
     def _token_count(x: Union[torch.Tensor, tuple, Fp4QuantizedTensor]) -> int:
-        if isinstance(x, (tuple, Fp4QuantizedTensor)):
-            return x[0].shape[0] if isinstance(x, tuple) else x.shape[0]
-        return x.reshape(-1,
-                         x.shape[-1]).shape[0] if x.dim() > 2 else x.shape[0]
+        # Row count M = product of leading dims (B*S for [B, S, ...]); keeps the
+        # fp4-out switch correct for rank-3 / pre-quantized inputs.
+        t = x.fp4_tensor if isinstance(
+            x, Fp4QuantizedTensor) else (x[0] if isinstance(x, tuple) else x)
+        return t.reshape(-1,
+                         t.shape[-1]).shape[0] if t.dim() > 2 else t.shape[0]
 
     def _fused_gelu(
         self,
@@ -205,6 +207,15 @@ def _fused_gelu(
         if not isinstance(x, (tuple, Fp4QuantizedTensor)) and x.dim() > 2:
             original_shape = x.shape
             x = x.reshape(-1, x.shape[-1])
+        elif isinstance(x, Fp4QuantizedTensor) and x.fp4_tensor.dim() > 2:
+            # Fused GEMM needs a 2D mat1: flatten a rank-3 fp4 activation's data
+            # to [M, D/2] (its SF is already flat-M, so reuse it); unflatten below.
+            original_shape = x.fp4_tensor.shape
+            x = Fp4QuantizedTensor(
+                fp4_tensor=x.fp4_tensor.reshape(-1, x.fp4_tensor.shape[-1]),
+                scaling_factor=x.scaling_factor,
+                is_sf_swizzled=x.is_sf_swizzled,
+            )
 
         act_fp4, act_sf, alpha = module.quant_method._input_prepare(module, x)
 

tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py

@@ -477,8 +477,11 @@ def forward_async(
             and self.qk_norm
         )
 
-        # SEPARATE_QKV self-attn 3x fp4_quantize dedup; see Attention.forward_async.
-        if self._maybe_share_qkv_quantize and getattr(self.to_q, "input_scale", None) is not None:
+        # SEPARATE_QKV self-attn shares ONE input quant across to_q/to_k/to_v (mirrors
+        # Attention.forward_async): reuse x if already fp4, else quantize once.
+        if isinstance(x, Fp4QuantizedTensor):
+            qkv_input = x
+        elif self._maybe_share_qkv_quantize and getattr(self.to_q, "input_scale", None) is not None:
             x_2d = x.reshape(-1, x.shape[-1])
             fp4, sf = torch.ops.trtllm.tunable_fp4_quantize(
                 x_2d, self.to_q.input_scale, self.to_q.scaling_vector_size, False

tensorrt_llm/_torch/visual_gen/utils.py

@@ -198,7 +198,9 @@ def shard(
         start = self._rank * chunk
         idx = [slice(None)] * tensor.ndim
         idx[dim] = slice(start, start + chunk)
-        return tensor[tuple(idx)]
+        # A dim>=1 block slice is non-contiguous when a leading dim is >1 (e.g.
+        # batched CFG B=2); fused DiT kernels need a dense buffer. No-op at B==1.
+        return tensor[tuple(idx)].contiguous()
 
     def shard_rope(
         self,

tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py

@@ -228,3 +228,57 @@ def _spy(x, fp4_out=False):
         assert seen["fp4_out"] == expected, (
             f"m={m}: expected fp4_out={expected}, got {seen['fp4_out']}"
         )
+
+
+# ---------------------------------------------------------------------------
+# The fused MLP flattens a rank-3 Fp4QuantizedTensor to 2D before the (2D-only)
+# GEMM. A fused rmsnorm hands it [B, S, D/2]; the path previously flattened only
+# plain tensors, tripping `assert inputs[0].dim() == 2`.
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("kind", ["gelu", "swiglu"])
+def test_fused_act_flattens_3d_fp4_input(kind):
+    """Fused NVFP4 MLP flattens a rank-3 Fp4QuantizedTensor to 2D before the GEMM."""
+    import types
+
+    from tensorrt_llm._torch.utils import Fp4QuantizedTensor, gelu_tanh
+
+    B, S, K = 2, 8, 64
+    if kind == "gelu":
+        from tensorrt_llm._torch.modules.mlp import MLP
+
+        mod = MLP(hidden_size=K, intermediate_size=128, bias=False, activation=gelu_tanh)
+        proj = mod.up_proj
+        run = lambda x: mod._fused_gelu(x)  # noqa: E731
+    else:
+        from tensorrt_llm._torch.modules.gated_mlp import GatedMLP
+
+        mod = GatedMLP(hidden_size=K, intermediate_size=128, bias=False)
+        proj = mod.gate_up_proj
+        run = lambda x: mod._fused_gate_up_swiglu(x)  # noqa: E731
+
+    captured = {}
+
+    class _Stop(Exception):
+        pass
+
+    def fake_input_prepare(module, x):
+        captured["dim"] = x.fp4_tensor.dim() if isinstance(x, Fp4QuantizedTensor) else x.dim()
+        captured["shape"] = (
+            tuple(x.fp4_tensor.shape) if isinstance(x, Fp4QuantizedTensor) else tuple(x.shape)
+        )
+        raise _Stop  # short-circuit before the real (HW-only) GEMM
+
+    proj.quant_method = types.SimpleNamespace(_input_prepare=fake_input_prepare)
+
+    fp4_3d = Fp4QuantizedTensor(
+        fp4_tensor=torch.zeros(B, S, K // 2, dtype=torch.uint8),
+        scaling_factor=torch.zeros(B * S, K // 16, dtype=torch.uint8),
+    )
+    with pytest.raises(_Stop):
+        run(fp4_3d)
+
+    assert captured["dim"] == 2, (
+        f"{kind}: rank-3 Fp4QuantizedTensor not flattened before the GEMM "
+        f"(kernel would see rank {captured['dim']})"
+    )
+    assert captured["shape"] == (B * S, K // 2)

tests/unittest/_torch/visual_gen/test_ltx2_transformer.py

@@ -74,6 +74,45 @@ def _init_all_weights(model: torch.nn.Module, std: float = 0.02):
                 torch.nn.init.normal_(p, mean=0.0, std=std)
 
 
+# float32 quant-scale params (filled to 1.0, not normal_, so dequant is well-behaved).
+_FP4_FLOAT_SCALE_SUFFIXES = (
+    ".input_scale",
+    ".inv_input_scale",
+    ".weight_scale_2",
+    ".alpha",
+    ".kv_scales",
+    ".inv_kv_scales",
+)
+
+
+def _init_weights_quant_safe(model: torch.nn.Module, std: float = 0.02):
+    """Like ``_init_all_weights`` but safe for NVFP4 modules.
+
+    NVFP4 Linears hold a packed-fp4 weight (``float4_e2m1x2``) and an FP8 scale
+    factor (``weight_scale``) whose dtypes ``torch.nn.init.normal_`` cannot fill.
+    Fill standard float params (norm weights and quant scales to 1.0, the rest
+    random); for the sub-byte/FP8 quant tensors write a benign finite byte.
+    """
+    with torch.no_grad():
+        for name, p in model.named_parameters():
+            if p.numel() == 0:
+                continue
+            if p.dtype in (torch.float32, torch.float16, torch.bfloat16):
+                if any(name.endswith(s) for s in _FP4_FLOAT_SCALE_SUFFIXES):
+                    p.fill_(1.0)
+                elif "norm" in name and "weight" in name:
+                    p.fill_(1.0)
+                else:
+                    torch.nn.init.normal_(p, mean=0.0, std=std)
+            else:
+                # Packed-fp4 weight / FP8 scale factor: normal_ is undefined on
+                # these dtypes, so write a benign finite byte via a uint8 view.
+                try:
+                    p.view(torch.uint8).fill_(1)
+                except Exception:
+                    pass
+
+
 def _make_video_positions(
     batch: int, n_patches: int, n_frames: int, grid_h: int, grid_w: int, device: torch.device
 ) -> torch.Tensor:
@@ -390,6 +429,153 @@ def test_video_only_input_to_audio_video_model(self):
         )
 
 
+# AudioVideo config with real fused-AdaLN hidden dims (video/audio both 32*64=2048,
+# in _FUSED_ADALN_SUPPORTED_DIMS). Smaller dims (e.g. 128) take the eager bf16 fallback,
+# so no rank-3 Fp4QuantizedTensor is ever produced and the smoke would be a trivial pass.
+FP4_SMOKE_AV_CONFIG = dict(
+    num_attention_heads=32,
+    attention_head_dim=64,  # video inner dim = 2048
+    in_channels=32,  # NVFP4 GEMM needs K (= in_channels patchify) divisible by 32
+    out_channels=32,
+    num_layers=1,
+    cross_attention_dim=2048,  # text cross-attn context_dim must equal inner_dim (32*64)
+    caption_channels=64,
+    norm_eps=1e-6,
+    positional_embedding_max_pos=[4, 32, 32],
+    timestep_scale_multiplier=1000,
+    use_middle_indices_grid=True,
+    audio_num_attention_heads=32,
+    audio_attention_head_dim=64,  # audio inner dim = 2048
+    audio_in_channels=32,
+    audio_out_channels=32,
+    audio_cross_attention_dim=2048,  # must equal audio inner_dim (32*64)
+    audio_positional_embedding_max_pos=[4],
+    av_ca_timestep_scale_multiplier=1,
+)
+
+
+class TestLTX2NVFP4ForwardSmoke(unittest.TestCase):
+    """NVFP4 forward smoke: a static-NVFP4 AudioVideo model must run end-to-end with
+    fabricated weights/inputs without raising.
+
+    Guards the rank-3 ``Fp4QuantizedTensor`` bug class (PR #15718): the fused
+    AdaLN+NVFP4 norm emits a rank-3 ``[B, S, D/2]`` fp4 activation that flows into the
+    fused MLP epilogue and the attention projections; consumers that ``reshape``/assert
+    2D crash on it. The forward-pre-hook below asserts the fp4 path is actually
+    exercised, so the test fails (rather than silently passing) if the model falls back
+    to bf16.
+
+    Requires Blackwell (NVFP4 fused kernels are SM100+) and a hidden dim in {2048, 4096}.
+    """
+
+    @pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0),
+        reason="NVFP4 fused kernels require Blackwell (SM100+)",
+    )
+    def test_audio_video_nvfp4_forward_smoke(self):
+        from types import SimpleNamespace
+
+        from tensorrt_llm._torch.utils import Fp4QuantizedTensor
+        from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality
+        from tensorrt_llm._torch.visual_gen.models.ltx2.transformer_ltx2 import (
+            LTXModel,
+            LTXModelType,
+        )
+        from tensorrt_llm.quantization.mode import QuantAlgo
+
+        torch.manual_seed(42)
+        device = "cuda"
+        dtype = torch.bfloat16
+
+        quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=16)
+        model_config = DiffusionModelConfig(
+            pretrained_config=SimpleNamespace(),
+            quant_config=quant_config,
+            mapping=Mapping(),
+            attention=AttentionConfig(backend="VANILLA"),
+            skip_create_weights_in_init=False,
+        )
+        # NVFP4 model: device-only .to() (a dtype cast would corrupt the packed fp4 weights).
+        model = (
+            LTXModel(
+                model_type=LTXModelType.AudioVideo,
+                model_config=model_config,
+                **FP4_SMOKE_AV_CONFIG,
+            )
+            .to(device)
+            .eval()
+        )
+        _init_weights_quant_safe(model)
+
+        batch = 1
+        v_frames, v_h, v_w = 1, 4, 4
+        v_patches = v_frames * v_h * v_w
+        a_patches = 8
+        in_channels = FP4_SMOKE_AV_CONFIG["in_channels"]
+        audio_in_channels = FP4_SMOKE_AV_CONFIG["audio_in_channels"]
+        caption_channels = FP4_SMOKE_AV_CONFIG["caption_channels"]
+        text_len = 8
+
+        v_context = (
+            torch.randn(batch, text_len, caption_channels, device=device, dtype=dtype) * 0.02
+        )
+        a_context = (
+            torch.randn(batch, text_len, caption_channels, device=device, dtype=dtype) * 0.02
+        )
+        v_positions = _make_video_positions(batch, v_patches, v_frames, v_h, v_w, device)
+        a_positions = _make_audio_positions(batch, a_patches, device)
+
+        video_modality = Modality(
+            latent=torch.randn(batch, v_patches, in_channels, device=device, dtype=dtype) * 0.02,
+            timesteps=torch.tensor([0.5], device=device),
+            positions=v_positions,
+            context=v_context,
+        )
+        audio_modality = Modality(
+            latent=torch.randn(batch, a_patches, audio_in_channels, device=device, dtype=dtype)
+            * 0.02,
+            timesteps=torch.tensor([0.5], device=device),
+            positions=a_positions,
+            context=a_context,
+        )
+        text_cache = model.prepare_text_cache(
+            video_context=v_context,
+            video_positions=v_positions,
+            audio_context=a_context,
+            audio_positions=a_positions,
+            dtype=dtype,
+        )
+
+        # Self-validation: record whether a rank-3 Fp4QuantizedTensor reaches any module,
+        # so a silent bf16 fallback (e.g. unsupported dim) fails the test instead of passing.
+        saw_rank3_fp4 = []
+
+        def _record(mod, args):
+            for a in args:
+                if isinstance(a, Fp4QuantizedTensor) and a.fp4_tensor.dim() > 2:
+                    saw_rank3_fp4.append(type(mod).__name__)
+
+        handles = [m.register_forward_pre_hook(_record) for m in model.modules()]
+        try:
+            with torch.no_grad():
+                video_out, audio_out = model(
+                    video=video_modality, audio=audio_modality, text_cache=text_cache
+                )
+        finally:
+            for h in handles:
+                h.remove()
+
+        self.assertEqual(video_out.shape, (batch, v_patches, FP4_SMOKE_AV_CONFIG["out_channels"]))
+        self.assertEqual(
+            audio_out.shape, (batch, a_patches, FP4_SMOKE_AV_CONFIG["audio_out_channels"])
+        )
+        self.assertTrue(
+            saw_rank3_fp4,
+            "no rank-3 Fp4QuantizedTensor reached any module -- the fused AdaLN+NVFP4 "
+            "path was not exercised (check hidden dim in {2048, 4096} and the NVFP4 build).",
+        )
+
+
 class TestLTX2QuantExcludeModuleRemapping(unittest.TestCase):
     """Tests for _remap_exclude_modules and _apply_quant_config_exclude_modules.