ggml: allow passing alloc dependencies in graph_optimize

[am17an]

Overview

Currently a backend doesn't have a good way to tell the allocator to not re-use memory for tensor operations which can be used for fusion or same-gpu parallelization (concurrent streams), this necessitates backends to work around this by implementing memory checks + also disabling fusion when memory checks fail. Some fusions never run because of this (e.g. #25952 (comment), #21897),

This PR introduces an extension to the graph_optimize API which can used by backends to add allocator dependencies in graph_optimize, e.g.

  static void ggml_backend_mybackend_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph,
                                                    ggml_backend_graph_optimize_params * params) {
      // ... detect fork/join regions, decide stream assignment ...

      for (const auto & branch : branches) {
          for (const ggml_tensor * n : branch) {
              // the node itself: another branch must not reuse its memory before the join
              params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(n), join_node);

              // its sources: e.g. attn_norm's output is read by all three branches concurrently,
              // and inp_pos is read by Qrope and Krope - none of them may be recycled mid-region
              for (int i = 0; i < GGML_MAX_SRC; ++i) {
                  if (n->src[i]) {
                      params->add_alloc_dep(params->user_data, n->src[i], join_node);
                  }
              }
          }
      }
  }

This can be also be used for fusion to avoid creating scratch buffers:

e.g. for MoE reduce in #25952

  params->add_alloc_dep(params->user_data, experts,      reduce_node);
  params->add_alloc_dep(params->user_data, weights,      reduce_node);
  params->add_alloc_dep(params->user_data, expert_scale, reduce_node);

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[anujj]

I tested #25952 locally with this PR, and this PR solves the lifetime issue cleanly. I registered experts, weights, and the optional expert_scale as allocation dependencies.

[anujj]

@ggerganov: did you get chance to review this PR ?

[ggerganov]

@anujj For my understanding, how did you test it? The PR currently does not provide any dependency information from the CUDA backend.

[ggerganov]

Before merging, can we get a sample fusion implementation based on this mechanism to understand better how it would work in the a real case?

[am17an]

I've put them in the PR desc earlier, do you want to me actually add one here?

[ggerganov]

Not needed to add to this PR. If you have a branch with #21897 for example using this new dependency information?

[am17an]

@ggerganov in 56339ee I have made the changes to use this PR and enable concurrent streams

[am17an]

cc @jeffbolznv if you have any thoughts before we merge

[max-krasnyansky]

Looks good to me. We might have a use for this soon in ggml-hexagon

[jeffbolznv]

It's not immediately obvious to me in which situations graph_optimize should call it vs when we can get away without doing it, or if we should just be calling it for all fusions. I also think the new view nodes have a chance of breaking fusions if they're not carefully placed. e.g. vulkan's graph_optimize goes to some effort to bunch view nodes together into groups so they aren't interleaved with non-empty ops, and using this callback will potentially scatter more views in between real ops. Like, if I have fusions for both rms_norm+mul and rms_norm+mul+rope, then I don't want to make mul depend on rms_norm and end up with a graph of rms_norm+mul+view+rope that breaks my rms_norm+mul+rope fusion.

I think when I had first suggested something like this in the past, my idea was to add the dependencies as additional operands to the existing nodes rather than additional nodes. That would avoid the issue of the additional views getting in the way. But it has drawbacks, too, like needing to modify tensors that should be treated as const.

[ggerganov]

Maybe an alternative approach is to carry the dependencies as part of ggml_cgraph instead of trying to modify or add new nodes to express those dependencies. The allocator can be adjusted to respect the dependencies from ggml_cgraph. (note I didn't look if it makes sense implementation-wise, just a thought to consider)

[am17an]

I think it can be refactored to do that if we find big issues with the current approach. For now as its opt-in I don't see issues doing this since it's already an established pattern in the scheduler.

I plan to add this for big perf improvement fusions like topk-moe and some others and let the rest just be as they are today.