Function our_method

benchmarking/hbm-access.py:7–15

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

def our_method(vocab_size: int, hidden_size: int, n_hidden_states: int):
    m = vocab_size
    n = n_hidden_states
    k = hidden_size
    # theory: fused matmul + reduction
    # we write only bsz results back to HBM
    reads = m * k + n * k
    writes = n
    return reads, writes


def naive_method(vocab_size: int, hidden_size: int, n_hidden_states: int):