patch.diff

diff --git a/src/transformers/integrations/metal_quantization.py b/src/transformers/integrations/metal_quantization.py
index 6ade40af53f5..57fd4444a005 100644
--- a/src/transformers/integrations/metal_quantization.py
+++ b/src/transformers/integrations/metal_quantization.py
@@ -47,21 +47,193 @@
 
 _metal_kernel = None
 
+# ---------------------------------------------------------------------------
+# Locally-compiled Metal fallback for affine_qmm_t
+# ---------------------------------------------------------------------------
+
+_AFFINE_QMM_T_METAL_SOURCE = """
+#include <metal_stdlib>
+using namespace metal;
+
+// Fused dequantize + matmul:  y = x @ dequant(w).T
+// x: [M, K], w: [N, K_packed] (uint32), scales/biases: [N, n_groups], out: [M, N]
+// Each thread computes one (m, n) output element.
+
+kernel void affine_qmm_t_float(
+    device const float* x        [[buffer(0)]],
+    device const uint*  w        [[buffer(1)]],
+    device const float* scales   [[buffer(2)]],
+    device const float* biases   [[buffer(3)]],
+    device float*       out      [[buffer(4)]],
+    constant uint& M             [[buffer(5)]],
+    constant uint& N             [[buffer(6)]],
+    constant uint& K             [[buffer(7)]],
+    constant uint& group_size    [[buffer(8)]],
+    constant uint& bits          [[buffer(9)]],
+    uint2 tid                    [[thread_position_in_grid]])
+{
+    uint m = tid.y;
+    uint n = tid.x;
+    if (m >= M || n >= N) return;
+
+    uint elems_per_int = 32 / bits;
+    uint mask = (1u << bits) - 1u;
+    uint K_packed = K / elems_per_int;
+    uint n_groups = K / group_size;
+
+    float acc = 0.0f;
+    for (uint k = 0; k < K; k++) {
+        uint packed_val = w[n * K_packed + k / elems_per_int];
+        float q = float((packed_val >> ((k % elems_per_int) * bits)) & mask);
+        uint g = k / group_size;
+        acc += x[m * K + k] * (q * scales[n * n_groups + g] + biases[n * n_groups + g]);
+    }
+    out[m * N + n] = acc;
+}
+
+kernel void affine_qmm_t_half(
+    device const half*  x        [[buffer(0)]],
+    device const uint*  w        [[buffer(1)]],
+    device const half*  scales   [[buffer(2)]],
+    device const half*  biases   [[buffer(3)]],
+    device half*        out      [[buffer(4)]],
+    constant uint& M             [[buffer(5)]],
+    constant uint& N             [[buffer(6)]],
+    constant uint& K             [[buffer(7)]],
+    constant uint& group_size    [[buffer(8)]],
+    constant uint& bits          [[buffer(9)]],
+    uint2 tid                    [[thread_position_in_grid]])
+{
+    uint m = tid.y;
+    uint n = tid.x;
+    if (m >= M || n >= N) return;
+
+    uint elems_per_int = 32 / bits;
+    uint mask = (1u << bits) - 1u;
+    uint K_packed = K / elems_per_int;
+    uint n_groups = K / group_size;
+
+    float acc = 0.0f;
+    for (uint k = 0; k < K; k++) {
+        uint packed_val = w[n * K_packed + k / elems_per_int];
+        float q = float((packed_val >> ((k % elems_per_int) * bits)) & mask);
+        uint g = k / group_size;
+        acc += float(x[m * K + k]) * (q * float(scales[n * n_groups + g]) + float(biases[n * n_groups + g]));
+    }
+    out[m * N + n] = half(acc);
+}
+
+kernel void affine_qmm_t_bfloat(
+    device const bfloat* x       [[buffer(0)]],
+    device const uint*   w       [[buffer(1)]],
+    device const bfloat* scales  [[buffer(2)]],
+    device const bfloat* biases  [[buffer(3)]],
+    device bfloat*       out     [[buffer(4)]],
+    constant uint& M             [[buffer(5)]],
+    constant uint& N             [[buffer(6)]],
+    constant uint& K             [[buffer(7)]],
+    constant uint& group_size    [[buffer(8)]],
+    constant uint& bits          [[buffer(9)]],
+    uint2 tid                    [[thread_position_in_grid]])
+{
+    uint m = tid.y;
+    uint n = tid.x;
+    if (m >= M || n >= N) return;
+
+    uint elems_per_int = 32 / bits;
+    uint mask = (1u << bits) - 1u;
+    uint K_packed = K / elems_per_int;
+    uint n_groups = K / group_size;
+
+    float acc = 0.0f;
+    for (uint k = 0; k < K; k++) {
+        uint packed_val = w[n * K_packed + k / elems_per_int];
+        float q = float((packed_val >> ((k % elems_per_int) * bits)) & mask);
+        uint g = k / group_size;
+        acc += float(x[m * K + k]) * (q * float(scales[n * n_groups + g]) + float(biases[n * n_groups + g]));
+    }
+    out[m * N + n] = bfloat(acc);
+}
+"""
+
+_compiled_shader_lib = None
+
+
+class _LocalMetalKernel:
+    """Wrapper that mimics the Hub kernel interface using ``torch.mps.compile_shader``."""
+
+    def __init__(self):
+        global _compiled_shader_lib
+        if _compiled_shader_lib is None:
+            _compiled_shader_lib = torch.mps.compile_shader(_AFFINE_QMM_T_METAL_SOURCE)
+        self._lib = _compiled_shader_lib
+
+    def affine_qmm_t(self, x, w, scales, biases, group_size, bits):
+        K_packed = w.shape[1]
+        N = w.shape[0]
+        elems_per_int = 32 // bits
+        K = K_packed * elems_per_int
+
+        x_2d = x.reshape(-1, K).contiguous()
+        M_total = x_2d.shape[0]
+        out = torch.empty(M_total, N, dtype=x.dtype, device=x.device)
+
+        M_t = torch.tensor(M_total, dtype=torch.uint32, device="mps")
+        N_t = torch.tensor(N, dtype=torch.uint32, device="mps")
+        K_t = torch.tensor(K, dtype=torch.uint32, device="mps")
+        gs_t = torch.tensor(group_size, dtype=torch.uint32, device="mps")
+        bits_t = torch.tensor(bits, dtype=torch.uint32, device="mps")
+
+        if x.dtype == torch.float32:
+            fn = self._lib.affine_qmm_t_float
+        elif x.dtype == torch.float16:
+            fn = self._lib.affine_qmm_t_half
+        elif x.dtype == torch.bfloat16:
+            fn = self._lib.affine_qmm_t_bfloat
+        else:
+            raise ValueError(f"Unsupported dtype {x.dtype} for Metal affine_qmm_t")
+
+        fn(x_2d, w, scales, biases, out, M_t, N_t, K_t, gs_t, bits_t, threads=[N, M_total, 1])
+
+        return out.reshape(*x.shape[:-1], N)
+
 
 def _get_metal_kernel():
-    """Lazily load the quantization-mlx kernel from Hugging Face Hub."""
+    """Lazily load the quantization-mlx kernel from Hugging Face Hub, falling back to a
+    locally-compiled Metal shader if the Hub kernel is unavailable or incompatible."""
     global _metal_kernel
     if _metal_kernel is None:
         try:
+            import os
+
             from .hub_kernels import get_kernel
 
-            _metal_kernel = get_kernel("kernels-community/mlx-quantization-metal-kernels")
-        except Exception as e:
-            raise ImportError(
-                f"Failed to load the quantization-mlx kernel from the Hub: {e}. "
-                "Make sure you have `kernels` installed (`pip install kernels`) "
-                "and are running on an Apple Silicon machine."
-            ) from e
+            hub_kernel = get_kernel("kernels-community/mlx-quantization-metal-kernels")
+            # Smoke-test: the pre-built metallib may target an MSL version newer
+            # than the current OS supports.  A tiny matmul catches this at init
+            # time rather than mid-inference.
+            # Suppress Metal runtime stderr noise ("Failed to create Metal library
+            # from embedded header") by temporarily redirecting fd 2 to /dev/null.
+            _x = torch.zeros(1, 64, dtype=torch.float32, device="mps")
+            _w = torch.zeros(1, 2, dtype=torch.uint32, device="mps")  # K=64 at 8-bit → 2 packed
+            _s = torch.ones(1, 1, dtype=torch.float32, device="mps")
+            _b = torch.zeros(1, 1, dtype=torch.float32, device="mps")
+            stderr_fd = os.dup(2)
+            devnull = os.open(os.devnull, os.O_WRONLY)
+            try:
+                os.dup2(devnull, 2)
+                hub_kernel.affine_qmm_t(_x, _w, _s, _b, 64, 8)
+            finally:
+                os.dup2(stderr_fd, 2)
+                os.close(stderr_fd)
+                os.close(devnull)
+            _metal_kernel = hub_kernel
+        except Exception:
+            logger.info(
+                "Hub kernel 'kernels-community/mlx-quantization-metal-kernels' unavailable; "
+                "using locally-compiled Metal shader fallback."
+            )
+            _metal_kernel = _LocalMetalKernel()
     return _metal_kernel
 
 
@@ -112,6 +284,14 @@ def __init__(
         else:
             self.register_parameter("bias", None)
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # MLX-quantized models store quantization biases as "biases" instead of "qbiases"
+        biases_key = prefix + "biases"
+        qbiases_key = prefix + "qbiases"
+        if biases_key in state_dict and qbiases_key not in state_dict:
+            state_dict[qbiases_key] = state_dict.pop(biases_key)
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         if self.weight.dtype != torch.uint32:
             return nn.functional.linear(input, self.weight, self.bias)
@@ -285,12 +465,16 @@ def convert(self, input_dict: dict, full_layer_name: str | None = None, **kwargs
         bits = self.hf_quantizer.quantization_config.bits
         group_size = self.hf_quantizer.quantization_config.group_size
 
+        # Use source_patterns from kwargs as dict keys (they are the keys in input_dict).
+        # Fall back to the default patterns for backward compatibility.
+        source_patterns = kwargs.get("source_patterns", ["weight$", "scales", "qbiases"])
+
         if len(input_dict) < 2:
-            return {full_layer_name: input_dict["weight$"]}
+            return {full_layer_name: input_dict[source_patterns[0]]}
 
-        quantized = input_dict["weight$"][0]
-        scales = input_dict["scales"][0]
-        qbiases = input_dict["qbiases"][0]
+        quantized = input_dict[source_patterns[0]][0]
+        scales = input_dict[source_patterns[1]][0]
+        qbiases = input_dict[source_patterns[2]][0]
 
         w_deq = _affine_dequantize_tensor(quantized, scales, qbiases, group_size, bits)
         return {full_layer_name: w_deq.to(scales.dtype)}
diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py
index 22a4dcb9da67..6be81c68f02f 100644
--- a/src/transformers/quantizers/auto.py
+++ b/src/transformers/quantizers/auto.py
@@ -137,9 +137,18 @@ def from_dict(cls, quantization_config_dict: dict):
             suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"
             quant_method = QuantizationMethod.BITS_AND_BYTES + suffix
         elif quant_method is None:
-            raise ValueError(
-                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"
-            )
+            # Recognize MLX-style affine quantization configs as Metal
+            if (
+                quantization_config_dict.get("mode") == "affine"
+                and "bits" in quantization_config_dict
+                and "group_size" in quantization_config_dict
+            ):
+                quantization_config_dict["quant_method"] = QuantizationMethod.METAL
+                quant_method = QuantizationMethod.METAL
+            else:
+                raise ValueError(
+                    "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"
+                )
 
         if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING:
             raise ValueError(
diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py
index 0c7515e5c6b1..2bfc4325d1e1 100644
--- a/src/transformers/quantizers/quantizer_metal.py
+++ b/src/transformers/quantizers/quantizer_metal.py
@@ -90,9 +90,16 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **
     def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
         from ..integrations.metal_quantization import replace_with_metal_linear
 
-        self.modules_to_not_convert = self.get_modules_to_not_convert(
-            model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules
-        )
+        self._model_type = getattr(model.config, "model_type", None)
+
+        skip_modules = self.quantization_config.modules_to_not_convert
+        if self.pre_quantized and skip_modules is None:
+            # Pre-quantized checkpoints (e.g. MLX) may have quantized the lm_head /
+            # output embedding too.  Don't auto-skip them; only honour explicit user
+            # overrides via modules_to_not_convert.
+            skip_modules = []
+
+        self.modules_to_not_convert = self.get_modules_to_not_convert(model, skip_modules, model._keep_in_fp32_modules)
 
         model = replace_with_metal_linear(
             model,
@@ -114,7 +121,7 @@ def get_quantize_ops(self):
         return MetalQuantize(self)
 
     def get_weight_conversions(self):
-        from ..core_model_loading import WeightConverter
+        from ..core_model_loading import WeightConverter, WeightRenaming
         from ..integrations.metal_quantization import MetalDequantize
 
         if self.pre_quantized and self.quantization_config.dequantize:
@@ -125,4 +132,34 @@ def get_weight_conversions(self):
                     operations=[MetalDequantize(self)],
                 )
             ]
+
+        if self.pre_quantized:
+            conversions = [
+                # MLX uses "biases", MetalLinear expects "qbiases"
+                WeightRenaming(source_patterns="biases", target_patterns="qbiases"),
+                # MLX quantizes embed_tokens but transformers keeps it as nn.Embedding (float);
+                # dequantize the embedding back to float so the standard Embedding layer can load it
+                WeightConverter(
+                    source_patterns=[r"embed_tokens\.weight$", r"embed_tokens\.scales", r"embed_tokens\.qbiases"],
+                    target_patterns="embed_tokens.weight",
+                    operations=[MetalDequantize(self)],
+                ),
+            ]
+
+            # MLX checkpoints may use different key prefixes than the model expects.
+            # These renamings are model-specific and only needed for pre-quantized MLX loads.
+            model_type = getattr(self, "_model_type", None)
+            if model_type == "qwen3_vl":
+                conversions.extend(
+                    [
+                        WeightRenaming(
+                            source_patterns="language_model.model.", target_patterns="model.language_model."
+                        ),
+                        WeightRenaming(source_patterns="language_model.lm_head.", target_patterns="lm_head."),
+                        WeightRenaming(source_patterns="vision_tower.", target_patterns="model.visual."),
+                    ]
+                )
+
+            return conversions
+
         return []
diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
index b3d5c19f2984..4493ffb264e9 100644
--- a/src/transformers/utils/hub.py
+++ b/src/transformers/utils/hub.py
@@ -829,6 +829,68 @@ def convert_file_size_to_int(size: int | str):
     raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")
 
 
+def _rebuild_shard_index_from_repo(
+    pretrained_model_name_or_path,
+    cache_dir=None,
+    force_download=False,
+    proxies=None,
+    local_files_only=False,
+    token=None,
+    user_agent=None,
+    revision=None,
+    subfolder="",
+    _commit_hash=None,
+):
+    """
+    When the shard index references files that don't exist (e.g. MLX repos that
+    copied the index from the original model), discover the actual safetensors
+    files on the Hub, download them, and rebuild the weight_map from their headers.
+    """
+    import struct
+
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+    all_files = api.list_repo_files(pretrained_model_name_or_path, revision=revision, token=token)
+    shard_names = sorted(f for f in all_files if f.endswith(".safetensors") and f != "model.safetensors.index.json")
+
+    if not shard_names:
+        raise OSError(
+            f"No .safetensors files found in repo '{pretrained_model_name_or_path}'. Cannot rebuild shard index."
+        )
+
+    # Download the actual shard files
+    cached_filenames = cached_files(
+        pretrained_model_name_or_path,
+        shard_names,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        proxies=proxies,
+        local_files_only=local_files_only,
+        token=token,
+        user_agent=user_agent,
+        revision=revision,
+        subfolder=subfolder,
+        _commit_hash=_commit_hash,
+    )
+
+    # Rebuild weight_map by reading safetensors headers
+    weight_map = {}
+    all_keys = []
+    for cached_path, shard_name in zip(cached_filenames, shard_names):
+        with open(cached_path, "rb") as f:
+            header_size = struct.unpack("<Q", f.read(8))[0]
+            header = json.loads(f.read(header_size))
+        for key in header:
+            if key == "__metadata__":
+                continue
+            weight_map[key] = shard_name
+            all_keys.append(key)
+
+    sharded_metadata = {"all_checkpoint_keys": all_keys, "weight_map": weight_map}
+    return cached_filenames, sharded_metadata
+
+
 def get_checkpoint_shard_files(
     pretrained_model_name_or_path,
     index_filename,
@@ -871,19 +933,36 @@ def get_checkpoint_shard_files(
 
     # At this stage pretrained_model_name_or_path is a model identifier on the Hub. Try to get everything from cache,
     # or download the files
-    cached_filenames = cached_files(
-        pretrained_model_name_or_path,
-        shard_filenames,
-        cache_dir=cache_dir,
-        force_download=force_download,
-        proxies=proxies,
-        local_files_only=local_files_only,
-        token=token,
-        user_agent=user_agent,
-        revision=revision,
-        subfolder=subfolder,
-        _commit_hash=_commit_hash,
-    )
+    try:
+        cached_filenames = cached_files(
+            pretrained_model_name_or_path,
+            shard_filenames,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            user_agent=user_agent,
+            revision=revision,
+            subfolder=subfolder,
+            _commit_hash=_commit_hash,
+        )
+    except OSError:
+        # The shard index may reference files that don't exist in this repo
+        # (e.g. MLX repos that copy the index from the original model).
+        # Fall back to discovering the actual safetensors files on the Hub.
+        cached_filenames, sharded_metadata = _rebuild_shard_index_from_repo(
+            pretrained_model_name_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            proxies=proxies,
+            local_files_only=local_files_only,
+            token=token,
+            user_agent=user_agent,
+            revision=revision,
+            subfolder=subfolder,
+            _commit_hash=_commit_hash,
+        )
 
     return cached_filenames, sharded_metadata