| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """FP8 accuracy tests for LoRA shrink and expand kernels. |
| 4 | |
| 5 | Tests the FP8 kernels by: |
| 6 | 1. Quantizing bf16 inputs/weights to FP8 |
| 7 | 2. Dequantizing them back to bf16 |
| 8 | 3. Running the bf16 reference (sgmv_shrink/sgmv_expand) with dequantized values |
| 9 | 4. Comparing FP8 kernel output against this dequantized reference |
| 10 | |
| 11 | This isolates kernel correctness from quantization precision loss, |
| 12 | allowing much tighter tolerances than comparing against the original bf16. |
| 13 | """ |
| 14 | |
| 15 | import math |
| 16 | from threading import Lock |
| 17 | |
| 18 | import pytest |
| 19 | import torch |
| 20 | import torch.nn.functional as F |
| 21 | |
| 22 | import vllm.lora.ops.torch_ops as torch_ops |
| 23 | import vllm.lora.ops.triton_ops as triton_ops |
| 24 | from vllm.lora.ops.triton_ops import LoRAKernelMeta |
| 25 | from vllm.lora.ops.triton_ops.lora_expand_fp8_op import ( |
| 26 | _EXPAND_LORA_SCALE_PTR_DICT, |
| 27 | ) |
| 28 | from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import ( |
| 29 | _SHRINK_LORA_SCALE_PTR_DICT, |
| 30 | ) |
| 31 | from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT |
| 32 | from vllm.platforms import current_platform |
| 33 | from vllm.utils.torch_utils import set_random_seed |
| 34 | |
| 35 | pytestmark = pytest.mark.skip_global_cleanup |
| 36 | |
| 37 | DEVICE_TYPE = current_platform.device_type |
| 38 | DEVICES = [f"{DEVICE_TYPE}:{0}"] |
| 39 | SEED = [0] |
| 40 | |
| 41 | _dict_lock = Lock() |
| 42 | |
| 43 | |
| 44 | @pytest.fixture(autouse=True) |
| 45 | def reset_device(reset_default_device): |
| 46 | pass |
| 47 | |
| 48 | |
| 49 | # ============================================================================ |
| 50 | # Reference implementations (bf16 baseline) |
| 51 | # ============================================================================ |
| 52 | |
| 53 | |
| 54 | def sgmv_shrink_for_nslices( |
| 55 | nslices, |
| 56 | inputs_tensor, |
| 57 | lora_weights_lst, |
| 58 | out_tensor, |
| 59 | b_seq_start_loc, |
| 60 | seq_len_tensor, |
| 61 | prompt_lora_mapping, |
| 62 | batches, |
| 63 | max_seq_length, |
| 64 | num_tokens, |
| 65 | scaling, |
| 66 | ): |
| 67 | """Wrapper around torch_ops.sgmv_shrink that handles any nslices.""" |
| 68 | for index in range(nslices): |
| 69 | torch_ops.sgmv_shrink( |
| 70 | inputs_tensor, |
| 71 | lora_weights_lst[index], |
| 72 | out_tensor[index], |
| 73 | b_seq_start_loc, |
| 74 | seq_len_tensor, |
| 75 | prompt_lora_mapping, |
| 76 | batches, |
| 77 | max_seq_length, |
| 78 | num_tokens, |
| 79 | scaling, |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | def sgmv_expand_for_nslices( |
| 84 | nslices, |
| 85 | hidden_size, |
| 86 | inputs_tensor, |
| 87 | lora_weights_lst, |
| 88 | out_tensor, |
| 89 | b_seq_start_loc, |
| 90 | seq_len_tensor, |
| 91 | prompt_lora_mapping, |
| 92 | batches, |
| 93 | max_seq_length, |
| 94 | num_tokens, |
| 95 | add_inputs, |
| 96 | ): |
| 97 | """Wrapper around torch_ops.sgmv_expand that handles any nslices.""" |
| 98 | if nslices == 1: |
| 99 | torch_ops.sgmv_expand( |
| 100 | inputs_tensor[0], |
| 101 | lora_weights_lst[0], |
| 102 | out_tensor, |
| 103 | b_seq_start_loc, |
| 104 | seq_len_tensor, |
| 105 | prompt_lora_mapping, |
| 106 | batches, |
| 107 | max_seq_length, |
| 108 | num_tokens, |
| 109 | add_inputs=add_inputs, |
| 110 | ) |
| 111 | else: |
| 112 | slice_offset = 0 |
| 113 | for index in range(nslices): |
| 114 | torch_ops.sgmv_expand_slice( |
| 115 | inputs_tensor[index], |
| 116 | lora_weights_lst[index], |
| 117 | out_tensor, |
| 118 | b_seq_start_loc, |
| 119 | seq_len_tensor, |
| 120 | prompt_lora_mapping, |
| 121 | batches, |
| 122 | max_seq_length, |
| 123 | num_tokens, |
| 124 | slice_offset, |
| 125 | hidden_size, |
| 126 | add_inputs=add_inputs, |
| 127 | ) |
| 128 | slice_offset += hidden_size |
| 129 | |
| 130 | |
| 131 | # ============================================================================ |
| 132 | # FP8 Quantization Helpers |
| 133 | # ============================================================================ |
| 134 | |
| 135 | FP8_DTYPE = torch.float8_e4m3fn |
| 136 | FP8_MAX = torch.finfo(FP8_DTYPE).max |
| 137 | FP8_MIN = torch.finfo(FP8_DTYPE).min |
| 138 | |
| 139 | |
| 140 | def quantize_to_fp8_per_tensor( |
| 141 | tensor: torch.Tensor, |
| 142 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 143 | """Quantize a tensor to FP8 with per-tensor scaling.""" |
| 144 | amax = tensor.abs().float().max().clamp(min=1e-12) |
| 145 | scale = (amax / FP8_MAX).to(torch.float32) |
| 146 | fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) |
| 147 | return fp8_tensor, scale.reshape(1) |
| 148 | |
| 149 | |
| 150 | def quantize_to_fp8_per_channel( |
| 151 | tensor: torch.Tensor, |
| 152 | channel_dim: int = 0, |
| 153 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 154 | """Quantize a tensor to FP8 with per-channel scaling. |
| 155 | |
| 156 | For shrink lora_a weights of shape (num_loras, rank, hidden_size): |
| 157 | channel_dim=1 gives per-rank scaling -> scale shape (num_loras, rank) |
| 158 | For expand lora_b weights of shape (num_loras, hidden_size, rank): |
| 159 | channel_dim=1 gives per-hidden scaling -> scale shape (num_loras, hidden_size) |
| 160 | """ |
| 161 | # Compute amax along all dims except the leading dims up to channel_dim+1 |
| 162 | reduce_dims = list(range(channel_dim + 1, tensor.ndim)) |
| 163 | if reduce_dims: |
| 164 | amax = tensor.abs().float().amax(dim=reduce_dims).clamp(min=1e-12) |
| 165 | else: |
| 166 | amax = tensor.abs().float().clamp(min=1e-12) |
| 167 | scale = (amax / FP8_MAX).to(torch.float32) |
| 168 | |
| 169 | # Expand scale for broadcasting |
| 170 | for _ in reduce_dims: |
| 171 | scale = scale.unsqueeze(-1) |
| 172 | fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) |
| 173 | scale = scale.squeeze() |
| 174 | if scale.ndim == 0: |
| 175 | scale = scale.unsqueeze(0) |
| 176 | return fp8_tensor, scale |
| 177 | |
| 178 | |
| 179 | def quantize_to_fp8_per_token( |
| 180 | tensor: torch.Tensor, |
| 181 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 182 | """Quantize a 2D tensor to FP8 with per-token (per-row) scaling. |
| 183 | |
| 184 | Input shape: (num_tokens, hidden_size) |
| 185 | Returns: (fp8_tensor, scale) where scale shape is (num_tokens, 1) |
| 186 | """ |
| 187 | assert tensor.ndim == 2 |
| 188 | amax = tensor.abs().float().amax(dim=1, keepdim=True).clamp(min=1e-12) |
| 189 | scale = (amax / FP8_MAX).to(torch.float32) |
| 190 | fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) |
| 191 | return fp8_tensor, scale |
| 192 | |
| 193 | |
| 194 | def quantize_to_fp8_blockwise( |
| 195 | tensor: torch.Tensor, |
| 196 | group_n: int, |
| 197 | group_k: int, |
| 198 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 199 | """Quantize a 2D or 3D tensor to FP8 with block-wise scaling. |
| 200 | |
| 201 | For a 2D tensor (num_tokens, hidden_size): |
| 202 | Blocks of size (1, group_k) -> |
| 203 | scale shape (num_tokens, ceil(hidden_size/group_k)) |
| 204 | |
| 205 | For a 3D tensor (num_loras, N, K): |
| 206 | Blocks of size (group_n, group_k) -> |
| 207 | scale shape (num_loras, ceil(N/group_n), ceil(K/group_k)) |
| 208 | """ |
| 209 | if tensor.ndim == 2: |
| 210 | M, K = tensor.shape |
| 211 | n_blocks_k = math.ceil(K / group_k) |
| 212 | # Zero padding leaves the per-block amax unchanged, so the padded |
| 213 | # columns are dropped again after quantizing. |
| 214 | padded = F.pad(tensor.float(), (0, n_blocks_k * group_k - K)) |
| 215 | blocks = padded.view(M, n_blocks_k, group_k) |
| 216 | scale = blocks.abs().amax(dim=-1).clamp(min=1e-12) / FP8_MAX |
| 217 | fp8_tensor = ( |
| 218 | (blocks / scale.unsqueeze(-1)).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) |
| 219 | ) |
| 220 | return fp8_tensor.view(M, -1)[:, :K].contiguous(), scale |
| 221 | elif tensor.ndim == 3: |
| 222 | L, N, K = tensor.shape |
| 223 | n_blocks_n = math.ceil(N / group_n) |
| 224 | n_blocks_k = math.ceil(K / group_k) |
| 225 | padded = F.pad( |
| 226 | tensor.float(), |
| 227 | (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N), |
| 228 | ) |
| 229 | blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k) |
| 230 | scale = blocks.abs().amax(dim=(2, 4)).clamp(min=1e-12) / FP8_MAX |
| 231 | fp8_tensor = ( |
| 232 | (blocks / scale[:, :, None, :, None]) |
| 233 | .clamp(FP8_MIN, FP8_MAX) |
| 234 | .to(FP8_DTYPE) |
| 235 | .view(L, n_blocks_n * group_n, n_blocks_k * group_k) |
| 236 | ) |
| 237 | return fp8_tensor[:, :N, :K].contiguous(), scale |
| 238 | else: |
| 239 | raise ValueError(f"Unsupported tensor ndim: {tensor.ndim}") |
| 240 | |
| 241 | |
| 242 | # ============================================================================ |
| 243 | # FP8 Dequantization Helpers |
| 244 | # ============================================================================ |
| 245 | |
| 246 | |
| 247 | def dequantize_fp8_per_tensor( |
| 248 | fp8_tensor: torch.Tensor, |
| 249 | scale: torch.Tensor, |
| 250 | output_dtype: torch.dtype = torch.bfloat16, |
| 251 | ) -> torch.Tensor: |
| 252 | """Dequantize FP8 tensor with per-tensor scale back to output_dtype.""" |
| 253 | return (fp8_tensor.float() * scale.float()).to(output_dtype) |
| 254 | |
| 255 | |
| 256 | def dequantize_fp8_per_channel( |
| 257 | fp8_tensor: torch.Tensor, |
| 258 | scale: torch.Tensor, |
| 259 | channel_dim: int, |
| 260 | output_dtype: torch.dtype = torch.bfloat16, |
| 261 | ) -> torch.Tensor: |
| 262 | """Dequantize FP8 tensor with per-channel scale back to output_dtype. |
| 263 | |
| 264 | For 3D tensor (num_loras, N, K) with channel_dim=1: |
| 265 | scale shape is (num_loras, N), broadcast over K. |
| 266 | """ |
| 267 | expand_scale = scale.float() |
| 268 | # Add trailing dims for broadcasting |
| 269 | for _ in range(channel_dim + 1, fp8_tensor.ndim): |
| 270 | expand_scale = expand_scale.unsqueeze(-1) |
| 271 | return (fp8_tensor.float() * expand_scale).to(output_dtype) |
| 272 | |
| 273 | |
| 274 | def dequantize_fp8_per_token( |
| 275 | fp8_tensor: torch.Tensor, |
| 276 | scale: torch.Tensor, |
| 277 | output_dtype: torch.dtype = torch.bfloat16, |
| 278 | ) -> torch.Tensor: |
| 279 | """Dequantize FP8 2D tensor with per-token scale back to output_dtype. |
| 280 | |
| 281 | fp8_tensor: (num_tokens, hidden_size), scale: (num_tokens, 1) |
| 282 | """ |
| 283 | return (fp8_tensor.float() * scale.float()).to(output_dtype) |
| 284 | |
| 285 | |
| 286 | def dequantize_fp8_blockwise( |
| 287 | fp8_tensor: torch.Tensor, |
| 288 | scale: torch.Tensor, |
| 289 | group_n: int, |
| 290 | group_k: int, |
| 291 | output_dtype: torch.dtype = torch.bfloat16, |
| 292 | ) -> torch.Tensor: |
| 293 | """Dequantize FP8 tensor with block-wise scale back to output_dtype.""" |
| 294 | if fp8_tensor.ndim == 2: |
| 295 | M, K = fp8_tensor.shape |
| 296 | n_blocks_k = math.ceil(K / group_k) |
| 297 | padded = F.pad(fp8_tensor.float(), (0, n_blocks_k * group_k - K)) |
| 298 | blocks = padded.view(M, n_blocks_k, group_k) |
| 299 | out = (blocks * scale.float().unsqueeze(-1)).to(output_dtype) |
| 300 | return out.view(M, -1)[:, :K].contiguous() |
| 301 | elif fp8_tensor.ndim == 3: |
| 302 | L, N, K = fp8_tensor.shape |
| 303 | n_blocks_n = math.ceil(N / group_n) |
| 304 | n_blocks_k = math.ceil(K / group_k) |
| 305 | padded = F.pad( |
| 306 | fp8_tensor.float(), |
| 307 | (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N), |
| 308 | ) |
| 309 | blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k) |
| 310 | out = (blocks * scale.float()[:, :, None, :, None]).to(output_dtype) |
| 311 | out = out.view(L, n_blocks_n * group_n, n_blocks_k * group_k) |
| 312 | return out[:, :N, :K].contiguous() |
| 313 | else: |
| 314 | raise ValueError(f"Unsupported tensor ndim: {fp8_tensor.ndim}") |
| 315 | |
| 316 | |
| 317 | # ============================================================================ |
| 318 | # FP8 Data Generation |
| 319 | # ============================================================================ |
| 320 | |
| 321 | |
| 322 | def generate_fp8_shrink_data( |
| 323 | batches: int, |
| 324 | hidden_size: int, |
| 325 | num_loras: int, |
| 326 | rank: int, |
| 327 | seq_length: int, |
| 328 | nslices: int, |
| 329 | dtype: torch.dtype, |
| 330 | device: str, |
| 331 | quant_mode: str, # "per_tensor", "per_channel", "blockwise" |
| 332 | group_k: int = 128, |
| 333 | group_n: int = 128, |
| 334 | ): |
| 335 | """Generate test data for FP8 shrink kernel. |
| 336 | |
| 337 | Shrink: output = input @ lora_a^T * scaling |
| 338 | input: (num_tokens, hidden_size) -> quantized to FP8 |
| 339 | lora_a: (num_loras, rank, hidden_size) -> quantized to FP8 |
| 340 | |
| 341 | Returns bf16 reference tensors, FP8 quantized tensors with scales, |
| 342 | and dequantized bf16 tensors for accurate reference computation. |
| 343 | """ |
| 344 | seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device) |
| 345 | b_seq_start_loc = torch.cumsum( |
| 346 | torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long), |
| 347 | dim=0, |
| 348 | ).to(device) |
| 349 | total_tokens = seq_len_tensor.sum().item() |
| 350 | |
| 351 | # Generate bf16 reference data |
| 352 | inputs_bf16 = torch.randn(total_tokens, hidden_size, dtype=dtype, device=device) |
| 353 | |
| 354 | lora_a_weights_bf16 = [] |
| 355 | for _ in range(nslices): |
| 356 | lora_a_weights_bf16.append( |
| 357 | torch.randn(num_loras, rank, hidden_size, dtype=dtype, device=device) |
| 358 | ) |
| 359 | |
| 360 | # Quantize inputs to FP8 and dequantize back for reference |
| 361 | if quant_mode == "blockwise": |
| 362 | inputs_fp8, a_scale = quantize_to_fp8_blockwise( |
| 363 | inputs_bf16, group_n=1, group_k=group_k |
| 364 | ) |
| 365 | inputs_dequant = dequantize_fp8_blockwise( |
| 366 | inputs_fp8, |
| 367 | a_scale, |
| 368 | group_n=1, |
| 369 | group_k=group_k, |
| 370 | output_dtype=dtype, |
| 371 | ) |
| 372 | elif quant_mode == "per_tensor": |
| 373 | # Per-tensor: kernel loads a single scalar from a_scale_ptr |
| 374 | inputs_fp8, a_scale = quantize_to_fp8_per_tensor(inputs_bf16) |
| 375 | inputs_dequant = dequantize_fp8_per_tensor( |
| 376 | inputs_fp8, |
| 377 | a_scale, |
| 378 | output_dtype=dtype, |
| 379 | ) |
| 380 | else: |
| 381 | # per_channel: kernel loads per-token a_scale via ram indexing |
| 382 | inputs_fp8, a_scale = quantize_to_fp8_per_token(inputs_bf16) |
| 383 | inputs_dequant = dequantize_fp8_per_token( |
| 384 | inputs_fp8, |
| 385 | a_scale, |
| 386 | output_dtype=dtype, |
| 387 | ) |
| 388 | |
| 389 | # Quantize lora_a weights to FP8 and dequantize back for reference |
| 390 | b_scales = [] |
| 391 | lora_a_weights_fp8 = [] |
| 392 | lora_a_weights_dequant = [] |
| 393 | for w in lora_a_weights_bf16: |
| 394 | if quant_mode == "per_tensor": |
| 395 | w_fp8, w_scale = quantize_to_fp8_per_tensor(w) |
| 396 | w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype) |
| 397 | # Scale shape: (1,) -> need (num_loras,) for the kernel |
| 398 | w_scale = w_scale.expand(num_loras).contiguous() |
| 399 | lora_a_weights_fp8.append(w_fp8) |
| 400 | b_scales.append(w_scale) |
| 401 | lora_a_weights_dequant.append(w_dequant) |
| 402 | elif quant_mode == "per_channel": |
| 403 | # Per-channel along rank dim: scale shape (num_loras, rank) |
| 404 | w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1) |
| 405 | w_dequant = dequantize_fp8_per_channel( |
| 406 | w_fp8, |
| 407 | w_scale, |
| 408 | channel_dim=1, |
| 409 | output_dtype=dtype, |
| 410 | ) |
| 411 | lora_a_weights_fp8.append(w_fp8) |
| 412 | b_scales.append(w_scale) |
| 413 | lora_a_weights_dequant.append(w_dequant) |
| 414 | elif quant_mode == "blockwise": |
| 415 | w_fp8, w_scale = quantize_to_fp8_blockwise( |
| 416 | w, group_n=group_n, group_k=group_k |
| 417 | ) |
| 418 | w_dequant = dequantize_fp8_blockwise( |
| 419 | w_fp8, |
| 420 | w_scale, |
| 421 | group_n=group_n, |
| 422 | group_k=group_k, |
| 423 | output_dtype=dtype, |
| 424 | ) |
| 425 | lora_a_weights_fp8.append(w_fp8) |
| 426 | b_scales.append(w_scale) |
| 427 | lora_a_weights_dequant.append(w_dequant) |
| 428 | |
| 429 | # Output tensor (float32 for shrink) |
| 430 | out_tensor = torch.zeros( |
| 431 | nslices, total_tokens, rank, dtype=torch.float32, device=device |
| 432 | ) |
| 433 | ref_out_tensor = out_tensor.clone() |
| 434 | |
| 435 | # Token-to-lora mapping |
| 436 | lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device) |
| 437 | token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device) |
| 438 | current_offset = 0 |
| 439 | for b_id in range(batches): |
| 440 | lora_index = lora_indices_tensor[b_id] |
| 441 | sl = seq_len_tensor[b_id].item() |
| 442 | token_lora_mapping[current_offset : current_offset + sl] = lora_index |
| 443 | current_offset += sl |
| 444 | |
| 445 | return { |
| 446 | "inputs_bf16": inputs_bf16, |
| 447 | "inputs_fp8": inputs_fp8, |
| 448 | "inputs_dequant": inputs_dequant, |
| 449 | "lora_a_bf16": lora_a_weights_bf16, |
| 450 | "lora_a_fp8": lora_a_weights_fp8, |
| 451 | "lora_a_dequant": lora_a_weights_dequant, |
| 452 | "a_scale": a_scale, |
| 453 | "b_scales": b_scales, |
| 454 | "out_tensor": out_tensor, |
| 455 | "ref_out_tensor": ref_out_tensor, |
| 456 | "token_lora_mapping": token_lora_mapping, |
| 457 | "seq_len_tensor": seq_len_tensor, |
| 458 | "b_seq_start_loc": b_seq_start_loc, |
| 459 | "lora_indices_tensor": lora_indices_tensor, |
| 460 | "total_tokens": total_tokens, |
| 461 | } |
| 462 | |
| 463 | |
| 464 | def generate_fp8_expand_data( |
| 465 | batches: int, |
| 466 | hidden_size: int, |
| 467 | num_loras: int, |
| 468 | rank: int, |
| 469 | seq_length: int, |
| 470 | nslices: int, |
| 471 | dtype: torch.dtype, |
| 472 | device: str, |
| 473 | quant_mode: str, # "per_tensor", "per_channel", "blockwise" |
| 474 | group_k: int = 128, |
| 475 | group_n: int = 128, |
| 476 | ): |
| 477 | """Generate test data for FP8 expand kernel (w8a8). |
| 478 | |
| 479 | Expand: output += input @ lora_b^T |
| 480 | input: (nslices, num_tokens, rank) -> quantized to FP8 (activations) |
| 481 | lora_b: (num_loras, hidden_size, rank) -> quantized to FP8 (weights) |
| 482 | |
| 483 | In w8a8 mode, both activations and weights are FP8. |
| 484 | Returns bf16 reference tensors, FP8 quantized tensors with scales, |
| 485 | and dequantized bf16 tensors for accurate reference computation. |
| 486 | """ |
| 487 | seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device) |
| 488 | b_seq_start_loc = torch.cumsum( |
| 489 | torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long), |
| 490 | dim=0, |
| 491 | ).to(device) |
| 492 | total_tokens = seq_len_tensor.sum().item() |
| 493 | |
| 494 | # Generate bf16 input (shrink output) and quantize to FP8 |
| 495 | inputs_bf16 = torch.randn(nslices, total_tokens, rank, dtype=dtype, device=device) |
| 496 | |
| 497 | # Quantize input to FP8 and dequantize back for reference |
| 498 | inputs_2d_all = inputs_bf16.reshape(-1, rank) |
| 499 | if quant_mode == "blockwise": |
| 500 | # For blockwise, the kernel indexes a_scale by token id (0..total_tokens-1) |
| 501 | # shared across slices. Compute shared scale across slices, then quantize. |
| 502 | # First compute per-token-per-block scale across all slices |
| 503 | n_blocks_k = math.ceil(rank / group_k) |
| 504 | padded = F.pad(inputs_bf16.float(), (0, n_blocks_k * group_k - rank)) |
| 505 | blocks = padded.view(nslices, total_tokens, n_blocks_k, group_k) |
| 506 | # Take the block amax across all slices so every slice shares the scale. |
| 507 | a_scale = blocks.abs().amax(dim=-1).clamp(min=1e-12).amax(dim=0) / FP8_MAX |
| 508 | fp8_blocks = ( |
| 509 | (blocks / a_scale[None, :, :, None]).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) |
| 510 | ) |
| 511 | inputs_fp8 = fp8_blocks.view(nslices, total_tokens, -1)[ |
| 512 | :, :, :rank |
| 513 | ].contiguous() |
| 514 | inputs_dequant = ( |
| 515 | (fp8_blocks.float() * a_scale[None, :, :, None]) |
| 516 | .to(dtype) |
| 517 | .view(nslices, total_tokens, -1)[:, :, :rank] |
| 518 | .contiguous() |
| 519 | ) |
| 520 | elif quant_mode == "per_tensor": |
| 521 | # Per-tensor: kernel loads a single scalar from a_scale_ptr |
| 522 | inputs_fp8_2d, a_scale = quantize_to_fp8_per_tensor(inputs_2d_all) |
| 523 | inputs_dequant_2d = dequantize_fp8_per_tensor( |
| 524 | inputs_fp8_2d, |
| 525 | a_scale, |
| 526 | output_dtype=dtype, |
| 527 | ) |
| 528 | inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank) |
| 529 | inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank) |
| 530 | else: |
| 531 | # per_channel: kernel loads per-token a_scale via ram indexing. |
| 532 | # The kernel uses the same a_scale for all slices (indexed by token |
| 533 | # id 0..total_tokens-1), so we compute a shared per-token scale |
| 534 | # across all slices, then quantize each slice with that shared scale. |
| 535 | per_slice_views = [inputs_bf16[s] for s in range(nslices)] |
| 536 | # (nslices, total_tokens, rank) -> max across slices per token |
| 537 | stacked = torch.stack(per_slice_views, dim=0) # (nslices, tokens, rank) |
| 538 | amax = stacked.abs().float().amax(dim=(0, 2), keepdim=False).clamp(min=1e-12) |
| 539 | # amax shape: (total_tokens,) |
| 540 | a_scale = (amax / FP8_MAX).to(torch.float32).unsqueeze(1) # (tokens, 1) |
| 541 | # Quantize all slices with the shared scale |
| 542 | inputs_fp8_2d = ( |
| 543 | (inputs_2d_all.float() / a_scale.repeat(nslices, 1)) |
| 544 | .clamp(FP8_MIN, FP8_MAX) |
| 545 | .to(FP8_DTYPE) |
| 546 | ) |
| 547 | inputs_dequant_2d = ( |
| 548 | inputs_fp8_2d.float() * a_scale.repeat(nslices, 1).float() |
| 549 | ).to(dtype) |
| 550 | inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank) |
| 551 | inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank) |
| 552 | |
| 553 | # Generate bf16 LoRA B weights |
| 554 | lora_b_weights_bf16 = [] |
| 555 | for _ in range(nslices): |
| 556 | lora_b_weights_bf16.append( |
| 557 | torch.randn(num_loras, hidden_size, rank, dtype=dtype, device=device) |
| 558 | ) |
| 559 | |
| 560 | # Quantize LoRA B weights to FP8 and dequantize back for reference |
| 561 | b_scales = [] |
| 562 | lora_b_weights_fp8 = [] |
| 563 | lora_b_weights_dequant = [] |
| 564 | for w in lora_b_weights_bf16: |
| 565 | if quant_mode == "per_tensor": |
| 566 | w_fp8, w_scale = quantize_to_fp8_per_tensor(w) |
| 567 | w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype) |
| 568 | w_scale = w_scale.expand(num_loras).contiguous() |
| 569 | lora_b_weights_fp8.append(w_fp8) |
| 570 | b_scales.append(w_scale) |
| 571 | lora_b_weights_dequant.append(w_dequant) |
| 572 | elif quant_mode == "per_channel": |
| 573 | # Per-channel along hidden_size dim: scale (num_loras, hidden_size) |
| 574 | w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1) |
| 575 | w_dequant = dequantize_fp8_per_channel( |
| 576 | w_fp8, |
| 577 | w_scale, |
| 578 | channel_dim=1, |
| 579 | output_dtype=dtype, |
| 580 | ) |
| 581 | lora_b_weights_fp8.append(w_fp8) |
| 582 | b_scales.append(w_scale) |
| 583 | lora_b_weights_dequant.append(w_dequant) |
| 584 | elif quant_mode == "blockwise": |
| 585 | w_fp8, w_scale = quantize_to_fp8_blockwise( |
| 586 | w, group_n=group_n, group_k=group_k |
| 587 | ) |
| 588 | w_dequant = dequantize_fp8_blockwise( |
| 589 | w_fp8, |
| 590 | w_scale, |
| 591 | group_n=group_n, |
| 592 | group_k=group_k, |
| 593 | output_dtype=dtype, |
| 594 | ) |
| 595 | lora_b_weights_fp8.append(w_fp8) |
| 596 | b_scales.append(w_scale) |
| 597 | lora_b_weights_dequant.append(w_dequant) |
| 598 | |
| 599 | # Output tensor (initialized randomly for add_inputs) |
| 600 | out_tensor = torch.randn( |
| 601 | total_tokens, hidden_size * nslices, dtype=dtype, device=device |
| 602 | ) |
| 603 | ref_out_tensor = out_tensor.clone() |
| 604 | |
| 605 | # Token-to-lora mapping |
| 606 | lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device) |
| 607 | token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device) |
| 608 | current_offset = 0 |
| 609 | for b_id in range(batches): |
| 610 | lora_index = lora_indices_tensor[b_id] |
| 611 | sl = seq_len_tensor[b_id].item() |
| 612 | token_lora_mapping[current_offset : current_offset + sl] = lora_index |
| 613 | current_offset += sl |
| 614 | |
| 615 | return { |
| 616 | "inputs_bf16": inputs_bf16, |
| 617 | "inputs_fp8": inputs_fp8, |
| 618 | "inputs_dequant": inputs_dequant, |
| 619 | "a_scale": a_scale, |
| 620 | "lora_b_bf16": lora_b_weights_bf16, |
| 621 | "lora_b_fp8": lora_b_weights_fp8, |
| 622 | "lora_b_dequant": lora_b_weights_dequant, |
| 623 | "b_scales": b_scales, |
| 624 | "out_tensor": out_tensor, |
| 625 | "ref_out_tensor": ref_out_tensor, |
| 626 | "token_lora_mapping": token_lora_mapping, |
| 627 | "seq_len_tensor": seq_len_tensor, |
| 628 | "b_seq_start_loc": b_seq_start_loc, |
| 629 | "lora_indices_tensor": lora_indices_tensor, |
| 630 | "total_tokens": total_tokens, |
| 631 | } |
| 632 | |
| 633 | |
| 634 | # ============================================================================ |
| 635 | # FP8 Shrink Kernel Check |
| 636 | # ============================================================================ |
| 637 | |
| 638 | |
| 639 | def check_lora_shrink_fp8_kernel( |
| 640 | batches: int, |
| 641 | num_loras: int, |
| 642 | rank: int, |
| 643 | hidden_size: int, |
| 644 | nslices: int, |
| 645 | dtype: torch.dtype, |
| 646 | device: str, |
| 647 | seq_length: int, |
| 648 | scaling: float, |
| 649 | quant_mode: str, |
| 650 | group_k: int = 128, |
| 651 | group_n: int = 128, |
| 652 | ): |
| 653 | """Test FP8 shrink kernel against dequantized bf16 reference. |
| 654 | |
| 655 | Instead of comparing FP8 kernel output against the original bf16 reference |
| 656 | (which conflates quantization error with kernel error), we: |
| 657 | 1. Quantize bf16 inputs/weights to FP8 |
| 658 | 2. Dequantize them back to bf16 |
| 659 | 3. Run the bf16 reference (sgmv_shrink) with the dequantized values |
| 660 | 4. Compare FP8 kernel output against this dequantized reference |
| 661 | |
| 662 | This isolates kernel correctness from quantization precision loss, |
| 663 | allowing much tighter tolerances. |
| 664 | """ |
| 665 | data = generate_fp8_shrink_data( |
| 666 | batches, |
| 667 | hidden_size, |
| 668 | num_loras, |
| 669 | rank, |
| 670 | seq_length, |
| 671 | nslices, |
| 672 | dtype, |
| 673 | device, |
| 674 | quant_mode, |
| 675 | group_k, |
| 676 | group_n, |
| 677 | ) |
| 678 | |
| 679 | total_tokens = data["total_tokens"] |
| 680 | |
| 681 | # Setup LoRA kernel metadata |
| 682 | lora_meta = LoRAKernelMeta.make( |
| 683 | max_loras=num_loras, max_num_tokens=total_tokens, device=device |
| 684 | ) |
| 685 | lora_meta.prepare_tensors(data["token_lora_mapping"]) |
| 686 | |
| 687 | out_tensor = data["out_tensor"] |
| 688 | |
| 689 | # Determine quantization params for the kernel |
| 690 | per_channel = quant_mode == "per_channel" |
| 691 | gk = group_k if quant_mode == "blockwise" else 0 |
| 692 | gn = group_n if quant_mode == "blockwise" else 0 |
| 693 | |
| 694 | with _dict_lock: |
| 695 | _LORA_A_PTR_DICT.clear() |
| 696 | _SHRINK_LORA_SCALE_PTR_DICT.clear() |
| 697 | triton_ops.lora_shrink_fp8( |
| 698 | data["inputs_fp8"], |
| 699 | data["lora_a_fp8"], |
| 700 | out_tensor, |
| 701 | *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False), |
| 702 | scaling, |
| 703 | data["b_scales"], |
| 704 | a_scale=data["a_scale"], |
| 705 | group_k=gk, |
| 706 | group_n=gn, |
| 707 | use_fp8_w8a8=True, |
| 708 | per_channel_quant=per_channel, |
| 709 | ) |
| 710 | |
| 711 | # Compute reference using dequantized (round-tripped) tensors. |
| 712 | # This means the reference sees the same quantization error as the kernel, |
| 713 | # so any difference is purely kernel error. |
| 714 | ref_out_tensor = data["ref_out_tensor"] |
| 715 | max_seq_length = data["seq_len_tensor"].max().item() |
| 716 | sgmv_shrink_for_nslices( |
| 717 | nslices, |
| 718 | data["inputs_dequant"], |
| 719 | data["lora_a_dequant"], |
| 720 | ref_out_tensor, |
| 721 | data["b_seq_start_loc"], |
| 722 | data["seq_len_tensor"], |
| 723 | data["lora_indices_tensor"], |
| 724 | batches, |
| 725 | max_seq_length, |
| 726 | total_tokens, |
| 727 | scaling, |
| 728 | ) |
| 729 | |
| 730 | # With dequantized reference, we can use much tighter tolerances |
| 731 | # since we're only measuring kernel error, not quantization error. |
| 732 | # Blockwise accumulation order differs from the bf16 reference, so |
| 733 | # allow a slightly larger margin for sporadic rounding outliers. |
| 734 | rtol, atol = 0.1, 0.25 |
| 735 | torch.testing.assert_close( |
| 736 | out_tensor.to(dtype), ref_out_tensor.to(dtype), rtol=rtol, atol=atol |
| 737 | ) |
| 738 | |
| 739 | |
| 740 | # ============================================================================ |
| 741 | # FP8 Expand Kernel Check |
| 742 | # ============================================================================ |
| 743 | |
| 744 | |
| 745 | def check_lora_expand_fp8_kernel( |
| 746 | batches: int, |
| 747 | num_loras: int, |
| 748 | rank: int, |
| 749 | hidden_size: int, |
| 750 | nslices: int, |
| 751 | dtype: torch.dtype, |
| 752 | device: str, |
| 753 | seq_length: int, |
| 754 | add_inputs: bool, |
| 755 | quant_mode: str, |
| 756 | group_k: int = 128, |
| 757 | group_n: int = 128, |
| 758 | ): |
| 759 | """Test FP8 expand kernel (w8a8) against dequantized bf16 reference. |
| 760 | |
| 761 | Instead of comparing FP8 kernel output against the original bf16 reference |
| 762 | (which conflates quantization error with kernel error), we: |
| 763 | 1. Quantize bf16 inputs/weights to FP8 |
| 764 | 2. Dequantize them back to bf16 |
| 765 | 3. Run the bf16 reference (sgmv_expand) with the dequantized values |
| 766 | 4. Compare FP8 kernel output against this dequantized reference |
| 767 | |
| 768 | This isolates kernel correctness from quantization precision loss, |
| 769 | allowing much tighter tolerances. |
| 770 | """ |
| 771 | data = generate_fp8_expand_data( |
| 772 | batches, |
| 773 | hidden_size, |
| 774 | num_loras, |
| 775 | rank, |
| 776 | seq_length, |
| 777 | nslices, |
| 778 | dtype, |
| 779 | device, |
| 780 | quant_mode, |
| 781 | group_k, |
| 782 | group_n, |
| 783 | ) |
| 784 | |
| 785 | total_tokens = data["total_tokens"] |
| 786 | |
| 787 | # Setup LoRA kernel metadata |
| 788 | lora_meta = LoRAKernelMeta.make( |
| 789 | max_loras=num_loras, max_num_tokens=total_tokens, device=device |
| 790 | ) |
| 791 | lora_meta.prepare_tensors(data["token_lora_mapping"]) |
| 792 | |
| 793 | out_tensor = data["out_tensor"] |
| 794 | |
| 795 | # Determine quantization params for the kernel |
| 796 | per_channel = quant_mode == "per_channel" |
| 797 | gk = group_k if quant_mode == "blockwise" else 0 |
| 798 | gn = group_n if quant_mode == "blockwise" else 0 |
| 799 | |
| 800 | with _dict_lock: |
| 801 | _LORA_B_PTR_DICT.clear() |
| 802 | _EXPAND_LORA_SCALE_PTR_DICT.clear() |
| 803 | triton_ops.lora_expand_fp8( |
| 804 | data["inputs_fp8"], |
| 805 | data["lora_b_fp8"], |
| 806 | out_tensor, |
| 807 | *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False), |
| 808 | data["b_scales"], |
| 809 | a_scale=data["a_scale"], |
| 810 | offset_start=0, |
| 811 | add_inputs=add_inputs, |
| 812 | group_k=gk, |
| 813 | group_n=gn, |
| 814 | use_fp8_w8a8=True, |
| 815 | per_channel_quant=per_channel, |
| 816 | ) |
| 817 | |
| 818 | # Compute reference using dequantized (round-tripped) tensors. |
| 819 | ref_out_tensor = data["ref_out_tensor"] |
| 820 | max_seq_length = data["seq_len_tensor"].max().item() |
| 821 | sgmv_expand_for_nslices( |
| 822 | nslices, |
| 823 | hidden_size, |
| 824 | data["inputs_dequant"], |
| 825 | data["lora_b_dequant"], |
| 826 | ref_out_tensor, |
| 827 | data["b_seq_start_loc"], |
| 828 | data["seq_len_tensor"], |
| 829 | data["lora_indices_tensor"], |
| 830 | batches, |
| 831 | max_seq_length, |
| 832 | total_tokens, |
| 833 | add_inputs=add_inputs, |
| 834 | ) |
| 835 | |
| 836 | # With dequantized reference, we can use much tighter tolerances |
| 837 | # since we're only measuring kernel error, not quantization error. |
| 838 | rtol, atol = 0.1, 0.15 |
| 839 | torch.testing.assert_close(out_tensor, ref_out_tensor, rtol=rtol, atol=atol) |
| 840 | |
| 841 | |
| 842 | # ============================================================================ |
| 843 | # FP8 Test Parameters |
| 844 | # ============================================================================ |
| 845 | |
| 846 | fp8_test_params = { |
| 847 | "hidden_sizes": [512, 1024, 2048], |
| 848 | "batches": [1, 4, 16], |
| 849 | "num_loras": [1, 4, 8], |
| 850 | "max_ranks": [8, 16, 32, 64], |
| 851 | } |
| 852 | |
| 853 | |
| 854 | # ============================================================================ |
| 855 | # FP8 Shrink Tests |
| 856 | # ============================================================================ |
| 857 | |
| 858 | |
| 859 | @pytest.mark.parametrize("batches", fp8_test_params["batches"]) |
| 860 | @pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"]) |
| 861 | @pytest.mark.parametrize("rank", fp8_test_params["max_ranks"]) |
| 862 | @pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"]) |
| 863 | @pytest.mark.parametrize("nslices", [1, 2, 3]) |
| 864 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 865 | @pytest.mark.parametrize("device", DEVICES) |
| 866 | @pytest.mark.parametrize("seed", SEED) |
| 867 | @pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"]) |
| 868 | def test_lora_shrink_fp8( |
| 869 | batches: int, |
| 870 | num_loras: int, |
| 871 | rank: int, |
| 872 | hidden_size: int, |
| 873 | nslices: int, |
| 874 | dtype: torch.dtype, |
| 875 | device: str, |
| 876 | seed: int, |
| 877 | quant_mode: str, |
| 878 | ): |
| 879 | """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise |
| 880 | quantization, comparing against the bf16 baseline.""" |
| 881 | torch.set_default_device(device) |
| 882 | set_random_seed(seed) |
| 883 | |
| 884 | # For blockwise, group sizes must divide evenly or be handled by the kernel |
| 885 | group_k = 128 |
| 886 | group_n = 128 |
| 887 | |
| 888 | # Adjust group sizes if they're larger than the dimensions |
| 889 | if quant_mode == "blockwise": |
| 890 | group_k = min(group_k, hidden_size) |
| 891 | group_n = min(group_n, rank) |
| 892 | |
| 893 | check_lora_shrink_fp8_kernel( |
| 894 | batches=batches, |
| 895 | num_loras=num_loras, |
| 896 | rank=rank, |
| 897 | hidden_size=hidden_size, |
| 898 | nslices=nslices, |
| 899 | dtype=dtype, |
| 900 | device=device, |
| 901 | seq_length=128, |
| 902 | scaling=0.5, |
| 903 | quant_mode=quant_mode, |
| 904 | group_k=group_k, |
| 905 | group_n=group_n, |
| 906 | ) |
| 907 | |
| 908 | |
| 909 | # ============================================================================ |
| 910 | # FP8 Expand Tests |
| 911 | # ============================================================================ |
| 912 | |
| 913 | |
| 914 | @pytest.mark.parametrize("batches", fp8_test_params["batches"]) |
| 915 | @pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"]) |
| 916 | @pytest.mark.parametrize("rank", fp8_test_params["max_ranks"]) |
| 917 | @pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"]) |
| 918 | @pytest.mark.parametrize("nslices", [1, 2, 3]) |
| 919 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 920 | @pytest.mark.parametrize("device", DEVICES) |
| 921 | @pytest.mark.parametrize("seed", SEED) |
| 922 | @pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"]) |
| 923 | def test_lora_expand_fp8( |
| 924 | batches: int, |
| 925 | num_loras: int, |
| 926 | rank: int, |
| 927 | hidden_size: int, |
| 928 | nslices: int, |
| 929 | dtype: torch.dtype, |
| 930 | device: str, |
| 931 | seed: int, |
| 932 | quant_mode: str, |
| 933 | ): |
| 934 | """Test FP8 expand kernel with per-tensor, per-channel, and block-wise |
| 935 | quantization, comparing against the bf16 baseline.""" |
| 936 | torch.set_default_device(device) |
| 937 | set_random_seed(seed) |
| 938 | |
| 939 | group_k = 128 |
| 940 | group_n = 128 |
| 941 | |
| 942 | # Adjust group sizes if they're larger than the dimensions |
| 943 | if quant_mode == "blockwise": |
| 944 | group_k = min(group_k, rank) |
| 945 | group_n = min(group_n, hidden_size) |
| 946 | |
| 947 | check_lora_expand_fp8_kernel( |
| 948 | batches=batches, |
| 949 | num_loras=num_loras, |
| 950 | rank=rank, |
| 951 | hidden_size=hidden_size, |
| 952 | nslices=nslices, |
| 953 | dtype=dtype, |
| 954 | device=device, |
| 955 | seq_length=128, |
| 956 | add_inputs=True, |
| 957 | quant_mode=quant_mode, |
| 958 | group_k=group_k, |
| 959 | group_n=group_n, |
| 960 | ) |