Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml : recurrent state rollback for ggml_ssm_scan (#26623)

Browse Source 
* Initial changes for Recurrent state rollback for nemotron for cpu and cuda

* Removing CPU RS rollback. Will enable it in subsequent PRs

* addition of test case

* Removing assert and calling runtime API to check if op is supported

* removing extra API and updating the call sites for K

* replace static cuda detection to runtime fused_op api

* address review comments and fallback when SSM rollback not supprted

* Adding changes for supporting RS-rollback in CPU. Also added test-backend-ops for cpu and cuda

* removing memory manipulation as rs rollback is now supported in CPU

* removing the static probe which is not needed now

* correcting the format

* address review comments

* enabling test for all the backends, unsupported backends will fallback to CPU

* Apply suggestions from code review

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

* choose different graph based on the result of fused_ssm_op is supported or not and also handled memory->n_rs_seq >1 case incase of op is not supported

* Support K > 1 in ssm_scan for all backends

* Fix CI Issues

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Gaurav Garg <gaugarg@nvidia.com>
This commit is contained in:
lnigam
2026-08-14 17:20:40 +03:00
committed by GitHub
co-authored by Georgi Gerganov Gaurav Garg
parent 4c1a0af40d
commit 1692f9e50b
25 changed files with 291 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+5 -2
View File
@@ -1861,6 +1861,7 @@ struct vk_op_ssm_scan_push_constants {
uint32_t nb42, nb43, nb52, nb53;
uint32_t s_off;
uint32_t n_head, d_head, n_group, n_tok;
uint32_t n_seq, K;
};
struct vk_op_ssm_conv_push_constants {
uint32_t nb01, nb02;
@@ -12731,7 +12732,8 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx,
(uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
(uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
(uint32_t)s_off,
n_head, head_dim, n_group, n_tok
n_head, head_dim, n_group, n_tok,
n_seq, (uint32_t) ggml_get_op_params_i32(dst, 0)
};
vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
@@ -19417,8 +19419,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
} else if (tensor->op == GGML_OP_ADD_ID) {
tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
} else if (tensor->op == GGML_OP_SSM_SCAN) {
const int32_t K = ggml_get_op_params_i32(tensor, 0);
tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
src_clone[3], src_clone[4], src_clone[5], src_clone[6], K);
} else if (tensor->op == GGML_OP_SSM_CONV) {
tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]);
} else if (tensor->op == GGML_OP_ROLL) {
ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp
+10
View File
@@ -33,6 +33,8 @@ layout(push_constant) uniform PushConstants {
uint d_head;
uint n_group;
uint n_tok;
uint n_seq;
uint K;
};
float softplus(float x) {
@@ -114,6 +116,14 @@ void main() {
if (lane == 0) {
d[y_base_idx + i * stride_y] = state_sum;
}
const uint slot = n_tok - 1u - i;
if (slot > 0u && slot < K) {
const uint snapshot_base_idx = s_base_idx + slot * n_seq * (nb03 / 4u);
[[unroll]] for (uint j = 0; j < c_factor; j++) {
d[snapshot_base_idx + SUBGROUP_SIZE * j + lane] = state[j];
}
}
}
// write back the state