const ggml_tensor * ffn_gate,
const ggml_tensor * glu,
const ggml_tensor * ffn_up_bias = nullptr,
- const ggml_tensor * ffn_gate_bias = nullptr) {
+ const ggml_tensor * ffn_gate_bias = nullptr,
+ const ggml_tensor * ffn_up_scale = nullptr,
+ const ggml_tensor * ffn_gate_scale = nullptr) {
const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr;
+ const bool has_scale = ffn_up_scale != nullptr || ffn_gate_scale != nullptr;
if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) {
return false;
}
+ if (has_scale && (!ffn_up_scale || !ffn_gate_scale)) {
+ return false;
+ }
const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU;
const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU;
}
const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID;
+ const ggml_tensor * ffn_up_bias_src = has_scale ? ffn_up_scale : ffn_up;
+ const ggml_tensor * ffn_gate_bias_src = has_scale ? ffn_gate_scale : ffn_gate;
+ const ggml_tensor * ffn_up_out = has_bias ? ffn_up_bias : ffn_up_bias_src;
+ const ggml_tensor * ffn_gate_out = has_bias ? ffn_gate_bias : ffn_gate_bias_src;
- if (has_bias) {
- if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) {
+ if (glu->src[0] != ffn_gate_out || glu->src[1] != ffn_up_out) {
+ return false;
+ }
+
+ if (has_scale) {
+ if (ffn_up_scale->op != GGML_OP_MUL || ffn_gate_scale->op != GGML_OP_MUL) {
+ return false;
+ }
+ const bool up_has_mm = ffn_up_scale->src[0] == ffn_up || ffn_up_scale->src[1] == ffn_up;
+ const bool gate_has_mm = ffn_gate_scale->src[0] == ffn_gate || ffn_gate_scale->src[1] == ffn_gate;
+ if (!up_has_mm || !gate_has_mm) {
return false;
}
+ }
- if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) {
+ if (has_bias) {
+ if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) {
return false;
}
if (expected_bias_op == GGML_OP_ADD) {
- const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up;
- const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate;
+ const bool up_has_mul = ffn_up_bias->src[0] == ffn_up_bias_src || ffn_up_bias->src[1] == ffn_up_bias_src;
+ const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate_bias_src || ffn_gate_bias->src[1] == ffn_gate_bias_src;
if (!up_has_mul || !gate_has_mul) {
return false;
}
} else { // GGML_OP_ADD_ID
- if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) {
+ if (ffn_up_bias->src[0] != ffn_up_bias_src || ffn_gate_bias->src[0] != ffn_gate_bias_src) {
return false;
}
if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) {
return false;
}
}
- } else {
- if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) {
- return false;
- }
}
if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) ||
return false;
}
- if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) {
+ if (is_mul_mat_id && ffn_up->src[2] != ffn_gate->src[2]) {
return false;
}
bool fused_mul_mat_vec = false;
int fused_node_count = 0;
- // gate + glu + up
+ auto get_mul_mat_scale = [](const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
+ const bool scale_lhs_mm = scale_node->src[0] == mm_node;
+ const bool scale_rhs_mm = scale_node->src[1] == mm_node;
+ if (!scale_lhs_mm && !scale_rhs_mm) {
+ return nullptr;
+ }
+
+ const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0];
+ if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
+ scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 ||
+ !ggml_are_same_shape(scale_node, mm_node)) {
+ return nullptr;
+ }
+
+ return scale;
+ };
+
+ auto get_mul_mat_id_scale = [](const ggml_tensor * reshape, const ggml_tensor * repeat, const ggml_tensor * getrows,
+ const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * {
+ if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm_node->src[2]) {
+ return nullptr;
+ }
+ if (!((scale_node->src[0] == mm_node && scale_node->src[1] == getrows) ||
+ (scale_node->src[0] == getrows && scale_node->src[1] == mm_node))) {
+ return nullptr;
+ }
+
+ const ggml_tensor * scale = reshape->src[0];
+ if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 ||
+ scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] ||
+ !ggml_are_same_shape(scale_node, mm_node)) {
+ return nullptr;
+ }
+
+ return scale;
+ };
+
+ auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) -> const ggml_tensor * {
+ if (op_bias == GGML_OP_ADD) {
+ if (bias_node->src[0] == mul_node) {
+ return bias_node->src[1];
+ }
+ if (bias_node->src[1] == mul_node) {
+ return bias_node->src[0];
+ }
+ return nullptr;
+ }
+ GGML_ASSERT(op_bias == GGML_OP_ADD_ID);
+ GGML_ASSERT(bias_node->src[0] == mul_node);
+ return bias_node->src[1];
+ };
+
+ // gate + glu + up, with optional scale/bias on both lanes.
for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
+ if (op == GGML_OP_MUL_MAT) {
+ for (const bool with_bias : { false, true }) {
+ const int gate_idx = i;
+ const int gate_scale_idx = i + 1;
+ const int gate_bias_idx = with_bias ? i + 2 : -1;
+ const int up_idx = with_bias ? i + 3 : i + 2;
+ const int up_scale_idx = up_idx + 1;
+ const int up_bias_idx = with_bias ? up_idx + 2 : -1;
+ const int glu_idx = with_bias ? up_idx + 3 : up_idx + 2;
+
+ const int out_nodes[] = { glu_idx };
+ ggml_op ops[7];
+ if (with_bias) {
+ ops[0] = op;
+ ops[1] = GGML_OP_MUL;
+ ops[2] = bias_op;
+ ops[3] = op;
+ ops[4] = GGML_OP_MUL;
+ ops[5] = bias_op;
+ ops[6] = GGML_OP_GLU;
+ } else {
+ ops[0] = op;
+ ops[1] = GGML_OP_MUL;
+ ops[2] = op;
+ ops[3] = GGML_OP_MUL;
+ ops[4] = GGML_OP_GLU;
+ }
+ const int n_ops = with_bias ? 7 : 5;
+
+ if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
+ !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
+ continue;
+ }
+
+ ggml_tensor * gate_n = cgraph->nodes[gate_idx];
+ ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx];
+ ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n;
+ ggml_tensor * up_n = cgraph->nodes[up_idx];
+ ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx];
+ ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n;
+ const ggml_tensor * glu = cgraph->nodes[glu_idx];
+
+ if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu,
+ with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) {
+ continue;
+ }
+
+ const ggml_tensor * gate_scale = get_mul_mat_scale(gate_scale_n, gate_n);
+ const ggml_tensor * up_scale = get_mul_mat_scale(up_scale_n, up_n);
+ if (!gate_scale || !up_scale) {
+ continue;
+ }
+
+ const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr;
+ const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr;
+ if (with_bias && (!ggml_are_same_shape(gate_out_n->src[0], gate_out_n->src[1]) ||
+ !ggml_are_same_shape(up_out_n->src[0], up_out_n->src[1]))) {
+ continue;
+ }
+
+ const ggml_tensor * src0 = up_n->src[0];
+ const ggml_tensor * src1 = up_n->src[1];
+ const ggml_tensor * ids = up_n->src[2];
+
+ ggml_cuda_mm_fusion_args_host fusion_data{};
+ fusion_data.gate = gate_n->src[0];
+ fusion_data.x_bias = up_bias;
+ fusion_data.gate_bias = gate_bias;
+ fusion_data.x_scale = up_scale;
+ fusion_data.gate_scale = gate_scale;
+ fusion_data.glu_op = ggml_get_glu_op(glu);
+
+ if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
+ ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
+ fused_mul_mat_vec = true;
+ fused_node_count = n_ops;
+ break;
+ }
+ }
+
+ if (fused_mul_mat_vec) {
+ break;
+ }
+ } else {
+ for (const bool with_bias : { false, true }) {
+ const int gate_idx = i;
+ const int gate_scale_idx = i + 4;
+ const int gate_bias_idx = with_bias ? i + 5 : -1;
+ const int up_idx = with_bias ? i + 6 : i + 5;
+ const int up_scale_idx = up_idx + 4;
+ const int up_bias_idx = with_bias ? up_idx + 5 : -1;
+ const int glu_idx = with_bias ? up_idx + 6 : up_idx + 5;
+
+ const int out_nodes[] = { glu_idx };
+ ggml_op ops[13];
+ if (with_bias) {
+ ops[0] = op;
+ ops[1] = GGML_OP_RESHAPE;
+ ops[2] = GGML_OP_REPEAT;
+ ops[3] = GGML_OP_GET_ROWS;
+ ops[4] = GGML_OP_MUL;
+ ops[5] = bias_op;
+ ops[6] = op;
+ ops[7] = GGML_OP_RESHAPE;
+ ops[8] = GGML_OP_REPEAT;
+ ops[9] = GGML_OP_GET_ROWS;
+ ops[10] = GGML_OP_MUL;
+ ops[11] = bias_op;
+ ops[12] = GGML_OP_GLU;
+ } else {
+ ops[0] = op;
+ ops[1] = GGML_OP_RESHAPE;
+ ops[2] = GGML_OP_REPEAT;
+ ops[3] = GGML_OP_GET_ROWS;
+ ops[4] = GGML_OP_MUL;
+ ops[5] = op;
+ ops[6] = GGML_OP_RESHAPE;
+ ops[7] = GGML_OP_REPEAT;
+ ops[8] = GGML_OP_GET_ROWS;
+ ops[9] = GGML_OP_MUL;
+ ops[10] = GGML_OP_GLU;
+ }
+ const int n_ops = with_bias ? 13 : 11;
+
+ if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
+ !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
+ continue;
+ }
+
+ ggml_tensor * gate_n = cgraph->nodes[gate_idx];
+ ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx];
+ ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n;
+ ggml_tensor * up_n = cgraph->nodes[up_idx];
+ ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx];
+ ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n;
+ const ggml_tensor * glu = cgraph->nodes[glu_idx];
+
+ if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu,
+ with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) {
+ continue;
+ }
+
+ const ggml_tensor * gate_scale = get_mul_mat_id_scale(cgraph->nodes[gate_idx + 1], cgraph->nodes[gate_idx + 2],
+ cgraph->nodes[gate_idx + 3], gate_scale_n, gate_n);
+ const ggml_tensor * up_scale = get_mul_mat_id_scale(cgraph->nodes[up_idx + 1], cgraph->nodes[up_idx + 2],
+ cgraph->nodes[up_idx + 3], up_scale_n, up_n);
+ if (!gate_scale || !up_scale) {
+ continue;
+ }
+
+ const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr;
+ const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr;
+
+ const ggml_tensor * src0 = up_n->src[0];
+ const ggml_tensor * src1 = up_n->src[1];
+ const ggml_tensor * ids = up_n->src[2];
+
+ ggml_cuda_mm_fusion_args_host fusion_data{};
+ fusion_data.gate = gate_n->src[0];
+ fusion_data.x_bias = up_bias;
+ fusion_data.gate_bias = gate_bias;
+ fusion_data.x_scale = up_scale;
+ fusion_data.gate_scale = gate_scale;
+ fusion_data.glu_op = ggml_get_glu_op(glu);
+
+ if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) {
+ ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data);
+ fused_mul_mat_vec = true;
+ fused_node_count = n_ops;
+ break;
+ }
+ }
+
+ if (fused_mul_mat_vec) {
+ break;
+ }
+ }
+
if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) {
ggml_tensor * glu = cgraph->nodes[i + 4];
ggml_tensor * gate_bias_n = glu->src[0];
continue;
}
- auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) {
- if (op_bias == GGML_OP_ADD) {
- if (bias_node->src[0] == mul_node) {
- return bias_node->src[1];
- }
- if (bias_node->src[1] == mul_node) {
- return bias_node->src[0];
- }
- return (ggml_tensor *) nullptr;
- }
- GGML_ASSERT(op_bias == GGML_OP_ADD_ID);
- GGML_ASSERT(bias_node->src[0] == mul_node);
- return bias_node->src[1];
- };
-
- ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op);
- ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op);
+ const ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op);
+ const ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op);
if (!up_bias_tensor || !gate_bias_tensor) {
continue;
fused_mul_mat_vec = false;
fused_node_count = 0;
- // gate + add + glu + up + add
+ // mul_mat + scale + optional bias
+ for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
+ const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
+
+ for (const bool with_bias : { false, true }) {
+ const int n_ops = op == GGML_OP_MUL_MAT ? (with_bias ? 3 : 2) : (with_bias ? 6 : 5);
+ const int out_nodes[] = { i + n_ops - 1 };
+ ggml_op ops[6];
+ if (op == GGML_OP_MUL_MAT) {
+ if (with_bias) {
+ ops[0] = op;
+ ops[1] = GGML_OP_MUL;
+ ops[2] = bias_op;
+ } else {
+ ops[0] = op;
+ ops[1] = GGML_OP_MUL;
+ }
+ } else {
+ if (with_bias) {
+ ops[0] = op;
+ ops[1] = GGML_OP_RESHAPE;
+ ops[2] = GGML_OP_REPEAT;
+ ops[3] = GGML_OP_GET_ROWS;
+ ops[4] = GGML_OP_MUL;
+ ops[5] = bias_op;
+ } else {
+ ops[0] = op;
+ ops[1] = GGML_OP_RESHAPE;
+ ops[2] = GGML_OP_REPEAT;
+ ops[3] = GGML_OP_GET_ROWS;
+ ops[4] = GGML_OP_MUL;
+ }
+ }
+
+ if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) ||
+ !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) {
+ continue;
+ }
+
+ ggml_tensor * mm_node = cgraph->nodes[i];
+ ggml_tensor * scale_node = op == GGML_OP_MUL_MAT ? cgraph->nodes[i + 1] : cgraph->nodes[i + 4];
+ ggml_tensor * out_node = with_bias ? cgraph->nodes[i + n_ops - 1] : scale_node;
+
+ const ggml_tensor * scale = nullptr;
+ if (op == GGML_OP_MUL_MAT) {
+ scale = get_mul_mat_scale(scale_node, mm_node);
+ } else {
+ scale = get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], scale_node, mm_node);
+ }
+ if (!scale) {
+ continue;
+ }
+
+ const ggml_tensor * bias = with_bias ? get_bias_tensor(out_node, scale_node, bias_op) : nullptr;
+ if (with_bias && !bias) {
+ continue;
+ }
+ if (with_bias && bias_op == GGML_OP_ADD && !ggml_are_same_shape(out_node->src[0], out_node->src[1])) {
+ continue;
+ }
+ if (with_bias && bias_op == GGML_OP_ADD_ID && out_node->src[2] != mm_node->src[2]) {
+ continue;
+ }
+
+ const ggml_tensor * src0 = mm_node->src[0];
+ const ggml_tensor * src1 = mm_node->src[1];
+ const ggml_tensor * ids = mm_node->src[2];
+
+ ggml_cuda_mm_fusion_args_host fusion_data{};
+ fusion_data.x_bias = bias;
+ fusion_data.x_scale = scale;
+
+ if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) {
+ ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data);
+ fused_mul_mat_vec = true;
+ fused_node_count = n_ops;
+ break;
+ }
+ }
+ if (fused_mul_mat_vec) {
+ break;
+ }
+ }
+
+ if (fused_mul_mat_vec) {
+ return fused_node_count - 1;
+ }
+
+ // mul_mat + add
for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) {
const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID;
}
}
-#ifdef GGML_CUDA_DEBUG
- const int nodes_fused = i - prev_i - 1;
- if (nodes_fused > 0) {
- GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused);
- }
-#endif
prev_i = i;
if (ggml_cuda_is_view_or_noop(node)) {
int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i);
if (nodes_to_skip != 0) {
+#ifdef GGML_CUDA_DEBUG
+ const int last_fused = i + nodes_to_skip;
+ GGML_LOG_INFO("nodes_fused: %d, first: %s (%s), last: %s (%s)\n",
+ nodes_to_skip + 1, ggml_op_name(node->op), node->name,
+ ggml_op_name(cgraph->nodes[last_fused]->op), cgraph->nodes[last_fused]->name);
+#endif
i += nodes_to_skip;
continue;
}
bool use_gate = false;
bool use_bias = false;
bool use_gate_bias = false;
+ bool use_scale = false;
+ bool use_gate_scale = false;
[[maybe_unused]] const void * vgate = nullptr;
const float * x_bias = nullptr;
const float * gate_bias = nullptr;
+ const float * x_scale = nullptr;
+ const float * gate_scale = nullptr;
ggml_glu_op active_glu;
if constexpr (has_fusion) {
x_bias = (const float *) fusion.x_bias;
gate_bias = (const float *) fusion.gate_bias;
active_glu = fusion.glu_op;
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ use_scale = fusion.x_scale != nullptr;
+ use_gate_scale = fusion.gate_scale != nullptr && use_gate;
+ x_scale = (const float *) fusion.x_scale;
+ gate_scale = (const float *) fusion.gate_scale;
+ }
}
[[maybe_unused]] float x_biases[ncols_dst] = { 0.0f };
[[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f };
+ [[maybe_unused]] float x_scales;
+ [[maybe_unused]] float gate_scales;
if constexpr (has_fusion) {
+ // 1. Hide latency by prefetching bias, gates and scales here
+ // 2. load only on threads that won't die after partial sum calculation
const uint32_t channel_bias = ids ? channel_x : channel_dst;
- if (use_bias) {
- x_bias = x_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
- // 1. Hide latency by prefetching bias and gate here
- // 2. load only on threads that won't die after partial sum calculation
- if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
- (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+ if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
+ (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+ if (use_bias) {
+ x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
#pragma unroll
for (int j = 0; j < ncols_dst; ++j) {
x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x];
}
}
- }
- if (use_gate_bias) {
- gate_bias = gate_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0;
- if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 &&
- (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+ if (use_gate_bias) {
+ gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0;
#pragma unroll
for (int j = 0; j < ncols_dst; ++j) {
gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x];
}
}
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ if (use_scale) {
+ x_scales = x_scale[ids ? channel_x : 0];
+ }
+ if (use_gate_scale) {
+ gate_scales = gate_scale[ids ? channel_x : 0];
+ }
+ }
}
}
if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
float result = tmp[j][threadIdx.x];
if constexpr (has_fusion) {
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ if (use_scale) {
+ result *= x_scales;
+ }
+ }
if (use_bias) {
result += x_biases[j];
}
if (use_gate) {
float gate_value = tmp_gate[j][threadIdx.x];
+ if constexpr (type == GGML_TYPE_NVFP4) {
+ if (use_gate_scale) {
+ gate_value *= gate_scales;
+ }
+ }
if (use_gate_bias) {
gate_value += gate_biases[j];
}
}
if constexpr (!has_fusion) {
- GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate);
+ GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate);
+ }
+ if constexpr (type != GGML_TYPE_NVFP4) {
+ GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales);
}
}
const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared,
const uint32_t ids_stride, cudaStream_t stream) {
- const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
+ const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
+ fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
if constexpr (c_ncols_dst == 1) {
if (has_fusion) {
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream);
const int warp_size = ggml_cuda_info().devices[device].warp_size;
const mmvq_parameter_table_id table_id = get_device_table_id(cc);
- const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr;
const bool has_ids = ids != nullptr;
const auto should_use_small_k = [&](int c_ncols_dst) {
GGML_ABORT("fatal error");
break;
}
-
- GGML_UNUSED(has_fusion);
}
static void mul_mat_vec_q_switch_type(
const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
if (fusion) {
GGML_ASSERT( !ids || dst->ne[2] == 1);
GGML_ASSERT( ids || dst->ne[1] == 1);
+ // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
+ // non-negligible for some models such as gpt-oss-20b
+ GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4);
if (fusion->x_bias) {
GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32);
GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
fusion_local.gate_bias = fusion->gate_bias->data;
}
+ if (fusion->x_scale) {
+ GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
+ GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
+ GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
+ fusion_local.x_scale = fusion->x_scale->data;
+ }
+ if (fusion->gate_scale) {
+ GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
+ GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
+ GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
+ fusion_local.gate_scale = fusion->gate_scale->data;
+ }
fusion_local.glu_op = fusion->glu_op;
}