Skip to content

cuda: 5120-wide fused norm + Hadamard with PDL, and two-row PQ2_0 decode tiles on Hopper #417

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

Merged
bri-prism merged 3 commits into from
Oct 30, 2026
Merged

cuda: 5120-wide fused norm + Hadamard with PDL, and two-row PQ2_0 decode tiles on Hopper #417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
66f8079
cuda: fuse the 5120-wide RMS norm and Hadamard transform with PDL
bri-prism Oct 29, 2026
20812af
cuda: two rows per warp for smaller PQ2_0 one-column mat-vecs on Hopper
bri-prism Oct 29, 2026
f73a2c9
Merge remote-tracking branch 'origin/merge/upstream-20261011' into cu…
bri-prism Oct 29, 2026
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
60 changes: 42 additions & 18 deletions ggml/src/ggml-cuda/mmvq-pq2_0.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
// permuted inside each 16-element group (position k*4+m holds element m*4+k), then one half2 per 32-block
// with d and the raw int16 sum of q. With that, (code_word >> 2k) & 0x03030303 pairs the raw 2-bit codes 0..3
// with the right activation bytes, so there is no per-weight decode. The digit bias is one integer subtraction:
// sum(q*(c-1)) = sum(q*c) - sum(q). Each warp handles 4 rows and reuses the activation slice of a 32-block for all of them.
// sum(q*(c-1)) = sum(q*c) - sum(q). Each warp handles up to 4 rows and reuses the activation slice of a 32-block for all of them.
#pragma once

#include "common.cuh"
Expand All @@ -18,7 +18,7 @@
// second weight matrix (gate, same layout and strides) is dotted against the same activations, gate_bias is added to
// it and the GLU op combines the two. Without fusion the epilogue is a plain store: a runtime bias check there costs
// the 8-column verify kernel ~12% on H200.
template <int ncols, bool has_fusion, bool has_gate>
template <int ncols, bool has_fusion, bool has_gate, int rows = PQ2_0_MC_ROWS>
__global__ void __launch_bounds__(PQ2_0_MC_WARPS * 32, 3)
mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst,
Expand All @@ -29,7 +29,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
const void * GGML_CUDA_RESTRICT vy = vy_;
float * GGML_CUDA_RESTRICT dst = dst_;
const int lane = threadIdx.x & 31;
const int row0 = ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * PQ2_0_MC_ROWS;
const int row0 = ((blockIdx.x * blockDim.x + threadIdx.x) >> 5) * rows;
if (row0 >= nrows_x) {
return;
}
Expand All @@ -42,17 +42,17 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
const size_t ds_off = (size_t) stride_col_y * QK8_1; // qs region of a column, padded length
const char * y = (const char *) vy;

float acc[PQ2_0_MC_ROWS][ncols] = {};
[[maybe_unused]] float accg[PQ2_0_MC_ROWS][ncols] = {};
float acc[rows][ncols] = {};
[[maybe_unused]] float accg[rows][ncols] = {};
[[maybe_unused]] const void * GGML_CUDA_RESTRICT vgate = vgate_;

for (int c = lane; c < nchunk; c += WARP_SIZE) {
uint32_t w0[PQ2_0_MC_ROWS], w1[PQ2_0_MC_ROWS];
float d2[PQ2_0_MC_ROWS];
[[maybe_unused]] uint32_t g0[PQ2_0_MC_ROWS], g1[PQ2_0_MC_ROWS];
[[maybe_unused]] float e2[PQ2_0_MC_ROWS];
uint32_t w0[rows], w1[rows];
float d2[rows];
[[maybe_unused]] uint32_t g0[rows], g1[rows];
[[maybe_unused]] float e2[rows];
#pragma unroll
for (int r = 0; r < PQ2_0_MC_ROWS; ++r) {
for (int r = 0; r < rows; ++r) {
const int row = min(row0 + r, nrows_x - 1); // clamped rows are computed but not stored
const block_pq2_0 * blk = (const block_pq2_0 *) vx + (int64_t) row*stride_row_x + (c >> 2);
const uint16_t * qp = (const uint16_t *) (blk->qs + (c & 3)*8); // blocks are only 2-byte aligned
Expand Down Expand Up @@ -82,7 +82,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
}

#pragma unroll
for (int r = 0; r < PQ2_0_MC_ROWS; ++r) {
for (int r = 0; r < rows; ++r) {
int t[8];
#pragma unroll
for (int k = 0; k < 4; ++k) {
Expand Down Expand Up @@ -129,7 +129,7 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
ggml_cuda_pdl_lc();

#pragma unroll
for (int r = 0; r < PQ2_0_MC_ROWS; ++r) {
for (int r = 0; r < rows; ++r) {
#pragma unroll
for (int j = 0; j < ncols; ++j) {
float v = warp_reduce_sum<WARP_SIZE>(acc[r][j]);
Expand Down Expand Up @@ -163,27 +163,51 @@ mul_mat_vec_pq2_0_mc(const void * vx_, const void * vy_, float * dst_,
}
}

template <int ncols>
static void mul_mat_vec_pq2_0_mc_launch(
template <int ncols, int rows = PQ2_0_MC_ROWS>
static void mul_mat_vec_pq2_0_mc_launch_rows(
const void * vx, const void * vy, float * dst, const int ncols_x, const int nrows_x,
const int stride_row_x, const int stride_col_y, const int stride_col_dst, cudaStream_t stream,
const ggml_cuda_mm_fusion_args_device & fusion) {
const int nwarps = (nrows_x + PQ2_0_MC_ROWS - 1) / PQ2_0_MC_ROWS;
const int nwarps = (nrows_x + rows - 1) / rows;
const dim3 block_nums((nwarps + PQ2_0_MC_WARPS - 1) / PQ2_0_MC_WARPS, 1, 1);
const dim3 block_dims(PQ2_0_MC_WARPS * WARP_SIZE, 1, 1);
const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
if (fusion.gate) {
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, true, true>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, true, true, rows>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
(const float *) fusion.x_bias, fusion.gate, (const float *) fusion.gate_bias, fusion.glu_op, fusion.glu_limit);
} else if (fusion.x_bias) {
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, true, false>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, true, false, rows>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
(const float *) fusion.x_bias, nullptr, nullptr, fusion.glu_op, fusion.glu_limit);
} else {
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, false, false>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
ggml_cuda_kernel_launch(mul_mat_vec_pq2_0_mc<ncols, false, false, rows>, lp, vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
nullptr, nullptr, nullptr, fusion.glu_op, fusion.glu_limit);
}
}

template <int ncols>
static void mul_mat_vec_pq2_0_mc_launch(const void * vx,
const void * vy,
float * dst,
const int ncols_x,
const int nrows_x,
const int stride_row_x,
const int stride_col_y,
const int stride_col_dst,
cudaStream_t stream,
const ggml_cuda_mm_fusion_args_device & fusion) {
if constexpr (ncols == 1) {
// Smaller Hopper projections need more CTAs; keep four-row reuse for the vocabulary head.
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
if (cc == GGML_CUDA_CC_HOPPER && nrows_x <= 17408) {
mul_mat_vec_pq2_0_mc_launch_rows<ncols, 2>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y,
stride_col_dst, stream, fusion);
return;
}
}
mul_mat_vec_pq2_0_mc_launch_rows<ncols>(vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst,
stream, fusion);
}

static void mul_mat_vec_pq2_0_mc_switch(
const void * vx, const void * vy, float * dst, const int ncols_x, const int nrows_x, const int ncols_dst,
const int stride_row_x, const int stride_col_y, const int stride_col_dst, cudaStream_t stream,
Expand Down
31 changes: 25 additions & 6 deletions ggml/src/ggml-cuda/norm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -737,6 +737,10 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32(
__shared__ float s_sum[32];
__shared__ float s_fwht[N];

if constexpr (N == 5120) {
ggml_cuda_pdl_lc();
}

const int64_t row = blockIdx.x;
const int tid = threadIdx.x;
const int lane = tid % ggml_cuda_get_physical_warp_size();
Expand All @@ -749,6 +753,10 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32(
normed += row*N;
}

if constexpr (N == 5120) {
ggml_cuda_pdl_sync();
}

float x[NE];
float tmp = 0.0f;
#pragma unroll
Expand Down Expand Up @@ -827,8 +835,8 @@ static __global__ void __launch_bounds__(1024, 1) add_rms_norm_fwht_f32(
// ncols: the row width; blk: the Hadamard block width (the transform matrix is blk x blk, applied to each blk-element
// segment of the row; blk == ncols for a full-row transform)
bool ggml_cuda_add_rms_norm_fwht_supported(const int ncols, const int blk) {
return (ncols == 1024 || ncols == 2048 || ncols == 4096) &&
(blk == 512 || blk == 1024 || blk == 2048 || blk == 4096) && blk <= ncols;
return (ncols == 1024 || ncols == 2048 || ncols == 4096 || ncols == 5120) &&
(blk == 512 || blk == 1024 || blk == 2048 || blk == 4096) && ncols % blk == 0;
}

void ggml_cuda_op_add_rms_norm_fwht_fused(
Expand All @@ -854,10 +862,18 @@ void ggml_cuda_op_add_rms_norm_fwht_fused(
cudaStream_t stream = ctx.stream();
float * normed_d = normed ? (float *) normed->data : nullptr;
void * out = q8 ? q8 : dst->data;
#define ARF_LAUNCH(NN, BB, ADD, Q8, L) \
add_rms_norm_fwht_f32<NN, BB, ADD, Q8, L><<<(unsigned) nrows, 1024, 0, stream>>>( \
(const float *) a->data, ADD ? (const float *) b->data : nullptr, (const float *) weight->data, \
(const float *) signs->data, ADD ? (float *) add->data : nullptr, normed_d, out, eps, fwht_scale)
#define ARF_ARGS(ADD) \
(const float *) a->data, ADD ? (const float *) b->data : nullptr, (const float *) weight->data, \
(const float *) signs->data, ADD ? (float *) add->data : nullptr, normed_d, out, eps, fwht_scale
#define ARF_LAUNCH(NN, BB, ADD, Q8, L) \
do { \
if constexpr (NN == 5120) { \
const ggml_cuda_kernel_launch_params lp(dim3((unsigned) nrows), dim3(1024), 0, stream); \
ggml_cuda_kernel_launch(add_rms_norm_fwht_f32<NN, BB, ADD, Q8, L>, lp, ARF_ARGS(ADD)); \
} else { \
add_rms_norm_fwht_f32<NN, BB, ADD, Q8, L><<<(unsigned) nrows, 1024, 0, stream>>>(ARF_ARGS(ADD)); \
} \
} while (0)
#define ARF_CASE_ADD(NN, BB, ADD) \
if (!q8) { \
ARF_LAUNCH(NN, BB, ADD, false, GGML_CUDA_Q8_1_AOS); \
Expand All @@ -880,12 +896,15 @@ void ggml_cuda_op_add_rms_norm_fwht_fused(
case 1024: ARF_CASE_B(1024, 512) ARF_CASE_B(1024, 1024) break;
case 2048: ARF_CASE_B(2048, 512) ARF_CASE_B(2048, 1024) ARF_CASE_B(2048, 2048) break;
case 4096: ARF_CASE_B(4096, 512) ARF_CASE_B(4096, 1024) ARF_CASE_B(4096, 2048) ARF_CASE_B(4096, 4096) break;
case 5120:
ARF_CASE_B(5120, 512) ARF_CASE_B(5120, 1024) break;
default: break;
}
GGML_ABORT("unsupported width %d / block %d", ncols, blk);
#undef ARF_CASE_B
#undef ARF_CASE_ADD
#undef ARF_LAUNCH
#undef ARF_ARGS
}

void ggml_cuda_op_add_rms_norm_scale_fused(
Expand Down
35 changes: 35 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6208,6 +6208,8 @@ struct test_add_rms_norm_fwht : public test_case {

bool run_whole_graph() override { return true; }

int n_eval() override { return n == 5120 ? 3 : 1; }

ggml_tensor * build_graph(ggml_context * ctx) override {
const int64_t bw = blk ? blk : n;
ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, bw, bw);
Expand Down Expand Up @@ -6302,6 +6304,8 @@ struct test_rms_norm_fwht_lookahead : public test_case {

bool run_whole_graph() override { return true; }

int n_eval() override { return n == 5120 ? 3 : 1; }

ggml_tensor * build_graph(ggml_context * ctx) override {
const int64_t bw = blk ? blk : n;
ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, bw, bw);
Expand Down Expand Up @@ -6358,6 +6362,13 @@ struct test_rms_norm_fwht_lookahead : public test_case {
}
}
ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float));
} else if (n == 5120 && strcmp(t->name, "w_alpha") == 0) {
// Keep the unrelated projection exact so it does not obscure transform errors.
std::vector<float> data(ggml_nelements(t), 0.0f);
for (int64_t row = 0; row < t->ne[1]; ++row) {
data[row * n + row] = 1.0f;
}
ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float));
} else if (strcmp(t->name, "s") == 0) {
std::vector<float> data(ggml_nelements(t));
for (size_t i = 0; i < data.size(); i++) {
Expand Down Expand Up @@ -12300,6 +12311,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_rms_norm_fwht_lookahead(2048, n_tokens, GGML_TYPE_PQ2_0, blk));
}
}
// Five 1024-element segments: the 27B row width is not a power of two.
for (int64_t blk : { 512, 1024 }) {
for (int64_t n_tokens : { 1, 2, 8, 9, 33 }) {
for (ggml_type type_mm : { GGML_TYPE_COUNT, GGML_TYPE_PQ2_0 }) {
test_cases.emplace_back(new test_add_rms_norm_fwht(5120, n_tokens, type_mm, blk));
test_cases.emplace_back(new test_rms_norm_fwht_lookahead(5120, n_tokens, type_mm, blk));
}
}
}
// the same with PQ2_0 consumers: q8_1 quantized in the fused kernel (1 column: AoS, 3-8 columns: PQ2 layout)
for (int64_t n : {1024, 4096}) {
for (int64_t n_tokens : {1, 2, 4, 8, 9}) {
Expand Down Expand Up @@ -12416,6 +12436,21 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
}
// Hopper two-row decode tiles: partial warps/CTAs and both sides of the row-count cutoff.
for (int64_t m : { 1, 7, 8, 9, 15, 16, 17, 17408, 17409 }) {
for (int64_t n : { 1, 2 }) {
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, m, n, 1024, { 1, 1 }, { 1, 1 }));
}
}
for (bool with_bias : { false, true }) {
for (int64_t m : { 67, 17408 }) {
test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PQ2_0, GGML_GLU_OP_SWIGLU, 1, m, 5120, false,
1, 1, false, with_bias, true, false, { 1, 1 }));
}
test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_PQ2_0, GGML_GLU_OP_SWIGLU, 1, 5120, 17408, false,
1, 1, false, with_bias, false, false, { 1, 1 }));
}

// PTQ1_0 / PQ2_0 integer-dot mat-vec: Bonsai-2 shapes, odd row counts (row tail), batches and multi-column B
for (int64_t n : {1, 2, 3, 4, 5, 6, 7, 8}) {
for (int64_t k : {1024, 5120, 6144, 17408}) {
Expand Down
Loading