update lama export DS specs to be more accurate.

Differential Revision: D83708583

Pull Request resolved: #14737

Author: laithsakka

extension/llm/export/builder.py

@@ -142,9 +142,14 @@ def __init__(
                     {1: torch.export.Dim("token_dim", max=self.max_seq_len - 1)},
                 )
             else:
-                # Two input arguments: tokens and input_pos but input_pos is static shape
+                # 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: torch.export.Dim("token_dim", max=self.max_seq_len - 1)},
                     {"input_pos": {0: 1}},
                 )
 

extension/llm/export/test/test_builder.py

@@ -88,7 +88,7 @@ def test_get_dynamic_shape_with_dynamic_shape_enabled_with_kv_cache(self) -> Non
         # Check first element (tokens dimension)
         self.assertIsInstance(result[0], dict)
         self.assertIn(1, result[0])
-        self.assertEqual(result[0][1].max, self.max_seq_len)
+        self.assertEqual(result[0][1].max, self.max_seq_len - 1)
 
         # Check second element (input_pos dimension)
         self.assertIsInstance(result[1], dict)