Add batch runner example and refactor single sequence runner to use text stream (pytorch#22627)

Co-authored-by: kiymetakdemir <kiymet@meta.com>

Author: metascroy

.github/workflows/mlx.yml

@@ -1037,10 +1037,13 @@ jobs:
         ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \
           -j$(( $(sysctl -n hw.ncpu) - 1 ))
         RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf
-        if [ ! -x "${RUNNER}" ]; then
-          echo "Failed: runner not found at ${RUNNER}"
-          exit 1
-        fi
+        BATCHED_RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_batched
+        for binary in "${RUNNER}" "${BATCHED_RUNNER}"; do
+          if [ ! -x "${binary}" ]; then
+            echo "Failed: runner not found at ${binary}"
+            exit 1
+          fi
+        done
         echo "::endgroup::"
 
         echo "::group::Install LLM requirements"
@@ -1085,3 +1088,28 @@ jobs:
           exit 1
         fi
         echo "::endgroup::"
+
+        echo "::group::Run ${MODEL_NAME} batched off-graph inference"
+        ${BATCHED_RUNNER} \
+          --pte /tmp/${MODEL_NAME}_offgraph.pte \
+          --tokenizer "${TOKENIZER}" \
+          --chat "${CHAT}" \
+          --max-session-tokens 1024 \
+          --max-new-tokens 50 \
+          --max-decode-sequences 2 \
+          --out-prefix /tmp/${MODEL_NAME}_batched \
+          "What is the capital of France?" \
+          "What color is grass?"
+        for check in "0:Paris" "1:green"; do
+          index="${check%%:*}"
+          expected="${check#*:}"
+          output="/tmp/${MODEL_NAME}_batched_${index}.txt"
+          if grep -iq "${expected}" "${output}"; then
+            echo "Success: '${expected}' found in ${output}"
+          else
+            echo "Failed: Expected '${expected}' not found in ${output}"
+            cat "${output}"
+            exit 1
+          fi
+        done
+        echo "::endgroup::"

backends/mlx/examples/llm/CMakeLists.txt

@@ -35,18 +35,10 @@ executorch_target_link_options_shared_lib(executorch)
 set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags)
 find_package(gflags REQUIRED)
 
-set(link_libraries
-    executorch
-    extension_module
-    extension_tensor
-    extension_llm_cache
-    extension_llm_runner
-    extension_llm_sampler
-    gflags
-)
+set(link_libraries executorch extension_llm_runner gflags)
 
 if(NOT TARGET mlxdelegate)
-  message(FATAL_ERROR "mlx_run_llm_hf requires the MLX backend (mlxdelegate)")
+  message(FATAL_ERROR "MLX LLM runners require the MLX backend (mlxdelegate)")
 endif()
 list(APPEND link_libraries mlxdelegate mlx)
 executorch_target_link_options_shared_lib(mlxdelegate)
@@ -67,7 +59,27 @@ target_include_directories(
 )
 target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries})
 
+add_executable(mlx_run_llm_batched run_llm_batched.cpp)
+target_include_directories(
+  mlx_run_llm_batched PUBLIC ${_common_include_directories} ${_json_include}
+                             ${_flatbuffers_include}
+)
+target_link_libraries(
+  mlx_run_llm_batched
+  PUBLIC executorch
+         extension_llm_batching_module
+         extension_llm_runner
+         gflags
+         mlxdelegate
+         mlx
+         tokenizers::tokenizers
+)
+if(TARGET optimized_native_cpu_ops_lib)
+  target_link_libraries(mlx_run_llm_batched PUBLIC optimized_native_cpu_ops_lib)
+endif()
+
 # The copy helper is gated on EXECUTORCH_BUILD_MLX, which the installed config
 # does not set; reaching here means mlxdelegate exists.
 set(EXECUTORCH_BUILD_MLX ON)
 executorch_target_copy_mlx_metallib(mlx_run_llm_hf)
+executorch_target_copy_mlx_metallib(mlx_run_llm_batched)