| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import random |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.utils.deep_gemm import ( |
| 10 | _ceil_to_ue8m0, |
| 11 | calc_diff, |
| 12 | fp8_fp4_mqa_logits, |
| 13 | fp8_fp4_paged_mqa_logits, |
| 14 | get_num_sms, |
| 15 | get_paged_mqa_logits_metadata, |
| 16 | native_next_n_supported, |
| 17 | ) |
| 18 | from vllm.utils.import_utils import has_deep_gemm |
| 19 | from vllm.utils.math_utils import cdiv |
| 20 | |
| 21 | |
| 22 | def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: |
| 23 | # x: (num_blocks, block_size, 1, head_dim) |
| 24 | num_blocks, block_size, num_heads, head_dim = x.shape |
| 25 | assert num_heads == 1 |
| 26 | x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) |
| 27 | sf = x_amax / 448.0 |
| 28 | x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) |
| 29 | x_fp8 = torch.empty( |
| 30 | (num_blocks, block_size * (head_dim + 4)), |
| 31 | device=x.device, |
| 32 | dtype=torch.uint8, |
| 33 | ) |
| 34 | x_fp8[:, : block_size * head_dim] = x_scaled.view( |
| 35 | num_blocks, block_size * head_dim |
| 36 | ).view(dtype=torch.uint8) |
| 37 | x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view( |
| 38 | dtype=torch.uint8 |
| 39 | ) |
| 40 | return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) |
| 41 | |
| 42 | |
| 43 | def per_custom_dims_cast_to_fp8( |
| 44 | x: torch.Tensor, dims: tuple, use_ue8m0: bool |
| 45 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 46 | excluded_dims = tuple([i for i in range(x.dim()) if i not in set(dims)]) |
| 47 | x_amax = x.abs().float().amax(dim=excluded_dims, keepdim=True).clamp(1e-4) |
| 48 | sf = x_amax / 448.0 |
| 49 | sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf |
| 50 | x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) |
| 51 | return x_scaled, sf.squeeze() |
| 52 | |
| 53 | |
| 54 | def _generate_cp_test_data(seq_len: int, seq_len_kv: int): |
| 55 | assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0 |
| 56 | chunk_size = seq_len // 2 |
| 57 | cp_size = seq_len_kv // seq_len |
| 58 | cp_id = cp_size // 3 |
| 59 | ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") |
| 60 | ke = torch.zeros(seq_len, dtype=torch.int, device="cuda") |
| 61 | for i in range(chunk_size): |
| 62 | ke[i] = cp_id * chunk_size + i |
| 63 | ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i |
| 64 | return ks, ke |
| 65 | |
| 66 | |
| 67 | def _ref_fp8_mqa_logits( |
| 68 | q: torch.Tensor, |
| 69 | kv: torch.Tensor, |
| 70 | weights: torch.Tensor, |
| 71 | cu_seqlen_ks: torch.Tensor, |
| 72 | cu_seqlen_ke: torch.Tensor, |
| 73 | ): |
| 74 | seq_len_kv = kv.shape[0] |
| 75 | |
| 76 | k = kv |
| 77 | q = q.float() |
| 78 | k = k.float() |
| 79 | |
| 80 | mask_lo = ( |
| 81 | torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None] |
| 82 | ) |
| 83 | mask_hi = ( |
| 84 | torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None] |
| 85 | ) |
| 86 | mask = mask_lo & mask_hi |
| 87 | score = torch.einsum("mhd,nd->hmn", q, k) |
| 88 | logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0) |
| 89 | logits = logits.masked_fill(~mask, float("-inf")) |
| 90 | |
| 91 | return logits |
| 92 | |
| 93 | |
| 94 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") |
| 95 | @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") |
| 96 | @pytest.mark.skipif( |
| 97 | not current_platform.has_device_capability(90), reason="SM90 and SM100 only" |
| 98 | ) |
| 99 | @pytest.mark.parametrize("clean_logits", [True, False]) |
| 100 | def test_deepgemm_fp8_mqa_logits(clean_logits: bool): |
| 101 | torch.manual_seed(0) |
| 102 | random.seed(0) |
| 103 | num_heads, head_dim = 32, 128 |
| 104 | for seq_len in (512,): |
| 105 | for seq_len_kv in (1024,): |
| 106 | for disable_cp in (False, True): |
| 107 | q = torch.randn( |
| 108 | seq_len, |
| 109 | num_heads, |
| 110 | head_dim, |
| 111 | device="cuda", |
| 112 | dtype=torch.bfloat16, |
| 113 | ) |
| 114 | kv = torch.randn( |
| 115 | seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16 |
| 116 | ) |
| 117 | weights = torch.randn( |
| 118 | seq_len, num_heads, device="cuda", dtype=torch.float32 |
| 119 | ) |
| 120 | |
| 121 | if disable_cp: |
| 122 | ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") |
| 123 | ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( |
| 124 | seq_len_kv - seq_len |
| 125 | ) |
| 126 | else: |
| 127 | ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) |
| 128 | |
| 129 | q_fp8 = q.to(torch.float8_e4m3fn) |
| 130 | kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) |
| 131 | logits = fp8_fp4_mqa_logits( |
| 132 | (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits |
| 133 | ) |
| 134 | |
| 135 | ref_logits = _ref_fp8_mqa_logits( |
| 136 | q=q, |
| 137 | kv=kv, |
| 138 | weights=weights, |
| 139 | cu_seqlen_ks=ks, |
| 140 | cu_seqlen_ke=ke, |
| 141 | ) |
| 142 | ref_neginf_mask = ref_logits == float("-inf") |
| 143 | |
| 144 | if clean_logits: |
| 145 | neginf_mask = logits == float("-inf") |
| 146 | assert torch.equal(neginf_mask, ref_neginf_mask) |
| 147 | |
| 148 | ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) |
| 149 | logits = logits.masked_fill(ref_neginf_mask, 0) |
| 150 | diff = calc_diff(logits, ref_logits) |
| 151 | assert diff < 1e-3, f"{diff=}" |
| 152 | |
| 153 | |
| 154 | def _ref_fp8_fp4_paged_mqa_logits( |
| 155 | q: torch.Tensor, |
| 156 | kv_cache: torch.Tensor, |
| 157 | weights: torch.Tensor, |
| 158 | context_lens: torch.Tensor, |
| 159 | block_tables: torch.Tensor, |
| 160 | max_model_len: int, |
| 161 | ): |
| 162 | batch_size, next_n, _, _ = q.size() |
| 163 | _, block_size, _, _ = kv_cache.size() |
| 164 | logits = torch.full( |
| 165 | [batch_size * next_n, max_model_len], |
| 166 | float("-inf"), |
| 167 | device=q.device, |
| 168 | dtype=torch.float32, |
| 169 | ) |
| 170 | context_lens_list = context_lens.tolist() |
| 171 | for i in range(batch_size): |
| 172 | context_len = context_lens_list[i] |
| 173 | q_offsets = torch.arange(context_len - next_n, context_len, device="cuda") |
| 174 | weight_slice = ( |
| 175 | weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous() |
| 176 | ) |
| 177 | for block_rk in range(cdiv(context_len, block_size)): |
| 178 | block_idx = block_tables[i][block_rk] |
| 179 | qx, kx = q[i], kv_cache[block_idx] |
| 180 | k_offsets = torch.arange( |
| 181 | block_rk * block_size, |
| 182 | (block_rk + 1) * block_size, |
| 183 | device="cuda", |
| 184 | ) |
| 185 | mask = (k_offsets[None, :] < context_len) & ( |
| 186 | k_offsets[None, :] <= q_offsets[:, None] |
| 187 | ) |
| 188 | s = torch.where( |
| 189 | mask[None, :, :], |
| 190 | (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to( |
| 191 | logits.dtype |
| 192 | ), |
| 193 | float("-inf"), |
| 194 | ) |
| 195 | s = torch.relu(s) * weight_slice[..., None] |
| 196 | s = s.sum(dim=0) |
| 197 | logits[ |
| 198 | i * next_n : (i + 1) * next_n, |
| 199 | block_rk * block_size : (block_rk + 1) * block_size, |
| 200 | ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf")) |
| 201 | return logits |
| 202 | |
| 203 | |
| 204 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") |
| 205 | @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") |
| 206 | @pytest.mark.skipif( |
| 207 | not current_platform.has_device_capability(90), reason="SM90 and SM100 only" |
| 208 | ) |
| 209 | # next_n = 1 + num_speculative_tokens, so next_n=4 is MTP=3 (issue #35878). |
| 210 | @pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (2, 4)]) |
| 211 | def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): |
| 212 | if not native_next_n_supported(next_n): |
| 213 | pytest.skip(f"next_n={next_n} has no native kernel on this architecture") |
| 214 | |
| 215 | # NOTE: clean_logits=True is incompatible with the 2D context_lens |
| 216 | # required by csrc/apis/attention.hpp; only the False path is exercised. |
| 217 | clean_logits = False |
| 218 | torch.manual_seed(0) |
| 219 | random.seed(0) |
| 220 | |
| 221 | max_model_len = 4096 |
| 222 | for heads, index_dim in [(32, 128)]: |
| 223 | for avg_kv in (2048,): |
| 224 | num_blocks, blocksize = max_model_len * 2, 64 |
| 225 | |
| 226 | q = torch.randn( |
| 227 | (batch_size, next_n, heads, index_dim), |
| 228 | device="cuda", |
| 229 | dtype=torch.bfloat16, |
| 230 | ) |
| 231 | kv_cache = torch.randn( |
| 232 | (num_blocks, blocksize, 1, index_dim), |
| 233 | device="cuda", |
| 234 | dtype=torch.bfloat16, |
| 235 | ) |
| 236 | weights = torch.randn( |
| 237 | (batch_size * next_n, heads), |
| 238 | device="cuda", |
| 239 | dtype=torch.float32, |
| 240 | ) |
| 241 | |
| 242 | context_lens = ( |
| 243 | torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,)) |
| 244 | .cuda() |
| 245 | .to(torch.int32) |
| 246 | ) |
| 247 | max_block_len = ( |
| 248 | (context_lens.max().item() + blocksize - 1) // blocksize * blocksize |
| 249 | ) |
| 250 | block_tables = torch.zeros( |
| 251 | (batch_size, max_block_len), |
| 252 | device="cuda", |
| 253 | dtype=torch.int32, |
| 254 | ) |
| 255 | |
| 256 | counter = 0 |
| 257 | block_idx_pool = list(range(num_blocks)) |
| 258 | random.shuffle(block_idx_pool) |
| 259 | for i in range(batch_size): |
| 260 | ctx_len = int(context_lens[i].item()) |
| 261 | for j in range((ctx_len + blocksize - 1) // blocksize): |
| 262 | block_tables[i][j] = block_idx_pool[counter] |
| 263 | counter += 1 |
| 264 | |
| 265 | q_fp8 = q.to(torch.float8_e4m3fn) |
| 266 | kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache) |
| 267 | |
| 268 | # deep_gemm paged MQA logits requires 2D context_lens of |
| 269 | # shape (B, next_n) (csrc/apis/attention.hpp:332-335); |
| 270 | # see indexer.py:607-608. For each batch/next_n token, the |
| 271 | # effective context length is context_lens[b] - next_n + j + 1. |
| 272 | next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32) |
| 273 | context_lens_2d = ( |
| 274 | context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange |
| 275 | ).contiguous() |
| 276 | schedule_metadata = get_paged_mqa_logits_metadata( |
| 277 | context_lens_2d, |
| 278 | blocksize, |
| 279 | get_num_sms(), |
| 280 | ) |
| 281 | logits = fp8_fp4_paged_mqa_logits( |
| 282 | (q_fp8, None), |
| 283 | kv_cache_fp8, |
| 284 | weights, |
| 285 | context_lens_2d, |
| 286 | block_tables, |
| 287 | schedule_metadata, |
| 288 | max_model_len, |
| 289 | clean_logits=clean_logits, |
| 290 | ) |
| 291 | |
| 292 | ref_logits = _ref_fp8_fp4_paged_mqa_logits( |
| 293 | q, |
| 294 | kv_cache, |
| 295 | weights, |
| 296 | context_lens, |
| 297 | block_tables, |
| 298 | max_model_len, |
| 299 | ) |
| 300 | |
| 301 | positions = ( |
| 302 | torch.arange(max_model_len, device="cuda") |
| 303 | .unsqueeze(0) |
| 304 | .expand(batch_size * next_n, -1) |
| 305 | ) |
| 306 | row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n |
| 307 | next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n |
| 308 | mask = positions <= ( |
| 309 | context_lens[row_indices] - next_n + next_n_offset |
| 310 | ).unsqueeze(1) |
| 311 | |
| 312 | logits = logits.masked_fill(~mask, 0) |
| 313 | ref_logits = ref_logits.masked_fill(~mask, 0) |
| 314 | diff = calc_diff(logits, ref_logits) |
| 315 | assert diff < 1e-3, f"{diff=}" |