[None][fix] visual_gen FLUX: enable fused DiT QK-norm + RoPE by default

tensorrt_llm/_torch/visual_gen/models/flux/attention.py

@@ -41,9 +41,11 @@ class FluxJointAttention(Attention):
     - FLUX-style RoPE on concatenated text+image tokens
     - pre_only mode for single-stream blocks (no output projection)
 
-    When fuse_qk_norm_rope is enabled (default), the fused CUDA kernel handles
-    QK norm + RoPE in a single pass. When disabled, falls back to separate
-    F.rms_norm + apply_rotary_emb calls.
+    FLUX enables fuse_qk_norm_rope by default when TP=1: the fused CUDA
+    kernel handles QK norm + RoPE in a single pass. It falls back to separate
+    F.rms_norm + apply_rotary_emb calls when disabled, and is disabled under
+    tensor parallelism because the fused op currently requires TP=1
+    (apply_packed_qk_norm_rope asserts tp_size == 1).
     """
 
     def __init__(
@@ -59,6 +61,11 @@ def __init__(
         layer_idx: int = 0,
         module_name: Optional[str] = None,
     ):
+        # Opt in to the fused DiT QK-norm + RoPE kernel (per-head template), but
+        # only when TP=1: the fused op asserts tp_size == 1
+        # (apply_packed_qk_norm_rope), so under TP>1 we fall back to the unfused
+        # F.rms_norm + apply_rotary_emb path. Mirrors WAN's gating.
+        tp_size = config.mapping.tp_size if config and config.mapping else 1
         super().__init__(
             hidden_size=hidden_size,
             num_attention_heads=num_attention_heads,
@@ -69,6 +76,7 @@ def __init__(
             eps=eps,
             bias=bias,
             interleave=True,
+            fuse_qk_norm_rope=(tp_size == 1),
             config=config,
             layer_idx=layer_idx,
             module_name=module_name,

tests/unittest/_torch/visual_gen/test_flux_attention.py

@@ -7,6 +7,7 @@
 - VANILLA backend (PyTorch SDPA)
 - TRTLLM backend (fmha_v2)
 - Backend equivalence comparison
+- Fused QK-norm + RoPE TP gating
 
 Note: With random weights, attention can produce NaN due to numerical instability.
       These tests use scaled inputs and primarily verify correct output shapes.
@@ -34,15 +35,29 @@ class TestFluxAttentionBackend(unittest.TestCase):
 
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-    def _create_config(self, backend: str) -> DiffusionModelConfig:
-        """Create DiffusionModelConfig with specified backend."""
+    def _create_config(self, backend: str, tp_size: int = 1) -> DiffusionModelConfig:
+        """Create DiffusionModelConfig with specified backend and TP size."""
         return DiffusionModelConfig(
             pretrained_config=SimpleNamespace(),
             quant_config=QuantConfig(),
-            mapping=Mapping(),
+            mapping=Mapping(world_size=tp_size, tp_size=tp_size),
             attention=AttentionConfig(backend=backend),
         )
 
+    def test_fused_qk_norm_rope_enabled_only_for_tp1(self) -> None:
+        """Test FLUX enables fused QK-norm + RoPE only when TP=1."""
+        from tensorrt_llm._torch.visual_gen.models.flux.attention import FluxJointAttention
+
+        for tp_size, expected in ((1, True), (2, False)):
+            with self.subTest(tp_size=tp_size):
+                attn = FluxJointAttention(
+                    hidden_size=128,
+                    num_attention_heads=2,
+                    head_dim=64,
+                    config=self._create_config("VANILLA", tp_size=tp_size),
+                )
+                self.assertEqual(attn.fuse_qk_norm_rope, expected)
+
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_vanilla_backend_sanity(self):
         """Test FLUX attention works with VANILLA backend."""