Use Float16Array instead of Uint16Array for kvcache when available (#1208)

Author: xenova

src/models.js

@@ -108,6 +108,7 @@ import {
     stack,
     std_mean,
     Tensor,
+    DataTypeMap,
 } from './utils/tensor.js';
 import { RawImage } from './utils/image.js';
 
@@ -1847,7 +1848,7 @@ export class PreTrainedModel extends Callable {
         } else {
             const session = this.sessions['decoder_model_merged'] ?? this.sessions['model'];
             const dtype = session?.config?.kv_cache_dtype ?? 'float32';
-            const empty = (dtype === 'float16') ? new Uint16Array() : [];
+            const empty = (dtype === 'float16') ? new DataTypeMap.float16() : [];
 
             const batch_size = (decoderFeeds[this.main_input_name] ?? decoderFeeds.attention_mask)?.dims?.[0] ?? 1;
             const shapes = getKeyValueShapes(this.config, { batch_size });

src/utils/tensor.js

@@ -20,9 +20,11 @@ import {
 
 import { TensorOpRegistry } from '../ops/registry.js';
 
-const DataTypeMap = Object.freeze({
+export const DataTypeMap = Object.freeze({
     float32: Float32Array,
-    float16: Uint16Array,
+    // @ts-expect-error ts(2552) Limited availability of Float16Array across browsers:
+    // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Float16Array
+    float16: typeof Float16Array !== "undefined" ? Float16Array: Uint16Array,
     float64: Float64Array,
     string: Array, // string[]
     int8: Int8Array,

tests/init.js

@@ -22,7 +22,7 @@ export function init() {
   let registerBackend = ONNX_COMMON.registerBackend;
 
   // Define the constructors to monkey-patch
-  const TYPED_ARRAYS_CONSTRUCTOR_NAMES = ["Int8Array", "Int16Array", "Int32Array", "BigInt64Array", "Uint8Array", "Uint8ClampedArray", "Uint16Array", "Uint32Array", "BigUint64Array", "Float32Array", "Float64Array"];
+  const TYPED_ARRAYS_CONSTRUCTOR_NAMES = ["Int8Array", "Int16Array", "Int32Array", "BigInt64Array", "Uint8Array", "Uint8ClampedArray", "Uint16Array", "Uint32Array", "BigUint64Array", "Float16Array", "Float32Array", "Float64Array"];
 
   // Keep a reference to the original initialization method
   const originalMethod = onnxruntimeBackend.init;
@@ -36,6 +36,7 @@ export function init() {
     for (const ctorName of TYPED_ARRAYS_CONSTRUCTOR_NAMES) {
       // Get the constructor from the current context
       const ctor = globalThis[ctorName];
+      if (ctor === undefined) continue; // If unavailable, skip the patching
 
       // Get the corresponding test function from the `util` module
       const value = types[`is${ctorName}`].bind(types);