# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import itertools import sys from pathlib import Path import jinja2 FILE_HEAD = """ // auto generated by generate_kernels.py // clang-format off #include "kernel.h" #include "marlin_template.h" namespace MARLIN_NAMESPACE_NAME { """.strip() FILE_TAIL = "}\n" TEMPLATE = ("template __global__ void Marlin<" "{{scalar_t}}, " "{{w_type_id}}, " "{{s_type_id}}, " "{{threads}}, " "{{thread_m_blocks}}, " "{{thread_n_blocks}}, " "{{thread_k_blocks}}, " "{{'true' if m_block_size_8 else 'false'}}, " "{{stages}}, " "{{group_blocks}}, " "{{'true' if is_zp_float else 'false'}}>" "( MARLIN_KERNEL_PARAMS );") # int8 with zero point case (vllm::kU8) is also supported, # we don't add it to reduce wheel size. SCALAR_TYPES = [ "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn", "vllm::kFE2M1f" ] THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] # group_blocks: # = 0 : act order case # = -1 : channelwise quantization # > 0 : group_size=16*group_blocks GROUP_BLOCKS = [0, 1, -1, 2, 4, 8] DTYPES = ["fp16", "bf16"] def _is_4bit_weight(scalar_type: str) -> bool: return scalar_type in { "vllm::kU4", "vllm::kU4B8", "vllm::kFE2M1f", } def remove_old_kernels() -> None: root = Path(__file__).parent for path in root.glob("kernel_*.cu"): path.unlink(missing_ok=True) def _kernel_output_path(root: Path, scalar_type: str, dtype: str) -> Path: scalar_suffix = scalar_type.split("::", 1)[1].lower() if "::" in scalar_type else scalar_type.lower() return root / f"kernel_{dtype}_{scalar_suffix}.cu" def _render_kernel_file_text(scalar_type: str, dtype: str, templates: list[str]) -> str: lines = [FILE_HEAD, "", f"// Instantiations for dtype={dtype}, weight={scalar_type}", ""] lines.append("\n".join(templates)) lines.append("") lines.append(FILE_TAIL) return "\n".join(lines) def _write_kernel_file(root: Path, scalar_type: str, dtype: str, templates: list[str]) -> Path: output_path = _kernel_output_path(root, scalar_type, dtype) output_path.write_text(_render_kernel_file_text(scalar_type, dtype, templates), encoding="utf-8") return output_path def build_expected_kernels(root: Path | None = None) -> dict[Path, str]: root = root or Path(__file__).parent expected: dict[Path, str] = {} for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES): templates = render_templates_for_combo(scalar_type, dtype) if not templates: continue output_path = _kernel_output_path(root, scalar_type, dtype) expected[output_path] = _render_kernel_file_text(scalar_type, dtype, templates) if not expected: raise RuntimeError("No marlin kernels were generated; check template configuration.") return expected def generated_kernels_are_current(root: Path | None = None) -> bool: root = root or Path(__file__).parent expected = build_expected_kernels(root) expected_names = {path.name for path in expected} existing_names = {path.name for path in root.glob("kernel_*.cu")} if existing_names != expected_names: return False for output_path, expected_text in expected.items(): try: current_text = output_path.read_text(encoding="utf-8") except OSError: return False if current_text != expected_text: return False return True def render_templates_for_combo(scalar_type: str, dtype: str) -> list[str]: results: list[str] = [] stage_values = ["pipe_stages"] if dtype == "fp16": # Turing uses a shorter pipeline depth than Ampere+. stage_values.insert(0, 2) for group_blocks, m_blocks, thread_configs in itertools.product( GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS): # act order case only support gptq-int4 and gptq-int8 if group_blocks == 0 and scalar_type not in [ "vllm::kU4B8", "vllm::kU8B128" ]: continue if thread_configs[2] == 256: # for small batch (m_blocks == 1), we only need (128, 128, 256) # for large batch (m_blocks > 1), we only need (64, 256, 256) if m_blocks <= 1 and thread_configs[0] != 128: continue if m_blocks > 1 and thread_configs[0] != 64: continue # FP8 weights support channelwise/group128 in both fp16 and bf16, and # group32 microscaling (MXFP8) in bf16. if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 2, 8]: continue # nvfp4 only supports group_size == 16 # mxfp4 only supports group_size == 32 if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]: continue # other quantization methods don't support group_size = 16 if scalar_type != "vllm::kFE2M1f" and group_blocks == 1: continue k_blocks = thread_configs[0] // 16 n_blocks = thread_configs[1] // 16 threads = thread_configs[2] c_dtype = "half" if dtype == "fp16" else "nv_bfloat16" is_zp_float_list = [False] if dtype == "fp16" and scalar_type == "vllm::kU4" and \ group_blocks == 4: # HQQ (is_zp_float = true) only supports # 4bit quantization and fp16 is_zp_float_list.append(True) if scalar_type == "vllm::kFE2M1f" and group_blocks == 1: s_type = "vllm::kFE4M3fn" elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2: s_type = "vllm::kFE8M0fnu" if dtype == "fp16": # we cannot safely dequantize e8m0 to fp16, so skip this continue elif scalar_type == "vllm::kFE4M3fn" and group_blocks == 2: s_type = "vllm::kFE8M0fnu" if dtype == "fp16": # MXFP8 is only supported with bf16 compute. continue elif dtype == "fp16": s_type = "vllm::kFloat16" elif dtype == "bf16": s_type = "vllm::kBFloat16" for is_zp_float in is_zp_float_list: for stage_value in stage_values: if ( stage_value == 2 and _is_4bit_weight(scalar_type) and max(m_blocks, 1) * 2 > k_blocks ): # Our dense Turing kernels need enough B-stage capacity to # cover the output tile. For 4-bit weights this rules out # the larger M tiles when thread_k_blocks == 4. continue template_str = jinja2.Template(TEMPLATE).render( scalar_t=c_dtype, w_type_id=scalar_type + ".id()", s_type_id=s_type + ".id()", threads=threads, thread_m_blocks=max(m_blocks, 1), thread_n_blocks=n_blocks, thread_k_blocks=k_blocks, m_block_size_8=m_blocks == 0.5, stages=stage_value, group_blocks=group_blocks, is_zp_float=is_zp_float, ) results.append(template_str) return results def generate_new_kernels(root: Path | None = None) -> None: root = root or Path(__file__).parent expected = build_expected_kernels(root) for path in root.glob("kernel_*.cu"): if path not in expected: path.unlink(missing_ok=True) for output_path, rendered in expected.items(): output_path.write_text(rendered, encoding="utf-8") if __name__ == "__main__": if "--check" in sys.argv[1:]: raise SystemExit(0 if generated_kernels_are_current() else 1) generate_new_kernels()