fix(qwen4_exp): read weights under the checkpoint's quant config

python/freetoken/layers/quantization/configs/base.py

@@ -40,10 +40,15 @@ def quantization_config_of(hf_config: Any) -> dict[str, Any] | None:
 
 @dataclass(frozen=True)
 class Stored:
-    """One checkpoint tensor behind a role: its suffix, and whether it holds the quant-side scale whose reciprocal the layer wants."""
+    """One checkpoint tensor behind a role: its suffix, other suffixes exports of the same dialect use for it, and whether it holds the quant-side scale whose reciprocal the layer wants."""
 
     name: str
     reciprocal: bool = False
+    aliases: tuple[str, ...] = ()
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return (self.name, *self.aliases)
 
 
 class QuantConfig(ABC):

python/freetoken/layers/quantization/configs/modelopt.py

@@ -27,7 +27,8 @@ class ModelOptConfig(QuantConfig):
     STORAGE: ClassVar[dict[QuantKind, dict[str, str | Stored]]] = {
         QuantKind.FP8_TENSOR: {"weight": "weight", "weight_scale": "weight_scale", "input_scale": "input_scale"},
         QuantKind.FP8_BLOCK: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"},
-        QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": "weight_scale_inv"},
+        # same e8m0 codes under either name: MiniMax-M3 exports weight_scale_inv, Qwen3.8-Flash-Next weight_scale
+        QuantKind.MXFP8: {"weight": "weight", "weight_scale_inv": Stored("weight_scale_inv", aliases=("weight_scale",))},
         QuantKind.NVFP4: {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2", "input_scale": "input_scale"},
     }
 

python/freetoken/models/qwen3_5_moe/weight.py

@@ -149,7 +149,7 @@ def add(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]]
         if module.rpartition(".")[2] not in _LINEAR_LEAVES:
             return None
         stored = self.stored(module)
-        roles = {"weight": "weight"} if stored is None else {e.name: r for r, e in self.quant.storage(stored).items()}
+        roles = {"weight": "weight"} if stored is None else {n: r for r, e in self.quant.storage(stored).items() for n in e.names}
         role = roles.get(suffix)
         if role is None:
             if suffix in _DROPPED_SUFFIXES:

python/freetoken/models/qwen4_exp/ple_disk.py

@@ -1,4 +1,4 @@
-"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant.
+"""Disk-backed PLE table (--ple-backend disk): the C++ store hashes n-gram windows and batch-reads rows from the checkpoint's fp8 or bf16 shard tensors into pinned staging; the captured ``lookup`` is a fixed-shape H2D copy + dequant.
 
 Hash windows are pure functions of ``req.input_ids`` + ``device_len`` (prefix hits, restores and COW forks need no bookkeeping); the decode input token lives device-side under overlap scheduling and is read back here.
 """
@@ -22,7 +22,7 @@
 from .weight import (
     _PLE_SCALE_SUFFIX,
     _PLE_SHARD_RE,
-    _PLE_ST_DTYPE,
+    _PLE_ST_DTYPES,
     _ple_table_files,
     _safetensors_header,
 )
@@ -50,9 +50,10 @@ class PleRowSource:
     row_bytes: int
     row_stride: int
     scale: float
-    # packed ggml rows (a GGUF table) of ``row_values`` values each; None for one fp8 byte per value
+    # packed ggml rows (a GGUF table) of ``row_values`` values each; None for float rows of ``row_dtype``
     ggml_type: int | None = None
     row_values: int | None = None
+    row_dtype: torch.dtype = torch.float8_e4m3fn
 
     @property
     def total_rows(self) -> int:
@@ -62,6 +63,7 @@ def total_rows(self) -> int:
 def source_from_safetensors(folder: str) -> PleRowSource:
     """Map the checkpoint's ``ngram_embedding.shard_<i>`` tensors in place: one extent per shard, no copy."""
     rows = cols = 0
+    dtype: torch.dtype | None = None
     scale: torch.Tensor | None = None
     paths: list[str] = []
     path_idx: dict[str, int] = {}
@@ -78,8 +80,10 @@ def source_from_safetensors(folder: str) -> PleRowSource:
             match = _PLE_SHARD_RE.search(key)
             if match is None:
                 continue
-            if meta["dtype"] != _PLE_ST_DTYPE:
-                raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected {_PLE_ST_DTYPE}")
+            shard_dtype = _PLE_ST_DTYPES.get(meta["dtype"])
+            if shard_dtype is None or dtype not in (None, shard_dtype):
+                raise ValueError(f"PLE shard {key} has dtype {meta['dtype']}, expected one of {sorted(_PLE_ST_DTYPES)} for every shard")
+            dtype = shard_dtype
             if rows and tuple(meta["shape"]) != (rows, cols):
                 raise ValueError(f"PLE shard {key} is {meta['shape']}, expected {[rows, cols]}")
             rows, cols = meta["shape"]
@@ -93,9 +97,13 @@ def source_from_safetensors(folder: str) -> PleRowSource:
     if sorted(shards) != list(range(len(shards))) or not shards:
         raise ValueError(f"PLE shard indices are not contiguous 0..N-1: {sorted(shards)[:8]}")
     if scale is None:
-        raise ValueError("PLE table has no weight_scale")
+        if dtype is torch.float8_e4m3fn:
+            raise ValueError("PLE table has no weight_scale")
+        scale = torch.tensor(1.0)
     order = [shards[i] for i in range(len(shards))]
-    return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, cols, cols, float(scale))
+    row_bytes = cols * dtype.itemsize
+    return PleRowSource(paths, [f for f, _ in order], [b for _, b in order], rows, row_bytes, row_bytes, float(scale),
+                        row_values=cols, row_dtype=dtype)
 
 
 def source_from_gguf(model_path: str) -> PleRowSource:
@@ -137,11 +145,11 @@ def __init__(
         from freetoken.layers.quantization.ggml_types import float_dtype
 
         self.num_rows = source.total_rows
-        self.head_dim = source.row_values or source.row_bytes  # fp8: one byte per element
+        self.head_dim = source.row_values or source.row_bytes
         self._ggml_type = source.ggml_type
         self._row_bytes = source.row_bytes
         # float rows (fp8, or F32 / F16 / BF16 in a GGUF) are only cast; quantized ggml rows go through ggml
-        self._row_dtype = torch.float8_e4m3fn if source.ggml_type is None else float_dtype(source.ggml_type)
+        self._row_dtype = source.row_dtype if source.ggml_type is None else float_dtype(source.ggml_type)
         self.dtype = dtype
         self.heads = int(hash_constants["num_ngram_heads"])
         self.scale = source.scale