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-webgpu/ggml-webgpu.cpp
+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 = {
ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl
+11
View File
@@ -41,6 +41,7 @@ struct Params {
n_seq_tokens: u32,
y_elems: u32,
K: u32,
};
@group(0) @binding(0) var<storage, read_write> s_in: array<f32>;
@@ -123,6 +124,7 @@ fn main(
let head_seq = wg_linear / params.d_inner;
let ir = head_seq % params.n_head;
let i3 = head_seq / params.n_head;
let n_seqs = params.y_elems / (params.n_seq_tokens * params.n_head * params.d_inner);
let state_slot = read_state_slot(i3);
let g = ir / (params.n_head / params.n_group);
@@ -179,6 +181,15 @@ fn main(
#endif
s_prev = s;
let slot = params.n_seq_tokens - 1u - token;
if (slot > 0u && slot < params.K) {
let snapshot_idx =
params.offset_dst + params.y_elems + tid + i1 * params.d_state +
ir * (params.d_state * params.d_inner) +
(slot * n_seqs + i3) * (params.d_state * params.d_inner * params.n_head);
dst[snapshot_idx] = s;
}
#ifdef USE_SUBGROUP_REDUCTION
#ifdef XBC_OVERLAP
let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx));