Optimize GDN state updates for decode

[quic-mkshirsa]

Summary

Generalizes Qwen3.5-MoE GDN retained-state updates to support full-batch kv_headpar decode for batch sizes 1, 2, 4, 8, and larger.

Problem

For kv_headpar decode, the compute batch covers the complete KV-cache batch. The gather/scatter operations are therefore unnecessary and cause compilation failures
for BS2 and BS4.

Solution

• Reorder full-batch inputs into physical cache-slot order once at the decoder boundary.
• Clear batch_index before cache updates so every layer replaces its complete retained state directly.
• Restore final outputs to logical request order.
• Preserve the existing gather/scatter path when the compute batch covers only part of the cache.
• Permit full-state decode when full_batch_size == kv_cache_batch_size.
• Update ONNX dynamic axes and dummy inputs for full-batch export.
• Preserve existing singleton prefill and kv_batch_fold behavior.

Validation

• Verified full-state and routed-state parity for decode BS1, BS2, and BS4 with permuted batch_index.
• Verified full-state compile modes for BS1, BS2, BS4, and BS8.
• Verified partial-batch updates preserve inactive cache rows.
• Verified prefill and decode ONNX Runtime parity.
• Verified exported full-state graphs contain no retained-state CB gather/scatter operations.
• 18 targeted tests passed.
• 33 broader Qwen3.5 tests passed.

[Rishin Raj (quic-rishinr)]

Please add a reproducer config test for this issue, use qwen3.6 dummy random model for the test

[Vinayak Baddi (vbaddi)]

CI-Ready