[Puzzletron] Fix KV-cache dimension calculation for models with custom head_dim (#2482) - #2484

[Puzzletron] Fix KV-cache dimension calculation for models with custom head_dim (#2482)#2484

musi22 wants to merge 2 commits into

NVIDIA:feature/puzzletron_v2NVIDIA/Model-Optimizer:feature/puzzletron_v2from

musi22:fix/puzzletron-kv-dim-2482musi22/Model-Optimizer:fix/puzzletron-kv-dim-2482Copy head branch name to clipboard