main.py

import time

from src.cli import build_parser

# from torch.profiler import profile, record_function, ProfilerActivity

def main():
    overall_start = time.time()

    parser = build_parser()
    args = parser.parse_args()

    import torch

    from src.app_config import ConfigError, resolve_app_config
    from src.engine import LLMEngine
    from src.runtime_loader import load_runtime

    try:
        config = resolve_app_config(args)
    except ConfigError as error:
        parser.error(str(error))

    if config.seed is not None:
        torch.manual_seed(config.seed)
    else:
        torch.seed()

    # Timing: CUDA initialization
    cuda_init_start = time.time()
    torch.cuda.init() if torch.cuda.is_available() else None
    cuda_init_time = time.time() - cuda_init_start

    print("="*70)
    print("INFERENCE CONFIGURATION")
    print("="*70)
    print(f"Device: {config.device}")
    print(f"CUDA Graph: {config.cuda_graph}")
    print(f"Quantization: {config.quantization or 'none'}")
    if cuda_init_time > 0.001:
        print(f"CUDA init: {cuda_init_time:.4f}s")

    system_prompt_display = config.system_prompt
    if len(system_prompt_display) > 100:
        system_prompt_display = system_prompt_display[:100] + "..."

    user_prompt_display = config.prompt
    if len(user_prompt_display) > 150:
        user_prompt_display = user_prompt_display[:150] + "..."

    print(f"System Prompt: {system_prompt_display}")
    print(f"User Prompt: {user_prompt_display}")
    print("="*70 + "\n")

    # Load tokenizer and model (overlapped when loading weights)
    if config.weights:
        print(f"\n[2] Model loading and tokenizer init:")
        runtime = load_runtime(config)
        print(f"    Tokenizer: {runtime.tokenizer_time_s:.4f}s")
        report = runtime.model_report
        print(f"    Model: loaded={report['loaded']}, missing={len(report['missing'])}, unexpected={len(report['unexpected'])}")
    else:
        print(f"\n[2] Model initialization (random):")
        runtime = load_runtime(config)
        print(f"    Tokenizer: {runtime.tokenizer_time_s:.4f}s")
        print(f"    Random init: {runtime.model_init_time_s:.4f}s")

    # Text generation
    print(f"\n[3] Generation:")
    torch.cuda.nvtx.range_push("Generation")
    engine = LLMEngine(
        model=runtime.model,
        tokenizer=runtime.tokenizer,
        model_args=config.model_args,
        device=config.device,
    )
    request = config.build_request()

    for event in engine.generate(request):
        if event.kind == "prefill" and event.stats is not None:
            print(
                f"    Prefill: {event.stats.prefill_time_s:.4f}s "
                f"({event.stats.prompt_tokens} tokens)"
            )
            print(f"\n{'='*70}")
            print("OUTPUT")
            print(f"{'='*70}")
        elif event.kind == "token" and event.text_delta:
            print(event.text_delta, end="", flush=True)
        elif event.kind == "finish" and event.stats is not None:
            print(f"\n{'='*70}")
            print(
                f"    Decode: {event.stats.decode_time_s:.4f}s "
                f"({event.stats.generated_tokens} tokens, "
                f"TPS: {event.stats.tokens_per_second:.2f} tokens/s)"
            )

    torch.cuda.nvtx.range_pop()

    overall_time = time.time() - overall_start

    print(f"\n{'='*70}")
    print(f"Total: {overall_time:.4f}s")
    print(f"{'='*70}")


if __name__ == "__main__":
    main()