GPU groundwork: standard-ops forecast graph and exact-order device linear kernels

CMakeLists.txt

@@ -18,6 +18,7 @@ set(TIMESTRAL_EVAL_RUNTIME_COMMIT "" CACHE STRING "Explicit evaluation runtime c
 set(TIMESTRAL_EVAL_SOURCE_SHA256 "" CACHE STRING "Explicit evaluation source snapshot SHA256")
 option(TIMESTRAL_BUILD_CLI    "Build timestral-cli" ON)
 option(TIMESTRAL_SHARED       "Build libtimestral as shared" OFF)
+option(TIMESTRAL_STANDARD_OPS "Default to the standard-ggml-op forecast graph (NORM-gated, GPU-capable) instead of the bit-exact CPU custom ops" OFF)
 option(TIMESTRAL_GGML_CUDA    "Forward GGML_CUDA" OFF)
 option(TIMESTRAL_GGML_METAL   "Forward GGML_METAL" OFF)
 option(TIMESTRAL_GGML_VULKAN  "Forward GGML_VULKAN" OFF)
@@ -64,7 +65,8 @@ set(TIMESTRAL_SRC
     src/encoder.cpp
     src/head.cpp
     src/predict_block.cpp
-    src/rope.cpp)
+    src/rope.cpp
+    src/standard_ops.cpp)
 
 # PyTorch materializes multiply/add separately at the scaler float32 boundaries.
 
@@ -82,6 +84,9 @@ set_target_properties(timestral PROPERTIES POSITION_INDEPENDENT_CODE ON)
 target_include_directories(timestral PUBLIC include $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/src>)
 target_compile_definitions(timestral PRIVATE TIMESTRAL_VERSION="${TIMESTRAL_VERSION}")
 target_link_libraries(timestral PUBLIC ggml)
+if(TIMESTRAL_STANDARD_OPS)
+  target_compile_definitions(timestral PRIVATE TS_STANDARD_OPS_DEFAULT=1)
+endif()
 if(TIMESTRAL_BUILD_TESTS)
   target_compile_definitions(timestral PRIVATE TS_CALLBACK_TESTING)
 endif()

docs/standard-ops.md

@@ -0,0 +1,79 @@
+# Standard-op forecast graph
+
+Status: experimental, milestone (a) of GPU support. Off by default.
+
+The default forecast graph uses `ggml_custom_4d` callbacks. They reproduce the
+reference rounding bit for bit on CPU and have no GPU kernel. The standard-op
+graph uses only built-in ggml operations, so a GPU backend can schedule every
+node. It is not bit-exact. It is held to the NORM gate
+(`err <= max(2 * floor, 16u)`, `docs/PLAN.md` D.1).
+
+## Switch
+
+- `ts::standard_ops()` is read when a graph is built.
+- Default: off. `-DTIMESTRAL_STANDARD_OPS=ON` makes the default on.
+- `TIMESTRAL_STANDARD_OPS=0|1` in the environment overrides the default.
+- `ts::set_standard_ops(bool)` changes it at run time.
+
+## Components
+
+| component | standard ops |
+|---|---|
+| linear layers, patch embedding, head | `ggml_mul_mat` per K panel of 192, `ggml_add` (panels in order, then bias) |
+| RMS norm | element-wise F32 cascade sum (`ggml_sqr`, views, `ggml_add`), `ggml_div`, `ggml_sqrt`, `ggml_mul` |
+| SiLU, sigmoid | `ggml_silu`, `ggml_sigmoid` |
+| attention | `ggml_mul_mat`, `ggml_soft_max_ext` (additive mask, scale `1/sqrt(head_dim)`), `ggml_mul_mat` |
+| RoPE | unchanged: stored tables, `mul`, `neg`, `concat`, `add` (no custom op) |
+
+Host steps stay as they are: scaler, refine, covariates, float64 quantile
+steps. The standalone MKL option is CPU only and is not changed.
+
+## Gates
+
+`tests/test_standard_ops.cpp` prints one `GATE` line per comparison. It is not
+registered with ctest. Inputs are the same as the real-model ladder:
+`TIMESTRAL_TEST_GGUF`, `TIMESTRAL_TEST_REAL_BASELINES`, `TIMESTRAL_TEST_REAL_CASES`.
+`TS_GATE_COMPONENTS` selects a comma-separated subset of
+`patch,layer,encoder,head,predict_block,forecast`.
+
+## Arithmetic choices
+
+Each choice below was found by running one operation on the reference's own
+input and comparing it with the stored reference (see the NORM gate in
+`tests/test_standard_ops.cpp`).
+
+- Linear layers. The reference F32 GEMM reduces K in panels of 192, with one
+  accumulator chain per panel and the panel sums added in order. One
+  `ggml_mul_mat` over all of K uses a different rounding path. Layers with
+  outlier activations amplify that: the error is a few units in the last place
+  of the largest output element, and the gate allows about one. Splitting K
+  into panels of 192 and adding the panel results in order removes most of it.
+  Dots shorter than 192 stay one `mul_mat`.
+- RMS norm. The reference sums x*x in F32 with a fixed cascade (32 lanes, power
+  of two merge blocks, fixed fold). `ggml_rms_norm` sums in double. Both are
+  accurate, but they round differently, and the last bit of 1/sqrt(mean + eps)
+  shows in the output of any row with large channels. The standard-op norm
+  rebuilds the cascade from element-wise nodes, which are exact IEEE F32 on
+  every backend. Scalars such as D and 1.0 come from `graph_constant`, which the
+  graph runner writes when the graph is allocated. If D is not a multiple of 32
+  the norm falls back to `ggml_rms_norm`.
+- SiLU and sigmoid use `ggml_silu` and `ggml_sigmoid`. Their effect on the gates
+  is below 0.05 of the tolerance in the cases checked.
+
+## Known residual
+
+On the 209M F32 model, 33 of 7899 NORM checks still fail (the first version of
+this graph failed 122). The worst ratio of error to tolerance is 3.5. The causes
+are not rounding bugs.
+
+- Attention scores. The reference accumulates q.k in double and rounds once.
+  `ggml_mul_mat` accumulates in F32. With the exact dot products and the standard
+  softmax and value product, the layer 0 attention output of `uni_one_patch`
+  passes at 0.20 of the tolerance; with F32 dots it fails at 1.95. No standard
+  operation computes a double-accumulated dot product.
+- Linear layers in layers with outlier activations (largest channel about 100
+  times the rms). A different F32 accumulation order moves the largest output by
+  one to four units in the last place, which is at or above the gate. A
+  compensated panel sum was tried (TwoSum, panels of 64). It made the dot
+  products closer to a double-precision sum but did not reduce the number of
+  failed gates, so it is not used.

src/backend.cpp

@@ -101,8 +101,14 @@ CallbackContext* register_callback(ggml_context* ctx,const std::vector<ggml_tens
  auto context=std::make_unique<CallbackContext>();context->status=&building_runner->callback_status_;context->split=split;context->tensors=tensors;
  auto result=context.get();building_runner->callbacks_.push_back(std::move(context));return result;
 }
+ggml_tensor* graph_constant(ggml_context* ctx,float value){
+ auto found=building_runners.find(ctx);auto runner=found==building_runners.end()?nullptr:found->second;
+ if(!runner||runner->ctx_!=ctx||!std::holds_alternative<lifecycle::Building>(runner->state_.phase))throw Error("graph constant requires owning GraphRunner");
+ for(auto& c:runner->constants_)if(std::memcmp(&c.second,&value,sizeof value)==0)return c.first;
+ auto t=ggml_new_tensor_1d(ctx,GGML_TYPE_F32,1);ggml_set_input(t);runner->constants_.emplace_back(t,value);return t;
+}
 ggml_context* GraphRunner::begin() {
-    transition(lifecycle::Event::Begin);callbacks_.clear();callback_status_.reset();if(ctx_)building_runners.erase(ctx_);
+    transition(lifecycle::Event::Begin);constants_.clear();callbacks_.clear();callback_status_.reset();if(ctx_)building_runners.erase(ctx_);
     if (ctx_) ggml_free(ctx_);
     ctx_ = nullptr; graph_ = nullptr; allocated_ = false; computed_ = false;
     ggml_init_params p{capacity_*ggml_tensor_overhead()+ggml_graph_overhead_custom(capacity_,false), nullptr, true};
@@ -122,6 +128,7 @@ void GraphRunner::allocate(const std::vector<ggml_tensor*>& outputs) {
     try {
       for(auto t:outputs){ggml_set_output(t);ggml_build_forward_expand(graph_,t);}
       if(!ggml_gallocr_reserve(allocator_,graph_)||!ggml_gallocr_alloc_graph(allocator_,graph_))throw Error("cannot allocate graph buffers");
+      for(auto& c:constants_)if(c.first->buffer)ggml_backend_tensor_set(c.first,&c.second,0,sizeof(float));
       transition(lifecycle::Event::Allocate);allocated_=true;
     }catch(...){state_=lifecycle::reduce(state_,lifecycle::Event::AllocationFail).state;throw;}
 }

src/backend.hpp

@@ -45,6 +45,11 @@ ggml_backend_buffer_t backend_buffer_from_mapping(MappedFile& file);
 #include "graph_lifecycle.hpp"
 #include "callback_context.hpp"
 namespace ts {
+class GraphRunner;
+// A one-element F32 tensor holding `value`, owned by the runner that is building `ctx`.
+// The runner writes it when the graph is allocated, so graph builders can use exact
+// scalar operands (for example a division by D) without a host upload step.
+ggml_tensor* graph_constant(ggml_context* ctx, float value);
 // Owns metadata per invocation and a persistent backend allocator across shapes.
 // Not reentrant. Graph tensors are invalidated by begin(); read() returns owned data.
 class GraphRunner {
@@ -60,6 +65,8 @@ class GraphRunner {
     std::vector<float> read(ggml_tensor* t) const;
 private:
     friend CallbackContext* register_callback(ggml_context*,const std::vector<ggml_tensor*>&,size_t);
+    friend ggml_tensor* graph_constant(ggml_context*,float);
+    std::vector<std::pair<ggml_tensor*,float>> constants_;
     lifecycle::State state_;
     CallbackStatus callback_status_;
     std::vector<std::unique_ptr<CallbackContext>> callbacks_;

src/encoder.cpp

@@ -5,8 +5,11 @@ namespace ts {
 static size_t encoder_capacity(const Model& model) {
     // Bound graph metadata before ggml construction (including views and leaves).
     const size_t n=model.config().n_layers;
-    if(!n || n>(1000000-4096)/512) throw Error("encoder graph capacity exceeded");
-    return 4096+512*n;
+    // The standard-op graph splits long contractions into panels and builds the
+    // RMS-norm reduction from element-wise nodes, so a layer needs more nodes.
+    constexpr size_t per_layer=1024;
+    if(!n || n>(1000000-4096)/per_layer) throw Error("encoder graph capacity exceeded");
+    return 4096+per_layer*n;
 }
 Encoder::Encoder(const Model& model) : model_(model), runner_(encoder_capacity(model)) {}
 Embedding Encoder::embed(const std::vector<float>& x,const std::vector<int8_t>& observed,size_t rows,size_t length) {

src/gpu/README.md

@@ -0,0 +1,46 @@
+# Exact-order F32 linear on a GPU
+
+Files:
+- `linear_exact_gpu.hpp`: one source for CUDA and HIP. One thread per output (row, column). The kernel and the host wrapper `linear_exact_gpu` (same signature as the CPU reference; allocate, copy in, launch, synchronize, copy out, free).
+- `linear_exact_ref.hpp`: scalar CPU reference with the exact order.
+- `test_linear_exact_gpu.cpp`: bitwise test over 880 cases (K in {768, 2048, 2304, 4096}, N in {1, 7, 8, 9, 15, 17, 768, 2047, 2048, 2304, 4096}, rows in {1, 3, 31, 32, 37}, with and without bias, with and without an outlier channel). Prints the mismatch count per case and a total. Exit code is nonzero on any mismatch.
+
+The default build does not use this directory.
+
+## Order of operations per output
+For K >= 192: for each 192-wide panel, ascending k, `panel = fmaf(w, x, panel)` from +0; then `sum += panel`, panels in ascending order; bias added last. For K < 192: one chain, no panel add. A chain is never split or reordered across threads.
+
+## Compile commands
+CUDA:
+
+    nvcc -std=c++17 -O2 -fmad=false -x cu test_linear_exact_gpu.cpp -o test_linear_exact_gpu
+
+HIP:
+
+    hipcc -std=c++17 -O2 -ffp-contract=off test_linear_exact_gpu.cpp -o test_linear_exact_gpu
+
+CPU emulation (no GPU toolchain needed):
+
+    g++ -std=c++17 -O2 -ffp-contract=off -DTS_GPU_EMULATE test_linear_exact_gpu.cpp -o test_linear_exact_emul
+
+Add the right `-arch` / `--offload-arch` flag for the target device.
+
+## Flags and why
+- Every multiply-add in the chain is an explicit `fmaf`, which is a single-rounding fused operation on both CUDA and HIP, so bit equality does not depend on the contraction flag.
+- We pick `-fmad=false` (CUDA) and `-ffp-contract=off` (HIP) anyway. The only other float operations are `sum += panel` and `sum + bias`, which have no multiply to fuse. Turning contraction off removes any chance that the compiler fuses or reorders something we did not write. `-fmad=true` would also give the same result here, but it adds nothing and leaves room for surprise.
+- No `--use_fast_math`, no `-ffast-math`, no flush-to-zero flags.
+
+## Test command
+Run the built binary with no arguments: `./test_linear_exact_gpu`. It needs about 2 GB of host memory for the largest case.
+
+## Verified so far
+- No `nvcc` or `hipcc` was available on the machine used to write this. The CUDA and HIP paths are NOT compiled and NOT run. The device arithmetic is unverified.
+- The CPU emulation build runs the same `ts_exact_output` function (with `std::fmaf`) in a host loop over all (row, column) pairs. It matched `linear_exact_ref` bit for bit on all 880 cases (18,836,480 outputs, 0 mismatches). This proves the loop structure and index math only. It does not prove anything about device fused-multiply-add behaviour.
+- The CPU reference itself was checked earlier against the default packed linear: 0 mismatches over the same 880 cases.
+
+## Remaining device-side checks
+1. Build with both toolchains (the macro layer has never been compiled).
+2. Run the test on hardware and record the real mismatch count (must be 0).
+3. Confirm `fmaf` compiles to a single fused instruction (inspect PTX `fma.rn.f32` or the ISA `v_fma_f32`), with no split multiply and add.
+4. Check subnormal handling: the CPU keeps subnormals, so the device must not flush them (CUDA without `-ftz=true`, HIP without flush-denormals options). Add a test case with subnormal inputs and products.
+5. Check that the rounding mode is round-to-nearest-even for the fused operation and the adds.

src/gpu/linear_exact_gpu.hpp

@@ -0,0 +1,101 @@
+// Exact-order F32 linear on a GPU: one thread per output (row, column).
+// One source for CUDA and HIP. Define TS_GPU_EMULATE to run the same
+// per-output function on the CPU instead (loop-structure check only).
+//
+// Order per output (identical to linear_exact_ref):
+//   k >= 192: for each 192-wide panel, ascending k, panel = fmaf(w, x, panel)
+//             starting at +0; then sum += panel (panels ascending, sum from +0).
+//   k <  192: one chain, sum = fmaf(w, x, sum).
+//   bias is added last. No chain is split or reordered across threads.
+#pragma once
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+
+#if defined(TS_GPU_EMULATE)
+  #define TS_DEV_FN inline
+  #define TS_FMAF(a, b, c) std::fmaf((a), (b), (c))
+#else
+  #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
+    #include <hip/hip_runtime.h>
+    #define TS_RT(x) hip##x
+  #else
+    #include <cuda_runtime.h>
+    #define TS_RT(x) cuda##x
+  #endif
+  #define TS_DEV_FN __host__ __device__ inline
+  // Explicit fused multiply-add; never left to the compiler's contraction.
+  #define TS_FMAF(a, b, c) fmaf((a), (b), (c))
+#endif
+
+constexpr int64_t TS_PANEL = 192;
+
+// The per-output chain. Shared by the device kernel and the CPU emulation.
+TS_DEV_FN float ts_exact_output(const float* w, const float* x, const float* b,
+                                int64_t k, int64_t n, int64_t row, int64_t col) {
+    const float* wr = w + col * k;
+    const float* xr = x + row * k;
+    float sum = 0.0f;
+    if (k >= TS_PANEL) {
+        for (int64_t s = 0; s < k; s += TS_PANEL) {
+            float panel = 0.0f;
+            const int64_t e = (s + TS_PANEL < k) ? (s + TS_PANEL) : k;
+            for (int64_t j = s; j < e; ++j) panel = TS_FMAF(wr[j], xr[j], panel);
+            sum += panel;
+        }
+    } else {
+        for (int64_t j = 0; j < k; ++j) sum = TS_FMAF(wr[j], xr[j], sum);
+    }
+    return b ? sum + b[col] : sum;
+}
+
+#if defined(TS_GPU_EMULATE)
+
+// Host emulation of the launch: visit every (row, col) thread once.
+inline void linear_exact_gpu(const float* w, const float* x, const float* b, float* y,
+                             int64_t k, int64_t n, int64_t rows) {
+    for (int64_t row = 0; row < rows; ++row)
+        for (int64_t col = 0; col < n; ++col)
+            y[row * n + col] = ts_exact_output(w, x, b, k, n, row, col);
+}
+
+#else
+
+__global__ void ts_linear_exact_kernel(const float* w, const float* x, const float* b,
+                                       float* y, int64_t k, int64_t n, int64_t rows) {
+    const int64_t col = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    const int64_t row = blockIdx.y;
+    if (col >= n || row >= rows) return;
+    y[row * n + col] = ts_exact_output(w, x, b, k, n, row, col);
+}
+
+#define TS_CHECK(call) do { auto e_ = (call); if (e_ != TS_RT(Success)) { \
+    std::fprintf(stderr, "%s:%d: %s failed: %s\n", __FILE__, __LINE__, #call, \
+                 TS_RT(GetErrorString)(e_)); std::exit(2); } } while (0)
+
+// Same signature as linear_exact_ref. Synchronous: allocate, copy, launch, copy back, free.
+inline void linear_exact_gpu(const float* w, const float* x, const float* b, float* y,
+                             int64_t k, int64_t n, int64_t rows) {
+    float *dw = nullptr, *dx = nullptr, *db = nullptr, *dy = nullptr;
+    const size_t bw = (size_t)n * k * 4, bx = (size_t)rows * k * 4;
+    const size_t bb = (size_t)n * 4, by = (size_t)rows * n * 4;
+    TS_CHECK(TS_RT(Malloc)((void**)&dw, bw));
+    TS_CHECK(TS_RT(Malloc)((void**)&dx, bx));
+    TS_CHECK(TS_RT(Malloc)((void**)&dy, by));
+    TS_CHECK(TS_RT(Memcpy)(dw, w, bw, TS_RT(MemcpyHostToDevice)));
+    TS_CHECK(TS_RT(Memcpy)(dx, x, bx, TS_RT(MemcpyHostToDevice)));
+    if (b) {
+        TS_CHECK(TS_RT(Malloc)((void**)&db, bb));
+        TS_CHECK(TS_RT(Memcpy)(db, b, bb, TS_RT(MemcpyHostToDevice)));
+    }
+    const int threads = 128;
+    dim3 grid((unsigned)((n + threads - 1) / threads), (unsigned)rows);
+    ts_linear_exact_kernel<<<grid, threads>>>(dw, dx, db, dy, k, n, rows);
+    TS_CHECK(TS_RT(GetLastError)());
+    TS_CHECK(TS_RT(DeviceSynchronize)());
+    TS_CHECK(TS_RT(Memcpy)(y, dy, by, TS_RT(MemcpyDeviceToHost)));
+    TS_RT(Free)(dw); TS_RT(Free)(dx); TS_RT(Free)(dy); if (db) TS_RT(Free)(db);
+}
+
+#endif

src/gpu/linear_exact_ref.hpp

@@ -0,0 +1,26 @@
+// Scalar exact-order CPU reference for the default packed F32 linear.
+// One output at a time. Ascending k inside a panel of 192, one fused
+// multiply-add per k, panel sums added in ascending panel order starting
+// from +0, bias last. For k < 192: one FMA chain, no panel boundary.
+#pragma once
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+inline void linear_exact_ref(const float* w, const float* x, const float* b, float* y,
+                             int64_t k, int64_t n, int64_t rows) {
+    for (int64_t r = 0; r < rows; ++r) for (int64_t c = 0; c < n; ++c) {
+        const float* wr = w + c * k; const float* xr = x + r * k;
+        float sum = 0;
+        if (k >= 192) {
+            for (int64_t s = 0; s < k; s += 192) {
+                float panel = 0;
+                const int64_t e = std::min<int64_t>(k, s + 192);
+                for (int64_t j = s; j < e; ++j) panel = std::fma(wr[j], xr[j], panel);
+                sum += panel;
+            }
+        } else {
+            for (int64_t j = 0; j < k; ++j) sum = std::fma(wr[j], xr[j], sum);
+        }
+        y[r * n + c] = b ? sum + b[c] : sum;
+    }
+}