CUDA: support non-F32 concat (I32/F16/BF16)

[0xSufi]

The CUDA concat op asserted F32-only; DeepSeek-V4 concats int32 (CSA indexer) and F16/BF16 (compressed KV), crashing at real context lengths. Make concat element-size-generic (raw byte copy), keeping the F32/I32 4-byte fast path.

Overview

Title

CUDA: support non-F32 concat (I32/F16/BF16) — fixes DeepSeek-V4 crash at real context lengths

Summary

The CUDA concat op hard-asserts F32, but the DeepSeek-V4 graph concatenates int32 (CSA indexer indices) and F16/BF16 (compressed-KV) tensors. As a result DS4 aborts on CUDA as soon as the compressed-attention path activates (i.e. any non-trivial context length). This makes concat element-size-generic so those concats work, while leaving the F32 fast path untouched.

Repro

Run any DeepSeek-V4-Flash GGUF on the CUDA backend with a prompt long enough to trigger the CSA compressed path (a few hundred+ tokens):

llama-cli -m DeepSeek-V4-Flash-*.gguf -ngl 999 -c 2048 -p "<~1k-token prompt>"

Abort:

ggml/src/ggml-cuda/concat.cu:165: GGML_ASSERT(src0->type == GGML_TYPE_F32) failed
  ... ggml_cuda_op_concat -> ggml_cuda_compute_forward -> llama_decode

A very short prompt does not trigger it (the compressed path doesn't engage), which is why it can pass smoke tests but fail on real input.

Root cause

ggml_cuda_op_concat and concat_f32_non_cont assume F32 (4-byte) elements: the assert restricts to GGML_TYPE_F32, the contiguous path uses float* with /4, and the non-contiguous kernel copies a single float. DeepSeek-V4 emits:

• I32 concat — CSA indexer (cat([window_idxs, compressed_idxs])), 4-byte.
• F16/BF16 concat — compressed-KV, 2-byte.

Fix

Concat is pure placement (no arithmetic), so any type can be moved as raw bytes:

• Relax the assert from "F32-only" to "all three tensors share one type."
• Keep the existing F32/I32 (4-byte) contiguous fast path unchanged.
• Route everything else (2-byte types, or non-contiguous) through the non-contiguous kernel, which now copies ggml_type_size(dst->type) bytes per element instead of a hardcoded float.

F32 behavior is bit-identical; I32/F16/BF16 are now supported. The 2-byte/non-contiguous case uses the existing "slow" non-contiguous kernel (acceptable — it was already the fallback path).

Testing

• Before: DS4-Flash (q2 GGUF) aborts at concat.cu:165 on first decode of a 2048-token window (2×RTX 3090, CUDA 12.4, sm_86).
• After: full forward completes — a 43-layer DS4-V4 forward over 28k tokens runs clean on CUDA (verified via an expert-routing eval-callback pass; all 43 layers produce output, no abort).
• Existing F32 concat paths are unchanged (fast path preserved; non-cont kernel produces identical bytes for es==4).
• Suggested follow-up (not in this diff): add I32/F16 cases to tests/test-backend-ops for concat so this is covered going forward — happy to add if you'd like.

Diff

@@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE)
           uint64_t   nb2,
-          uint64_t   nb3){
+          uint64_t   nb3,
+               int   es){
     static_assert(dim >= 0 && dim <= 3, "dim must be in [0, 3]");
 ...
-    const float * x;
+    const char * x;
     for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) {
         if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
-            x = (const float *)(src0 + (i3)*nb03 + (i2)*nb02 + (i1)*nb01 + (i0)*nb00);
+            x = (const char *)(src0 + (i3)*nb03 + (i2)*nb02 + (i1)*nb01 + (i0)*nb00);
         } else {
-            if constexpr (dim == 0) { x = (const float *) (src1 + ... (i0 - ne00) * nb10); }
-            else if (dim == 1) { x = (const float *) (src1 + ... (i1 - ne01) * nb11 ...); }
-            else if (dim == 2) { x = (const float *) (src1 + ... (i2 - ne02) * nb12 ...); }
-            else if (dim == 3) { x = (const float *) (src1 + (i3 - ne03) * nb13 ...); }
+            if constexpr (dim == 0) { x = (const char *) (src1 + ... (i0 - ne00) * nb10); }
+            else if (dim == 1) { x = (const char *) (src1 + ... (i1 - ne01) * nb11 ...); }
+            else if (dim == 2) { x = (const char *) (src1 + ... (i2 - ne02) * nb12 ...); }
+            else if (dim == 3) { x = (const char *) (src1 + (i3 - ne03) * nb13 ...); }
         }
-        float * y = (float *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
-        *y = *x;
+        char * y = (char *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
+        for (int b = 0; b < es; ++b) y[b] = x[b];   // element-size-generic copy (I32/F16/BF16/F32)
     }
 }

@@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT(src1->type == GGML_TYPE_F32);
-    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    // DS4: the CSA path concats I32 (indices) and F16/BF16 (compressed KV), not just F32.
+    // Concat is pure placement, so copy raw element-size bytes. Require matching types;
+    // keep the F32/I32 (4-byte) contiguous fast path, route the rest through the generic kernel.
+    GGML_ASSERT(src1->type == src0->type);
+    GGML_ASSERT(dst->type  == src0->type);

-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
+    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_type_size(dst->type) == 4) {
         ... existing 4-byte fast path (unchanged) ...
     } else {
         ... non-contiguous launch, now passing element size:
-                dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3]);
+                dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3],
+                (int) ggml_type_size(dst->type));
     }

(Abbreviated for readability — the authoritative diff is git diff ggml/src/ggml-cuda/concat.cu in the working tree.)

Additional information

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: this change was developed with AI assistance (Claude Code / Claude Opus). The AI diagnosed the abort, wrote the patch and this PR description, and ran the before/after tests on 2× RTX 3090 box. I've reviewed the diff, understand the change, confirmed it builds, and verified DeepSeek-V4-Flash runs on CUDA with it. I'll own and address any review feedback.