diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py index 2373a31ff..0b19d9f34 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/server/accuracy_config.py b/python/sglang/multimodal_gen/test/server/accuracy_config.py index 955c84ad1..46b20a23a 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_config.py +++ b/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}, diff --git a/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py b/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py index 83268b989..ec8654b3a 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index a0adacefa..46e6cd0a2 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/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, - "ssim_threshold": 0.52, - "psnr_threshold": 9.3, - "mean_abs_diff_threshold": 68.0 + "clip_threshold": 0.94, + "ssim_threshold": 0.78, + "psnr_threshold": 17.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, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 75a7857c1..93ad2f55f 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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( diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 9772146ae..dacfe4508 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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,