f24b2de3d3 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -19,8 +19,13 @@ from tests.v1.attention.utils import (

19 )

20 from vllm import _custom_ops as ops

21 from vllm.config.vllm import set_current_vllm_config

22 -from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport

22 +from vllm.model_executor.layers.attention.mla_attention import (

23 + QueryLenSupport,

24 + _DecodeConcatQuantFP8,

25 +)

26 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

27 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

28 +from vllm.platforms import current_platform

29 from vllm.utils.math_utils import cdiv

30 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE

31 from vllm.v1.attention.backend import CommonAttentionMetadata

@@ -50,6 +55,7 @@ if not flash_attn_supports_mla():

55 if not is_flashmla_dense_supported()[0]:

56 BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)

57

58 +

59 SPEC_DECODE_BACKENDS = []

60 for backend in BACKENDS_TO_TEST:

61 builder_cls, _ = try_get_attention_backend(backend)

@@ -144,9 +150,8 @@ def create_and_prepopulate_kv_cache(

150 common_attn_metadata: Common attention metadata

151 randomize_blocks: Whether to randomly permute blocks

152 or use sequential order

147 - kv_cache_dtype: Optional kv cache dtype string. When set to

148 - "fp8_ds_mla" the cache is populated using the

149 - fp8 DeepSeek MLA layout via concat_and_cache_mla.

153 + kv_cache_dtype: Optional kv cache dtype string. For fp8 cache dtype,

154 + the cache is populated via concat_and_cache_mla.

155 scale: Scaling factor forwarded to concat_and_cache_mla when the

156 fp8 cache layout is requested.

157

@@ -163,18 +168,21 @@ def create_and_prepopulate_kv_cache(

168 block_table = common_attn_metadata.block_table_tensor

169 slot_mapping = common_attn_metadata.slot_mapping

170

171 + fp8_attention = kv_cache_dtype and kv_cache_dtype.startswith("fp8")

172 use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla"

173

168 - if use_fp8_ds_mla:

169 - if not kv_c_contexts:

170 - raise ValueError(

171 - "kv_c_contexts cannot be empty when using fp8_ds_mla cache dtype"

172 - )

173 - kv_lora_rank = kv_c_contexts[0].shape[-1]

174 - rope_dim = k_pe_contexts[0].shape[-1]

175 - entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim

174 + if fp8_attention:

175 + if use_fp8_ds_mla:

176 + kv_lora_rank = kv_c_contexts[0].shape[-1]

177 + rope_dim = k_pe_contexts[0].shape[-1]

178 + # 4 * 4: 4 float32 scale values for 128-element tiles

179 + # 2 * rope_dim: 16-bit RoPE values

180 + kv_entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim

181 + else:

182 + kv_entry_size = head_size

183 +

184 kv_cache = torch.zeros(

177 - num_blocks, block_size, entry_size, dtype=torch.uint8, device=device

185 + num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device

186 )

187 scale_tensor = (

188 scale

@@ -201,14 +209,14 @@ def create_and_prepopulate_kv_cache(

209

210 start = start_block_idx * block_size

211

204 - if use_fp8_ds_mla:

212 + if fp8_attention:

213 slots = torch.arange(context_len, device=device, dtype=torch.long) + start

214 ops.concat_and_cache_mla(

215 kv_c_context,

216 k_pe_context.squeeze(1),

217 kv_cache,

218 slots,

211 - kv_cache_dtype="fp8_ds_mla",

219 + kv_cache_dtype=kv_cache_dtype,

220 scale=scale_tensor,

221 )

222 else:

@@ -329,8 +337,9 @@ class MockSparseMLAAttentionLayer:

337 output: torch.Tensor,

338 ) -> torch.Tensor:

339 """Forward for sparse MLA - uses forward_mqa for all tokens."""

332 - # Write to KV cache

340 kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")

341 +

342 + # Write to KV cache

343 if kv_cache.numel() > 0:

344 ops.concat_and_cache_mla(

345 kv_c,

@@ -426,6 +435,12 @@ class MockMLAAttentionLayer(AttentionLayerBase):

435 self._k_scale_float = 1.0

436 self._v_scale_float = 1.0

437

438 + self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

439 + static=True,

440 + group_shape=GroupShape.PER_TENSOR,

441 + compile_native=True,

442 + )

443 +

444 def get_attn_backend(self):

445 raise NotImplementedError

446

@@ -443,16 +458,21 @@ class MockMLAAttentionLayer(AttentionLayerBase):

458 ) -> torch.Tensor:

459 """Replicates MLAAttention.forward_impl logic for testing."""

460 # Write to KV cache

461 + kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")

462 + fp8_attention = kv_cache_dtype.startswith("fp8")

463 if kv_cache.numel() > 0:

464 ops.concat_and_cache_mla(

465 kv_c,

466 k_pe.squeeze(1),

467 kv_cache,

468 attn_metadata.slot_mapping.flatten(),

452 - kv_cache_dtype="auto",

469 + kv_cache_dtype=kv_cache_dtype,

470 scale=self._k_scale,

471 )

472

473 + if fp8_attention and kv_cache_dtype != "fp8_ds_mla":

474 + kv_cache = kv_cache.view(current_platform.fp8_dtype())

475 +

476 # Determine decode vs prefill split

477 num_decode_tokens = attn_metadata.num_decode_tokens or 0

478 has_decode = (attn_metadata.num_decodes or 0) > 0

@@ -491,8 +511,14 @@ class MockMLAAttentionLayer(AttentionLayerBase):

511 # Convert from (N, B, L) to (B, N, L)

512 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

513

494 - # Pass as tuple to forward_mqa

495 - mqa_q = (mqa_ql_nope, mqa_q_pe)

514 + if fp8_attention and self.impl.supports_quant_query_input:

515 + assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]

516 + assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]

517 + mqa_q = self._decode_concat_quant_fp8_op(

518 + mqa_ql_nope, mqa_q_pe, self._q_scale

519 + )

520 + else:

521 + mqa_q = (mqa_ql_nope, mqa_q_pe)

522

523 attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)

524

@@ -526,6 +552,7 @@ def run_attention_backend(

552 qk_rope_head_dim: int,

553 v_head_dim: int,

554 mock_kv_b_proj,

555 + kv_cache_dtype: str = "auto",

556 ) -> torch.Tensor:

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

558

@@ -550,7 +577,7 @@ def run_attention_backend(

577 num_kv_heads=num_kv_heads,

578 alibi_slopes=None,

579 sliding_window=None,

553 - kv_cache_dtype="auto",

580 + kv_cache_dtype=kv_cache_dtype,

581 logits_soft_cap=None,

582 attn_type="decoder",

583 kv_sharing_target_layer_name=None,

@@ -630,12 +657,14 @@ def run_attention_backend(

657 )

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

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

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

661 def test_backend_correctness(

662 default_vllm_config,

663 dist_init,

664 batch_spec_name: str,

665 model: str,

666 tensor_parallel_size: int,

667 + kv_cache_dtype: str,

668 ):

669 """

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

@@ -658,9 +687,18 @@ def test_backend_correctness(

687 head counts.

688 """

689

690 + # Filter backends to those that support the requested kv_cache_dtype

691 + backends_to_test = [

692 + b

693 + for b in BACKENDS_TO_TEST

694 + if kv_cache_dtype in b.get_class().supported_kv_cache_dtypes

695 + ]

696 + if not backends_to_test:

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

698 +

699 batch_spec = BATCH_SPECS[batch_spec_name]

700 is_spec_decode_test = batch_spec_name.startswith("spec_decode")

663 - unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES.values()))

701 + unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))

702 default_block_size = unique_block_sizes[0]

703 required_blocks = sum(

704 (seq_len + default_block_size - 1) // default_block_size

@@ -694,6 +732,7 @@ def test_backend_correctness(

732 block_size=default_block_size,

733 hf_config_override=hf_config_override,

734 )

735 + vllm_config.cache_config.cache_dtype = kv_cache_dtype

736

737 # For spec decode tests, add a speculative_config to set the reorder_batch_threshold

738 if is_spec_decode_test:

@@ -751,7 +790,7 @@ def test_backend_correctness(

790

791 kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)

792

754 - for i, backend in enumerate(BACKENDS_TO_TEST):

793 + for i, backend in enumerate(backends_to_test):

794 all_sdpa_outputs.append([])

795

796 for i in range(batch_size):

@@ -785,7 +824,7 @@ def test_backend_correctness(

824 # pipeline (MHA-style). This ensures the reference implementation

825 # matches each backend's actual decode/prefill pipeline path.

826 is_decode = []

788 - for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

827 + for backend_idx, backend in enumerate(backends_to_test):

828 builder_cls, _ = try_get_attention_backend(backend)

829 if is_spec_decode_test:

830 query_len_support = getattr(

@@ -885,7 +924,7 @@ def test_backend_correctness(

924 sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)

925 sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)

926

888 - for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

927 + for backend_idx, backend in enumerate(backends_to_test):

928 if is_decode[backend_idx]:

929 all_sdpa_outputs[backend_idx].append(sdpa_out_i_decode)

930 else:

@@ -905,7 +944,7 @@ def test_backend_correctness(

944 kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)

945 k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)

946 sdpa_outputs = {}

908 - for backend_idx, backend in enumerate(BACKENDS_TO_TEST):

947 + for backend_idx, backend in enumerate(backends_to_test):

948 sdpa_outputs[backend] = torch.cat(all_sdpa_outputs[backend_idx], dim=0)

949

950 # Create mock kv_b_proj using the same weights as reference implementation

@@ -973,12 +1012,13 @@ def test_backend_correctness(

1012 num_blocks=num_blocks_for_size,

1013 common_attn_metadata=common_attn_metadata,

1014 randomize_blocks=True,

1015 + kv_cache_dtype=kv_cache_dtype,

1016 )

1017 kv_cache_per_block_size[block_size] = kv_cache

1018

1019 # 4. Run vLLM backends and compare

1020 failures = []

981 - for backend_idx, backend_name in enumerate(BACKENDS_TO_TEST):

1021 + for backend_idx, backend_name in enumerate(backends_to_test):

1022 # Skip backends that don't support spec decode for spec decode tests

1023 if is_spec_decode_test and backend_name not in SPEC_DECODE_BACKENDS:

1024 continue

@@ -997,7 +1037,7 @@ def test_backend_correctness(

1037 head_size=vllm_config.model_config.get_head_size(),

1038 dtype=vllm_config.model_config.dtype,

1039 sliding_window=vllm_config.model_config.get_sliding_window(),

1000 - cache_dtype_str=vllm_config.cache_config.cache_dtype,

1040 + cache_dtype_str=kv_cache_dtype,

1041 )

1042

1043 backend_output = run_attention_backend(

@@ -1016,6 +1056,7 @@ def test_backend_correctness(

1056 qk_rope_head_dim,

1057 v_head_dim,

1058 mock_kv_b_proj,

1059 + kv_cache_dtype=kv_cache_dtype,

1060 )

1061

1062 # Use backend_idx to get the correct SDPA output for this backend

[Test] Add FP8 KV Cache Testing for MLA Backends (#34473)