wehub-resource-sync
346 lines
12 KiB
Python
346 lines
12 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
"""FlashInfer JIT compilation module for CUDA backend"""
|
|
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import tvm
|
|
from tvm.target import Target
|
|
|
|
|
|
def _rename_exported_func_names(source_paths: list[Path], prefix: str):
|
|
"""Rename the ffi-exported function names in the source files to the given prefix."""
|
|
pattern = re.compile(r"^(\s*TVM_FFI_DLL_EXPORT_TYPED_FUNC\()([A-Za-z0-9_]+)(,.*)$")
|
|
for source_path in source_paths:
|
|
if not source_path.name.endswith("_binding.cu"):
|
|
continue
|
|
|
|
original_text = source_path.read_text(encoding="utf-8")
|
|
lines = original_text.splitlines(keepends=True)
|
|
updated = False
|
|
for idx, line in enumerate(lines):
|
|
line_body = line.rstrip("\r\n")
|
|
line_ending = line[len(line_body) :]
|
|
match = pattern.match(line_body)
|
|
if not match:
|
|
continue
|
|
new_body = f"{match.group(1)}{prefix}_{match.group(2)}{match.group(3)}"
|
|
lines[idx] = new_body + line_ending
|
|
updated = True
|
|
|
|
if updated:
|
|
source_path.write_text("".join(lines), encoding="utf-8")
|
|
|
|
|
|
def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm.runtime.Module]:
|
|
return [
|
|
tvm.runtime.load_static_library(str(obj_path.absolute()), func_names=[])
|
|
for obj_path in object_files
|
|
]
|
|
|
|
|
|
def gen_flashinfer_prefill_module(
|
|
dtype_q: str,
|
|
dtype_kv: str,
|
|
dtype_o: str,
|
|
qk_head_dim: int,
|
|
v_head_dim: int,
|
|
enable_inline_rope: bool,
|
|
return_static_libs: bool = False,
|
|
) -> list[tvm.runtime.Module]:
|
|
"""Generate a FlashInfer module for prefill.
|
|
|
|
Parameters
|
|
----------
|
|
dtype_q : str
|
|
The data type of the query tensor.
|
|
dtype_kv : str
|
|
The data type of the key/value tensors.
|
|
dtype_o : str
|
|
The data type of the output tensor.
|
|
qk_head_dim : int
|
|
The head dimension of the query and key tensors.
|
|
v_head_dim : int
|
|
The head dimension of the value tensor.
|
|
enable_inline_rope : bool
|
|
Whether to enable inline rotary positional embedding.
|
|
return_static_libs : bool
|
|
Whether to return static library modules instead of compiled modules.
|
|
When it is False, it returns the loaded shared library that links all the object files.
|
|
When it is True, it returns the static libraries of each compiled object files.
|
|
|
|
Returns
|
|
-------
|
|
A list of compiled static library modules for FlashInfer prefill kernels.
|
|
"""
|
|
try:
|
|
from flashinfer.jit import ( # pylint: disable=import-outside-toplevel
|
|
gen_customize_batch_prefill_module,
|
|
)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"FlashInfer is not installed. Please follow instructions "
|
|
"in https://docs.flashinfer.ai to install FlashInfer."
|
|
)
|
|
try:
|
|
import torch # pylint: disable=import-outside-toplevel
|
|
except ImportError:
|
|
raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
|
|
|
|
if enable_inline_rope and qk_head_dim != v_head_dim:
|
|
raise ValueError("Inline rope mode is not supported when qk_head_dim == v_head_dim")
|
|
|
|
torch_dtype_q = getattr(torch, dtype_q)
|
|
torch_dtype_kv = getattr(torch, dtype_kv)
|
|
torch_dtype_o = getattr(torch, dtype_o)
|
|
# Todo(tvm-team): decide which backend ("fa2/fa3") to use
|
|
backend = "fa2"
|
|
variant_name = (
|
|
"DefaultAttention<false, false, false, false>"
|
|
if backend == "fa2"
|
|
else "DefaultAttention<false>"
|
|
)
|
|
variant_decl = (
|
|
"#include <flashinfer/attention/variants.cuh>"
|
|
if backend == "fa2"
|
|
else "#include <flashinfer/attention/hopper/variants.cuh>"
|
|
)
|
|
jit_spec = gen_customize_batch_prefill_module(
|
|
backend=backend,
|
|
uri=f"batch_prefill_tvm_dtype_q_{dtype_q}_"
|
|
+ f"dtype_kv_{dtype_kv}_"
|
|
+ f"dtype_o_{dtype_o}_"
|
|
+ f"qk_head_dim_{qk_head_dim}_"
|
|
+ f"v_head_dim_{v_head_dim}_"
|
|
+ f"enable_inline_rope_{enable_inline_rope}",
|
|
dtype_q=torch_dtype_q,
|
|
dtype_kv=torch_dtype_kv,
|
|
dtype_o=torch_dtype_o,
|
|
idtype=torch.int32,
|
|
head_dim_qk=qk_head_dim,
|
|
head_dim_vo=v_head_dim,
|
|
pos_encoding_mode=int(enable_inline_rope),
|
|
additional_tensor_names=[],
|
|
additional_tensor_dtypes=[],
|
|
additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
|
|
additional_scalar_dtypes=["double", "double", "double"],
|
|
variant_name=variant_name,
|
|
variant_decl=variant_decl,
|
|
)
|
|
_rename_exported_func_names(jit_spec.sources, "batch_prefill")
|
|
if return_static_libs:
|
|
jit_spec.build(verbose=False)
|
|
return _load_flashinfer_modules(jit_spec.get_object_paths())
|
|
return [jit_spec.build_and_load()]
|
|
|
|
|
|
def gen_flashinfer_decode_module(
|
|
dtype_q: str,
|
|
dtype_kv: str,
|
|
dtype_o: str,
|
|
qk_head_dim: int,
|
|
v_head_dim: int,
|
|
enable_inline_rope: bool,
|
|
return_static_libs: bool = False,
|
|
) -> list[tvm.runtime.Module]:
|
|
"""Generate a FlashInfer module for decode.
|
|
|
|
Parameters
|
|
----------
|
|
dtype_q : str
|
|
The data type of the query tensor.
|
|
dtype_kv : str
|
|
The data type of the key/value tensors.
|
|
dtype_o : str
|
|
The data type of the output tensor.
|
|
qk_head_dim : int
|
|
The head dimension of the query and key tensors.
|
|
v_head_dim : int
|
|
The head dimension of the value tensor.
|
|
enable_inline_rope : bool
|
|
Whether to enable inline rotary positional embedding.
|
|
return_static_libs : bool
|
|
Whether to return static library modules instead of compiled modules.
|
|
When it is False, it returns the loaded shared library that links all the object files.
|
|
When it is True, it returns the static libraries of each compiled object files.
|
|
|
|
Returns
|
|
-------
|
|
A list of compiled static library modules for FlashInfer decode kernels.
|
|
"""
|
|
try:
|
|
from flashinfer.jit import ( # pylint: disable=import-outside-toplevel
|
|
gen_customize_batch_decode_module,
|
|
)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"FlashInfer is not installed. Please follow instructions "
|
|
"in https://docs.flashinfer.ai to install FlashInfer."
|
|
)
|
|
try:
|
|
import torch # pylint: disable=import-outside-toplevel
|
|
except ImportError:
|
|
raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
|
|
|
|
torch_dtype_q = getattr(torch, dtype_q)
|
|
torch_dtype_kv = getattr(torch, dtype_kv)
|
|
torch_dtype_o = getattr(torch, dtype_o)
|
|
jit_spec = gen_customize_batch_decode_module(
|
|
uri=f"batch_decode_tvm_dtype_q_{dtype_q}_"
|
|
+ f"dtype_kv_{dtype_kv}_"
|
|
+ f"dtype_o_{dtype_o}_"
|
|
+ f"qk_head_dim_{qk_head_dim}_"
|
|
+ f"v_head_dim_{v_head_dim}_"
|
|
+ f"enable_inline_rope_{enable_inline_rope}",
|
|
dtype_q=torch_dtype_q,
|
|
dtype_kv=torch_dtype_kv,
|
|
dtype_o=torch_dtype_o,
|
|
idtype=torch.int32,
|
|
head_dim_qk=qk_head_dim,
|
|
head_dim_vo=v_head_dim,
|
|
pos_encoding_mode=int(enable_inline_rope),
|
|
additional_tensor_names=[],
|
|
additional_tensor_dtypes=[],
|
|
additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
|
|
additional_scalar_dtypes=["double", "double", "double"],
|
|
variant_name="DefaultAttention<false, false, false, false>",
|
|
variant_decl="#include <flashinfer/attention/variants.cuh>",
|
|
)
|
|
_rename_exported_func_names(jit_spec.sources, "batch_decode")
|
|
if return_static_libs:
|
|
jit_spec.build(verbose=False)
|
|
return _load_flashinfer_modules(jit_spec.get_object_paths())
|
|
return [jit_spec.build_and_load()]
|
|
|
|
|
|
def gen_flashinfer_mla_module(
|
|
dtype_q: str,
|
|
dtype_kv: str,
|
|
dtype_o: str,
|
|
head_dim_ckv: int,
|
|
head_dim_kpe: int,
|
|
return_static_libs: bool = False,
|
|
) -> list[tvm.runtime.Module]:
|
|
"""Generate a FlashInfer module for MLA.
|
|
|
|
Parameters
|
|
----------
|
|
dtype_q : str
|
|
The data type of the query tensor.
|
|
dtype_kv : str
|
|
The data type of the key/value tensors.
|
|
dtype_o : str
|
|
The data type of the output tensor.
|
|
head_dim_ckv : int
|
|
The head dimension of the compressed key/value tensors.
|
|
head_dim_kpe : int
|
|
The head dimension of the query/key positional embedding.
|
|
target : Target
|
|
The target device to compile for.
|
|
num_threads : int
|
|
The number of threads to use for compilation.
|
|
return_static_libs : bool
|
|
Whether to return static library modules instead of compiled modules.
|
|
When it is False, it returns the loaded shared library that links all the object files.
|
|
When it is True, it returns the static libraries of each compiled object files.
|
|
|
|
Returns
|
|
-------
|
|
A list of compiled static library modules for FlashInfer MLA kernels.
|
|
"""
|
|
try:
|
|
from flashinfer.jit import ( # pylint: disable=import-outside-toplevel
|
|
gen_batch_mla_module,
|
|
)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"FlashInfer is not installed. Please follow instructions "
|
|
"in https://docs.flashinfer.ai to install FlashInfer."
|
|
)
|
|
try:
|
|
import torch # pylint: disable=import-outside-toplevel
|
|
except ImportError:
|
|
raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
|
|
|
|
torch_dtype_q = getattr(torch, dtype_q)
|
|
torch_dtype_kv = getattr(torch, dtype_kv)
|
|
torch_dtype_o = getattr(torch, dtype_o)
|
|
jit_spec = gen_batch_mla_module(
|
|
backend="fa2",
|
|
dtype_q=torch_dtype_q,
|
|
dtype_kv=torch_dtype_kv,
|
|
dtype_o=torch_dtype_o,
|
|
dtype_idx=torch.int32,
|
|
head_dim_ckv=head_dim_ckv,
|
|
head_dim_kpe=head_dim_kpe,
|
|
use_profiler=False,
|
|
)
|
|
_rename_exported_func_names(jit_spec.sources, "batch_mla")
|
|
if return_static_libs:
|
|
jit_spec.build(verbose=False)
|
|
return _load_flashinfer_modules(jit_spec.get_object_paths())
|
|
return [jit_spec.build_and_load()]
|
|
|
|
|
|
def gen_grouped_gemm_module(
|
|
target: Target, return_static_libs: bool = False
|
|
) -> list[tvm.runtime.Module]:
|
|
"""Generate a FlashInfer module for FP8 grouped GEMM.
|
|
|
|
Parameters
|
|
----------
|
|
target : Target
|
|
The target device to compile for.
|
|
return_static_libs : bool
|
|
Whether to return static library modules instead of compiled modules.
|
|
When it is False, it returns the loaded shared library that links all the object files.
|
|
When it is True, it returns the static libraries of each compiled object files.
|
|
|
|
Returns
|
|
-------
|
|
List[tvm.runtime.Module]
|
|
A list of compiled static library modules for FlashInfer FP8 grouped GEMM kernels.
|
|
|
|
Note
|
|
_____
|
|
when apply grouped gemm on A: (total_m, k), B: (batch_size, n, k), m_indptr: (batch_size, )
|
|
requires all m in m_indptr to be multiple of 4
|
|
"""
|
|
# NOTE: This function is still under development,
|
|
# and we currently only support SM100 grouped gemm
|
|
try:
|
|
from flashinfer.gemm import ( # pylint: disable=import-outside-toplevel
|
|
gen_gemm_sm100_module,
|
|
)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"FlashInfer is not installed. Please follow instructions "
|
|
"in https://docs.flashinfer.ai to install FlashInfer."
|
|
)
|
|
|
|
compute_version = "".join(tvm.support.nvcc.get_target_compute_version(target).split("."))
|
|
if compute_version == "100":
|
|
jit_spec = gen_gemm_sm100_module()
|
|
else:
|
|
raise ValueError(f"Unsupported compute version: {compute_version}")
|
|
if return_static_libs:
|
|
jit_spec.build(verbose=False)
|
|
return _load_flashinfer_modules(jit_spec.get_object_paths())
|
|
return [jit_spec.build_and_load()]
|