| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | import os |
| 5 | from collections.abc import Callable |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | import vllm.envs as envs |
| 11 | from vllm.model_executor.determinism.batch_invariant_configs import ( |
| 12 | _get_descriptor_matmul_config, |
| 13 | _get_matmul_config, |
| 14 | resolve_tuned_matmul_configs, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.triton_utils import tl, triton |
| 18 | from vllm.utils.mem_utils import get_max_shared_memory_bytes |
| 19 | from vllm.utils.nccl import pin_nccl_env |
| 20 | from vllm.utils.platform_utils import num_compute_units |
| 21 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 22 | |
| 23 | |
| 24 | def _matmul_launch_metadata( |
| 25 | grid: Callable[..., Any], kernel: Any, args: dict[str, Any] |
| 26 | ) -> dict[str, Any]: |
| 27 | ret = {} |
| 28 | m, n, k = args["M"], args["N"], args["K"] |
| 29 | ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]" |
| 30 | |
| 31 | bytes_per_elem = args["c_ptr"].element_size() |
| 32 | ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k |
| 33 | ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n) |
| 34 | return ret |
| 35 | |
| 36 | |
| 37 | @triton.jit |
| 38 | def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M): |
| 39 | group_id = tile_id // num_pid_in_group |
| 40 | first_pid_m = group_id * GROUP_SIZE_M |
| 41 | group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) |
| 42 | pid_m = first_pid_m + (tile_id % group_size_m) |
| 43 | pid_n = (tile_id % num_pid_in_group) // group_size_m |
| 44 | return pid_m, pid_n |
| 45 | |
| 46 | |
| 47 | @triton.jit(launch_metadata=_matmul_launch_metadata) |
| 48 | def matmul_kernel_persistent( |
| 49 | a_ptr, |
| 50 | b_ptr, |
| 51 | c_ptr, # |
| 52 | bias_ptr, |
| 53 | M, |
| 54 | N, |
| 55 | K, # |
| 56 | stride_am, |
| 57 | stride_ak, |
| 58 | stride_bk, |
| 59 | stride_bn, |
| 60 | stride_cm, |
| 61 | stride_cn, |
| 62 | BLOCK_SIZE_M: tl.constexpr, # |
| 63 | BLOCK_SIZE_N: tl.constexpr, # |
| 64 | BLOCK_SIZE_K: tl.constexpr, # |
| 65 | GROUP_SIZE_M: tl.constexpr, # |
| 66 | NUM_SMS: tl.constexpr, # |
| 67 | A_LARGE: tl.constexpr, |
| 68 | B_LARGE: tl.constexpr, |
| 69 | C_LARGE: tl.constexpr, |
| 70 | HAS_BIAS: tl.constexpr, |
| 71 | ): |
| 72 | start_pid = tl.program_id(axis=0) |
| 73 | num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) |
| 74 | num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) |
| 75 | k_tiles = tl.cdiv(K, BLOCK_SIZE_K) |
| 76 | num_tiles = num_pid_m * num_pid_n |
| 77 | |
| 78 | tile_id_c = start_pid - NUM_SMS |
| 79 | |
| 80 | offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) |
| 81 | num_pid_in_group = GROUP_SIZE_M * num_pid_n |
| 82 | |
| 83 | for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): |
| 84 | pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M) |
| 85 | start_m = pid_m * BLOCK_SIZE_M |
| 86 | start_n = pid_n * BLOCK_SIZE_N |
| 87 | offs_am = start_m + tl.arange(0, BLOCK_SIZE_M) |
| 88 | offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N) |
| 89 | if A_LARGE: |
| 90 | offs_am = offs_am.to(tl.int64) |
| 91 | if B_LARGE: |
| 92 | offs_bn = offs_bn.to(tl.int64) |
| 93 | offs_am = tl.where(offs_am < M, offs_am, 0) |
| 94 | offs_bn = tl.where(offs_bn < N, offs_bn, 0) |
| 95 | offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M) |
| 96 | offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N) |
| 97 | |
| 98 | accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 99 | for ki in range(k_tiles): |
| 100 | if A_LARGE or B_LARGE: |
| 101 | offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) |
| 102 | else: |
| 103 | offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) |
| 104 | a_ptrs = a_ptr + ( |
| 105 | offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak |
| 106 | ) |
| 107 | b_ptrs = b_ptr + ( |
| 108 | offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn |
| 109 | ) |
| 110 | |
| 111 | a = tl.load( |
| 112 | a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0 |
| 113 | ) |
| 114 | b = tl.load( |
| 115 | b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0 |
| 116 | ) |
| 117 | accumulator = tl.dot(a, b, accumulator) |
| 118 | |
| 119 | tile_id_c += NUM_SMS |
| 120 | pid_m, pid_n = _compute_pid( |
| 121 | tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M |
| 122 | ) |
| 123 | offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 124 | offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 125 | if C_LARGE: |
| 126 | offs_cm = offs_cm.to(tl.int64) |
| 127 | offs_cn = offs_cn.to(tl.int64) |
| 128 | c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] |
| 129 | c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) |
| 130 | if HAS_BIAS: |
| 131 | bias_ptrs = bias_ptr + offs_cn |
| 132 | bias = tl.load(bias_ptrs, mask=offs_cn < N, other=0.0).to(tl.float32) |
| 133 | accumulator += bias |
| 134 | c = accumulator.to(c_ptr.dtype.element_ty) |
| 135 | tl.store(c_ptrs, c, mask=c_mask) |
| 136 | |
| 137 | |
| 138 | @triton.jit(launch_metadata=_matmul_launch_metadata) |
| 139 | def matmul_kernel_descriptor_persistent( |
| 140 | a_ptr, |
| 141 | b_ptr, |
| 142 | c_ptr, # |
| 143 | bias_ptr, |
| 144 | M, |
| 145 | N, |
| 146 | K, # |
| 147 | BLOCK_SIZE_M: tl.constexpr, # |
| 148 | BLOCK_SIZE_N: tl.constexpr, # |
| 149 | BLOCK_SIZE_K: tl.constexpr, # |
| 150 | GROUP_SIZE_M: tl.constexpr, # |
| 151 | NUM_SMS: tl.constexpr, # |
| 152 | HAS_BIAS: tl.constexpr, |
| 153 | ): |
| 154 | """Persistent matmul using tensor descriptors for 2D block I/O. |
| 155 | |
| 156 | Expects b_ptr to point to a transposed B matrix of shape [N, K] with |
| 157 | row-major (K-contiguous) layout. The dot product transposes each loaded |
| 158 | B-tile back: dot(A_tile, B_tile.T). |
| 159 | |
| 160 | ~3x faster than the pointer-based persistent kernel on Intel XPU because |
| 161 | tensor descriptors leverage hardware 2D block load/store with automatic |
| 162 | bounds checking (no explicit masks needed). |
| 163 | """ |
| 164 | dtype = c_ptr.dtype.element_ty |
| 165 | start_pid = tl.program_id(axis=0) |
| 166 | num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) |
| 167 | num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) |
| 168 | k_tiles = tl.cdiv(K, BLOCK_SIZE_K) |
| 169 | num_tiles = num_pid_m * num_pid_n |
| 170 | |
| 171 | a_desc = tl.make_tensor_descriptor( |
| 172 | a_ptr, |
| 173 | shape=[M, K], |
| 174 | strides=[K, 1], |
| 175 | block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], |
| 176 | ) |
| 177 | b_desc = tl.make_tensor_descriptor( |
| 178 | b_ptr, |
| 179 | shape=[N, K], |
| 180 | strides=[K, 1], |
| 181 | block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K], |
| 182 | ) |
| 183 | c_desc = tl.make_tensor_descriptor( |
| 184 | c_ptr, |
| 185 | shape=[M, N], |
| 186 | strides=[N, 1], |
| 187 | block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N], |
| 188 | ) |
| 189 | |
| 190 | tile_id_c = start_pid - NUM_SMS |
| 191 | num_pid_in_group = GROUP_SIZE_M * num_pid_n |
| 192 | |
| 193 | for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): |
| 194 | pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M) |
| 195 | offs_am = pid_m * BLOCK_SIZE_M |
| 196 | offs_bn = pid_n * BLOCK_SIZE_N |
| 197 | |
| 198 | accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 199 | for ki in range(k_tiles): |
| 200 | offs_k = ki * BLOCK_SIZE_K |
| 201 | a = a_desc.load([offs_am, offs_k]) |
| 202 | b = b_desc.load([offs_bn, offs_k]) |
| 203 | accumulator = tl.dot(a, b.T, accumulator) |
| 204 | |
| 205 | tile_id_c += NUM_SMS |
| 206 | pid_m, pid_n = _compute_pid( |
| 207 | tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M |
| 208 | ) |
| 209 | offs_cm = pid_m * BLOCK_SIZE_M |
| 210 | offs_cn = pid_n * BLOCK_SIZE_N |
| 211 | |
| 212 | if HAS_BIAS: |
| 213 | bias_offs = offs_cn + tl.arange(0, BLOCK_SIZE_N) |
| 214 | bias_vals = tl.load(bias_ptr + bias_offs, mask=bias_offs < N, other=0.0).to( |
| 215 | tl.float32 |
| 216 | ) |
| 217 | accumulator += bias_vals[None, :] |
| 218 | |
| 219 | c = accumulator.to(dtype) |
| 220 | c_desc.store([offs_cm, offs_cn], c) |
| 221 | |
| 222 | |
| 223 | def matmul_descriptor_persistent( |
| 224 | a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None |
| 225 | ) -> torch.Tensor: |
| 226 | """Persistent matmul using tensor descriptors (Intel XPU fast path). |
| 227 | |
| 228 | Args: |
| 229 | a: Input matrix [M, K], must be contiguous. |
| 230 | b: Weight matrix [K, N] (standard layout — transposed internally). |
| 231 | bias: Optional 1D bias vector [N]. |
| 232 | |
| 233 | Returns: |
| 234 | Output matrix [M, N] with dtype matching the inputs. |
| 235 | |
| 236 | """ |
| 237 | assert a.shape[1] == b.shape[0], "Incompatible dimensions" |
| 238 | assert a.dtype == b.dtype, "Incompatible dtypes" |
| 239 | assert bias is None or bias.dim() == 1, "bias must be 1D" |
| 240 | |
| 241 | M, K = a.shape |
| 242 | _, N = b.shape |
| 243 | dtype = a.dtype |
| 244 | |
| 245 | # Tensor descriptors require contiguous row-major layout. |
| 246 | a = a.contiguous() |
| 247 | # Descriptor kernel expects B in [N, K] layout (K-contiguous). |
| 248 | b_t = b.t().contiguous() |
| 249 | |
| 250 | c = torch.empty((M, N), device=a.device, dtype=dtype) |
| 251 | # This path is XPU-only. _NUM_SMS is set by enable_batch_invariant_mode() |
| 252 | # before torch.compile traces this function, so Dynamo can lift it |
| 253 | # as a compile-time constant. |
| 254 | NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(a.device.index) |
| 255 | |
| 256 | def grid(META): |
| 257 | return ( |
| 258 | min( |
| 259 | NUM_SMS, |
| 260 | triton.cdiv(M, META["BLOCK_SIZE_M"]) |
| 261 | * triton.cdiv(N, META["BLOCK_SIZE_N"]), |
| 262 | ), |
| 263 | ) |
| 264 | |
| 265 | cfg = _get_descriptor_matmul_config(M, N, K, dtype) |
| 266 | |
| 267 | matmul_kernel_descriptor_persistent[grid]( |
| 268 | a, |
| 269 | b_t, |
| 270 | c, |
| 271 | bias, |
| 272 | M, |
| 273 | N, |
| 274 | K, |
| 275 | NUM_SMS=NUM_SMS, |
| 276 | HAS_BIAS=bias is not None, |
| 277 | **cfg, |
| 278 | ) |
| 279 | return c |
| 280 | |
| 281 | |
| 282 | def matmul_persistent( |
| 283 | a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None |
| 284 | ): |
| 285 | # Check constraints. |
| 286 | assert a.shape[1] == b.shape[0], "Incompatible dimensions" |
| 287 | assert a.dtype == b.dtype, "Incompatible dtypes" |
| 288 | assert bias is None or bias.dim() == 1, ( |
| 289 | "Currently assuming bias is 1D, let Horace know if you run into this" |
| 290 | ) |
| 291 | NUM_SMS = num_compute_units(a.device.index) |
| 292 | M, K = a.shape |
| 293 | K, N = b.shape |
| 294 | dtype = a.dtype |
| 295 | # Allocates output. |
| 296 | c = torch.empty((M, N), device=a.device, dtype=dtype) |
| 297 | |
| 298 | # 1D launch kernel where each block gets its own program. |
| 299 | def grid(META): |
| 300 | return ( |
| 301 | min( |
| 302 | NUM_SMS, |
| 303 | triton.cdiv(M, META["BLOCK_SIZE_M"]) |
| 304 | * triton.cdiv(N, META["BLOCK_SIZE_N"]), |
| 305 | ), |
| 306 | ) |
| 307 | |
| 308 | configs = { |
| 309 | torch.bfloat16: { |
| 310 | "BLOCK_SIZE_M": 128, |
| 311 | "BLOCK_SIZE_N": 128, |
| 312 | "BLOCK_SIZE_K": 64, |
| 313 | "GROUP_SIZE_M": 8, |
| 314 | "num_stages": 3, |
| 315 | "num_warps": 8, |
| 316 | }, |
| 317 | torch.float16: { |
| 318 | "BLOCK_SIZE_M": 128, |
| 319 | "BLOCK_SIZE_N": _fp16_block_size_n, |
| 320 | "BLOCK_SIZE_K": 64, |
| 321 | "GROUP_SIZE_M": 8, |
| 322 | "num_stages": 3, |
| 323 | "num_warps": 8, |
| 324 | }, |
| 325 | torch.float32: { |
| 326 | "BLOCK_SIZE_M": 128, |
| 327 | "BLOCK_SIZE_N": _fp32_block_size_n if N == 1 else 128, |
| 328 | "BLOCK_SIZE_K": 32, |
| 329 | "GROUP_SIZE_M": 8, |
| 330 | "num_stages": _fp32_num_stages if N == 1 else 3, |
| 331 | "num_warps": 8, |
| 332 | }, |
| 333 | } |
| 334 | matmul_kernel_persistent[grid]( |
| 335 | a, |
| 336 | b, |
| 337 | c, # |
| 338 | bias, |
| 339 | M, |
| 340 | N, |
| 341 | K, # |
| 342 | a.stride(0), |
| 343 | a.stride(1), # |
| 344 | b.stride(0), |
| 345 | b.stride(1), # |
| 346 | c.stride(0), |
| 347 | c.stride(1), # |
| 348 | NUM_SMS=NUM_SMS, # |
| 349 | A_LARGE=a.numel() > 2**31, |
| 350 | B_LARGE=b.numel() > 2**31, |
| 351 | C_LARGE=c.numel() > 2**31, |
| 352 | HAS_BIAS=bias is not None, |
| 353 | **_get_matmul_config(M, N, K, dtype, configs[dtype]), |
| 354 | ) |
| 355 | return c |
| 356 | |
| 357 | |
| 358 | @triton.jit |
| 359 | def bmm_kernel( |
| 360 | a_ptr, # (*, ) pointer to A, (B, M, K) |
| 361 | b_ptr, # (*, ) pointer to B, (B, K, N) |
| 362 | c_ptr, # (*, ) pointer to C, (B, M, N) |
| 363 | B, # int, batch size |
| 364 | M, # int, output rows |
| 365 | N, # int, output cols |
| 366 | K, # int, reduction dim |
| 367 | stride_ab, |
| 368 | stride_am, |
| 369 | stride_ak, |
| 370 | stride_bb, |
| 371 | stride_bk, |
| 372 | stride_bn, |
| 373 | stride_cb, |
| 374 | stride_cm, |
| 375 | stride_cn, |
| 376 | BLOCK_SIZE_M: tl.constexpr, |
| 377 | BLOCK_SIZE_N: tl.constexpr, |
| 378 | BLOCK_SIZE_K: tl.constexpr, |
| 379 | A_LARGE: tl.constexpr, |
| 380 | B_LARGE: tl.constexpr, |
| 381 | C_LARGE: tl.constexpr, |
| 382 | ): |
| 383 | """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N). |
| 384 | |
| 385 | Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating |
| 386 | along K in a fixed order to preserve batch invariance. |
| 387 | """ |
| 388 | pid_b = tl.program_id(0) |
| 389 | pid = tl.program_id(1) |
| 390 | |
| 391 | if pid_b >= B: |
| 392 | return |
| 393 | |
| 394 | # number of tiles along M / N |
| 395 | num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) |
| 396 | num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) |
| 397 | |
| 398 | pid_m = pid // num_pid_n |
| 399 | pid_n = pid % num_pid_n |
| 400 | |
| 401 | if pid_m >= num_pid_m or pid_n >= num_pid_n: |
| 402 | return |
| 403 | |
| 404 | # offs_m / offs_n: raw global row/col indices for this tile |
| 405 | offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) |
| 406 | offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) |
| 407 | # masks for valid logical rows/cols within (M, N) |
| 408 | mask_m = offs_m < M # [BLOCK_SIZE_M] |
| 409 | mask_n = offs_n < N # [BLOCK_SIZE_N] |
| 410 | |
| 411 | if A_LARGE or B_LARGE or C_LARGE: |
| 412 | offs_m = offs_m.to(tl.int64) |
| 413 | offs_n = offs_n.to(tl.int64) |
| 414 | |
| 415 | offs_m = tl.where(mask_m, offs_m, 0) |
| 416 | offs_n = tl.where(mask_n, offs_n, 0) |
| 417 | |
| 418 | # hint for triton contiguous memory |
| 419 | offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M) |
| 420 | offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N) |
| 421 | |
| 422 | # base pointers for current batch, shape-wise: |
| 423 | # a_batch_ptr points to A[pid_b, 0, 0] |
| 424 | # b_batch_ptr points to B[pid_b, 0, 0] |
| 425 | # c_batch_ptr points to C[pid_b, 0, 0] |
| 426 | a_batch_ptr = a_ptr + pid_b * stride_ab |
| 427 | b_batch_ptr = b_ptr + pid_b * stride_bb |
| 428 | c_batch_ptr = c_ptr + pid_b * stride_cb |
| 429 | |
| 430 | accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) |
| 431 | # number of K-blocks this tile iterates over |
| 432 | k_tiles = tl.cdiv(K, BLOCK_SIZE_K) |
| 433 | offs_k_mask = tl.arange(0, BLOCK_SIZE_K) |
| 434 | |
| 435 | for ki in range(k_tiles): |
| 436 | if A_LARGE or B_LARGE: |
| 437 | # offs_k: [BLOCK_SIZE_K], global K indices |
| 438 | offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) |
| 439 | else: |
| 440 | offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) |
| 441 | |
| 442 | # a_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_K] |
| 443 | # element (i, j) points to A[pid_b, offs_m[i], offs_k[j]] |
| 444 | a_ptrs = a_batch_ptr + ( |
| 445 | offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak |
| 446 | ) |
| 447 | # b_ptrs: [BLOCK_SIZE_K, BLOCK_SIZE_N] |
| 448 | # element (i, j) points to B[pid_b, offs_k[i], offs_n[j]] |
| 449 | b_ptrs = b_batch_ptr + ( |
| 450 | offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn |
| 451 | ) |
| 452 | |
| 453 | # valid K lanes for this block |
| 454 | k_valid = offs_k_mask < (K - ki * BLOCK_SIZE_K) |
| 455 | # A mask within (M, K): [BLOCK_SIZE_M, BLOCK_SIZE_K] |
| 456 | a_mask = mask_m[:, None] & k_valid[None, :] |
| 457 | # B mask within (K, N): [BLOCK_SIZE_K, BLOCK_SIZE_N] |
| 458 | b_mask = k_valid[:, None] & mask_n[None, :] |
| 459 | |
| 460 | # a: [BLOCK_SIZE_M, BLOCK_SIZE_K] from A[offs_m, offs_k] |
| 461 | a = tl.load( |
| 462 | a_ptrs, |
| 463 | mask=a_mask, |
| 464 | other=0.0, |
| 465 | ) |
| 466 | # b: [BLOCK_SIZE_K, BLOCK_SIZE_N] from B[offs_k, offs_n] |
| 467 | b = tl.load( |
| 468 | b_ptrs, |
| 469 | mask=b_mask, |
| 470 | other=0.0, |
| 471 | ) |
| 472 | accumulator = tl.dot(a, b, accumulator) |
| 473 | |
| 474 | # c_m / c_n: [BLOCK_SIZE_M] / [BLOCK_SIZE_N], row/col indices for C |
| 475 | c_m = offs_m |
| 476 | c_n = offs_n |
| 477 | if C_LARGE: |
| 478 | c_m = c_m.to(tl.int64) |
| 479 | c_n = c_n.to(tl.int64) |
| 480 | |
| 481 | # c_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_N] |
| 482 | # element (i, j) points to C[pid_b, c_m[i], c_n[j]] |
| 483 | c_ptrs = c_batch_ptr + stride_cm * c_m[:, None] + stride_cn * c_n[None, :] |
| 484 | # mask out elements that fall outside logical (M, N) range |
| 485 | c_mask = mask_m[:, None] & mask_n[None, :] |
| 486 | # cast FP32 accumulator back to original dtype of C |
| 487 | c = accumulator.to(c_ptr.dtype.element_ty) |
| 488 | tl.store(c_ptrs, c, mask=c_mask) |
| 489 | |
| 490 | |
| 491 | @triton.jit |
| 492 | def _log_softmax_kernel( |
| 493 | input_ptr, |
| 494 | output_ptr, |
| 495 | input_row_stride, |
| 496 | output_row_stride, |
| 497 | n_cols, |
| 498 | BLOCK_SIZE: tl.constexpr, |
| 499 | ): |
| 500 | """Compute log_softmax along the last dimension of a 2D tensor. |
| 501 | Each block handles one row of the input tensor. |
| 502 | """ |
| 503 | # Get the row index for this block |
| 504 | row_idx = tl.program_id(0).to(tl.int64) |
| 505 | |
| 506 | # Compute base pointers for input and output rows |
| 507 | row_start_ptr = input_ptr + row_idx * input_row_stride |
| 508 | output_row_start_ptr = output_ptr + row_idx * output_row_stride |
| 509 | |
| 510 | # Step 1: Find maximum value in the row for numerical stability |
| 511 | max_val = -float("inf") |
| 512 | for col_offset in range(0, n_cols, BLOCK_SIZE): |
| 513 | col_idx = col_offset + tl.arange(0, BLOCK_SIZE) |
| 514 | mask = col_idx < n_cols |
| 515 | |
| 516 | # Load values |
| 517 | vals = tl.load(row_start_ptr + col_idx, mask=mask, other=-float("inf")) |
| 518 | |
| 519 | # Update maximum |
| 520 | max_val = tl.max(tl.maximum(vals, max_val)) |
| 521 | |
| 522 | # Step 2: Compute sum of exp(x - max_val) |
| 523 | sum_exp = 0.0 |
| 524 | for col_offset in range(0, n_cols, BLOCK_SIZE): |
| 525 | col_idx = col_offset + tl.arange(0, BLOCK_SIZE) |
| 526 | mask = col_idx < n_cols |
| 527 | |
| 528 | # Load values |
| 529 | vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) |
| 530 | |
| 531 | # Compute exp(x - max_val) and accumulate |
| 532 | exp_vals = tl.exp(vals - max_val) |
| 533 | sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0)) |
| 534 | |
| 535 | # Compute log(sum_exp) |
| 536 | log_sum_exp = tl.log(sum_exp) |
| 537 | |
| 538 | # Step 3: Compute final log_softmax values: x - max_val - log_sum_exp |
| 539 | for col_offset in range(0, n_cols, BLOCK_SIZE): |
| 540 | col_idx = col_offset + tl.arange(0, BLOCK_SIZE) |
| 541 | mask = col_idx < n_cols |
| 542 | |
| 543 | # Load values |
| 544 | vals = tl.load(row_start_ptr + col_idx, mask=mask) |
| 545 | |
| 546 | # Compute log_softmax |
| 547 | output = vals - max_val - log_sum_exp |
| 548 | |
| 549 | # Store results |
| 550 | tl.store(output_row_start_ptr + col_idx, output, mask=mask) |
| 551 | |
| 552 | |
| 553 | def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 554 | """Compute log_softmax using Triton kernel. |
| 555 | |
| 556 | Args: |
| 557 | input: Input tensor |
| 558 | dim: Dimension along which to compute log_softmax |
| 559 | (only -1 or last dim supported) |
| 560 | |
| 561 | Returns: |
| 562 | Tensor with log_softmax applied along the specified dimension |
| 563 | |
| 564 | """ |
| 565 | if dim != -1 and dim != input.ndim - 1: |
| 566 | raise ValueError( |
| 567 | "This implementation only supports log_softmax along the last dimension" |
| 568 | ) |
| 569 | |
| 570 | # Flatten all dimensions except the last one |
| 571 | original_shape = input.shape |
| 572 | input_2d = input.reshape(-1, input.shape[-1]) |
| 573 | input_2d = input_2d.contiguous() |
| 574 | |
| 575 | n_rows, n_cols = input_2d.shape |
| 576 | |
| 577 | # Allocate output tensor |
| 578 | output = torch.empty_like(input_2d) |
| 579 | |
| 580 | # Choose block size based on the number of columns |
| 581 | BLOCK_SIZE = 1024 |
| 582 | |
| 583 | # Launch kernel with one block per row |
| 584 | grid = (n_rows,) |
| 585 | _log_softmax_kernel[grid]( |
| 586 | input_2d, |
| 587 | output, |
| 588 | input_2d.stride(0), |
| 589 | output.stride(0), |
| 590 | n_cols, |
| 591 | BLOCK_SIZE=BLOCK_SIZE, |
| 592 | ) |
| 593 | # Reshape output back to original shape |
| 594 | return output.reshape(original_shape) |
| 595 | |
| 596 | |
| 597 | @triton.jit |
| 598 | def mean_kernel( |
| 599 | input_ptr, |
| 600 | output_ptr, |
| 601 | input_stride0, |
| 602 | input_stride1, |
| 603 | input_stride2, |
| 604 | output_stride0, |
| 605 | output_stride1, |
| 606 | M, # size before reduction dim |
| 607 | N, # size of reduction dim |
| 608 | K, # size after reduction dim |
| 609 | BLOCK_SIZE: tl.constexpr, |
| 610 | ): |
| 611 | """Kernel for computing mean along a single dimension. |
| 612 | Input is viewed as (M, N, K) where N is the dimension being reduced. |
| 613 | """ |
| 614 | # Program ID gives us which output element we're computing |
| 615 | pid = tl.program_id(0) |
| 616 | |
| 617 | # Compute output indices |
| 618 | m_idx = pid // K |
| 619 | k_idx = pid % K |
| 620 | |
| 621 | # Bounds check |
| 622 | if m_idx >= M or k_idx >= K: |
| 623 | return |
| 624 | |
| 625 | # Accumulate sum across reduction dimension |
| 626 | acc = 0.0 |
| 627 | for n_start in range(0, N, BLOCK_SIZE): |
| 628 | n_offsets = n_start + tl.arange(0, BLOCK_SIZE) |
| 629 | mask = n_offsets < N |
| 630 | |
| 631 | # Calculate input indices |
| 632 | input_idx = ( |
| 633 | m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 |
| 634 | ) |
| 635 | |
| 636 | # Load and accumulate |
| 637 | vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) |
| 638 | acc += tl.sum(vals) |
| 639 | |
| 640 | # Compute mean and store |
| 641 | mean_val = acc / N |
| 642 | output_idx = m_idx * output_stride0 + k_idx * output_stride1 |
| 643 | tl.store(output_ptr + output_idx, mean_val) |
| 644 | |
| 645 | |
| 646 | def mean_dim( |
| 647 | input: torch.Tensor, |
| 648 | dim: int, |
| 649 | keepdim: bool = False, |
| 650 | dtype: torch.dtype | None = None, |
| 651 | ) -> torch.Tensor: |
| 652 | """Triton implementation of torch.mean with single dimension reduction. |
| 653 | |
| 654 | Args: |
| 655 | input: Input tensor |
| 656 | dim: Single dimension along which to compute mean |
| 657 | keepdim: Whether to keep the reduced dimension |
| 658 | dtype: Output dtype. If None, uses input dtype |
| 659 | (or float32 for integer inputs) |
| 660 | |
| 661 | Returns: |
| 662 | Tensor with mean values along specified dimension |
| 663 | |
| 664 | """ |
| 665 | # Validate inputs |
| 666 | assert -input.ndim <= dim < input.ndim, ( |
| 667 | f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" |
| 668 | ) |
| 669 | |
| 670 | # Handle negative dim |
| 671 | if dim < 0: |
| 672 | dim = dim + input.ndim |
| 673 | |
| 674 | # Handle dtype |
| 675 | if dtype is None: |
| 676 | if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]: |
| 677 | dtype = torch.float32 |
| 678 | else: |
| 679 | dtype = input.dtype |
| 680 | |
| 681 | # Convert input to appropriate dtype if needed |
| 682 | if input.dtype != dtype: |
| 683 | input = input.to(dtype) |
| 684 | |
| 685 | # Get input shape and strides |
| 686 | shape = list(input.shape) |
| 687 | |
| 688 | # Calculate dimensions for kernel |
| 689 | M = 1 |
| 690 | for i in range(dim): |
| 691 | M *= shape[i] |
| 692 | |
| 693 | N = shape[dim] |
| 694 | |
| 695 | K = 1 |
| 696 | for i in range(dim + 1, len(shape)): |
| 697 | K *= shape[i] |
| 698 | |
| 699 | # Reshape input to 3D view (M, N, K) |
| 700 | input_3d = input.reshape(M, N, K) |
| 701 | |
| 702 | # Create output shape |
| 703 | if keepdim: |
| 704 | output_shape = shape.copy() |
| 705 | output_shape[dim] = 1 |
| 706 | else: |
| 707 | output_shape = shape[:dim] + shape[dim + 1 :] |
| 708 | |
| 709 | # Create output tensor |
| 710 | output = torch.empty(output_shape, dtype=dtype, device=input.device) |
| 711 | |
| 712 | # Reshape output for kernel |
| 713 | output_2d = output.reshape(M, 1, K).squeeze(1) if keepdim else output.reshape(M, K) |
| 714 | |
| 715 | # Launch kernel |
| 716 | grid = (M * K,) |
| 717 | BLOCK_SIZE = 1024 |
| 718 | |
| 719 | mean_kernel[grid]( |
| 720 | input_3d, |
| 721 | output_2d, |
| 722 | input_3d.stride(0), |
| 723 | input_3d.stride(1), |
| 724 | input_3d.stride(2), |
| 725 | output_2d.stride(0), |
| 726 | output_2d.stride(1) if output_2d.ndim > 1 else 0, |
| 727 | M, |
| 728 | N, |
| 729 | K, |
| 730 | BLOCK_SIZE, |
| 731 | ) |
| 732 | |
| 733 | return output |
| 734 | |
| 735 | |
| 736 | def mm_batch_invariant(a, b, bias=None): |
| 737 | if a.device.type == "xpu": |
| 738 | return matmul_descriptor_persistent(a, b, bias=bias) |
| 739 | return matmul_persistent(a, b, bias=bias) |
| 740 | |
| 741 | |
| 742 | def matmul_batch_invariant(a, b, *, out=None): |
| 743 | # torch.matmul can handle various dimensions |
| 744 | # For 2D x 2D, it's the same as mm |
| 745 | if a.ndim == 2 and b.ndim == 2: |
| 746 | result = mm_batch_invariant(a, b) |
| 747 | if out is not None: |
| 748 | out.copy_(result) |
| 749 | return out |
| 750 | return result |
| 751 | elif b.ndim == 2: |
| 752 | # Handle ND x 2D: Common for linear layers |
| 753 | # (..., batch, seq, hidden) @ (hidden, out) -> (..., batch, seq, out) |
| 754 | batch_dims = a.shape[:-1] |
| 755 | hidden = a.shape[-1] |
| 756 | out_dim = b.shape[-1] |
| 757 | a_2d = a.reshape(-1, hidden) |
| 758 | result_2d = mm_batch_invariant(a_2d, b) |
| 759 | result = result_2d.reshape(batch_dims + (out_dim,)) |
| 760 | if out is not None: |
| 761 | out.copy_(result) |
| 762 | return out |
| 763 | return result |
| 764 | elif a.ndim >= 2 and b.ndim >= 3: |
| 765 | # Generic handler for 2D x ND and ND x ND (except 1D) |
| 766 | # Broadcast dims to ensure both matrices have the same shape |
| 767 | # If 2D x ND, then unsqueeze to add a dim to a |
| 768 | if a.ndim == 2: |
| 769 | a = a.unsqueeze(0) |
| 770 | broadcast_shape = torch.broadcast_shapes(a.shape[:-2], b.shape[:-2]) |
| 771 | a = a.expand(broadcast_shape + a.shape[-2:]) |
| 772 | b = b.expand(broadcast_shape + b.shape[-2:]) |
| 773 | batch_dim = math.prod(broadcast_shape) |
| 774 | # Reuse broadcast shape to get all dims except mm dims |
| 775 | a_3d = a.reshape(batch_dim, a.shape[-2], a.shape[-1]) |
| 776 | b_3d = b.reshape(batch_dim, b.shape[-2], b.shape[-1]) |
| 777 | # Do batched matmul |
| 778 | result_3d = bmm_batch_invariant(a_3d, b_3d) |
| 779 | # Reshape back to [broadcast_shape, seq_a, seq_b] |
| 780 | result = result_3d.reshape(broadcast_shape + (a.shape[-2], b.shape[-1])) |
| 781 | if out is not None: |
| 782 | out.copy_(result) |
| 783 | return out |
| 784 | return result |
| 785 | else: |
| 786 | raise ValueError( |
| 787 | f"matmul_batch_invariant requires both inputs be at least 2D " |
| 788 | f"got shapes {a.shape} and {b.shape}" |
| 789 | ) |
| 790 | |
| 791 | |
| 792 | def bmm_batch_invariant(a, b, *, out=None): |
| 793 | # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N) |
| 794 | if not (a.ndim == 3 and b.ndim == 3): |
| 795 | raise ValueError( |
| 796 | f"bmm_batch_invariant expects 3D tensors, " |
| 797 | f"got shapes {a.shape} and {b.shape}" |
| 798 | ) |
| 799 | |
| 800 | if a.shape[0] != b.shape[0]: |
| 801 | raise ValueError( |
| 802 | f"Batch dimensions of tensors must match, " |
| 803 | f"but got {a.shape[0]} and {b.shape[0]}." |
| 804 | ) |
| 805 | if a.shape[2] != b.shape[1]: |
| 806 | raise ValueError( |
| 807 | f"Incompatible inner dimensions for matmul: got {a.shape} and {b.shape}." |
| 808 | ) |
| 809 | if a.dtype != b.dtype: |
| 810 | raise ValueError(f"Incompatible dtypes: got {a.dtype} and {b.dtype}.") |
| 811 | |
| 812 | B, M, K = a.shape |
| 813 | _, _, N = b.shape |
| 814 | dtype = a.dtype |
| 815 | |
| 816 | if out is None: |
| 817 | c = torch.empty((B, M, N), device=a.device, dtype=dtype) |
| 818 | else: |
| 819 | assert out.shape == (B, M, N), "out tensor has incorrect shape" |
| 820 | assert out.dtype == dtype and out.device == a.device, "out tensor mismatch" |
| 821 | c = out |
| 822 | |
| 823 | configs = { |
| 824 | torch.bfloat16: { |
| 825 | "BLOCK_SIZE_M": 128, |
| 826 | "BLOCK_SIZE_N": 128, |
| 827 | "BLOCK_SIZE_K": 64, |
| 828 | "num_stages": 3, |
| 829 | "num_warps": 8, |
| 830 | }, |
| 831 | torch.float16: { |
| 832 | "BLOCK_SIZE_M": 128, |
| 833 | "BLOCK_SIZE_N": _fp16_block_size_n, |
| 834 | "BLOCK_SIZE_K": 64, |
| 835 | "num_stages": 3, |
| 836 | "num_warps": 8, |
| 837 | }, |
| 838 | torch.float32: { |
| 839 | "BLOCK_SIZE_M": 128, |
| 840 | "BLOCK_SIZE_N": _fp32_block_size_n if N == 1 else 128, |
| 841 | "BLOCK_SIZE_K": 32, |
| 842 | "num_stages": _fp32_num_stages if N == 1 else 3, |
| 843 | "num_warps": 8, |
| 844 | }, |
| 845 | } |
| 846 | |
| 847 | cfg = configs[dtype] |
| 848 | # grid = (B, num_tiles_per_matrix) |
| 849 | grid = ( |
| 850 | B, |
| 851 | triton.cdiv(M, cfg["BLOCK_SIZE_M"]) * triton.cdiv(N, cfg["BLOCK_SIZE_N"]), |
| 852 | ) |
| 853 | |
| 854 | bmm_kernel[grid]( |
| 855 | a, |
| 856 | b, |
| 857 | c, |
| 858 | B, |
| 859 | M, |
| 860 | N, |
| 861 | K, |
| 862 | a.stride(0), |
| 863 | a.stride(1), |
| 864 | a.stride(2), |
| 865 | b.stride(0), |
| 866 | b.stride(1), |
| 867 | b.stride(2), |
| 868 | c.stride(0), |
| 869 | c.stride(1), |
| 870 | c.stride(2), |
| 871 | A_LARGE=a.numel() > 2**31, |
| 872 | B_LARGE=b.numel() > 2**31, |
| 873 | C_LARGE=c.numel() > 2**31, |
| 874 | **cfg, |
| 875 | ) |
| 876 | |
| 877 | return c |
| 878 | |
| 879 | |
| 880 | def addmm_batch_invariant(bias, a, b): |
| 881 | return mm_batch_invariant(a, b, bias=bias) |
| 882 | |
| 883 | |
| 884 | def _log_softmax_batch_invariant(input, dim, _half_to_float): |
| 885 | if _half_to_float: |
| 886 | return log_softmax(input.float(), dim=dim) |
| 887 | return log_softmax(input, dim=dim) |
| 888 | |
| 889 | |
| 890 | def softmax_batch_invariant(input, dim, dtype=None): |
| 891 | # Compute softmax in a deterministic way |
| 892 | # First subtract max for numerical stability (standard practice) |
| 893 | input_max = torch.amax(input, dim=dim, keepdim=True) |
| 894 | input = input - input_max |
| 895 | exp_x = torch.exp(input) |
| 896 | sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True) |
| 897 | return exp_x / sum_exp_x |
| 898 | |
| 899 | |
| 900 | def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None): |
| 901 | assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}" |
| 902 | |
| 903 | result = input.to(torch.float32) |
| 904 | |
| 905 | if len(dim) == 0: |
| 906 | dim = [i for i in range(len(input.shape))] |
| 907 | |
| 908 | # Sort dimensions to reduce from largest to smallest to handle shifting dims |
| 909 | # during iterative reduction. |
| 910 | sorted_dims = sorted([d % input.ndim for d in dim], reverse=True) |
| 911 | |
| 912 | # Iteratively apply a deterministic mean. |
| 913 | for d in sorted_dims: |
| 914 | result = mean_dim(result, dim=d, keepdim=True) |
| 915 | |
| 916 | if not keepdim: |
| 917 | # Squeeze the reduced dimensions. |
| 918 | for d in sorted_dims: |
| 919 | result = result.squeeze(d) |
| 920 | |
| 921 | return result |
| 922 | |
| 923 | |
| 924 | @triton.jit |
| 925 | def _rms_norm_kernel( |
| 926 | input_ptr, |
| 927 | weight_ptr, |
| 928 | output_ptr, |
| 929 | input_row_stride, |
| 930 | output_row_stride, |
| 931 | n_cols, |
| 932 | eps, |
| 933 | BLOCK_SIZE: tl.constexpr, |
| 934 | HAS_WEIGHT: tl.constexpr, |
| 935 | ): |
| 936 | """Compute RMS normalization along the last dimension of a 2D tensor. |
| 937 | RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight |
| 938 | Each block handles one row of the input tensor. |
| 939 | """ |
| 940 | row_idx = tl.program_id(0).to(tl.int64) |
| 941 | row_start_ptr = input_ptr + row_idx * input_row_stride |
| 942 | output_row_start_ptr = output_ptr + row_idx * output_row_stride |
| 943 | |
| 944 | # Step 1: Compute sum of squares in float32 to avoid overflow |
| 945 | sum_sq = tl.zeros([1], dtype=tl.float32) |
| 946 | for col_offset in range(0, n_cols, BLOCK_SIZE): |
| 947 | col_idx = col_offset + tl.arange(0, BLOCK_SIZE) |
| 948 | mask = col_idx < n_cols |
| 949 | |
| 950 | vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) |
| 951 | # Convert to float32 for accumulation to prevent overflow |
| 952 | vals_f32 = vals.to(tl.float32) |
| 953 | sq_vals = vals_f32 * vals_f32 |
| 954 | sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) |
| 955 | |
| 956 | # Step 2: Compute RMS (root mean square) in float32 |
| 957 | mean_sq = sum_sq / n_cols |
| 958 | rms = tl.sqrt(mean_sq + eps) |
| 959 | inv_rms = 1.0 / rms |
| 960 | |
| 961 | # Step 3: Normalize and apply weight |
| 962 | for col_offset in range(0, n_cols, BLOCK_SIZE): |
| 963 | col_idx = col_offset + tl.arange(0, BLOCK_SIZE) |
| 964 | mask = col_idx < n_cols |
| 965 | vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) |
| 966 | # Compute in float32 then convert back to input dtype |
| 967 | vals_f32 = vals.to(tl.float32) |
| 968 | output_f32 = vals_f32 * inv_rms |
| 969 | if HAS_WEIGHT: |
| 970 | weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) |
| 971 | output_f32 = output_f32 * weight.to(tl.float32) |
| 972 | output = output_f32.to(vals.dtype) |
| 973 | tl.store(output_row_start_ptr + col_idx, output, mask=mask) |
| 974 | |
| 975 | |
| 976 | def rms_norm_batch_invariant( |
| 977 | input: torch.Tensor, |
| 978 | weight: torch.Tensor | None, |
| 979 | eps: float = 1e-6, |
| 980 | residual: torch.Tensor | None = None, |
| 981 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 982 | """Compute RMS normalization using Triton kernel. |
| 983 | |
| 984 | Args: |
| 985 | input: Input tensor of shape (..., hidden_size) |
| 986 | weight: Weight tensor of shape (hidden_size,), or None to skip the |
| 987 | per-channel multiply (``RMSNorm(has_weight=False)``) |
| 988 | eps: Small constant for numerical stability |
| 989 | residual: Optional residual tensor fused into the normalization path |
| 990 | |
| 991 | Returns: |
| 992 | RMS normalized tensor, or ``(output, residual_out)`` when ``residual`` |
| 993 | is provided |
| 994 | |
| 995 | """ |
| 996 | if residual is not None: |
| 997 | assert input.shape == residual.shape, ( |
| 998 | f"Input shape {input.shape} must match residual shape {residual.shape}" |
| 999 | ) |
| 1000 | import vllm._custom_ops as ops |
| 1001 | |
| 1002 | ops.fused_add_rms_norm(input, residual, weight, eps) |
| 1003 | return input, residual |
| 1004 | |
| 1005 | if weight is not None: |
| 1006 | assert weight.dim() == 1, "Weight must be 1-dimensional" |
| 1007 | assert input.shape[-1] == weight.shape[0], ( |
| 1008 | f"Input last dimension ({input.shape[-1]}) must match " |
| 1009 | f"weight dimension ({weight.shape[0]})" |
| 1010 | ) |
| 1011 | weight = weight.contiguous() |
| 1012 | |
| 1013 | # Flatten all dimensions except the last one |
| 1014 | original_shape = input.shape |
| 1015 | input_2d = input.reshape(-1, input.shape[-1]) |
| 1016 | input_2d = input_2d.contiguous() |
| 1017 | |
| 1018 | n_rows, n_cols = input_2d.shape |
| 1019 | |
| 1020 | output = torch.empty_like(input_2d) |
| 1021 | BLOCK_SIZE = 1024 |
| 1022 | grid = (n_rows,) |
| 1023 | _rms_norm_kernel[grid]( |
| 1024 | input_2d, |
| 1025 | weight if weight is not None else input_2d, |
| 1026 | output, |
| 1027 | input_2d.stride(0), |
| 1028 | output.stride(0), |
| 1029 | n_cols, |
| 1030 | eps, |
| 1031 | BLOCK_SIZE=BLOCK_SIZE, |
| 1032 | HAS_WEIGHT=weight is not None, |
| 1033 | ) |
| 1034 | return output.reshape(original_shape) |
| 1035 | |
| 1036 | |
| 1037 | def linear_batch_invariant(input, weight, bias=None): |
| 1038 | output = matmul_batch_invariant(input, weight.t()) |
| 1039 | |
| 1040 | if bias is not None: |
| 1041 | output = output + bias |
| 1042 | return output |
| 1043 | |
| 1044 | |
| 1045 | _batch_invariant_MODE = False |
| 1046 | _batch_invariant_LIB = None |
| 1047 | _fp16_block_size_n = 256 |
| 1048 | _fp32_block_size_n = 128 |
| 1049 | _fp32_num_stages = 3 |
| 1050 | _NUM_SMS: int = 0 |
| 1051 | |
| 1052 | |
| 1053 | def enable_batch_invariant_mode(): |
| 1054 | global _batch_invariant_MODE, _batch_invariant_LIB |
| 1055 | global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages |
| 1056 | global _NUM_SMS |
| 1057 | |
| 1058 | if _batch_invariant_MODE: |
| 1059 | return |
| 1060 | |
| 1061 | _batch_invariant_MODE = True |
| 1062 | _batch_invariant_LIB = torch.library.Library("aten", "IMPL") |
| 1063 | |
| 1064 | key = current_platform.dispatch_key |
| 1065 | if current_platform.is_cuda(): |
| 1066 | if current_platform.is_device_capability_family(80): |
| 1067 | # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the |
| 1068 | # triton persistent matmul overrides for mm/addmm/matmul/linear. |
| 1069 | _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) |
| 1070 | _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) |
| 1071 | _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) |
| 1072 | _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) |
| 1073 | else: |
| 1074 | # Hopper (SM90) and Blackwell (SM100): the only source of batch |
| 1075 | # variance is split-k, which we disable via the cuBLAS workspace |
| 1076 | # config. |
| 1077 | os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" |
| 1078 | os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" |
| 1079 | |
| 1080 | # Query the shared memory size and set block size |
| 1081 | # accordingly to avoid triton OutOfResources. |
| 1082 | if get_max_shared_memory_bytes() > 106496: |
| 1083 | _fp16_block_size_n = 256 |
| 1084 | _fp32_block_size_n = 128 |
| 1085 | _fp32_num_stages = 3 |
| 1086 | else: |
| 1087 | _fp16_block_size_n = 128 |
| 1088 | # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376). |
| 1089 | # Only N=1 OORs; keep 128/stages=3 for every other N. |
| 1090 | _fp32_block_size_n = 32 |
| 1091 | _fp32_num_stages = 2 |
| 1092 | elif current_platform.is_xpu(): |
| 1093 | # Tensor descriptors need a global-memory allocator; must be set outside |
| 1094 | # torch.compile regions (triton.set_allocator modifies global state). |
| 1095 | def _triton_alloc_fn(size: int, alignment: int, stream: int | None): |
| 1096 | return torch.empty(size, device="xpu", dtype=torch.int8) |
| 1097 | |
| 1098 | triton.set_allocator(_triton_alloc_fn) |
| 1099 | |
| 1100 | _NUM_SMS = num_compute_units(0) |
| 1101 | _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key) |
| 1102 | _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key) |
| 1103 | _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) |
| 1104 | _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) |
| 1105 | |
| 1106 | _fp16_block_size_n = 128 |
| 1107 | |
| 1108 | # Native mean can change reduction order with the row count on XPU too. |
| 1109 | _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key) |
| 1110 | |
| 1111 | # Softmax, log_softmax, and bmm are already batch-invariant on XPU |
| 1112 | # (oneDNN backend produces bitwise-identical results regardless of batch |
| 1113 | # context). Only register these overrides on CUDA where they are needed. |
| 1114 | if not current_platform.is_xpu(): |
| 1115 | _batch_invariant_LIB.impl( |
| 1116 | "aten::_log_softmax", _log_softmax_batch_invariant, key |
| 1117 | ) |
| 1118 | _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key) |
| 1119 | _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key) |
| 1120 | # torch 2.12+ registers a built-in Triton bmm kernel for CUDA |
| 1121 | # (torch._native.ops.bmm_outer_product), so we need allow_override |
| 1122 | # to replace it at the dispatcher level. |
| 1123 | _batch_invariant_LIB.impl( |
| 1124 | "aten::bmm", bmm_batch_invariant, key, allow_override=True |
| 1125 | ) |
| 1126 | torch.bmm = bmm_batch_invariant |
| 1127 | |
| 1128 | reduced_precision_val = ( |
| 1129 | (False, False) if is_torch_equal_or_newer("2.10.0") else False |
| 1130 | ) |
| 1131 | torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = ( |
| 1132 | reduced_precision_val |
| 1133 | ) |
| 1134 | torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = ( |
| 1135 | reduced_precision_val |
| 1136 | ) |
| 1137 | if current_platform.is_cuda(): |
| 1138 | torch.backends.cuda.preferred_blas_library(backend="cublaslt") |
| 1139 | |
| 1140 | |
| 1141 | def override_envs_for_invariance(): |
| 1142 | os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" |
| 1143 | # Only the 1-stage kernel has a size- and rank-independent reduction order. |
| 1144 | os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = "1stage" |
| 1145 | |
| 1146 | os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" |
| 1147 | |
| 1148 | # NCCL determinism settings. NCCL keeps the launch mode process-wide, so, |
| 1149 | # unlike the pins below, it also applies to communicators shared with other |
| 1150 | # processes. |
| 1151 | os.environ["NCCL_LAUNCH_MODE"] = "GROUP" |
| 1152 | pin_nccl_env( |
| 1153 | { |
| 1154 | "NCCL_COLLNET_ENABLE": "0", |
| 1155 | "NCCL_NVLS_ENABLE": "0", |
| 1156 | "NCCL_P2P_NET_DISABLE": "1", |
| 1157 | "NCCL_MIN_NCHANNELS": "1", |
| 1158 | "NCCL_MAX_NCHANNELS": "1", |
| 1159 | "NCCL_PROTO": "Simple", |
| 1160 | # NCCL >= 2.31 zero-fills the algorithm table of every collective |
| 1161 | # when NCCL_ALGO is set and re-enables only the named ones; together |
| 1162 | # with the NCCL_PROTO above, collectives not named here end up with |
| 1163 | # no algorithm and fail with ncclInvalidUsage. Re-enable Ring and |
| 1164 | # Tree for all collectives, then pin AllReduce to Tree. |
| 1165 | "NCCL_ALGO": "ring,tree;allreduce:tree", |
| 1166 | "NCCL_NTHREADS": "1", |
| 1167 | "NCCL_SOCKET_NTHREADS": "1", |
| 1168 | } |
| 1169 | ) |
| 1170 | |
| 1171 | # torch.compile settings |
| 1172 | os.environ["VLLM_USE_AOT_COMPILE"] = "0" |
| 1173 | |
| 1174 | |
| 1175 | def init_batch_invariance(): |
| 1176 | # this will hit all the csrc overrides as well |
| 1177 | if envs.VLLM_BATCH_INVARIANT: |
| 1178 | resolve_tuned_matmul_configs() |
| 1179 | override_envs_for_invariance() |
| 1180 | enable_batch_invariant_mode() |
| 1181 | |
| 1182 | # Disable TF32 for batch invariance - it causes non-deterministic rounding |
| 1183 | torch.backends.cuda.matmul.fp32_precision = "ieee" |
| 1184 | torch.backends.cudnn.conv.fp32_precision = "ieee" |
| 1185 | torch.backends.cudnn.rnn.fp32_precision = "ieee" |