class ModelRunner:

    def __init__(self, config: Config, rank: int, event: Event | list[Event]):
        self.config = config
        hf_config = config.hf_config
        self.block_size = config.kvcache_block_size
        self.enforce_eager = config.enforce_eager
        self.world_size = config.tensor_parallel_size
        self.rank = rank
        self.event = event

        dist.init_process_group("hccl", "tcp://localhost:3330", world_size=self.world_size, rank=rank)
        torch.npu.set_device(rank)
        default_dtype = torch.get_default_dtype()
        torch.set_default_dtype(hf_config.torch_dtype)
        torch.set_default_device("npu")
        self.model = model_dict[hf_config.architectures[0]](hf_config)
        load_model(self.model, config.model)
        self.sampler = Sampler()
        torch.npu.empty_cache()
        self.allocate_kv_cache()
        if not self.enforce_eager:
            self.compiler_config = CompilerConfig()
            self.compile_decode = torchair.inference.cache_compile(self.model.forward, config=self.compiler_config,
                                                                   dynamic=False,
                                                                   ge_cache=True)

@torch.inference_mode()
    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool, is_graph_warmup: bool):
        logger.info(f"{'prefill' if is_prefill else 'decode'} execute tokens: {len(input_ids)}")
        if is_prefill or self.enforce_eager or input_ids.size(0) > 512:
            return self.model.compute_logits(self.model(input_ids, positions))
        else:
            
            return self.model.compute_logits(self.compile_decode(input_ids, positions))