Repository: mlc-ai/mlc-llm
File: python/mlc_chat/interface/compiler_flags.py (Python)

Code before:
        # - 2) Require `prefill_chunk_size` to present
        # - 3) Set `attention_sink_size` to default (4)
        if self.sliding_window_size is not None:
            self.context_window_size = -1
            logger.info(
                "Setting %s to -1 (disabled), because %s is already set",
                bold("context_window_size"),
                bold("sliding_window_size"),
            )
            if self.prefill_chunk_size is None:
                logger.info(
                    "Default %s to %s (%d) because it is not provided",
                    bold("prefill_chunk_size"),
                    bold("sliding_window_size"),
                    self.sliding_window_size,
                )
                self.prefill_chunk_size = self.sliding_window_size
            elif self.prefill_chunk_size > self.sliding_window_size:
                logger.info(
                    "Lower %s (%d) to %s (%d) as the upper bound",
                    bold("prefill_chunk_size"),
                    self.prefill_chunk_size,
                    bold("sliding_window_size"),
                    self.sliding_window_size,
                )
                self.prefill_chunk_size = self.sliding_window_size
            if self.attention_sink_size is None:
                logger.info(
                    "Default %s to %d because it is not provided",
                    bold("attention_sink_size"),
                    4,
                )
                self.attention_sink_size = 4
        elif self.context_window_size is not None:
            if self.prefill_chunk_size is None:
                logger.info(
                    "Default %s to %s (%d) because it is not provided",
                    bold("prefill_chunk_size"),
                    bold("context_window_size"),
                    self.context_window_size,
                )
                self.prefill_chunk_size = self.context_window_size
            elif self.prefill_chunk_size > self.context_window_size:
                logger.info(
                    "Lower %s (%d) to %s (%d) as the upper bound",
                    bold("prefill_chunk_size"),
                    self.prefill_chunk_size,
                    bold("context_window_size"),
                    self.context_window_size,
                )
                self.prefill_chunk_size = self.context_window_size

Review comment (suggestion):
CC: @CharlieFRuan @davidpissarra I updated a few pieces in the code to make sure `prefill_chunk_size` is upper-bounded by `sliding_window_size` and `context_window_size`

Code after:
        for field in dataclasses.fields(self):
            key = field.name
            value = getattr(self, key)
            if value is None:
                continue
            if key not in updated:
                logger.warning(
                    "%s: Cannot override %s, because %s does not have this field",
                    red("Warning"),
                    bold(key),
                    bold(type(model_config).__name__),
                )
            else:
                logger.info(
                    "Overriding %s from %d to %d",
                    bold(key),
                    updated[key],
                    value,
                )
                updated[key] = value
        return type(model_config).from_dict(updated)

    @staticmethod
    def from_str(source: str) -> "ModelConfigOverride":
        """Parse model config override values from a string."""
        parser = argparse.ArgumentParser(description="model config override values")
        parser.add_argument("--context_window_size", type=int, default=None)
        parser.add_argument("--sliding_window_size", type=int, default=None)
        parser.add_argument("--prefill_chunk_size", type=int, default=None)
        parser.add_argument("--attention_sink_size", type=int, default=None)
        parser.add_argument("--max_batch_size", type=int, default=None)
        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
        return ModelConfigOverride(
            context_window_size=results.context_window_size,
            sliding_window_size=results.sliding_window_size,
            prefill_chunk_size=results.prefill_chunk_size,
            attention_sink_size=results.attention_sink_size,
            max_batch_size=results.max_batch_size,
            tensor_parallel_shards=results.tensor_parallel_shards,
        )


OPT_FLAG_PRESET = {
    "O0": OptimizationFlags(
        flashinfer=False,
        cublas_gemm=False,
        cudagraph=False,
    ),
    "O1": OptimizationFlags(
        flashinfer=False,