"""The compilation pipeline for LLM applications.""" from pathlib import Path from typing import Any, Dict, List, Optional # noqa: UP035 import tvm from tvm import IRModule from tvm.relax import register_pipeline from tvm.relax.frontend import nn from tvm.s_tir import dlight as dl from mlc_llm.interface.compiler_flags import IPCAllReduceStrategyType from mlc_llm.support import logging from .attach_cuda_graph_alloc_init_func import AttachCUDAGraphAllocInitFunc from .attach_embedding_allocator import AttachAllocEmbeddingTensorFunc from .attach_logit_processor import AttachLogitProcessFunc from .attach_sampler import AttachGPUSamplingFunc from .attach_softmax_with_temperature import AttachSoftmaxWithTemperature from .attach_spec_decode_aux_funcs import AttachSpecDecodeAuxFuncs from .attach_support_info import ( AttachAdditionalPrimFuncs, AttachCUDAGraphSymbolicCaptureHints, AttachMemoryPlanAttr, AttachPipelineParallelStages, AttachSequenceLengthPaddingFactor, AttachVariableBounds, ) from .blas_dispatch import BLASDispatch from .clean_up_tir_attrs import CleanUpTIRAttrs from .dispatch_kv_cache_creation import DispatchKVCacheCreation from .dispatch_triton_kernel import DispatchTritonKernel from .estimate_memory_usage import AttachMetadataWithMemoryUsage from .fuse_add_norm import FuseAddRMSNorm from .fuse_dequantize_matmul_ewise import FuseDequantizeMatmulEwise from .fuse_dequantize_take import FuseDequantizeTake from .fuse_dequantize_transpose import FuseDequantizeTranspose from .fuse_ft_dequantize_matmul_epilogue import FuseFTDequantizeEpilogue from .fuse_transpose_matmul import FuseTransposeMatmul from .lift_global_buffer_alloc import LiftTIRGlobalBufferAlloc from .low_batch_specialization import LowBatchGemvSpecialize from .pipeline_parallel_rewrite import PipelineParallelRewrite from .scatter_tuple_get_item import ScatterTupleGetItem logger = logging.getLogger(__name__) @tvm.transform.module_pass(opt_level=0, name="_LogProgress") class _LogProgress: """A dummy compiler pass that does nothing but logging.""" def __init__(self, *args): self.args = args def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: """A dummy transformation""" logger.info(*self.args) return mod @tvm.transform.module_pass(opt_level=0, name="DebugDump") class _DebugDump: """A dummy compiler pass that does nothing but logging. Only enabled when debug_dump is not None""" def __init__(self, file_name: str, file_path: Optional[Path], show_meta: bool = False): self.file_name = file_name self.file_path = file_path self.show_meta = show_meta def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: """A dummy transformation that dumps the module to file""" if self.file_path is not None: # NOTE: We use debug level here to avoid spamming the console logger.debug("Dumping IR to %s", self.file_path / self.file_name) with open(self.file_path / self.file_name, "w", encoding="utf-8") as f: f.write(mod.script(show_meta=self.show_meta)) return mod @register_pipeline("mlc_llm") def _mlc_llm_pipeline( target: tvm.target.Target, flashinfer: bool = False, cublas_gemm: bool = False, faster_transformer: bool = False, allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE, variable_bounds: Optional[Dict[str, int]] = None, # noqa: UP006 cuda_graph_symbolic_capture_hints: Optional[Dict[str, List[str]]] = None, # noqa: UP006 additional_tirs: Optional[Dict[str, tvm.tirx.PrimFunc]] = None, # noqa: UP006 metadata: Optional[Dict[str, Any]] = None, # noqa: UP006 ext_mods: Optional[List[nn.ExternModule]] = None, # noqa: UP006 debug_dump: Optional[Path] = None, ): variable_bounds = variable_bounds or {} cuda_graph_symbolic_capture_hints = cuda_graph_symbolic_capture_hints or {} additional_tirs = additional_tirs or {} metadata = metadata or {} ext_mods = ext_mods or [] tensor_parallel_shards = metadata.get("tensor_parallel_shards", 1) @tvm.transform.module_pass(opt_level=0) def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule: seq = tvm.transform.Sequential( [ # Phase 0. Add additional information for compilation and remove unused Relax func DispatchKVCacheCreation(target, flashinfer, metadata), AttachSoftmaxWithTemperature(target, metadata), AttachVariableBounds(variable_bounds), AttachCUDAGraphSymbolicCaptureHints(cuda_graph_symbolic_capture_hints), AttachPipelineParallelStages(metadata["pipeline_parallel_stages"]), AttachLogitProcessFunc(target), AttachAdditionalPrimFuncs(additional_tirs), AttachAllocEmbeddingTensorFunc(metadata), AttachGPUSamplingFunc(target, variable_bounds), AttachSpecDecodeAuxFuncs(tensor_parallel_shards), AttachMemoryPlanAttr(), AttachSequenceLengthPaddingFactor(target, metadata), tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)), _DebugDump("debug-phase0.py", debug_dump, show_meta=False), # Phase 1. Passes on high-level operator graph _LogProgress("Running TVM Relax graph-level optimizations"), DispatchTritonKernel(target), FuseFTDequantizeEpilogue(), FuseDequantizeTranspose(), BLASDispatch(target) if cublas_gemm else tvm.transform.Sequential([]), ( FuseAddRMSNorm(target=target) if target.kind.name != "llvm" else tvm.transform.Sequential([]) ), FuseTransposeMatmul(), _DebugDump("debug-phase1.py", debug_dump, show_meta=False), # Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline _LogProgress("Lowering to TVM TIR kernels"), tvm.relax.backend.DispatchSampling(), tvm.relax.backend.DispatchSortScan(), tvm.relax.transform.LegalizeOps(), tvm.relax.transform.AnnotateTIROpPattern(), tvm.relax.transform.FoldConstant(), tvm.relax.transform.FuseOps(), tvm.relax.transform.FuseTIR(), _DebugDump("debug-phase2.py", debug_dump, show_meta=False), # Phase 3. Passes on TIR _LogProgress("Running TVM TIR-level optimizations"), FuseDequantizeMatmulEwise(), FuseDequantizeTake(), tvm.relax.transform.DeadCodeElimination(), CleanUpTIRAttrs(["op_pattern"]), _DebugDump("debug-phase3.py", debug_dump, show_meta=False), # Phase 4. Low-level Optimizations _LogProgress("Running TVM Dlight low-level optimizations"), LowBatchGemvSpecialize(), ( dl.ApplyDefaultSchedule( dl.gpu.Matmul(), dl.gpu.GEMV(), dl.gpu.Reduction(), dl.gpu.GeneralReduction(), dl.gpu.Fallback(), ) if target.kind.name != "llvm" else dl.ApplyDefaultSchedule( dl.cpu.GEMV(), ) ), _DebugDump("debug-phase4.py", debug_dump, show_meta=False), _LogProgress("Lowering to VM bytecode"), ( LiftTIRGlobalBufferAlloc() if target.kind.name != "llvm" else tvm.transform.Sequential([]) ), ( tvm.tirx.transform.ForceNarrowIndexToInt32() if target.kind.name != "cuda" else tvm.transform.Sequential([]) ), ScatterTupleGetItem(), PipelineParallelRewrite(), tvm.relax.transform.RewriteDataflowReshape(), tvm.relax.transform.ToNonDataflow(), tvm.relax.transform.RemovePurityChecking(), tvm.relax.transform.CallTIRRewrite(), ( tvm.relax.transform.IPCAllReduceRewrite(allreduce_strategy) if allreduce_strategy != IPCAllReduceStrategyType.NONE else tvm.transform.Sequential([]) ), tvm.relax.transform.StaticPlanBlockMemory(), AttachMetadataWithMemoryUsage(metadata), _DebugDump("debug-phase5.py", debug_dump, show_meta=False), tvm.relax.transform.RewriteCUDAGraph(), AttachCUDAGraphAllocInitFunc(), tvm.relax.transform.LowerGPUIPCAllocStorage(), tvm.relax.transform.LowerAllocTensor(), tvm.relax.transform.KillAfterLastUse(), tvm.relax.transform.LowerRuntimeBuiltin(), tvm.relax.transform.VMShapeLower(), tvm.relax.transform.AttachGlobalSymbol(), _LogProgress("Compiling external modules"), tvm.relax.transform.AttachExternModules(ext_mods), _LogProgress("Compilation complete! Exporting to disk"), ] ) mod = seq(mod) return mod return _pipeline