crash in ggml-cuda with rms_norm and mul fusion

[Codes4Fun]

I don't know exactly what this code does, but it looks like it is trying to access the second src of an rms_norm tensor, which only has one src:

ggml/src/ggml-cuda/ggml-cuda.cu

Line 2904 in fcc2a5c

if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm->src[1])) {

[slaren]

Yes, the intention was probably to check mul->src[1] instead. @am17an can you take a look?

[am17an]

Can you share the exact steps to reproduce this crash? This code is copied from the vulkan backend, so that will also need fixing in case this is an issue. cc: @jeffbolznv

[jeffbolznv]

The Vulkan code is different and not accessing rms_norm->src[1]. But it looks like the opencl backend has the same bug as the cuda backend.

[am17an]

This is what I was referring to, since this code is exactly copied in the CUDA backend.

https://github.com/ggml-org/llama.cpp/blob/e60f241eacec42d3bd7c9edd37d236ebf35132a8/ggml/src/ggml-vulkan/ggml-vulkan.cpp#L11952

[jeffbolznv]

Look again, the code is different and as far as I can tell the Vulkan backend never had this bug.

[am17an]

Apologies, I was looking at mul->src[1] being the culprit

[Codes4Fun]

I tested the change and it fixes it, and if a test case is still wanted:

void test_rms_norm() {
    // init backend
    //auto backend = ggml_backend_cpu_init();
    auto backend = ggml_backend_cuda_init(0);
    assert( backend );
    auto ctx = ggml_init({
        /*.mem_size   =*/ ggml_tensor_overhead()*GGML_DEFAULT_GRAPH_SIZE + ggml_graph_overhead(),
        /*.mem_buffer =*/ NULL,
        /*.no_alloc   =*/ true,
    });
    // create tensors
    std::vector<float> data = { 1.f, 2.f };
    auto input = ggml_new_tensor_1d( ctx, GGML_TYPE_F32, data.size() );
    auto result = ggml_mul( ctx, input, ggml_rms_norm( ctx, input, 1e-8 ) );
    // graph
    auto gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, result);
    // allocate tensors
    auto buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
    // initialize tensors
    ggml_backend_tensor_set(input, data.data(), 0, ggml_nbytes(input));
    // compute
    ggml_backend_graph_compute(backend, gf);
    // output results, should be 0.632456 2.529822
    ggml_backend_tensor_get(result, data.data(), 0, ggml_nbytes(result));
    printf("%f %f\npress enter to continue", data[0], data[1]);
    getchar();
    // cleanup
    ggml_backend_buffer_free(buffer);
    ggml_free(ctx);
    ggml_backend_free(backend);
}