sm graph: reshape folded-head MLA tensors to 3D before the per-head split (#2130)

Loading a deepseek2 model with -sm graph across 2 or more GPUs segfaults
during model load, in the per-head split of the MLA attn_k_b/attn_v_b tensors.

In the ik-converted GGUFs these ship 2D with the heads folded into ne[1]
(DSV2-Lite attn_k_b is [128, 8192]). distribute_mla_tensors_for_split_mode_graph
splits them with prepare_split_tensors(2, ...), which takes the head count
from ne[2]. With the heads still in ne[1] each split comes out n_head times
too large and the upload reads past the end of the host buffer. It is
host-side (compute-sanitizer is clean) and happens at any device count >= 2.

Fix: reshape the 2D layout to 3D [head_dim, inner, n_head] before the split,
the same head-outermost interpretation ggml_reshape_3d already uses on the
non-split path. Metadata-only, guarded on ne[2]==1 so it is a no-op for GGUFs
that already ship these tensors 3D.

Tested on DeepSeek-V2-Lite-Chat Q4_K_M (only what fits across the GPUs here):
-sm graph PPL matches -sm layer and CPU within stderr (7.26 / 7.24 / 7.24,
wikitext-2 n_ctx 2048), coherent generation at 2, 3 and 4 GPU splits, and a
REAP-pruned DSV2-Lite in parity too. GLM-DSA, Mistral4 and larger deepseek
GGUFs share the split path so I expect the same fix to cover them, untested here.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/llama-load-tensors.cpp

@@ -4179,6 +4179,24 @@ static void prepare_split_tensors(int split_dim, ggml_context * ctx, ggml_tensor
     }
 }
 
+// GGUF may store attn_k_b/attn_v_b as a 2D tensor [head_dim, n_head*inner] (the ik DSV2-Lite
+// converter folds all heads into ne[1]) instead of the 3D per-head form [head_dim, inner, n_head]
+// used by the non-lite / GLM-DSA create paths. The MLA distribution below splits these per head
+// (split_dim=2), which requires the head count to live in ne[2]. If it is left folded into ne[1],
+// prepare_split_tensors builds [head_dim, n_head*inner, split_heads] split shapes — n_head times
+// too large — and the split_dim==2 set_tensor loop reads n_head*ggml_nbytes(tensor) bytes off the
+// host source, running past the end of the loaded tensor (host OOB -> SIGSEGV during model load).
+// Reinterpret the contiguous 2D layout as 3D in place: heads are the outermost stride of the folded
+// ne[1], so this is a pure metadata reshape and ggml_nbytes is unchanged.
+static void reshape_folded_heads_to_3d(ggml_tensor * t, int n_head) {
+    if (!t || t->ne[2] != 1 || t->ne[3] != 1) return;   // already 3D (non-lite path) or N/A
+    GGML_ASSERT(n_head > 0 && t->ne[1] % n_head == 0);
+    t->ne[1] /= n_head;
+    t->ne[2]  = n_head;
+    t->nb[2]  = t->nb[1] * t->ne[1];
+    t->nb[3]  = t->nb[2] * t->ne[2];
+}
+
 // MLA tensor distribution for -sm graph / -sm attn.
 // q_a/wkv_a_mqa/norms replicated; q_b row-split by Q head; wo row-split.
 // wk_b/wv_b are per-head split (split_dim=2) — loaded directly when present
@@ -4231,14 +4249,17 @@ static void distribute_mla_tensors_for_split_mode_graph(
         prepare_split_tensors(1, ctx_split, layer.wq, layer.split_wq, split_wq_cols, mem_used);
     }
 
-    // wkv_a_mqa, wk_b, wv_b replicated: the per-head 3D batched mul_mat can't read a split src0.
+    // wkv_a_mqa is replicated (mirror): its per-head 3D batched mul_mat can't read a split src0.
+    // wk_b/wv_b are split per head (split_dim=2); reshape a 2D folded-head GGUF layout to 3D first.
     if (layer.wkv_a_mqa) {
         prepare_split_tensors(-1, ctx_split, layer.wkv_a_mqa, layer.split_wkv_a_mqa, mirror, mem_used);
     }
     if (layer.wk_b) {
+        reshape_folded_heads_to_3d(layer.wk_b, n_head);
         prepare_split_tensors( 2, ctx_split, layer.wk_b, layer.split_wk_b, split_heads, mem_used);
     }
     if (layer.wv_b) {
+        reshape_folded_heads_to_3d(layer.wv_b, n_head);
         prepare_split_tensors( 2, ctx_split, layer.wv_b, layer.split_wv_b, split_heads, mem_used);
     }