diff --git a/litedram/frontend/bist.py b/litedram/frontend/bist.py
index 8556c27..099887f 100644
--- a/litedram/frontend/bist.py
+++ b/litedram/frontend/bist.py
@@ -186,13 +186,18 @@ class _LiteDRAMBISTGenerator(Module):
                 addr_gen.ce.eq(1),
                 NextValue(cmd_counter, cmd_counter + 1),
                 If(cmd_counter == (self.length[ashift:] - 1),
-                    NextState("DONE")
+                    NextState("AWAIT_FIFO_EMPTY")
                 ).Elif(~self.run_cascade_in,
                     NextState("WAIT")
                 )
             ),
             NextValue(self.ticks, self.ticks + 1)
         )
+        fsm.act("AWAIT_FIFO_EMPTY",
+            If(~dma.fifo.source.valid,
+                NextState("DONE"),
+            ),
+        )
         fsm.act("DONE",
             self.run_cascade_out.eq(1),
             self.done.eq(1)
diff --git a/litedram/frontend/dma.py b/litedram/frontend/dma.py
index e36a8dc..0b28e1b 100644
--- a/litedram/frontend/dma.py
+++ b/litedram/frontend/dma.py
@@ -200,8 +200,7 @@ class LiteDRAMDMAWriter(Module, AutoCSR):
             raise NotImplementedError
 
         # FIFO -------------------------------------------------------------------------------------
-        fifo = stream.SyncFIFO([("data", port.data_width)], fifo_depth, fifo_buffered)
-        self.submodules += fifo
+        self.submodules.fifo = fifo = stream.SyncFIFO([("data", port.data_width)], fifo_depth, fifo_buffered)
 
         if is_native:
             self.comb += cmd.we.eq(1)