Add transformer-gpu state manager that keeps transformer states on the GPU by larissakl · Pull Request #282 · rwth-i6/rasr

larissakl · 2026-10-09T10:52:42Z

This PR adds a transformer state manager that keeps the states, e.g. the KV cache, in GPU memory. It behaves like the transformer state manager, including the parameters max-history and always-include-first-token-state. The search, label scorers and LMs are unchanged.

Each state is stored as a DeviceVector, a CompressedVector whose data lives in a block of a GPU memory pool. This keeps the HistoryState interface unchanged. mergeStates assembles the batched state inputs with a small CUDA kernel and passes them to the session as device tensors. With #281, the session then keeps the state outputs on the GPU, and splitStates copies them into new pool blocks. Outputs that arrive in host memory, e.g. from a model without state inputs, are uploaded first. This removes the host-side merge and split and the copies of all states between host and GPU in every search step.
The pool only frees its memory when the state manager is destroyed, freed blocks are reused. GPU memory therefore follows the peak number of states that are alive at the same time. For label scorers, this is bounded by their caches (max-cached-score-vectors) and for the OnnxRecurrentLanguageModel by the number of live LM histories (as with the transformer state manager in host memory).

Depends on #281.

larissakl · 2026-10-27T12:07:05Z

My proposal to make RASR use the GPU state manager automatically whenever the session runs with the CUDA execution provider, such that setting state-manager.type = transformer in the config is enough:

In Onnx::Session (in PR ONNX: support tensors in CUDA device memory, IO binding and CUDA provider options #281), add a getter like useCudaExecutionProvider(), so others can "ask" whether the session runs on CUDA
Add an extra argument to Nn::Module::createStateManager() saying whether the session uses CUDA and construct either the TransformerGpuStateManager or the OnnxTransformerStateManager accordingly
In StateManagedOnnxLabelScorer, the state manager then needs to be created in the constructor body after the model is loaded
In OnnxRecurrentLanguageModel, the state manager can be passed to the base class Lm::RecurrentLanguageModel in its constructor and e.g. a protected setter is added there

What do you think about that? @SimBe195 @curufinwe

SimBe195 · 2026-11-02T11:00:19Z

+    }
+}
+
+std::vector<float const*> StateCopier::deviceData(std::vector<Value> const& tensors) {


The upload of host tensors into stagingBuffer_ never runs because mergeStates always feeds device tensors and both users (StateManagedOnnxLabelScorer and OnnxRecurrentLanguageModel) run the session with default output locations. So Session::run already puts the state outputs on the device. Maybe remove the staging code and stagingBuffer_ and instead just require that the tensors are on device.

You're right that at the moment the staging code is never used and I will therefore remove it. However, the case does occur for label scorers with an initializer model that has no state inputs, e.g. the prefix LM prefill in the PrefixSpeechLmOnnxLabelScorer. We will need to run the initializer there with device output locations for the state outputs when the state manager keeps its states on device.

SimBe195 · 2026-11-02T11:16:07Z

+}
+
+// Abort with a message if `status` (a cudaError_t) is an error
+void checkCudaStatus(int status, char const* what);


This header is anyway only compiled with MODULE_CUDA active, so you can include <cuda_runtime.h> and use cudaError_t directly. This also allows getting rid of the checkCuda forwarding wrapper in DeviceVector.cc and the casts.

SimBe195 · 2026-11-02T11:17:44Z

+ * For each task copy `stepSize` floats between block and tensor. Block element i is tensor element
+ * `blockOffsets[i / blockSize] + i % blockSize`. `tasks`/`blockOffsets` are device pointers, returns a cudaError_t.
+ */
+int launchBlockCopy(BlockCopyTask const* tasks,


Here you could also include <cuda_runtime.h> and use cudaError_t directly as return type instead of an int (similar as checkCudaStatus in DeviceVector.hh).

SimBe195 · 2026-11-02T11:22:37Z

+                    size_t               blockSize,
+                    size_t               stepSize,
+                    BlockCopyDirection   direction,
+                    void*                stream);


This can be a cudaStream_t instead of void*. For Session::sharedCudaStream() it makes sense to return void* since that must also be compilable without MODULE_CUDA but the cast back to cudaStream_t is repeated in three local sharedStream helper functions (DeviceVector.cc:33, StateCopier.cc:28 and TransformerGpuStateManger.cc:87). This could be unified via one inline sharedStream helper in DeviceVector.hh.

SimBe195 · 2026-11-02T11:36:35Z

+                if (not alwaysIncludeFirstTokenState_ or p != 0ul) {
+                    idx += prefix_offset + p;
+                }
+                auto const* state = dynamic_cast<DeviceVector const*>(prefix_states[idx]->at(v).get());


This dynamic cast runs in the innermost loop which can come out to hundreds of thousands of casts per step with large history or batch size. Prefer static_cast instead. If you want the check, maybe do it only on the first state.

SimBe195 · 2026-11-02T12:18:23Z

+        return cudaSuccess;
+    }
+    bool const toTensor = direction == BlockCopyDirection::BlockToTensor;
+    for (size_t start = 0ul; start < numTasks; start += kMaxGridSize) {


This loop is never going to run multiple iterations. That would mean that a single variable has over 2^30 tasks which is unrealistic. You can just get rid of the loop and do a single launch.

SimBe195 · 2026-11-02T12:20:45Z

+
+float DeviceVector::get(size_t pos) const {
+    require_lt(pos, size_);
+    return toHost()[pos];


This copies the full vector to host just to read one float. Do a single-element cudaMemcpy from data_ + pos instead.

SimBe195 · 2026-11-02T12:27:15Z

+        StateVariables const&                     vars,
+        std::vector<size_t>&                      suffix_lengths,
+        std::vector<Value> const&                 state_tensors,
+        Nn::CompressedVectorFactory<float> const& vector_factory) {


vector_factory is not used, so you could leave it unnamed here (like in initialState).

SimBe195 · 2026-11-02T12:28:37Z

+        auto&        tasks = batches.back().tasks;
+        float const* base  = bases[v];
+
+        size_t const timeAxis   = timeAxisOf(vars[v]);


timeAxisOf(vars[v]) is computed twice per variable (first time in line 233). It could be saved in StepLayout instead.

SimBe195 · 2026-11-02T12:35:43Z

+#ifdef MODULE_CUDA
+            return std::unique_ptr<AbstractStateManager<Onnx::Value, Onnx::OnnxStateVariable>>(
+                    new Onnx::TransformerGpuStateManager(config));
+#else
+            Core::Application::us()->criticalError("State manager type transformer-gpu requires RASR compiled with MODULE_CUDA");
+            return nullptr;
+#endif


Instead of the criticalError + nullptr return here, you can conditionally only register "transformer-gpu" if MODULE_CUDA is enabled. Then the value is rejected by the choice parser already. But this anyway should go away when we adopt the backend auto-selection.

SimBe195 · 2026-11-02T12:39:55Z

The code in this PR contains a wild mix of snake_case and camelCase naming, probably because some parts were copied from the CPU state manager and other parts written from scratch. It would be better to make it consistently camelCase since that's the usual convention in RASR.

Also, there are several usages of the term "block" in different contexts which mean different things. It can be a pool block, a contiguous span of data in a tensor or a CUDA thread block. This can make it confusing to follow. I'd prefer some more distinctive names.

SimBe195 · 2026-11-02T12:46:08Z

+    HistoryState result;
+    result.reserve(vars.size());
+    for (size_t v = 0ul; v < vars.size(); ++v) {
+        result.emplace_back(new DeviceVector(nullptr, 0ul));


Please prefer smart pointers instead of raw new. Same thing also further down on line 268.

SimBe195 · 2026-11-02T12:48:10Z

In RecurrentLanguageModel if dump-inputs is enabled it will try to save device tensors which will lead to a crash.

SimBe195 · 2026-11-02T13:04:38Z

My proposal to make RASR use the GPU state manager automatically whenever the session runs with the CUDA execution provider, such that setting state-manager.type = transformer in the config is enough:
...
What do you think about that? @SimBe195 @curufinwe

I like the approach of having a getter function in Onnx::Session. This can even be static, so just static bool usesCudaExecutionProvider(Core::Configuration const& config) which reads the value from config (and from whether MODULE_CUDA is active). This way, the decision can be made before creating a concrete session object. We could also think about using an execution provider type enum instead of just a bool for CUDA to account for different providers in the future.

I would also keep explicit choice values which force the state manager For example one plain "transformer" which auto-selects and "transformer-cpu"/"transformer-cuda" which force a selection. And make sure that "-16bit" and "-8bit" variants don't switch since they are not currently supported on GPU. In fact, maybe a warning should be logged if state-compression is configured since that gets ignored by the GPU state manager.

larissakl added 2 commits October 9, 2026 12:46

Add transformer-gpu state manager

675eb5d

Move device helpers into shared classes

1487523

larissakl mentioned this pull request Oct 13, 2026

Add lstm-gpu state manager #286

Open

larissakl marked this pull request as ready for review October 16, 2026 12:13

SimBe195 requested changes Nov 2, 2026

View reviewed changes

SimBe195 reviewed Nov 2, 2026

View reviewed changes

larissakl added 4 commits November 2, 2026 14:23

Merge branch 'onnx-device-memory' into transformer-gpu-state-manager

42de7f8

Implement suggestions from code review

7a63990

Merge all state variables in one kernel launch without memset

8c795ea

Choose CPU or CUDA version by the session's execution provider

dc266f0

larissakl requested review from SimBe195 and curufinwe November 4, 2026 16:08

Conversation

larissakl commented Oct 9, 2026

Uh oh!

larissakl commented Oct 27, 2026

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!

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!

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!

SimBe195 commented Nov 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SimBe195 commented Nov 2, 2026

Uh oh!

SimBe195 commented Nov 2, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants