ReplaySSM: Cache SSM Inputs, Not State

Current SSM decoding updates and writes recurrent state back to HBM every step. ReplaySSM instead caches recent inputs, reconstructing the state only when needed and otherwise computing the output directly from the cache. Figure 1a. ReplaySSM caches recent SSM inputs instead of storing the recurrent state every step, reconstructing states on the fly and writing back only when the buffer is full.…