Skip to content

Commit 91373e1

Browse files
mudlermudler
committed
feat(paged): P1 bf16-stream bf16 residual-add + rope op-variants
Round out the op-variant set for the bf16-resident stream: bf16 branches in binbcast.cu (residual add) and bf16 rope instantiations (asserts widened only). Standing infra; Option-A keeps f32 at segment boundaries so these are not on the current measured path. Existing f32 paths untouched. Assisted-by: Claude:opus-4.8 [Claude Code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 1271488 commit 91373e1

2 files changed

Lines changed: 37 additions & 4 deletions

File tree

‎ggml/src/ggml-cuda/binbcast.cu‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -413,7 +413,7 @@ static void ggml_cuda_op_bin_bcast(
413413
const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
414414
const void * src0_dd, const void * src1_dd, void * dst_dd, cudaStream_t stream) {
415415

416-
GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
416+
GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16);
417417

418418
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
419419
op()(src0, src1, dst, (const float *)src0_dd, (const float *)src1_dd, (float *)dst_dd, stream);
@@ -423,6 +423,14 @@ static void ggml_cuda_op_bin_bcast(
423423
op()(src0, src1, dst, (const half *) src0_dd, (const float *)src1_dd, (half *) dst_dd, stream);
424424
} else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) {
425425
op()(src0, src1, dst, (const half *) src0_dd, (const float *)src1_dd, (float *)dst_dd, stream);
426+
// [P1 bf16-stream] bf16 residual-add variants, so a bf16-resident segment can keep
427+
// its residual add in bf16 (half the memory traffic) rather than widening to f32.
428+
} else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) {
429+
op()(src0, src1, dst, (const nv_bfloat16 *) src0_dd, (const nv_bfloat16 *) src1_dd, (nv_bfloat16 *) dst_dd, stream);
430+
} else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) {
431+
op()(src0, src1, dst, (const nv_bfloat16 *) src0_dd, (const float *) src1_dd, (nv_bfloat16 *) dst_dd, stream);
432+
} else if (src0->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) {
433+
op()(src0, src1, dst, (const nv_bfloat16 *) src0_dd, (const float *) src1_dd, (float *) dst_dd, stream);
426434
} else {
427435
fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__,
428436
ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type));

‎ggml/src/ggml-cuda/rope.cu‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -528,11 +528,16 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
528528
}
529529
cudaStream_t stream = ctx.stream();
530530

531-
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
532-
GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16);
531+
// [P1 bf16-stream] bf16 is accepted so a bf16-resident attention segment can rope
532+
// its Q/K in bf16 (the norm/neox kernels are float-internal, so the bf16 arms just
533+
// add T/D = nv_bfloat16 instantiations below).
534+
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
535+
GGML_ASSERT( dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_BF16);
533536
// When not fused, src0 and dst types must match
534537
// When fused (ROPE+VIEW+SET_ROWS), src0 may be F32 and dst may be F16
535-
GGML_ASSERT(src0->type == dst->type || (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16));
538+
GGML_ASSERT(src0->type == dst->type ||
539+
(src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) ||
540+
(src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16));
536541

537542
const int64_t ne00 = src0->ne[0]; // head dims
538543
const int64_t ne01 = src0->ne[1]; // num heads
@@ -610,6 +615,16 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
610615
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
611616
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
612617
set_rows_stride, stream);
618+
} else if (src0->type == GGML_TYPE_BF16 && dst_type == GGML_TYPE_BF16) {
619+
rope_neox_cuda<forward, nv_bfloat16, nv_bfloat16>((const nv_bfloat16 *) src0_d, (nv_bfloat16 *) dst_d, ne00, ne01, ne02, s01, s02,
620+
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
621+
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
622+
set_rows_stride, stream);
623+
} else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_BF16) {
624+
rope_neox_cuda<forward, float, nv_bfloat16>((const float *) src0_d, (nv_bfloat16 *) dst_d, ne00, ne01, ne02, s01, s02,
625+
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
626+
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
627+
set_rows_stride, stream);
613628
} else {
614629
GGML_ABORT("fatal error");
615630
}
@@ -653,6 +668,16 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx,
653668
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
654669
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
655670
set_rows_stride, stream);
671+
} else if (src0->type == GGML_TYPE_BF16 && dst_type == GGML_TYPE_BF16) {
672+
rope_norm_cuda<forward, nv_bfloat16, nv_bfloat16>((const nv_bfloat16 *) src0_d, (nv_bfloat16 *) dst_d, ne00, ne01, ne02, s01, s02,
673+
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
674+
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
675+
set_rows_stride, stream);
676+
} else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_BF16) {
677+
rope_norm_cuda<forward, float, nv_bfloat16>((const float *) src0_d, (nv_bfloat16 *) dst_d, ne00, ne01, ne02, s01, s02,
678+
s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base,
679+
ext_factor, attn_factor, corr_dims, freq_factors, row_indices,
680+
set_rows_stride, stream);
656681
} else {
657682
GGML_ABORT("fatal error");
658683
}

0 commit comments

Comments
 (0)