[v1] Add FlashAttention selection and implement normal / padding-free / dynamic batching (#10469)

src/llamafactory/v1/utils/objects.py

@@ -33,6 +33,10 @@ class StatefulBuffer:
     def size(self) -> int:
         return self._buffer_size
 
+    @property
+    def samples(self) -> list[ModelInput]:
+        return self._buffer
+
     def put(self, samples: list[ModelInput]) -> None:
         """Add samples to the buffer."""
         num_tokens = sum(len(sample["input_ids"]) for sample in samples)

src/llamafactory/v1/utils/types.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Iterator
+from enum import StrEnum, unique
 from typing import TYPE_CHECKING, Any, Literal, NamedTuple, NotRequired, TypedDict, Union
 
 
@@ -54,6 +54,13 @@ else:
     ProcessGroup = None
 
 
+@unique
+class AttentionFunction(StrEnum):
+    EAGER = "eager"
+    SDPA = "sdpa"
+    FLASH_ATTENTION_2 = "flash_attention_2"
+
+
 class DatasetInfo(TypedDict, total=False):
     path: str
     """Local file path."""
@@ -171,8 +178,6 @@ class BatchInfo(TypedDict):
     """Number of micro batches."""
     cutoff_len: int
     """Cutoff length."""
-    data_iter: Iterator[list[ModelInput]]
-    """Data iterator."""
 
 
 class ModelOutput(NamedTuple):