Initial release

engine/patches/README.md

@@ -0,0 +1,15 @@
+# engine/patches
+
+`bf16-out-prod.patch` teaches ggml-cuda's `out_prod` to accept a **BF16 `src0`**, which is what lets the DiT trainer's `--mirror bf16` keep frozen trainable-layer weights in BF16 instead of promoting the whole mirror to F32. `engine/ggml/` is a git submodule, so a submodule update silently reverts it.
+
+`mm-backward.patch` adds an **env-gated alternative formulation** for the `MUL_MAT` activation gradient in `ggml.c`'s `ggml_compute_backward`. Upstream emits `out_prod(src0, transpose(grad))`, and ggml-cuda's `OUT_PROD` is F32-only — which forces the frozen weight to F32 and drags the *forward* `mul_mat` onto TF32 tensor cores too. With `GGML_BACKWARD_MM=1` set, the backward becomes `mul_mat(cont(transpose(src0)), grad)` instead: mathematically and shape-wise identical, but dtype-agnostic, so a BF16 weight rides real BF16 tensor cores in both directions with no dequant. Measured on an RTX 5090 at ~1.7–1.8× per layer per step. With the env var unset the emitted graph is byte-identical to upstream.
+
+`ace-train`'s `--bwd <outprod|mm>` sets that env var; the Training Studio defaults both trainers to `mm`.
+
+Reapply from the repo root — **apply all of them**, they touch disjoint files (`ggml-cuda/*.cu` vs `ggml.c`) so order does not matter:
+
+```sh
+for p in engine/patches/*.patch; do git apply --verbose "$p"; done
+```
+
+Verify they are still in place: `powershell -File engine\verify-hooks.ps1` (Hook 7 greps `out-prod.cu`, Hook 8 greps `ggml.c`, both for their HOT-Step marker comments). CI reapplies them in the "Apply engine patches" step of every build job, using the same glob loop.

engine/patches/bf16-out-prod.patch

@@ -0,0 +1,84 @@
+diff --git a/engine/ggml/src/ggml-cuda/ggml-cuda.cu b/engine/ggml/src/ggml-cuda/ggml-cuda.cu
+index 3e11b456..aff788f5 100644
+--- a/engine/ggml/src/ggml-cuda/ggml-cuda.cu
++++ b/engine/ggml/src/ggml-cuda/ggml-cuda.cu
+@@ -5189,7 +5189,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
+                 }
+             } break;
+         case GGML_OP_OUT_PROD:
+-            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
++            // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
++            return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
++                   (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16);
+         case GGML_OP_GET_ROWS:
+             {
+                 switch (op->src[0]->type) {
+diff --git a/engine/ggml/src/ggml-cuda/out-prod.cu b/engine/ggml/src/ggml-cuda/out-prod.cu
+index 499903d0..9b48c8b5 100644
+--- a/engine/ggml/src/ggml-cuda/out-prod.cu
++++ b/engine/ggml/src/ggml-cuda/out-prod.cu
+@@ -1,4 +1,5 @@
+ #include "out-prod.cuh"
++#include "convert.cuh" // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
+ 
+ #include <cstdint>
+ 
+@@ -8,7 +9,14 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+ 
+     GGML_TENSOR_BINARY_OP_LOCALS
+ 
+-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
++    // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
++    // ggml computes the gradient w.r.t. a mul_mat's ACTIVATION input as
++    // out_prod(weight, transpose(grad)), so the frozen weight lands in src0. An
++    // F32-only assert here is what forces the DiT trainer to mirror every
++    // trainable-layer weight to F32. Accepting BF16 src0 (dequantized once into
++    // an F32 workspace below) halves that mirror. src1/dst stay F32 and the F32
++    // path is byte-identical.
++    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
+     GGML_ASSERT(src1->type == GGML_TYPE_F32);
+     GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+ 
+@@ -22,19 +30,37 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+     GGML_ASSERT(ne2 == src1->ne[2]);
+     GGML_ASSERT(ne3 == src1->ne[3]);
+ 
+-    const float * src0_d = (const float *) src0->data;
+-    const float * src1_d = (const float *) src1->data;
+-    float       *  dst_d = (float       *)  dst->data;
+-
+     cudaStream_t   stream = ctx.stream();
+     cublasHandle_t handle = ctx.cublas_handle();
+ 
++    // HOT-Step patch: BF16 out_prod — see engine/patches/bf16-out-prod.patch
++    // Dequantize a BF16 src0 into a contiguous F32 workspace, so every GEMM path
++    // below is the shipped F32 one with lda == ne00. Restricted to a 2-D src0
++    // (the trainer's case: a frozen weight matrix), which also makes the dim-2/3
++    // workspace strides moot — with ne02 == ne03 == 1 the (i2/dps2) and (i3/dps3)
++    // src0 offsets below are identically zero.
++    ggml_cuda_pool_alloc<float> src0_f32(ctx.pool());
++    if (src0->type == GGML_TYPE_BF16) {
++        GGML_ASSERT(ne02 == 1 && ne03 == 1);
++        GGML_ASSERT(ggml_is_contiguous(src0));
++        const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16);
++        GGML_ASSERT(to_fp32_cuda != nullptr);
++        src0_f32.alloc((size_t) ne00 * ne01);
++        to_fp32_cuda(src0->data, src0_f32.get(), ne00 * ne01, stream);
++    }
++
++    const float * src0_d = src0_f32.get() ? src0_f32.get() : (const float *) src0->data;
++    const float * src1_d = (const float *) src1->data;
++    float       *  dst_d = (float       *)  dst->data;
++
+     const float alpha = 1.0f;
+     const float beta = 0.0f;
+ 
+     CUBLAS_CHECK(cublasSetStream(handle, stream));
+ 
+-    const int64_t lda = nb01 / sizeof(float);
++    // HOT-Step patch: BF16 out_prod — the workspace is contiguous, so its leading
++    // dimension is ne00 rather than src0's own (BF16-sized) row stride.
++    const int64_t lda = src0_f32.get() ? ne00 : (int64_t) (nb01 / sizeof(float));
+     const int64_t ldc = nb1  / sizeof(float);
+ 
+     const bool src1_T = ggml_is_transposed(src1);

engine/patches/mm-backward.patch

@@ -0,0 +1,65 @@
+diff --git a/engine/ggml/src/ggml.c b/engine/ggml/src/ggml.c
+index b43016c8..55d50ad8 100644
+--- a/engine/ggml/src/ggml.c
++++ b/engine/ggml/src/ggml.c
+@@ -6597,6 +6597,52 @@ static void ggml_compute_backward(
+                 ggml_add_or_set(ctx, cgraph, isrc0, tmp);
+             }
+             if (src1_needs_grads) {
++                // HOT-Step patch: mm-backward — see engine/patches/mm-backward.patch
++                //
++                // The default arm below emits the ACTIVATION gradient as OUT_PROD.
++                // ggml-cuda implements OUT_PROD F32-only (cublasSgemm), so the frozen
++                // weight src0 must be F32, which drags the FORWARD mul_mat onto TF32
++                // tensor cores too. The mul_mat formulation upstream left commented
++                // out here is dtype-agnostic, so a BF16 src0 rides real BF16 tensor
++                // cores with no dequant and no F32 window.
++                //
++                // The two are provably shape-identical:
++                //   out_prod(src0[n,m,q1,r1], transpose(grad)[p,m,qq,rr]) -> [n,p,qq,rr]
++                //   mul_mat(cont(transpose(src0))[m,n,q1,r1], grad[m,p,qq,rr]) -> [n,p,qq,rr]
++                // and their broadcast preconditions are the same pair of
++                // (b->ne[2] % a->ne[2] == 0, b->ne[3] % a->ne[3] == 0) checks.
++                //
++                // Measured on an RTX 5090 (engine/src/train/spike-gemmbench.h): 1.67-1.82x
++                // per layer per step, parity vs the TF32 out_prod arm cosine 0.999996 /
++                // max_rel ~4e-3. Per-use `cont` is free at these shapes, so no
++                // pre-transposed weight cache is needed.
++                //
++                // Env-gated so an unset environment is byte-identical to upstream.
++                // Set GGML_BACKWARD_MM=1 (ace-train's `--bwd mm`) to take it. Read once:
++                // this runs once per mul_mat per graph build, and getenv is not free.
++                //
++                // CONTIGUITY GUARD (measured, not theoretical): `grad` becomes
++                // mul_mat's src1, and the CUDA mul_mat kernels require src1 to be
++                // row-contiguous — ggml-cuda/mmf.cu:28 asserts `nb10 == ts_src1`
++                // and ABORTS otherwise. out_prod has no such requirement, which is
++                // why upstream can hand it an arbitrary transposed view. A
++                // non-contiguous grad therefore keeps the out_prod arm rather than
++                // crashing (ggml_cont on grad is not the answer: grad is an
++                // ACTIVATION-sized tensor, so copying it would cost more than the
++                // GEMM saves). Found by ace-train train-dit --self-test under
++                // GGML_BACKWARD_MM=1: the LoKR rung SC3 aborted here.
++                static int hs_bwd_mm = -1;
++                if (hs_bwd_mm < 0) {
++                    const char * hs_e = getenv("GGML_BACKWARD_MM");
++                    hs_bwd_mm = (hs_e && hs_e[0] && strcmp(hs_e, "0") != 0) ? 1 : 0;
++                }
++                if (hs_bwd_mm && ggml_is_contiguous(grad)) {
++                    ggml_add_or_set(ctx, cgraph, isrc1,
++                            ggml_mul_mat(ctx,                   // [n,p,qq,rr]
++                                ggml_cont(ctx,                  // [m,n,q1,r1]
++                                    ggml_transpose(ctx, src0)), // [m,n,q1,r1]
++                                grad));                         // [m,p,qq,rr]
++                } else {
+                 ggml_add_or_set(ctx, cgraph, isrc1,
+                         // ggml_mul_mat(ctx,                   // [n,p,qq,rr]
+                         //     ggml_cont(ctx,                  // [m,n,q1,r1]
+@@ -6610,6 +6656,7 @@ static void ggml_compute_backward(
+                             src0,               // [n,m,q1,r1]
+                             ggml_transpose(ctx, // [p,m,qq,rr]
+                                 grad)));        // [m,p,qq,rr]
++                }
+             }
+         } break;
+         case GGML_OP_SCALE: {