model : add LLADA 2.0 diffusion support

[wsbagnsv1]

I recently created feature request #16973, but decided to implement it myself.

Status: Work in Progress (WIP)

This is my first PR here. I aimed for backwards compatibility, but the implementation currently contains some temporary workarounds (such as the ubatch one in the diffusion CLI) which I will remove once the underlying issues are resolved.

Current State:

• Success: Inference seems to work generally well and the output is coherent.
• Known Issues: I am currently tracking down bugs regarding unexpected tokens and an issue with max length.

Any feedback or assistance would be appreciated!

[wsbagnsv1]

This is the output with this prompt "what is the meaning of life?", I tested it on cuda and cpu and both work generally fine.

[wsbagnsv1]

okay looked over the code a bit again, the weird tokens seem to be an issue with diffusion cli rendering itself on windows and the n_ubatch stuff seems to be a bit weird and also an issue with diffusion cli itself, ill take a close look tomorrow (;

[am17an]

I didn't review all of this. But preferably there should be no LLaDA 2.0 specific stuff in diffusion-cli, it should all be covered by diffusion-cli parameters and maybe gguf parameters. Since it already supports Dream and LLaDA 1.0 without any specific case handling, it should be possible to do this for LLaDA 2.0 as well. As such I don't find it the sampling to be very different from block-based sampling used in LLaDA 1.0

[deniaud]

Is there a chance we can use LLADA in a masking mode with multiple masks within the input context, rather than just as a standard auto-regressive model?

[wsbagnsv1]

I didn't review all of this. But preferably there should be no LLaDA 2.0 specific stuff in diffusion-cli, it should all be covered by diffusion-cli parameters and maybe gguf parameters. Since it already supports Dream and LLaDA 1.0 without any specific case handling, it should be possible to do this for LLaDA 2.0 as well. As such I don't find it the sampling to be very different from block-based sampling used in LLaDA 1.0

Thanks for the feedback! Ill give it a try (;

[wsbagnsv1]

I didn't review all of this. But preferably there should be no LLaDA 2.0 specific stuff in diffusion-cli, it should all be covered by diffusion-cli parameters and maybe gguf parameters. Since it already supports Dream and LLaDA 1.0 without any specific case handling, it should be possible to do this for LLaDA 2.0 as well. As such I don't find it the sampling to be very different from block-based sampling used in LLaDA 1.0

Also ill try to implement a general model independent eos stop for such models with blocks, since i couldnt find any functionality that would do the same, should help future models too and ill use a gguf parameter to mark llada2.0 for it (;

[wsbagnsv1]

Okay i think ill need to do the same for the threshold feature as well as the way it truncates blocks to save time, so ill create parameters in the ggufs that will trigger and allow it to use threshold as well as truncate instead of full standard one. That way it should be fully backwards compatible and model independent (;

[am17an]

Please keep all perf improvements like hybrid diffusion (which IMO should be called use_kv_cache or something which can be toggled via hparams) for a separate PR. This PR should just focus on getting the model in correctly

[wsbagnsv1]

Okay then i probably will use that flag from the other pr as a toggle for the truncate stuff as well, so this would solve that problem too (;

[wsbagnsv1]

Since that one is also a performance improvement by not calculating the whole block

[wsbagnsv1]

one question though should that also be a cli arg in the other pr?

[wsbagnsv1]

Hmm, i think there was an issue when not using truncate for some reason the output becomes gibberish? Like with the --diffusion-hybrid flag (which i now use to toggle the truncate as well) it works fine but without it it stops with a "ΓÇ£" token 🤔

[wsbagnsv1]

@wsbagnsv1 as I said earlier, you should keep any kv-cache related stuff for a subsequent PR. Also I would suggest formatting your comments to keep them short and precise.

well the issue with that part is that i dont see a way to get rid of the truncate stuff without changing that part or a custom flag with truncate toggle. I might be missing something but i have no idea what /:

[am17an]

I'm not entirely sure what you're saying, but you can ask a question in Q&A as I'm not too familiar with the kv cache. Nevertheless, your original implementation without any "hybrid diffusion" stuff was closer to being in a state where it could be merged. So my suggestion to you would be to get that in first, before tackling optimisations.

[am17an]

well the issue with that part is that i dont see a way to get rid of the truncate stuff without changing that part or a custom flag with truncate toggle. I might be missing something but i have no idea what /:

Leave a comment explaining why you need it and it should be fine to add truncate

[wsbagnsv1]

Well in that case i would need to do the masking stuff in diffusion cli and use a truncate cli arg which would also suck /:

[wsbagnsv1]

Ill gonna ask in Q&A thanks for your feedback though! Maybe they have an idea how to get this done the right way (;