Comment by mungoman2

Ohh this is really cool.

So one could pack all of common state, every question, every answer in the same prefill, using attention mask to only let them attend to their logical parent.

Then additionally do position encoding for token based on their logical position rather than physical.

Then the diffusion step also applies an attention mask to prevent bidirectional attention between answers.