[林SO (Linxiushen)]

Describe your changes

TransformersDummyDataset.get_past_kv_inputs() calculates each KV head's width after dividing the attention-head count by world_size. For hidden_size=32, four attention heads and world_size=2, it produces width 16 instead of 8. It also ignores an explicit model head_dim; a local Llama configured with head_dim=12 receives width-8 cached tensors and fails when extending the cache.

Use the already-resolved ModelWrapper.head_dim for the width. The existing KV-head sharding remains in place. Regression coverage includes MHA/GQA, missing KV-head configuration, single-rank MQA, empty/populated caches, float32/float16, explicit/derived widths, and a randomly initialized two-layer Llama that extends the resulting cache in a real CPU forward pass.

Validation on base 1ef7d129570ea0faafeb2554ef1208de1080d12d:

• Unmodified production source: 29 failed, 16 passed in the new regression file, including the Llama cache-concatenation error.
• Fixed source: 65 passed with python -m pytest -q test/data_container/test_transformers_dummy_dataset.py test/data_container/test_dataset.py test/model/test_kv_cache_config.py (normal repository fixtures, fresh local --basetemp, Hugging Face offline mode).
• All applicable checks passed with lintrunner --verbose olive/data/component/dataset.py test/data_container/test_transformers_dummy_dataset.py, including Pylint, Ruff, formatting, whitespace and EditorConfig; git diff --check also passed. No further formatting fixes were needed.
• Windows / Python 3.12.14 / PyTorch 2.14.1+cpu / Transformers 5.20.1. The isolated environment uses NumPy 1.26.4, ONNX 1.20.1, ONNXScript 0.6.0, Datasets 4.4.1 and PyArrow 21.0.0 to satisfy the repository's test constraints. An initial dependency/fixture failure was resolved before the unchanged-source reproduction above. The successful run emitted one existing OpenTelemetry deprecation warning.

The full suite, ONNX export, distributed execution and GPU execution were not run. No pretrained weights were downloaded; this is cache-shape/compatibility validation, not model-quality or performance evidence. KV-head counts still use the existing floor division; KV-head replication and validation of unsupported sharding configurations are outside this change.

AI assistance: this patch and tests were developed and reviewed with Codex; the reported commands were executed locally.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass. The focused 65-test selection passes; the full suite was not run.
• Update documents if necessary. No documentation change is needed for this correction.
• Lint the changed files with lintrunner (all applicable rules passed; no further fixes required).
• Is this a user-facing change? Release note: Correct dummy KV-cache head widths for tensor-parallel inputs and models with an explicit attention head dimension.

(Optional) Issue link

Found by inspecting the dummy-cache generator; no new issue was opened.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟢 Approval recommended

The narrow correction preserves sharding behavior and has targeted regression coverage, with no blocking issues identified.

No open issues 🎉

What changed in this PR

Corrects dummy KV-cache widths to honor model head dimensions while preserving existing KV-head sharding.

Changes:

• Uses the resolved ModelWrapper.head_dim instead of recalculating width from sharded attention heads.
• Adds regression coverage for cache shapes, dtypes, sharding, and a local Llama forward pass.

File	Description
test/​data_container/​test_transformers_dummy_dataset.py	Adds shape and cache-extension regression tests.
olive/​data/​component/​dataset.py	Uses the model’s resolved head dimension for dummy caches.

_{Review effort: Balanced}