Skip to content

cuda: native PTQ1_0 MMQ tile loader (closes the dequantize-to-cuBLAS prefill fallback) #160

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bri-prism wants to merge 1 commit into
base: prism
Choose a base branch
from
Open

cuda: native PTQ1_0 MMQ tile loader (closes the dequantize-to-cuBLAS prefill fallback) #160

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 93 additions & 0 deletions ggml/src/ggml-cuda/convert.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,87 @@ static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t
}
}

#if !defined(GGML_USE_HIP)
template <typename dst_t>
static __device__
__forceinline__ void dequantize_ptq1_0_qs4(uint32_t packed, float d, dst_t * __restrict__ y, int base, int stride) {
uint32_t v_lo = __byte_perm(packed, 0, 0x4140);
uint32_t v_hi = __byte_perm(packed, 0, 0x4342);

# pragma unroll
for (int t = 0; t < 5; ++t) {
const uint32_t w_lo = v_lo * 3;
const uint32_t w_hi = v_hi * 3;
v_lo = w_lo & 0x00FF00FF;
v_hi = w_hi & 0x00FF00FF;

const uint32_t q = __vsub4(__byte_perm(w_lo, w_hi, 0x7531), 0x01010101);
# pragma unroll
for (int b = 0; b < 4; ++b) {
const int trit = (int8_t) (q >> (8 * b));
y[base + t * stride + b] = ggml_cuda_cast<dst_t>(d * trit);
}
}
}

template <typename dst_t>
static __global__ void dequantize_block_ptq1_0(const block_ptq1_0 * __restrict__ x,
dst_t * __restrict__ y,
int64_t nb) {
constexpr int threads_per_quant_block = 8;
constexpr int quant_blocks_per_cuda_block = CUDA_DEQUANTIZE_BLOCK_SIZE / threads_per_quant_block;
// The shared transpose keeps packed-byte ownership while making the CTA stores contiguous.
__shared__ dst_t dequantized[quant_blocks_per_cuda_block * QK_PTQ1_0];

const int64_t ib0 = int64_t(blockIdx.x) * quant_blocks_per_cuda_block;
const int64_t ib = ib0 + threadIdx.x / threads_per_quant_block;
const int lane = threadIdx.x % threads_per_quant_block;
const bool valid = ib < nb;

const block_ptq1_0 * bq = valid ? x + ib : x;
float d = valid && lane == 7 ? (float) bq->d : 0.0f;
d = __shfl_sync(__activemask(), d, 7, threads_per_quant_block);

dst_t * out = dequantized + (threadIdx.x / threads_per_quant_block) * QK_PTQ1_0;
if (valid && lane < 4) {
const uint32_t packed = ((const uint32_t *) bq->qs)[lane];
dequantize_ptq1_0_qs4(packed, d, out, 4 * lane, 16);
} else if (valid && lane < 6) {
const int g = lane - 4;
const uint32_t packed = ((const uint32_t *) (bq->qs + 16))[g];
dequantize_ptq1_0_qs4(packed, d, out, 80 + 4 * g, 8);
} else if (valid && lane == 6) {
uint32_t v = (uint32_t) bq->qh[0] | ((uint32_t) bq->qh[1] << 16);
# pragma unroll
for (int t = 0; t < 4; ++t) {
const uint32_t w = v * 3;
v = w & 0x00FF00FF;
out[120 + 2 * t + 0] = ggml_cuda_cast<dst_t>(d * ((int) ((w >> 8) & 0xFF) - 1));
out[120 + 2 * t + 1] = ggml_cuda_cast<dst_t>(d * ((int) ((w >> 24) & 0xFF) - 1));
}
}

__syncthreads();
const int nblocks = nb - ib0 < quant_blocks_per_cuda_block ? nb - ib0 : quant_blocks_per_cuda_block;
const int nelements = nblocks * QK_PTQ1_0;
for (int i = threadIdx.x; i < nelements; i += CUDA_DEQUANTIZE_BLOCK_SIZE) {
y[ib0 * QK_PTQ1_0 + i] = dequantized[i];
}
}

template <typename dst_t>
static void dequantize_row_ptq1_0_cuda(const void * __restrict__ vx,
dst_t * __restrict__ y,
const int64_t k,
cudaStream_t stream) {
GGML_ASSERT(k % QK_PTQ1_0 == 0);
constexpr int quant_blocks_per_cuda_block = CUDA_DEQUANTIZE_BLOCK_SIZE / 8;
const int64_t nb = k / QK_PTQ1_0;
const int num_blocks = (nb + quant_blocks_per_cuda_block - 1) / quant_blocks_per_cuda_block;
dequantize_block_ptq1_0<<<num_blocks, CUDA_DEQUANTIZE_BLOCK_SIZE, 0, stream>>>((const block_ptq1_0 *) vx, y, nb);
}
#endif

//================================== k-quants

template<typename dst_t>
Expand Down Expand Up @@ -464,7 +545,11 @@ to_bf16_cuda_t ggml_get_to_bf16_cuda(ggml_type type) {
case GGML_TYPE_PQ2_0:
return dequantize_block_cont_cuda<QK_PQ2_0, QR_PQ2_0, dequantize_pq2_0>;
case GGML_TYPE_PTQ1_0:
#if !defined(GGML_USE_HIP)
return dequantize_row_ptq1_0_cuda;
#else
return dequantize_block_cont_cuda<QK_PTQ1_0, QR_PTQ1_0, dequantize_ptq1_0>;
#endif
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
Expand Down Expand Up @@ -525,7 +610,11 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) {
case GGML_TYPE_PQ2_0:
return dequantize_block_cont_cuda<QK_PQ2_0, QR_PQ2_0, dequantize_pq2_0>;
case GGML_TYPE_PTQ1_0:
#if !defined(GGML_USE_HIP)
return dequantize_row_ptq1_0_cuda;
#else
return dequantize_block_cont_cuda<QK_PTQ1_0, QR_PTQ1_0, dequantize_ptq1_0>;
#endif
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
Expand Down Expand Up @@ -589,7 +678,11 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) {
case GGML_TYPE_PQ2_0:
return dequantize_block_cont_cuda<QK_PQ2_0, QR_PQ2_0, dequantize_pq2_0>;
case GGML_TYPE_PTQ1_0:
#if !defined(GGML_USE_HIP)
return dequantize_row_ptq1_0_cuda;
#else
return dequantize_block_cont_cuda<QK_PTQ1_0, QR_PTQ1_0, dequantize_ptq1_0>;
#endif
case GGML_TYPE_Q4_0:
return dequantize_row_q4_0_cuda;
case GGML_TYPE_Q4_1:
Expand Down
12 changes: 12 additions & 0 deletions ggml/src/ggml-cuda/mmq-config-ampere.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,18 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf
CASE(GGML_TYPE_PQ2_0, 256, 1, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PQ2_0, 256, 1, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);

CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);
CASE(GGML_TYPE_PTQ1_0, 256, 1, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, false);

CASE(GGML_TYPE_Q4_0, 256, 1, 128, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q4_0, 256, 1, 128, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
CASE(GGML_TYPE_Q4_0, 256, 1, 128, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, true, true);
Expand Down
100 changes: 100 additions & 0 deletions ggml/src/ggml-cuda/mmq-load-tiles.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,106 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
}

#if !defined(GGML_USE_HIP)
static __device__
__forceinline__ void ggml_cuda_mmq_decode_ptq1_0_qs4(uint32_t packed, int * __restrict__ dst, int stride) {
uint32_t v_lo = __byte_perm(packed, 0, 0x4140);
uint32_t v_hi = __byte_perm(packed, 0, 0x4342);

# pragma unroll
for (int t = 0; t < 5; ++t) {
const uint32_t w_lo = v_lo * 3;
const uint32_t w_hi = v_hi * 3;
v_lo = w_lo & 0x00FF00FF;
v_hi = w_hi & 0x00FF00FF;
dst[t * stride] = __vsub4(__byte_perm(w_lo, w_hi, 0x7531), 0x01010101);
}
}

template <ggml_type type, int J, bool fallback>
static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_ptq1_0(const char * __restrict__ x,
int * __restrict__ x_tile,
const int kbx0,
const int i_max,
const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback);
constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback);

# if defined(TURING_MMA_AVAILABLE)
int * x_qs = (int *) x_tile;
float * x_df = (float *) (x_qs + 2 * MMQ_TILE_NE_K);
# else
constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, I);
int * x_qs = (int *) x_tile;
float * x_df = (float *) (x_qs + txs.qs);
# endif

constexpr int blocks_per_iter = MMQ_ITER_K / QK_PTQ1_0;
constexpr int threads_per_block = 8;
constexpr int threads_per_row = blocks_per_iter * threads_per_block;
constexpr int nrows = warp_size / threads_per_row;

const int txi = threadIdx.x % threads_per_row;
const int kbx = txi / threads_per_block;
const int lane = txi % threads_per_block;

# pragma unroll
for (int i0 = 0; i0 < I; i0 += nrows * nwarps) {
int i = i0 + threadIdx.y * nrows + threadIdx.x / threads_per_row;
if (fallback) {
i = min(i, i_max);
}

const block_ptq1_0 * bxi = (const block_ptq1_0 *) x + kbx0 + i * stride + kbx;
# if defined(TURING_MMA_AVAILABLE)
int * row = x_qs + i * sram_stride + kbx * (QK_PTQ1_0 / 4);
# else
int * row = x_qs + i * (2 * MMQ_TILE_NE_K + 1) + kbx * (QK_PTQ1_0 / 4);
# endif

if (lane < 4) {
ggml_cuda_mmq_decode_ptq1_0_qs4(get_int_b4(bxi->qs, lane), row + lane, 4);
} else if (lane < 6) {
const int g = lane - 4;
ggml_cuda_mmq_decode_ptq1_0_qs4(get_int_b4(bxi->qs + 16, g), row + 20 + g, 2);
} else if (lane == 6) {
uint32_t v = (uint32_t) bxi->qh[0] | ((uint32_t) bxi->qh[1] << 16);
# pragma unroll
for (int t = 0; t < 4; t += 2) {
const uint32_t w0 = v * 3;
v = w0 & 0x00FF00FF;
const uint32_t w1 = v * 3;
v = w1 & 0x00FF00FF;
row[30 + t / 2] = __vsub4(__byte_perm(w0, w1, 0x7531), 0x01010101);
}
}
}

constexpr int scale_entries_per_block = QK_PTQ1_0 / QK8_1;
constexpr int scale_entries_per_row = blocks_per_iter * scale_entries_per_block;
constexpr int rows_per_warp = warp_size / scale_entries_per_row;
const int ksx = threadIdx.x % scale_entries_per_row;
const int scale_block = ksx / scale_entries_per_block;

# pragma unroll
for (int i0 = 0; i0 < I; i0 += nwarps * rows_per_warp) {
int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / scale_entries_per_row;
if (fallback) {
i = min(i, i_max);
}

const block_ptq1_0 * bxi = (const block_ptq1_0 *) x + kbx0 + i * stride + scale_block;
# if defined(TURING_MMA_AVAILABLE)
x_df[i * sram_stride + ksx] = bxi->d;
# else
x_df[i * (2 * MMQ_TILE_NE_K / QI8_0) + i / (QI8_0 / 2) + ksx] = bxi->d;
# endif
}
}
#endif

template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
Expand Down
16 changes: 16 additions & 0 deletions ggml/src/ggml-cuda/mmq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
case GGML_TYPE_PQ2_0:
mul_mat_q_case<GGML_TYPE_PQ2_0>(ctx, args, stream);
break;
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
mul_mat_q_case<GGML_TYPE_PTQ1_0>(ctx, args, stream);
break;
#endif
case GGML_TYPE_Q4_0:
mul_mat_q_case<GGML_TYPE_Q4_0>(ctx, args, stream);
break;
Expand Down Expand Up @@ -267,6 +272,11 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
bool mmq_supported;

switch (type) {
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
mmq_supported = turing_mma_available(cc);
break;
#endif
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_PQ2_0:
Expand Down Expand Up @@ -313,6 +323,12 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
}
}

#if !defined(GGML_USE_HIP)
if (type == GGML_TYPE_PTQ1_0) {
return ne11 <= MMQ_PTQ1_0_MAX_BATCH_SIZE;
}
#endif

if (turing_mma_available(cc)) {
return true;
}
Expand Down
35 changes: 30 additions & 5 deletions ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@
#include <climits>
#include <cstdint>

#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available.
#define MMQ_ITER_K 256
#define MMQ_ITER_K_FP4 512
#define MMQ_NWARPS 8
#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available.
#define MMQ_PTQ1_0_MAX_BATCH_SIZE 64
#define MMQ_ITER_K 256
#define MMQ_ITER_K_FP4 512
#define MMQ_NWARPS 8

typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
Expand Down Expand Up @@ -62,6 +63,9 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_PQ2_0:
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
#endif
return MMQ_Q8_1_DS_LAYOUT_D4;
case GGML_TYPE_Q4_0:
case GGML_TYPE_Q4_1:
Expand Down Expand Up @@ -388,7 +392,12 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml
switch (type) {
case GGML_TYPE_Q1_0: return MMQ_DP4A_TXS_Q8_0;
case GGML_TYPE_Q2_0: return MMQ_DP4A_TXS_Q8_0;
case GGML_TYPE_PQ2_0: return MMQ_DP4A_TXS_Q8_0;
case GGML_TYPE_PQ2_0:
return MMQ_DP4A_TXS_Q8_0;
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
return MMQ_DP4A_TXS_Q8_0;
#endif
case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0;
case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1;
case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0;
Expand Down Expand Up @@ -558,6 +567,12 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
ggml_cuda_mmq_load_tiles_pq2_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
return ggml_cuda_mmq_util_funcs(VDR_PTQ1_0_Q8_1_MMQ, ggml_cuda_mmq_load_tiles_ptq1_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback>);
#endif
case GGML_TYPE_Q4_0:
return ggml_cuda_mmq_util_funcs(
VDR_Q4_0_Q8_1_MMQ,
Expand Down Expand Up @@ -728,6 +743,13 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
ggml_cuda_mmq_load_tiles_pq2_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
ggml_cuda_mmq_write_back_mma<type, J, fallback>);
#if !defined(GGML_USE_HIP)
case GGML_TYPE_PTQ1_0:
return ggml_cuda_mmq_util_funcs(
-1, ggml_cuda_mmq_load_tiles_ptq1_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
ggml_cuda_mmq_write_back_mma<type, J, fallback>);
#endif
case GGML_TYPE_Q4_0:
return ggml_cuda_mmq_util_funcs(
-1,
Expand Down Expand Up @@ -1580,6 +1602,9 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda
extern DECL_MMQ_CASE(GGML_TYPE_Q1_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q2_0);
extern DECL_MMQ_CASE(GGML_TYPE_PQ2_0);
#if !defined(GGML_USE_HIP)
extern DECL_MMQ_CASE(GGML_TYPE_PTQ1_0);
#endif
extern DECL_MMQ_CASE(GGML_TYPE_Q4_0);
extern DECL_MMQ_CASE(GGML_TYPE_Q4_1);
extern DECL_MMQ_CASE(GGML_TYPE_Q5_0);
Expand Down
5 changes: 5 additions & 0 deletions ggml/src/ggml-cuda/mmvq.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,11 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) {
if (!ggml_is_quantized(type)) {
return false;
}
#if !defined(GGML_USE_HIP)
if (type == GGML_TYPE_PTQ1_0 && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) {
return ne11 <= 7;
}
#endif
// k-quants cost more to decode and mvq redoes that per column, so MMQ wins sooner.
// Only list quant-types MMQ supports, others would fall back to cuBLAS.
if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_ADA_LOVELACE) {
Expand Down
Loading
Loading