create architecture neutral forward for channel/group-wise embedding quantization operator

Author: Michael Gschwind

.github/workflows/compile.yml

@@ -65,12 +65,12 @@ jobs:
           # echo "******************************************"
           # echo "******** Emb: group-wise quantized *******"
           # echo "******************************************"
-          # python generate.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --temperature 0 > ./output_eager
-          # cat ./output_eager
-          # python generate.py --compile --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --temperature 0 > ./output_compiled
-          # cat ./output_compiled
-          # python export.py --quant "embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --output-dso-path ${MODEL_DIR}/${MODEL_NAME}.so
-          # python generate.py --checkpoint-path ${MODEL_PATH} --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so  > ./output_aoti
+          python generate.py --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --temperature 0 > ./output_eager
+          cat ./output_eager
+          python generate.py --compile --quant '{"embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --temperature 0 > ./output_compiled
+          cat ./output_compiled
+          python export.py --quant "embedding" : {"bitwidth": 8, "group_size": 8}}' --checkpoint-path ${MODEL_PATH} --output-dso-path ${MODEL_DIR}/${MODEL_NAME}.so
+          python generate.py --checkpoint-path ${MODEL_PATH} --temperature 0 --dso-path ${MODEL_DIR}/${MODEL_NAME}.so  > ./output_aoti
           cat ./output_aoti
 
           # echo "******************************************"

quantize.py

@@ -461,11 +461,25 @@ def forward(self, indices: torch.Tensor) -> torch.Tensor:
                 self.weight, self.scales, None, 0, 0, indices, dtype=self.dtype
             )
 
-        result_weights = self.weight.index_select(0, indices.view(-1))
-        result_scales = self.scales.index_select(0, indices.view(-1))
+        
+        # result_weights = self.weight.index_select(0, indices.view(-1))
+        # result_scales = self.scales.index_select(0, indices.view(-1))
+
+        weight = self.weight
+        scales = self.scales.view(weight.shape[0], -1)
+        
+        result_weights = F.embedding(indices, weight)
+        result_scales = F.embedding(indices, scales)
+
+        rw_view = result_weights.to(dtype=result_scales.dtype).view(tuple(result_weights.shape[:-1] + (scales.shape[1], -1, )))
+        rs_view = result_scales.view(tuple(result_scales.shape[:-1]) + (scales.shape[1], 1, ))
+        # print(f"rw_view {rw_view.shape}")
+        # print(f"rs_view {rs_view.shape}")
 
-        r = result_weights.to(dtype=result_scales.dtype) * result_scales
+        r = rw_view * rs_view
         return r.view(indices.size() + (-1,))
+        
+        # r = result_weights.to(dtype=result_scales.dtype).view(list(result_weights.shape[:-1] + (scales.shape[1], -1, )) * result_scales.view(scales.shape[-1] + (scales.shape[1], 1, ))
 
 ##################################################################
 ##### weight only int4 per channel groupwise quantized code ######