Repository navigation
GgufLinear: inference-time GGUF matmul on Apple Silicon — llama.cpp parity - #45977
Closed
ArthurZucker wants to merge 32 commits into
Closed
ArthurZucker wants to merge 32 commits into
ArthurZucker wants to merge 32 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
ArthurZucker
force-pushed
the
gguf-matmul-kernels
branch
from
May 14, 2026 11:30
56d3847 to
cb6ba16
Compare
ArthurZucker
added a commit
that referenced
this pull request
May 17, 2026
``run_one_warmup`` wraps the warmup forward in ``torch.cuda.stream(self.inputs_and_outputs.compute_stream)``. On MPS and CPU ``compute_stream`` is ``None`` (set explicitly in ``input_outputs.py:141``), and ``torch.cuda.stream(None)`` falls through a code path that calls ``torch.mps.current_device()`` — which doesn't exist — and raises ``AttributeError``. The except clause turns this into a ``Failed to warm up: module 'torch.mps' has no attribute 'current_device'`` warning, and warmup silently fails (graph pool fragments under load; CB+compile never gets a primed graph). Mirror the existing ``nullcontext`` fallback used elsewhere (``input_outputs.py:248``, ``model_runner.py:108``) so the no-stream path is a clean no-op on MPS/CPU. Unlocks ``generate_batch`` + ``use_default_compile_configs=True`` on Apple Silicon (CB N=1 + block-table + compile cells in PR #45977's final speed sweep).
Opt-in path for running inference directly on GGUF block bytes via the
Metal kernels from `ArthurZ/gguf-kernels` (the kernels llama.cpp ships).
Weights stay at native quant — same memory footprint as llama.cpp, no
dequant on the forward path.
Supported quant types: Q4_0 / Q5_0 / Q5_1 / Q8_0 / Q4_K / Q5_K / Q6_K /
IQ4_NL / IQ4_XS for dense Linear; Q4_K / Q5_K / Q6_K / Q8_0 for
fused-expert MoE via a quantized `Qwen2MoeExperts` swap.
API:
```python
# Native GGUF (default when a .gguf is loaded)
AutoModelForCausalLM.from_pretrained(repo, gguf_file="...gguf")
# Dequantize-on-load (explicit dtype override)
AutoModelForCausalLM.from_pretrained(repo, gguf_file="...gguf", dtype=torch.bfloat16)
# On-the-fly quantization of an fp16/bf16 model via gguf-py
AutoModelForCausalLM.from_pretrained(
repo,
quantization_config=GgufQuantizeConfig(
quant_type="Q4_0",
modules_to_convert=["model.layers.*.mlp.*"],
),
)
```
Auto-detects a single `.gguf` in the target repo so `gguf_file=` can
often be omitted. `GgufQuantizeConfig` is registered through the standard
`AUTO_QUANTIZER_MAPPING`, mirroring the bnb / torchao / quanto pattern.
ArthurZucker
force-pushed
the
gguf-matmul-kernels
branch
from
May 18, 2026 05:00
5134799 to
69d0f97
Compare
Major shrink + restructure of the GGUF integration so it mirrors the
shape of every other HF quantizer (FP8 / bnb / torchao / ...). No more
bespoke setup_for_compile / fast_greedy_decode / save_pretrained_gguf /
register_gguf_moe / _process_model_after_weight_loading machinery.
Files reshaped:
* `integrations/gguf_linear.py` — 846 → 365 LOC. Owns `GgufLinear` +
`GgufExperts` (renamed from `GgufQwen2MoeExperts`) + the
`MODEL_TYPE_TO_GGUF_EXPERTS` registry + the new `replace_with_gguf_linear`
meta-time swap. Drops: `TRANSFORMERS_GGUF_USE_METAL_KERNELS` env var,
`_dequant_forward`, `_gguf_eager_experts_forward`, `_gather_dequant`,
`_row_permute_attn_k_bytes`, `_prepare_id_kernel_refs`,
`register_gguf_moe`, `replace_qwen2_moe_experts`, `fast_greedy_decode`,
`setup_for_compile`. The MoE bmm forward moved in from `moe.py`.
* `integrations/gguf_kernels.py` (new, 142 LOC) — owns the metal-kernel
loader + bind helper. No env-var feature flag. Kernels-unavailable on
MPS → RuntimeError; there is no slow fallback.
* `integrations/gguf_writer.py` (new, 145 LOC) — replaces the 369-line
`gguf_save.py`. Walks the quantizer's reverse rename map and uses
the `GGUFQuantize` op (the new `GGUFDequantize.reverse_op`) for
re-quantization on save.
* `integrations/gguf_save.py` — deleted.
* `gguf_conversion_ops.py` — adds `GGUFQuantize` mirroring
`Fp8Quantize`, wires `GGUFDequantize.reverse_op` to return it.
* `quantizers/quantizer_gguf.py` — drops `_process_model_after_weight_loading`,
`_swap_moe_experts`, `_quantize_on_the_fly`. Swap now happens at
`_process_model_before_weight_loading` (FP8 pattern).
`postprocess_model` applies generation defaults (`cache_implementation="static"`
+ `compile_config=CompileConfig(mode="reduce-overhead", dynamic=False)`)
with an INFO log when the GgufLinear swap is in effect.
* `integrations/moe.py` — `gguf_bmm_experts_forward` is now imported
from `gguf_linear`; this file only registers it in `ExpertsInterface`.
* `modeling_utils.py` — gguf_file branch now device-aware. On MPS with
`use_kernels=True` (default for the GGUF path) → swap to GgufLinear.
On CUDA / CPU → force dequant (`linear_mode=False`) until matching
GGUF kernels exist for those backends (TODO). Explicit `dtype=` always
forces dequant.
* `tests/quantization/ggml/test_gguf_linear.py` — rewritten to match the
new behavioural contract; structural tests run on CI hosts, forward
tests gated behind MPS + kernels.
Behavioural contract:
from_pretrained(repo, gguf_file=...) on MPS → GgufLinear, kernels
from_pretrained(repo, gguf_file=..., dtype=) any dev → dequant → nn.Linear
from_pretrained(repo, gguf_file=...) non-MPS → dequant → nn.Linear
from_pretrained(repo, quantization_config=GgufQuantizeConfig(...))
→ on-the-fly via GGUFQuantize
Use the same kwarg pattern other conversion ops already rely on
(``ReversePermuteAttn`` reads ``config.num_attention_heads`` from the
``kwargs`` threaded through ``WeightConverter.convert`` / ``WeightRenaming.convert``):
* ``GGUFDequantize.convert`` now reads ``model`` + ``full_layer_name``
out of kwargs, looks up the live target module, and:
- swapped :class:`GgufLinear` → return the GGUFQuantizedTensor's raw
uint8 bytes (flattened) so they land directly in ``GgufLinear.weight``;
- anything else → dequantize as before.
* As a result the ``GGUFQuantizer.postprocess_model`` hook no longer
copies bytes for dense Linears — the rename pipeline does it on its
own. The fix-up step shrinks to ``_copy_experts_bytes`` (MoE only),
which is still needed because the GGUF rename rules ship a merge
``WeightConverter`` (``ffn_gate_exps + ffn_up_exps → gate_up_proj``)
whose target doesn't exist on the swapped ``GgufExperts``.
Also decouple ``use_kernels`` from the GGUF kernel decision: the
model-wide kernel resolver (attention / RMSNorm via the ``kernels``
package) is unrelated to the GGUF metal kernels we load in
``integrations.gguf_kernels``. ``linear_mode`` now keys solely off
``target_device == "mps"`` + no-explicit-dtype.
Make :class:`GGUFDequantize` recognize swapped :class:`GgufExperts` modules
as pass-through targets (same target-aware pattern other conversion ops
already use via the ``model`` / ``full_layer_name`` kwargs threaded
through ``ConversionOp.convert``). Source ``GGUFQuantizedTensor`` bytes
get reshaped to the live target buffer's shape, so the same op handles
both the dense ``GgufLinear.weight`` 1D buffer and the
``GgufExperts.{gate,up,down}_proj_q`` 2D buffers.
The merge ``WeightConverter`` (``ffn_gate_exps + ffn_up_exps →
mlp.experts.gate_up_proj``) shipped by the GGUF rename rules doesn't
apply when the swapped target is ``GgufExperts``; rewrite it in
:meth:`GGUFQuantizer.update_weight_conversions` into two plain
``WeightRenaming`` entries (``…ffn_gate_exps → mlp.experts.gate_proj_q``
and ``…ffn_up_exps → mlp.experts.up_proj_q``). The existing
``ffn_down_exps → mlp.experts.down_proj`` rename gets its target
rewritten to ``mlp.experts.down_proj_q``.
Net effect: the GGUF byte path for the entire pre-quantized linear_mode
load (dense + MoE) is now one machinery — the rename pipeline. The
post-load ``_copy_experts_bytes`` fix-up is gone. ``postprocess_model``
now only resolves kernel refs for the swapped expert modules and sets
generation defaults — no byte plumbing.
Both fuse their per-expert weights via the standard GGUF merge converter (`ffn_gate_exps + ffn_up_exps → gate_up_proj`) onto a 3D Parameter — `MixtralExperts` does it directly, `DeepseekV3NaiveMoe` via subclass. The byte layout is identical to `Qwen2MoeExperts`, so a single entry per arch in the registry suffices.
…t works Restructure ``GgufExperts`` to expose the same buffer names as ``MixtralExperts`` / ``Qwen2MoeExperts`` / ``DeepseekV3NaiveMoe``: ``gate_up_proj`` (single uint8 buffer, shape ``(num_experts, 2*intermediate, bytes_per_row)``) + ``down_proj`` (single uint8 buffer). That's the natural output of the GGUF rename pipeline's merge ``WeightConverter`` (``ffn_gate_exps + ffn_up_exps → gate_up_proj``) running ``Concatenate(dim=1)`` on uint8 byte tensors — exactly the FP8 pattern: don't rewrite the converters, make your swap module match the target shape so they just work. * Drop ``_rewrite_experts_mapping`` and its 30 lines of converter surgery. * ``update_weight_conversions`` is now ~15 lines, identical in shape to ``Fp8Quantizer.update_weight_conversions``. * ``_build_quant_info_from_gguf_file`` rewrites: walk the model once, look up quant types from a single ``hf_target → quant_type`` map. ~20 lines instead of ~50. * ``gguf_bmm_experts_forward`` splits the gate / up halves of ``gate_up_proj`` at kernel-call time. Only one kernel ref needed (``_id_op_gate_up``) since gate and up share a quant type (the merge converter requires it; Q4_K_M ships both as Q4_K). * ``bind_id_kernel_refs`` resolves two ops (gate_up + down) instead of three. Quantizer shrinks 356 → 288 LOC. Forward kernel cleaner. Verified on dense Llama-3.2-1B (MPS): 112 GgufLinears with bytes loaded via the rename pipeline. 8/8 structural unit tests pass.
…rs + polish #1 Fix ``GgufExperts`` config-attr lookups for arch variation - Read ``num_experts`` / ``num_local_experts`` / ``n_routed_experts`` (Qwen2MoE, Mixtral, DeepSeek-V3 respectively). - Read ``moe_intermediate_size`` / ``intermediate_size`` (Qwen2MoE+DSV3 vs Mixtral). - New ``_read_first(config, *attrs)`` helper. #2 Fold gguf_file loading into the quantizer - ``GgufQuantizeConfig`` gains ``gguf_file`` + ``dequantize`` fields. - ``modeling_utils.from_pretrained``: when ``gguf_file=`` is passed, build a ``GgufQuantizeConfig(gguf_file=..., dequantize=dtype_was_explicit)`` and let the standard ``get_hf_quantizer`` flow construct the quantizer — no bespoke gguf branch. - New ``GGUFQuantizer.load_checkpoint_state(gguf_path)`` hook: ``modeling_utils`` calls it after file resolution; the quantizer owns the GGUF parsing + state-dict return. - ``_resolve_gguf_target_device`` moved into the quantizer (renamed ``_resolve_target_device``). - ``GGUFQuantizer.__init__`` shrinks to ``(quantization_config, **kwargs)`` — no more ``weight_mapping`` / ``gguf_tensors`` / ``gguf_kv`` / ``linear_mode`` kwargs. - ``linear_mode`` finalized in ``_process_model_before_weight_loading`` once the live ``device_map`` is available. #3 Drop ``_load_kernels_direct`` fallback - Originally meant to drop entirely, but ``ArthurZ/gguf-kernels`` isn't on the kernels-package publisher allowlist so ``kernels.get_kernel`` 404s. Keep the ``snapshot_download`` path as the canonical loader (TODO: migrate the kernels repo under ``kernels-community/`` so we can switch to ``get_kernel``). #4-7 Polishes - ``GgufLinear.forward``: matvec when N==1, matmul when N%32==0, per-row matvec otherwise — single y allocation, no per-iteration ``torch.empty``. - ``GGUFQuantize._resolve_quant_type`` inlined (one line at the call site). - ``_build_quant_info`` dispatcher + two helpers collapsed into one method with an ``if self.on_the_fly:`` branch. - ``gguf_writer._ARCH_BY_MODEL_TYPE`` → ``_ARCH_OVERRIDES`` (only the actually aliased entries; default falls back to ``model.config.model_type``). Bug fix: ``GGUFDequantize.convert`` was reshaping uint8 bytes to the target buffer's flat shape, which broke ``ReversePermuteAttnQ`` / ``ReversePermuteAttnK`` for rope-permuted attention projections in Llama-style GGUFs (they need the source 2D ``(out, bytes_per_row)`` to row-permute byte-wise). Preserve the source shape instead — ``ReversePermuteAttn*`` ops work on uint8 byte rows the same way they work on fp32 rows (the reshape only touches ``shape[0]``). Verified on MPS: dense Llama-3.2-1B forward produces non-NaN logits at N=1 (matvec), N=32 (matmul), and N=33 (per-row matvec). All three load paths (dequant, MPS native, on-the-fly) green. 8/8 structural unit tests pass.
… only Per request: the GGUF fast path goes through the standard `kernels` package exclusively. If the repo isn't registered with the kernels-package backend (today: ArthurZ/gguf-kernels — needs migration under kernels-community/), loading fails loudly. No more bespoke .so discovery.
…ONVERTERS Adds ``gpt_oss`` (also emits the ``gate_up_proj`` merge converter); drops ``mixtral`` / ``deepseek_v3`` because they have no entries in ``_GGUF_ARCH_CONVERTERS`` — no GGUF rename rules means no .gguf file routes through them today, so registering them in the swap map is dead code. Source of truth: the four MoE rows under the ``# MoE`` comment in ``modeling_gguf_pytorch_utils._GGUF_ARCH_CONVERTERS``.
Registry now covers every popular MoE arch (whether or not it already has ``_GGUF_ARCH_CONVERTERS`` rules today — the rules can be added later and the swap will Just Work): * ``GgufExperts`` (Mixtral-style layout, fused ``gate_up_proj`` rows = output): qwen2_moe, qwen3_moe, minimax_m2, mixtral, deepseek_v3. * ``GgufExpertsTransposed`` (new): gpt_oss-style ``GptOssExperts`` layout with ``gate_up_proj`` shape ``(num_experts, hidden, 2*intermediate)`` and per-expert biases on both projections. ``GgufExpertsTransposed`` subclasses ``GgufExperts`` and reuses the ``ExpertsInterface`` forward dispatch + the bind-kernel-refs helper. The fused-MoE forward kernel for the transposed layout is a TODO — buffers populate correctly via the rename pipeline, ``forward`` needs a transposed-aware kernel path before it yields correct logits.
FP8 pattern: ``FP8ExpertsInterface(ExpertsInterface)`` carries the FP8 impls;
the base ``ExpertsInterface`` is generic. GGUF was breaking that pattern by
shoving ``gguf_bmm`` / ``gguf_grouped_mm`` into the base ``_global_mapping``
at module-init time.
* Drop ``gguf_bmm`` / ``gguf_grouped_mm`` from ``ExpertsInterface._global_mapping``.
* New ``GgufExpertsInterface(ExpertsInterface)`` in ``gguf_linear`` with its own
``_global_mapping = {"bmm": ..., "grouped_mm": ...}``. Singleton:
``ALL_GGUF_EXPERTS_FUNCTIONS``.
* ``GgufExperts.forward`` dispatches through that singleton (default ``"bmm"``)
instead of the base interface. No more cross-module coupling at import time.
* ``moe.py`` no longer imports from ``gguf_linear``.
* drop ``_resolve_target_device`` — single ``torch.backends.mps.is_available()``
in ``_process_model_before_weight_loading``. No more device_map parsing in
the quantizer.
* drop ``renaming_dequantize_op`` property — the renaming pipeline picks up
GGUFDequantize through ``update_weight_conversions`` anyway, so the extra
hook was dead weight.
* drop ``_bind_kernels`` on GgufLinear — replace with a module-level
``_LINEAR_OPS`` cache keyed by quant_type. ``GgufLinear.forward`` calls
``_linear_ops(self.quant_type)`` to get the (matvec, matmul) pair. One
lookup per quant type per process, shared across every GgufLinear instance.
* drop ``_first_device`` — meta-time swap pattern (FP8-style): wrap the swap
in ``with torch.device("meta"):``, drop the ``device=`` kwargs on
GgufLinear / GgufExperts / GgufExpertsTransposed. Real bytes flow in via
the rename pipeline at weight-load time.
* drop ``_read_first`` — accept that configs spell MoE dims differently and
rely on PretrainedConfig's ``attribute_map`` to alias the names. Archs
whose maps are incomplete (DSV3 missing num_experts, Mixtral missing
moe_intermediate_size) are config-level fixes, not quantizer workarounds.
GgufExpertsTransposed reads ``config.intermediate_size`` directly since
gpt_oss has no moe/dense split.
* drop the ``_make_gguf_experts_interface`` factory — declare
``GgufExpertsInterface`` + ``ALL_GGUF_EXPERTS_FUNCTIONS`` at module level
exactly like FP8 (``moe.py`` no longer imports from gguf_linear, so the
deferred construction was unjustified).
* clean up ``gguf_bmm_experts_forward``: rename ``I_dim`` → ``intermediate_dim``,
``S_dec`` → ``decode_size``, ``S`` annotated. ``bind_id_kernel_refs`` docs
the scratch-buffer policy properly (why we pre-allocate at swap time vs
per-forward, what triggers dynamo recompiles, prefill vs decode).
``_KERNEL_FMT`` consolidated under ``gguf_kernels`` so the
``mul_mat_<fmt>_f32`` name table has one source of truth.
* strip ``\`\`xxxx\`\``` Sphinx-style backtick markers from comments / docstrings.
No behavior changes: dense Llama MPS load → 112 GgufLinears, dequant path →
0, on-the-fly Q4_0 on SmolLM2 → 90. 8/8 unit tests pass. Base
ExpertsInterface = {batched_mm, grouped_mm, sonicmoe}; GgufExpertsInterface
= {bmm, grouped_mm}.
* GgufLinear / GgufExperts / GgufExpertsTransposed each resolve their own kernel OpOverload refs in __init__ and stash them as ``_mv_op`` / ``_mat_op`` / ``_id_op_gate_up`` / ``_id_op_down`` directly on the module. Eager resolution (kernel names are deterministic from quant_type); when the kernels package isn't loadable we leave the attrs ``None`` and forward raises a clear error. * Drop ``bind_id_kernel_refs`` from ``gguf_kernels`` and ``_bind_experts_kernels`` from ``quantizer_gguf.postprocess_model`` — redundant now that the module owns its refs. * Drop the module-level ``_LINEAR_OPS`` cache in ``gguf_linear`` — each GgufLinear keeps its own ops; same getattr cost per process as before, one fewer indirection. * Drop the lazy-bind branch in ``gguf_bmm_experts_forward``; the kernel refs are guaranteed present (or the forward raises). * Scratch buffers for the decode hot path register as meta tensors at __init__; ``.to(device)`` materialises them when the model lands on its real device. No change to the dynamo behaviour. Also: revert an accidental garble in models/llama/modeling_llama.py where stray ``renaming_dequantize_op`` comments had been pasted into ``LlamaModel.__init__`` and broken its indentation.
save_pretrained -> from_pretrained now reconstructs the GGUF swap layout
from the standard transformers config, no .gguf file needed:
* GgufQuantizeConfig gains a ``module_quant_types`` field carrying the
per-module quant plan (Linear -> "Q4_K"; experts ->
{gate_up_quant, down_quant}). Populated by GGUFQuantizer at swap time so
it lands in config.json on save.
* On reload, when ``module_quant_types`` is set the quantizer drives the
swap from it directly (no gguf_tensors walk, no rename pipeline lookup).
``linear_mode`` is forced True for this path — the saved buffers are
uint8, only GgufLinear / GgufExperts can host them.
* ``GgufQuantizeConfig.to_dict()`` drops ``gguf_file`` from the serialised
form: the field is a load-time pointer to an external .gguf, not state
worth re-persisting after the bytes have moved into the models own
safetensors.
* ``GGUFQuantizer.is_serializable`` flips True: buffers + biases + norms
are all standard safetensors dtypes, and the swap plan is replayable
via the config.
Also:
* ``GgufQuantizeConfig.post_init`` validation expanded to the full kernel
suite (Q4_0 / Q5_0 / Q5_1 / Q8_0 / Q4_K / Q5_K / Q6_K / IQ4_NL / IQ4_XS).
The narrower Q4_0/Q8_0 limit (gguf-py quantizer subset) moves into
``GGUFQuantize.convert`` where it actually applies — only the on-the-fly
path cares; loading existing GGUFs supports all kernel-backed quants.
Verified: on-the-fly quantize -> save_pretrained -> from_pretrained
reproduces 90 GgufLinears on SmolLM2 with byte-identical buffers. Unit
tests pass.
* Fix MPS UB on non-writable mmap by `np.copy`ing tensor data at the GGUFReader source, and drop the `GGUFQuantizedTensor` subclass via `as_subclass(torch.Tensor)` so the loader's `.to(device,dtype)` cast no longer corrupts storage. * Restrict the head-pair Q/K permute to `llama` only — every other arch (mistral, phi3, cohere, qwen2/3, deci, nemotron, gemma2/3, stablelm, lfm2, starcoder2, gpt_oss, minimax_m2) stores Q/K in HF layout and was being row-shuffled out of correctness. * Fix arch-specific renames: Bloom `input_layernorm`/`post_attention_layernorm` (not `ln_attn`/`ln_mlp`); Mamba `backbone.embeddings` (plural); StableLM/Nemotron LayerNorm `.bias` renames. * Infer `rope_scaling = llama3` for Llama-3.x GGUFs (arch=llama + freq_base=500000 + max_pos ≥ 16384); GGUF format doesn't store the scaling fields directly. * Apply `SubtractOne` on Gemma2/3 / Nemotron norm weights in fp32 in `load_checkpoint_state`, before the loader's dtype cast — doing it after the cast costs 1 ULP near `w = 1` and breaks weights-conv tests. * `gemma2`/`gemma3` `output_norm` also stores `w + 1` — add the matching `SubtractOne` converter and strip the duplicate plain rename from the shared list. * F32/F16/BF16 GGUF tensors skip the dequant byte-view roundtrip; unsupported i-quants (IQ1_S/M, IQ2_*, IQ3_*) fall back to gguf-py's numpy dequant. * Pre-dequant tensors whose target isn't a swappable Linear/Experts (embeddings, unsupported quants); auto-disable `linear_mode` when `kernels` package isn't loadable so non-MPS hosts run the dequant path. * `update_weight_conversions` now just splices the GGUF rename map in front of the model-side conversions — no `GGUFDequantize` injection (every dequant decision lives in `load_checkpoint_state`). Tests: drop the `pass_through` story, pin the load-time dequant matrix, and add an `arch_coverage` file that ensures every supported model_type stays registered + has a kernel/dequant entry.
`quantizers/__init__.py` imports `GGUFQuantizer` to register it; the top-level `import torch` broke transformers' "no-torch / PIL-only" import test. Move torch into the methods that actually use it. Drops the `pr-ci / Check repository consistency` failure on the PR.
- Drop GGUF save support (gguf_writer.py, save_gguf) — ship in a separate PR. - Remove metal_kernels_available probe; GgufLinear/GgufExperts now let ensure_metal_kernels() raise instead of swallowing into a dead module. - Move gguf_file= config construction into GgufQuantizeConfig.from_gguf_file and the disk-offload guard into GGUFQuantizer.validate_environment, so from_pretrained carries no GGUF special-case. - use_kernels defaults to False directly (no None dance). - Rename linear_mode -> keep_quantized, centralized in _resolve_keep_quantized. - Build the module-swap plan from GGUF header metadata (tensor_quant_types) instead of the materialized tensors. - Fix rename_source_key(prefix=...) -> base_model_prefix=...; the wrong kwarg raised TypeError for every tensor, silently caught, so the gguf_file swap plan was always empty (verified end-to-end: 169 GgufLinear modules now swap and generate correctly on MPS). Remove the masking try/except. - Single-regex skip in replace_with_gguf_linear; use model.set_submodule. - Add CPU coverage for the metadata-driven swap plan. SubtractOne stays a no-op pass-through with the de-offset pre-applied in load_checkpoint_state: slow weights-conversion tests confirm the loader casts to the target dtype before the converter chain, so the fp32 subtraction must happen up-front.
…iven Found via the slow GGUF suite on MPS. - dequantize_gguf_tensor now unwraps a GGUFQuantizedTensor input to a plain tensor up-front. Previously the dequant ops preserved the subclass, so the fp32 output still carried quant_type; the subclass's `.to` override then reinterpreted the storage as the target dtype during the loader's device cast on MPS (fp16/bf16 weights came out as raw-byte garbage like 28512.0; fp32 happened to survive). Fixes garbage generation on the MPS dequant path. - The Gemma/Nemotron `w + 1` norm de-offset still has to be pre-applied on the fp32 source (the loader casts to the target dtype before the converter chain runs — verified SubtractOne sees fp16, so it can't subtract safely). But it is no longer keyed off a hardcoded arch list: load_checkpoint_state derives the offset source patterns from the SubtractOne converters already present in weight_mapping, so any arch whose converter graph declares the offset works for free. SubtractOne is documented as the declarative marker it now is. Validated on MPS: gemma2/gemma3_text/nemotron weights-conversion and qwen2/stablelm generation tests pass; fp16/bf16/fp32 dequant all correct.
Contributor
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, ggml |
Merged
SubtractOne is a real fp32 op again. Rather than pre-subtracting on the fp32 source in load_checkpoint_state, the quantizer forces the `w + 1` norms to stay fp32 for the load via `_keep_in_fp32_modules_strict` (covers fp16 and bf16), keyed off the SubtractOne converter targets in weight_mapping — no arch list. The converter then subtracts in fp32 and the result is exact. Those norms end up fp32 on the GGUF load, so the weights-conversion tests compare at the original dtype. Validated on MPS: nemotron(fp16)/gemma2(fp16)/ gemma3_text(bf16) conversion + stablelm control + qwen2 generation.
load_gguf_checkpoint now computes `tensor_quant_types` ({name: quant_type}) and
`weight_mapping` unconditionally — they are read straight off the GGUF header
(no tensor data), so a `return_tensors=False` call returns them cheaply. Only the
eager `np.copy` of tensor bytes stays behind `return_tensors=True`.
This lets the module-swap plan be built from metadata + renamings alone (pure
name resolution, no tensor load / no conversion). Verified: return_tensors=False
yields 291 quant types + 12 rules with no `tensors`; full load and AutoConfig via
gguf unchanged; 63 fast tests pass.
Replace the rename-based swap plan with a uniform meta-time swap and post-load binding. The gguf_name -> hf_name mapping is no longer computed twice. - GgufLinear / GgufExperts take an optional quant_type: the swap creates uint8 placeholders sized from the Q4_K_M average (good enough for device_map memory estimation), and the exact quant type + metal kernels are bound post-load from the bytes that arrive (bind_after_load). Forward stays branch-free (compile). - GGUFQuantizedTensor.__torch_function__ now propagates quant_type through every subclass-preserving op (the loader's .to, the Q/K permute, the gate/up cat) so bind_after_load can read it back; these subclasses only exist during load. - replace_with_gguf_linear is a uniform swap (modules_to_not_convert + tied output embeddings skipped); the on-the-fly path passes its single config quant_type. load_checkpoint_state drops per-tensor work: a coarse metadata-only "all kernel-supported?" check gates the byte path vs full dequant. - postprocess_model reconciles: bind each module's kernels (from the bytes, or module_quant_types on reload), revert float-source Linears (MoE routers) to nn.Linear, and record module_quant_types for the save round-trip. It also drops the GGUF load-renames so save_pretrained writes HF names (the swapped model is now standard safetensors, re-swapped on reload) — fixes the round-trip. Validated on MPS: dense (qwen2) and MoE (qwen2moe: router revert + experts) swap generation, gemma2/nemotron/gemma3 de-offset conversion, on-the-fly swap, dequant-on-load, and save_pretrained -> from_pretrained round-trip; 64 fast tests.
4 of 6 tasks
model_manager only received a model id, so serving a GGUF failed. Auto-resolve a single .gguf (via _maybe_auto_detect_gguf) and pass gguf_file to AutoConfig + from_pretrained and to the processor load, converting a local .gguf path to dir+filename. Fall back to AutoModelForCausalLM when a GGUF-derived config has architectures=None (routes by model_type).
Collaborator
Author
|
Closing this one: inference-time GGUF matmul landed through #47779 ( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Opt-in GGUF inference on Apple Silicon. Keeps weights at native quant
(Q4_0 / Q5_0 / Q5_1 / Q8_0 / Q4_K / Q5_K / Q6_K / IQ4_NL / IQ4_XS) and
runs forward through the same Metal kernels llama.cpp ships (packaged
as
ArthurZ/gguf-kernels).Enabled via
from_pretrained(..., gguf_file=..., gguf_linear=True).Default off.
Memory
Qwen1.5-MoE-A2.7B Q4_K_M on M3 Max:
dtype=torch.bfloat16gguf_linear=TrueSpeed sweep — M3 Max, MAX_NEW=128
generate()generate()generate_batchN=1generate_batchN=1generate_batchN=1generate_batchN=8generate_batchN=8generate_batchN=8generate_batchN=1generate_batchN=8tg256Aggregate tok/s for N>1.
What made it fast
GgufLinearrunning native Q4_K / Q5_K / Q6_K matvec & matmul Metal kernels — bandwidth-bound decode wins 1.27× vs bf16 matvec.tiisg; second simdgroup recomputed the same rows. Addedsgitgrow split. Expert BMMmul_mat_id+12 % / +19 %._fast_decode_loopingeneration/utils.py— pre-allocated output buffer,argmax(out=), in-place Gumbel-max sampling, single_SYNC_EVERY=16CPU↔MPS sync.setup_for_compilehelper —cache_implementation="static"+dynamo.cache_size_limit=512+torch.compile(forward, mode="reduce-overhead", dynamic=False).PagedAttentionCache.update→index_put_(wasindex_copy_): 55× per call on MPS.paged_decode_attention_f32Metal kernel — Flash-Decoding online softmax reading K/V through the block-table indirection. No gather. 134 µs/call vs ~16 ms/layer for the gather-then-SDPA fallback.kv_paged_write_f32Metal kernel — single-dispatch paged KV write replacing the 5-op Python chain. 6.7 µs/call.nullcontext(torch.cuda.stream(None)was falling throughtorch.mps.current_device()which doesn't exist) — unlocksgenerate_batch+ compile on MPS.reduce-overheadon MPS (wasmax-autotune-no-cudagraphs, autotuning Triton kernels that don't exist on Metal).