# How to Implement Custom Sampling Strategies in GPULlama3.java: A Complete Guide > Learn how to implement custom sampling strategies in GPULlama3.java by extending the Sampler interface. Override default behavior for unique text generation. - Repository: [Beehive lab/gpullama3.java](https://github.com/beehive-lab/gpullama3.java) - Tags: how-to-guide - Published: 2026-02-26 --- **Implement the `Sampler` functional interface from `org.beehive.gpullama3.inference.sampler` and pass your implementation to `Model.generateTokens()` to override the default sampling behavior.** GPULlama3.java is a GPU-accelerated inference engine for Llama 3 models. While it ships with built-in samplers like greedy argmax and nucleus sampling, you can implement custom sampling strategies by leveraging the `Sampler` interface. This guide shows you exactly how to create, integrate, and deploy custom samplers without modifying the core inference engine. ## Understanding the Sampler Architecture ### The Sampler Interface All token generation in GPULlama3 funnels through the `Sampler` functional interface defined in [`org.beehive.gpullama3.inference.sampler.Sampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/org.beehive.gpullama3.inference.sampler.Sampler.java). This interface declares a single method: ```java int sampleToken(Object logits) ``` The `logits` parameter is polymorphic: it arrives as either a `FloatTensor` (CPU path) or a `FloatArray` (TornadoVM GPU path). Your implementation must handle both types to function across execution backends. The method returns the integer token ID selected by your sampling strategy. ### Built-in Reference Implementations GPULlama3 provides three built-in implementations that serve as reference patterns: - **`TENSOR_ARGMAX`** – Greedy decoding that selects the highest logit - **`CategoricalSampler`** – Applies temperature scaling and softmax, then performs a weighted random draw - **`ToppSampler`** – Nucleus (top-p) sampling that filters logits by cumulative probability These classes reside in [`CategoricalSampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/CategoricalSampler.java) and [`ToppSampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/ToppSampler.java) respectively, demonstrating how to support both `FloatTensor` and `FloatArray` within the same sampler. ## Implementing a Custom Top-K Sampler ### Handling Both CPU and GPU Tensor Types Here is a complete implementation of a **Top-K** sampler that restricts sampling to the K most probable tokens. This example follows the architectural pattern used by the built-in samplers: ```java // TopKSampler.java package org.beehive.gpullama3.inference.sampler; import org.beehive.gpullama3.tensor.standard.FloatTensor; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import java.util.Random; /** * Keeps only the K highest‑probability tokens, then samples uniformly among them. */ public record TopKSampler(int k, Random rng) implements Sampler { @Override public int sampleToken(Object logits) { if (logits instanceof FloatTensor ft) { return sampleFromTensor(ft); } else if (logits instanceof FloatArray fa) { return sampleFromArray(fa); } throw new IllegalArgumentException( "Unsupported logits type: " + (logits == null ? "null" : logits.getClass())); } private int sampleFromTensor(FloatTensor t) { // 1. Find K highest indices int[] topIndices = new int[k]; float[] topValues = new float[k]; for (int i = 0; i < t.size(); i++) { float v = t.getFloat(i); // simple insertion sort for the K‑size window for (int j = k - 1; j >= 0; j--) { if (j == 0 || v > topValues[j - 1]) { if (j < k - 1) { topValues[j + 1] = topValues[j]; topIndices[j + 1] = topIndices[j]; } topValues[j] = v; topIndices[j] = i; break; } } } // 2. Sample uniformly from the K candidates return topIndices[rng.nextInt(k)]; } private int sampleFromArray(FloatArray a) { // Same logic as above but using FloatArray API int[] topIndices = new int[k]; float[] topValues = new float[k]; for (int i = 0; i < a.getSize(); i++) { float v = a.get(i); for (int j = k - 1; j >= 0; j--) { if (j == 0 || v > topValues[j - 1]) { if (j < k - 1) { topValues[j + 1] = topValues[j]; topIndices[j + 1] = topIndices[j]; } topValues[j] = v; topIndices[j] = i; break; } } } return topIndices[rng.nextInt(k)]; } } ``` ### Integrating Temperature and Softmax The built-in `CategoricalSampler` applies temperature scaling and softmax before sampling. If your custom sampler needs these transformations but you want to keep the sampling logic separate, use a wrapper function that preprocesses the logits: ```java public static Sampler withTemperatureAndSoftmax(Sampler inner, float temperature) { return logits -> { if (logits instanceof FloatTensor ft) { ft.divideInPlace(0, ft.size(), temperature); ft.softmaxInPlace(0, ft.size()); } else if (logits instanceof FloatArray fa) { FloatArrayUtils.divideInPlace(fa, 0, fa.getSize(), temperature); FloatArrayUtils.softmaxInPlace(fa, 0, fa.getSize()); } else { throw new IllegalArgumentException("Unsupported logits type"); } return inner.sampleToken(logits); }; } // Usage Sampler topK = new TopKSampler(100, new Random(123)); Sampler temperedTopK = withTemperatureAndSoftmax(topK, 0.7f); List output = model.generateTokens(state, 0, prompt, stop, 128, temperedTopK, true, null); ``` This pattern mirrors the logic found in `Sampler.selectSampler` (lines 99-119 of [`Sampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Sampler.java)) and allows you to compose behaviors without modifying your core sampler implementation. ## Wiring Custom Samplers into the Inference Pipeline ### Direct Injection via Model.generateTokens() The simplest integration path is to bypass the default factory entirely. The `Model` class (in [`Model.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Model.java)) provides overloads of `generateTokens` that accept a `Sampler` instance directly: ```java Model model = new Llama(...); TopKSampler customSampler = new TopKSampler(50, new Random(42)); List tokens = model.generateTokens( state, 0, prompt, Set.of(), 256, customSampler, // Your implementation here false, null ); ``` When you provide a custom sampler this way, the engine ignores the `temperature` and `topp` fields in `Options` because your sampler has full control over the sampling strategy. ### CLI Integration with Reflection To expose custom samplers to command-line users, extend the `Options` class (in [`Options.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Options.java)) to accept a class name, then use reflection to instantiate it in the CLI entry point ([`LlamaApp.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaApp.java) or [`LlamaTornadoCli.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaTornadoCli.java)): ```java // In Options.java – add a new field private final String customSamplerClass; // e.g. "org.beehive.gpullama3.inference.sampler.TopKSampler" // In LlamaApp.java – modify sampler creation: Sampler sampler; if (options.customSamplerClass() != null) { Class cls = Class.forName(options.customSamplerClass()); // Assume a constructor (int k, Random) – adapt as needed Constructor ctor = cls.getConstructor(int.class, Random.class); sampler = (Sampler) ctor.newInstance(options.topK(), new Random(options.seed())); } else { sampler = Sampler.createSampler(model, options); } ``` Now users can invoke your custom strategy without code changes: ```bash java -jar gpullama3.jar --model llama --custom-sampler org.beehive.gpullama3.inference.sampler.TopKSampler \ --topk 64 --seed 2024 ``` Because [`InferenceEngine.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/InferenceEngine.java) only requires the `sampleToken` method, no kernel recompilation or engine modification is necessary when adding new sampling algorithms. ## Key Source Files for Custom Sampling | File | Why it matters for custom sampling | |------|------------------------------------| | **[`Sampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Sampler.java)** – definition of the functional interface and the default factory (`selectSampler`). | Entry point for any sampler you write. | | **[`CategoricalSampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/CategoricalSampler.java)** – reference implementation for a standard categorical draw. | Shows how to support both `FloatTensor` and `FloatArray`. | | **[`ToppSampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/ToppSampler.java)** – reference for a more complex strategy (top-p). | Demonstrates extra bookkeeping (cumulative probability). | | **[`InferenceEngine.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/InferenceEngine.java)** – the loop that calls `sampler.sampleToken`. | Guarantees that any `Sampler` implementation will be used without further changes. | | **[`Model.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Model.java)** – high-level generate-tokens methods that accept a `Sampler`. | Provides the overloads you'll call from application code. | | **[`Options.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Options.java)** – holds CLI-level configuration. | Useful if you want to expose custom sampler selection to end-users. | | **[`LlamaApp.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaApp.java) / [`LlamaTornadoCli.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaTornadoCli.java)** – command-line entry points. | Places to hook the custom-sampler factory if you need a user-friendly flag. | ## Summary - **The `Sampler` interface** in `org.beehive.gpullama3.inference.sampler` is the sole integration point for custom sampling strategies in GPULlama3.java. - **Implement `sampleToken(Object logits)`** to handle both `FloatTensor` (CPU) and `FloatArray` (GPU/TornadoVM) tensor types. - **Inject custom samplers** directly via `Model.generateTokens()` overloads or wire them into the CLI by extending `Options` and using reflection in [`LlamaApp.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaApp.java). - **Compose behaviors** by wrapping samplers with temperature scaling and softmax using the pattern found in `Sampler.selectSampler`. - **No engine modifications required** – [`InferenceEngine.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/InferenceEngine.java) treats all samplers uniformly, so you can experiment with Top-K, Beam Search, or penalized sampling without touching the GPU kernels. ## Frequently Asked Questions ### How do I ensure my custom sampler works with both CPU and GPU inference modes? Your `sampleToken` implementation must check the runtime type of the `logits` parameter using `instanceof`. If it is a `FloatTensor`, use the CPU tensor API (`getFloat`, `size`). If it is a `FloatArray` (from TornadoVM), use the GPU array API (`get`, `getSize`). The built-in [`CategoricalSampler.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/CategoricalSampler.java) demonstrates this dual-path pattern. ### Can I combine temperature scaling with my custom sampling logic? Yes. You can either apply temperature and softmax inside your sampler's `sampleToken` method before your selection logic, or use a wrapper function that preprocesses the logits and delegates to your inner sampler. The wrapper approach mirrors the implementation in `Sampler.selectSampler` and keeps your core sampling logic clean. ### Where should I register a custom sampler to use it from the command line? Extend [`Options.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/Options.java) to add a new field (e.g., `customSamplerClass`) and parse it from a new CLI flag like `--custom-sampler`. Then modify [`LlamaApp.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaApp.java) or [`LlamaTornadoCli.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/LlamaTornadoCli.java) to instantiate your class via reflection and pass it to `Model.generateTokens()`. This avoids hard-coding sampler logic while exposing it to end users. ### Do I need to recompile the GPU kernels when adding a new sampler? No. The [`InferenceEngine.java`](https://github.com/beehive-lab/gpullama3.java/blob/main/InferenceEngine.java) token generation loop is completely decoupled from sampling implementations. It only requires the `Sampler` interface's `sampleToken` method. Because sampling happens after the forward pass completes and logits are materialized in host memory, you can add, modify, or swap samplers without touching TornadoVM kernel code or recompiling GPU binaries.