Initial release

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);