support TP

Signed-off-by: Yu Chin Fabian Lim <fabian.lim@gmail.com>

Author: fabianlim

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -132,7 +132,10 @@ def forward_cuda(
             return x * nn.functional.silu(gate.to(
                 torch.float32)).to(input_dtype)
 
-        if self.tp_size > 1 or self.n_groups != 1:
+        if (
+            ((self.n_groups % self.tp_size) != 0) or 
+            self.n_groups != 1
+        ):
             return self.forward_native(x, gate)
 
         return rms_norm_gated(x,