From 387cba8d30c3eebcf578e82f7c9bab2ef2a8b679 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Wed, 19 Feb 2025 04:49:26 +0000 Subject: [PATCH 1/4] shape infer dispatcher for GatherBlockQuantized --- onnxruntime/python/tools/symbolic_shape_infer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index c25e5dc087614..f469633b7236c 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -144,6 +144,7 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""): "Equal": self._infer_symbolic_compute_ops, "Floor": self._infer_symbolic_compute_ops, "Gather": self._infer_Gather, + "GatherBlockQuantized": self._infer_Gather, "GatherElements": self._infer_GatherElements, "GatherND": self._infer_GatherND, "Identity": self._pass_on_shape_and_type, From 96f92ff6c1453facd2cdb66cab97aa1109d2cab0 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Wed, 19 Feb 2025 05:28:12 +0000 Subject: [PATCH 2/4] move to contrib ops --- onnxruntime/python/tools/symbolic_shape_infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index f469633b7236c..f4da8d1de039c 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -144,7 +144,6 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""): "Equal": self._infer_symbolic_compute_ops, "Floor": self._infer_symbolic_compute_ops, "Gather": self._infer_Gather, - "GatherBlockQuantized": self._infer_Gather, "GatherElements": self._infer_GatherElements, "GatherND": self._infer_GatherND, "Identity": self._pass_on_shape_and_type, @@ -203,6 +202,7 @@ def __init__(self, int_max, auto_merge, guess_output_rank, verbose, prefix=""): "EmbedLayerNormalization": self._infer_EmbedLayerNormalization, "FastGelu": self._infer_FastGelu, "GatedRelativePositionBias": self._infer_GatedRelativePositionBias, + "GatherBlockQuantized": self._infer_Gather, "Gelu": self._infer_Gelu, "GemmFastGelu": self._infer_GemmFastGelu, "GemmFloat8": self._infer_GemmFloat8, From cc13b14483b65e4454f855f1a596f6e8c9e8c169 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Thu, 20 Feb 2025 19:45:28 +0000 Subject: [PATCH 3/4] skip in _onnx_infer_single_node --- onnxruntime/python/tools/symbolic_shape_infer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index f4da8d1de039c..86e0939e90543 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -460,6 +460,7 @@ def _onnx_infer_single_node(self, node): "BiasGelu", "EmbedLayerNormalization", "FastGelu", + "GatherBlockQuantized", "Gelu", "GemmFastGelu", "LayerNormalization", From d462c5d4a31188b248c228801c2601749c8f9d13 Mon Sep 17 00:00:00 2001 From: Jambay Kinley Date: Thu, 20 Feb 2025 20:07:18 +0000 Subject: [PATCH 4/4] elem type from scales --- onnxruntime/python/tools/symbolic_shape_infer.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 86e0939e90543..fda604d3f729b 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1120,10 +1120,17 @@ def _infer_Gather(self, node): # noqa: N802 axis = handle_negative_axis(get_attribute(node, "axis", 0), len(data_shape)) indices_shape = self._get_shape(node, 1) vi = self.known_vi_[node.output[0]] + if node.op_type == "Gather": + elem_type = self.known_vi_[node.input[0]].type.tensor_type.elem_type + elif node.op_type == "GatherBlockQuantized": + # scales + elem_type = self.known_vi_[node.input[2]].type.tensor_type.elem_type + else: + raise ValueError(f"Unsupported Gather op_type: {node.op_type}") vi.CopyFrom( helper.make_tensor_value_info( node.output[0], - self.known_vi_[node.input[0]].type.tensor_type.elem_type, + elem_type, data_shape[:axis] + indices_shape + data_shape[axis + 1 :], ) )