ggml-cuda: GEMM weight row padding + one-time K-padded f16 dequant for prefill

ggml/include/ggml.h

@@ -649,6 +649,7 @@ extern "C" {
         GGML_TENSOR_FLAG_PARAM   =  4, // ...contains trainable parameters
         GGML_TENSOR_FLAG_LOSS    =  8, // ...defines loss for numerical optimization (multiple loss tensors add up)
         GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed
+        GGML_TENSOR_FLAG_PAD_ROWS = 32, // ...is a GEMM weight whose row stride a backend may pad to avoid cache-set aliasing
     };
 
     enum ggml_tri_type {

ggml/src/ggml-cuda/common.cuh

@@ -177,6 +177,15 @@ static int ggml_cuda_highest_compiled_arch(const int arch) {
 
 #define GGML_CUDA_MAX_STREAMS 8
 
+// True when each row is internally contiguous and the higher dimensions are tightly nested,
+// while the row stride nb[1] may exceed the packed row size. Such a tensor can be fed to a
+// strided GEMM (leading dimension = nb[1] / type_size) without repacking to contiguous.
+static inline bool ggml_cuda_is_contiguous_rows(const ggml_tensor * t) {
+    return t->nb[0] == ggml_type_size(t->type)
+        && t->nb[2] == t->nb[1]*t->ne[1]
+        && t->nb[3] == t->nb[2]*t->ne[2];
+}
+
 [[noreturn]]
 void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg);
 

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -663,17 +663,33 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() {
 
 // cuda buffer
 
+// A persistent f16 copy of a quantized dense weight, built once at load time (see
+// ggml_backend_cuda_buffer_set_tensor) and consumed by the prefill cuBLAS path instead of
+// dequantizing on the fly. `lda` is the row stride in f16 elements (K-padded when it avoids
+// cache-set aliasing). Owned by the buffer context so it frees with the model.
+struct ggml_cuda_f16_shadow {
+    half *  ptr = nullptr;
+    int64_t lda = 0;
+};
+
 struct ggml_backend_cuda_buffer_context {
     int device;
     void * dev_ptr = nullptr;
     std::string name;
+    std::vector<ggml_cuda_f16_shadow *> shadows;
 
     ggml_backend_cuda_buffer_context(int device, void * dev_ptr) :
         device(device), dev_ptr(dev_ptr),
         name(GGML_CUDA_NAME + std::to_string(device)) {
     }
 
     ~ggml_backend_cuda_buffer_context() {
+        for (ggml_cuda_f16_shadow * s : shadows) {
+            if (s->ptr) {
+                CUDA_CHECK(cudaFree(s->ptr));
+            }
+            delete s;
+        }
         CUDA_CHECK(cudaFree(dev_ptr));
     }
 };
@@ -692,6 +708,39 @@ static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) {
     return ctx->dev_ptr;
 }
 
+// One cache line, added to a weight's row stride when the packed row size is a multiple of
+// 2048 bytes to avoid cache-set aliasing in matmul. Enabled by default; disable with
+// GGML_CUDA_NO_PAD_WEIGHTS.
+#define GGML_CUDA_CACHE_LINE 128
+#define GGML_CUDA_ROW_ALIAS_STRIDE 2048
+
+static bool ggml_cuda_pad_weights_enabled() {
+    static const bool enabled = getenv("GGML_CUDA_NO_PAD_WEIGHTS") == nullptr;
+    return enabled;
+}
+
+// A float GEMM weight (tagged GGML_TENSOR_FLAG_PAD_ROWS by the loader) whose packed row size
+// aliases at GGML_CUDA_ROW_ALIAS_STRIDE and therefore benefits from a one-cache-line row-stride
+// pad. Quantized weights are excluded: a 128-byte pad is not a multiple of their block size and
+// would misalign the block-indexed matmul kernels.
+static bool ggml_cuda_should_pad_weight(const ggml_tensor * tensor) {
+    if (!ggml_cuda_pad_weights_enabled()) {
+        return false;
+    }
+    if (tensor->view_src != nullptr || !(tensor->flags & GGML_TENSOR_FLAG_PAD_ROWS)
+            || ggml_is_quantized(tensor->type) || tensor->ne[1] <= 1) {
+        return false;
+    }
+    return ggml_row_size(tensor->type, tensor->ne[0]) % GGML_CUDA_ROW_ALIAS_STRIDE == 0;
+}
+
+// Row stride (nb[1]) in bytes for a weight tensor: the packed row size, plus one cache line
+// when padding applies.
+static size_t ggml_cuda_padded_nb1(const ggml_tensor * tensor) {
+    const size_t packed = ggml_row_size(tensor->type, tensor->ne[0]);
+    return ggml_cuda_should_pad_weight(tensor) ? packed + GGML_CUDA_CACHE_LINE : packed;
+}
+
 static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
     ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
 
@@ -700,6 +749,19 @@ static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer
         return GGML_STATUS_SUCCESS;
     }
 
+    if (ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE
+            && ggml_cuda_should_pad_weight(tensor)) {
+        tensor->nb[1] = ggml_cuda_padded_nb1(tensor);
+        tensor->nb[2] = tensor->nb[1] * tensor->ne[1];
+        tensor->nb[3] = tensor->nb[2] * tensor->ne[2];
+
+        // zero the allocation so the inter-row padding gaps hold no NaNs
+        ggml_cuda_set_device(ctx->device);
+        CUDA_CHECK(cudaMemset((char *) tensor->data, 0,
+            ggml_backend_buft_get_alloc_size(buffer->buft, tensor)));
+        return GGML_STATUS_SUCCESS;
+    }
+
     if (ggml_is_quantized(tensor->type) && tensor->view_src == nullptr && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
         // initialize padding to 0 to avoid possible NaN values
         const size_t original_size = ggml_nbytes(tensor);
@@ -721,12 +783,72 @@ static void ggml_backend_cuda_buffer_memset_tensor(ggml_backend_buffer_t buffer,
     CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
 }
 
+// Build a one-time f16 shadow of a freshly-uploaded quantized dense weight (opt-in via
+// GGML_PREFILL_DEQUANT). Runs after the quantized bytes are resident in VRAM, so the prefill
+// cuBLAS path (ggml_cuda_mul_mat_cublas_impl) can consume tensor->extra instead of
+// dequantizing on every call. Dense 2D weights only; MoE expert stacks (ne2>1) are excluded.
+static void ggml_cuda_maybe_build_f16_shadow(
+        ggml_backend_cuda_buffer_context * ctx, ggml_backend_buffer_t buffer,
+        ggml_tensor * tensor, size_t offset, size_t size) {
+    static const bool enabled = getenv("GGML_PREFILL_DEQUANT") != nullptr;
+    if (!enabled || tensor->extra != nullptr || tensor->view_src != nullptr) {
+        return;
+    }
+    if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE
+            || !ggml_is_quantized(tensor->type)
+            || tensor->ne[1] <= 1 || tensor->ne[2] != 1 || tensor->ne[3] != 1
+            || !ggml_is_contiguously_allocated(tensor)
+            || offset != 0 || size != ggml_nbytes(tensor)) {
+        return;
+    }
+    const int cc = ggml_cuda_info().devices[ctx->device].cc;
+    if (!fast_fp16_hardware_available(cc)) {
+        return;
+    }
+    const to_fp16_cuda_t to_fp16 = ggml_get_to_fp16_cuda(tensor->type);
+    if (to_fp16 == nullptr) {
+        return;
+    }
+
+    const int64_t ne00 = tensor->ne[0];
+    const int64_t ne01 = tensor->ne[1];
+    // pad the f16 row stride by one cache line only when the f16 row size aliases at 2048 B
+    const int64_t pad  = (ggml_row_size(GGML_TYPE_F16, ne00) % GGML_CUDA_ROW_ALIAS_STRIDE == 0)
+        ? (int64_t) (GGML_CUDA_CACHE_LINE / sizeof(half)) : 0;
+    const int64_t lda  = ne00 + pad;
+
+    ggml_cuda_set_device(ctx->device);
+    ggml_cuda_f16_shadow * sh = new ggml_cuda_f16_shadow();
+    sh->lda = lda;
+    CUDA_CHECK(cudaMalloc(&sh->ptr, (size_t) ne01 * lda * sizeof(half)));
+
+    if (pad == 0) {
+        to_fp16(tensor->data, sh->ptr, ggml_nelements(tensor), cudaStreamPerThread);
+    } else {
+        half * tmp = nullptr;
+        CUDA_CHECK(cudaMalloc(&tmp, (size_t) ne01 * ne00 * sizeof(half)));
+        to_fp16(tensor->data, tmp, ggml_nelements(tensor), cudaStreamPerThread);
+        CUDA_CHECK(cudaMemcpy2DAsync(sh->ptr, lda * sizeof(half),
+                                     tmp, ne00 * sizeof(half),
+                                     ne00 * sizeof(half), ne01,
+                                     cudaMemcpyDeviceToDevice, cudaStreamPerThread));
+        CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+        CUDA_CHECK(cudaFree(tmp));
+    }
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+
+    tensor->extra = sh;
+    ctx->shadows.push_back(sh);
+}
+
 static void ggml_backend_cuda_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context;
 
     ggml_cuda_set_device(ctx->device);
     CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cudaStreamPerThread));
     CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+
+    ggml_cuda_maybe_build_f16_shadow(ctx, buffer, tensor, offset, size);
 }
 
 static void ggml_backend_cuda_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
@@ -856,6 +978,16 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t
         }
     }
 
+    // K-dimension cache-line padding: reserve the padded per-row stride for weight matrices.
+    if (ggml_cuda_should_pad_weight(tensor)) {
+        const size_t   padded_nb1 = ggml_cuda_padded_nb1(tensor);
+        const int64_t  nrows      = tensor->ne[1] * tensor->ne[2] * tensor->ne[3];
+        const size_t   padded_sz  = padded_nb1 * nrows;
+        if (padded_sz > size) {
+            size = padded_sz;
+        }
+    }
+
     return size;
 }
 
@@ -1370,7 +1502,25 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
     bool is_src0_cont_2 = ggml_is_contiguous_2(src0);
     bool is_src1_cont_2 = ggml_is_contiguous_2(src1);
 
-    if (src0->type == compute_type) {
+    bool src0_cached = false;
+    if constexpr (compute_type == GGML_TYPE_F16) {
+        // GGML_PREFILL_DEQUANT: consume the persistent f16 shadow built once at load time
+        // (ggml_cuda_maybe_build_f16_shadow), instead of dequantizing on every call. Dense 2D
+        // weights only; view_src != nullptr excludes mul_mat_id expert slices.
+        if (ggml_is_quantized(src0->type) && src0->view_src == nullptr
+                && src0->extra != nullptr && ne02 == 1 && ne03 == 1) {
+            const ggml_cuda_f16_shadow * sh = (const ggml_cuda_f16_shadow *) src0->extra;
+            src0_ptr = (const cuda_t *) sh->ptr;
+            s01 = sh->lda;
+            s02 = ne01*s01;
+            s03 = ne02*s02;
+            src0_cached = true;
+        }
+    }
+
+    if (src0_cached) {
+        // src0_ptr / strides already set from the load-time f16 shadow
+    } else if (src0->type == compute_type) {
         src0_ptr = (const cuda_t *) src0->data;
     } else {
         src0_alloc.alloc(ggml_nelements(src0));
@@ -1759,6 +1909,19 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
     const int cc        = ggml_cuda_info().devices[ctx.device].cc;
     const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
 
+    // GGML_PREFILL_DEQUANT=<min_tokens>: for large-batch (prefill) quantized GEMMs that have a
+    // load-time f16 shadow (ggml_cuda_maybe_build_f16_shadow), skip MMQ/MMVQ and fall through to
+    // the hipBLAS path, which consumes the shadow (see ggml_cuda_mul_mat_cublas_impl). Weights
+    // without a shadow (e.g. mul_mat_id expert slices) stay on MMQ.
+    static const int prefill_dequant_min = [] {
+        const char * env = getenv("GGML_PREFILL_DEQUANT");
+        if (env == nullptr) return 0;
+        const int v = atoi(env);
+        return v > 1 ? v : 32;
+    }();
+    const bool prefill_dequant = prefill_dequant_min > 0 && ggml_is_quantized(src0->type)
+        && ne11 >= prefill_dequant_min && src0->view_src == nullptr && src0->extra != nullptr;
+
     if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) {
         // The custom F16 vector kernel can be used over batched cuBLAS GEMM.
         // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention)
@@ -1790,11 +1953,11 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
         ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst);
         return;
     }
-    if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
+    if (!prefill_dequant && ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
         ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
         return;
     }
-    if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
+    if (!prefill_dequant && ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
         ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
         return;
     }

ggml/src/ggml-cuda/mmvf.cu

@@ -740,8 +740,10 @@ void ggml_cuda_op_mul_mat_vec_f(
     const int cc = ggml_cuda_info().devices[id].cc;
     const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec(dst->op_params[0]) : GGML_PREC_F32;
 
-    // ggml_cuda_op provides single, contiguous matrices
-    const int64_t stride_row         = ne00;
+    // honor a padded row stride (nb[1]) when src0 is passed through un-copied; equals ne00 for
+    // the packed case
+    const int64_t stride_row         = ggml_cuda_is_contiguous_rows(src0)
+        ? (int64_t) (src0->nb[1] / ggml_type_size(src0->type)) : ne00;
     const int64_t stride_col_y       = ne10;
     const int64_t stride_col_dst     = id == ctx.device ? ne0 : row_diff; // main device has larger memory buffer
     const int64_t nchannels_x        = 1;

src/llama-model-loader.cpp

@@ -1284,6 +1284,20 @@ struct ggml_tensor * llama_model_loader::create_tensor(
     struct ggml_tensor * tensor = ggml_dup_tensor(ctx, cur);
     ggml_set_name(tensor, ggml_get_name(cur));
 
+    // Mark GEMM weights (used as mul_mat / mul_mat_id src0) so backends may pad their row
+    // stride to avoid cache-set aliasing. Only real weights ("weight" suffix), not bias/scale.
+    {
+        llm_tensor weight_tensor = tn.tensor;
+        if (tn.tensor == LLM_TENSOR_TOKEN_EMBD && (flags & TENSOR_DUPLICATED)) {
+            weight_tensor = LLM_TENSOR_OUTPUT;
+        }
+        const bool is_weight = tn.suffix == nullptr || strcmp(tn.suffix, "weight") == 0;
+        const ggml_op weight_op = llm_tensor_info_for(weight_tensor).op;
+        if (is_weight && (weight_op == GGML_OP_MUL_MAT || weight_op == GGML_OP_MUL_MAT_ID)) {
+            tensor->flags |= GGML_TENSOR_FLAG_PAD_ROWS;
+        }
+    }
+
     if (duplicated) {
         size_data += ggml_nbytes(cur);
     } else {
@@ -1542,6 +1556,16 @@ bool llama_model_loader::load_all_data(
 
         size_t n_size = ggml_nbytes(cur);
 
+        // K-dim cache-line padding: the on-disk data is packed, but the destination tensor
+        // may carry a padded per-row stride (nb[1]). Read the packed size and upload the
+        // rows with a strided 2D copy.
+        const size_t  packed_row = ggml_row_size(cur->type, cur->ne[0]);
+        const int64_t n_rows     = ggml_nelements(cur) / cur->ne[0];
+        const bool    row_padded = (size_t) cur->nb[1] != packed_row;
+        if (row_padded) {
+            n_size = packed_row * n_rows;
+        }
+
         if (use_mmap) {
             const auto & mapping = mappings.at(weight->idx);
             ggml_backend_buffer_t buf_mmap = nullptr;
@@ -1567,6 +1591,8 @@ bool llama_model_loader::load_all_data(
                 auto & mmap_used = mmaps_used[weight->idx];
                 mmap_used.first  = std::min(mmap_used.first,  weight->offs);
                 mmap_used.second = std::max(mmap_used.second, weight->offs + n_size);
+            } else if (row_padded) {
+                ggml_backend_tensor_set_2d(cur, data, 0, packed_row, n_rows, cur->nb[1], packed_row);
             } else {
                 ggml_backend_tensor_set(cur, data, 0, n_size);
             }
@@ -1583,7 +1609,7 @@ bool llama_model_loader::load_all_data(
                 }
             } else {
                 // If upload_backend is valid load the tensor in chunks to pinned memory and upload the buffers asynchronously to the GPU.
-                if (upload_backend) {
+                if (upload_backend && !row_padded) {
                     size_t offset = weight->offs;
                     alignment = file->read_alignment();
                     size_t aligned_offset = offset & ~(alignment - 1);
@@ -1639,7 +1665,11 @@ bool llama_model_loader::load_all_data(
                     read_buf.resize(n_size);
                     file->seek(weight->offs, SEEK_SET);
                     file->read_raw(read_buf.data(), n_size);
-                    ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
+                    if (row_padded) {
+                        ggml_backend_tensor_set_2d(cur, read_buf.data(), 0, packed_row, n_rows, cur->nb[1], packed_row);
+                    } else {
+                        ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
+                    }
                     if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
                         throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
                     }