@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 210 lines 6.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright (c) 2024, Tri Dao, Albert Gu.
5 # Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/ssd_bmm.py
6
7 # ruff: noqa: E501,SIM102
8
9 import torch
10
11 from vllm.triton_utils import tl, triton
12
13
14 @triton.autotune(
15 configs=[
16 triton.Config(
17 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},
18 num_stages=3,
19 num_warps=8,
20 ),
21 triton.Config(
22 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32},
23 num_stages=4,
24 num_warps=4,
25 ),
26 triton.Config(
27 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
28 num_stages=4,
29 num_warps=4,
30 ),
31 triton.Config(
32 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
33 num_stages=4,
34 num_warps=4,
35 ),
36 triton.Config(
37 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32},
38 num_stages=4,
39 num_warps=4,
40 ),
41 triton.Config(
42 {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
43 num_stages=4,
44 num_warps=4,
45 ),
46 triton.Config(
47 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32},
48 num_stages=5,
49 num_warps=2,
50 ),
51 triton.Config(
52 {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
53 num_stages=5,
54 num_warps=2,
55 ),
56 triton.Config(
57 {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32},
58 num_stages=4,
59 num_warps=2,
60 ),
61 ],
62 key=["chunk_size", "K", "IS_CAUSAL"],
63 )
64 @triton.jit
65 def _bmm_chunk_fwd_kernel(
66 # Pointers to matrices
67 a_ptr,
68 b_ptr,
69 out_ptr,
70 cu_chunk_seqlens_ptr,
71 # Matrix dimensions
72 chunk_size: tl.constexpr,
73 K: tl.constexpr,
74 ngroups: tl.constexpr,
75 stride_a_seqlen: tl.int64,
76 stride_a_head: tl.int64,
77 stride_ak: tl.constexpr,
78 stride_b_seqlen: tl.int64,
79 stride_b_head: tl.int64,
80 stride_bk: tl.constexpr,
81 stride_out_chunk: tl.int64,
82 stride_out_head: tl.int64,
83 stride_outm: tl.int64,
84 stride_outn: tl.constexpr,
85 # Meta-parameters
86 IS_CAUSAL: tl.constexpr,
87 dot_dtype: tl.constexpr,
88 BLOCK_SIZE_M: tl.constexpr,
89 BLOCK_SIZE_N: tl.constexpr,
90 BLOCK_SIZE_K: tl.constexpr,
91 ):
92 pid_ch = tl.program_id(axis=1).to(tl.int64)
93 pid_c = pid_ch // ngroups
94 pid_h = pid_ch - pid_c * ngroups
95 num_pid_n = tl.cdiv(chunk_size, BLOCK_SIZE_N)
96 pid_m = tl.program_id(axis=0) // num_pid_n
97 pid_n = tl.program_id(axis=0) % num_pid_n
98 if IS_CAUSAL:
99 if pid_n * BLOCK_SIZE_N >= (pid_m + 1) * BLOCK_SIZE_M:
100 return
101
102 chunk_seqlen_start = tl.load(cu_chunk_seqlens_ptr + pid_c)
103 chunk_seqlen_end = tl.load(cu_chunk_seqlens_ptr + pid_c + 1)
104
105 a_ptr += chunk_seqlen_start * stride_a_seqlen + pid_h * stride_a_head
106 b_ptr += chunk_seqlen_start * stride_b_seqlen + pid_h * stride_b_head
107
108 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
109 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
110 offs_k = tl.arange(0, BLOCK_SIZE_K)
111 a_ptrs = a_ptr + (offs_m[:, None] * stride_a_seqlen + offs_k[None, :] * stride_ak)
112 b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_b_seqlen)
113 chunk_size_limit = chunk_seqlen_end - chunk_seqlen_start
114
115 acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
116
117 # compute a * b.T
118 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
119 a = tl.load(
120 a_ptrs,
121 mask=(offs_m[:, None] < chunk_size_limit)
122 & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
123 other=0.0,
124 ).to(dot_dtype)
125 b = tl.load(
126 b_ptrs,
127 mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K)
128 & (offs_n[None, :] < chunk_size_limit),
129 other=0.0,
130 ).to(dot_dtype)
131 acc += tl.dot(a, b)
132 a_ptrs += BLOCK_SIZE_K * stride_ak
133 b_ptrs += BLOCK_SIZE_K * stride_bk
134
135 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
136 offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
137
138 out = acc.to(out_ptr.dtype.element_ty)
139 out_ptr += pid_c * stride_out_chunk + pid_h * stride_out_head
140 out_ptrs = out_ptr + (stride_outm * offs_m[:, None] + offs_n[None, :] * stride_outn)
141 tl.store(
142 out_ptrs,
143 out,
144 mask=(offs_m[:, None] < chunk_size) & (offs_n[None, :] < chunk_size),
145 )
146
147
148 def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtype=None):
149 """Argument:
150 a: (seqlen, ngroups, k)
151 b: (seqlen, ngroups, k)
152 chunk_size: int
153 cu_chunk_seq_lens: (nchunks+1,)
154 causal: if True, then out[i, j] for i > j will be arbitrary, only out[i, j] for i <= j are
155 guaranteed to be correct.
156
157 Return:
158 out: (nchunks, ngroups, chunk_size, chunk_size)
159
160 """
161 seqlen, ngroups, k = a.shape
162 assert b.shape == a.shape
163 if a.stride(-1) != 1 and a.stride(0) != 1:
164 a = a.contiguous()
165 if b.stride(-1) != 1 and b.stride(0) != 1:
166 b = b.contiguous()
167
168 nchunks = len(cu_chunk_seqlens) - 1
169 # Allocates output.
170 out_dtype = a.dtype if output_dtype is None else output_dtype
171 out = torch.empty(
172 (nchunks, ngroups, chunk_size, chunk_size), device=a.device, dtype=out_dtype
173 )
174 dot_dtype = (
175 tl.bfloat16
176 if a.dtype == torch.bfloat16 or b.dtype == torch.bfloat16
177 else (
178 tl.float16
179 if a.dtype == torch.float16 or b.dtype == torch.float16
180 else tl.float32
181 )
182 )
183 grid = lambda META: (
184 triton.cdiv(chunk_size, META["BLOCK_SIZE_M"])
185 * triton.cdiv(chunk_size, META["BLOCK_SIZE_N"]),
186 nchunks * ngroups,
187 )
188 with torch.accelerator.device_index(a.device.index):
189 _bmm_chunk_fwd_kernel[grid](
190 a_ptr=a,
191 b_ptr=b,
192 out_ptr=out,
193 cu_chunk_seqlens_ptr=cu_chunk_seqlens,
194 chunk_size=chunk_size,
195 K=k,
196 ngroups=ngroups,
197 stride_a_seqlen=a.stride(0),
198 stride_a_head=a.stride(1),
199 stride_ak=a.stride(2),
200 stride_b_seqlen=b.stride(0),
201 stride_b_head=b.stride(1),
202 stride_bk=b.stride(2),
203 stride_out_chunk=out.stride(0),
204 stride_out_head=out.stride(1),
205 stride_outm=out.stride(-2),
206 stride_outn=out.stride(-1),
207 IS_CAUSAL=causal,
208 dot_dtype=dot_dtype,
209 )
210 return out