fix: adjust typos and docker build

Dockerfile

@@ -154,17 +154,17 @@ COPY server/Makefile-vllm Makefile
 # Build specific version of vllm
 RUN make build-vllm-cuda
 
-# Build megablocks
-FROM kernel-builder as megablocks-builder
-
-RUN pip install git+https://github.com/OlivierDehaene/megablocks@181709df192de9a941fdf3a641cdc65a0462996e
-
 # Build mamba kernels
 FROM kernel-builder as mamba-builder
 WORKDIR /usr/src
 COPY server/Makefile-selective-scan Makefile
 RUN make build-all
 
+# Build megablocks
+FROM kernel-builder as megablocks-builder
+
+RUN pip install git+https://github.com/OlivierDehaene/megablocks@181709df192de9a941fdf3a641cdc65a0462996e
+
 # Text Generation Inference base image
 FROM nvidia/cuda:12.1.0-base-ubuntu20.04 as base
 

server/text_generation_server/models/custom_modeling/mamba_modeling.py

@@ -137,7 +137,7 @@ def step(self, hidden_states, conv_state, ssm_state):
             z = _z[i:i+1, -1, :]
             x_db = self.x_proj(x)
             dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1)
-            df = self.dt_proj_no_bias(x)
+            dt = F.linear(dt, self.dt_proj.weight)
             y = selective_state_update(
                 ssm_state[i:i+1,:,:], x, dt, self.negA, B, C, self.D, z=z, dt_bias=self.dt_proj.bias, dt_softplus=True
             )