modilify_mlx/profile.py

# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Optional per-denoise phase timers. Enabled only with --profile."""

from __future__ import annotations

from collections import defaultdict
import time

import mlx.core as mx

PHASES = (
    "latent",
    "attn",
    "moe",
    "lm_head",
    "softmax",
    "commit",
    "sync",
    "update_cache",
)


class DenoiseProfiler:
    def __init__(self) -> None:
        self.totals = defaultdict(float)
        self.steps = 0

    def add(self, phase: str, seconds: float) -> None:
        self.totals[phase] += float(seconds)

    def finish_step(self) -> None:
        self.steps += 1

    def measure(self, phase: str, *arrays: mx.array):
        mx.eval(*arrays)
        started = time.perf_counter()

        class _Span:
            def __init__(self, profiler: DenoiseProfiler, name: str) -> None:
                self.profiler = profiler
                self.name = name
                self.started = started

            def done(self, *outputs: mx.array) -> None:
                if outputs:
                    mx.eval(*outputs)
                self.profiler.add(self.name, time.perf_counter() - self.started)

        return _Span(self, phase)

    def summary(self) -> str:
        counted = sum(self.totals[name] for name in PHASES)
        lines = [
            f"[profile] steps={self.steps} accounted={counted:.3f}s",
        ]
        for name in PHASES:
            value = self.totals[name]
            share = (100.0 * value / counted) if counted else 0.0
            per = (value / self.steps) if self.steps else 0.0
            lines.append(
                f"[profile] {name:12s}  {value:7.3f}s  {share:5.1f}%  {per*1000:6.1f} ms/step"
            )
        return "\n".join(lines)