`Mamba2.step()` handles `D` incorrectly when `D_has_dim=True`

[GiftedNovaHD]

Hi, I think there's a bug in Mamba2.step() when we set D_has_dim = True.

In __init__, self.D is initialized as

self.D = nn.Parameter(torch.ones(self.d_ssm if self.D_has_hdim else self.nheads, device=device))

Thus, when D_has_hdim=True, self.D has shape (d_ssm,) = (nheads * headdim,).

I believe the forward path appear to handle this correctly by reshaping D as (h, p)

So when D_has_hdim=True, self.D has shape (d_ssm,) = (nheads * headdim,).

But I don't think step() makes the same distinction because for this line, it implicitly assumes that self.D is per head, with shape (nheads,), and not per head-dim: (nheads * headdim,).

So for D_has_hdim = True, step() seems inconsistent with forward(), as forward() reshapes self.D from (h * p,) -> (h,p), whereas step() instead treats it as it were (h,), as mentioned earlier

I can contribute a PR to fix this

[ccs1112]

I hit this too and traced it to the same root cause. Both step() paths assume self.D is (nheads,) — the "h -> h 1" rearrange at line 319 and the repeat("h -> h p") at line 327 both break when D_has_hdim=True gives shape (nheads * headdim,).

I have a fix + regression test at #893 if you'd like to take a look.