Skip to content

test(quantization): expand externalized composite ops coverage to rope and gather_mm - #131

Open
rohith500 wants to merge 1 commit into
apple:mainfrom
rohith500:feat/test-externalize-composite-ops
Open

rohith500 wants to merge 1 commit into
apple:mainfrom
rohith500:feat/test-externalize-composite-ops

Conversation

@rohith500

Copy link
Copy Markdown
Contributor

Expands test coverage in tests/quantization/test_composite_op_externalize.py to cover RoPE and GatherMM from coreai_torch.composite_ops, which were previously only exercised with SDPA and RMSNorm.

  • Added CompositeRoPEModel and CompositeGatherMMModel along with their ExternalizeSpec helpers in tests/models/composite.py.
  • Parameterized test_composite_op_survives_prepare_and_finalize across SDPA, RoPE, and GatherMM for both weight-only and w8a8 quantization.
  • Parameterized test_composite_boundary_quantized and test_composite_boundary_input_index_selects_those_args to ensure multi-tensor inputs handle selective input port boundary quantization correctly.
  • Extracted shared tensor sample helpers (_single_tensor_sample, _rope_sample, _gather_mm_sample) to keep test fixtures clean.

Tested with pytest tests/quantization/test_composite_op_externalize.py (20 passed) and make check.

…e and gather_mm

- Add rope_externalize_spec and gathermm_externalize_spec with CompositeRoPEModel and CompositeGatherMMModel
- Parameterize test_composite_op_survives_prepare_and_finalize across SDPA, RoPE, and GatherMM for w8 and w8a8 quantization modes
- Parameterize test_composite_boundary_quantized and test_composite_boundary_input_index_selects_those_args for multi-input composite ops
- Extract reusable sample tensor generators to keep test tables clean and maintainable
@aseemw
aseemw requested a review from pkmandke October 5, 2026 18:03

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants