Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes by ammbra · Pull Request #1085 · openjdk/babylon

ammbra · 2026-07-03T08:38:19Z

I apologize for such a big PR, but most of its changes are result of regeneration for native bindings, operators, prototufs etc. So, this PR contains many changes that relate to upgrading the ONNX runtime library (1.26.0) and its native GenAI (0.14.0) counterpart:

The changes from oracle.code.onnx.foreign are generated via the cr-examples/onnx/opgen/setup.sh script and its helper symbols files to generate only the absolute necessary contents.
The operators and protobufs have been generated via their respective Java generators.
OnnxProtoBuilder needed to support also shapes because I kept getting this native error once rebuilding Babylon:

:bomb: Program crashed: Bad pointer dereference at 0x0000000000000008

Thread 37 "Java: oracle.code.onnx.llm.LlamaDemo.main()" crashed:

 0 0x000000012aa83f8c Generators::Model::IsPruned() const + 60 in libonnxruntime-genai.dylib
 1 0x000000012aa76cf4 Generators::Logits::Logits(Generators::State&) + 268 in libonnxruntime-genai.dylib
 2 0x000000012aa41848 Generators::DecoderOnly_State::DecoderOnly_State(Generators::DecoderOnly_Model const&, Generators::DeviceSpan<int>, Generators::GeneratorParams const&) + 88 in libonnxruntime-genai.dylib
 3 0x000000012aa4178c Generators::DecoderOnly_Model::CreateState(Generators::DeviceSpan<int>, Generators::GeneratorParams const&) const + 104 in libonnxruntime-genai.dylib
 4 0x000000012aa36e90 Generators::Generator::Generator(Generators::Model const&, Generators::GeneratorParams const&) + 432 in libonnxruntime-genai.dylib
 5 0x000000012aa36c90 Generators::CreateGenerator(Generators::Model const&, Generators::GeneratorParams const&) + 52 in libonnxruntime-genai.dylib
 6 0x000000012aade410 OgaCreateGenerator + 28 in libonnxruntime-genai.dylib
 7 0x000000013f6a2b00
 8 0x000000013f53cbc8
 9 0x000000013f53cbc8
10 0x000000013f53cbc8
11 0x000000013f53ca70
12 0x000000013f53ca70
13 0x000000013807e010
14 0x00007fffffffffff

I managed to identify that the exported model was the source for such an error (both local debug and by also comparing in Netron). To my understanding, Generators::Model::IsPruned() const method checks if the underlying model has been compressed using model pruning. Pruning seems to reduce model size and latency by zeroing out or permanently removing unnecessary weights and parameters, but I think in this case the generated model was aggressively pruned from the beginning of traversing the graph, thus resulting in the error above. The aggressive pruning was performed due the lack of shapes being there. This error occurred only for LLMs, the rest of the ML models being unaffected.

I also adjusted the license headers and cleaned a few leftovers in variable or patterns not really used by the library.
I updated the LlamaDemo to be more interactive.
To test my changes beyond the existing demo, I created a new one. The AllMiniLML6V2EmbeddingModel is inspired by https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX and by running the OnnxLift over its default (model.onnx). As a result, I could obtain a quite good model to start. Then I cleaned it locally based on how the official config.json looked like. The BertTokenizer is just a helper for working with the tokenizer and its constants are from tokenizer_config.json and config.json , while the EmbeddingDemo is just a Java version of the code from the landing page of the model.

I confirm that I make this contribution in accordance with the OpenJDK Interim AI Policy.

Progress

Change must not contain extraneous whitespace

Reviewing

Using git

Checkout this PR locally:
$ git fetch https://git.openjdk.org/babylon.git pull/1085/head:pull/1085
$ git checkout pull/1085

Update a local copy of the PR:
$ git checkout pull/1085
$ git pull https://git.openjdk.org/babylon.git pull/1085/head

Using Skara CLI tools

Checkout this PR locally:
$ git pr checkout 1085

View PR using the GUI difftool:
$ git pr show -t 1085

Using diff file

Download this PR as a diff file:
https://git.openjdk.org/babylon/pull/1085.diff

Using Webrev

Link to Webrev Comment

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

bridgekeeper · 2026-07-03T08:40:31Z

👋 Welcome back amihalceanu! A progress list of the required criteria for merging this PR into code-reflection will be added to the body of your pull request. There are additional pull request commands available for use with this pull request.

openjdk · 2026-07-03T08:40:52Z

@ammbra This change now passes all automated pre-integration checks.

ℹ️ This project also has non-automated pre-integration requirements. Please see the file CONTRIBUTING.md for details.

After integration, the commit message for the final commit will be:

Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes

You can use pull request commands such as /summary, /contributor and /issue to adjust it as needed.

At the time when this comment was updated there had been 2 new commits pushed to the code-reflection branch:

8e278d7: Align Op.ofElement with ReflectMethods::visitMethodDef
5b95d66: Turn unsupported error into warnings

Please see this link for an up-to-date comparison between the source branch of this pull request and the code-reflection branch.
As there are no conflicts, your changes will automatically be rebased on top of these commits when integrating. If you prefer to avoid this automatic rebasing, please check the documentation for the /integrate command for further details.

As you do not have Committer status in this project an existing Committer must agree to sponsor your change.

➡️ To flag this PR as ready for integration with the above commit message, type /integrate in a new comment. (Afterwards, your sponsor types /sponsor in a new comment to perform the integration).

mlbridge · 2026-07-03T08:46:02Z

Webrevs

asotona · 2026-07-03T13:00:03Z

    }

-    static TensorProto tensorProto(String name, oracle.code.onnx.Tensor tensor, Function<Tensor, ExternalTensorDataInfo> tensorDataExternalizer) {
+    private static List<Object> determineLlmShape(String name) {


Could we please avoid hardcoding tensor dimensions in OnnxProtoBuilder? This would work only for one concrete model.
If the shape information is missing from the ONNX maybe we can pass this metadata from the GenAI/demo layer now or find better solution in a follow-up PR.

As a temporary solution we can use for example an interface with concrete shapes:

interface OnnxShapeHints { Map<String, double[]> onnxShapeHints(); }

A rock solid solution would be to implement shape inference.

I added a factory for shape inference. Can you please take a look at the new way to handle it?

asotona · 2026-07-03T13:08:48Z

-                        <|start_header_id|>user<|end_header_id|>Hello, tell me a joke.<|eot_id|>
-                        <|start_header_id|>assistant<|end_header_id|>
-                        """, System.out::print);
+                Reader inreader = new InputStreamReader(System.in);


Why did the Llama demo change?

Is more interactive, you can have a chat conversation in command line. But if you believe would be best to keep it simple, I can revert it.

It might be interactive, however it requires user to wrap simple text prompt into Llama-specific form manually.

Oh, I read more here and I understand what you mean now. While the easiest way would be to go back, would prompt formatting here be ok to keep it interactive? One drawback I find to prompt formatting is that sanitising the input would be needed. 🙏

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona

I’m still working through all the proposal, so please correct me if I misunderstand something.

There is clearly a need for named inputs and outputs in the graph. This is already handled inside OnnxTransformer, where ModuleAndInitializers holds the name map, allowing method parameter names and return record component names to be forwarded to the ONNX model.

The proposal seems to go in the opposite direction: to apply names parsed from the config file. I’m a bit concerned about this, because it moves us away from the idea that Java should remain the source of truth for the model.

There is also a clear need to provide shape information for some ONNX operations. This is new and would require passing additional information into the protobuf-building process. There are two parts to this issue:

identify the op input that requires an explicit shape
provide the actual shape value

If I understand the proposal correctly, the tensor would be identified by its name coming from the config, and the shape value would be represented as a list of objects, mixing literal values with references to the configuration file or even references to internal runtime values. My concern is that none of this information comes directly from the Java source.

If the shape definitions are primarily needed for method parameters, then the mapping key should probably be the tensor method parameter, or more generally some code-reflection value. If we had code reflection support for value attributes, or if we had support for parameter annotations, this could be handled more cleanly. Since we are not there yet, some workaround may be necessary. For example using parameter names as keys for the tensor shape resolver.

For the second part (the shape itself) I think it should still be Java-sourced. In practice, that would mean resolving it to an array of longs. Otherwise we risk introducing unclear references to “magic” constants coming from external configuration.

When I refer to shape inference, I mean that if shape information is required deeper in the model tree, it should ideally be inferred from the parameters and weights through the operations in the tree. Since that is a more complex task, we could temporarily rely on a Java variable-name-based mapping, even for values deeper in the model. Another option would be to identify the ops that require shape information and add synthetic shape attributes to them, in addition to their specs-generated attributes, so the shape constants are still passed explicitly from Java.

asotona · 2026-07-07T07:47:32Z

-                        <|start_header_id|>user<|end_header_id|>Hello, tell me a joke.<|eot_id|>
-                        <|start_header_id|>assistant<|end_header_id|>
-                        """, System.out::print);
+                Reader inreader = new InputStreamReader(System.in);


It might be interactive, however it requires user to wrap simple text prompt into Llama-specific form manually.

asotona · 2026-07-07T11:39:37Z

+            case String shapeName when (shapeName.matches("past(Key|Value)\\.\\d+")) ->  List.of("batch_size", keyValueHeads, "past_sequence_length", headSize);
+            case String shapeName when (shapeName.matches("present(Key|Value)\\.\\d+")) -> List.of("batch_size", keyValueHeads, "total_sequence_length", headSize);
+            case String _ -> null;
+        };


This should break into if/else blocks, pattern matching does not make sense here.

asotona · 2026-07-07T11:40:43Z

+
+import java.util.List;
+
+public record NonTensorShapeResolver() implements TensorShapeResolver {


Why there is a custom resolver resolving to nothing?

Upgrade ONNX runtime, gen ai and add bert model.

ec7ce5f

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

ammbra mentioned this pull request Jul 3, 2026

Upgrade ONNX runtime, ONNX Gen AI runtime, add a new model to test changes #1075

Closed

2 tasks

openjdk Bot added ready Pull request is ready to be integrated rfr Pull request is ready for review labels Jul 3, 2026

asotona reviewed Jul 3, 2026

View reviewed changes

ammbra added 4 commits July 4, 2026 19:11

Add a refactoring to how shapes are interpreted.

5ee2c38

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Move constant to OnnxProtoBuilder

33a1227

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Adjust package name upon refactoring

cc2308f

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

Replace Copyright with proper years.

48344f3

Signed-off-by: Ana-Maria Mihalceanu <ana-maria.m.mihalceanu@oracle.com>

asotona reviewed Jul 7, 2026

View reviewed changes


		import java.util.List;

		public record NonTensorShapeResolver() implements TensorShapeResolver {

Uh oh!

Conversation

ammbra commented Jul 3, 2026 • edited by openjdk Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Progress

Reviewing

Uh oh!

bridgekeeper Bot commented Jul 3, 2026

Uh oh!

openjdk Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mlbridge Bot commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Webrevs

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

asotona left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Milestone

Development

Uh oh!

2 participants

ammbra commented Jul 3, 2026 •

edited by openjdk Bot

Loading

openjdk Bot commented Jul 3, 2026 •

edited

Loading

mlbridge Bot commented Jul 3, 2026 •

edited

Loading

asotona left a comment •

edited

Loading