@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 204 lines 5.65 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4 from cutlass import (
5 BFloat16,
6 Float8E4M3FN,
7 Float16,
8 Float32,
9 Int32,
10 Int64,
11 Uint32,
12 cute,
13 )
14 from cutlass._mlir import ir
15 from cutlass._mlir.dialects import llvm, vector
16 from cutlass.cute.nvgpu import cpasync
17 from cutlass.cutlass_dsl import T, dsl_user_op
18
19 _TORCH_TO_CUTE_DTYPE = {
20 torch.bfloat16: BFloat16,
21 torch.float8_e4m3fn: Float8E4M3FN,
22 torch.float32: Float32,
23 }
24
25
26 def torch_to_cute_dtype(dtype: torch.dtype):
27 try:
28 return _TORCH_TO_CUTE_DTYPE[dtype]
29 except KeyError as exc:
30 raise TypeError(f"Unsupported CuTe dtype {dtype}") from exc
31
32
33 _CUTE_TO_PTX_DTYPE = {
34 BFloat16: "bf16",
35 Float16: "f16",
36 Float8E4M3FN: "e4m3",
37 Float32: "f32",
38 }
39
40 # https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197
41 EVICT_NORMAL = Int64(0x1000000000000000)
42 EVICT_FIRST = Int64(0x12F0000000000000)
43 EVICT_LAST = Int64(0x14F0000000000000)
44
45
46 @dsl_user_op
47 def recast_val(x, dtype, *, loc=None, ip=None):
48 return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
49
50
51 @dsl_user_op
52 def to_cta0_smem(ptr: cute.Pointer, *, loc=None, ip=None):
53 return cute.make_ptr(
54 ptr.dtype,
55 ptr.toint(loc=loc, ip=ip) & 0xFEFF_FFFF,
56 cute.AddressSpace.smem,
57 assumed_align=8,
58 loc=loc,
59 ip=ip,
60 )
61
62
63 def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
64 """A simple helper that wraps group_modes() and tma_partition()
65 NOTE: this should be called WITHOUT cute.elect_one()
66 """
67 if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
68 gmem = src
69 smem = dst
70 elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
71 smem = src
72 gmem = dst
73 else:
74 raise ValueError
75
76 s_part, g_part = cpasync.tma_partition(
77 atom,
78 0,
79 cute.make_layout(1),
80 cute.group_modes(smem, 0),
81 cute.group_modes(gmem, 0),
82 )
83
84 if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
85 cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy)
86 elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
87 cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
88 else:
89 raise ValueError
90
91
92 # can't find the equivalent in nvvm
93 @dsl_user_op
94 def fence_before_tma_store(*, loc=None, ip=None):
95 llvm.inline_asm(
96 T.i32(),
97 [],
98 "mov.u32 $0, 0;\n\t"
99 "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;",
100 "=r",
101 has_side_effects=True,
102 is_align_stack=False,
103 loc=loc,
104 ip=ip,
105 )
106
107
108 @dsl_user_op
109 def mma_sync(a, b, c: cute.Tensor, *, loc=None, ip=None):
110 a_ty = _CUTE_TO_PTX_DTYPE[a.element_type]
111 b_ty = _CUTE_TO_PTX_DTYPE[b.element_type]
112 c_ty = _CUTE_TO_PTX_DTYPE[c.element_type]
113 mlir_ty = c.element_type.mlir_type
114 K = 256 // a.element_type.width # 32B
115
116 # Recast expects tensor-backed fragments; materialize SSA fragments here so
117 # callsites can pass converted FP8 fragments directly.
118 if isinstance(a, cute.TensorSSA):
119 a_ = cute.make_rmem_tensor_like(a)
120 a_.store(a, loc=loc, ip=ip)
121 a = a_
122 if isinstance(b, cute.TensorSSA):
123 b_ = cute.make_rmem_tensor_like(b)
124 b_.store(b, loc=loc, ip=ip)
125 b = b_
126
127 a = cute.recast_tensor(a, Int32, loc=loc, ip=ip)
128 b = cute.recast_tensor(b, Int32, loc=loc, ip=ip)
129 out = llvm.inline_asm(
130 llvm.StructType.get_literal([mlir_ty] * 4),
131 [a[i].ir_value(loc=loc, ip=ip) for i in range(4)]
132 + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)]
133 + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)],
134 f"mma.sync.aligned.m16n8k{K}.row.col.{c_ty}.{a_ty}.{b_ty}.{c_ty} "
135 "{$0, $1, $2, $3}, "
136 "{$4, $5, $6, $7}, "
137 "{$8, $9}, "
138 "{$10, $11, $12, $13};",
139 "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
140 has_side_effects=False,
141 is_align_stack=False,
142 loc=loc,
143 ip=ip,
144 )
145 vec = vector.from_elements(
146 ir.VectorType.get([4], mlir_ty, loc=loc),
147 [llvm.extractvalue(mlir_ty, out, [i], loc=loc, ip=ip) for i in range(4)],
148 loc=loc,
149 ip=ip,
150 )
151 return cute.TensorSSA(vec, 4, c.element_type)
152
153
154 def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32:
155 out = llvm.inline_asm(
156 T.i32(),
157 [a.ir_value(loc=loc, ip=ip)],
158 f"{asm}.bf16x2 $0, $1;",
159 "=r,r",
160 has_side_effects=False,
161 is_align_stack=False,
162 loc=loc,
163 ip=ip,
164 )
165 return Uint32(out)
166
167
168 def _bf16x2_binary(asm: str, a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
169 out = llvm.inline_asm(
170 T.i32(),
171 [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
172 f"{asm}.bf16x2 $0, $1, $2;",
173 "=r,r,r",
174 has_side_effects=False,
175 is_align_stack=False,
176 loc=loc,
177 ip=ip,
178 )
179 return Uint32(out)
180
181
182 @dsl_user_op
183 def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
184 return _bf16x2_unary("abs", a, loc=loc, ip=ip)
185
186
187 @dsl_user_op
188 def _bf16x2_neg(a: Uint32, *, loc=None, ip=None) -> Uint32:
189 return _bf16x2_unary("neg", a, loc=loc, ip=ip)
190
191
192 @dsl_user_op
193 def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
194 return _bf16x2_binary("max", a, b, loc=loc, ip=ip)
195
196
197 @dsl_user_op
198 def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
199 return _bf16x2_binary("mul.rn", a, b, loc=loc, ip=ip)
200
201
202 @dsl_user_op
203 def _bf16x2_sub(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
204 return _bf16x2_binary("sub.rn", a, b, loc=loc, ip=ip)