| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | import random |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | import vllm._custom_ops as ops |
| 10 | from vllm.platforms import current_platform |
| 11 | from vllm.triton_utils import triton |
| 12 | from vllm.utils.platform_utils import num_compute_units |
| 13 | |
| 14 | pytestmark = pytest.mark.skip_global_cleanup |
| 15 | |
| 16 | |
| 17 | def cal_diff( |
| 18 | x: torch.Tensor, |
| 19 | y: torch.Tensor, |
| 20 | name: str, |
| 21 | use_fp8: bool = False, |
| 22 | diff_threshold: float | None = None, |
| 23 | ) -> None: |
| 24 | x, y = x.double(), y.double() |
| 25 | cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12) |
| 26 | if diff_threshold is not None: |
| 27 | # directly compare the cos_diff with the threshold |
| 28 | assert cos_diff < diff_threshold |
| 29 | else: |
| 30 | # use the default threshold |
| 31 | if use_fp8: |
| 32 | assert cos_diff < 1e-4 |
| 33 | else: |
| 34 | assert cos_diff < 1e-5 |
| 35 | |
| 36 | |
| 37 | CUTLASS_MLA_UNSUPPORTED_REASON = ( |
| 38 | "Cutlass MLA Requires compute capability of 100 or above." |
| 39 | if not current_platform.is_device_capability_family(100) |
| 40 | else "Cutlass MLA is supported" |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | @pytest.mark.skipif( |
| 45 | not current_platform.is_device_capability_family(100), |
| 46 | reason=CUTLASS_MLA_UNSUPPORTED_REASON, |
| 47 | ) |
| 48 | @pytest.mark.parametrize("b", [128]) |
| 49 | @pytest.mark.parametrize("s_q", [1]) |
| 50 | @pytest.mark.parametrize("mean_sk", [4096, 8192, 16384]) |
| 51 | @pytest.mark.parametrize("h_q", [16, 32, 64, 128]) |
| 52 | @pytest.mark.parametrize("h_kv", [1]) |
| 53 | @pytest.mark.parametrize("d", [576]) |
| 54 | @pytest.mark.parametrize("dv", [512]) |
| 55 | @pytest.mark.parametrize("block_size", [64]) |
| 56 | @pytest.mark.parametrize("causal", [True]) |
| 57 | @pytest.mark.parametrize("varlen", [False, True]) |
| 58 | @pytest.mark.parametrize( |
| 59 | "torch_dtype", |
| 60 | [ |
| 61 | torch.bfloat16, |
| 62 | # fp8 can have occasional precision-related failures. |
| 63 | pytest.param(torch.float8_e4m3fn, marks=pytest.mark.flaky(reruns=2)), |
| 64 | ], |
| 65 | ) |
| 66 | @torch.inference_mode() |
| 67 | def test_cutlass_mla_decode( |
| 68 | b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype |
| 69 | ): |
| 70 | device = torch.device("cuda:0") |
| 71 | init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype |
| 72 | torch.set_default_dtype(init_dtype) |
| 73 | torch.set_default_device(device) |
| 74 | torch.accelerator.set_device_index(device) |
| 75 | torch.manual_seed(42) |
| 76 | random.seed(42) |
| 77 | |
| 78 | print( |
| 79 | f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, " |
| 80 | f"{d=}, {dv=}, {causal=}, {varlen=}, {torch_dtype=}" |
| 81 | ) |
| 82 | |
| 83 | use_fp8 = torch_dtype == torch.float8_e4m3fn |
| 84 | scale = math.sqrt(d) ** (-1) |
| 85 | cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32) |
| 86 | if varlen: |
| 87 | for i in range(b): |
| 88 | cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q) |
| 89 | total_seqlens = cache_seqlens.sum().item() |
| 90 | max_seqlen = cache_seqlens.max().item() |
| 91 | max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 |
| 92 | |
| 93 | q = torch.randn(b, s_q, h_q, d) |
| 94 | block_table = torch.arange( |
| 95 | b * max_seqlen_pad // block_size, dtype=torch.int32 |
| 96 | ).view(b, max_seqlen_pad // block_size) |
| 97 | blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d) |
| 98 | blocked_v = blocked_k[..., :dv] |
| 99 | |
| 100 | init_dtype = q.dtype |
| 101 | if use_fp8: |
| 102 | fp8_dtype = torch.float8_e4m3fn |
| 103 | descale_q = torch.ones((1), dtype=torch.float32) |
| 104 | descale_k = torch.ones((1), dtype=torch.float32) |
| 105 | |
| 106 | q = q.to(fp8_dtype) |
| 107 | blocked_k = blocked_k.to(fp8_dtype) |
| 108 | blocked_v = blocked_v.to(fp8_dtype) |
| 109 | else: |
| 110 | descale_q = None |
| 111 | descale_k = None |
| 112 | |
| 113 | def cutlass_mla(): |
| 114 | MAX_HEADS = 128 |
| 115 | |
| 116 | q_reshaped = q.squeeze(1) |
| 117 | q_nope = q_reshaped[:, :, :dv].clone() |
| 118 | q_pe = q_reshaped[:, :, dv:].clone() |
| 119 | |
| 120 | if h_q < MAX_HEADS: |
| 121 | q_nope_padded = q_nope.new_empty((b, MAX_HEADS, dv)) |
| 122 | q_nope_padded[:, :h_q] = q_nope |
| 123 | q_nope = q_nope_padded |
| 124 | |
| 125 | q_pe_padded = q_pe.new_empty((b, MAX_HEADS, d - dv)) |
| 126 | q_pe_padded[:, :h_q] = q_pe |
| 127 | q_pe = q_pe_padded |
| 128 | |
| 129 | kv_cache_flat = blocked_k.squeeze(2) |
| 130 | sm_count = num_compute_units(device.index) |
| 131 | workspace_size = ops.sm100_cutlass_mla_get_workspace_size( |
| 132 | max_seqlen * block_size, b, sm_count, num_kv_splits=1 |
| 133 | ) |
| 134 | workspace = torch.empty(workspace_size, device="cuda", dtype=torch.uint8) |
| 135 | |
| 136 | out_ans = torch.empty(b, MAX_HEADS, dv, dtype=init_dtype) |
| 137 | output_lse = torch.empty( |
| 138 | (b, MAX_HEADS), dtype=torch.float32, device=q_nope.device |
| 139 | ) |
| 140 | ops.sm100_cutlass_mla_decode( |
| 141 | out_ans, |
| 142 | output_lse, |
| 143 | q_nope, |
| 144 | q_pe, |
| 145 | kv_cache_flat, |
| 146 | cache_seqlens, |
| 147 | block_table, |
| 148 | workspace, |
| 149 | scale, |
| 150 | 1, |
| 151 | ) |
| 152 | return out_ans[:, :h_q].contiguous(), output_lse[:, :h_q].contiguous() |
| 153 | |
| 154 | def scaled_dot_product_attention(query, key, value, is_causal=False): |
| 155 | query = query.float() |
| 156 | key = key.float() |
| 157 | value = value.float() |
| 158 | key = key.repeat_interleave(h_q // h_kv, dim=0) |
| 159 | value = value.repeat_interleave(h_q // h_kv, dim=0) |
| 160 | attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1)) |
| 161 | if is_causal: |
| 162 | s_q = query.shape[-2] |
| 163 | s_k = key.shape[-2] |
| 164 | attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype) |
| 165 | temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q) |
| 166 | attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) |
| 167 | attn_bias.to(query.dtype) |
| 168 | attn_weight += attn_bias |
| 169 | lse = attn_weight.logsumexp(dim=-1) |
| 170 | attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32) |
| 171 | return attn_weight @ value, lse |
| 172 | |
| 173 | def ref_mla(): |
| 174 | q_ = (q.to(torch.float) * descale_q).to(init_dtype) if use_fp8 else q |
| 175 | blocked_k_ = ( |
| 176 | (blocked_k.to(torch.float) * descale_k).to(init_dtype) |
| 177 | if use_fp8 |
| 178 | else blocked_k |
| 179 | ) |
| 180 | blocked_v_ = ( |
| 181 | (blocked_v.to(torch.float) * descale_k).to(init_dtype) |
| 182 | if use_fp8 |
| 183 | else blocked_v |
| 184 | ) |
| 185 | out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32) |
| 186 | lse = torch.empty(b, h_q, s_q, dtype=torch.float32) |
| 187 | for i in range(b): |
| 188 | begin = i * max_seqlen_pad |
| 189 | end = begin + cache_seqlens[i] |
| 190 | out_i, lse_i = scaled_dot_product_attention( |
| 191 | q_[i].transpose(0, 1), |
| 192 | blocked_k_.view(-1, h_kv, d)[begin:end].transpose(0, 1), |
| 193 | blocked_v_.view(-1, h_kv, dv)[begin:end].transpose(0, 1), |
| 194 | is_causal=causal, |
| 195 | ) |
| 196 | out[i] = out_i.transpose(0, 1) |
| 197 | lse[i] = lse_i |
| 198 | return out, lse |
| 199 | |
| 200 | out_cutlass, lse_cutlass = cutlass_mla() |
| 201 | out_torch, lse_torch = ref_mla() |
| 202 | # Extract the single token (s_q=1) slice to match cutlass output shape |
| 203 | out_torch_slice = out_torch[:, 0, :, :] # [b, h_q, dv] |
| 204 | lse_torch_slice = lse_torch[:, 0, :] # [b, h_q] |
| 205 | cal_diff(out_cutlass, out_torch_slice, "out", use_fp8) |
| 206 | # lse has larger numerical error, so use a larger threshold |
| 207 | cal_diff(lse_cutlass, lse_torch_slice, "lse", use_fp8, diff_threshold=1e-3) |
| 208 | |
| 209 | t = triton.testing.do_bench(cutlass_mla) |
| 210 | FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2 |
| 211 | bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d) * ( |
| 212 | torch.finfo(torch_dtype).bits // 8 |
| 213 | ) + (b * s_q * h_q * dv) * (torch.finfo(init_dtype).bits // 8) |
| 214 | print( |
| 215 | f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s" |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | @pytest.mark.skipif( |
| 220 | not current_platform.is_device_capability_family(100), |
| 221 | reason=CUTLASS_MLA_UNSUPPORTED_REASON, |
| 222 | ) |
| 223 | @torch.inference_mode() |
| 224 | def test_cutlass_mla_decode_cross_layer_view(): |
| 225 | """The kernel must read the cache's page-dim stride instead of assuming |
| 226 | pages are packed back-to-back. A per-layer view into a cross-layer |
| 227 | (block-major) cache has stride(0) inflated by num_layers; outputs must |
| 228 | match a contiguous cache holding the same data exactly.""" |
| 229 | device = torch.device("cuda:0") |
| 230 | torch.set_default_dtype(torch.bfloat16) |
| 231 | torch.set_default_device(device) |
| 232 | torch.manual_seed(42) |
| 233 | |
| 234 | b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64 |
| 235 | num_layers, layer_idx = 3, 1 |
| 236 | scale = math.sqrt(d) ** (-1) |
| 237 | |
| 238 | num_pages = b * (mean_sk // block_size) |
| 239 | cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32) |
| 240 | block_table = torch.arange(num_pages, dtype=torch.int32).view( |
| 241 | b, mean_sk // block_size |
| 242 | ) |
| 243 | |
| 244 | kv_contig = torch.randn(num_pages, block_size, d) |
| 245 | # Neighbor layers hold random data so packed-pages addressing reads |
| 246 | # garbage rather than zeros. |
| 247 | kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d) |
| 248 | kv_view = kv_cross_layer[:, layer_idx] |
| 249 | kv_view.copy_(kv_contig) |
| 250 | assert kv_view.stride(0) == num_layers * block_size * d |
| 251 | |
| 252 | q_nope = torch.randn(b, 128, dv) |
| 253 | q_pe = torch.randn(b, 128, d - dv) |
| 254 | sm_count = num_compute_units(device.index) |
| 255 | workspace_size = ops.sm100_cutlass_mla_get_workspace_size( |
| 256 | mean_sk, b, sm_count, num_kv_splits=1 |
| 257 | ) |
| 258 | workspace = torch.empty(workspace_size, dtype=torch.uint8) |
| 259 | |
| 260 | def run(cache): |
| 261 | out = torch.empty(b, 128, dv) |
| 262 | lse = torch.empty(b, 128, dtype=torch.float32) |
| 263 | ops.sm100_cutlass_mla_decode( |
| 264 | out, |
| 265 | lse, |
| 266 | q_nope, |
| 267 | q_pe, |
| 268 | cache, |
| 269 | cache_seqlens, |
| 270 | block_table, |
| 271 | workspace, |
| 272 | scale, |
| 273 | 1, |
| 274 | ) |
| 275 | return out, lse |
| 276 | |
| 277 | out_contig, lse_contig = run(kv_contig) |
| 278 | out_view, lse_view = run(kv_view) |
| 279 | |
| 280 | # Same data and same compute order; only addressing differs. |
| 281 | assert torch.equal(out_contig, out_view) |
| 282 | assert torch.equal(lse_contig, lse_view) |