diff --git a/server/text_generation_server/layers/tensor_parallel.py b/server/text_generation_server/layers/tensor_parallel.py
index ec4ecf38..34b9c51e 100644
--- a/server/text_generation_server/layers/tensor_parallel.py
+++ b/server/text_generation_server/layers/tensor_parallel.py
@@ -1,7 +1,7 @@
 import torch
 from torch.nn import functional as F
 from typing import List
-from text_generation_server.layers.linear import get_linear
+from text_generation_server.layers.linear import get_linear, FastLinear
 
 
 class SuperLayer(torch.nn.Module):
diff --git a/server/text_generation_server/models/custom_modeling/mamba_modeling.py b/server/text_generation_server/models/custom_modeling/mamba_modeling.py
index ed638804..293051c2 100644
--- a/server/text_generation_server/models/custom_modeling/mamba_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/mamba_modeling.py
@@ -11,9 +11,9 @@ import torch.nn.functional as F
 from text_generation_server.layers import (
     SpeculativeHead,
     TensorParallelEmbedding,
-    FastRMSNorm,
     FastLinear,
 )
+from text_generation_server.layers.layernorm import FastRMSNorm
 
 from einops import rearrange
 from causal_conv1d import causal_conv1d_fn, causal_conv1d_update