Upload worker_image_quant.py with huggingface_hub

worker_image_quant.py

@@ -350,14 +350,25 @@ class _PleQuantTable:
         )
         n_shards = meta["shards"]
         assert n_shards * self.ROWS_PER_SHARD == total_rows, "non-uniform shards"
-        self._q, self._s = [], []
+        # Order matters: the e2m1 layout string contains "e4m3" (its scale dtype).
+        if "e2m1" in self.layout:
+            key = "weight_e2m1"
+        elif "e4m3" in self.layout:
+            key = "weight_fp8"
+        else:
+            key = "weight_i4"
+        self._q, self._s, self._s2 = [], [], []
         for n in range(n_shards):
             f = safe_open(os.path.join(quant_dir, f"shard_{n}.safetensors"),
                           framework="pt")
-            key = "weight_fp8" if "e4m3" in self.layout else "weight_i4"
             self._q.append(f.get_tensor(key))
             self._s.append(f.get_tensor("weight_scale"))
+            self._s2.append(
+                f.get_tensor("weight_scale_2").item()
+                if "weight_scale_2" in f.keys() else 1.0
+            )
         self.width = width
+        self._lut = None
         logger.info("PLE quant table: %s, %d shards mmapped from %s",
                     self.layout, n_shards, quant_dir)
 
@@ -376,6 +387,20 @@ class _PleQuantTable:
             pos += c
 
     def _dequant(self, s: int, sel: torch.Tensor) -> torch.Tensor:
+        if "e2m1" in self.layout:
+            if self._lut is None:
+                mags = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+                self._lut = torch.tensor(mags + [-m for m in mags],
+                                         dtype=torch.float32)
+            packed = self._q[s].index_select(0, sel)
+            lo = (packed & 0xF).long()
+            hi = (packed >> 4).long()
+            nib = torch.stack((lo, hi), dim=-1).view(packed.shape[0], self.width)
+            scale = self._s[s].index_select(0, sel).to(torch.float32)
+            g = self.width // scale.shape[1]
+            return (self._lut[nib]
+                    * scale.repeat_interleave(g, dim=1)
+                    * self._s2[s])
         if "e4m3" in self.layout:
             q = self._q[s].index_select(0, sel).to(torch.float32)
             return q * self._s[s].index_select(0, sel)[:, None]