Skip to content

[tornadovm] Replace the dummy forward pass in forceCopyInReadOnlyData() with TornadoExecutionPlan.transferToDevice() - #148

Open
mikepapadim wants to merge 1 commit into
mainfrom
feat/plan-transfer-to-device
Open

[tornadovm] Replace the dummy forward pass in forceCopyInReadOnlyData() with TornadoExecutionPlan.transferToDevice()#148
mikepapadim wants to merge 1 commit into
mainfrom
feat/plan-transfer-to-device

Conversation

@mikepapadim

Copy link
Copy Markdown
Member

Depends on beehive-lab/TornadoVM#1036, which adds TornadoExecutionPlan.transferToDevice(). This PR cannot build until that lands and ships in a TornadoVM release; the numbers below were measured against a local TornadoVM build containing it.

What this changes

forceCopyInReadOnlyData() exists to get the model weights onto the GPU. It did that by running a full forward pass on zeroed state — the activation graph, then every layer graph, then the logits graph — purely for the copy-in side effect, because there was no way to ask for a transfer without running something:

public void forceCopyInReadOnlyData() {
    state.wrapX.clear();
    state.positionHolder.init(0);
    executionPlan.withGraph(taskGraphLayout.activationIdx())...execute();
    for (int layer = 0; layer < config.numberOfLayers(); layer++) {
        executionPlan.withGraph(taskGraphLayout.layerIdx(layer))...execute();
    }
    executionPlan.withGraph(taskGraphLayout.logitsIdx())...execute();
}

TornadoVM now has the operation this method was faking, so the whole thing becomes:

executionPlan.transferToDevice();

Applied to all three plan variants — TornadoVMMasterPlanSingleToken, TornadoVMMasterPlanPrefillDecode, TornadoVMMasterPlanBatchPrefillDecode.

With CUDA graphs enabled the dummy pass is also the capture, so it is kept as-is in that case. The change is a guard clause; the --cuda-graphs path is untouched, and the benchmarks below confirm that.

Benefits

1. The copy-in phase gets faster. It no longer runs N+2 graphs' worth of kernels it does not want. Llama-3.2-3B-Instruct-Q8_0, RTX 4090, JDK 21, CUDA backend, 3 runs each:

copy-in phase end-to-end wall
main (dummy forward pass) 1035.9 / 1046.0 / 1036.7 ms — avg 1039.5 3.74 / 3.85 / 3.79 s
this PR (transferToDevice()) 926.9 / 929.2 / 933.3 ms — avg 929.8 3.69 / 3.78 / 3.78 s

~10.6% off the copy-in phase, and it is the more honest kind of saving: the phase is dominated by the upload itself (host-register + PCIe), so removing the kernels removes essentially all of what was not the transfer.

2. Inference performance is unchanged — which is the point. The change only moves when data arrives, not how the model runs. llama-bench-style, -p 128 -n 128 -r 5 --no-warmup:

config pp128 t/s tg128 t/s
default main 104.84 ± 0.90 105.24 ± 0.18
default this PR 103.17 ± 4.47 105.61 ± 0.06
--with-prefill-decode main 104.46 ± 1.04 105.29 ± 0.12
--with-prefill-decode this PR 103.46 ± 4.74 105.97 ± 0.10
--cuda-graphs (path untouched) main 133.00 ± 1.49 135.06 ± 0.17
--cuda-graphs (path untouched) this PR 132.88 ± 1.66 135.08 ± 0.28

Token generation is a tie in every configuration, and the --cuda-graphs numbers are identical to three significant figures, as they should be for a code path this PR does not touch.

3. The method now says what it means. A method whose name is "copy in read-only data" is a copy-in, not a forward pass on zeroed state. It also stops depending on the forward pass being safe to run on garbage input — a constraint nobody was tracking, which quietly rules out anything that would divide by a zeroed norm or index from a zeroed position.

One honest caveat

For a short single-shot generation the reported tok/s drops, because the dummy pass used to install and first-launch every kernel, and the API path leaves that for the first real token. It is the same work moving, not new work:

generation length main this PR
30 tokens 110.8 t/s 81.6 t/s
~400 tokens 69.69 t/s 69.23 t/s (−0.7%)

By a few hundred tokens it is parity, and end-to-end wall clock is the same or marginally better at every length measured (the startup saving offsets the first-launch cost). Nothing gets slower overall — the cost just shows up on a different counter. Fixing that properly means a warm-up that installs code without transferring, which is worth doing separately and is easier to build now that the transfer exists as its own operation.

Testing

  • mvn test — 16/16 pass.
  • Correct output on all three plan variants (default, --with-prefill-decode, --cuda-graphs) with Llama-3.2-3B-Instruct-Q8_0.
  • Benchmarks above: --bench -p 128 -n 128 -r 5 --no-warmup, run against main and this branch built identically, same machine, same TornadoVM build.
  • --batch-prefill-size 128 fails identically on main and on this branch in the dev build used here (a TornadoVM codegen error unrelated to this change), so that path is unverified either way rather than regressed.

Environment: RTX 4090, Ubuntu 24.04, JDK 21, TornadoVM CUDA backend.

Copilot AI lite review requested due to automatic review settings August 26, 2026 13:39

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants