diff --git a/BATCHED_DECODE.md b/BATCHED_DECODE.md index cd3a0cd3..fa8628d5 100644 --- a/BATCHED_DECODE.md +++ b/BATCHED_DECODE.md @@ -171,6 +171,15 @@ capped lower): | Llama-3.2-1B | 16 · 2048 · 128256 | 101 tok/s | 128 | 4175 tok/s | 41× | | Qwen3-1.7B | 28 · 2048 · 151936 | 48 tok/s | 64 | 1433 tok/s | 30× | | Qwen3-4B | 36 · 2560 · 151936 | 39 tok/s | 32 | 405 tok/s | 10× | +| Mistral-7B-v0.3 | 32 · 4096 · 32768 | 27 tok/s | 32 | 331 tok/s | 12× | + +**Supported models: LLaMA, Qwen3, Mistral (FP16).** Mistral runs on the LLaMA decode path with +no new kernels — the only generalizations were typing the layer graph to the base +`Configuration` and parameterizing the RoPE theta (`config.ropeTheta()`, e.g. Mistral 1e6). Any +LLaMA-family model (RMSNorm + SwiGLU + GQA + RoPE, no QK-norm, no sliding window) drops in the +same way; the full serving stack (continuous + paging + prefix + on-device sampling) is +model-agnostic and applies unchanged (Mistral-7B: 128-request continuous+paged+prefix → +prefix-consistent, ~11.6× less KV, 82.8% fewer prefill tokens). All bit-exact vs the single-stream greedy reference (`all B streams identical: true`) and coherent. diff --git a/src/main/java/org/beehive/gpullama3/bench/BatchedDecodeEngine.java b/src/main/java/org/beehive/gpullama3/bench/BatchedDecodeEngine.java index 46ef3989..0650e6ca 100644 --- a/src/main/java/org/beehive/gpullama3/bench/BatchedDecodeEngine.java +++ b/src/main/java/org/beehive/gpullama3/bench/BatchedDecodeEngine.java @@ -158,10 +158,10 @@ public static void main(String[] args) throws Exception { } else { LlamaFP16LayersBatchDecodeMMA l = paged ? new LlamaFP16LayersBatchDecodeMMA((LlamaState) state, (LlamaTornadoWeights) weights, - (LlamaConfiguration) config, B, decodeCtx, keyCacheBatch, valueCacheBatch, seqPositions, + config, B, decodeCtx, keyCacheBatch, valueCacheBatch, seqPositions, blockTable, blockSize, maxBlocksPerSlot) : new LlamaFP16LayersBatchDecodeMMA((LlamaState) state, (LlamaTornadoWeights) weights, - (LlamaConfiguration) config, B, decodeCtx, keyCacheBatch, valueCacheBatch, seqPositions); + config, B, decodeCtx, keyCacheBatch, valueCacheBatch, seqPositions); layerITGs = l.getLayerImmutableTaskGraphs(); lastLayerId = l.getLastLayerTaskGraphID(); updateLayerSched = l::updateGridScheduler; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 50e744d0..7d568df9 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -1658,7 +1658,7 @@ public static void batchedDecodeRopeWithKVCachePacked(KernelContext context, FloatArray wrapValueCache, int kvDim, int headSize, int layerIndex, int numLayers, - int contextLength, int dim) { + int contextLength, int dim, float ropeTheta) { int globalIdx = context.globalIdx; int halfDim = dim / 2; int batchIdx = globalIdx / halfDim; @@ -1673,7 +1673,7 @@ public static void batchedDecodeRopeWithKVCachePacked(KernelContext context, if (i + 1 < dim) { int head_dim = i % headSize; - float freq = 1.0f / TornadoMath.pow(50000.0f, head_dim / (float) headSize); + float freq = 1.0f / TornadoMath.pow(ropeTheta, head_dim / (float) headSize); float val = pos * freq; float fcr = TornadoMath.cos(val); float fci = TornadoMath.sin(val); @@ -1830,7 +1830,7 @@ public static void batchedDecodePagedRopeWithKVCachePacked(KernelContext context FloatArray valuePool, int kvDim, int headSize, int layerIndex, int numLayers, - int blockCfg, int dim) { + int blockCfg, int dim, float ropeTheta) { int blockSize = blockCfg & 0xFFFF; int maxBlocksPerSlot = blockCfg >>> 16; int globalIdx = context.globalIdx; @@ -1847,7 +1847,7 @@ public static void batchedDecodePagedRopeWithKVCachePacked(KernelContext context if (i + 1 < dim) { int head_dim = i % headSize; - float freq = 1.0f / TornadoMath.pow(50000.0f, head_dim / (float) headSize); + float freq = 1.0f / TornadoMath.pow(ropeTheta, head_dim / (float) headSize); float val = pos * freq; float fcr = TornadoMath.cos(val); float fci = TornadoMath.sin(val); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16LayersBatchDecodeMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16LayersBatchDecodeMMA.java index 811155f8..3691e119 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16LayersBatchDecodeMMA.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16LayersBatchDecodeMMA.java @@ -2,7 +2,7 @@ import org.beehive.gpullama3.inference.state.LlamaState; import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; -import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.model.Configuration; import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; import uk.ac.manchester.tornado.api.GridScheduler; @@ -42,7 +42,7 @@ public class LlamaFP16LayersBatchDecodeMMA { private final LlamaState state; private final LlamaTornadoWeights weights; - private final LlamaConfiguration config; + private final Configuration config; private final KernelContext context = new KernelContext(); private final int batchSize; private final int paddedBatch; @@ -59,7 +59,7 @@ public class LlamaFP16LayersBatchDecodeMMA { private String lastLayerTaskGraphID; public LlamaFP16LayersBatchDecodeMMA(LlamaState state, LlamaTornadoWeights weights, - LlamaConfiguration config, int batchSize, int decodeCtx, + Configuration config, int batchSize, int decodeCtx, FloatArray keyCacheBatch, FloatArray valueCacheBatch, IntArray seqPositions) { this(state, weights, config, batchSize, decodeCtx, keyCacheBatch, valueCacheBatch, seqPositions, @@ -68,7 +68,7 @@ public LlamaFP16LayersBatchDecodeMMA(LlamaState state, LlamaTornadoWeights weigh /** Paged constructor: keyCacheBatch/valueCacheBatch are the shared block pools. */ public LlamaFP16LayersBatchDecodeMMA(LlamaState state, LlamaTornadoWeights weights, - LlamaConfiguration config, int batchSize, int decodeCtx, + Configuration config, int batchSize, int decodeCtx, FloatArray keyCacheBatch, FloatArray valueCacheBatch, IntArray seqPositions, IntArray blockTable, int blockSize, int maxBlocksPerSlot) { @@ -172,7 +172,7 @@ private TaskGraph createLayerTaskGraph(int layerIndex) { context, seqPositions, blockTable, state.qkvResultBatch, keyCacheBatch, valueCacheBatch, kvDim, config.headSize(), layerIndex, config.numberOfLayers(), - blockCfg, dim); + blockCfg, dim, config.ropeTheta()); g.task("batch_attention", TransformerBatchPrefillKernels::batchedDecodePagedAttentionFP16Out, @@ -188,7 +188,7 @@ private TaskGraph createLayerTaskGraph(int layerIndex) { context, seqPositions, state.qkvResultBatch, keyCacheBatch, valueCacheBatch, - kvDim, config.headSize(), layerIndex, config.numberOfLayers(), decodeCtx, dim); + kvDim, config.headSize(), layerIndex, config.numberOfLayers(), decodeCtx, dim, config.ropeTheta()); g.task("batch_attention", TransformerBatchPrefillKernels::batchedDecodeAttentionFP16Out,