One assistant draft step over B rows: the last token's embedding and the hidden it follows [B, 1, H] (borrowed), the donors, and the rows' position (one shared, or [B]); owned tokens and next hidden out.
One assistant draft step over B rows: the last token's embedding and the hidden it follows [B, 1, H] (borrowed), the donors, and the rows' position (one shared, or [B]); owned tokens and next hidden out.