chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,243 @@
+"""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