| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """K3's fused chunked-context prefill must match the generic MLA impl. |
| 4 | |
| 5 | The layer owns its own context loop so it can fuse the per-chunk K/V pack and |
| 6 | skip re-quantizing an already-quantized query. That is only safe if it feeds the |
| 7 | prefill backend exactly what ``MLACommonBaseImpl._compute_prefill_context`` |
| 8 | would, chunk for chunk. |
| 9 | """ |
| 10 | |
| 11 | from types import SimpleNamespace |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 17 | MLACommonBaseImpl, |
| 18 | MLACommonPrefillMetadata, |
| 19 | build_mla_chunked_context_metadata, |
| 20 | ) |
| 21 | from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention |
| 22 | from vllm.platforms import current_platform |
| 23 | |
| 24 | pytestmark = pytest.mark.skipif( |
| 25 | not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA" |
| 26 | ) |
| 27 | |
| 28 | _KV_LORA_RANK = 512 |
| 29 | _QK_NOPE = 128 |
| 30 | _QK_ROPE = 64 |
| 31 | _V_HEAD_DIM = 128 |
| 32 | _ENTRY = _KV_LORA_RANK + _QK_ROPE |
| 33 | _NUM_HEADS = 2 |
| 34 | _BLOCK_SIZE = 16 |
| 35 | _WORKSPACE_TOKENS = 128 |
| 36 | # Splits one long request across chunks, packs short ones together, and leaves |
| 37 | # the last request without any context. |
| 38 | _CONTEXT_LENS = [200, 48, 32, 0] |
| 39 | _QUERY_LENS = [8, 4, 6, 5] |
| 40 | |
| 41 | |
| 42 | class _RecordingPrefillBackend: |
| 43 | """Records what each chunk is asked to attend over. |
| 44 | |
| 45 | ``honors_out`` mimics the two backend families: one that writes into a |
| 46 | caller-provided ``out`` (trtllm_ragged, flashinfer, tokenspeed) and one that |
| 47 | always returns its own buffer (flash_attn with a padded V, aiter). |
| 48 | """ |
| 49 | |
| 50 | def __init__(self, honors_out: bool = False) -> None: |
| 51 | self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] |
| 52 | self.out_destinations: list[torch.Tensor | None] = [] |
| 53 | self._honors_out = honors_out |
| 54 | |
| 55 | @staticmethod |
| 56 | def get_name() -> str: |
| 57 | return "recording" |
| 58 | |
| 59 | def supports_out(self) -> bool: |
| 60 | return self._honors_out |
| 61 | |
| 62 | def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None): |
| 63 | self.calls.append((q.float().clone(), k.float().clone(), v.float().clone())) |
| 64 | self.out_destinations.append(out) |
| 65 | assert out is None or self._honors_out |
| 66 | # Fold K/V into the partial so any packing difference shows up in the |
| 67 | # merged context output, not just in the recorded calls. |
| 68 | digest = (k.float().mean() + v.float().mean()).item() |
| 69 | num_q = q.shape[0] |
| 70 | if out is None: |
| 71 | out = torch.empty( |
| 72 | (num_q, _NUM_HEADS, _V_HEAD_DIM), device=q.device, dtype=torch.bfloat16 |
| 73 | ) |
| 74 | else: |
| 75 | assert out.shape == (num_q, _NUM_HEADS, _V_HEAD_DIM) |
| 76 | out.fill_(digest) |
| 77 | lse = torch.full( |
| 78 | (_NUM_HEADS, num_q), |
| 79 | 1.0 + chunk.index, |
| 80 | device=q.device, |
| 81 | dtype=torch.float32, |
| 82 | ) |
| 83 | return out, lse |
| 84 | |
| 85 | |
| 86 | class _KVBProj(torch.nn.Module): |
| 87 | """Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple). |
| 88 | |
| 89 | Enforces the same input contract as the real linear methods, because that is |
| 90 | what decides whether the gathered latent needs a cast: |
| 91 | |
| 92 | * an fp8 weight consumes the fp8 latent directly and dequantizes internally, |
| 93 | so it takes fp8 or bf16; |
| 94 | * a bf16 weight -- what a stock K3 checkpoint carries -- is a plain |
| 95 | ``F.linear`` and rejects anything but bf16, exactly as torch does. |
| 96 | |
| 97 | Either weight dtype produces a bf16 output. |
| 98 | """ |
| 99 | |
| 100 | def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: |
| 101 | super().__init__() |
| 102 | weight = ( |
| 103 | torch.randn( |
| 104 | _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM), |
| 105 | _KV_LORA_RANK, |
| 106 | device=device, |
| 107 | dtype=torch.bfloat16, |
| 108 | ) |
| 109 | * 0.05 |
| 110 | ) |
| 111 | self.register_buffer("weight", weight.to(weight_dtype)) |
| 112 | |
| 113 | def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: |
| 114 | if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16: |
| 115 | raise RuntimeError( |
| 116 | "a bfloat16 kv_b_proj cannot consume the gathered latent as " |
| 117 | f"{x.dtype}; it must be cast first" |
| 118 | ) |
| 119 | return torch.nn.functional.linear( |
| 120 | x.to(torch.bfloat16), self.weight.to(torch.bfloat16) |
| 121 | ), None |
| 122 | |
| 123 | |
| 124 | class _FusedLayer: |
| 125 | """Only the attributes K3's context loop reads.""" |
| 126 | |
| 127 | _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context |
| 128 | _gather_context_latent = MultiHeadLatentAttention._gather_context_latent |
| 129 | _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache |
| 130 | |
| 131 | def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: |
| 132 | self.kv_b_proj = kv_b_proj |
| 133 | self.kv_cache = kv_cache |
| 134 | self.kv_cache_dtype = kv_cache_dtype |
| 135 | self._k_scale = k_scale |
| 136 | self.kv_lora_rank = _KV_LORA_RANK |
| 137 | self.num_local_heads = _NUM_HEADS |
| 138 | self.qk_nope_head_dim = _QK_NOPE |
| 139 | self.v_head_dim = _V_HEAD_DIM |
| 140 | |
| 141 | |
| 142 | class _ReferenceImpl: |
| 143 | """Only the attributes the generic context loop reads.""" |
| 144 | |
| 145 | _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context |
| 146 | _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe |
| 147 | _use_flashinfer_concat_mla_k = False |
| 148 | |
| 149 | def __init__(self, kv_b_proj, kv_cache_dtype) -> None: |
| 150 | self.kv_b_proj = kv_b_proj |
| 151 | self.kv_cache_dtype = kv_cache_dtype |
| 152 | self.kv_lora_rank = _KV_LORA_RANK |
| 153 | self.num_heads = _NUM_HEADS |
| 154 | self.qk_nope_head_dim = _QK_NOPE |
| 155 | self.qk_rope_head_dim = _QK_ROPE |
| 156 | self.v_head_dim = _V_HEAD_DIM |
| 157 | |
| 158 | |
| 159 | def _build_prefill_metadata( |
| 160 | device: torch.device, |
| 161 | workspace_dtype: torch.dtype, |
| 162 | q_data_type: torch.dtype, |
| 163 | backend: _RecordingPrefillBackend, |
| 164 | ) -> MLACommonPrefillMetadata: |
| 165 | query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32) |
| 166 | query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0) |
| 167 | workspace = torch.empty( |
| 168 | (_WORKSPACE_TOKENS, _ENTRY), dtype=workspace_dtype, device=device |
| 169 | ) |
| 170 | chunked_context = build_mla_chunked_context_metadata( |
| 171 | context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32), |
| 172 | prefill_query_start_loc_cpu=query_start_loc_cpu, |
| 173 | chunked_prefill_workspace=workspace, |
| 174 | chunked_prefill_workspace_size=_WORKSPACE_TOKENS, |
| 175 | block_size=_BLOCK_SIZE, |
| 176 | align_chunk_to_block=True, |
| 177 | device=device, |
| 178 | dcp_world_size=1, |
| 179 | dcp_local_block_size=1, |
| 180 | dcp_virtual_block_size=1, |
| 181 | ) |
| 182 | assert chunked_context is not None |
| 183 | assert len(chunked_context.chunks) > 1, "the batch must exercise accumulation" |
| 184 | |
| 185 | max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1 |
| 186 | num_blocks = max_blocks * len(_CONTEXT_LENS) |
| 187 | block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view( |
| 188 | len(_CONTEXT_LENS), max_blocks |
| 189 | ) |
| 190 | return MLACommonPrefillMetadata( |
| 191 | block_table=block_table, |
| 192 | query_start_loc=query_start_loc_cpu.to(device), |
| 193 | max_query_len=max(_QUERY_LENS), |
| 194 | chunked_context=chunked_context, |
| 195 | q_data_type=q_data_type, |
| 196 | output_dtype=torch.bfloat16, |
| 197 | prefill_backend=backend, |
| 198 | ), num_blocks |
| 199 | |
| 200 | |
| 201 | @pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"]) |
| 202 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 203 | @pytest.mark.parametrize( |
| 204 | "kv_b_proj_quantized", [True, False], ids=["fp8_kv_b_proj", "bf16_kv_b_proj"] |
| 205 | ) |
| 206 | @torch.inference_mode() |
| 207 | def test_fused_context_matches_generic_impl( |
| 208 | kv_b_proj_quantized: bool, kv_cache_dtype: str, honors_out: bool |
| 209 | ) -> None: |
| 210 | """Cache dtype and ``kv_b_proj`` dtype vary independently. |
| 211 | |
| 212 | A stock K3 checkpoint pairs a bf16 ``kv_b_proj`` with an fp8 cache, so the |
| 213 | fused loop cannot assume the gathered latent is already in the dtype |
| 214 | ``kv_b_proj`` accepts. |
| 215 | """ |
| 216 | torch.manual_seed(0) |
| 217 | device = torch.device("cuda") |
| 218 | fp8 = current_platform.fp8_dtype() |
| 219 | quantized = kv_cache_dtype == "fp8" |
| 220 | # An fp8 cache is read by an fp8 query, so the workspace keeps the fp8 |
| 221 | # layout; a bf16 cache dequantizes into a bf16 workspace. |
| 222 | q_data_type = fp8 if quantized else torch.bfloat16 |
| 223 | workspace_dtype = q_data_type |
| 224 | |
| 225 | kv_b_proj = _KVBProj( |
| 226 | device, weight_dtype=fp8 if kv_b_proj_quantized else torch.bfloat16 |
| 227 | ) |
| 228 | k_scale = torch.ones(1, dtype=torch.float32, device=device) |
| 229 | |
| 230 | backend_fused = _RecordingPrefillBackend(honors_out=honors_out) |
| 231 | # The reference impl never passes `out`, so it always allocates its own. |
| 232 | backend_ref = _RecordingPrefillBackend() |
| 233 | prefill_fused, num_blocks = _build_prefill_metadata( |
| 234 | device, workspace_dtype, q_data_type, backend_fused |
| 235 | ) |
| 236 | prefill_ref, _ = _build_prefill_metadata( |
| 237 | device, workspace_dtype, q_data_type, backend_ref |
| 238 | ) |
| 239 | |
| 240 | cache = torch.randn( |
| 241 | (num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16 |
| 242 | ) |
| 243 | kv_cache = cache.to(fp8) if quantized else cache |
| 244 | q = ( |
| 245 | torch.randn( |
| 246 | (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), |
| 247 | device=device, |
| 248 | dtype=torch.bfloat16, |
| 249 | ) |
| 250 | * 0.2 |
| 251 | ).to(q_data_type) |
| 252 | |
| 253 | layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale) |
| 254 | fused_out, fused_lse = layer._compute_prefill_context( |
| 255 | q, SimpleNamespace(prefill=prefill_fused) |
| 256 | ) |
| 257 | |
| 258 | impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype) |
| 259 | ref_out, ref_lse = impl._compute_prefill_context( |
| 260 | q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale |
| 261 | ) |
| 262 | |
| 263 | assert len(backend_fused.calls) == len(backend_ref.calls) |
| 264 | for chunk_idx, (fused_call, ref_call) in enumerate( |
| 265 | zip(backend_fused.calls, backend_ref.calls, strict=True) |
| 266 | ): |
| 267 | for name, fused_t, ref_t in zip( |
| 268 | ("q", "k", "v"), fused_call, ref_call, strict=True |
| 269 | ): |
| 270 | torch.testing.assert_close( |
| 271 | fused_t, |
| 272 | ref_t, |
| 273 | atol=0, |
| 274 | rtol=0, |
| 275 | msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}", |
| 276 | ) |
| 277 | torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0) |
| 278 | torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0) |
| 279 | |
| 280 | # Every chunk but the continuation should have been written in place, i.e. |
| 281 | # straight into the returned accumulator, with no intermediate copy. |
| 282 | wrote_in_place = [ |
| 283 | out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr() |
| 284 | for out, chunk in zip( |
| 285 | backend_fused.out_destinations, |
| 286 | prefill_fused.chunked_context.chunks, |
| 287 | strict=True, |
| 288 | ) |
| 289 | ] |
| 290 | continuations = [c.is_continuation for c in prefill_fused.chunked_context.chunks] |
| 291 | assert any(continuations), "the batch must exercise a continuation chunk" |
| 292 | if honors_out: |
| 293 | assert wrote_in_place == [not c for c in continuations] |
| 294 | else: |
| 295 | assert not any(wrote_in_place) |
| 296 | |
| 297 | |
| 298 | @torch.inference_mode() |
| 299 | def test_fused_context_rejects_an_unquantized_query() -> None: |
| 300 | """The fp8 query is produced by the new-token epilogue, not re-cast here.""" |
| 301 | torch.manual_seed(0) |
| 302 | device = torch.device("cuda") |
| 303 | fp8 = current_platform.fp8_dtype() |
| 304 | prefill, num_blocks = _build_prefill_metadata( |
| 305 | device, fp8, fp8, _RecordingPrefillBackend() |
| 306 | ) |
| 307 | layer = _FusedLayer( |
| 308 | _KVBProj(device, weight_dtype=fp8), |
| 309 | torch.zeros((num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=fp8), |
| 310 | "fp8", |
| 311 | torch.ones(1, dtype=torch.float32, device=device), |
| 312 | ) |
| 313 | q = torch.zeros( |
| 314 | (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), |
| 315 | device=device, |
| 316 | dtype=torch.bfloat16, |
| 317 | ) |
| 318 | with pytest.raises(AssertionError, match="new-token epilogue"): |
| 319 | layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) |