e489ad7a21
- **Add SPDX license headers to python source files**
- **Check for SPDX headers using pre-commit**
commit 9d7ef44c3cfb72ca4c32e1c677d99259d10d4745
Author: Russell Bryant <rbryant@redhat.com>
Date: Fri Jan 31 14:18:24 2025 -0500
Add SPDX license headers to python source files
This commit adds SPDX license headers to python source files as
recommended to
the project by the Linux Foundation. These headers provide a concise way
that is
both human and machine readable for communicating license information
for each
source file. It helps avoid any ambiguity about the license of the code
and can
also be easily used by tools to help manage license compliance.
The Linux Foundation runs license scans against the codebase to help
ensure
we are in compliance with the licenses of the code we use, including
dependencies. Having these headers in place helps that tool do its job.
More information can be found on the SPDX site:
- https://spdx.dev/learn/handling-license-info/
Signed-off-by: Russell Bryant <rbryant@redhat.com>
commit 5a1cf1cb3b80759131c73f6a9dddebccac039dea
Author: Russell Bryant <rbryant@redhat.com>
Date: Fri Jan 31 14:36:32 2025 -0500
Check for SPDX headers using pre-commit
Signed-off-by: Russell Bryant <rbryant@redhat.com>
---------
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Machete (Mixed Precision Cutlass-Based GEMM)
Machete is a spiritual successor to the Marlin kernel but optimized for Hopper architectures and based on Cutlass. Being based on Cutlass, new type pairs and epilogues are easier to add compared to Marlin.
Overview
Machete effectively performs
scale_type = w_s.dtype
compute_type = a.dtype
out = (w_q.to(scale_type) * w_s - w_z.to(scale_type)) @ a
Where w_q is a quantized weight matrix, w_s is the quantization scales, and
w_z is the quantization zeropoints.
NOTE:
w_zis added after the scales so we can use FMA operations, but this means they must have the scales pre-applied if the supplied zeropoints assume that they will be subtracted before the scales are applied.
API
The main optimization within Machete is prepacking the weight matrix to more closely match the tensor core layouts, allowing for wider shared memory loads when loading the weight matrix. This means that the weight matrix must be prepacked before calling machete_gemm. The flow looks something like:
from vllm import _custom_ops as ops
...
W_q_packed = ops.machete_prepack_B(w_q, wtype)
output = ops.machete_gemm(
a,
b_q=W_q_packed,
b_type=wtype,
b_scales=w_s,
b_group_size=group_size
)
Code Generation
Since Machete is based on Cutlass, we can generate multiple type pairs and different tile shapes using the same kernel template. We generate multiple instantiations of this template using generate.py.
New type pairs (TypeConfigs) can be appended to impl_configs (in generate()), and these will get automatically generated (assuming they can be supported without issues). For each TypeConfig, you must also provide an ImplConfig, which bundles a TypeConfig with a list of ScheduleConfigs, Specializations, and a default heuristic. The ScheduleConfigs (which contain info on tile shapes, tile scheduler, etc.) can perform differently for different problem shapes, and there is almost never one ScheduleConfig that works well for all problem shapes, so it is generally beneficial to generate different ScheduleConfigs for different potential problem shapes. This is where the heuristic comes in. For each TypeConfig, a default heuristic should be provided. This maps different problem shapes to different ScheduleConfigs and is used when the user does not provide the schedule parameter to machete_gemm. The Specializations define what feature combinations to generate, i.e., with_zeropoints, with_scales, etc. We can reduce compile times and the final binary size by limiting the set of feature combinations we generate.