mlc-llm: [Bug] Codellama 34B cannot handle inputs > 4k tokens despite context window of 16k

🐛 Bug

Codellama 34B model does not behave as expected when passing inputs with more than 4k tokens. The model seems to lose the context and start spitting out random code imports.

Notes: Codellama 13B does not exhibit the same behavior and can handle inputs up to 16K tokens fine. Codellama 34B fp16, int8, and int4 all have the same issue with tokens >4K.

To Reproduce

Steps to reproduce the behavior:

1. Download Codellama 34B
2. Compile the model using MLC-LLM
3. Create a request with >4k tokens
4. Run the model with the request.

Expected behavior

Model should respect the context of the input and respond accordingly.

Environment

• Platform (e.g. WebGPU/Vulkan/IOS/Android/CUDA): CUDA
• Operating system (e.g. Ubuntu/Windows/MacOS/…): Ubuntu
• Device (e.g. iPhone 12 Pro, PC+RTX 3090, …) A100
• How you installed MLC-LLM (conda, source): Yes
• How you installed TVM-Unity (pip, source): Yes
• Python version (e.g. 3.10): 3.10
• GPU driver version (if applicable):
• CUDA/cuDNN version (if applicable): CUDA 12.0
• TVM Unity Hash Tag (python -c "import tvm; print('\n'.join(f'{k}: {v}' for k, v in tvm.support.libinfo().items()))", applicable if you compile models): 631f37b6bf8b101d16ecc55de7e6a749a3588570
• Any other relevant information:

Additional context