Fix GQA shape inference for shape-derived int32 sequence lengths

.github/workflows/react_native.yml

@@ -284,13 +284,26 @@ jobs:
         run: npm run detox -- build --configuration ios.sim_18_6.release
         working-directory: ${{ github.workspace }}/js/react_native/e2e
 
-      - name: Reset iOS Simulators before Detox
+      - name: Prepare Detox iOS Simulator
+        timeout-minutes: 10
         run: |
-          set -e -x
-          xcrun simctl shutdown all || true
-          xcrun simctl erase all || true
-          killall Simulator || true
-          killall com.apple.CoreSimulator.CoreSimulatorService || true
+          set -euo pipefail
+          simulator_type=$(node -p "require('./.detoxrc.js').devices.simulator_ios_18_6.device.type")
+          simulator_os=$(node -p "require('./.detoxrc.js').devices.simulator_ios_18_6.device.os")
+          device=$(applesimutils --list --byType "$simulator_type" --byOS "$simulator_os" |
+            jq -e 'if length == 1 then .[0] else error("Expected exactly one simulator matching the Detox configuration") end')
+          device_id=$(jq -er '.udid' <<< "$device")
+          device_state=$(jq -er '.state' <<< "$device")
+          case "$device_state" in
+            Booted) xcrun simctl shutdown "$device_id" ;;
+            Shutdown) ;;
+            *) echo "Unexpected state for simulator $device_id: $device_state" >&2; exit 1 ;;
+          esac
+          xcrun simctl erase "$device_id"
+          xcrun simctl boot "$device_id"
+          xcrun simctl bootstatus "$device_id" -b
+          xcrun simctl listapps "$device_id" > /dev/null
+        working-directory: ${{ github.workspace }}/js/react_native/e2e
 
       - name: Run Detox iOS e2e Tests
         run: |

js/react_native/README.md

@@ -58,8 +58,16 @@ From `js/react_native`, run `npm run test:cpp` for the native option-policy test
 `npm run bootstrap-no-pods` runs both checks before the Android and iOS Detox tests in CI.
 
 The device checks exercise in-flight aborts, pre-dispatch aborts, and isolated Env-listener teardown.
+They verify native worker completion hooks run before success or failure settlement is dispatched
+to JavaScript. EPContext Detox navigation targets the button's unique `testID`, avoiding child text
+views and accessibility labels duplicated across native wrappers on iOS.
 The `__testEpContextDataReadCallback` native hook and its `__testWorker` handles are internal test APIs.
 
+The [macOS CI workflow](../../.github/workflows/react_native.yml) resets only the simulator
+selected by the iOS 18.6 Detox configuration and waits for boot and app-service readiness before
+starting the device tests. Detox still uninstalls and reinstalls the app, and its 240-second
+setup deadline and device assertions remain unchanged.
+
 ### License
 
 License information can be found [here](https://github.com/microsoft/onnxruntime/blob/main/README.md#license).

js/react_native/cpp/AsyncWorker.h

@@ -57,14 +57,22 @@ class AsyncWorker : public HostObject, public std::enable_shared_from_this<Async
                   cancel_ = false;
                   worker_ = std::thread([this]() {
                     if (!cancel_) {
+                      bool finished = false;
                       try {
                         execute();
+                        // Dispatch can release the last worker reference on this thread.
+                        onFinished();
+                        finished = true;
                         dispatchResolve();
                       } catch (const std::exception& e) {
+                        if (!finished) {
+                          onFinished();
+                        }
                         dispatchReject(e.what());
                       }
+                    } else {
+                      onFinished();
                     }
-                    onFinished();
                   });
                   return Value::undefined();
                 }));

js/react_native/e2e/src/App.tsx

@@ -112,6 +112,7 @@ export default class App extends React.PureComponent<{}, State> {
                 title="EPContext Data Read Test"
                 onPress={() => this.navigateTo('ep-context-data-read')}
                 accessibilityLabel="ep-context-data-read-test-button"
+                testID="ep-context-data-read-test-button"
               />
             </View>
             <Text style={styles.description}>Test the epContextDataRead session option and its lifetime</Text>

js/react_native/e2e/src/EPContextDataReadTest.tsx

@@ -341,7 +341,7 @@ export default class EPContextDataReadTest extends React.PureComponent<{}, State
     this.updateTestResult(index, { status: 'running' });
     try {
       const expectedName = 'context/data.bin';
-      const result = await getOrtApi().__testEpContextDataReadCallback(
+      const pendingRead = getOrtApi().__testEpContextDataReadCallback(
         (name) => {
           if (name !== expectedName) {
             throw new Error(`Unexpected callback name: ${name}`);
@@ -351,6 +351,10 @@ export default class EPContextDataReadTest extends React.PureComponent<{}, State
         3,
         expectedName,
       );
+      const result = await pendingRead;
+      if (!pendingRead.__testWorker.isFinished) {
+        throw new Error('Callback worker resolved before its native completion hook');
+      }
       if (result.length !== 3 || result[0] !== 1 || result[1] !== 2 || result[2] !== 3) {
         throw new Error(`Unexpected callback bytes: ${Array.from(result).join(',')}`);
       }
@@ -377,15 +381,19 @@ export default class EPContextDataReadTest extends React.PureComponent<{}, State
           throw new Error('test callback failure');
         },
       ]) {
+        const pendingRead = getOrtApi().__testEpContextDataReadCallback(callback, 1, 'failure.bin');
         let rejected = false;
         try {
-          await getOrtApi().__testEpContextDataReadCallback(callback, 1, 'failure.bin');
+          await pendingRead;
         } catch {
           rejected = true;
         }
         if (!rejected) {
           throw new Error('Invalid callback result unexpectedly succeeded');
         }
+        if (!pendingRead.__testWorker.isFinished) {
+          throw new Error('Callback worker rejected before its native completion hook');
+        }
       }
       const empty = await getOrtApi().__testEpContextDataReadCallback(() => new Uint8Array(0), 1, 'empty.bin');
       if (empty.length !== 0) {

js/react_native/e2e/test/OnnxruntimeModuleExample.test.js

@@ -60,11 +60,7 @@ describe('OnnxruntimeModuleExample', () => {
 
   it('EPContext data read test should run successfully', async () => {
     // Tap EPContext data read test button
-    if (platform === 'android') {
-      await element(by.label('ep-context-data-read-test-button')).tap();
-    } else {
-      await element(by.text('EPContext Data Read Test')).tap();
-    }
+    await element(by.id('ep-context-data-read-test-button')).tap();
 
     await delay(500);
 

js/react_native/e2e/unit-test/detox-navigation.test.js

@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { readFileSync } from 'node:fs';
+import { join } from 'node:path';
+import { runInNewContext } from 'node:vm';
+import * as React from 'react';
+import { Button } from 'react-native';
+import renderer, { act } from 'react-test-renderer';
+import App from '../src/App';
+
+jest.mock('../src/MNISTTest', () => 'MNISTTest');
+jest.mock('../src/BasicTypesTest', () => 'BasicTypesTest');
+jest.mock('../src/EPContextDataReadTest', () => 'EPContextDataReadTest');
+
+test('EPContext navigation identifier belongs to exactly one native button', () => {
+  let tree;
+  act(() => {
+    tree = renderer.create(React.createElement(App));
+  });
+  const nativeTargets = tree.root.findAll(
+    (node) => typeof node.type === 'string' && node.props.testID === 'ep-context-data-read-test-button',
+  );
+  expect(nativeTargets).toHaveLength(1);
+  const button = tree.root.find(
+    (node) => node.type === Button && node.props.testID === 'ep-context-data-read-test-button',
+  );
+  act(() => {
+    button.props.onPress();
+  });
+  expect(tree.root.findByType('EPContextDataReadTest')).toBeDefined();
+  act(() => {
+    tree.unmount();
+  });
+});
+
+test.each(['android', 'ios'])('EPContext navigation taps the unique button on %s', async (platform) => {
+  const scenarios = new Map();
+  let initialize;
+  const element = jest.fn((matcher) => {
+    if (platform === 'ios' && matcher.label === 'ep-context-data-read-test-button') {
+      throw new Error('Multiple native elements share the EPContext accessibility label');
+    }
+    return { tap: jest.fn().mockResolvedValue(undefined) };
+  });
+  const waitFor = jest.fn(() => ({
+    toHaveText: jest.fn(() => ({ withTimeout: jest.fn().mockResolvedValue(undefined) })),
+  }));
+  runInNewContext(readFileSync(join(__dirname, '../test/OnnxruntimeModuleExample.test.js'), 'utf8'), {
+    describe: (_name, register) => register(),
+    beforeAll: (callback) => {
+      initialize = callback;
+    },
+    beforeEach: jest.fn(),
+    it: (name, callback) => scenarios.set(name, callback),
+    device: { getPlatform: () => platform, launchApp: jest.fn().mockResolvedValue(undefined) },
+    element,
+    by: { id: (id) => ({ id }), label: (label) => ({ label }), text: (text) => ({ text }) },
+    waitFor,
+    setTimeout: (callback) => callback(),
+  });
+
+  await initialize();
+  await scenarios.get('EPContext data read test should run successfully')();
+
+  expect(element).toHaveBeenNthCalledWith(1, { id: 'ep-context-data-read-test-button' });
+  expect(element.mock.results[0].value.tap).toHaveBeenCalledTimes(1);
+  expect(waitFor).toHaveBeenCalledWith({ tap: expect.any(Function) });
+  expect(waitFor.mock.results[0].value.toHaveText).toHaveBeenCalledWith('14/14 checks passed');
+});

js/react_native/e2e/unit-test/ep-context-data-read.test.js

@@ -114,3 +114,42 @@ test('cancellation checks distinguish in-flight, queued, and Env-listener teardo
     workers[2].invalidateEnv.mock.invocationCallOrder[0],
   );
 });
+
+test('successful callback checks reject settlement before native worker completion', async () => {
+  const component = createComponent();
+  const pendingRead = Promise.resolve(new Uint8Array([1, 2, 3]));
+  pendingRead.__testWorker = { isFinished: false };
+  globalThis.OrtApi = {
+    __testEpContextDataReadCallback: jest.fn(() => pendingRead),
+  };
+
+  await component.runCallbackBridgeCheck(11);
+
+  expect(component.state.testResults[11]).toEqual(
+    expect.objectContaining({
+      status: 'error',
+      message: 'Callback worker resolved before its native completion hook',
+    }),
+  );
+});
+
+test('failed callback checks reject settlement before native worker completion', async () => {
+  const component = createComponent();
+  globalThis.OrtApi = {
+    __testEpContextDataReadCallback: jest.fn((_callback, _maxDataSize, name) => {
+      const pendingRead =
+        name === 'empty.bin' ? Promise.resolve(new Uint8Array(0)) : Promise.reject(new Error('callback failure'));
+      pendingRead.__testWorker = { isFinished: name === 'empty.bin' };
+      return pendingRead;
+    }),
+  };
+
+  await component.runCallbackBridgeFailureCheck(12);
+
+  expect(component.state.testResults[12]).toEqual(
+    expect.objectContaining({
+      status: 'error',
+      message: 'Callback worker rejected before its native completion hook',
+    }),
+  );
+});

onnxruntime/core/graph/contrib_ops/bert_defs.cc

@@ -389,11 +389,16 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte
     const auto* total_sequence_length_data =
         total_sequence_length_index >= 0 ? ctx.getInputData(total_sequence_length_index) : nullptr;
     if (total_sequence_length_data != nullptr) {
-      const auto& data = ParseData<int32_t>(total_sequence_length_data);
-      if (data.size() != 1) {
-        fail_shape_inference("total_sequence_length input must contain a single element");
-      }
-      total_sequence_length_value = static_cast<int64_t>(data[0]);
+      const auto read_single_value = [](const auto& data) -> int64_t {
+        if (data.size() != 1) {
+          fail_shape_inference("total_sequence_length input must contain a single element");
+        }
+        return static_cast<int64_t>(data[0]);
+      };
+      // Propagated shape values use int64 even when the input type is int32.
+      total_sequence_length_value = total_sequence_length_data->data_type() == TensorProto_DataType_INT64
+                                        ? read_single_value(ParseData<int64_t>(total_sequence_length_data))
+                                        : read_single_value(ParseData<int32_t>(total_sequence_length_data));
     }
 
     if (past_key_index >= 0 && hasInputShape(ctx, past_key_index)) {

onnxruntime/test/contrib_ops/group_query_attention_shape_inference_test.cc

@@ -150,6 +150,82 @@ TEST(GroupQueryAttentionShapeInferenceTest, UnshapedValueInputIsNotPacked) {
   ASSERT_STATUS_OK(status);
 }
 
+TEST(GroupQueryAttentionShapeInferenceTest, Int32TotalSequenceLengthFromStaticShape) {
+  for (bool from_shape : {false, true}) {
+    for (bool scalar : {true, false}) {
+      for (int64_t past_length : {int64_t{-1}, int64_t{3}, int64_t{9}}) {
+        SCOPED_TRACE(testing::Message() << "from_shape=" << from_shape << ", scalar=" << scalar
+                                        << ", past_length=" << past_length);
+        const std::unordered_map<std::string, int> domain_to_version{{kOnnxDomain, 20}, {kMSDomain, 1}};
+        Model model("gqa_static_total_sequence_length", false, ModelMetaData(), PathString(),
+                    IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {},
+                    DefaultLoggingManager().DefaultLogger());
+        Graph& graph = model.MainGraph();
+        ModelTestBuilder builder(graph);
+        NodeArg* empty = builder.MakeEmptyInput();
+        NodeArg* query = builder.MakeInput<float>(std::vector<int64_t>{1, 7, 64});
+        NodeArg* key = builder.MakeInput<float>(std::vector<int64_t>{1, 7, 32});
+        NodeArg* value = builder.MakeInput<float>(std::vector<int64_t>{1, 7, 32});
+        NodeArg* past_key = empty;
+        NodeArg* past_value = empty;
+        if (past_length >= 0) {
+          past_key = builder.MakeInput<float>(std::vector<int64_t>{1, 2, past_length, 16});
+          past_value = builder.MakeInput<float>(std::vector<int64_t>{1, 2, past_length, 16});
+        }
+
+        NodeArg* total_sequence_length;
+        if (from_shape) {
+          NodeArg* length64 = builder.MakeIntermediate();
+          Node& shape = builder.AddNode("Shape", {query}, {length64});
+          shape.AddAttribute("start", int64_t{1});
+          shape.AddAttribute("end", int64_t{2});
+          NodeArg* length32 = builder.MakeIntermediate();
+          Node& cast = builder.AddNode("Cast", {length64}, {length32});
+          cast.AddAttribute("to", int64_t{ONNX_NAMESPACE::TensorProto_DataType_INT32});
+          total_sequence_length = length32;
+          if (scalar) {
+            total_sequence_length = builder.MakeIntermediate();
+            builder.AddNode("Squeeze", {length32}, {total_sequence_length});
+          }
+        } else {
+          total_sequence_length = builder.MakeInitializer<int32_t>(
+              scalar ? std::vector<int64_t>{} : std::vector<int64_t>{1}, {7});
+        }
+
+        NodeArg* seqlens_k = builder.MakeInput<int32_t>({1}, {6});
+        NodeArg* output = builder.MakeOutput();
+        NodeArg* present_key = builder.MakeOutput();
+        NodeArg* present_value = builder.MakeOutput();
+        Node& gqa = builder.AddNode(
+            "GroupQueryAttention", {query, key, value, past_key, past_value, seqlens_k, total_sequence_length},
+            {output, present_key, present_value}, kMSDomain);
+        gqa.AddAttribute("num_heads", int64_t{4});
+        gqa.AddAttribute("kv_num_heads", int64_t{2});
+        builder.SetGraphOutputs();
+
+        ASSERT_STATUS_OK(graph.Resolve());
+        EXPECT_EQ(total_sequence_length->TypeAsProto()->tensor_type().elem_type(),
+                  ONNX_NAMESPACE::TensorProto_DataType_INT32);
+        const auto* output_shape = output->Shape();
+        ASSERT_NE(output_shape, nullptr);
+        ASSERT_EQ(output_shape->dim_size(), 3);
+        EXPECT_EQ(output_shape->dim(0).dim_value(), 1);
+        EXPECT_EQ(output_shape->dim(1).dim_value(), 7);
+        EXPECT_EQ(output_shape->dim(2).dim_value(), 64);
+        for (const NodeArg* present : {present_key, present_value}) {
+          const auto* present_shape = present->Shape();
+          ASSERT_NE(present_shape, nullptr);
+          ASSERT_EQ(present_shape->dim_size(), 4);
+          EXPECT_EQ(present_shape->dim(0).dim_value(), 1);
+          EXPECT_EQ(present_shape->dim(1).dim_value(), 2);
+          EXPECT_EQ(present_shape->dim(2).dim_value(), past_length > 7 ? past_length : 7);
+          EXPECT_EQ(present_shape->dim(3).dim_value(), 16);
+        }
+      }
+    }
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace onnxruntime