Misc. bug: CUDA ADD/GELU slower on B200 since PDL commit #30004

// Times single ggml ops on the CUDA backend at the shapes of a small
// vision-transformer world model (588 tokens, batch 200), through the public
// API only, so the same file builds against any ggml version.
//
//   repro <op> [reps]     op: add_bias | add_res | gelu_erf | concat | flash_attn | all
//
// Prints one JSON line per op: median and min milliseconds over reps, after
// warmup. Inputs are fixed pseudo-random values.

#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml-cuda.h"
#include "ggml.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <functional>
#include <string>
#include <vector>

namespace {

constexpr int64_t kTokens = 196, kFrames = 3, kS = kTokens * kFrames, kB = 200;
constexpr int64_t kHeads = 16, kHeadDim = 64, kInner = kHeads * kHeadDim, kMlp = 2048, kDim = 28;

void fill(ggml_tensor* t) {
    if (t->type == GGML_TYPE_F16) {
        return;  // set by the caller (the mask)
    }
    std::vector<float> v(ggml_nelements(t));
    uint32_t s = 12345u;
    for (auto& x : v) {
        s = s * 1664525u + 1013904223u;
        x = float(s >> 8) / float(1u << 24) - 0.5f;
    }
    ggml_backend_tensor_set(t, v.data(), 0, ggml_nbytes(t));
}

double run(ggml_backend_t backend, const char* name, int reps,
           const std::function<ggml_tensor*(ggml_context*, std::vector<ggml_tensor*>&)>& build,
           const std::function<void(ggml_tensor*)>& set = nullptr) {
    ggml_init_params ip = {64 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true};
    ggml_context* ctx = ggml_init(ip);
    std::vector<ggml_tensor*> inputs;
    ggml_tensor* out = build(ctx, inputs);
    ggml_cgraph* gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, out);
    ggml_backend_buffer_t inbuf = ggml_backend_alloc_ctx_tensors(ctx, backend);  // inputs (leafs)
    ggml_gallocr_t ga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
    ggml_gallocr_alloc_graph(ga, gf);
    for (ggml_tensor* t : inputs) {
        if (set && t->type == GGML_TYPE_F16) {
            set(t);
        } else {
            fill(t);
        }
    }
    for (int i = 0; i < 3; ++i) {
        ggml_backend_graph_compute(backend, gf);
    }
    ggml_backend_synchronize(backend);
    std::vector<double> ms;
    for (int i = 0; i < reps; ++i) {
        auto t0 = std::chrono::steady_clock::now();
        ggml_backend_graph_compute(backend, gf);
        ggml_backend_synchronize(backend);
        ms.push_back(std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count());
    }
    std::sort(ms.begin(), ms.end());
    const double median = ms[ms.size() / 2];
    std::printf("{\"op\":\"%s\",\"median_ms\":%.4f,\"min_ms\":%.4f,\"reps\":%d}\n", name, median, ms[0], reps);
    std::fflush(stdout);
    ggml_gallocr_free(ga);
    ggml_backend_buffer_free(inbuf);
    ggml_free(ctx);
    return median;
}

ggml_tensor* input(ggml_context* ctx, std::vector<ggml_tensor*>& in, ggml_type type, int64_t a, int64_t b = 1,
                   int64_t c = 1, int64_t d = 1) {
    ggml_tensor* t = ggml_new_tensor_4d(ctx, type, a, b, c, d);
    in.push_back(t);
    return t;
}

}  // namespace

int main(int argc, char** argv) {
    const std::string op = argc > 1 ? argv[1] : "all";
    const int reps = argc > 2 ? std::atoi(argv[2]) : 50;
    ggml_backend_t backend = ggml_backend_cuda_init(0);
    if (backend == nullptr) {
        std::fprintf(stderr, "no CUDA device\n");
        return 2;
    }
    const bool all = op == "all";
    double last = 0;

    // the FFN bias add: [mlp, S, B] + [mlp]
    if (all || op == "add_bias") {
        last = run(backend, "add_bias", reps, [](ggml_context* c, std::vector<ggml_tensor*>& in) {
            return ggml_add(c, input(c, in, GGML_TYPE_F32, kMlp, kS, kB), input(c, in, GGML_TYPE_F32, kMlp));
        });
    }
    // the residual add: [dim, S, B] + [dim, S, B]
    if (all || op == "add_res") {
        last = run(backend, "add_res", reps, [](ggml_context* c, std::vector<ggml_tensor*>& in) {
            return ggml_add(c, input(c, in, GGML_TYPE_F32, kDim, kS, kB), input(c, in, GGML_TYPE_F32, kDim, kS, kB));
        });
    }
    // the FFN activation: GELU (erf) over [mlp, S, B]
    if (all || op == "gelu_erf") {
        last = run(backend, "gelu_erf", reps, [](ggml_context* c, std::vector<ggml_tensor*>& in) {
            return ggml_gelu_erf(c, input(c, in, GGML_TYPE_F32, kMlp, kS, kB));
        });
    }
    // the latent and action-token concat: [18, tokens, frames, B] ++ [10, tokens, frames, B] on dim 0
    if (all || op == "concat") {
        last = run(backend, "concat", reps, [](ggml_context* c, std::vector<ggml_tensor*>& in) {
            return ggml_concat(c, input(c, in, GGML_TYPE_F32, 18, kTokens, kFrames, kB),
                               input(c, in, GGML_TYPE_F32, 10, kTokens, kFrames, kB), 0);
        });
    }
    // flash attention over strided Q/K/V views of one fused [3*inner, S, B]
    // projection, with a block-causal F16 mask repeated over the batch
    if (all || op == "flash_attn") {
        last = run(
            backend, "flash_attn", reps,
            [](ggml_context* c, std::vector<ggml_tensor*>& in) {
                ggml_tensor* qkv = input(c, in, GGML_TYPE_F32, 3 * kInner, kS, kB);
                auto split = [&](int64_t part) {
                    ggml_tensor* v = ggml_view_4d(c, qkv, kHeadDim, kHeads, kS, kB, kHeadDim * qkv->nb[0],
                                                  qkv->nb[1], qkv->nb[2], part * kInner * qkv->nb[0]);
                    return ggml_permute(c, v, 0, 2, 1, 3);
                };
                ggml_tensor* mask = input(c, in, GGML_TYPE_F16, kS, kS, 1, kB);
                return ggml_flash_attn_ext(c, split(0), split(1), split(2), mask, 1.0f / std::sqrt(float(kHeadDim)),
                                           0.0f, 0.0f);
            },
            [](ggml_tensor* mask) {
                std::vector<ggml_fp16_t> m(ggml_nelements(mask));
                for (int64_t b = 0; b < kB; ++b) {
                    for (int64_t q = 0; q < kS; ++q) {
                        for (int64_t k = 0; k < kS; ++k) {
                            const bool blocked = k / kTokens > q / kTokens;
                            m[(b * kS + q) * kS + k] = ggml_fp32_to_fp16(blocked ? -INFINITY : 0.0f);
                        }
                    }
                }
                ggml_backend_tensor_set(mask, m.data(), 0, ggml_nbytes(mask));
            });
    }
    ggml_backend_free(backend);
    return last > 0 ? 0 : 1;
}