@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 191 lines 5.85 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 #
4 # Modified by Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 #
6 # Copied from https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats
7 #
8 # Copyright (c) Meta Platforms, Inc. and affiliates.
9 # All rights reserved.
10
11 # This source code is licensed under the license found in the
12 # LICENSE file in the root directory of this source tree.
13
14 from typing import Literal
15
16 import torch
17 from torch.library import wrap_triton
18
19 from vllm.triton_utils import tl, triton
20 from vllm.utils.math_utils import cdiv
21
22
23 @triton.jit
24 def triton_scale_swizzle(
25 scale_ptr: torch.Tensor,
26 scale_rows: int,
27 scale_cols: int,
28 output_ptr: torch.Tensor,
29 input_row_stride: int,
30 output_block_stride: int,
31 BLOCK_ROWS: tl.constexpr,
32 BLOCK_COLS: tl.constexpr,
33 ):
34 """Rearranges tensor data from row-major to block-scaled swizzle format.
35
36 Args:
37 scale_ptr: Pointer to the input scale tensor
38 scale_rows: Number of rows in the scale tensor
39 scale_cols: Number of columns in the scale tensor
40 output_ptr: Pointer to the output tensor
41 input_row_stride: Stride between rows in the input tensor
42 output_block_stride: Stride between blocks in the output tensor
43 BLOCK_ROWS: Number of rows in a tile (compile-time constant)
44 BLOCK_COLS: Number of columns in a tile (compile-time constant)
45
46 """
47 pid_row = tl.program_id(0)
48 pid_col = tl.program_id(1)
49
50 rows = tl.arange(0, BLOCK_ROWS)[:, None]
51 cols = tl.arange(0, BLOCK_COLS)[None, :]
52
53 # Calculate starting row and column for this tile
54 start_row = pid_row * BLOCK_ROWS
55 start_col = pid_col * BLOCK_COLS
56 global_rows = start_row + rows
57 global_cols = start_col + cols
58
59 mask = (global_rows < scale_rows) & (global_cols < scale_cols)
60
61 input_scales = tl.load(
62 scale_ptr + global_rows * input_row_stride + global_cols,
63 mask=mask,
64 other=0.0,
65 )
66
67 r_div_32 = rows // 32
68 r_mod_32 = rows % 32
69
70 # 2) Rearrange to (32, 4, 4) then to final (32, 16) coordinates
71 dest_indices = r_mod_32 * 16 + r_div_32 * 4 + cols
72
73 # Flatten
74 dest_indices_flat = tl.reshape(dest_indices, (BLOCK_ROWS * BLOCK_COLS))
75 scales_flat = tl.reshape(input_scales, (BLOCK_ROWS * BLOCK_COLS))
76
77 # Calculate block offset using provided output block stride
78 LOCAL_NUMEL = BLOCK_ROWS * BLOCK_COLS
79 block_offset = pid_col * LOCAL_NUMEL + (pid_row * output_block_stride)
80
81 tl.store(
82 output_ptr + block_offset + dest_indices_flat,
83 scales_flat,
84 )
85
86
87 @torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=())
88 def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
89 """Rearranges an E8M0 tensor scale from row-major format to
90 block-scaled swizzle format.
91
92 This format is suitable for Tmem as described in NVIDIA documentation:
93 https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
94
95 Args:
96 scale_tensor: Input tensor in row-major format with 8-bit elements
97
98 Returns:
99 Rearranged tensor in block-scaled swizzle format
100
101 """
102 assert scale_tensor.element_size() == 1, (
103 "Expected element size to be 1 byte (8 bits)"
104 )
105 assert scale_tensor.is_contiguous(), "Input tensor must be contiguous"
106
107 rows, cols = scale_tensor.shape
108
109 # Calculate blocks needed
110 n_row_blocks = triton.cdiv(rows, 128)
111 n_col_blocks = triton.cdiv(cols, 4)
112 padded_rows = n_row_blocks * 128
113 padded_cols = n_col_blocks * 4
114
115 out = scale_tensor.new_empty((padded_rows, padded_cols))
116
117 # Input stride (for row-major format)
118 input_row_stride = cols
119
120 # We probably want handle multiple blocks per tile but
121 # for now keep it simple
122 BLOCK_ROWS, BLOCK_COLS = 128, 4
123
124 # Output block stride for the rearranged format
125 output_block_stride = BLOCK_ROWS * BLOCK_COLS * (padded_cols // BLOCK_COLS)
126
127 grid = lambda META: (
128 triton.cdiv(padded_rows, BLOCK_ROWS),
129 triton.cdiv(padded_cols, BLOCK_COLS),
130 )
131
132 wrap_triton(triton_scale_swizzle)[grid](
133 scale_tensor.view(torch.uint8),
134 rows,
135 cols,
136 out.view(torch.uint8),
137 input_row_stride,
138 output_block_stride,
139 BLOCK_ROWS=BLOCK_ROWS,
140 BLOCK_COLS=BLOCK_COLS,
141 )
142
143 return out
144
145
146 @triton_mx_block_rearrange.register_fake
147 def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor:
148 rows, cols = scale_tensor.shape
149 padded_rows = cdiv(rows, 128) * 128
150 padded_cols = cdiv(cols, 4) * 4
151 return scale_tensor.new_empty((padded_rows, padded_cols))
152
153
154 def to_blocked(
155 input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"
156 ) -> torch.Tensor:
157 """Rearrange a large matrix by breaking it into blocks and applying
158 the rearrangement pattern.
159
160 See:
161 https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
162
163 Args:
164 input_matrix: Input tensor of shape (H, W)
165 backend: "torch" (PyTorch path) or "triton" (Triton kernel)
166
167 Returns:
168 Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4))
169
170 """
171 if backend == "triton":
172 return triton_mx_block_rearrange(input_matrix).flatten()
173 elif backend != "torch":
174 raise ValueError(f'backend must be "torch" or "triton", got {backend!r}')
175
176 rows, cols = input_matrix.shape
177 n_row_blocks = cdiv(rows, 128)
178 n_col_blocks = cdiv(cols, 4)
179
180 # Calculate the padded shape
181 padded_rows = n_row_blocks * 128
182 padded_cols = n_col_blocks * 4
183
184 padded = input_matrix
185 assert (rows, cols) == (padded_rows, padded_cols)
186
187 # Rearrange the blocks
188 blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
189 rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
190
191 return rearranged.flatten()