GPT-NeoX allocating full-length KV cache (octoml#179)

This PR changes the GPT-NeoX KV cache creation function to create to
full size at the beginning, so no memory allocation will be required
when running on the fly.

mlc_llm/relax_model/gpt_neox.py

@@ -593,7 +593,7 @@ def create_kv_cache_func(
 ) -> None:
     init_shape = relax.ShapeExpr(
         (
-            1,
+            config.max_sequence_length,
             config.num_attention_heads,
             config.hidden_size // config.num_attention_heads,
         )