| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.model_executor.layers.fusion.quant_activation import ( |
| 7 | QuantizedActivation, |
| 8 | get_input_quant_key, |
| 9 | ) |
| 10 | from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | |
| 14 | |
| 15 | @triton.jit |
| 16 | def _norm_row(x, weight, row, stride, eps, SIZE: tl.constexpr, BLOCK: tl.constexpr): |
| 17 | cols = tl.arange(0, BLOCK) |
| 18 | w = tl.load(weight + cols, cols < SIZE, 0).to(tl.float32) |
| 19 | values = tl.load(x + row * stride + cols, cols < SIZE, 0).to(tl.float32) |
| 20 | rrms = tl.rsqrt(tl.sum(values * values, 0) / SIZE + eps) |
| 21 | return values * rrms * w |
| 22 | |
| 23 | |
| 24 | @triton.jit(do_not_specialize=["num_tokens"]) |
| 25 | def _q_kv_norm_quant_kernel( |
| 26 | q, |
| 27 | kv, |
| 28 | qw, |
| 29 | kvw, |
| 30 | qo, |
| 31 | kvo, |
| 32 | scales, |
| 33 | num_tokens, |
| 34 | q_stride, |
| 35 | kv_stride, |
| 36 | eps, |
| 37 | Q_SIZE: tl.constexpr, |
| 38 | KV_SIZE: tl.constexpr, |
| 39 | BLOCK: tl.constexpr, |
| 40 | LAUNCH_PDL: tl.constexpr, |
| 41 | ): |
| 42 | row = tl.program_id(0).to(tl.int64) |
| 43 | cols = tl.arange(0, BLOCK) |
| 44 | if LAUNCH_PDL: |
| 45 | tl.extra.cuda.gdc_wait() |
| 46 | tl.extra.cuda.gdc_launch_dependents() |
| 47 | if tl.program_id(1) == 0: |
| 48 | groups = tl.arange(0, BLOCK // 32) |
| 49 | if row < num_tokens: |
| 50 | y = _norm_row(q, qw, row, q_stride, eps, Q_SIZE, BLOCK) |
| 51 | # Preserve the materialized normalization output's rounding boundary. |
| 52 | y = y.to(q.dtype.element_ty).to(tl.float32) |
| 53 | grouped = tl.reshape(y, (BLOCK // 32, 32)) |
| 54 | amax = tl.max(tl.abs(grouped), 1) |
| 55 | normalized = amax * (1.0 / 448.0) |
| 56 | bits = normalized.to(tl.uint32, bitcast=True) |
| 57 | exponent = (bits >> 23) & 255 |
| 58 | mantissa = bits & 0x7FFFFF |
| 59 | bump = (mantissa != 0) & ~((exponent == 0) & (mantissa <= 0x400000)) |
| 60 | sf = tl.minimum(exponent + bump, 254) |
| 61 | sf = tl.where(normalized <= 0, 0, sf) |
| 62 | # Match FlashInfer's UE8M0 conversion, including zero/subnormal scales. |
| 63 | inv_bits = tl.where(sf == 0, 0, (254 - sf) << 23) |
| 64 | inv_scale = inv_bits.to(tl.float32, bitcast=True) |
| 65 | quantized = tl.reshape(grouped * inv_scale[:, None], (BLOCK,)) |
| 66 | tl.store(qo + row * Q_SIZE + cols, quantized, cols < Q_SIZE) |
| 67 | else: |
| 68 | sf = tl.full((BLOCK // 32,), 0, tl.uint32) |
| 69 | padded_groups: tl.constexpr = triton.cdiv(Q_SIZE // 32, 4) * 4 |
| 70 | sf = tl.where(groups < Q_SIZE // 32, sf, 0) |
| 71 | # F8_128x4: [row/128, group/4, row%32, row%128/32, group%4]. |
| 72 | offsets = ( |
| 73 | row // 128 * (128 * padded_groups) |
| 74 | + groups // 4 * 512 |
| 75 | + row % 32 * 16 |
| 76 | + row % 128 // 32 * 4 |
| 77 | + groups % 4 |
| 78 | ) |
| 79 | tl.store(scales + offsets, sf, groups < padded_groups) |
| 80 | elif row < num_tokens: |
| 81 | y = _norm_row(kv, kvw, row, kv_stride, eps, KV_SIZE, BLOCK) |
| 82 | tl.store(kvo + row * KV_SIZE + cols, y, cols < KV_SIZE) |
| 83 | |
| 84 | |
| 85 | def fused_q_kv_rmsnorm_quant( |
| 86 | qr: torch.Tensor, |
| 87 | kv: torch.Tensor, |
| 88 | q_weight: torch.Tensor, |
| 89 | kv_weight: torch.Tensor, |
| 90 | eps: float, |
| 91 | ) -> tuple[QuantizedActivation, torch.Tensor]: |
| 92 | """Normalize Q/KV and quantize Q with FlashInfer's swizzled MXFP8 scales.""" |
| 93 | assert qr.ndim == kv.ndim == 2 and qr.shape[0] == kv.shape[0] |
| 94 | assert qr.stride(-1) == kv.stride(-1) == 1 |
| 95 | assert q_weight.is_contiguous() and kv_weight.is_contiguous() |
| 96 | assert qr.shape[1] % 32 == 0 |
| 97 | tokens, q_size = qr.shape |
| 98 | kv_size = kv.shape[1] |
| 99 | qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device) |
| 100 | kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device) |
| 101 | padded_tokens = triton.cdiv(tokens, 128) * 128 |
| 102 | padded_groups = triton.cdiv(q_size // 32, 4) * 4 |
| 103 | scales = torch.empty( |
| 104 | padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device |
| 105 | ) |
| 106 | if tokens: |
| 107 | block = triton.next_power_of_2(max(q_size, kv_size)) |
| 108 | launch_pdl = current_platform.is_arch_support_pdl() |
| 109 | _q_kv_norm_quant_kernel[(padded_tokens, 2)]( |
| 110 | qr, |
| 111 | kv, |
| 112 | q_weight, |
| 113 | kv_weight, |
| 114 | qo, |
| 115 | kvo, |
| 116 | scales, |
| 117 | tokens, |
| 118 | qr.stride(0), |
| 119 | kv.stride(0), |
| 120 | eps, |
| 121 | q_size, |
| 122 | kv_size, |
| 123 | block, |
| 124 | launch_pdl, |
| 125 | num_warps=8 if block >= 2048 else 4, |
| 126 | launch_pdl=launch_pdl, |
| 127 | ) |
| 128 | return QuantizedActivation(qo, scales, qr.dtype, qr.shape, kMxfp8Dynamic), kvo |
| 129 | |
| 130 | |
| 131 | def can_fuse_query_quant(linears: list[torch.nn.Module]) -> bool: |
| 132 | """Require both local projections to use the same MXFP8 activation ABI.""" |
| 133 | if not current_platform.is_cuda(): |
| 134 | return False |
| 135 | from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( |
| 136 | FlashInferCutedslMxfp8LinearKernel, |
| 137 | FlashInferCutlassMxfp8LinearKernel, |
| 138 | ) |
| 139 | |
| 140 | # QuantKey does not encode scale layout; restrict this producer to the |
| 141 | # consumers that accept F8_128x4 swizzled scales. |
| 142 | return all( |
| 143 | get_input_quant_key(linear) == kMxfp8Dynamic |
| 144 | and type(getattr(getattr(linear, "quant_method", None), "kernel", None)) |
| 145 | in (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel) |
| 146 | for linear in linears |
| 147 | ) |