| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import itertools |
| 5 | import math |
| 6 | import os |
| 7 | import shutil |
| 8 | from collections.abc import Iterable |
| 9 | from copy import deepcopy |
| 10 | from dataclasses import dataclass, fields |
| 11 | from functools import reduce |
| 12 | |
| 13 | import jinja2 |
| 14 | from vllm_cutlass_library_extension import ( |
| 15 | DataType, |
| 16 | EpilogueScheduleTag, |
| 17 | EpilogueScheduleType, |
| 18 | MixedInputKernelScheduleType, |
| 19 | TileSchedulerTag, |
| 20 | TileSchedulerType, |
| 21 | VLLMDataType, |
| 22 | VLLMDataTypeNames, |
| 23 | VLLMDataTypeSize, |
| 24 | VLLMDataTypeTag, |
| 25 | VLLMDataTypeTorchDataTypeTag, |
| 26 | VLLMDataTypeVLLMScalarTypeTag, |
| 27 | VLLMKernelScheduleTag, |
| 28 | ) |
| 29 | |
| 30 | # |
| 31 | # Generator templating |
| 32 | # |
| 33 | |
| 34 | DISPATCH_TEMPLATE = """ |
| 35 | #include "../machete_mm_launcher.cuh" |
| 36 | |
| 37 | namespace machete { |
| 38 | |
| 39 | {% for impl_config in impl_configs %} |
| 40 | {% set type_sig = gen_type_sig(impl_config.types) -%} |
| 41 | {% for s in impl_config.schedules %} |
| 42 | extern torch::stable::Tensor impl_{{type_sig}}_sch_{{gen_sch_sig(s)}}(MMArgs); |
| 43 | {%- endfor %} |
| 44 | |
| 45 | torch::stable::Tensor mm_dispatch_{{type_sig}}(MMArgs args) { |
| 46 | [[maybe_unused]] auto M = args.A.size(0); |
| 47 | [[maybe_unused]] auto N = args.B.size(1); |
| 48 | [[maybe_unused]] auto K = args.A.size(1); |
| 49 | |
| 50 | if (!args.maybe_schedule) { |
| 51 | {%- for cond, s in impl_config.heuristic %} |
| 52 | {%if cond is not none%}if ({{cond}}) |
| 53 | {%- else %}else |
| 54 | {%- endif %} |
| 55 | return impl_{{type_sig}}_sch_{{ gen_sch_sig(s) }}(args);{% endfor %} |
| 56 | } |
| 57 | |
| 58 | {%- for s in impl_config.schedules %} |
| 59 | if (*args.maybe_schedule == "{{ gen_sch_sig(s) }}") |
| 60 | return impl_{{type_sig}}_sch_{{ gen_sch_sig(s) }}(args); |
| 61 | {%- endfor %} |
| 62 | STD_TORCH_CHECK_NOT_IMPLEMENTED(false, "machete_gemm(..) is not implemented for " |
| 63 | "schedule = ", *args.maybe_schedule); |
| 64 | } |
| 65 | {%- endfor %} |
| 66 | |
| 67 | |
| 68 | static inline std::optional<torch::headeronly::ScalarType> maybe_scalartype( |
| 69 | std::optional<torch::stable::Tensor> const& t) { |
| 70 | if (!t) { |
| 71 | return std::nullopt; |
| 72 | } else { |
| 73 | return t->scalar_type(); |
| 74 | }; |
| 75 | } |
| 76 | |
| 77 | torch::stable::Tensor mm_dispatch(MMArgs args) { |
| 78 | auto out_type = args.maybe_out_type.value_or(args.A.scalar_type()); |
| 79 | auto a_type = args.A.scalar_type(); |
| 80 | auto maybe_g_scales_type = maybe_scalartype(args.maybe_group_scales); |
| 81 | auto maybe_g_zeros_type = maybe_scalartype(args.maybe_group_zeros); |
| 82 | auto maybe_ch_scales_type = maybe_scalartype(args.maybe_channel_scales); |
| 83 | auto maybe_tok_scales_type = maybe_scalartype(args.maybe_token_scales); |
| 84 | |
| 85 | {% for impl_config in impl_configs %} |
| 86 | {% set t = impl_config.types -%} |
| 87 | {% set type_sig = gen_type_sig(t) -%} |
| 88 | if (args.b_type == {{VLLMScalarTypeTag[t.b]}} |
| 89 | && a_type == {{TorchTypeTag[t.a]}} |
| 90 | && out_type == {{TorchTypeTag[t.out]}} |
| 91 | && {%if t.b_group_scale != void -%} |
| 92 | maybe_g_scales_type == {{TorchTypeTag[t.b_group_scale]}} |
| 93 | {%- else %}!maybe_g_scales_type{%endif%} |
| 94 | && {%if t.b_group_zeropoint != void -%} |
| 95 | maybe_g_zeros_type == {{TorchTypeTag[t.b_group_zeropoint]}} |
| 96 | {%- else %}!maybe_g_zeros_type{%endif%} |
| 97 | && {%if t.b_channel_scale != void -%} |
| 98 | maybe_ch_scales_type == {{TorchTypeTag[t.b_channel_scale]}} |
| 99 | {%- else %}!maybe_ch_scales_type{%endif%} |
| 100 | && {%if t.a_token_scale != void -%} |
| 101 | maybe_tok_scales_type == {{TorchTypeTag[t.a_token_scale]}} |
| 102 | {%- else %}!maybe_tok_scales_type{%endif%} |
| 103 | ) { |
| 104 | return mm_dispatch_{{type_sig}}(args); |
| 105 | } |
| 106 | {%- endfor %} |
| 107 | |
| 108 | STD_TORCH_CHECK_NOT_IMPLEMENTED( |
| 109 | false, "machete_mm(..) is not implemented for " |
| 110 | "a_type=", torch::headeronly::toString(args.A.scalar_type()), |
| 111 | ", b_type=", args.b_type.str(), |
| 112 | ", out_type=", torch::headeronly::toString(out_type), |
| 113 | ", with_group_scale_type=", maybe_g_scales_type |
| 114 | ? torch::headeronly::toString(*maybe_g_scales_type) : "None", |
| 115 | ", with_group_zeropoint_type=", maybe_g_zeros_type |
| 116 | ? torch::headeronly::toString(*maybe_g_zeros_type) : "None", |
| 117 | ", with_channel_scale_type=", maybe_ch_scales_type |
| 118 | ? torch::headeronly::toString(*maybe_ch_scales_type) : "None", |
| 119 | ", with_token_scale_type=", maybe_tok_scales_type |
| 120 | ? torch::headeronly::toString(*maybe_tok_scales_type) : "None", |
| 121 | "; implemented types are: \\n", |
| 122 | {%- for impl_config in impl_configs %} |
| 123 | {% set t = impl_config.types -%} |
| 124 | "\\t{{gen_type_option_name(t)}}\\n", |
| 125 | {%- endfor %} |
| 126 | ""); |
| 127 | } |
| 128 | |
| 129 | std::vector<std::string> supported_schedules_dispatch( |
| 130 | SupportedSchedulesArgs args) { |
| 131 | auto out_type = args.maybe_out_type.value_or(args.a_type); |
| 132 | |
| 133 | {% for impl_config in impl_configs %} |
| 134 | {% set t = impl_config.types -%} |
| 135 | {% set schs = impl_config.schedules -%} |
| 136 | if (args.b_type == {{VLLMScalarTypeTag[t.b]}} |
| 137 | && args.a_type == {{TorchTypeTag[t.a]}} |
| 138 | && out_type == {{TorchTypeTag[t.out]}} |
| 139 | && {%if t.b_group_scale != void -%} |
| 140 | args.maybe_group_scales_type == {{TorchTypeTag[t.b_group_scale]}} |
| 141 | {%- else %}!args.maybe_group_scales_type{%endif%} |
| 142 | && {%if t.b_group_zeropoint != void-%} |
| 143 | args.maybe_group_zeros_type == {{TorchTypeTag[t.b_group_zeropoint]}} |
| 144 | {%- else %}!args.maybe_group_zeros_type{%endif%} |
| 145 | ) { |
| 146 | return { |
| 147 | {%- for s in impl_config.schedules %} |
| 148 | "{{gen_sch_sig(s)}}"{% if not loop.last %},{% endif %} |
| 149 | {%- endfor %} |
| 150 | }; |
| 151 | } |
| 152 | {%- endfor %} |
| 153 | |
| 154 | return {}; |
| 155 | }; |
| 156 | |
| 157 | }; // namespace machete |
| 158 | """ |
| 159 | |
| 160 | IMPL_TEMPLATE = """ |
| 161 | #include "../machete_mm_launcher.cuh" |
| 162 | |
| 163 | namespace machete { |
| 164 | |
| 165 | {% for sch in unique_schedules(impl_configs) %} |
| 166 | {% set sch_sig = gen_sch_sig(sch) -%} |
| 167 | struct sch_{{sch_sig}} { |
| 168 | using TileShapeNM = Shape<{{ |
| 169 | to_cute_constant(sch.tile_shape_mn)|join(', ')}}>; |
| 170 | using ClusterShape = Shape<{{ |
| 171 | to_cute_constant(sch.cluster_shape_mnk)|join(', ')}}>; |
| 172 | // TODO: Reimplement |
| 173 | // using KernelSchedule = {{KernelScheduleTag[sch.kernel_schedule]}}; |
| 174 | using EpilogueSchedule = {{EpilogueScheduleTag[sch.epilogue_schedule]}}; |
| 175 | using TileScheduler = {{TileSchedulerTag[sch.tile_scheduler]}}; |
| 176 | using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; |
| 177 | }; |
| 178 | {% endfor %} |
| 179 | |
| 180 | {% for impl_config in impl_configs %} |
| 181 | {% set t = impl_config.types -%} |
| 182 | {% set schs = impl_config.schedules -%} |
| 183 | {% set type_sig = gen_type_sig(t) -%} |
| 184 | |
| 185 | template<typename Sch> |
| 186 | using Kernel_{{type_sig}} = MacheteKernelTemplate< |
| 187 | {{DataTypeTag[t.a]}}, // ElementA |
| 188 | {{DataTypeTag[t.b]}}, // ElementB |
| 189 | {{DataTypeTag[t.out]}}, // ElementD |
| 190 | {{DataTypeTag[t.accumulator]}}, // Accumulator |
| 191 | {{DataTypeTag[t.b_group_scale]}}, // GroupScaleT |
| 192 | {{DataTypeTag[t.b_group_zeropoint]}}, // GroupZeroT |
| 193 | {{DataTypeTag[t.b_channel_scale]}}, // ChannelScaleT |
| 194 | {{DataTypeTag[t.a_token_scale]}}, // TokenScaleT |
| 195 | cutlass::gemm::KernelTmaWarpSpecializedCooperative, |
| 196 | Sch>; |
| 197 | |
| 198 | {% for sch in schs %} |
| 199 | {% set sch_sig = gen_sch_sig(sch) -%} |
| 200 | torch::stable::Tensor |
| 201 | impl_{{type_sig}}_sch_{{sch_sig}}(MMArgs args) { |
| 202 | return run_impl<Kernel_{{type_sig}}<sch_{{sch_sig}}>>(args); |
| 203 | } |
| 204 | {%- endfor %} |
| 205 | {%- endfor %} |
| 206 | |
| 207 | }; // namespace machete |
| 208 | """ |
| 209 | |
| 210 | PREPACK_TEMPLATE = """ |
| 211 | #include "../machete_prepack_launcher.cuh" |
| 212 | |
| 213 | namespace machete { |
| 214 | |
| 215 | torch::stable::Tensor prepack_B_dispatch(PrepackBArgs args) { |
| 216 | auto convert_type = args.maybe_group_scales_type.value_or(args.a_type); |
| 217 | {%- for t in types %} |
| 218 | {% set b_type = unsigned_type_with_bitwidth(t.b_num_bits) %} |
| 219 | if (args.a_type == {{TorchTypeTag[t.a]}} |
| 220 | && args.b_type.size_bits() == {{t.b_num_bits}} |
| 221 | && convert_type == {{TorchTypeTag[t.convert]}}) { |
| 222 | return prepack_impl< |
| 223 | PrepackedLayoutBTemplate< |
| 224 | {{DataTypeTag[t.a]}}, // ElementA |
| 225 | {{DataTypeTag[b_type]}}, // ElementB |
| 226 | {{DataTypeTag[t.convert]}}, // ElementConvert |
| 227 | {{DataTypeTag[t.accumulator]}}, // Accumulator |
| 228 | cutlass::layout::ColumnMajor, |
| 229 | cutlass::gemm::KernelTmaWarpSpecializedCooperative> |
| 230 | >(args.B); |
| 231 | } |
| 232 | {%- endfor %} |
| 233 | |
| 234 | STD_TORCH_CHECK_NOT_IMPLEMENTED(false, |
| 235 | "prepack_B_dispatch(..) is not implemented for " |
| 236 | "atype = ", torch::headeronly::toString(args.a_type), |
| 237 | ", b_type = ", args.b_type.str(), |
| 238 | ", with_group_scales_type= ", args.maybe_group_scales_type ? |
| 239 | torch::headeronly::toString(*args.maybe_group_scales_type) : "None"); |
| 240 | } |
| 241 | |
| 242 | }; // namespace machete |
| 243 | """ |
| 244 | |
| 245 | TmaMI = MixedInputKernelScheduleType.TmaWarpSpecializedCooperative |
| 246 | TmaCoop = EpilogueScheduleType.TmaWarpSpecializedCooperative |
| 247 | |
| 248 | |
| 249 | @dataclass(frozen=True) |
| 250 | class ScheduleConfig: |
| 251 | tile_shape_mn: tuple[int, int] |
| 252 | cluster_shape_mnk: tuple[int, int, int] |
| 253 | kernel_schedule: MixedInputKernelScheduleType |
| 254 | epilogue_schedule: EpilogueScheduleType |
| 255 | tile_scheduler: TileSchedulerType |
| 256 | |
| 257 | |
| 258 | @dataclass(frozen=True) |
| 259 | class TypeConfig: |
| 260 | a: DataType |
| 261 | b: DataType | VLLMDataType |
| 262 | b_group_scale: DataType |
| 263 | b_group_zeropoint: DataType |
| 264 | b_channel_scale: DataType |
| 265 | a_token_scale: DataType |
| 266 | out: DataType |
| 267 | accumulator: DataType |
| 268 | |
| 269 | |
| 270 | @dataclass(frozen=True) |
| 271 | class PrepackTypeConfig: |
| 272 | a: DataType |
| 273 | b_num_bits: int |
| 274 | convert: DataType |
| 275 | accumulator: DataType |
| 276 | |
| 277 | |
| 278 | @dataclass |
| 279 | class ImplConfig: |
| 280 | types: TypeConfig |
| 281 | schedules: list[ScheduleConfig] |
| 282 | heuristic: list[tuple[str | None, ScheduleConfig]] |
| 283 | |
| 284 | |
| 285 | def generate_sch_sig(schedule_config: ScheduleConfig) -> str: |
| 286 | tile_shape = ( |
| 287 | f"{schedule_config.tile_shape_mn[0]}x{schedule_config.tile_shape_mn[1]}" |
| 288 | ) |
| 289 | cluster_shape = ( |
| 290 | f"{schedule_config.cluster_shape_mnk[0]}" |
| 291 | f"x{schedule_config.cluster_shape_mnk[1]}" |
| 292 | f"x{schedule_config.cluster_shape_mnk[2]}" |
| 293 | ) |
| 294 | kernel_schedule = VLLMKernelScheduleTag[schedule_config.kernel_schedule].split( |
| 295 | "::" |
| 296 | )[-1] |
| 297 | epilogue_schedule = EpilogueScheduleTag[schedule_config.epilogue_schedule].split( |
| 298 | "::" |
| 299 | )[-1] |
| 300 | tile_scheduler = TileSchedulerTag[schedule_config.tile_scheduler].split("::")[-1] |
| 301 | |
| 302 | return ( |
| 303 | f"{tile_shape}_{cluster_shape}_{kernel_schedule}" |
| 304 | f"_{epilogue_schedule}_{tile_scheduler}" |
| 305 | ) |
| 306 | |
| 307 | |
| 308 | # mostly unique shorter sch_sig |
| 309 | def generate_terse_sch_sig(schedule_config: ScheduleConfig) -> str: |
| 310 | kernel_terse_names_replace = { |
| 311 | "KernelTmaWarpSpecializedCooperative": "TmaMI_", |
| 312 | "TmaWarpSpecializedCooperative_": "TmaCoop_", |
| 313 | "StreamKScheduler": "streamK", |
| 314 | } |
| 315 | |
| 316 | sch_sig = generate_sch_sig(schedule_config) |
| 317 | for orig, terse in kernel_terse_names_replace.items(): |
| 318 | sch_sig = sch_sig.replace(orig, terse) |
| 319 | return sch_sig |
| 320 | |
| 321 | |
| 322 | # unique type_name |
| 323 | def generate_type_signature(kernel_types: TypeConfig): |
| 324 | return str( |
| 325 | "".join( |
| 326 | [ |
| 327 | VLLMDataTypeNames[getattr(kernel_types, field.name)] |
| 328 | for field in fields(TypeConfig) |
| 329 | ] |
| 330 | ) |
| 331 | ) |
| 332 | |
| 333 | |
| 334 | def generate_type_option_name(kernel_types: TypeConfig): |
| 335 | return ", ".join( |
| 336 | [ |
| 337 | f"{field.name.replace('b_', 'with_') + '_type'}=" |
| 338 | + VLLMDataTypeNames[getattr(kernel_types, field.name)] |
| 339 | for field in fields(TypeConfig) |
| 340 | ] |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | def is_power_of_two(n): |
| 345 | return (n != 0) and (n & (n - 1) == 0) |
| 346 | |
| 347 | |
| 348 | def to_cute_constant(value: list[int]): |
| 349 | def _to_cute_constant(value: int): |
| 350 | if is_power_of_two(value): |
| 351 | return f"_{value}" |
| 352 | else: |
| 353 | return f"Int<{value}>" |
| 354 | |
| 355 | if isinstance(value, Iterable): |
| 356 | return [_to_cute_constant(value) for value in value] |
| 357 | else: |
| 358 | return _to_cute_constant(value) |
| 359 | |
| 360 | |
| 361 | def unique_schedules(impl_configs: list[ImplConfig]): |
| 362 | # Use dict over set for deterministic ordering |
| 363 | return list( |
| 364 | { |
| 365 | sch: None for impl_config in impl_configs for sch in impl_config.schedules |
| 366 | }.keys() |
| 367 | ) |
| 368 | |
| 369 | |
| 370 | def unsigned_type_with_bitwidth(num_bits): |
| 371 | return { |
| 372 | 4: DataType.u4, |
| 373 | 8: DataType.u8, |
| 374 | 16: DataType.u16, |
| 375 | 32: DataType.u32, |
| 376 | 64: DataType.u64, |
| 377 | }[num_bits] |
| 378 | |
| 379 | |
| 380 | template_globals = { |
| 381 | "void": DataType.void, |
| 382 | "DataTypeTag": VLLMDataTypeTag, |
| 383 | "VLLMScalarTypeTag": VLLMDataTypeVLLMScalarTypeTag, |
| 384 | "TorchTypeTag": VLLMDataTypeTorchDataTypeTag, |
| 385 | "KernelScheduleTag": VLLMKernelScheduleTag, |
| 386 | "EpilogueScheduleTag": EpilogueScheduleTag, |
| 387 | "TileSchedulerTag": TileSchedulerTag, |
| 388 | "to_cute_constant": to_cute_constant, |
| 389 | "gen_sch_sig": generate_terse_sch_sig, |
| 390 | "gen_type_sig": generate_type_signature, |
| 391 | "unique_schedules": unique_schedules, |
| 392 | "unsigned_type_with_bitwidth": unsigned_type_with_bitwidth, |
| 393 | "gen_type_option_name": generate_type_option_name, |
| 394 | } |
| 395 | |
| 396 | |
| 397 | def create_template(template_str): |
| 398 | template = jinja2.Template(template_str) |
| 399 | template.globals.update(template_globals) |
| 400 | return template |
| 401 | |
| 402 | |
| 403 | mm_dispatch_template = create_template(DISPATCH_TEMPLATE) |
| 404 | mm_impl_template = create_template(IMPL_TEMPLATE) |
| 405 | prepack_dispatch_template = create_template(PREPACK_TEMPLATE) |
| 406 | |
| 407 | |
| 408 | def create_sources(impl_configs: list[ImplConfig], num_impl_files=8): |
| 409 | sources = [] |
| 410 | |
| 411 | sources.append( |
| 412 | ( |
| 413 | "machete_mm_dispatch", |
| 414 | mm_dispatch_template.render(impl_configs=impl_configs), |
| 415 | ) |
| 416 | ) |
| 417 | |
| 418 | prepack_types = [] |
| 419 | for impl_config in impl_configs: |
| 420 | convert_type = ( |
| 421 | impl_config.types.a |
| 422 | if impl_config.types.b_group_scale == DataType.void |
| 423 | else impl_config.types.b_group_scale |
| 424 | ) |
| 425 | prepack_types.append( |
| 426 | PrepackTypeConfig( |
| 427 | a=impl_config.types.a, |
| 428 | b_num_bits=VLLMDataTypeSize[impl_config.types.b], |
| 429 | convert=convert_type, |
| 430 | accumulator=impl_config.types.accumulator, |
| 431 | ) |
| 432 | ) |
| 433 | |
| 434 | def prepacked_type_key(prepack_type: PrepackTypeConfig): |
| 435 | # For now, we can just use the first accumulator type seen since |
| 436 | # the tensor core shapes/layouts don't vary based on accumulator |
| 437 | # type so we can generate less code this way |
| 438 | return (prepack_type.a, prepack_type.b_num_bits, prepack_type.convert) |
| 439 | |
| 440 | unique_prepack_types = [] |
| 441 | prepack_types_seen = set() |
| 442 | for prepack_type in prepack_types: |
| 443 | key = prepacked_type_key(prepack_type) |
| 444 | if key not in prepack_types_seen: |
| 445 | unique_prepack_types.append(prepack_type) |
| 446 | prepack_types_seen.add(key) |
| 447 | |
| 448 | sources.append( |
| 449 | ( |
| 450 | "machete_prepack", |
| 451 | prepack_dispatch_template.render( |
| 452 | types=unique_prepack_types, |
| 453 | ), |
| 454 | ) |
| 455 | ) |
| 456 | |
| 457 | # Split up impls across files |
| 458 | num_impls = reduce(lambda x, y: x + len(y.schedules), impl_configs, 0) |
| 459 | num_impls_per_file = math.ceil(num_impls / num_impl_files) |
| 460 | |
| 461 | files_impls: list[list[ImplConfig]] = [[]] |
| 462 | |
| 463 | curr_num_impls_assigned = 0 |
| 464 | curr_impl_in_file = 0 |
| 465 | curr_impl_configs = deepcopy(list(reversed(impl_configs))) |
| 466 | |
| 467 | while curr_num_impls_assigned < num_impls: |
| 468 | room_left_in_file = num_impls_per_file - curr_impl_in_file |
| 469 | if room_left_in_file == 0: |
| 470 | files_impls.append([]) |
| 471 | room_left_in_file = num_impls_per_file |
| 472 | curr_impl_in_file = 0 |
| 473 | |
| 474 | curr_ic = curr_impl_configs[-1] |
| 475 | if len(curr_ic.schedules) >= room_left_in_file: |
| 476 | # Break apart the current impl config |
| 477 | tmp_ic = deepcopy(curr_ic) |
| 478 | tmp_ic.schedules = curr_ic.schedules[:room_left_in_file] |
| 479 | curr_ic.schedules = curr_ic.schedules[room_left_in_file:] |
| 480 | files_impls[-1].append(tmp_ic) |
| 481 | else: |
| 482 | files_impls[-1].append(curr_ic) |
| 483 | curr_impl_configs.pop() |
| 484 | curr_num_impls_assigned += len(files_impls[-1][-1].schedules) |
| 485 | curr_impl_in_file += len(files_impls[-1][-1].schedules) |
| 486 | |
| 487 | for part, file_impls in enumerate(files_impls): |
| 488 | sources.append( |
| 489 | ( |
| 490 | f"machete_mm_impl_part{part + 1}", |
| 491 | mm_impl_template.render(impl_configs=file_impls), |
| 492 | ) |
| 493 | ) |
| 494 | |
| 495 | return sources |
| 496 | |
| 497 | |
| 498 | def generate(): |
| 499 | # See csrc/quantization/machete/Readme.md, the Codegeneration for more info |
| 500 | # about how this works |
| 501 | SCRIPT_DIR = os.path.dirname(__file__) |
| 502 | |
| 503 | sch_common_params = dict( |
| 504 | kernel_schedule=TmaMI, |
| 505 | epilogue_schedule=TmaCoop, |
| 506 | tile_scheduler=TileSchedulerType.StreamK, |
| 507 | ) |
| 508 | |
| 509 | # Stored as "condition": ((tile_shape_mn), (cluster_shape_mnk)) |
| 510 | default_tile_heuristic_config = { |
| 511 | #### M = 257+ |
| 512 | "M > 256 && K <= 16384 && N <= 4096": ((128, 128), (2, 1, 1)), |
| 513 | "M > 256": ((128, 256), (2, 1, 1)), |
| 514 | #### M = 129-256 |
| 515 | "M > 128 && K <= 4096 && N <= 4096": ((128, 64), (2, 1, 1)), |
| 516 | "M > 128 && K <= 8192 && N <= 8192": ((128, 128), (2, 1, 1)), |
| 517 | "M > 128": ((128, 256), (2, 1, 1)), |
| 518 | #### M = 65-128 |
| 519 | "M > 64 && K <= 4069 && N <= 4069": ((128, 32), (2, 1, 1)), |
| 520 | "M > 64 && K <= 4069 && N <= 8192": ((128, 64), (2, 1, 1)), |
| 521 | "M > 64 && K >= 8192 && N >= 12288": ((256, 128), (2, 1, 1)), |
| 522 | "M > 64": ((128, 128), (2, 1, 1)), |
| 523 | #### M = 33-64 |
| 524 | "M > 32 && K <= 6144 && N <= 6144": ((128, 16), (1, 1, 1)), |
| 525 | "M > 32 && K >= 16384 && N >= 12288": ((256, 64), (2, 1, 1)), |
| 526 | "M > 32": ((128, 64), (2, 1, 1)), |
| 527 | #### M = 17-32 |
| 528 | "M > 16 && K <= 12288 && N <= 8192": ((128, 32), (2, 1, 1)), |
| 529 | "M > 16": ((256, 32), (2, 1, 1)), |
| 530 | #### M = 1-16 |
| 531 | "N >= 26624": ((256, 16), (1, 1, 1)), |
| 532 | None: ((128, 16), (1, 1, 1)), |
| 533 | } |
| 534 | |
| 535 | # For now we use the same heuristic for all types |
| 536 | # Heuristic is currently tuned for H100s |
| 537 | default_heuristic = [ |
| 538 | (cond, ScheduleConfig(*tile_config, **sch_common_params)) # type: ignore |
| 539 | for cond, tile_config in default_tile_heuristic_config.items() |
| 540 | ] |
| 541 | |
| 542 | def get_unique_schedules(heuristic: dict[str, ScheduleConfig]): |
| 543 | # Do not use schedules = list(set(...)) because we need to make sure |
| 544 | # the output list is deterministic; otherwise the generated kernel file |
| 545 | # will be non-deterministic and causes ccache miss. |
| 546 | schedules = [] |
| 547 | for _, schedule_config in heuristic: |
| 548 | if schedule_config not in schedules: |
| 549 | schedules.append(schedule_config) |
| 550 | return schedules |
| 551 | |
| 552 | impl_configs = [] |
| 553 | |
| 554 | GPTQ_kernel_type_configs = list( |
| 555 | TypeConfig( |
| 556 | a=a, |
| 557 | b=b, |
| 558 | b_group_scale=a, |
| 559 | b_group_zeropoint=DataType.void, |
| 560 | b_channel_scale=DataType.void, |
| 561 | a_token_scale=DataType.void, |
| 562 | out=a, |
| 563 | accumulator=DataType.f32, |
| 564 | ) |
| 565 | for b in (VLLMDataType.u4b8, VLLMDataType.u8b128) |
| 566 | for a in (DataType.f16, DataType.bf16) |
| 567 | ) |
| 568 | |
| 569 | impl_configs += [ |
| 570 | ImplConfig(x[0], x[1], x[2]) |
| 571 | for x in zip( |
| 572 | GPTQ_kernel_type_configs, |
| 573 | itertools.repeat(get_unique_schedules(default_heuristic)), |
| 574 | itertools.repeat(default_heuristic), |
| 575 | ) |
| 576 | ] |
| 577 | |
| 578 | AWQ_kernel_type_configs = list( |
| 579 | TypeConfig( |
| 580 | a=a, |
| 581 | b=b, |
| 582 | b_group_scale=a, |
| 583 | b_group_zeropoint=a, |
| 584 | b_channel_scale=DataType.void, |
| 585 | a_token_scale=DataType.void, |
| 586 | out=a, |
| 587 | accumulator=DataType.f32, |
| 588 | ) |
| 589 | for b in (DataType.u4, DataType.u8) |
| 590 | for a in (DataType.f16, DataType.bf16) |
| 591 | ) |
| 592 | |
| 593 | impl_configs += [ |
| 594 | ImplConfig(x[0], x[1], x[2]) |
| 595 | for x in zip( |
| 596 | AWQ_kernel_type_configs, |
| 597 | itertools.repeat(get_unique_schedules(default_heuristic)), |
| 598 | itertools.repeat(default_heuristic), |
| 599 | ) |
| 600 | ] |
| 601 | |
| 602 | # TODO: Support W4A8 when ready |
| 603 | # # Stored as "condition": ((tile_shape_mn), (cluster_shape_mnk)) |
| 604 | # # TODO (LucasWilkinson): Further tuning required |
| 605 | # qqq_tile_heuristic_config = { |
| 606 | # #### M = 257+ |
| 607 | # # ((128, 256), (2, 1, 1)) Broken for QQQ types |
| 608 | # # TODO (LucasWilkinson): Investigate further |
| 609 | # # "M > 256 && K <= 16384 && N <= 4096": ((128, 128), (2, 1, 1)), |
| 610 | # # "M > 256": ((128, 256), (2, 1, 1)), |
| 611 | # "M > 256": ((128, 128), (2, 1, 1)), |
| 612 | # #### M = 129-256 |
| 613 | # "M > 128 && K <= 4096 && N <= 4096": ((128, 64), (2, 1, 1)), |
| 614 | # "M > 128 && K <= 8192 && N <= 8192": ((128, 128), (2, 1, 1)), |
| 615 | # # ((128, 256), (2, 1, 1)) Broken for QQQ types |
| 616 | # # TODO (LucasWilkinson): Investigate further |
| 617 | # # "M > 128": ((128, 256), (2, 1, 1)), |
| 618 | # "M > 128": ((128, 128), (2, 1, 1)), |
| 619 | # #### M = 65-128 |
| 620 | # "M > 64 && K <= 4069 && N <= 4069": ((128, 32), (2, 1, 1)), |
| 621 | # "M > 64 && K <= 4069 && N <= 8192": ((128, 64), (2, 1, 1)), |
| 622 | # "M > 64 && K >= 8192 && N >= 12288": ((256, 128), (2, 1, 1)), |
| 623 | # "M > 64": ((128, 128), (2, 1, 1)), |
| 624 | # #### M = 33-64 |
| 625 | # "M > 32 && K <= 6144 && N <= 6144": ((128, 16), (1, 1, 1)), |
| 626 | # # Broken for QQQ types |
| 627 | # # TODO (LucasWilkinson): Investigate further |
| 628 | # #"M > 32 && K >= 16384 && N >= 12288": ((256, 64), (2, 1, 1)), |
| 629 | # "M > 32": ((128, 64), (2, 1, 1)), |
| 630 | # #### M = 17-32 |
| 631 | # "M > 16 && K <= 12288 && N <= 8192": ((128, 32), (2, 1, 1)), |
| 632 | # "M > 16": ((256, 32), (2, 1, 1)), |
| 633 | # #### M = 1-16 |
| 634 | # "N >= 26624": ((256, 16), (1, 1, 1)), |
| 635 | # None: ((128, 16), (1, 1, 1)), |
| 636 | # } |
| 637 | |
| 638 | # # For now we use the same heuristic for all types |
| 639 | # # Heuristic is currently tuned for H100s |
| 640 | # qqq_heuristic = [ |
| 641 | # (cond, ScheduleConfig(*tile_config, |
| 642 | # **sch_common_params)) # type: ignore |
| 643 | # for cond, tile_config in qqq_tile_heuristic_config.items() |
| 644 | # ] |
| 645 | |
| 646 | # QQQ_kernel_types = [ |
| 647 | # *(TypeConfig( |
| 648 | # a=DataType.s8, |
| 649 | # b=VLLMDataType.u4b8, |
| 650 | # b_group_scale=b_group_scale, |
| 651 | # b_group_zeropoint=DataType.void, |
| 652 | # b_channel_scale=DataType.f32, |
| 653 | # a_token_scale=DataType.f32, |
| 654 | # out=DataType.f16, |
| 655 | # accumulator=DataType.s32, |
| 656 | # ) for b_group_scale in (DataType.f16, DataType.void)), |
| 657 | # *(TypeConfig( |
| 658 | # a=DataType.e4m3, |
| 659 | # b=VLLMDataType.u4b8, |
| 660 | # b_group_scale=b_group_scale, |
| 661 | # b_group_zeropoint=DataType.void, |
| 662 | # b_channel_scale=DataType.f32, |
| 663 | # a_token_scale=DataType.f32, |
| 664 | # out=DataType.f16, |
| 665 | # accumulator=DataType.f32, |
| 666 | # ) for b_group_scale in (DataType.f16, DataType.void)), |
| 667 | # ] |
| 668 | |
| 669 | # impl_configs += [ |
| 670 | # ImplConfig(x[0], x[1], x[2]) |
| 671 | # for x in zip(QQQ_kernel_types, |
| 672 | # itertools.repeat(get_unique_schedules(qqq_heuristic)), |
| 673 | # itertools.repeat(qqq_heuristic)) |
| 674 | # ] |
| 675 | |
| 676 | output_dir = os.path.join(SCRIPT_DIR, "generated") |
| 677 | |
| 678 | # Delete the "generated" directory if it exists |
| 679 | if os.path.exists(output_dir): |
| 680 | shutil.rmtree(output_dir) |
| 681 | |
| 682 | # Create the "generated" directory |
| 683 | os.makedirs(output_dir) |
| 684 | |
| 685 | # Render each group of configurations into separate files |
| 686 | for filename, code in create_sources(impl_configs): |
| 687 | filepath = os.path.join(output_dir, f"{filename}.cu") |
| 688 | with open(filepath, "w") as output_file: |
| 689 | output_file.write(code) |
| 690 | print(f"Rendered template to {filepath}") |
| 691 | |
| 692 | |
| 693 | if __name__ == "__main__": |
| 694 | generate() |