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:
co-authored by
Georgi Gerganov
Gaurav Garg
parent
4c1a0af40d
commit
1692f9e50b
@@ -2064,6 +2064,7 @@ bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_te
|
||||
params.src5 = *node->src[5];
|
||||
params.src6 = *node->src[6];
|
||||
params.dst = *node;
|
||||
params.K = ggml_get_op_params_i32(node, 0);
|
||||
|
||||
bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", ¶ms, sizeof(params), 0xFFFFFFFF);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user