Onboard a streaming mode model similar to Moshi.

🚀 The feature, motivation and pitch

Using Moshi as a case study here but I think the MoE style probably applies to any model with time step reasoning.

Moshi tracks the current start pos as state, it uses 2 values to do this. 1 tensor and 1 int. Its not obvious why 2 variables are needed. I have an open Issue asking about it https://github.com/kyutai-labs/moshi/issues/255. I think the int one could just be rewritten as an IO tensor though if needed. It seems pretty equivalent to start_pos from llama.

The next problem is that Moshi has something similar to MoE but for different time steps. As seen here. https://github.com/kyutai-labs/moshi/blob/main/moshi/moshi/modules/transformer.py#L278

I think that will need to be rewritten like how MoE was in gpt-fast where the weights of all the linear layers are combined into one tensor and then sliced out. https://github.com/pytorch-labs/gpt-fast/blob/main/mixtral-moe/model.py#L187-L201

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @larryliu0820 @mergennachin @cccclai @helunwencser @jackzhxng

Start by reading Moshi's moshi/moshi/modules/transformer.py around line 278 and the referenced gpt-fast mixtral model.py lines 187-201. Compare how streaming state and time-step-specific linear layers are represented, then trace the relevant Executorch model onboarding path. Done means a Moshi-like streaming model is supported with its state handling and time-step weights integrated.