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
@@ -1376,9 +1376,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
ggml_is_contiguous_rows(op->src[1]) &&
|
||||
ggml_is_contiguous_rows(op->src[2]) &&
|
||||
ggml_is_contiguous_rows(op->src[3]);
|
||||
case GGML_OP_SSM_CONV:
|
||||
case GGML_OP_SSM_SCAN:
|
||||
return has_simdgroup_reduction;
|
||||
case GGML_OP_SSM_CONV:
|
||||
return has_simdgroup_reduction;
|
||||
case GGML_OP_RWKV_WKV6:
|
||||
case GGML_OP_RWKV_WKV7:
|
||||
return true;
|
||||
|
||||
Reference in New Issue
Block a user