Enforced kv layer ordering

[oandreeva-nv]

Overview:

Besides kv layer order enforcement, here's what I do:

• Eliminated redundant operations:

• Replaced list(kv_caches.values())[0] with next(iter(kv_caches.values())) - more efficient for getting the first tensor
• Removed the unnecessary tensors = list(kv_caches.values()) variable

• The original code iterated over kv_caches.values() three times, now it only does it once (in the validation check)

Details:

Where should the reviewer start?

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.