annotated-mamba icon indicating copy to clipboard operation
annotated-mamba copied to clipboard

associative ssm op order

Open kpich opened this issue 9 months ago • 1 comments

i might be wrong, but in the colab/blog, i think the $\oplus$ op used to do the associative scan for the selective state space model should have, as the value of its first output, $a_2 a_1$ (rather than $a_1 a_2$), reflecting the fact that the leftmost $A$ transform gets applied first.

(Since the Mamba matrices are diagonal and therefore commutative it doesn't actually matter here I guess, I just found this initially confusing in the presentation).

It looks to be correct in the triton first_order_op but is I think reversed in the reference pytorch op and latex above it.

Thanks for this terrific writeup! It really clarified some things for me, thanks

kpich avatar Apr 30 '24 21:04 kpich