ggml : recurrent state rollback for ggml_ssm_scan (#26623)

* 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
+1
View File
@@ -1327,6 +1327,7 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx,
(uint32_t) src4->ne[1],
(uint32_t) src1->ne[2],
(uint32_t) ggml_nelements(src1),
(uint32_t) ggml_get_op_params_i32(dst, 0),
};
std::vector<wgpu::BindGroupEntry> entries = {