@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 76 lines 2.36 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import enum
5
6 from cutlass_library import *
7
8 #
9 # Extend cutlass library with custom types, and missing values
10 #
11
12
13 class VLLMDataType(enum.Enum):
14 u4b8 = enum_auto()
15 u8b128 = enum_auto()
16
17
18 class MixedInputKernelScheduleType(enum.Enum):
19 TmaWarpSpecialized = enum_auto()
20 TmaWarpSpecializedPingpong = enum_auto()
21 TmaWarpSpecializedCooperative = enum_auto()
22
23
24 VLLMDataTypeNames: dict[VLLMDataType | DataType, str] = {
25 **DataTypeNames, # type: ignore
26 **{
27 VLLMDataType.u4b8: "u4b8",
28 VLLMDataType.u8b128: "u8b128",
29 },
30 }
31
32 VLLMDataTypeTag: dict[VLLMDataType | DataType, str] = {
33 **DataTypeTag, # type: ignore
34 **{
35 VLLMDataType.u4b8: "cutlass::vllm_uint4b8_t",
36 VLLMDataType.u8b128: "cutlass::vllm_uint8b128_t",
37 },
38 }
39
40 VLLMDataTypeSize: dict[VLLMDataType | DataType, int] = {
41 **DataTypeSize, # type: ignore
42 **{
43 VLLMDataType.u4b8: 4,
44 VLLMDataType.u8b128: 8,
45 },
46 }
47
48 VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {
49 VLLMDataType.u4b8: "vllm::kU4B8",
50 VLLMDataType.u8b128: "vllm::kU8B128",
51 DataType.u4: "vllm::kU4",
52 DataType.u8: "vllm::kU8",
53 DataType.s4: "vllm::kS4",
54 DataType.s8: "vllm::kS8",
55 DataType.f16: "vllm::kFloat16",
56 DataType.bf16: "vllm::kBfloat16",
57 }
58
59 VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {
60 DataType.u8: "torch::headeronly::ScalarType::Byte",
61 DataType.s8: "torch::headeronly::ScalarType::Char",
62 DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn",
63 DataType.s32: "torch::headeronly::ScalarType::Int",
64 DataType.f16: "torch::headeronly::ScalarType::Half",
65 DataType.bf16: "torch::headeronly::ScalarType::BFloat16",
66 DataType.f32: "torch::headeronly::ScalarType::Float",
67 }
68
69 VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {
70 **KernelScheduleTag, # type: ignore
71 **{
72 MixedInputKernelScheduleType.TmaWarpSpecialized: "cutlass::gemm::KernelTmaWarpSpecialized", # noqa: E501
73 MixedInputKernelScheduleType.TmaWarpSpecializedPingpong: "cutlass::gemm::KernelTmaWarpSpecializedPingpong", # noqa: E501
74 MixedInputKernelScheduleType.TmaWarpSpecializedCooperative: "cutlass::gemm::KernelTmaWarpSpecializedCooperative", # noqa: E501
75 },
76 }