Inference FLOPs with Long Contexts

The familiar "2N FLOPs per token" ignores attention, which grows with context length. Kaplan et al. (2020) give the forward cost per token at context position n as

C(n) \approx 2N + 2\,n_{layers}\, n\, d_{model}

Implement:

flops_per_token(N, layers, d_model, ctx)
prefill_flops(N, layers, d_model, prompt_len): the sum of flops_per_token for positions 1 … prompt_len. Use the closed form $2NP + n_{layers}\,d_{model}\,P(P + 1)$ .
decode_flops(N, layers, d_model, prompt_len, gen_len): the sum for positions prompt_len + 1 … prompt_len + gen_len
attention_share(N, layers, d_model, ctx): the fraction of the per-token cost spent on the attention term
crossover_ctx(N, layers, d_model): the context length at which the attention term equals the 2N term