[Performance] Inference and data copy are slow when running Dolphin-1.5 in Java

[ningpp]

Describe the issue

When running Dolphin-1.5 model in Java with ONNX Runtime, the inference process is slow and data copy operations are also slow. This impacts overall application performance. Please investigate possible optimizations for inference speed and data transfer efficiency when using Java API.

To reproduce

decoder sample code:

        for (int i = 0; i < maxNewTokens; i++) {

            Map<String, OnnxTensor> inputs = new HashMap<>(2);
            long[] flat = ArrayUtil.flat(inputIds);
            LongBuffer buffer = LongBuffer.wrap(flat);
            OnnxTensor input_ids_tensor = OnnxTensor.createTensor(env, buffer,
                    new long[] {inputIds.length, inputIds[0].length});
            inputs.put("input_ids", input_ids_tensor);
            inputs.put("encoder_hidden_states", encoder_hidden_states_tensor);

            OrtSession.Result onnxResult = session.run(inputs, outputNames);
            Optional<OnnxValue> optinalResult = onnxResult.get(List.copyOf(outputNames).get(0));
            if (optinalResult.isPresent()) {
                float[][][] decoderResultFloats = (float[][][]) optinalResult.get().getValue();
                long[][] nextIds = new long[batchSize][1];
                ......
            }
    }

Urgency

No response

Platform

Windows

OS Version

windows 11

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

f217402

ONNX Runtime API

Java

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA 12.6

Model File

Dolphin-1.5

Is this a quantized model?

No

[ningpp]

original model: https://huggingface.co/ByteDance/Dolphin-1.5

[Craigacp]

You should allocate a single direct ByteBuffer per thread, wrap that as a LongBuffer then slice it to the size you need for the current input. This will prevent allocations and copies on the input side. On the output side don't use getValue to get a 3d float array out, 3d arrays in Java are very slow due to a lot of pointer chasing, you should get a FloatBuffer back out and then use that. You can also preallocate any tensors used by the outputs which can be reused for each output after appropriate slicing, this will remove allocations and copying in the output.

If you're using a GPU then there is unavoidably some copying between Java and GPU memory as that is not easily accessible from Java.

As you're iteratively decoding tokens out of the model, then you need to figure out how the ONNX model wants to store the kv cache, and make sure you supply it at each time step. If you don't then it'll recalculate the kv projection for all tokens when generating each token, which makes it very slow.

[ningpp]

@Craigacp thank you , i will try it

[microsoft-github-policy-service]

Applying stale label due to no activity in 30 days