Merge pull request #78 from AMD-Ecosystem/rogarcia.padd-gemm-weight-rows-to-avoid-cache-set

ggml-cuda: pad matmul weight rows to avoid cache-set aliasing

Author: roberteg16

ggml/include/ggml.h

@@ -648,11 +648,12 @@ extern "C" {
 
     // this tensor...
     enum ggml_tensor_flag {
-        GGML_TENSOR_FLAG_INPUT   =  1, // ...is an input for the GGML compute graph
-        GGML_TENSOR_FLAG_OUTPUT  =  2, // ...is an output for the GGML compute graph
-        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_INPUT    =  1, // ...is an input for the GGML compute graph
+        GGML_TENSOR_FLAG_OUTPUT   =  2, // ...is an output for the GGML compute graph
+        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 matmul weight whose row stride a backend may pad to avoid cache-set aliasing
     };
 
     enum ggml_tri_type {

ggml/src/ggml-cuda/common.cuh

@@ -1126,6 +1126,15 @@ struct ggml_cuda_type_traits<GGML_TYPE_IQ3_S> {
 
 //////////////////////
 
+// Row-stride padding for matrix multiplication weights. A weight whose packed row size is a
+// multiple of alias_stride puts every row in the same cache sets, so pad bytes are added to the
+// row stride to break the aliasing. Both values follow from the cache geometry (line size times
+// number of sets), which neither CUDA nor HIP report, so they are tabulated per architecture.
+struct ggml_cuda_row_pad_params {
+    size_t alias_stride; // packed row sizes that are a multiple of this alias in the cache
+    size_t pad;          // bytes added to the row stride; 0 disables the padding
+};
+
 struct ggml_cuda_device_info {
     int device_count;           // number of (possibly virtual) devices exposed to the rest of ggml
     int physical_device_count;  // number of physical CUDA devices actually present
@@ -1144,6 +1153,7 @@ struct ggml_cuda_device_info {
         int     physical_device;                // backing physical CUDA device for this (virtual) device
         int     physical_share_count;           // number of (virtual) devices sharing this device's physical GPU
         int     virtual_index;                  // index of this (virtual) device among those sharing its physical GPU
+        ggml_cuda_row_pad_params row_pad;       // weight row-stride padding
     };
 
     cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};

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

@@ -220,6 +220,71 @@ static int ggml_cuda_parse_id(char devName[]) {
 }
 #endif // defined(GGML_USE_HIP)
 
+// Weight row-stride padding per architecture. An architecture is opted in only once the padding
+// has been measured on it; every other one gets {0, 0} and keeps packed rows.
+static ggml_cuda_row_pad_params ggml_cuda_row_pad_params_for(int cc) {
+    if (GGML_CUDA_CC_IS_RDNA3_5(cc)) {
+        return { 2048, 128 }; // one cache line added to rows that alias every 2 KiB
+    }
+
+    return { 0, 0 };
+}
+
+// Returns the environment override for name, or -1 when it is unset or not a number.
+static int64_t ggml_cuda_row_pad_env(const char * name) {
+    const char * val = getenv(name);
+    if (!val || !val[0]) {
+        return -1;
+    }
+
+    char * end = nullptr;
+    const int64_t parsed = strtoll(val, &end, 0);
+    if (*end != '\0' || parsed < 0) {
+        GGML_LOG_WARN("%s: ignoring %s=%s, expected a non-negative integer\n", __func__, name, val);
+        return -1;
+    }
+
+    return parsed;
+}
+
+// Resolves the padding for one device: the architecture default, then the environment overrides,
+// then the sanity checks. Runs once per device so an invalid setting warns once, not per tensor.
+static ggml_cuda_row_pad_params ggml_cuda_resolve_row_pad(int cc) {
+    if (getenv("GGML_CUDA_NO_PAD_WEIGHTS")) {
+        return { 0, 0 };
+    }
+
+    ggml_cuda_row_pad_params params = ggml_cuda_row_pad_params_for(cc);
+
+    const int64_t alias_stride_override = ggml_cuda_row_pad_env("GGML_CUDA_ROW_ALIAS_STRIDE");
+    const int64_t pad_override          = ggml_cuda_row_pad_env("GGML_CUDA_ROW_PAD");
+    if (alias_stride_override >= 0) {
+        params.alias_stride = alias_stride_override;
+    }
+    if (pad_override >= 0) {
+        params.pad = pad_override;
+    }
+
+    if (params.pad == 0) {
+        return { 0, 0 };
+    }
+
+    if (params.alias_stride == 0 || (params.alias_stride & (params.alias_stride - 1)) != 0) {
+        GGML_LOG_WARN("%s: disabling weight row padding, alias stride %zu is not a power of two\n",
+                      __func__, params.alias_stride);
+        return { 0, 0 };
+    }
+
+    // a pad that is itself a multiple of the alias stride leaves the rows aliasing
+    if (params.pad % params.alias_stride == 0) {
+        GGML_LOG_WARN("%s: disabling weight row padding, pad %zu is a multiple of alias stride %zu\n",
+                      __func__, params.pad, params.alias_stride);
+        return { 0, 0 };
+    }
+
+    return params;
+}
+
 static ggml_cuda_device_info ggml_cuda_init() {
     ggml_cuda_device_info info = {};
 
@@ -376,6 +441,8 @@ static ggml_cuda_device_info ggml_cuda_init() {
         }
 
 #endif  // defined(GGML_USE_HIP)
+
+        info.devices[id].row_pad = ggml_cuda_resolve_row_pad(info.devices[id].cc);
     }
 
     if (ggml_cuda_highest_compiled_arch(GGML_CUDA_CC_TURING) >= GGML_CUDA_CC_TURING && !turing_devices_without_mma.empty()) {
@@ -761,6 +828,32 @@ static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) {
     return ctx->dev_ptr;
 }
 
+// Quantized weights are excluded: the matmul kernels read the row stride as nb[1]/type_size, an
+// exact division that a pad which is not a multiple of the block size would break.
+static bool ggml_cuda_should_pad_weight(const ggml_tensor * tensor, int device) {
+    const ggml_cuda_row_pad_params & row_pad = ggml_cuda_info().devices[device].row_pad;
+
+    if (row_pad.pad == 0 || !(tensor->flags & GGML_TENSOR_FLAG_PAD_ROWS)) {
+        return false;
+    }
+
+    if (tensor->view_src != nullptr || ggml_is_quantized(tensor->type) || tensor->ne[1] <= 1) {
+        return false;
+    }
+
+    // ggml_cuda_should_use_mmf and ggml_cuda_should_use_mmvf reject a src0 whose strides are not a
+    // multiple of 2*type_size, so a pad that misses this would silently cost those kernels
+    if (row_pad.pad % (2*ggml_type_size(tensor->type)) != 0) {
+        return false;
+    }
+
+    return ggml_row_size(tensor->type, tensor->ne[0]) % row_pad.alias_stride == 0;
+}
+
+static size_t ggml_cuda_padded_row_size(const ggml_tensor * tensor, int device) {
+    return ggml_row_size(tensor->type, tensor->ne[0]) + ggml_cuda_info().devices[device].row_pad.pad;
+}
+
 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;
 
@@ -769,6 +862,17 @@ static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer
         return GGML_STATUS_SUCCESS;
     }
 
+    if (ggml_cuda_should_pad_weight(tensor, ctx->device)) {
+        tensor->nb[1] = ggml_cuda_padded_row_size(tensor, ctx->device);
+        tensor->nb[2] = tensor->nb[1]*tensor->ne[1];
+        tensor->nb[3] = tensor->nb[2]*tensor->ne[2];
+
+        // the gaps between rows are never written, initialize them to 0 to avoid possible NaN values
+        ggml_cuda_set_device(ctx->device);
+        CUDA_CHECK(cudaMemset(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);
@@ -929,6 +1033,11 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t
         }
     }
 
+    // reserve room for the row stride that ggml_backend_cuda_buffer_init_tensor will assign
+    if (ggml_cuda_should_pad_weight(tensor, buft_ctx->device)) {
+        size = ggml_cuda_padded_row_size(tensor, buft_ctx->device)*ggml_nrows(tensor);
+    }
+
     return size;
 }
 

src/llama-model-loader.cpp

@@ -1038,6 +1038,17 @@ static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hpara
     return nullptr;
 }
 
+// some models use the token embedding tensor as the output, but since these are used in different layers and with different ops
+// the tensor is duplicated
+// to handle this, we check if the tensor is duplicated, and if so, we assume that it is being loaded as the output tensor
+static llm_tensor resolve_tn_tensor(const LLM_TN_IMPL & tn, int flags) {
+    if (tn.tensor == LLM_TENSOR_TOKEN_EMBD && (flags & llama_model_loader::TENSOR_DUPLICATED)) {
+        return LLM_TENSOR_OUTPUT;
+    }
+
+    return tn.tensor;
+}
+
 struct ggml_tensor * llama_model_loader::create_tensor(
         const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output,
         const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
@@ -1079,17 +1090,9 @@ struct ggml_tensor * llama_model_loader::create_tensor(
             throw std::runtime_error(format("missing tensor '%s'", tn.str().c_str()));
         }
 
-        // some models use the token embedding tensor as the output, but since these are used in different layers and with different ops
-        // the tensor is duplicated
-        // to handle this, we check if the tensor is duplicated, and if so, we assume that it is being loaded as the output tensor
-        llm_tensor tn_tensor = tn.tensor;
-        if (tn.tensor == LLM_TENSOR_TOKEN_EMBD && (flags & TENSOR_DUPLICATED)) {
-            tn_tensor = LLM_TENSOR_OUTPUT;
-        }
-
         llm_tensor_info info;
         try {
-            info = llm_tensor_info_for(tn_tensor);
+            info = llm_tensor_info_for(resolve_tn_tensor(tn, flags));
         } catch (const std::out_of_range & e) {
             throw std::runtime_error(format("missing tensor info mapping for %s", tn.str().c_str()));
         }
@@ -1271,6 +1274,16 @@ 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));
 
+    // tensors with a "weight" suffix are used as the src0 of the op that they map to; flag the ones
+    // that feed a matrix multiplication so that backends may pad their row stride
+    const bool is_weight = tn.suffix == nullptr || strcmp(tn.suffix, "weight") == 0;
+    if (is_weight) {
+        const ggml_op op = llm_tensor_info_for(resolve_tn_tensor(tn, flags)).op;
+        if (op == GGML_OP_MUL_MAT || op == GGML_OP_MUL_MAT_ID) {
+            tensor->flags |= GGML_TENSOR_FLAG_PAD_ROWS;
+        }
+    }
+
     if (duplicated) {
         size_data += ggml_nbytes(cur);
     } else {
@@ -1527,7 +1540,13 @@ bool llama_model_loader::load_all_data(
             }
         }
 
-        size_t n_size = ggml_nbytes(cur);
+        // the data in the file is packed, while the destination tensor may have a padded row
+        // stride, in which case the rows are uploaded with a strided 2D copy
+        const size_t  packed_row_size = ggml_row_size(cur->type, cur->ne[0]);
+        const int64_t n_rows          = ggml_nelements(cur) / cur->ne[0];
+        const bool    row_padded      = cur->nb[1] != packed_row_size;
+
+        const size_t n_size = row_padded ? packed_row_size*n_rows : ggml_nbytes(cur);
 
         if (use_mmap) {
             const auto & mapping = mappings.at(weight->idx);
@@ -1554,6 +1573,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_size, n_rows, cur->nb[1], packed_row_size);
             } else {
                 ggml_backend_tensor_set(cur, data, 0, n_size);
             }
@@ -1570,7 +1591,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);
@@ -1626,7 +1647,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_size, n_rows, cur->nb[1], packed_row_size);
+                    } 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)));
                     }