From e23efe7a0441346da360b069faf6130eb16ce6ff Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Tue, 14 Jul 2026 09:51:12 +0100 Subject: [PATCH] Hybrid-library logits matvec experiment (llama.logitsLib=jit|gemmEx|lt) + findings: cuBLAS/cuBLASLt at parity with JIT on n=1 decode (bandwidth wall, launch-bound regime); CUDA graphs remain the single-token lever --- HYBRID_DECODE_LIBS.md | 101 ++++++++++++++++++ llama-tornado | 2 +- pom.xml | 5 + .../layers/type/fp16/LogitsFP16Layer.java | 76 ++++++++++--- 4 files changed, 170 insertions(+), 14 deletions(-) create mode 100644 HYBRID_DECODE_LIBS.md diff --git a/HYBRID_DECODE_LIBS.md b/HYBRID_DECODE_LIBS.md new file mode 100644 index 00000000..9d1f848c --- /dev/null +++ b/HYBRID_DECODE_LIBS.md @@ -0,0 +1,101 @@ +# Hybrid CUDA libraries for single-token decode — findings + +Can the TornadoVM hybrid library API (cuBLAS / cuBLASLt / cuDNN / cuFFT library tasks +riding the task-graph pipeline) speed up **single-request, single-token** decode? +Measured on the CUDA backend, RTX 4090, TornadoVM 5.0.1-jdk21-dev, JDK 21; +Llama-3.2-1B FP16 and Mistral-7B FP16 (both verified coherent on every change). + +**TL;DR: no.** The n=1 decode matvecs are memory-bound and the JIT matvec already runs +at ~95% of peak DRAM bandwidth — cuBLAS and cuBLASLt land on the *same* bandwidth wall +(within 0.6% per-kernel). Worse, decode is **launch-bound** (GPU busy only ~33–38% of +wall), and library tasks de-fuse TornadoVM's fused kernels, *adding* launches. The +single-token levers are CUDA graphs (+8–12%, `--cuda-graphs`) and kernel fusion — both +already in the code. Hybrid GEMM libraries pay off from n≳16: batch prefill (PR #127) +and batched decode (PR #129), where the same projections become compute-bound +tensor-core GEMMs. + +## What was tried + +`-Dllama.logitsLib={jit,gemmEx,lt}` switches the biggest single matvec — the logits +vocabulary projection (`Wcls(vocab×dim)·x`, 525 MB FP16 read/token on Llama-1B = ~23% +of all decode bytes) — between: + +- `jit` — stock `matrixVectorGeneric` (TornadoVM JIT). +- `gemmEx` — `CuBlas::cublasGemmExFP16FP32` library task (FP16 in, FP32 out; + `gemm(OP_T, OP_N, m=vocab, n=1, k=dim)` — cuBLAS is column-major, row-major `Wcls` + is a `dim×vocab` column-major matrix, `OP_T` recovers it). +- `lt` — `CuBlasLt::ltMatmulFP16` (heuristic kernel selection) + a 1.2 µs + `halfToFloat` copy task (Lt FP16 writes half; the sampler reads FP32). + +All three produce coherent output; library tasks run in the same task graph, on the +same stream, and are CUDA-graph-capturable. + +## Per-kernel truth (nsys, 101 logits calls, Llama-1B) + +| logits GEMV 128256×2048 FP16 | avg / call | effective BW | +|------------------------------|-----------:|-------------:| +| JIT `matrixVectorGeneric` | 553.8 µs | ~948 GB/s (94% peak) | +| cuBLAS `gemvx::kernel` (gemmEx) | 550.6 µs | ~953 GB/s | +| cuBLASLt `gemvx::kernel` (+halfToFloat 1.2 µs) | 550.9 µs | ~953 GB/s | + +Both libraries dispatch a GEMV-specialized `internal::gemvx::kernel` — 0.6% faster than +the JIT kernel, i.e. **parity**. There is nothing to win: the kernel is a +DRAM-bandwidth measurement device. + +## End-to-end (RTX 4090, 3 runs each, tok/s) + +| Llama-3.2-1B FP16 | graphs off | graphs on | +|--------------------|-----------|-----------| +| jit | 88.7–91.3 | 98.1–102.4 | +| lt | 88.3–90.5 | 99.9–102.2 | +| gemmEx | 92.1 | 105.3 | + +| Mistral-7B FP16 (graphs on) | tok/s | +|------------------------------|------:| +| jit | 23.5 | +| lt | 23.7 | + +Statistical ties everywhere (run-to-run spread ±2% exceeds any config delta). +CUDA graphs, by contrast, are worth +8–12% on every config. + +## Why libraries can't win at n=1 (and what does) + +Decode kernel profile (nsys, Llama-1B, % of GPU time): `fusedRmsNormFFNGateUp` 36.6% +(805 GB/s), `matrixVectorGenericWithResidual` (wo, w2+residual) 21.8%, +`processHeadsFlashAttention` 16.4%, logits GEMV 15.2% (948 GB/s), `fusedQKVMatmulX` +6.7% (823 GB/s). Every projection is a bandwidth-bound GEMV at 80–95% of peak. + +1. **Bandwidth wall.** At n=1 each projection reads its full weight matrix for one + output vector (arithmetic intensity ~1). Tensor cores never engage; the best any + kernel can do is stream weights at DRAM speed — which JIT, cuBLAS, and cuBLASLt all + already do. +2. **Launch-bound regime.** GPU busy is ~3.6 ms/token vs ~11.2 ms wall (no graphs) — + ~65 µs of host overhead per kernel launch. TornadoVM's fused JIT kernels + (rms+gate+up in one, qkv in one, matvec+residual in one) exist to *remove* + launches. Replacing `fusedRmsNormFFNGateUp` with library GEMVs would de-fuse it + into ≥3 launches per layer (48/step on Llama-1B) to save ~13 µs/layer of GPU time — + a net loss even under graph replay. +3. **Attention** (16%) is already a fused flash kernel; cuDNN `sdpaForward` was + assessed earlier — no single-request win (pays off for batched/ragged attention). + +**When hybrid libraries do win:** n≥16 turns the same projections into compute-bound +GEMMs — batch prefill via tensor-core MMA (PR #127), batched/continuous decode at +200–4200 tok/s aggregate (PR #129), and cuBLAS/cuBLASLt GEMMs generally. The library +integration itself is sound: correct results, same stream, graph-capturable — it is +the n=1 arithmetic intensity that leaves no room. + +## Reproduce + +```bash +# build (JDK 21, TornadoVM 5.0.1-jdk21-dev CUDA backend with tornado-cublas) +mvn -Pjdk21 -Dtornadovm.base.version=5.0.1 -Djdk.version.suffix=-jdk21-dev clean package -DskipTests + +# stock vs library logits (verify text is coherent every run) +JAVA_TOOL_OPTIONS="-Dllama.logitsLib=jit" llama-tornado --gpu --cuda [--cuda-graphs] --model --prompt "..." -n 200 --instruct +JAVA_TOOL_OPTIONS="-Dllama.logitsLib=lt" ... +JAVA_TOOL_OPTIONS="-Dllama.logitsLib=gemmEx" ... + +# per-kernel comparison +nsys profile -o out --trace=cuda python3 llama-tornado --gpu --cuda ... +nsys stats --report cuda_gpu_kern_sum out.nsys-rep # matrixVectorGeneric vs gemvx::kernel +``` diff --git a/llama-tornado b/llama-tornado index 78388295..dbaa856b 100755 --- a/llama-tornado +++ b/llama-tornado @@ -184,7 +184,7 @@ class LlamaRunner: [ f"@{self.tornado_sdk}/etc/exportLists/cuda-exports", "--add-modules", - "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.cuda", + "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.cuda,tornado.cublas", ] ) elif args.backend == Backend.METAL: diff --git a/pom.xml b/pom.xml index f5dd084f..c1d27828 100644 --- a/pom.xml +++ b/pom.xml @@ -68,6 +68,11 @@ tornado-runtime ${tornadovm.version} + + io.github.beehive-lab + tornado-cublas + ${tornadovm.version} + diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java index 1295db09..f106bafc 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java @@ -13,15 +13,39 @@ import uk.ac.manchester.tornado.api.GridScheduler; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.annotations.Parallel; import uk.ac.manchester.tornado.api.enums.DataTransferMode; +import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; +import uk.ac.manchester.tornado.cublas.CuBlas; +import uk.ac.manchester.tornado.cublas.CuBlasLt; +import uk.ac.manchester.tornado.cublas.enums.CuBlasOperation; public class LogitsFP16Layer extends AbstractLogitsTaskGraph { + /** + * Vocabulary-projection backend for the logits matvec (hybrid-library experiment): + * {@code jit} (default: TornadoVM JIT matvec), {@code gemmEx} (cuBLAS GemmEx, FP16 in / + * FP32 out) or {@code lt} (cuBLASLt heuristic-selected FP16 matmul + FP16→FP32 copy task). + */ + static final String LOGITS_LIB = System.getProperty("llama.logitsLib", "jit"); + + /** FP16 logits buffer for the cuBLASLt path (Lt FP16 matmul writes half precision). */ + private HalfFloatArray wrapLogitsFP16; + public LogitsFP16Layer(String name, State state, Weights weights, Configuration config, String lastTaskGraphID, SchedulerType schedulerType) { super(name, state, weights, config, lastTaskGraphID, schedulerType); } + /** Elementwise FP16→FP32 copy (cuBLASLt path: the sampler reads FP32 logits on the host). */ + public static void halfToFloat(FloatArray out, HalfFloatArray in, int size) { + for (@Parallel int i = 0; i < size; i++) { + float v = in.get(i).getFloat32(); + out.set(i, v); + } + } + /** * Hook called before any data transfers or tasks. Override to prepend * {@code consumeFromDevice} declarations that must precede the bytecode @@ -82,15 +106,39 @@ protected TaskGraph setupLogitsTaskGraph(TornadoWeights weights, Configuration c state.tempLogits); // scale factor from reduction // === Vocabulary Projection === - logits.task("vocab_proj", - TransformerComputeKernelsLayered::matrixVectorGeneric, - context, - state.wrapXbFP16, // input (FP16) - state.wrapLogits, // output - weights.wclsByteArray.asHalfFloatArray(), // vocabulary weights - config.dim(), // input dimension - config.vocabularySize(), // output dimension - LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS); + // Row-major logits(vocab) = Wcls(vocab x dim) · x(dim). cuBLAS is column-major, so + // Wcls stored row-major is a (dim x vocab) column-major matrix and op(A)=TRANSPOSE + // recovers Wcls: gemm(OP_T, OP_N, m=vocab, n=1, k=dim, A=Wcls lda=dim, B=x ldb=dim, + // C=logits ldc=vocab). + switch (LOGITS_LIB) { + case "gemmEx" -> logits.libraryTask("vocab_proj", CuBlas::cublasGemmExFP16FP32, + CuBlasOperation.CUBLAS_OP_T.operation(), CuBlasOperation.CUBLAS_OP_N.operation(), + config.vocabularySize(), 1, config.dim(), + 1.0f, weights.wclsByteArray.asHalfFloatArray(), config.dim(), + state.wrapXbFP16, config.dim(), + 0.0f, state.wrapLogits, config.vocabularySize()); + case "lt" -> { + wrapLogitsFP16 = new HalfFloatArray(config.vocabularySize()); + logits.transferToDevice(DataTransferMode.FIRST_EXECUTION, wrapLogitsFP16); + logits.libraryTask("vocab_proj", CuBlasLt::ltMatmulFP16, + CuBlasOperation.CUBLAS_OP_T.operation(), CuBlasOperation.CUBLAS_OP_N.operation(), + config.vocabularySize(), 1, config.dim(), + 1.0f, weights.wclsByteArray.asHalfFloatArray(), config.dim(), + state.wrapXbFP16, config.dim(), + 0.0f, wrapLogitsFP16, config.vocabularySize()); + logits.task("logits_f32", LogitsFP16Layer::halfToFloat, + state.wrapLogits, wrapLogitsFP16, config.vocabularySize()); + } + default -> logits.task("vocab_proj", + TransformerComputeKernelsLayered::matrixVectorGeneric, + context, + state.wrapXbFP16, // input (FP16) + state.wrapLogits, // output + weights.wclsByteArray.asHalfFloatArray(), // vocabulary weights + config.dim(), // input dimension + config.vocabularySize(), // output dimension + LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS); + } // === Transfer Results to Host === logits.transferToHost(DataTransferMode.EVERY_EXECUTION, state.wrapLogits); @@ -102,12 +150,14 @@ protected TaskGraph setupLogitsTaskGraph(TornadoWeights weights, Configuration c @Override public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) { var logitsRMS = WorkerGridFactory.createRmsNormWorker(config.dim(), rmsLocalSize()); - var vocabSizeRowMajor = config.vocabularySize() * LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS; - var vocabWorker = new WorkerGrid1D(vocabSizeRowMajor); - vocabWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS, 1, 1); tornadoForwardScheduler.addWorkerGrid("logits.rms_reduce", logitsRMS); tornadoForwardScheduler.addWorkerGrid("logits.rms_apply_fp16", logitsRMS); - tornadoForwardScheduler.addWorkerGrid("logits.vocab_proj", vocabWorker); + if (LOGITS_LIB.equals("jit")) { + var vocabSizeRowMajor = config.vocabularySize() * LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS; + var vocabWorker = new WorkerGrid1D(vocabSizeRowMajor); + vocabWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS, 1, 1); + tornadoForwardScheduler.addWorkerGrid("logits.vocab_proj", vocabWorker); + } return tornadoForwardScheduler; }