From 4254aeb56f280609e57e6b0134b3d6268d2fa87f Mon Sep 17 00:00:00 2001 From: Carl Y <4531192+carlyou@users.noreply.github.com> Date: Tue, 21 Apr 2026 23:29:58 -0700 Subject: [PATCH] [fix] flaky test_mla_attn_quant_fusion.py (#40530) Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com> --- tests/compile/passes/test_mla_attn_quant_fusion.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index 8a575909612..0a38ffca483 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -83,10 +83,6 @@ class MLAAttentionQuantPatternModel(torch.nn.Module): self.vllm_config = vllm_config self.dtype = vllm_config.model_config.dtype - # Create kv_b_proj (ColumnParallelLinear) on device. - # Reuse weights from prior model instance when available, because - # ColumnParallelLinear may get NaN from recycled CUDA memory after - # torch.compile runs in the same process. kv_b_proj = ColumnParallelLinear( input_size=kv_lora_rank, output_size=num_heads * (qk_nope_head_dim + v_head_dim), @@ -96,8 +92,7 @@ class MLAAttentionQuantPatternModel(torch.nn.Module): kv_b_proj_weight = kwargs.get("kv_b_proj_weight") if kv_b_proj_weight is not None: kv_b_proj.weight.data.copy_(kv_b_proj_weight) - elif kv_b_proj.weight.data.isnan().any(): - # Sanitize NaN from recycled CUDA memory + else: kv_b_proj.weight.data.normal_() # Create MLAAttention