Oh this is super cool! I think maybe the new RoPE scaling method Llama 3.1 uses isn't yet added in? It's some weird one time scaling mechanism found by a grid search to enable 128K context. Essentially the model was trained on 15.6T tokens on 8K context, then iteratively extended to 128K context with 800B tokens.
Can open a PR if people want :) [Edit: Just opened a PR! Apologies my C is very rusty! https://github.com/trholding/llama2.c/pull/14]
https://github.com/trholding/llama2.c/blob/master/runq.c#L65... needs to be scaled with some weird formula like in https://github.com/unslothai/unsloth/blob/main/unsloth/model...