[diffusion] fix: align encoder of flux klein with official (#24008)

python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py

@@ -158,6 +158,7 @@ class Qwen3Attention(nn.Module):
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
+        attention_lengths: tuple[int, ...] | None = None,
     ) -> torch.Tensor:
         # QKV projection
         qkv, _ = self.qkv_proj(hidden_states)
@@ -185,13 +186,61 @@ class Qwen3Attention(nn.Module):
         k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
 
         # Attention
-        attn_output = self.attn(q, k, v)
+        attn_output = self._masked_causal_attention(q, k, v, attention_lengths)
         attn_output = attn_output.reshape(batch_size, seq_len, -1)
 
         # Output projection
         output, _ = self.o_proj(attn_output)
         return output
 
+    def _masked_causal_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_lengths: tuple[int, ...] | None,
+    ) -> torch.Tensor:
+        if attention_lengths is None:
+            return self.attn(q, k, v)
+
+        seq_len = q.shape[1]
+        if all(valid_len == seq_len for valid_len in attention_lengths):
+            return self.attn(q, k, v)
+
+        outputs: list[torch.Tensor] = []
+        for batch_index, valid_len in enumerate(attention_lengths):
+            q_item = q[batch_index : batch_index + 1]
+            k_item = k[batch_index : batch_index + 1]
+            v_item = v[batch_index : batch_index + 1]
+
+            real_output = self.attn(
+                q_item[:, :valid_len],
+                k_item[:, :valid_len],
+                v_item[:, :valid_len],
+            )
+            if valid_len == seq_len:
+                outputs.append(real_output)
+                continue
+
+            pad_q = q_item[:, valid_len:].transpose(1, 2)
+            real_k = k_item[:, :valid_len].transpose(1, 2)
+            real_v = v_item[:, :valid_len].transpose(1, 2)
+            if self.num_heads != self.num_kv_heads:
+                repeat_factor = self.num_heads // self.num_kv_heads
+                real_k = real_k.repeat_interleave(repeat_factor, dim=1)
+                real_v = real_v.repeat_interleave(repeat_factor, dim=1)
+            pad_output = torch.nn.functional.scaled_dot_product_attention(
+                pad_q,
+                real_k,
+                real_v,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+            ).transpose(1, 2)
+            outputs.append(torch.cat([real_output, pad_output], dim=1))
+
+        return torch.cat(outputs, dim=0)
+
 
 class Qwen3DecoderLayer(nn.Module):
     """Qwen3 transformer decoder layer."""
@@ -241,6 +290,7 @@ class Qwen3DecoderLayer(nn.Module):
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
         residual: torch.Tensor | None,
+        attention_lengths: tuple[int, ...] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Self Attention
         if residual is None:
@@ -249,7 +299,11 @@ class Qwen3DecoderLayer(nn.Module):
         else:
             hidden_states, residual = self.input_layernorm(hidden_states, residual)
 
-        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            attention_lengths=attention_lengths,
+        )
 
         # MLP
         hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
@@ -335,6 +389,13 @@ class Qwen3ForCausalLM(TextEncoder):
                 0, hidden_states.shape[1], device=hidden_states.device
             ).unsqueeze(0)
 
+        attention_lengths = None
+        if attention_mask is not None:
+            attention_lengths = tuple(
+                int(valid_len)
+                for valid_len in attention_mask.sum(dim=-1).detach().cpu().tolist()
+            )
+
         all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
 
         for layer in self.layers:
@@ -344,7 +405,9 @@ class Qwen3ForCausalLM(TextEncoder):
                     if residual is None
                     else (hidden_states + residual,)
                 )
-            hidden_states, residual = layer(position_ids, hidden_states, residual)
+            hidden_states, residual = layer(
+                position_ids, hidden_states, residual, attention_lengths
+            )
 
         hidden_states, _ = self.norm(hidden_states, residual)
 

python/sglang/multimodal_gen/test/server/accuracy_config.py

@@ -42,7 +42,6 @@ CASE_THRESHOLDS: Dict[str, Dict[ComponentType, float]] = {
     "flux_2_image_t2i": {ComponentType.TRANSFORMER: 0.99},
     "flux_2_image_t2i_layerwise_offload": {ComponentType.TRANSFORMER: 0.99},
     "flux_2_image_t2i_2_gpus": {ComponentType.TRANSFORMER: 0.99},
-    "flux_2_klein_ti2i_2_gpus": {ComponentType.TRANSFORMER: 0.975},
     "flux_2_ti2i": {ComponentType.TRANSFORMER: 0.99},
     "flux_2_t2i_customized_vae_path": {ComponentType.TRANSFORMER: 0.99},
     "fast_hunyuan_video": {ComponentType.TRANSFORMER: 0.99},

python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py

@@ -66,7 +66,6 @@ ACCURACY_TWO_GPU_CASE_IDS = (
     "zimage_image_t2i_2_gpus_non_square",
     "flux_image_t2i_2_gpus",
     "flux_2_image_t2i_2_gpus",
-    "flux_2_klein_ti2i_2_gpus",
 )
 
 ACCURACY_ONE_GPU_CASES = _select_accuracy_cases(

python/sglang/multimodal_gen/test/server/consistency_threshold.json

@@ -14,10 +14,10 @@
             "mean_abs_diff_threshold": 8.0
         },
         "flux_2_klein_image_t2i": {
-            "clip_threshold": 0.83,
+            "clip_threshold": 0.94,
-            "ssim_threshold": 0.52,
+            "ssim_threshold": 0.78,
-            "psnr_threshold": 9.3,
+            "psnr_threshold": 17.0,
-            "mean_abs_diff_threshold": 68.0
+            "mean_abs_diff_threshold": 17.0
         },
         "zimage_image_t2i": {
             "clip_threshold": 0.92,
@@ -193,12 +193,6 @@
             "psnr_threshold": 18.7,
             "mean_abs_diff_threshold": 8.0
         },
-        "flux_2_klein_ti2i_2_gpus": {
-            "clip_threshold": 0.92,
-            "ssim_threshold": 0.77,
-            "psnr_threshold": 18.4,
-            "mean_abs_diff_threshold": 18.0
-        },
         "wan2_2_t2v_a14b_teacache_2gpu": {
             "clip_threshold": 0.90,
             "ssim_threshold": 0.72,

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -616,13 +616,6 @@ TWO_GPU_CASES = [
         ),
         T2I_sampling_params,
     ),
-    DiffusionTestCase(
-        "flux_2_klein_ti2i_2_gpus",
-        DiffusionServerArgs(
-            model_path="black-forest-labs/FLUX.2-klein-4B",
-        ),
-        TI2I_sampling_params,
-    ),
     DiffusionTestCase(
         "ltx_2.3_one_stage_ti2v",
         DiffusionServerArgs(

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -2119,27 +2119,6 @@
             "expected_median_denoise_ms": 151.72,
             "estimated_full_test_time_s": 129.3
         },
-        "flux_2_klein_ti2i_2_gpus": {
-            "stages_ms": {
-                "InputValidationStage": 35.39,
-                "DecodingStage": 6.68,
-                "TextEncodingStage": 160.46,
-                "TimestepPreparationStage": 27.82,
-                "LatentPreparationStage": 0.4,
-                "DenoisingStage": 364.99,
-                "ImageVAEEncodingStage": 65.69
-            },
-            "denoise_step_ms": {
-                "0": 30.19,
-                "1": 63.04,
-                "2": 89.37,
-                "3": 89.7
-            },
-            "expected_e2e_ms": 827.17,
-            "expected_avg_denoise_ms": 68.93,
-            "expected_median_denoise_ms": 77.89,
-            "estimated_full_test_time_s": 120.8
-        },
         "flux_2_image_t2i_upscaling_4x": {
             "stages_ms": {
                 "TextEncodingStage": 494.65,