@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 577df69b26

Fork 0

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 6, 10:45 UTC

[Bugfix] Fix KV scales inconsistency in fp8 MLA & FlashInfer kv_cache_dtype "auto" leading to gibberish (#37054)

Signed-off-by: Andy Lo <andy@mistral.ai>

Andy Lo committed Mar 18, 2026 at 23:07 UTC 577df69b26491aaa8f3fef2ea44d6ac256172032

8 files changed +71 -43

tests/v1/attention/test_mla_backends.py

+31 -28

@@ -266,22 +266,6 @@ def create_and_prepopulate_kv_cache(

266 return kv_cache

267

268

269 -class MockAttentionLayer:

270 - """A mock attention layer for testing."""

271 -

272 - def __init__(self, device: torch.device):

273 - self._q_scale = torch.tensor(1.0, device=device)

274 - self._k_scale = torch.tensor(1.0, device=device)

275 - self._v_scale = torch.tensor(1.0, device=device)

276 - self._prob_scale = torch.tensor(1.0, device=device)

277 - self._q_scale_float = 1.0

278 - self._k_scale_float = 1.0

279 - self._v_scale_float = 1.0

280 -

281 - def forward(self, *_args, **_kwargs):

282 - raise NotImplementedError

283 -

284 -

269 class MockSparseMLAAttentionLayer:

270 """A mock sparse MLA attention layer for testing.

271

@@ -304,6 +288,8 @@ class MockSparseMLAAttentionLayer:

288 device: torch.device,

289 W_UK: torch.Tensor,

290 W_UV: torch.Tensor,

291 + q_scale: float,

292 + k_scale: float,

293 ):

294 self.impl = impl

295 self.num_heads = num_heads

@@ -319,13 +305,13 @@ class MockSparseMLAAttentionLayer:

305 self.W_UV = W_UV.transpose(0, 1)

306

307 # Scale attributes needed by attention backends

322 - self._q_scale = torch.tensor(1.0, device=device)

323 - self._k_scale = torch.tensor(1.0, device=device)

324 - self._v_scale = torch.tensor(1.0, device=device)

308 + self._q_scale = torch.tensor(q_scale, device=device)

309 + self._k_scale = torch.tensor(k_scale, device=device)

310 + self._v_scale = torch.tensor(float("nan"), device=device)

311 self._prob_scale = torch.tensor(1.0, device=device)

326 - self._q_scale_float = 1.0

327 - self._k_scale_float = 1.0

328 - self._v_scale_float = 1.0

312 + self._q_scale_float = q_scale

313 + self._k_scale_float = k_scale

314 + self._v_scale_float = float("nan")

315

316 self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

317 static=True,

@@ -420,6 +406,8 @@ class MockMLAAttentionLayer(AttentionLayerBase):

406 kv_lora_rank: int,

407 device: torch.device,

408 kv_b_proj,

409 + q_scale: float,

410 + k_scale: float,

411 ):

412 self.impl = impl

413 self.num_heads = num_heads

@@ -443,13 +431,13 @@ class MockMLAAttentionLayer(AttentionLayerBase):

431 self.W_UK_T = W_UK.permute(1, 2, 0)

432

433 # Scale attributes needed by attention backends

446 - self._q_scale = torch.tensor(1.0, device=device)

447 - self._k_scale = torch.tensor(1.0, device=device)

448 - self._v_scale = torch.tensor(1.0, device=device)

434 + self._q_scale = torch.tensor(q_scale, device=device)

435 + self._k_scale = torch.tensor(k_scale, device=device)

436 + self._v_scale = torch.tensor(float("nan"), device=device)

437 self._prob_scale = torch.tensor(1.0, device=device)

450 - self._q_scale_float = 1.0

451 - self._k_scale_float = 1.0

452 - self._v_scale_float = 1.0

438 + self._q_scale_float = q_scale

439 + self._k_scale_float = k_scale

440 + self._v_scale_float = float("nan")

441

442 self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

443 static=True,

@@ -568,6 +556,8 @@ def run_attention_backend(

556 qk_rope_head_dim: int,

557 v_head_dim: int,

558 mock_kv_b_proj,

559 + q_scale: float,

560 + k_scale: float,

561 kv_cache_dtype: str = "auto",

562 ) -> torch.Tensor:

563 """Run attention computation using the specified backend's AttentionImpl."""

@@ -625,6 +615,8 @@ def run_attention_backend(

615 kv_lora_rank=kv_lora_rank,

616 device=device,

617 kv_b_proj=mock_kv_b_proj,

618 + q_scale=q_scale,

619 + k_scale=k_scale,

620 )

621

622 # Populate static_forward_context with mock attention layers

@@ -674,6 +666,7 @@ def run_attention_backend(

666 @pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])

667 @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])

668 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])

669 +@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])

670 def test_backend_correctness(

671 default_vllm_config,

672 dist_init,

@@ -681,6 +674,8 @@ def test_backend_correctness(

674 model: str,

675 tensor_parallel_size: int,

676 kv_cache_dtype: str,

677 + q_scale: float,

678 + k_scale: float,

679 ):

680 """

681 Test that all backends produce similar outputs to a reference implementation

@@ -709,6 +704,11 @@ def test_backend_correctness(

704 for b in BACKENDS_TO_TEST

705 if kv_cache_dtype in b.get_class().supported_kv_cache_dtypes

706 ]

707 + if (

708 + q_scale != 1.0 or k_scale != 1.0

709 + ) and AttentionBackendEnum.CUTLASS_MLA in backends_to_test:

710 + # CUTLASS_MLA does not support non-1 Q/K scales

711 + backends_to_test.remove(AttentionBackendEnum.CUTLASS_MLA)

712 if not backends_to_test:

713 pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")

714

@@ -1029,6 +1029,7 @@ def test_backend_correctness(

1029 common_attn_metadata=common_attn_metadata,

1030 randomize_blocks=True,

1031 kv_cache_dtype=kv_cache_dtype,

1032 + scale=k_scale,

1033 )

1034 kv_cache_per_block_size[block_size] = kv_cache

1035

@@ -1072,6 +1073,8 @@ def test_backend_correctness(

1073 qk_rope_head_dim,

1074 v_head_dim,

1075 mock_kv_b_proj,

1076 + q_scale=q_scale,

1077 + k_scale=k_scale,

1078 kv_cache_dtype=kv_cache_dtype,

1079 )

1080

tests/v1/attention/test_sparse_mla_backends.py

+9 -2

@@ -178,6 +178,7 @@ def _quantize_dequantize_fp8_ds_mla(

178 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])

179 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])

180 @pytest.mark.parametrize("block_size", [32, 64])

181 +@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])

182 def test_sparse_backend_decode_correctness(

183 default_vllm_config,

184 dist_init,

@@ -187,6 +188,8 @@ def test_sparse_backend_decode_correctness(

188 tensor_parallel_size,

189 block_size,

190 workspace_init,

191 + q_scale: float,

192 + k_scale: float,

193 ):

194 if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:

195 pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")

@@ -332,7 +335,7 @@ def test_sparse_backend_decode_correctness(

335 kv_c_contexts, k_pe_contexts = [], []

336 reference_outputs = []

337

335 - kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

338 + kv_cache_scale = torch.tensor(k_scale, dtype=torch.float32, device=device)

339 global_token_idx = 0

340

341 for i in range(batch_spec.batch_size):

@@ -490,6 +493,8 @@ def test_sparse_backend_decode_correctness(

493 device=device,

494 W_UK=W_UK,

495 W_UV=W_UV,

496 + q_scale=q_scale,

497 + k_scale=k_scale,

498 )

499

500 out_buffer = torch.empty(

@@ -513,7 +518,9 @@ def test_sparse_backend_decode_correctness(

518 # FP8 quantization introduces some error, but should be within reasonable bounds

519 # BF16 (auto) should be very accurate, FP8 allows slightly more tolerance

520 if kv_cache_dtype.startswith("fp8"):

516 - torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05)

521 + torch.testing.assert_close(

522 + backend_output, sdpa_reference, rtol=0.065, atol=0.05

523 + )

524 else:

525 torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)

526

tests/v1/attention/test_trtllm_attention_integration.py

+6 -6

@@ -43,12 +43,12 @@ class MockAttentionLayer:

43 """Minimal mock of an attention layer for testing."""

45 def __init__(self, device: torch.device):

46 - self._q_scale = torch.tensor(1.0, device=device)

47 - self._k_scale = torch.tensor(1.0, device=device)

48 - self._v_scale = torch.tensor(1.0, device=device)

49 - self._q_scale_float = 1.0

50 - self._k_scale_float = 1.0

51 - self._v_scale_float = 1.0

46 + self._q_scale = torch.tensor(2.0, device=device)

47 + self._k_scale = torch.tensor(3.0, device=device)

48 + self._v_scale = torch.tensor(4.0, device=device)

49 + self._q_scale_float = 2.0

50 + self._k_scale_float = 3.0

51 + self._v_scale_float = 4.0

52 self._o_scale_float = None

vllm/v1/attention/backends/flashinfer.py

+6 -2

@@ -1319,10 +1319,14 @@ class FlashInferImpl(AttentionImpl):

1319 )

1320

1321 if self.bmm1_scale is None:

1322 - self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale

1322 + self.bmm1_scale = self.scale

1323 + if self.kv_cache_dtype.startswith("fp8"):

1324 + self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

1325

1326 if self.bmm2_scale is None:

1325 - self.bmm2_scale = layer._v_scale_float

1327 + self.bmm2_scale = 1.0

1328 + if self.kv_cache_dtype.startswith("fp8"):

1329 + self.bmm2_scale *= layer._v_scale_float

1330

1331 prefill_use_trtllm = isinstance(attn_metadata.prefill, TRTLLMPrefill)

1332 decode_use_trtllm = isinstance(attn_metadata.decode, TRTLLMDecode)

vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -255,6 +255,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

255 assert kv_c_and_k_pe_cache.numel() > 0

256 assert attn_metadata.decode is not None

257

258 + if layer._q_scale_float != 1.0 or layer._k_scale_float != 1.0:

259 + raise NotImplementedError(

260 + "CutlassMLAImpl does not support scaling for q and kv_latent yet"

261 + )

262 +

263 if type(q) is tuple:

264 q_nope, q_pe = q

265 else:

vllm/v1/attention/backends/mla/flashinfer_mla.py

+7 -2

@@ -177,9 +177,14 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

177 q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])

178

179 if self.bmm1_scale is None:

180 - self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale

180 + self.bmm1_scale = self.scale

181 + if self.kv_cache_dtype.startswith("fp8"):

182 + self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

183 +

184 if self.bmm2_scale is None:

182 - self.bmm2_scale = layer._v_scale_float

185 + self.bmm2_scale = 1.0

186 + if self.kv_cache_dtype.startswith("fp8"):

187 + self.bmm2_scale *= layer._k_scale_float

188

189 o = trtllm_batch_decode_with_kv_cache_mla(

190 query=q,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

+6 -2

@@ -340,9 +340,13 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

340 self._workspace_buffer = _get_workspace_buffer(q.device)

341

342 if self.bmm1_scale is None:

343 - self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale

343 + self.bmm1_scale = self.scale

344 + if self.kv_cache_dtype.startswith("fp8"):

345 + self.bmm1_scale *= layer._q_scale_float * layer._k_scale_float

346 if self.bmm2_scale is None:

345 - self.bmm2_scale = layer._v_scale_float

347 + self.bmm2_scale = 1.0

348 + if self.kv_cache_dtype.startswith("fp8"):

349 + self.bmm2_scale *= layer._k_scale_float

350

351 o = trtllm_batch_decode_with_kv_cache_mla(

352 query=q.unsqueeze(1),

vllm/v1/attention/backends/mla/triton_mla.py

+1 -1

@@ -187,7 +187,7 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

187 self.scale,

188 PAGE_SIZE,

189 k_scale=layer._k_scale,

190 - v_scale=layer._v_scale,

190 + v_scale=layer._k_scale,

191 )

192

193 return o, lse