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,39 @@
+"""The pass that attaches embedding allocation function to the IRModule."""
+
+from typing import Any, Dict  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachAllocEmbeddingTensorFunc")
+class AttachAllocEmbeddingTensorFunc:
+    """Attach embedding tensor allocation Relax function to IRModule."""
+
+    def __init__(self, metadata: Dict[str, Any]):  # noqa: UP006
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        embed_func = None
+        for gv, func in mod.functions_items():
+            if gv.name_hint == "embed":
+                embed_func = func
+
+        if embed_func is None:
+            return mod
+
+        hidden_size = embed_func.ret_ty.shape[-1]
+        dtype = relax.DataTypeImm(embed_func.ret_ty.dtype.dtype)
+        bb = relax.BlockBuilder(mod)
+        with bb.function("alloc_embedding_tensor", []):
+            bb.emit_func_output(
+                bb.emit(
+                    relax.op.builtin.alloc_tensor(
+                        relax.ShapeExpr([self.metadata["prefill_chunk_size"], hidden_size]),
+                        dtype,
+                        runtime_device_index=0,
+                    )
+                )
+            )
+        return bb.finalize()