Implementation:Mlc ai Mlc llm Compiler Flags

Overview

The Compiler Flags module defines configuration dataclasses for controlling model compilation and optimization in MLC LLM. Located at python/mlc_llm/interface/compiler_flags.py, this file provides three main components: the IPCAllReduceStrategyType enum, the OptimizationFlags dataclass, and the ModelConfigOverride dataclass. It also defines the OPT_FLAG_PRESET dictionary that maps preset names (O0 through O3) to predefined optimization configurations.

Purpose

When compiling models with MLC LLM, users can control various GPU optimization strategies (FlashInfer attention kernels, cuBLAS GEMM, CUDA graphs, CUTLASS, etc.) and override model configuration parameters (context window size, batch size, tensor parallelism, etc.). This module provides the structured types and parsing logic for these compiler and model configuration flags.

File Location

python/mlc_llm/interface/compiler_flags.py

Imports and Dependencies

import dataclasses
import enum
from io import StringIO
from typing import Optional

from mlc_llm.support import argparse, logging
from mlc_llm.support.config import ConfigOverrideBase

IPCAllReduceStrategyType Enum

Defines the inter-process communication all-reduce strategies available for multi-GPU tensor parallelism.

class IPCAllReduceStrategyType(enum.IntEnum):
    """The all-reduce strategy."""
    NONE = 0
    ONESHOT = 1
    TWOSHOT = 2
    AUTO = 3

Value	Name	Description
0	NONE	No IPC all-reduce (single GPU or default communication)
1	ONESHOT	Single-pass all-reduce
2	TWOSHOT	Two-pass all-reduce
3	AUTO	Automatically select the best strategy

OptimizationFlags Dataclass

This dataclass holds boolean flags and strategy settings that control which GPU optimizations are enabled during compilation.

@dataclasses.dataclass
class OptimizationFlags:
    """Optimization flags"""
    flashinfer: bool = False
    cublas_gemm: bool = False
    faster_transformer: bool = False
    cudagraph: bool = False
    cutlass: bool = False
    ipc_allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE

Fields

Field	Type	Default	Description
flashinfer	bool	False	Enable FlashInfer attention kernels (requires CUDA arch >= 80)
cublas_gemm	bool	False	Use cuBLAS for general matrix multiplication (CUDA/ROCm only, specific quantizations)
faster_transformer	bool	False	Enable FasterTransformer optimizations (CUDA only)
cudagraph	bool	False	Enable CUDA graph capture for reduced kernel launch overhead (CUDA only)
cutlass	bool	False	Enable CUTLASS kernels (CUDA only)
ipc_allreduce_strategy	IPCAllReduceStrategyType	NONE	IPC all-reduce strategy for multi-GPU inference

String Representation

The __repr__ method produces a semicolon-delimited string of all flag values:

def __repr__(self) -> str:
    out = StringIO()
    print(f"flashinfer={int(self.flashinfer)}", file=out, end="")
    print(f";cublas_gemm={int(self.cublas_gemm)}", file=out, end="")
    print(f";faster_transformer={int(self.faster_transformer)}", file=out, end="")
    print(f";cudagraph={int(self.cudagraph)}", file=out, end="")
    print(f";cutlass={int(self.cutlass)}", file=out, end="")
    print(f";ipc_allreduce_strategy={self.ipc_allreduce_strategy.name}", file=out, end="")
    return out.getvalue().rstrip()

This format is also the input format for from_str, making the representation round-trippable.

Parsing from String

The from_str static method first checks if the input matches a preset name (e.g., "O2"). If not, it parses the semicolon-delimited key=value string using argparse.

@staticmethod
def from_str(source: str) -> "OptimizationFlags":
    if source in OPT_FLAG_PRESET:
        return OPT_FLAG_PRESET[source]
    # ... argparse-based parsing of semicolon-delimited flags

Target-Aware Update

The update method adjusts optimization flags based on the actual compilation target and quantization scheme. Each flag has a nested validation function:

def update(self, target, quantization) -> None:

The validation logic includes:

• flashinfer: Disabled if the target is not CUDA or if the CUDA architecture is below SM80 (Ampere).
• cublas_gemm: Only enabled on CUDA/ROCm targets with specific quantization formats (q0f16, q0bf16, q0f32, e4m3, e5m2).
• faster_transformer, cutlass, cudagraph: Only enabled on CUDA targets.

ModelConfigOverride Dataclass

This dataclass allows users to override model configuration values at compile time or serve time.

@dataclasses.dataclass
class ModelConfigOverride(ConfigOverrideBase):
    """Flags for overriding model config."""
    context_window_size: Optional[int] = None
    sliding_window_size: Optional[int] = None
    prefill_chunk_size: Optional[int] = None
    attention_sink_size: Optional[int] = None
    max_batch_size: Optional[int] = None
    tensor_parallel_shards: Optional[int] = None
    pipeline_parallel_stages: Optional[int] = None
    disaggregation: Optional[bool] = None

Fields

Field	Type	Description
context_window_size	Optional[int]	Maximum sequence length supported by the model
sliding_window_size	Optional[int]	Sliding window size for sliding window attention (SWA)
prefill_chunk_size	Optional[int]	Chunk size during prefilling
attention_sink_size	Optional[int]	Number of stored attention sinks
max_batch_size	Optional[int]	Maximum batch size for KV cache
tensor_parallel_shards	Optional[int]	Number of shards for tensor parallelism
pipeline_parallel_stages	Optional[int]	Number of stages for pipeline parallelism
disaggregation	Optional[bool]	Whether to enable disaggregation during compilation

All fields default to None, meaning no override is applied unless explicitly set. The from_str static method parses a semicolon-delimited string (e.g., "context_window_size=1024;prefill_chunk_size=128") using argparse.

Optimization Presets

The OPT_FLAG_PRESET dictionary provides four named optimization levels:

OPT_FLAG_PRESET = {
    "O0": OptimizationFlags(
        flashinfer=False, cublas_gemm=False, cudagraph=False,
    ),
    "O1": OptimizationFlags(
        flashinfer=False, cublas_gemm=True, faster_transformer=True,
        cudagraph=False, cutlass=True,
    ),
    "O2": OptimizationFlags(
        flashinfer=True, cublas_gemm=True, faster_transformer=False,
        cudagraph=True, cutlass=True,
        ipc_allreduce_strategy=IPCAllReduceStrategyType.NONE,
    ),
    "O3": OptimizationFlags(
        flashinfer=True, cublas_gemm=True, faster_transformer=True,
        cudagraph=True, cutlass=True,
        ipc_allreduce_strategy=IPCAllReduceStrategyType.AUTO,
    ),
}

Preset	FlashInfer	cuBLAS GEMM	FasterTransformer	CUDA Graph	CUTLASS	IPC AllReduce
O0	Off	Off	Off	Off	Off	NONE
O1	Off	On	On	Off	On	NONE
O2	On	On	Off	On	On	NONE
O3	On	On	On	On	On	AUTO

• O0 disables all optimizations -- useful for debugging and correctness verification.
• O1 enables computational kernel optimizations (cuBLAS, FasterTransformer, CUTLASS) without FlashInfer or CUDA graphs.
• O2 enables FlashInfer and CUDA graphs but disables FasterTransformer. This is the recommended level for most production workloads.
• O3 enables all optimizations including automatic IPC all-reduce strategy selection. The source code documentation notes this could "potentially break the system."

Relationship to Other Modules

• CLI Help (mlc_llm.interface.help) -- The opt and overrides help text entries describe how to use these flags from the command line.
• ConfigOverrideBase (mlc_llm.support.config) -- The base class for ModelConfigOverride that provides the common override application mechanism.
• auto_target (mlc_llm.support.auto_target) -- Used by OptimizationFlags.update() to detect CUDA architecture and validate FlashInfer compatibility.