[Llama3] Debug accuracy drop on sequence lengths > 64k tokens

[mtairum]

Status:

Prefill seqlen up to 128k supported in Llama3 codebase (1b, 3B, 8B, 11B, 70B), but bad accuracy in some corner cases:

• ✅ All models show good output for prefill seqlen up to 32K tokens.
• ❌ All models showing bad (repetitive) output for 128K tokens.
• ✅ With top-p sampling, prefill up to 64K is good in most models.
  • ❌ Exception 1B/3B

Things we've tried so far:

• TODO

[mtairum]

This is now supported in main.

Added a P2 followup for small Llama models here: #15737