| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Adapted from FlashInfer's flashinfer/comm/mnnvl_cutedsl/cute_dsl_primitives.py |
| 4 | # and runtime.py (flashinfer-ai/flashinfer@139af6f6). |
| 5 | |
| 6 | """PTX building blocks and launch helpers for the DSV4.1 LL all-reduce.""" |
| 7 | |
| 8 | import cuda.bindings.driver as cuda |
| 9 | import cutlass |
| 10 | import cutlass.cute as cute |
| 11 | import torch |
| 12 | from cutlass import Float32, Int32, Int64, Uint16, Uint32 |
| 13 | from cutlass._mlir import ir |
| 14 | from cutlass._mlir.dialects import llvm, vector |
| 15 | from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor |
| 16 | from cutlass.cutlass_dsl import T, dsl_user_op |
| 17 | |
| 18 | WARP_SIZE = 32 |
| 19 | VEC_BF16 = 8 |
| 20 | QUAD_BF16 = 4 |
| 21 | NEGATIVE_ZERO_BF16_BITS = 0x8000 |
| 22 | NEGATIVE_ZERO_BF16_PAIR = 0x80008000 |
| 23 | |
| 24 | |
| 25 | def _asm(result, operands, text, constraints, *, side_effects, loc, ip): |
| 26 | return llvm.inline_asm( |
| 27 | result, |
| 28 | operands, |
| 29 | text, |
| 30 | constraints, |
| 31 | has_side_effects=side_effects, |
| 32 | is_align_stack=False, |
| 33 | asm_dialect=llvm.AsmDialect.AD_ATT, |
| 34 | loc=loc, |
| 35 | ip=ip, |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | @dsl_user_op |
| 40 | def load_global_u32x4( |
| 41 | pointer: cute.Pointer, |
| 42 | *, |
| 43 | volatile: cutlass.Constexpr[bool] = False, |
| 44 | loc=None, |
| 45 | ip=None, |
| 46 | ): |
| 47 | address = pointer.toint(loc=loc, ip=ip) |
| 48 | opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32" |
| 49 | loaded = _asm( |
| 50 | llvm.StructType.get_literal([T.i32()] * 4), |
| 51 | [address.ir_value(loc=loc, ip=ip)], |
| 52 | f"{opcode} {{$0, $1, $2, $3}}, [$4];", |
| 53 | "=r,=r,=r,=r,l", |
| 54 | side_effects=volatile, |
| 55 | loc=loc, |
| 56 | ip=ip, |
| 57 | ) |
| 58 | packed = vector.from_elements( |
| 59 | ir.VectorType.get([4], T.i32(), loc=loc), |
| 60 | [ |
| 61 | llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) |
| 62 | for index in range(4) |
| 63 | ], |
| 64 | loc=loc, |
| 65 | ip=ip, |
| 66 | ) |
| 67 | return cute.TensorSSA(packed, 4, Uint32) |
| 68 | |
| 69 | |
| 70 | @dsl_user_op |
| 71 | def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None): |
| 72 | address = pointer.toint(loc=loc, ip=ip) |
| 73 | loaded = _asm( |
| 74 | llvm.StructType.get_literal([T.i32()] * 2), |
| 75 | [address.ir_value(loc=loc, ip=ip)], |
| 76 | "ld.global.v2.u32 {$0, $1}, [$2];", |
| 77 | "=r,=r,l", |
| 78 | side_effects=False, |
| 79 | loc=loc, |
| 80 | ip=ip, |
| 81 | ) |
| 82 | packed = vector.from_elements( |
| 83 | ir.VectorType.get([2], T.i32(), loc=loc), |
| 84 | [ |
| 85 | llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip) |
| 86 | for index in range(2) |
| 87 | ], |
| 88 | loc=loc, |
| 89 | ip=ip, |
| 90 | ) |
| 91 | return cute.TensorSSA(packed, 2, Uint32) |
| 92 | |
| 93 | |
| 94 | @dsl_user_op |
| 95 | def load_global_bf16_as_f32(address: Int64, *, loc=None, ip=None) -> Float32: |
| 96 | return Float32( |
| 97 | _asm( |
| 98 | T.f32(), |
| 99 | [address.ir_value(loc=loc, ip=ip)], |
| 100 | "{\n\t.reg .b16 bits;\n\tld.global.b16 bits, [$1];\n\t" |
| 101 | "cvt.f32.bf16 $0, bits;\n\t}", |
| 102 | "=f,l", |
| 103 | side_effects=False, |
| 104 | loc=loc, |
| 105 | ip=ip, |
| 106 | ) |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | @dsl_user_op |
| 111 | def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None: |
| 112 | words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)] |
| 113 | _asm( |
| 114 | None, |
| 115 | [address.ir_value(loc=loc, ip=ip), *words], |
| 116 | "st.global.v4.u32 [$0], {$1, $2, $3, $4};", |
| 117 | "l,r,r,r,r", |
| 118 | side_effects=True, |
| 119 | loc=loc, |
| 120 | ip=ip, |
| 121 | ) |
| 122 | |
| 123 | |
| 124 | @dsl_user_op |
| 125 | def store_lamport_sentinel_u32x4(address: Int64, *, loc=None, ip=None) -> None: |
| 126 | sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip) |
| 127 | _asm( |
| 128 | None, |
| 129 | [address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel], |
| 130 | "st.global.v4.u32 [$0], {$1, $2, $3, $4};", |
| 131 | "l,r,r,r,r", |
| 132 | side_effects=True, |
| 133 | loc=loc, |
| 134 | ip=ip, |
| 135 | ) |
| 136 | |
| 137 | |
| 138 | @dsl_user_op |
| 139 | def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32: |
| 140 | address = pointer.toint(loc=loc, ip=ip) |
| 141 | return Uint32( |
| 142 | _asm( |
| 143 | T.i32(), |
| 144 | [address.ir_value(loc=loc, ip=ip)], |
| 145 | "ld.volatile.global.u32 $0, [$1];", |
| 146 | "=r,l", |
| 147 | side_effects=True, |
| 148 | loc=loc, |
| 149 | ip=ip, |
| 150 | ) |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | @dsl_user_op |
| 155 | def store_global_u32( |
| 156 | pointer: cute.Pointer, value: Uint32, *, loc=None, ip=None |
| 157 | ) -> None: |
| 158 | address = pointer.toint(loc=loc, ip=ip) |
| 159 | _asm( |
| 160 | None, |
| 161 | [address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], |
| 162 | "st.global.u32 [$0], $1;", |
| 163 | "l,r", |
| 164 | side_effects=True, |
| 165 | loc=loc, |
| 166 | ip=ip, |
| 167 | ) |
| 168 | |
| 169 | |
| 170 | @dsl_user_op |
| 171 | def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None): |
| 172 | values = llvm.bitcast( |
| 173 | ir.VectorType.get([VEC_BF16], cutlass.BFloat16.mlir_type, loc=loc), |
| 174 | packed.ir_value(loc=loc, ip=ip), |
| 175 | loc=loc, |
| 176 | ip=ip, |
| 177 | ) |
| 178 | return cute.TensorSSA(values, VEC_BF16, cutlass.BFloat16) |
| 179 | |
| 180 | |
| 181 | @dsl_user_op |
| 182 | def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None): |
| 183 | values = llvm.bitcast( |
| 184 | ir.VectorType.get([QUAD_BF16], cutlass.BFloat16.mlir_type, loc=loc), |
| 185 | packed.ir_value(loc=loc, ip=ip), |
| 186 | loc=loc, |
| 187 | ip=ip, |
| 188 | ) |
| 189 | return cute.TensorSSA(values, QUAD_BF16, cutlass.BFloat16) |
| 190 | |
| 191 | |
| 192 | @dsl_user_op |
| 193 | def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None): |
| 194 | packed = llvm.bitcast( |
| 195 | ir.VectorType.get([2], T.i32(), loc=loc), |
| 196 | values.ir_value(loc=loc, ip=ip), |
| 197 | loc=loc, |
| 198 | ip=ip, |
| 199 | ) |
| 200 | return cute.TensorSSA(packed, 2, Uint32) |
| 201 | |
| 202 | |
| 203 | @dsl_user_op |
| 204 | def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None): |
| 205 | packed = llvm.bitcast( |
| 206 | ir.VectorType.get([4], T.i32(), loc=loc), |
| 207 | values.ir_value(loc=loc, ip=ip), |
| 208 | loc=loc, |
| 209 | ip=ip, |
| 210 | ) |
| 211 | return cute.TensorSSA(packed, 4, Uint32) |
| 212 | |
| 213 | |
| 214 | @cute.jit |
| 215 | def sanitize_negative_zero_u32(word: Uint32) -> Uint32: |
| 216 | low = Uint16(word & Uint32(0xFFFF)) |
| 217 | high = Uint16(word >> Uint32(16)) |
| 218 | if low == Uint16(NEGATIVE_ZERO_BF16_BITS): |
| 219 | word = word & Uint32(0xFFFF0000) |
| 220 | if high == Uint16(NEGATIVE_ZERO_BF16_BITS): |
| 221 | word = word & Uint32(0x0000FFFF) |
| 222 | return word |
| 223 | |
| 224 | |
| 225 | @cute.jit |
| 226 | def sanitize_negative_zero_u32x4(packed): |
| 227 | sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32) |
| 228 | for index in cutlass.range_constexpr(4): |
| 229 | sanitized[index] = sanitize_negative_zero_u32(packed[index]) |
| 230 | return sanitized.load() |
| 231 | |
| 232 | |
| 233 | @cute.jit |
| 234 | def sanitize_negative_zero_u32x2(packed): |
| 235 | sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32) |
| 236 | for index in cutlass.range_constexpr(2): |
| 237 | sanitized[index] = sanitize_negative_zero_u32(packed[index]) |
| 238 | return sanitized.load() |
| 239 | |
| 240 | |
| 241 | @cute.jit |
| 242 | def fragment_has_negative_zero(packed): |
| 243 | dirty = False |
| 244 | for index in cutlass.range_constexpr(4): |
| 245 | word = packed[index] |
| 246 | dirty = ( |
| 247 | dirty |
| 248 | | (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) |
| 249 | | (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS)) |
| 250 | ) |
| 251 | return dirty |
| 252 | |
| 253 | |
| 254 | @dsl_user_op |
| 255 | def map_shared_to_peer( |
| 256 | smem_pointer: cute.Pointer, peer_rank: Int32, *, loc=None, ip=None |
| 257 | ) -> Int32: |
| 258 | address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) |
| 259 | return Int32( |
| 260 | _asm( |
| 261 | T.i32(), |
| 262 | [address, peer_rank.ir_value(loc=loc, ip=ip)], |
| 263 | "mapa.shared::cluster.u32 $0, $1, $2;", |
| 264 | "=r,r,r", |
| 265 | side_effects=False, |
| 266 | loc=loc, |
| 267 | ip=ip, |
| 268 | ) |
| 269 | ) |
| 270 | |
| 271 | |
| 272 | @dsl_user_op |
| 273 | def store_shared_cluster_f32( |
| 274 | remote_address: Int32, value: Float32, *, loc=None, ip=None |
| 275 | ) -> None: |
| 276 | _asm( |
| 277 | None, |
| 278 | [remote_address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)], |
| 279 | "st.shared::cluster.f32 [$0], $1;", |
| 280 | "r,f", |
| 281 | side_effects=True, |
| 282 | loc=loc, |
| 283 | ip=ip, |
| 284 | ) |
| 285 | |
| 286 | |
| 287 | @dsl_user_op |
| 288 | def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32: |
| 289 | return Uint32( |
| 290 | _asm( |
| 291 | T.i32(), |
| 292 | [value.ir_value(loc=loc, ip=ip)], |
| 293 | "{\n\t.reg .b16 bits;\n\tcvt.rn.bf16.f32 bits, $1;\n\t" |
| 294 | "cvt.u32.u16 $0, bits;\n\t}", |
| 295 | "=r,f", |
| 296 | side_effects=False, |
| 297 | loc=loc, |
| 298 | ip=ip, |
| 299 | ) |
| 300 | ) |
| 301 | |
| 302 | |
| 303 | @dsl_user_op |
| 304 | def shuffle_sync_idx_u32( |
| 305 | value: Uint32, source_lane: Int32, *, loc=None, ip=None |
| 306 | ) -> Uint32: |
| 307 | return Uint32( |
| 308 | _asm( |
| 309 | T.i32(), |
| 310 | [value.ir_value(loc=loc, ip=ip), source_lane.ir_value(loc=loc, ip=ip)], |
| 311 | "shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;", |
| 312 | "=r,r,r", |
| 313 | # Preserve full-warp execution across later divergent consumers. |
| 314 | side_effects=True, |
| 315 | loc=loc, |
| 316 | ip=ip, |
| 317 | ) |
| 318 | ) |
| 319 | |
| 320 | |
| 321 | @dsl_user_op |
| 322 | def stmc_bf16x2(address: Int64, packed: Uint32, *, loc=None, ip=None) -> None: |
| 323 | _asm( |
| 324 | None, |
| 325 | [address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)], |
| 326 | "multimem.st.relaxed.sys.global.bf16x2 [$0], $1;", |
| 327 | "l,r", |
| 328 | side_effects=True, |
| 329 | loc=loc, |
| 330 | ip=ip, |
| 331 | ) |
| 332 | |
| 333 | |
| 334 | @dsl_user_op |
| 335 | def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None: |
| 336 | words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)] |
| 337 | _asm( |
| 338 | None, |
| 339 | [address.ir_value(loc=loc, ip=ip), *words], |
| 340 | "multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};", |
| 341 | "l,r,r", |
| 342 | side_effects=True, |
| 343 | loc=loc, |
| 344 | ip=ip, |
| 345 | ) |
| 346 | |
| 347 | |
| 348 | @dsl_user_op |
| 349 | def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None: |
| 350 | words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)] |
| 351 | _asm( |
| 352 | None, |
| 353 | [address.ir_value(loc=loc, ip=ip), *words], |
| 354 | "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};", |
| 355 | "l,r,r,r,r", |
| 356 | side_effects=True, |
| 357 | loc=loc, |
| 358 | ip=ip, |
| 359 | ) |
| 360 | |
| 361 | |
| 362 | class _GraphSafeDLPack: |
| 363 | __slots__ = ("tensor",) |
| 364 | |
| 365 | def __init__(self, tensor: torch.Tensor) -> None: |
| 366 | self.tensor = tensor |
| 367 | |
| 368 | def __dlpack__(self, stream=None): |
| 369 | # stream=-1 skips producer sync; CuTe launches on the current stream, |
| 370 | # which is the capturing one under CUDA graphs. |
| 371 | return self.tensor.__dlpack__(stream=-1) |
| 372 | |
| 373 | def __dlpack_device__(self): |
| 374 | return self.tensor.__dlpack_device__() |
| 375 | |
| 376 | |
| 377 | def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor: |
| 378 | return from_dlpack(_GraphSafeDLPack(tensor.detach()), assumed_align=alignment) |
| 379 | |
| 380 | |
| 381 | def to_cute_dynamic( |
| 382 | tensor: torch.Tensor, alignment: int, *, divisibility: int |
| 383 | ) -> cute.Tensor: |
| 384 | return to_cute(tensor, alignment).mark_compact_shape_dynamic( |
| 385 | mode=0, divisibility=divisibility |
| 386 | ) |
| 387 | |
| 388 | |
| 389 | def make_fake_dynamic_compact_tensor(dtype, *, alignment: int, divisibility: int): |
| 390 | return make_fake_compact_tensor( |
| 391 | dtype, (cute.sym_int32(divisibility=divisibility),), assumed_align=alignment |
| 392 | ) |
| 393 | |
| 394 | |
| 395 | def current_cu_stream() -> cuda.CUstream: |
| 396 | return cuda.CUstream(torch.cuda.current_stream().cuda_stream) |