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 9, 23:00 UTC
[BUGFIX] Fix `test_mla_backends.py`. Scale MLA projection weights to prevent numerical instability (#32529)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Jan 19, 2026 at 22:49 UTC
0727cc9ecf2ee511f952f5eed0f749e0e92cf57a
1 file changed
+8
tests/v1/attention/test_mla_backends.py
+8
@@ -504,6 +504,14 @@ def test_backend_correctness(
504
W_UV = torch.randn(
505
kv_lora_rank, num_q_heads, v_head_dim, dtype=dtype, device=device
506
)
507
+
508
+ # Scale weights to produce realistic magnitude outputs.
509
+ # Without scaling, projection output has std ~sqrt(kv_lora_rank) ≈ 22.6,
510
+ # causing extreme attention scores and numerical instability in LSE merging.
511
+ weight_scale = 1.0 / (kv_lora_rank**0.5)
512
+ W_UK = W_UK * weight_scale
513
+ W_UV = W_UV * weight_scale
514
+
515
kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
516
517
for i, backend in enumerate(BACKENDS_TO_TEST):