| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher |
| 9 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 10 | LaunchSpec, |
| 11 | TritonWarmupTensor, |
| 12 | VllmTritonJitKernel, |
| 13 | ) |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | |
| 16 | |
| 17 | class PackSeqTritonKernel(VllmTritonJitKernel["PackSeqTritonKernel.CompileKey"]): |
| 18 | @dataclass(frozen=True) |
| 19 | class CompileKey: |
| 20 | dtype: torch.dtype |
| 21 | pad_value: float | int |
| 22 | pad_is_uint8: bool |
| 23 | block_t: int |
| 24 | block_d: int |
| 25 | |
| 26 | # N is unused; D and Lmax only affect runtime addressing and masks. The |
| 27 | # compile-time tile shape remains controlled by BLOCK_T and BLOCK_D. |
| 28 | @staticmethod |
| 29 | @triton.jit(do_not_specialize=["N", "D", "Lmax"]) |
| 30 | def kernel( |
| 31 | x_ptr, # [N, D] |
| 32 | out_ptr, # [B, Lmax, D] |
| 33 | lengths_ptr, # *i32, [B] |
| 34 | N, |
| 35 | D, |
| 36 | Lmax, |
| 37 | PAD_VALUE: tl.constexpr, |
| 38 | PAD_IS_UINT8: tl.constexpr, |
| 39 | BLOCK_T: tl.constexpr, # timesteps per program |
| 40 | BLOCK_D: tl.constexpr, # features per program |
| 41 | ): |
| 42 | pid_b = tl.program_id(0) # batch id |
| 43 | pid_t = tl.program_id(1) # block over time dimension |
| 44 | pid_d = tl.program_id(2) # block over feature dimension |
| 45 | off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T] |
| 46 | off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D] |
| 47 | |
| 48 | # Compute start index and sequence length from cumulative lengths |
| 49 | in_start = 0 |
| 50 | for i in range(pid_b): |
| 51 | in_start += tl.load(lengths_ptr + i) |
| 52 | seq_len = tl.load(lengths_ptr + pid_b) |
| 53 | |
| 54 | # valid time positions for this block |
| 55 | t_mask = off_t < Lmax |
| 56 | |
| 57 | # compute input row indices for valid (b, t) |
| 58 | in_row = in_start + off_t |
| 59 | valid_row = (off_t < seq_len) & t_mask |
| 60 | |
| 61 | # Pointers |
| 62 | # x_ptr: row-major [N, D] |
| 63 | x_row_ptr = x_ptr + in_row[:, None] * D + off_d[None, :] |
| 64 | |
| 65 | # out_ptr: row-major [B, Lmax, D] |
| 66 | out_row_ptr = out_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :] |
| 67 | |
| 68 | # Initialize with PAD. PAD_IS_UINT8 selects the pad tensor's dtype so |
| 69 | # integer-typed outputs (e.g. MXFP4 packed nibbles, ue8m0 scale bytes) |
| 70 | # get an exact-byte pad rather than going through an fp32→uint8 cast |
| 71 | # that's implementation-defined outside of value 0. |
| 72 | d_mask = off_d[None, :] < D |
| 73 | if PAD_IS_UINT8: |
| 74 | pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.uint8) |
| 75 | else: |
| 76 | pad_vals = tl.full([BLOCK_T, BLOCK_D], PAD_VALUE, tl.float32) |
| 77 | tl.store(out_row_ptr, pad_vals, mask=t_mask[:, None] & d_mask) |
| 78 | |
| 79 | # Load & write only where within seq_len |
| 80 | x_vals = tl.load(x_row_ptr, mask=valid_row[:, None] & d_mask) |
| 81 | tl.store(out_row_ptr, x_vals, mask=valid_row[:, None] & d_mask) |
| 82 | |
| 83 | def dispatch( # type: ignore[override] |
| 84 | self, |
| 85 | *, |
| 86 | dtype: torch.dtype, |
| 87 | pad_value: float | int, |
| 88 | **compile_key_fields: int, |
| 89 | ) -> CompileKey: |
| 90 | is_uint8 = dtype == torch.uint8 |
| 91 | return self.CompileKey( |
| 92 | **compile_key_fields, |
| 93 | dtype=dtype, |
| 94 | pad_value=int(pad_value) if is_uint8 else float(pad_value), |
| 95 | pad_is_uint8=is_uint8, |
| 96 | ) |
| 97 | |
| 98 | def get_warmup_keys( |
| 99 | self, |
| 100 | *, |
| 101 | dtype: torch.dtype, |
| 102 | pad_value: float | int, |
| 103 | ) -> list[CompileKey]: |
| 104 | return self._trace_dispatch(self.dispatch)( |
| 105 | dtype=dtype, |
| 106 | pad_value=pad_value, |
| 107 | block_t=64, |
| 108 | block_d=64, |
| 109 | ) |
| 110 | |
| 111 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 112 | return dict( |
| 113 | x_reshaped=TritonWarmupTensor(compile_key.dtype), |
| 114 | out=TritonWarmupTensor(compile_key.dtype), |
| 115 | lengths=TritonWarmupTensor(torch.int32), |
| 116 | N=1, |
| 117 | D=1, |
| 118 | Lmax=1, |
| 119 | pad_value=compile_key.pad_value, |
| 120 | pad_is_uint8=compile_key.pad_is_uint8, |
| 121 | block_t=compile_key.block_t, |
| 122 | block_d=compile_key.block_d, |
| 123 | ) |
| 124 | |
| 125 | @kernel_launcher |
| 126 | def __call__( |
| 127 | self, |
| 128 | x_reshaped: torch.Tensor, |
| 129 | out: torch.Tensor, |
| 130 | lengths: torch.Tensor, |
| 131 | *, |
| 132 | N: int, |
| 133 | D: int, |
| 134 | Lmax: int, |
| 135 | pad_value: float | int, |
| 136 | pad_is_uint8: bool, |
| 137 | block_t: int, |
| 138 | block_d: int, |
| 139 | ) -> LaunchSpec: |
| 140 | grid = ( |
| 141 | lengths.shape[0], |
| 142 | triton.cdiv(Lmax, block_t), |
| 143 | triton.cdiv(D, block_d), |
| 144 | ) |
| 145 | return grid, dict( |
| 146 | x_ptr=x_reshaped, |
| 147 | lengths_ptr=lengths, |
| 148 | PAD_VALUE=pad_value, |
| 149 | PAD_IS_UINT8=pad_is_uint8, |
| 150 | BLOCK_T=block_t, |
| 151 | BLOCK_D=block_d, |
| 152 | num_warps=4, |
| 153 | num_stages=2, |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | def pack_seq_triton( |
| 158 | x: torch.Tensor, |
| 159 | lengths: torch.Tensor, |
| 160 | pad_value: float | int = -float("inf"), |
| 161 | block_t: int = 64, |
| 162 | block_d: int = 64, |
| 163 | ) -> torch.Tensor: |
| 164 | """Pack sequences of different lengths into a batched tensor. |
| 165 | |
| 166 | Supports float dtypes (any, via fp32 pad) and ``torch.uint8`` (exact-byte |
| 167 | pad — e.g. MXFP4 packed nibbles or ue8m0 scale bytes). For uint8 inputs |
| 168 | ``pad_value`` must be an integer in ``[0, 255]``. |
| 169 | |
| 170 | Args: |
| 171 | x: [N, ...] — input tensor where N is total number of tokens. |
| 172 | lengths: [B] — sequence lengths for each batch. |
| 173 | pad_value: value to use for padding. Defaults to ``-inf`` which is |
| 174 | only sensible for float dtypes; pass ``0`` (or any byte) for |
| 175 | uint8 inputs. |
| 176 | block_t: block size for time dimension. |
| 177 | block_d: block size for feature dimension. |
| 178 | |
| 179 | Returns: |
| 180 | packed: [B, Lmax, ...] — packed tensor. |
| 181 | |
| 182 | """ |
| 183 | is_uint8 = x.dtype == torch.uint8 |
| 184 | if is_uint8: |
| 185 | assert isinstance(pad_value, int) and 0 <= pad_value <= 255, ( |
| 186 | f"uint8 pack requires an integer pad in [0, 255], got {pad_value!r}" |
| 187 | ) |
| 188 | pad_constexpr: int | float = int(pad_value) |
| 189 | else: |
| 190 | pad_constexpr = float(pad_value) |
| 191 | |
| 192 | # Handle multi-dimensional input by reshaping to (N, -1) |
| 193 | original_shape = x.shape |
| 194 | if len(original_shape) > 2: |
| 195 | N = original_shape[0] |
| 196 | x_reshaped = x.reshape(N, -1) |
| 197 | D = x_reshaped.shape[1] |
| 198 | else: |
| 199 | N, D = x.shape |
| 200 | x_reshaped = x |
| 201 | |
| 202 | B = lengths.numel() |
| 203 | Lmax = int(lengths.max().item()) |
| 204 | |
| 205 | out = torch.empty((B, Lmax, D), device=x.device, dtype=x.dtype) |
| 206 | |
| 207 | lengths = lengths.int() |
| 208 | _PACK_SEQ_TRITON_KERNEL( |
| 209 | x_reshaped, |
| 210 | out, |
| 211 | lengths, |
| 212 | N=N, |
| 213 | D=D, |
| 214 | Lmax=Lmax, |
| 215 | pad_value=pad_constexpr, |
| 216 | pad_is_uint8=is_uint8, |
| 217 | block_t=block_t, |
| 218 | block_d=block_d, |
| 219 | ) |
| 220 | |
| 221 | if len(original_shape) > 2: |
| 222 | out = out.reshape((B, Lmax) + original_shape[1:]) |
| 223 | |
| 224 | return out |
| 225 | |
| 226 | |
| 227 | class UnpackSeqTritonKernel(VllmTritonJitKernel["UnpackSeqTritonKernel.CompileKey"]): |
| 228 | @dataclass(frozen=True) |
| 229 | class CompileKey: |
| 230 | dtype: torch.dtype |
| 231 | block_t: int |
| 232 | block_d: int |
| 233 | |
| 234 | # B, Lmax, and D only affect runtime addressing and masks. The compile-time |
| 235 | # tile shape remains controlled by BLOCK_T and BLOCK_D. |
| 236 | @staticmethod |
| 237 | @triton.jit(do_not_specialize=["B", "Lmax", "D"]) |
| 238 | def kernel( |
| 239 | packed_ptr, # [B, Lmax, D] |
| 240 | out_ptr, # [N, D] |
| 241 | lengths_ptr, # *i32, [B] |
| 242 | B, |
| 243 | Lmax, |
| 244 | D, |
| 245 | BLOCK_T: tl.constexpr, # timesteps per program |
| 246 | BLOCK_D: tl.constexpr, # features per program |
| 247 | ): |
| 248 | pid_b = tl.program_id(0) # batch id |
| 249 | pid_t = tl.program_id(1) # block over time dimension |
| 250 | pid_d = tl.program_id(2) # block over feature dimension |
| 251 | off_t = pid_t * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T] |
| 252 | off_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) # [BLOCK_D] |
| 253 | |
| 254 | # bounds: compute start from cumulative lengths |
| 255 | in_start = 0 |
| 256 | for i in range(pid_b): |
| 257 | in_start += tl.load(lengths_ptr + i) |
| 258 | seq_len = tl.load(lengths_ptr + pid_b) |
| 259 | |
| 260 | # valid time positions for this block |
| 261 | t_mask = off_t < Lmax |
| 262 | valid_row = (off_t < seq_len) & t_mask |
| 263 | |
| 264 | # compute output row indices for valid (b, t) |
| 265 | out_row = in_start + off_t |
| 266 | |
| 267 | # Pointers |
| 268 | # packed_ptr: row-major [B, Lmax, D] |
| 269 | packed_row_ptr = ( |
| 270 | packed_ptr + (pid_b * Lmax + off_t)[:, None] * D + off_d[None, :] |
| 271 | ) |
| 272 | |
| 273 | # out_ptr: row-major [N, D] |
| 274 | out_row_ptr = out_ptr + out_row[:, None] * D + off_d[None, :] |
| 275 | |
| 276 | # Load from packed tensor and store to output |
| 277 | d_mask = off_d[None, :] < D |
| 278 | packed_vals = tl.load(packed_row_ptr, mask=valid_row[:, None] & d_mask) |
| 279 | tl.store(out_row_ptr, packed_vals, mask=valid_row[:, None] & d_mask) |
| 280 | |
| 281 | def dispatch( # type: ignore[override] |
| 282 | self, |
| 283 | *, |
| 284 | dtype: torch.dtype, |
| 285 | **compile_key_fields: int, |
| 286 | ) -> CompileKey: |
| 287 | return self.CompileKey(**compile_key_fields, dtype=dtype) |
| 288 | |
| 289 | def get_warmup_keys(self) -> list[CompileKey]: |
| 290 | return self._trace_dispatch(self.dispatch)( |
| 291 | dtype=torch.int32, |
| 292 | block_t=64, |
| 293 | block_d=64, |
| 294 | ) |
| 295 | |
| 296 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 297 | return dict( |
| 298 | packed_reshaped=TritonWarmupTensor(compile_key.dtype), |
| 299 | out=TritonWarmupTensor(compile_key.dtype), |
| 300 | lengths=TritonWarmupTensor(torch.int32), |
| 301 | B=1, |
| 302 | Lmax=1, |
| 303 | D=1, |
| 304 | block_t=compile_key.block_t, |
| 305 | block_d=compile_key.block_d, |
| 306 | ) |
| 307 | |
| 308 | @kernel_launcher |
| 309 | def __call__( |
| 310 | self, |
| 311 | packed_reshaped: torch.Tensor, |
| 312 | out: torch.Tensor, |
| 313 | lengths: torch.Tensor, |
| 314 | *, |
| 315 | B: int, |
| 316 | Lmax: int, |
| 317 | D: int, |
| 318 | block_t: int, |
| 319 | block_d: int, |
| 320 | ) -> LaunchSpec: |
| 321 | grid = (B, triton.cdiv(Lmax, block_t), triton.cdiv(D, block_d)) |
| 322 | return grid, dict( |
| 323 | packed_ptr=packed_reshaped, |
| 324 | lengths_ptr=lengths, |
| 325 | BLOCK_T=block_t, |
| 326 | BLOCK_D=block_d, |
| 327 | num_warps=4, |
| 328 | num_stages=2, |
| 329 | ) |
| 330 | |
| 331 | |
| 332 | def unpack_seq_triton( |
| 333 | packed_tensor: torch.Tensor, |
| 334 | lengths: torch.Tensor, |
| 335 | block_t: int = 64, |
| 336 | block_d: int = 64, |
| 337 | ) -> torch.Tensor: |
| 338 | """Unpack a packed decode query tensor back to the original format. |
| 339 | Efficient Triton implementation. |
| 340 | |
| 341 | Args: |
| 342 | packed_tensor: [B, Lmax, ...] - packed tensor from pack_seq_triton |
| 343 | lengths: [B] - sequence lengths for each batch |
| 344 | block_t: block size for time dimension |
| 345 | block_d: block size for feature dimension |
| 346 | |
| 347 | Returns: |
| 348 | unpacked_tensor: [N, ...] where N = sum(lengths) |
| 349 | |
| 350 | """ |
| 351 | # Handle multi-dimensional input by reshaping to (B, Lmax, -1) |
| 352 | original_shape = packed_tensor.shape |
| 353 | if len(original_shape) > 3: |
| 354 | B, Lmax = original_shape[:2] |
| 355 | packed_reshaped = packed_tensor.reshape(B, Lmax, -1) |
| 356 | D = packed_reshaped.shape[2] |
| 357 | else: |
| 358 | B, Lmax, D = packed_tensor.shape |
| 359 | packed_reshaped = packed_tensor |
| 360 | |
| 361 | # Calculate total number of elements |
| 362 | N = int(lengths.sum().item()) |
| 363 | |
| 364 | out = torch.empty((N, D), device=packed_tensor.device, dtype=packed_tensor.dtype) |
| 365 | |
| 366 | lengths = lengths.int() |
| 367 | _UNPACK_SEQ_TRITON_KERNEL( |
| 368 | packed_reshaped, |
| 369 | out, |
| 370 | lengths, |
| 371 | B=B, |
| 372 | Lmax=Lmax, |
| 373 | D=D, |
| 374 | block_t=block_t, |
| 375 | block_d=block_d, |
| 376 | ) |
| 377 | |
| 378 | # Reshape output back to original dimensions (except first dimension) |
| 379 | if len(original_shape) > 3: |
| 380 | output_shape = (N,) + original_shape[2:] |
| 381 | out = out.reshape(output_shape) |
| 382 | |
| 383 | return out |
| 384 | |
| 385 | |
| 386 | _PACK_SEQ_TRITON_KERNEL = PackSeqTritonKernel() |
| 387 | _UNPACK_SEQ_TRITON_KERNEL = UnpackSeqTritonKernel() |