Skip to content

Add TornadoVM CUDA backend support and tensor-core (MMA) accelerated batch prefill (FP16 & Q8_0, Llama & Qwen3)#127

Open
mairooni wants to merge 19 commits into
mainfrom
feat/mma_cuda
Open

Add TornadoVM CUDA backend support and tensor-core (MMA) accelerated batch prefill (FP16 & Q8_0, Llama & Qwen3)#127
mairooni wants to merge 19 commits into
mainfrom
feat/mma_cuda

Conversation

@mairooni

@mairooni mairooni commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

This PR enables GPULlama3.java to run on TornadoVM's new CUDA backend and adds a tensor-core (MMA) accelerated batched-prefill pipeline on top of it. It incorporates and supersedes #124 (this branch is based on it, the CUDA backend changes are carried here so backend support and the MMA path that requires it land together).

All five prefill projections (QKV, Wo, W1/W3, W2) are executed as 128×128-tile GEMMs on tensor cores (m16n8k16, FP32 accumulate) via TornadoVM's MMA kernel API, replacing the per-token matvec formulation. Key elements:

  • Fused projections: Q/K/V in a single launch writing a packed [q|k|v] buffer (fixes grid starvation on the skinny GQA projections), and W1/W3 fused into a packed [gate|up] buffer.
  • Software-pipelined GEMMs: per-thread register staging of the next K-step's tiles overlaps global loads with tensor-core compute.
  • Q8_0 as W8A16: quantized weights stay in memory and are dequantized to FP16 in the GEMM staging registers — roughly half the weight traffic of FP16 with the same tensor-core pipeline.
  • Rewritten flash attention: register-partitioned P·V accumulation (removes redundant per-thread work), coalesced K/V tile loads, direct FP16 output.
  • Parallel RMS reductions with fused residual adds; packed RoPE/SwiGLU; 15 → 12 kernels per layer (13 for Qwen3, including per-head QK-norm).
  • Model coverage: Llama 3.x and Qwen3 (split-half RoPE, QK-norm, qDim ≠ dim handled), FP16 and Q8_0.
  • Arbitrary batch sizes via M-padding: The GEMMs operate on fixed 128-row tiles and the MMA API has no predicated store to mask partial tiles, so activation buffers are padded to the next 128-row multiple; padded rows are computed but never read (each output row depends only on its own input row, and non-GEMM kernels use the true batch size). Multiples of 128 give full GEMM efficiency; other sizes work correctly at proportionally reduced efficiency, reported at startup. This is the same tail-tile strategy that dense GEMM libraries use internally (cuBLAS computes full tiles for ragged M too) just visible at the allocation level here.
  • Graceful fallback: on backends without MMA support, the batch-prefill planners automatically fall back to the portable matvec pipeline

Example usage:
./llama-tornado --gpu --cuda --model Q8_0-Llama-3.2-1B-Instruct.gguf --prompt "$LONG_PROMPT" --max-tokens 2048 --with-prefill-decode --batch-prefill-size 128

mikepapadim and others added 17 commits June 20, 2026 14:46
Add a --cuda flag to llama-tornado that selects the TornadoVM CUDA backend,
mirroring the existing --opencl/--ptx/--metal plumbing: it loads the
tornado.drivers.cuda module and the cuda-exports export list. Also disambiguate
--ptx help text (was 'PTX/CUDA').
The CUDA backend is only available in a dev build of TornadoVM (PR #861), so
point the JDK21 build at 4.0.2-jdk21-dev. The project's own version is unchanged.
List CUDA among the supported backends, add a --cuda usage example, and note
that the CUDA backend requires a TornadoVM build with the CUDA backend from
PR #861 (beehive-lab/TornadoVM#861).
Add a cuda variant to the build, standalone-inference, and quarkus-integration
backend matrices. The setup-tornadovm action now builds the CUDA backend from
the cuda2 branch (TornadoVM PR #861) until it is merged to master; other
backends still build from master. Shared inference steps run on CUDA via the
matrix; the PTX-only CUDA-graph steps remain gated to ptx.
@mairooni mairooni self-assigned this Jul 7, 2026
@mairooni mairooni added enhancement New feature or request prefill-decode labels Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request prefill-decode

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants