Function model_bytes

benchmarking/plot-triton-bench.py:236–243

Minimum bytes transferred for fused matmul+sampling (read W + read X + write indices).

(vocab_size: int, hidden_size: int, n_hidden_states: float)

def model_bytes(vocab_size: int, hidden_size: int, n_hidden_states: float) -> float:
    """Minimum bytes transferred for fused matmul+sampling (read W + read X + write indices)."""
    H = int(n_hidden_states)  # noqa: N806
    return (
        vocab_size * hidden_size * BYTES_PER_ELEMENT  # read weights [V, D]
        + H * hidden_size * BYTES_PER_ELEMENT  # read hidden_states [H, D]
        + H * BYTES_PER_INDEX  # write sampled indices [H]
    )


def model_flops(vocab_size: int, hidden_size: int, n_hidden_states: float) -> float: