ggml : add ggml_build_forward_select

[ggerganov]

target #18547
alt #18549

• Add GGML_TENSOR_FLAG_COMPUTE flag indicating that a tensor in the graph must be computed
• Add new ggml_build_forward_select() call:

    GGML_API struct ggml_tensor * ggml_build_forward_select(
            struct ggml_cgraph  * cgraph,
            struct ggml_tensor ** tensors,
            int                   n_tensors,
            int                   idx);

All provided tensors are built forward into the graph. Only tensors[idx] and it's ancestry are marked for computing via the new flag value.

This new logic allows us to construct graphs that compute different things, but at the same time have the same topology. This is needed to avoid unwanted graph reallocations (#17617).

TODOs:

• Check flag for CUDA graph updates (ggml : add ggml_build_forward_select #18550 (comment))
• Update fusion logic? (ggml : add ggml_build_forward_select #18550 (comment))
• Enable -DGGML_SCHED_NO_REALLOC=ON for server CI
• Fix vulkan command submission logic ggml : add ggml_build_forward_select #18550 (comment)

[jeffbolznv]

Just want to make sure I understand how this is used - it would still be two separate graphs, they'd just be able to reuse allocations (i.e. ggml-alloc would decide they match)?

I think ggml_can_fuse and ggml_can_fuse_subgroup would need to be updated to make sure all nodes are computed. And any backend-specific fusion logic.

[ggerganov]

Just want to make sure I understand how this is used - it would still be two separate graphs, they'd just be able to reuse allocations (i.e. ggml-alloc would decide they match)?

Yes, for example the graph when the input is token ids (batch.token != null) and the graph when the input is directly embedding vectors (batch.embd != null) are still different, but with this extra logic the scheduler will not need to reallocate them because all nodes remain the same. It's just a different subset of the nodes being marked for computing.

I think ggml_can_fuse and ggml_can_fuse_subgroup would need to be updated to make sure all nodes are computed. And any backend-specific fusion logic.

Not yet sure that it's really necessary to do so - at least I can't think of a fail case so far. Note that the GGML_TENSOR_FLAG_COMPUTE flag is controlled only through ggml_build_forward_select().

[max-krasnyansky]

Looks good to me.

[am17an]

We would need to check how this behaves with CUDA graphs, since inherently the computation is changing

[taronaeo]

cc: @AlekseiNikiforovIBM @Andreas-Krebbel

Give us a week or so to check on this :)

[am17an]

For CUDA graphs I think adding a check for flags in ggml_graph_node_has_matching_properties should be enough. This would trigger an update to the graph

[AlekseiNikiforovIBM]

cc: @AlekseiNikiforovIBM @Andreas-Krebbel

Give us a week or so to check on this :)

LGTM

[taronaeo]

Ack for IBM zDNN backend :)

[0cc4m]

If the last node or nodes are not flagged, the loop would end without the final command submission. This would need some way to ensure a final submit if submitted_nodes > 0.

[reeselevine]

WebGPU update looks good to me, we always do a final submission if commands > 0 so there shouldn't be a problem like the comment about the Vulkan backend above.