| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused QK-RMSNorm + (partial) RoPE + gate copy Triton kernel. |
| 4 | |
| 5 | Currently used by the Qwen3.5 attention path (``attn_output_gate`` with |
| 6 | NeoX-style partial RoPE). The unfused reference sequence is |
| 7 | ``split -> GemmaRMSNorm -> RoPE -> gate chunk``; this collapses it into a |
| 8 | single Triton launch. See :func:`fused_qk_rmsnorm_rope_gate`. |
| 9 | """ |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.triton_utils import tl, triton |
| 14 | |
| 15 | |
| 16 | @triton.jit |
| 17 | def _fused_qk_rmsnorm_rope_gate_kernel( |
| 18 | q_gate_ptr, |
| 19 | k_ptr, |
| 20 | q_out_ptr, |
| 21 | k_out_ptr, |
| 22 | gate_out_ptr, |
| 23 | q_weight_ptr, |
| 24 | k_weight_ptr, |
| 25 | cos_sin_cache_ptr, |
| 26 | positions_ptr, |
| 27 | q_gate_stride_t, |
| 28 | k_stride_t, |
| 29 | q_out_stride_t, |
| 30 | k_out_stride_t, |
| 31 | gate_out_stride_t, |
| 32 | cache_stride_p, |
| 33 | positions_stride_m, |
| 34 | positions_stride_t, |
| 35 | num_q_heads: tl.constexpr, |
| 36 | num_kv_heads: tl.constexpr, |
| 37 | head_dim: tl.constexpr, |
| 38 | rotary_dim: tl.constexpr, |
| 39 | half_rotary: tl.constexpr, |
| 40 | eps: tl.constexpr, |
| 41 | norm_beta: tl.constexpr, |
| 42 | INPUT_DTYPE: tl.constexpr, |
| 43 | HEAD_BLOCK: tl.constexpr, |
| 44 | ROT_HALF_BLOCK: tl.constexpr, |
| 45 | HAS_PASS: tl.constexpr, |
| 46 | HAS_MROPE: tl.constexpr, |
| 47 | MROPE_SECTION_H: tl.constexpr, |
| 48 | MROPE_SECTION_W: tl.constexpr, |
| 49 | ): |
| 50 | token = tl.program_id(0) |
| 51 | head = tl.program_id(1) |
| 52 | is_k = head >= num_q_heads |
| 53 | local_head = tl.where(is_k, head - num_q_heads, head) |
| 54 | |
| 55 | if is_k: |
| 56 | in_base = k_ptr + token * k_stride_t + local_head * head_dim |
| 57 | w_ptr = k_weight_ptr |
| 58 | out_base = k_out_ptr + token * k_out_stride_t + local_head * head_dim |
| 59 | else: |
| 60 | in_base = q_gate_ptr + token * q_gate_stride_t + local_head * 2 * head_dim |
| 61 | w_ptr = q_weight_ptr |
| 62 | out_base = q_out_ptr + token * q_out_stride_t + local_head * head_dim |
| 63 | |
| 64 | # --- RMSNorm: variance over the full head_dim --- |
| 65 | head_offs = tl.arange(0, HEAD_BLOCK) |
| 66 | head_mask = head_offs < head_dim |
| 67 | x = tl.load(in_base + head_offs, mask=head_mask, other=0.0).to(tl.float32) |
| 68 | var = tl.sum(x * x, axis=0) / head_dim |
| 69 | inv_rms = tl.rsqrt(var + eps) |
| 70 | w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32) + norm_beta |
| 71 | # Round-trip through INPUT_DTYPE so the RoPE input matches the bf16-storage |
| 72 | # behavior of the unfused (qk_rmsnorm -> memory -> apply_rope) reference path. |
| 73 | x_norm = (x * inv_rms * w).to(INPUT_DTYPE).to(tl.float32) |
| 74 | |
| 75 | # --- Pass-through tail [rotary_dim, head_dim): RMSNorm-only, no rotation --- |
| 76 | # The rotary head [0, rotary_dim) will be overwritten by the RoPE store below. |
| 77 | if HAS_PASS: |
| 78 | pass_mask = head_mask & (head_offs >= rotary_dim) |
| 79 | tl.store(out_base + head_offs, x_norm, mask=pass_mask) |
| 80 | |
| 81 | # --- Partial RoPE on the first rotary_dim elements --- |
| 82 | # Triton lacks easy sub-vector slicing of x_norm, so we recompute the |
| 83 | # normalized rotary halves on a smaller block (next_pow2(half_rotary)). |
| 84 | # The extra ~rotary_dim element reload hits L1, so the cost is negligible. |
| 85 | rot_offs = tl.arange(0, ROT_HALF_BLOCK) |
| 86 | rot_mask = rot_offs < half_rotary |
| 87 | x_rot1 = tl.load(in_base + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) |
| 88 | x_rot2 = tl.load(in_base + half_rotary + rot_offs, mask=rot_mask, other=0.0).to( |
| 89 | tl.float32 |
| 90 | ) |
| 91 | w_rot1 = ( |
| 92 | tl.load(w_ptr + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + norm_beta |
| 93 | ) |
| 94 | w_rot2 = ( |
| 95 | tl.load(w_ptr + half_rotary + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) |
| 96 | + norm_beta |
| 97 | ) |
| 98 | x_rot1 = (x_rot1 * inv_rms * w_rot1).to(INPUT_DTYPE).to(tl.float32) |
| 99 | x_rot2 = (x_rot2 * inv_rms * w_rot2).to(INPUT_DTYPE).to(tl.float32) |
| 100 | |
| 101 | # Always use int64 for positions to avoid overflow in address computation. |
| 102 | pos_t = tl.load(positions_ptr + token * positions_stride_t).to(tl.int64) |
| 103 | if HAS_MROPE: |
| 104 | pos_h = tl.load( |
| 105 | positions_ptr + positions_stride_m + token * positions_stride_t |
| 106 | ).to(tl.int64) |
| 107 | pos_w = tl.load( |
| 108 | positions_ptr + 2 * positions_stride_m + token * positions_stride_t |
| 109 | ).to(tl.int64) |
| 110 | is_h = (rot_offs % 3 == 1) & (rot_offs < 3 * MROPE_SECTION_H) |
| 111 | is_w = (rot_offs % 3 == 2) & (rot_offs < 3 * MROPE_SECTION_W) |
| 112 | pos = tl.where(is_h, pos_h, tl.where(is_w, pos_w, pos_t)) |
| 113 | else: |
| 114 | pos = pos_t |
| 115 | cache_offset = pos * cache_stride_p |
| 116 | cos = tl.load( |
| 117 | cos_sin_cache_ptr + cache_offset + rot_offs, mask=rot_mask, other=0.0 |
| 118 | ).to(tl.float32) |
| 119 | sin = tl.load( |
| 120 | cos_sin_cache_ptr + cache_offset + half_rotary + rot_offs, |
| 121 | mask=rot_mask, |
| 122 | other=0.0, |
| 123 | ).to(tl.float32) |
| 124 | |
| 125 | o1 = x_rot1 * cos - x_rot2 * sin |
| 126 | o2 = x_rot2 * cos + x_rot1 * sin |
| 127 | tl.store(out_base + rot_offs, o1, mask=rot_mask) |
| 128 | tl.store(out_base + half_rotary + rot_offs, o2, mask=rot_mask) |
| 129 | |
| 130 | # --- Gate copy (q heads only, verbatim) --- |
| 131 | if not is_k: |
| 132 | gate_in_base = in_base + head_dim |
| 133 | gate_out_base = gate_out_ptr + token * gate_out_stride_t + local_head * head_dim |
| 134 | g = tl.load(gate_in_base + head_offs, mask=head_mask, other=0.0) |
| 135 | tl.store(gate_out_base + head_offs, g, mask=head_mask) |
| 136 | |
| 137 | |
| 138 | def fused_qk_rmsnorm_rope_gate( |
| 139 | q_gate: torch.Tensor, |
| 140 | k: torch.Tensor, |
| 141 | q_weight: torch.Tensor, |
| 142 | k_weight: torch.Tensor, |
| 143 | cos_sin_cache: torch.Tensor, |
| 144 | positions: torch.Tensor, |
| 145 | eps: float, |
| 146 | num_q_heads: int, |
| 147 | num_kv_heads: int, |
| 148 | head_dim: int, |
| 149 | rotary_dim: int, |
| 150 | mrope_section: list[int] | tuple[int, int, int] | None = None, |
| 151 | norm_beta: float = 0.0, |
| 152 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 153 | """Fused split + QK-RMSNorm + (partial) RoPE + gate copy for Qwen attn. |
| 154 | |
| 155 | Args: |
| 156 | q_gate: (n_tokens, num_q_heads * 2 * head_dim) -- per head: [q|gate] |
| 157 | k: (n_tokens, num_kv_heads * head_dim) |
| 158 | q_weight: (head_dim,) RMSNorm weight |
| 159 | k_weight: (head_dim,) RMSNorm weight |
| 160 | cos_sin_cache: (max_pos, rotary_dim) packed [cos|sin] |
| 161 | positions: (n_tokens,) or (3, n_tokens) int32 or int64 |
| 162 | eps: RMSNorm epsilon |
| 163 | num_q_heads: number of Q heads (after TP split) |
| 164 | num_kv_heads: number of KV heads (after TP split) |
| 165 | head_dim: per-head dimension |
| 166 | rotary_dim: rotary dimension; must be even and <= head_dim |
| 167 | mrope_section: interleaved T/H/W frequency counts for 2D positions |
| 168 | norm_beta: scalar added to the RMSNorm weight |
| 169 | |
| 170 | Returns: |
| 171 | (q_out, k_out, gate_out) -- all contiguous (n_tokens, heads * head_dim). |
| 172 | ``gate_out`` is the raw (pre-sigmoid) gate. |
| 173 | |
| 174 | """ |
| 175 | if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0: |
| 176 | raise ValueError( |
| 177 | f"rotary_dim must be a positive even integer <= head_dim, " |
| 178 | f"got rotary_dim={rotary_dim}, head_dim={head_dim}" |
| 179 | ) |
| 180 | if q_gate.dtype not in (torch.float16, torch.bfloat16) or k.dtype != q_gate.dtype: |
| 181 | raise ValueError( |
| 182 | "q_gate and k must have the same FP16 or BF16 dtype, " |
| 183 | f"got {q_gate.dtype} and {k.dtype}" |
| 184 | ) |
| 185 | for name, tensor in ( |
| 186 | ("q_gate", q_gate), |
| 187 | ("k", k), |
| 188 | ("q_weight", q_weight), |
| 189 | ("k_weight", k_weight), |
| 190 | ("cos_sin_cache", cos_sin_cache), |
| 191 | ): |
| 192 | if tensor.stride(-1) != 1: |
| 193 | raise ValueError(f"{name} must be contiguous in its last dimension") |
| 194 | |
| 195 | if positions.ndim not in (1, 2): |
| 196 | raise ValueError(f"positions must be 1D or 2D, got shape={positions.shape}") |
| 197 | if positions.shape[-1] != q_gate.shape[0]: |
| 198 | raise ValueError( |
| 199 | "positions token dimension must match q_gate, " |
| 200 | f"got {positions.shape[-1]} and {q_gate.shape[0]}" |
| 201 | ) |
| 202 | |
| 203 | has_mrope = positions.ndim == 2 |
| 204 | if has_mrope: |
| 205 | if positions.shape[0] != 3: |
| 206 | raise ValueError( |
| 207 | f"MRoPE positions must have shape (3, n_tokens), got {positions.shape}" |
| 208 | ) |
| 209 | if mrope_section is None or len(mrope_section) != 3: |
| 210 | raise ValueError("mrope_section must contain the T/H/W frequency counts") |
| 211 | if sum(mrope_section) != rotary_dim // 2: |
| 212 | raise ValueError( |
| 213 | "mrope_section must sum to rotary_dim // 2, " |
| 214 | f"got {mrope_section} and rotary_dim={rotary_dim}" |
| 215 | ) |
| 216 | mrope_section_h = mrope_section[1] |
| 217 | mrope_section_w = mrope_section[2] |
| 218 | positions_stride_m, positions_stride_t = positions.stride() |
| 219 | else: |
| 220 | if mrope_section is not None: |
| 221 | raise ValueError("mrope_section requires 2D MRoPE positions") |
| 222 | mrope_section_h = 0 |
| 223 | mrope_section_w = 0 |
| 224 | positions_stride_m = 0 |
| 225 | positions_stride_t = positions.stride(0) |
| 226 | |
| 227 | n_tokens = q_gate.shape[0] |
| 228 | q_out = torch.empty( |
| 229 | (n_tokens, num_q_heads * head_dim), dtype=q_gate.dtype, device=q_gate.device |
| 230 | ) |
| 231 | k_out = torch.empty( |
| 232 | (n_tokens, num_kv_heads * head_dim), dtype=k.dtype, device=k.device |
| 233 | ) |
| 234 | gate_out = torch.empty_like(q_out) |
| 235 | if n_tokens == 0: |
| 236 | return q_out, k_out, gate_out |
| 237 | |
| 238 | half_rotary = rotary_dim // 2 |
| 239 | head_block = triton.next_power_of_2(head_dim) |
| 240 | rot_half_block = triton.next_power_of_2(half_rotary) |
| 241 | num_warps = max(1, head_block // 64) |
| 242 | |
| 243 | grid = (n_tokens, num_q_heads + num_kv_heads) |
| 244 | _fused_qk_rmsnorm_rope_gate_kernel[grid]( |
| 245 | q_gate, |
| 246 | k, |
| 247 | q_out, |
| 248 | k_out, |
| 249 | gate_out, |
| 250 | q_weight, |
| 251 | k_weight, |
| 252 | cos_sin_cache, |
| 253 | positions, |
| 254 | q_gate.stride(0), |
| 255 | k.stride(0), |
| 256 | q_out.stride(0), |
| 257 | k_out.stride(0), |
| 258 | gate_out.stride(0), |
| 259 | cos_sin_cache.stride(0), |
| 260 | positions_stride_m, |
| 261 | positions_stride_t, |
| 262 | num_q_heads, |
| 263 | num_kv_heads, |
| 264 | head_dim, |
| 265 | rotary_dim, |
| 266 | half_rotary, |
| 267 | eps, |
| 268 | norm_beta=norm_beta, |
| 269 | INPUT_DTYPE=tl.bfloat16 if q_gate.dtype == torch.bfloat16 else tl.float16, |
| 270 | HEAD_BLOCK=head_block, |
| 271 | ROT_HALF_BLOCK=rot_half_block, |
| 272 | HAS_PASS=rotary_dim < head_dim, |
| 273 | HAS_MROPE=has_mrope, |
| 274 | MROPE_SECTION_H=mrope_section_h, |
| 275 | MROPE_SECTION_W=mrope_section_w, |
| 276 | num_warps=num_warps, |
| 277 | num_stages=2, |
| 278 | ) |
| 279 | return q_out, k_out, gate_out |