Fix max seq length bug

Summary:
update cache op was forcing incorrect constraint on sequence length.
Fixing that alongw with fixing export allows us to correctly export the model

Differential Revision: D84562463

Author: kimishpatel

extension/llm/custom_ops/custom_ops.py

@@ -207,8 +207,8 @@ def _validate_update_cache_params(
             1
         ), f"Start position {start_pos} must be less than sequence length {cache.size(1)}"
 
-        torch._check((start_pos + seq_len) < cache.size(1))
-        assert (start_pos + seq_len) < cache.size(
+        torch._check((start_pos + seq_len) <= cache.size(1))
+        assert (start_pos + seq_len) <= cache.size(
             1
         ), f"Start position  + length = {start_pos + seq_len} must be less than sequence length {cache.size(1)}"
 

extension/llm/export/builder.py

@@ -144,12 +144,8 @@ def __init__(
             else:
                 # Two input arguments: tokens and input_pos but input_pos is static shape.
 
-                # A runtime assertion is added by torch.ops.llama.update_cache requires that
-                # L['tokens'].size()[1] + input_pos[0].item() < self.max_seq_len
-                # This consttaint L['tokens'].size()[1] to be elf.max_seq_len-1
-                # run with TORCH_LOGS=+dynamic for details
                 self.dynamic_shapes = (
-                    {1: torch.export.Dim("token_dim", max=self.max_seq_len - 1)},
+                    {1: torch.export.Dim("token_dim", max=self.max_seq_len)},
                     {"input_pos": {0: 1}},
                 )