See what happens if we export with max_seq_len (#11611)

Summary:
See if any CI is broken by this.


Differential Revision: D76530379

Pulled By: larryliu0820

Author: larryliu0820

extension/llm/export/builder.py

@@ -180,7 +180,7 @@ def _get_dynamic_shape(self) -> Any:
         if self.dynamic_shapes:
             return self.dynamic_shapes
 
-        dim = torch.export.Dim("token_dim", max=self.max_seq_len - 1)
+        dim = torch.export.Dim("token_dim", max=self.max_seq_len)
         if self.enable_dynamic_shape:
             if not self.use_kv_cache:
                 # Only one input argument: tokens

extension/llm/export/test/test_builder.py

@@ -63,7 +63,7 @@ def test_get_dynamic_shape_with_dynamic_shape_enabled_no_kv_cache(self) -> None:
         self.assertIsInstance(result[0], dict)
         self.assertIn(1, result[0])
         # Check that the value at key 1 is a torch.export.Dim with the correct max value
-        self.assertEqual(result[0][1].max, self.max_seq_len - 1)
+        self.assertEqual(result[0][1].max, self.max_seq_len)
 
     def test_get_dynamic_shape_with_dynamic_shape_enabled_with_kv_cache(self) -> None:
         """Test _get_dynamic_shape when enable_dynamic_shape=True and use_kv_cache=True."""
@@ -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 - 1)
+        self.assertEqual(result[0][1].max, self.max_seq_len)
 
         # Check second element (input_pos dimension)
         self.assertIsInstance(result[1], dict)