From a09542acc8fe83386bcbaecee022659d44eaa11c Mon Sep 17 00:00:00 2001 From: aseem wadhwa <7546478+aseemw@users.noreply.github.com> Date: Wed, 21 Oct 2026 06:23:43 -0700 Subject: [PATCH 1/2] update the qspec sharing logic for concat op --- .../_graph/_annotation_pattern_registry.py | 26 +++++++----- .../test_annotation_pattern_registry.py | 42 +++++++++++++++++++ 2 files changed, 57 insertions(+), 11 deletions(-) diff --git a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py index aa0c53da..6d422926 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py +++ b/src/coreai_opt/quantization/_graph/_annotation_pattern_registry.py @@ -886,9 +886,12 @@ def generate_patterns(cls) -> list[torch.fx.GraphModule]: def generate_qspec_sharing_constraints( cls, node: torch.fx.Node, qspecs: ProvisionalQSpecMap ) -> list[Constraint]: - """Inputs must observe together unless the output is per-channel along the - concat axis, in which case each keeps its own scale and they need only - agree on ``DTYPE``. + """Return the constraints tying the cat's input and output quantizers together. + + They share one quantizer only if every slot (inputs and output) is per-tensor, or + per-channel along a non-concat axis. Otherwise each input keeps its own quantizer, + since the inputs can differ in size along the concat axis, and they only agree on + ``DTYPE``. """ output_slot, input_slots = _shared_op_boundary_slots(node) if not _any_input_slot_populated(input_slots, qspecs): @@ -897,14 +900,15 @@ def generate_qspec_sharing_constraints( all_slots = frozenset([output_slot, *input_slots]) concat_dim = _concat_dim(node) - output_fields = qspecs.get(output_slot, ProvisionalQSpec()).fields - output_granularity = output_fields.get(FieldName.GRANULARITY) - - # A per-tensor granularity has ``axis is None``. - granularity_axis = output_granularity.value.axis if output_granularity is not None else None - per_channel_along_concat_axis = granularity_axis is not None and _same_axis( - granularity_axis, concat_dim, node - ) + per_channel_along_concat_axis = False + for slot in all_slots: + granularity = qspecs.get(slot, ProvisionalQSpec()).fields.get(FieldName.GRANULARITY) + if granularity is None or granularity.value is None: + continue + # A per-tensor granularity has ``axis is None``. + granularity_axis = granularity.value.axis + if granularity_axis is not None and _same_axis(granularity_axis, concat_dim, node): + per_channel_along_concat_axis = True if not per_channel_along_concat_axis: # Sharing an instance subsumes sharing fields, since the merged group # reconciles all of them. diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index 2a6ef05f..a49be969 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -419,6 +419,25 @@ def forward(self, x1, x2): return x +class CatOfParametersModel(nn.Module): + """Concat of two parameters with different sizes along the concat axis. + + Similar to token concat done in models such as SAM 2's mask decoder: + ``torch.cat([iou_token.weight, mask_tokens.weight], dim=0)``. + """ + + def __init__(self): + super().__init__() + self.token1 = nn.Embedding(1, 4) + self.token2 = nn.Embedding(3, 4) + self.linear = nn.Linear(4, 4) + + def forward(self, x): + tokens = torch.cat([self.token1.weight, self.token2.weight], dim=0) + out = self.linear(tokens + x) + return out + + class SimpleConcatModel(nn.Module): """Simple model with concat operation.""" @@ -2045,6 +2064,29 @@ def test_concat(self): # Final linear output quantizer is not associated with the others assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 + def test_cat_inputs_per_channel_along_concat_axis(self): + """ + Concat inputs that are per-channel along the concat axis keep separate quantizers. + + The inputs are parameters named ``weight``, so the default config gives them the + per-channel (axis 0) weight spec, while the output is per-tensor. The inputs differ + in size along axis 0 ([1, 4] and [3, 4]), so one shared per-channel quantizer can't + observe both. + """ + model = CatOfParametersModel().eval() + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),)) + + cat_nodes = [ + node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + token1_quantizer, token2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + assert token1_quantizer is not token2_quantizer + @pytest.mark.parametrize( "granularity", [ From 99e0774e3f40177ade24418e006d9124312029de Mon Sep 17 00:00:00 2001 From: aseem wadhwa <7546478+aseemw@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:28:47 -0700 Subject: [PATCH 2/2] update concat annotation tests --- .../test_annotation_pattern_registry.py | 159 +++++++++--------- 1 file changed, 80 insertions(+), 79 deletions(-) diff --git a/tests/quantization/test_annotation_pattern_registry.py b/tests/quantization/test_annotation_pattern_registry.py index a49be969..f94aaf7e 100644 --- a/tests/quantization/test_annotation_pattern_registry.py +++ b/tests/quantization/test_annotation_pattern_registry.py @@ -402,23 +402,6 @@ def forward(self, x): return torch.cat([x, xt], dim=0) # concat along batch, not channel/spatial -class SimpleCatModel(nn.Module): - """Simple model with concat operation.""" - - def __init__(self): - super().__init__() - self.linear1 = torch.nn.Linear(2, 2, bias=False) - self.linear2 = torch.nn.Linear(2, 2, bias=False) - self.linear3 = torch.nn.Linear(2, 2, bias=False) - - def forward(self, x1, x2): - x1 = self.linear1(x1) - x2 = self.linear2(x2) - x = torch.cat([x1, x2]) - x = self.linear3(x) - return x - - class CatOfParametersModel(nn.Module): """Concat of two parameters with different sizes along the concat axis. @@ -438,23 +421,6 @@ def forward(self, x): return out -class SimpleConcatModel(nn.Module): - """Simple model with concat operation.""" - - def __init__(self): - super().__init__() - self.linear1 = torch.nn.Linear(2, 2, bias=False) - self.linear2 = torch.nn.Linear(2, 2, bias=False) - self.linear3 = torch.nn.Linear(2, 2, bias=False) - - def forward(self, x1, x2): - x1 = self.linear1(x1) - x2 = self.linear2(x2) - x = torch.concat([x1, x2]) - x = self.linear3(x) - return x - - class AddLinearMulModel(nn.Module): def __init__(self): super().__init__() @@ -2000,69 +1966,71 @@ def test_shared_observer_preserves_per_block_when_pool_preserves_shape( expected_scale_shape = (2, 4, *spatial[:-1], spatial[-1] // 2) assert tuple(input_fq.qparams_calculator.scale.shape) == expected_scale_shape - def test_cat(self): + @pytest.mark.parametrize("cat_fn", [torch.cat, torch.concat], ids=["cat", "concat"]) + def test_cat_per_tensor_shares_quantizers(self, cat_fn): """ Given a model with - input1 -> relu1 - \ - -> Cat -> Linear - input2 -> relu2 / + input1 -> linear1 + \ + -> cat -> linear3 + input2 -> linear2 / - Check that concat inputs and outputs share the same quantizers. - """ - model = SimpleCatModel() - inp = (torch.randn(1, 2), torch.randn(1, 2)) + Check that with per-tensor quantization, the cat's inputs and output share one + quantizer, while linear3's output has its own. - config = QuantizerConfig( - global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None), - ) - - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare(example_inputs=inp) - quantizers = [ - module - for module in prepared_model.modules() - if isinstance(module, _DefaultFakeQuantizeImpl) - ] - assert len(quantizers) == 2 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2 - - # Final linear output quantizer is not associated with the others - assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 - - def test_concat(self): + torch.cat and torch.concat export to different aten ops (aten.cat, aten.concat), + each matched by its own entry in ConcatPattern. """ - Given a model with - input1 -> relu1 - \ - -> Concat -> Linear - input2 -> relu2 / + class CatModel(nn.Module): + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(2, 2, bias=False) + self.linear2 = nn.Linear(2, 2, bias=False) + self.linear3 = nn.Linear(2, 2, bias=False) - Check that concat inputs and outputs share the same quantizers. - """ - model = SimpleConcatModel() - inp = (torch.randn(1, 2), torch.randn(1, 2)) + def forward(self, x1, x2): + x = cat_fn([self.linear1(x1), self.linear2(x2)]) + out = self.linear3(x) + return out config = QuantizerConfig( - global_config=ModuleQuantizerConfig(op_state_spec=None, op_input_spec=None), + global_config=ModuleQuantizerConfig( + op_state_spec=None, + op_input_spec=None, + op_output_spec={"*": per_tensor_int8_qspec()}, + ), ) + quantizer = Quantizer(CatModel(), config) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(1, 2), torch.randn(1, 2))) - quantizer = Quantizer(model, config) - prepared_model = quantizer.prepare(example_inputs=inp) + cat_nodes = [ + node + for node in prepared_model.graph.nodes + if node.target in (torch.ops.aten.cat.default, torch.ops.aten.concat.default) + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + cat_input1_quantizer, cat_input2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + (cat_output_quantizer_node,) = cat.users + cat_output_quantizer = getattr(prepared_model, cat_output_quantizer_node.target) + (linear3,) = cat_output_quantizer_node.users + (linear3_output_quantizer_node,) = linear3.users + linear3_output_quantizer = getattr(prepared_model, linear3_output_quantizer_node.target) + + assert cat_input1_quantizer is cat_input2_quantizer + assert cat_input1_quantizer is cat_output_quantizer + assert linear3_output_quantizer is not cat_output_quantizer + # No other quantizers: the one the cat shares, and linear3's output quantizer. quantizers = [ module for module in prepared_model.modules() if isinstance(module, _DefaultFakeQuantizeImpl) ] assert len(quantizers) == 2 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_1 - assert prepared_model.activation_post_process_0 is prepared_model.activation_post_process_2 - - # Final linear output quantizer is not associated with the others - assert prepared_model.activation_post_process_0 != prepared_model.activation_post_process_3 def test_cat_inputs_per_channel_along_concat_axis(self): """ @@ -2087,6 +2055,39 @@ def test_cat_inputs_per_channel_along_concat_axis(self): ) assert token1_quantizer is not token2_quantizer + def test_cat_inputs_per_channel_along_other_axis(self): + """ + Concat inputs that are per-channel along a non-concat axis share one quantizer. + + The inputs ([1, 4] and [3, 4], concatenated along axis 0) are quantized per-channel + along axis 1, where they have the same size, so one quantizer can observe both. + """ + model = CatOfParametersModel().eval() + weight_spec = QuantizationSpec( + dtype=torch.int8, + qscheme=QuantizationScheme.SYMMETRIC, + granularity=PerChannelGranularity(axis=1), + ) + config = QuantizerConfig( + global_config=ModuleQuantizerConfig( + op_state_spec={"weight": weight_spec}, + op_input_spec={"*": per_tensor_int8_qspec()}, + op_output_spec={"*": per_tensor_int8_qspec()}, + ), + ) + quantizer = Quantizer(model, config) + prepared_model = quantizer.prepare(example_inputs=(torch.randn(4, 4),)) + + cat_nodes = [ + node for node in prepared_model.graph.nodes if node.target == torch.ops.aten.cat.default + ] + assert len(cat_nodes) == 1 + cat = cat_nodes[0] + token1_quantizer, token2_quantizer = ( + getattr(prepared_model, input_node.target) for input_node in cat.args[0] + ) + assert token1_quantizer is token2_quantizer + @pytest.mark.parametrize( "granularity", [