wehub-resource-sync
244 lines
9.2 KiB
Python
244 lines
9.2 KiB
Python
"""A pass that rewrites KV cache creation functions in IRModule."""
|
|
|
|
import json
|
|
from typing import Any, Dict, List # noqa: UP035
|
|
|
|
import tvm
|
|
from tvm import IRModule, relax
|
|
from tvm.relax.frontend.nn.llm import kv_cache
|
|
from tvm.relax.frontend.nn.llm.kv_cache import RopeMode
|
|
|
|
from mlc_llm.support import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def extract_creation_args(func: relax.Function) -> Dict[str, Any]: # noqa: UP006
|
|
"""Extract the KV cache creation args from the given generic creation func."""
|
|
assert isinstance(func.body, relax.SeqExpr)
|
|
assert len(func.body.blocks) == 1
|
|
assert isinstance(func.body.blocks[0], relax.DataflowBlock)
|
|
assert isinstance(func.body.blocks[0].bindings[0], relax.VarBinding)
|
|
assert isinstance(func.body.blocks[0].bindings[0].value, relax.Call)
|
|
assert func.body.blocks[0].bindings[0].value.op == tvm.ir.Op.get("relax.call_pure_packed")
|
|
call_args = func.body.blocks[0].bindings[0].value.args
|
|
assert isinstance(call_args[0], relax.ExternFunc)
|
|
assert call_args[0].global_symbol == "mlc.create_paged_kv_cache_generic"
|
|
args = call_args[1:]
|
|
assert len(args) == 18
|
|
assert isinstance(args[0], (relax.StringImm, relax.Tuple))
|
|
# Check if attn_kind is a single value or a list with length of hidden layers
|
|
if isinstance(args[0], relax.StringImm):
|
|
assert args[0].value in ["mha", "mla"]
|
|
attn_kind = args[0].value
|
|
else:
|
|
assert len(args[0].fields) == args[3].value
|
|
for i, attention_type in enumerate(args[0].fields):
|
|
assert isinstance(attention_type, relax.StringImm)
|
|
assert attention_type.value in ["mha", "mla", "mha_sliding"]
|
|
attn_kind = [args[0].fields[i].value for i in range(len(args[0]))]
|
|
assert isinstance(args[1], relax.ShapeExpr)
|
|
assert len(args[1].values) == 5
|
|
assert isinstance(args[2], relax.ShapeExpr)
|
|
for i in range(3, 18):
|
|
if i in [13, 14, 17]:
|
|
continue
|
|
# PrimValue wrappers were phased out of Relax: scalar args are now bare
|
|
# tirx PrimExprs (IntImm/FloatImm) directly.
|
|
assert isinstance(args[i], (tvm.tirx.IntImm, tvm.tirx.FloatImm)), (
|
|
f"args[{i}] is {type(args[i])}"
|
|
)
|
|
assert isinstance(args[13], relax.StringImm)
|
|
assert isinstance(args[16], (relax.Constant, tvm.tirx.IntImm, tvm.tirx.FloatImm))
|
|
assert isinstance(args[17], relax.DataTypeImm)
|
|
|
|
return {
|
|
"attn_kind": attn_kind,
|
|
"max_batch_size": args[1].values[0],
|
|
"max_total_seq_len": args[1].values[1],
|
|
"prefill_chunk_size": args[1].values[2],
|
|
"page_size": args[1].values[3],
|
|
"support_sliding_window": args[1].values[4],
|
|
"layer_partition": args[2],
|
|
"num_hidden_layers": args[3].value,
|
|
"num_attention_heads": args[4].value,
|
|
"num_key_value_heads": args[5].value,
|
|
"qk_head_dim": args[6].value,
|
|
"v_head_dim": args[7].value,
|
|
"mla_original_qk_head_dim": args[8].value,
|
|
"mla_original_v_head_dim": args[9].value,
|
|
"rope_mode": args[10].value,
|
|
"rope_scale": args[11].value,
|
|
"rope_theta": args[12].value,
|
|
"rope_scaling": json.loads(args[13].value),
|
|
"rope_ext_factors": args[14],
|
|
"rotary_dim": args[15].value,
|
|
"enable_disaggregation": bool(args[16].value),
|
|
"dtype": args[17].value,
|
|
}
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="DispatchKVCacheCreation")
|
|
class DispatchKVCacheCreation:
|
|
"""Rewrite KV cache creation functions to IRModule."""
|
|
|
|
def __init__(
|
|
self,
|
|
target: tvm.target.Target,
|
|
flashinfer: bool,
|
|
metadata: Dict[str, Any], # noqa: UP006
|
|
) -> None:
|
|
"""Initializer.
|
|
|
|
Parameters
|
|
----------
|
|
target : tvm.target.Target
|
|
The target of the model compilation.
|
|
|
|
flashinfer : bool
|
|
A boolean indicating if flashinfer is enabled.
|
|
|
|
metadata : Dict[str, Any]
|
|
The model's metadata for KV cache creation.
|
|
Note that the metadata will be updated in this pass -- the
|
|
KV cache metadata will be attached.
|
|
"""
|
|
self.target = target
|
|
self.flashinfer = flashinfer
|
|
self.metadata = metadata
|
|
|
|
def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
|
|
"""Entrypoint"""
|
|
func_dict = {}
|
|
creation_func = None
|
|
for g_var, func in mod.functions_items():
|
|
# Try to find the `create_paged_kv_cache` func.
|
|
if g_var.name_hint == "create_paged_kv_cache":
|
|
creation_func = func
|
|
else:
|
|
func_dict[g_var] = func
|
|
|
|
if creation_func is None:
|
|
return mod
|
|
|
|
new_mod = IRModule(func_dict)
|
|
if mod.attrs is not None:
|
|
new_mod = new_mod.with_attrs(mod.attrs)
|
|
|
|
kwargs = extract_creation_args(creation_func)
|
|
self.attach_kv_cache_metadata(kwargs)
|
|
|
|
bb = relax.BlockBuilder(new_mod)
|
|
extern_mods = []
|
|
extern_mods += self.create_tir_paged_kv_cache(bb, kwargs)
|
|
extern_mods += self.create_flashinfer_paged_kv_cache(bb, kwargs)
|
|
|
|
mod = bb.finalize()
|
|
mod_attrs = dict(mod.attrs) if mod.attrs else {}
|
|
mod = mod.with_attr("external_mods", mod_attrs.get("external_mods", []) + extern_mods)
|
|
return mod
|
|
|
|
def attach_kv_cache_metadata(self, kwargs: Dict[str, Any]): # noqa: UP006
|
|
"""Attach the KV cache metadata to model metadata."""
|
|
self.metadata["kv_cache"] = {
|
|
"num_hidden_layers": kwargs["num_hidden_layers"],
|
|
"num_attention_heads": kwargs["num_attention_heads"],
|
|
"num_key_value_heads": kwargs["num_key_value_heads"],
|
|
"head_dim": kwargs["qk_head_dim"],
|
|
}
|
|
|
|
def create_tir_paged_kv_cache(
|
|
self,
|
|
bb: relax.BlockBuilder,
|
|
kwargs: Dict[str, Any], # noqa: UP006
|
|
) -> List[tvm.runtime.Module]: # noqa: UP006
|
|
"""Create the TIR-based PagedKVCache"""
|
|
max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
|
|
max_total_seq_len = relax.Var(
|
|
"max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
|
|
)
|
|
prefill_chunk_size = relax.Var(
|
|
"prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
|
|
)
|
|
page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
|
|
support_sliding_window = relax.Var(
|
|
"support_sliding_window_",
|
|
relax.ShapeType([kwargs["support_sliding_window"]]),
|
|
)
|
|
|
|
# Ensure 'enable_disaggregation' is optional
|
|
enable_disaggregation = kwargs.pop("enable_disaggregation", False)
|
|
kwargs["enable_disaggregation"] = enable_disaggregation
|
|
|
|
with bb.function(
|
|
name="create_tir_paged_kv_cache",
|
|
params=[
|
|
max_batch_size,
|
|
max_total_seq_len,
|
|
prefill_chunk_size,
|
|
page_size,
|
|
support_sliding_window,
|
|
],
|
|
):
|
|
cache = kv_cache.TIRPagedKVCache(target=self.target, **kwargs)
|
|
bb.emit_func_output(cache._expr)
|
|
|
|
return cache.extern_mods
|
|
|
|
def create_flashinfer_paged_kv_cache(
|
|
self,
|
|
bb: relax.BlockBuilder,
|
|
kwargs: Dict[str, Any], # noqa: UP006
|
|
) -> List[tvm.runtime.Module]: # noqa: UP006
|
|
"""Create the FlashInfer-based PagedKVCache"""
|
|
# Filter the cases which FlashInfer does not support.
|
|
if (
|
|
not self.flashinfer
|
|
or self.target.kind.name != "cuda"
|
|
or str(kwargs["dtype"]) not in ["float16", "bfloat16"]
|
|
or (
|
|
kwargs["rope_mode"] == RopeMode.INLINE
|
|
and (
|
|
kwargs["rotary_dim"] != kwargs["qk_head_dim"]
|
|
or kwargs["qk_head_dim"] != kwargs["v_head_dim"]
|
|
)
|
|
)
|
|
):
|
|
return []
|
|
|
|
max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
|
|
max_total_seq_len = relax.Var(
|
|
"max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
|
|
)
|
|
prefill_chunk_size = relax.Var(
|
|
"prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
|
|
)
|
|
page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
|
|
support_sliding_window = relax.Var(
|
|
"support_sliding_window_",
|
|
relax.ShapeType([kwargs["support_sliding_window"]]),
|
|
)
|
|
|
|
try:
|
|
with bb.function(
|
|
name="create_flashinfer_paged_kv_cache",
|
|
params=[
|
|
max_batch_size,
|
|
max_total_seq_len,
|
|
prefill_chunk_size,
|
|
page_size,
|
|
support_sliding_window,
|
|
],
|
|
):
|
|
cache = kv_cache.FlashInferPagedKVCache(target=self.target, **kwargs)
|
|
bb.emit_func_output(cache._expr)
|
|
except Exception as e:
|
|
logger.info(
|
|
"Error caught when creating FlashInfer PagedKVCache: %s\n"
|
|
"The model will fallback to TIR-based KV cache.",
|
|
e,
|
|
)
|
|
return []
|
|
|
|
return cache.extern_mods
|