Ok I tried it. And it's faster. Though I'm still in verifying quality phase.

But only works for Gemma4, or other attention-only models. (i.e. not the Qwen3.8 models I had working with the other way)

I'm getting 63ms per answer for 1 question, 79 for 3, 93.7ms for 5. So scales nicely, too. ~55ms fixed cost + ~7-8ms per additional question.

Much better. Thanks. I'll commit and share work after testing it some more.

EDIT: how can I credit you in the commit body?

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.

I think you're right. Better. I will have to think through if it would be faster or slower. If understand what you're getting at with this.. broken analogy...

What I described was -- we have a bunch of orders to the kitchen, all of which have the same "base" meal but different topics.

> Cook the "base" meal once, divide servings onto multiple plates, then add different toppings to each plate.

vs what you suggest:

> Put the base meal and every topping through the kitchen together, but use some kind of dividers so the toppings never mix up together.

Except.. ok, that analogy is confusing lol.