Architecture design of the DiT module (update soon)