CUDA: update captured graph when a source tensor's shape changes#2096
CUDA: update captured graph when a source tensor's shape changes#2096fibrahimov wants to merge 1 commit into
Conversation
ggml_graph_node_has_matching_properties only compared each node's own ne/nb, the source data pointers, and op_params (for SCALE only). It never looked at the sources' ne/nb. During decode n_kv grows, so the KV-cache views feeding attention change shape while keeping the same data pointer. For the V*softmax matmul the output shape is independent of n_kv, so nothing the check looked at changed and the stale captured graph got replayed with the old n_kv baked into the kernel, producing garbage output (seen with MiniMax-M2). Store each source's ne/nb in the node properties and compare them, and compare op_params for every op instead of only SCALE. This forces a graph update only when a source shape or op params actually change (i.e. when the KV length crosses a padding boundary), so steady-state decode still reuses the graph. The CPY/VIEW data-pointer exclusion is kept so copy indirection is unaffected.
| // n_kv grows). The consuming node's own ne/nb may be n_kv-independent (the V*softmax | ||
| // "kqv" matmul is the canonical case), so relying on the node's own ne/nb is not | ||
| // enough: we must also detect changes in each source's ne/nb, otherwise a stale | ||
| // captured kernel (with the old n_kv baked in) gets replayed and produces garbage. |
There was a problem hiding this comment.
Interesting observations. They do apply to all model architectures, don't they? Hence, we should be observing "stale captured kernels" being used all the time. Which means that we should be able to reproduce the reported gibberish with any model, not just MiniMax-M2 (which I happen to not have lying around on my disk, and don't feel like downloading due to the large size). Can you provide a repro with a smaller model? This should be really easy, given the convincing arguments made in the comment.
There was a problem hiding this comment.
I was focused on fixing MiniMaxM2 and assumed the same would apply to any model that was a guess, not something I checked. I tried small dense models (Qwen2.5-3B, Llama-3.2-3B) and can't reproduce it there graphs on/off give identical greedy output, even across the 256-token KV boundary. So the attention/n_kv story isn't the real cause. On MiniMaxM2 it does reproduce (graphs on -> gibberish, graphs off -> clean, patch -> clean). The stale reuse is on the fused MoE combine GGML_OP_MUL_MULTI_ADD. Its output stays [n_embd, n_tokens], but a source shape changes 256 -> 8 (n_expert -> n_expert_used) between warmup and the first decode, and the old matcher didn't check source ne/nb. The 256 is from warmup — llama-build-context.cpp uses n_expert_used = n_expert there. So it captures an all experts graph, then reuses it for a top-8 decode. --no-warmup makes it go away, so does GGML_CUDA_DISABLE_GRAPHS=1. So it's narrower than I claimed not dense attention, and not even every MoE (Granite-3.0-3B-A800M didn't repro). MiniMaxM2 is the one I can reproduce reliably. Fix still looks right though checking source ne/nb + op_params for all ops catches the warmup/decode mismatch.
ggml_graph_node_has_matching_properties only compared each node's own
ne/nb, the source data pointers, and op_params (for SCALE only). It never
looked at the sources' ne/nb. During decode n_kv grows, so the KV-cache
views feeding attention change shape while keeping the same data pointer.
For the V*softmax matmul the output shape is independent of n_kv, so
nothing the check looked at changed and the stale captured graph got
replayed with the old n_kv baked into the kernel, producing garbage
output (seen with MiniMax-M2).
Store each source's ne/nb in the node properties and compare them, and
compare op_params for every op instead of only SCALE. This forces a graph
update only when a source shape or op params actually change (i.e. when
the KV length crosses a padding boundary), so steady-state decode still
reuses the graph. The CPY/VIEW data-pointer exclusion is kept so copy
indirection is unaffected.
Tested on 3x RTX A6000, CUDA 13.1, MiniMax-M2 Q8_0, experts split across the 3 GPUs + CPU, FA on, q8_0 KV:
The capital of France is->a,8 conmem conmem...(garbage), only usable with GGML_CUDA_DISABLE_GRAPHS=1Paris., coherent output incl. non-English, 320-token generations stay clean across padding boundariesNot specific to MiniMax — it affects any model where an attention matmul's output shape is independent of the growing KV length.
Likely fixes the gibberish reported in #1642.