public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-10-12 01:05:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

cuda: support Nemotron 3 Puzzle state size 96 for ssm scan (llama/28717)

Browse Source
This commit is contained in:
Anav Prasad
2026-10-06 10:38:04 +03:00
committed by Georgi Gerganov
parent dfe8fbd7e7
commit fc1ebfabf2
2 changed files with 22 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+3 -2
View File
@@ -5487,8 +5487,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
if (op->src[3]->ne[0] == 1) {
// Mamba2
// (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
// (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
const int64_t d_state = op->src[0]->ne[0];
return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
} else {
if (K > 1) {
return false;
ggml/src/ggml-cuda/ssm-scan.cu
+19 -3
View File
@@ -163,13 +163,19 @@ __global__ void __launch_bounds__(d_state, 1)
const int lane = threadIdx.x % WARP_SIZE;
const int warp_idx = blockIdx.x * c_factor + warp;
ggml_cuda_pdl_sync();
// the last block can have unused warps when n_head*d_head is not a multiple of c_factor
if (warp_idx >= n_head * d_head) {
return;
}
const int head_idx = warp_idx / d_head;
const int head_off = (warp_idx % d_head) * sizeof(float);
const int seq_idx = blockIdx.y;
const int group_off = (head_idx / (n_head / n_group)) * d_state * sizeof(float);
ggml_cuda_pdl_sync();
// TODO: refactor strides to be in elements/floats instead of bytes to be cleaner and consistent with the rest of the codebase
const float * s0_warp = (const float *) ((const char *) src0 + src6[seq_idx] * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
const float * x_warp = (const float *) ((const char *) src1 + (seq_idx * src1_nb3) + (warp_idx * sizeof(float)));
@@ -246,7 +252,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
// NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
if (src3_nb1 == sizeof(float)) {
// Mamba-2
if (d_state == 128) {
if (d_state == 96) {
constexpr int threads = 96;
constexpr int num_warps = threads/WARP_SIZE;
const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
ggml_cuda_kernel_launch(ssm_scan_f32_group<96/WARP_SIZE, 96>, launch_params,
src0, src1, src2, src3, src4, src5, src6, dst,
src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
} else if (d_state == 128) {
constexpr int threads = 128;
constexpr int num_warps = threads/WARP_SIZE;
@@ -267,7 +283,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
} else {
GGML_ABORT("doesn't support d_state!=(128 or 256).");
GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
}
} else {
// Mamba-1