Some problem

[xiaoshideta]

Hello, I have reviewed your source code and found that your generation approach is somewhat similar to the dual cache in Fast-DLLM. Since you conducted testing on OpenCompass, I migrated your code to the LLM-eval architecture of Fast-DLLM. I noticed a decrease in speed compared to the original Fast-DLLM, even when I do not perform pruning operations. The storage of the KV cache is slower than the past KV in the original Fast-DLLM.

A possible reason might be the dimensions of K and V? In the original Fast-DLLM, the dimensions are (batch size, sequence length, N), whereas in your approach, they are (batch size, number of heads, sequence length, dimension)?

[xiaoshideta]

1. I examined the code and noticed that it does not use parallel decoding similar to fast-dllm. Would implementing this further enhance the current performance?

2. Additionally, I observed that there are two complete forward passes of the input in each block, while the subsequent decoding uses truncated input. Compared to fast-dllm, this adds an extra full forward pass overhead. For a task generating 256 tokens, this would result in 8 additional passes, which might be unacceptable.

3. For each generation task, an extra attention calculation is required per block for pruning. I am particularly concerned about this because while SDPA (Scaled Dot-Product Attention) accelerates inference, additional non-accelerated computations for pruning might have a negative impact?

[sunnyremm]

Thanks for the careful reading and great questions — and sorry for the late reply!

1. Parallel decoding
   You're right that we don’t explicitly implement parallel decoding like fast-dLLM. That said, our method is compatible with the standard diffusion decoding pipeline, so in principle it can be combined with step-skipping / parallel decoding. We haven’t tried it yet, but it’s definitely an interesting direction.
2. Extra full forward pass
   Yes, we do add one extra full forward pass per block. This is mainly to get a more stable saliency estimate — otherwise pruning can easily hurt quality. In practice, the cost is largely offset since pruning shortens the sequence and makes later steps cheaper. We still see overall speedup, though for very short generations this overhead can be more noticeable.
3. Pruning overhead
   While pruning does add some extra attention computation, we reduce it a lot by doing token pooling within each block. So instead of attending over all tokens, we only use a pooled representation, and the overhead is quite small compared to the full attention (especially with SDPA).

If you have further questions, feel free to reach out via email or WeChat — happy to discuss more :)