Repository navigation
test(quantization): expand externalized composite ops coverage to rope and gather_mm - #131
Conversation
…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
There was a problem hiding this comment.
Hi @rohith500 thanks for helping expand test coverage for externalized composite ops!
My main comment is that we can restrict the structural tests (such as test_composite_boundary_quantized and test_composite_boundary_input_index_selects_those_args) to unique cases- and RoPE seems to cover the exact same structural guarantees as SDPA.
Having said that, I like the idea of having extensive coverage in the E2E test_composite_op_survives_prepare_and_finalize test which you have added already.
|
Thanks for the review @pkmandke! Updated. Removed |
Expands test coverage in
tests/quantization/test_composite_op_externalize.pyto coverRoPEandGatherMMfromcoreai_torch.composite_ops, which were previously only exercised withSDPAandRMSNorm.CompositeRoPEModelandCompositeGatherMMModelalong with theirExternalizeSpechelpers intests/models/composite.py.test_composite_op_survives_prepare_and_finalizeacross SDPA, RoPE, and GatherMM for both weight-only and w8a8 quantization.test_composite_boundary_quantizedandtest_composite_boundary_input_index_selects_those_argsto ensure multi-tensor inputs handle selective input port boundary quantization correctly._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) andmake check.