A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 20:00 UTC
[fix] flaky test_mla_attn_quant_fusion.py (#40530)
Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com>
Carl Y committed
Apr 21, 2026 at 23:29 UTC
4254aeb56f280609e57e6b0134b3d6268d2fa87f
1 file changed
+1
-6
tests/compile/passes/test_mla_attn_quant_fusion.py
+1
-6
@@ -83,10 +83,6 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
83
self.vllm_config = vllm_config
84
self.dtype = vllm_config.model_config.dtype
85
86
- # Create kv_b_proj (ColumnParallelLinear) on device.
87
- # Reuse weights from prior model instance when available, because
88
- # ColumnParallelLinear may get NaN from recycled CUDA memory after
89
- # torch.compile runs in the same process.
86
kv_b_proj = ColumnParallelLinear(
87
input_size=kv_lora_rank,
88
output_size=num_heads * (qk_nope_head_dim + v_head_dim),
@@ -96,8 +92,7 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
92
kv_b_proj_weight = kwargs.get("kv_b_proj_weight")
93
if kv_b_proj_weight is not None:
94
kv_b_proj.weight.data.copy_(kv_b_proj_weight)
99
- elif kv_b_proj.weight.data.isnan().any():
100
- # Sanitize NaN from recycled CUDA memory
95
+ else:
96
kv_b_proj.weight.data.normal_()
97
98
# Create MLAAttention