Question regarding fixed token length

[xiaoBOSS97]

Hi:
Thank you for your awesome work!
I have a question regarding PI05 and the fixed token length in your implementation. In PI05, the robot state should also be treated as part of the prompt, combined with language instruction and image encoding before tokenization. So the token length after the tokenizer can vary due to changes in the number in the robot state. But in your implementation, the token length must be fixed after set_prompt() if I understand correctly. Otherwise, it has to recapture the graph. And in your PI05 implementation, the robot state does not seem to be given as an input during inference. Can you explain to me why the state is missing, and do you have any plan or idea to deal with variant prompt length?
Thanks a lot!

[LiangSu8899]

Hi! Thanks for the detailed report!

This should be handled on the current main branch now. The Pi0.5 state prompt path was reworked in #50 and follow-up updates, and it should work on both RTX and Thor backends.

Please update to the latest code and take a look at the “Pi0.5 State Prompts” section in USAGE.md.

The recommended public API is:

model.predict(images=..., prompt="...", state=robot_state)

or, if using the lower-level API:

model.set_prompt(prompt, state=robot_state)

State is part of the prompt by design, following the OpenPI / Pi0.5 behavior. Variable token lengths are also supported now: on RTX, different token lengths are cached as separate pipeline buckets, and repeated lengths reuse the captured graph; same length with different state only updates embeddings in place. Thor also supports in-place update for the same token length, while new lengths rebuild/capture once.

You can also pre-warm expected state prompt lengths with:

model.warm_state_prompt_buckets(images, prompt, states=[...])

Please try again with the latest main and the documented API path. If you still see repeated recapture or very slow inference after that, feel free to share the exact call path / logs and I’ll take another look.

Thx!!

[xiaoBOSS97]

Thanks for the quick reply! I will definitely try.

[kc-shilin]

Hi,
Thanks for the amazing work! For handling variable token lengths, would adding a mask tensor with a static input shape be a good alternative approach here? It might help avoid creating separate pipeline buckets for different lengths.

[LiangSu8899]

Yeah, that’s a good point.

My design usually tries to control resources as tightly as possible, which matters a lot for edge deployment. So the current main path is basically “don’t pay for padding”: each real token length gets its own cached pipeline/graph. The downside is exactly what you mentioned — if a new length misses the cache, it needs to capture once.

Your static-shape + mask approach is also a good design. It clearly avoids repeated captures, and I can probably add it as another entry point.

One question: for your real deployment use case, which path do you think would be more common / more useful — the exact-length fastest path, or the static-shape masked path with more stable graph reuse?

I’d like to use that feedback to decide which one should be the preferred default.

[kc-shilin]

Thanks for the thoughtful response! It makes total sense that controlling resources tightly is a priority, especially for edge deployment.

For the trade-offs, I think both paths have very strong merits depending on the user's constraints.

If the primary goal of FlashRT is ultra-lean edge deployment, keeping the exact-length path as the default makes a lot of sense, with the static-shape approach as a highly recommended configuration for users needing latency stability. Personally, for my current use case, I lean toward the static-shape path for the peace of mind of stable graph reuse, as avoiding sudden latency spikes is really important for my workload.

Super excited to see this feature added as another entry point!

[LiangSu8899]

Got it! thanks for ur feeback
Will add this as an explicit optional path

[xiaoBOSS97]

Hi, thank you both for your feedback and contribution.
I have already tried your method of adding caches. But unfortunately, the token length varies a lot in my case. Therefore, I tried to use a static max token length with padding to make the pipeline adaptable to different input token lengths. This increases the inference time slightly from ~28ms to ~32ms, which I think is a good trade-off. In conclusion, I think adding an optional path could be quite useful. Hope this experience helps.
Best Regard

[LiangSu8899]

Hi! Actually I updated a split-kv :#87
and just need to pay for ~1ms

You could create a small lookup table based on your max length range in your deployment pipeline, so that the overhead can be controlled at around 1ms