bitsandbytes-cpu-fork 0.50.2.dev2__py3-none-win_amd64.whl → 0.50.2.dev4__py3-none-win_amd64.whl

bitsandbytes/backends/cpu/ops.py

@@ -1,805 +1,805 @@
-from collections.abc import Sequence
-import ctypes as ct
-import logging
-import math
-from math import prod
-from typing import Optional
-
-import torch
-
-from bitsandbytes.functional import get_ptr, has_avx512bf16
-
-from ..._ops import register_kernel
-from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib
-
-logger = logging.getLogger(__name__)
-
-# 256 项 fp32 8-bit 码表缓存（见 _gemm_8bit_cpu 的说明）
-_GEMM_8BIT_CODE = None
-
-_has_avx512 = torch.backends.cpu.get_cpu_capability() == "AVX512"
-
-# torch._int_mm for s8@s8->s32 is supported on CPU from torch 2.4+.
-# However, we can overflow if we use this without AVX512_VNNI support.
-# This is fixed in torch 2.6+, so we set this as the minimum to be safe.
-# For more information: https://github.com/pytorch/pytorch/pull/136942
-#
-# Without AVX-512 (including aarch64), torch._int_mm uses a scalar fallback
-# that is much slower than fp32 matmul. Only use it when AVX-512 is available.
-if torch.__version__ >= (2, 6) and _has_avx512:
-
-    @register_kernel("bitsandbytes::int8_linear_matmul", "cpu")
-    def _(A: torch.Tensor, B: torch.Tensor):
-        return torch._int_mm(
-            A.reshape(-1, A.shape[-1]),
-            B.t(),
-        ).reshape(*A.shape[:-1], B.shape[0])
-
-
-if not isinstance(lib, ErrorHandlerMockBNBNativeLibrary):
-
-    @register_kernel("bitsandbytes::quantize_blockwise", "cpu")
-    def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
-        A = A.contiguous()
-        n = A.numel()
-        blocks = -(n // -blocksize)
-
-        absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
-        out = torch.empty(A.shape, device=A.device, dtype=torch.uint8)
-
-        if A.dtype == torch.float32:
-            lib.cquantize_blockwise_cpu_fp32(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(n),
-            )
-        elif A.dtype == torch.bfloat16:
-            lib.cquantize_blockwise_cpu_bf16(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(n),
-            )
-        elif A.dtype == torch.float16:
-            lib.cquantize_blockwise_cpu_fp16(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(n),
-            )
-        else:
-            # Generic fallback for other dtypes
-            A_flat = A.reshape(n).float()
-            rem = n % blocksize
-            has_rem = rem > 0
-            A_com = A_flat[: n - rem]
-            A_com_reshaped = A_com.reshape(n // blocksize, blocksize)
-            absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0]
-            scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1)
-            scaled_A = scaled_A.reshape(-1)
-            if has_rem:
-                absmax[-1] = torch.abs(A_flat[n - rem :]).max()
-                scaled_A_rem = torch.clamp(A_flat[n - rem :] * (1 / absmax[-1]), -1, 1)
-                scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0)
-
-            diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device))
-            out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape)
-
-        return out, absmax
-
-    @register_kernel("bitsandbytes::dequantize_blockwise", "cpu")
-    def _(
-        A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype
-    ) -> torch.Tensor:
-        A = A.contiguous()
-        out = torch.empty_like(A, dtype=dtype)
-        if dtype == torch.float32:
-            lib.cdequantize_blockwise_cpu_fp32(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(A.numel()),
-            )
-        elif dtype == torch.bfloat16:
-            lib.cdequantize_blockwise_cpu_bf16(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(A.numel()),
-            )
-        elif dtype == torch.float16:
-            lib.cdequantize_blockwise_cpu_fp16(
-                get_ptr(code),
-                get_ptr(A),
-                get_ptr(absmax),
-                get_ptr(out),
-                ct.c_longlong(blocksize),
-                ct.c_longlong(A.numel()),
-            )
-        else:
-            out = code[A.reshape(-1).int()]
-            blocks = out.shape[-1] // blocksize
-            res = out.shape[-1] % blocksize
-            if res != 0:
-                out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0)
-            out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1)
-            out = out[: blocks * blocksize + res]
-            out = out.reshape(A.shape)
-
-        return out
-
-    @register_kernel("bitsandbytes::dequantize_4bit", "cpu")
-    def _(
-        A: torch.Tensor,
-        absmax: torch.Tensor,
-        blocksize: int,
-        quant_type: str,
-        shape: Sequence[int],
-        dtype: torch.dtype,
-    ) -> torch.Tensor:
-        # Fallback as AVX512 implementation has accuracy issues with blocksize >= 2048.
-        # Note: this is not a common use case.
-        avx512_fallback = _has_avx512 and blocksize >= 2048
-
-        # Odd shape is not supported by this kernel; fallback to generic implementation
-        shape_fallback = shape[-1] % 2 != 0
-
-        if avx512_fallback or shape_fallback:
-            from ..default.ops import _dequantize_4bit_compute
-            from ..utils import _get_4bit_code
-
-            if A.dtype != torch.uint8:
-                A = A.view(torch.uint8)
-            code = _get_4bit_code(quant_type, A.device)
-            return _dequantize_4bit_compute(A.reshape(-1), absmax, code, blocksize, shape, dtype)
-
-        # Enable non uint8 dtype
-        if A.dtype != torch.uint8:
-            A = A.view(torch.uint8)
-
-        # TODO: support half precision absmax
-        if absmax.dtype != torch.float32:
-            absmax = absmax.float()
-
-        if len(shape) == 1:
-            shape = (1, shape[0])
-
-        m = prod(shape[:-1])
-        n = shape[-1]
-
-        A = A.reshape(m, n // 2)
-        out = torch.empty(shape, dtype=dtype, device=A.device)
-
-        if quant_type == "fp4":
-            if dtype == torch.float32:
-                lib.cdequantize_blockwise_cpu_fp4_fp32(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-            elif dtype == torch.bfloat16:
-                lib.cdequantize_blockwise_cpu_fp4_bf16(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-            elif dtype == torch.float16:
-                lib.cdequantize_blockwise_cpu_fp4_fp16(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-        elif quant_type == "nf4":
-            if dtype == torch.float32:
-                lib.cdequantize_blockwise_cpu_nf4_fp32(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-            elif dtype == torch.bfloat16:
-                lib.cdequantize_blockwise_cpu_nf4_bf16(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-            elif dtype == torch.float16:
-                lib.cdequantize_blockwise_cpu_nf4_fp16(
-                    get_ptr(A),
-                    get_ptr(absmax),
-                    get_ptr(out),
-                    ct.c_longlong(blocksize),
-                    ct.c_longlong(m),
-                    ct.c_longlong(n),
-                )
-        else:
-            raise ValueError
-
-        return out
-
-    # NOTE 2026-09-22: the registration below used to live INSIDE `if has_avx512bf16():`,
-    # which meant AVX2-only CPUs (Zen2 4500U / Comet Lake i5 / EPYC Zen2-3) never
-    # registered a CPU kernel for gemv_4bit at all. They silently fell through to the
-    # default backend, which dequantizes the whole weight matrix and then runs a dense
-    # linear -- measured 12.3 ms per call vs 1.16 ms for the fused kernel below (10.6x).
-    # The guard is therefore narrowed to cover only the AVX512-BF16-only HuggingFace
-    # kernel load; the fused registration is unconditional.
-    gemm_4bit_forward_kernel = None
-    if has_avx512bf16():
-        gemm_4bit_forward_kernel = None
-        try:
-            from kernels import get_kernel
-
-            gemm_4bit_forward_kernel = get_kernel(
-                "kernels-community/quantization-bitsandbytes", version=1
-            ).gemm_4bit_forward
-        except Exception as exc:  # pragma: no cover - best effort fallback
-            gemm_4bit_forward_kernel = None
-            logger.warning(
-                "Failed to load CPU gemm_4bit_forward from kernels-community: %s. Please make sure you already `pip install kernels` and the kernels >= 0.11.1",
-                exc,
-            )
-
-    @register_kernel("bitsandbytes::gemv_4bit", "cpu")
-    def _(
-        A: torch.Tensor,
-        B: torch.Tensor,
-        shapeB: Sequence[int],
-        absmax: torch.Tensor,
-        code: torch.Tensor,
-        blocksize: int,
-    ) -> torch.Tensor:
-        if B.dtype != torch.uint8:
-            B = B.contiguous().view(torch.uint8)
-        dtype = A.dtype
-        A_orig = A
-        absmax_orig = absmax
-        quant_type = "fp4" if code[1] > 0 else "nf4"
-        # cpu fused op only support bf16 for now.
-        if dtype != torch.bfloat16:
-            A = A.to(torch.bfloat16)
-        if absmax.dtype != torch.bfloat16:
-            absmax = absmax.to(torch.bfloat16)
-
-        final_out_shape = (*A.shape[:-1], shapeB[0])
-        A = A.reshape(-1, A.shape[-1])
-        out_shape = (*A.shape[:-1], shapeB[0])
-        if gemm_4bit_forward_kernel is not None:
-            quant_type_num = 1 if quant_type == "fp4" else 0
-            # C++ kernel expects weight shape (N, K_packed), ensure 2D contiguous
-            B_2d = B.reshape(shapeB[0], -1).contiguous()
-            out = gemm_4bit_forward_kernel(A, B_2d, absmax, blocksize, quant_type_num)
-        else:
-            # AVX2 回退：用所有 CPU 构建都导出的 cgemv_4bit_inference_cpu_* 融合内核
-            # （lib.gemv_4bit_inference_cpu_fp4/nf4_bf16 只存在于 AVX512+BF16 构建，
-            #  AVX2 机器走到会 AttributeError；data_type: FP4=1 / NF4=2）
-            A_c = A_orig.reshape(-1, A_orig.shape[-1]).contiguous()
-            B_c = B.reshape(shapeB[0], -1).contiguous()
-            absmax_c = absmax_orig.float().contiguous()
-            data_type_num = 1 if quant_type == "fp4" else 2
-            M, K = A_c.shape[0], A_c.shape[1]
-            N = shapeB[0]
-            out = torch.empty(out_shape, dtype=A_c.dtype, device=A_c.device)
-            if dtype == torch.float32:
-                lib.cgemv_4bit_inference_cpu_fp32(
-                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
-                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
-                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
-                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
-                    ct.c_int(data_type_num),
-                )
-            elif dtype == torch.bfloat16:
-                lib.cgemv_4bit_inference_cpu_bf16(
-                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
-                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
-                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
-                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
-                    ct.c_int(data_type_num),
-                )
-            else:
-                lib.cgemv_4bit_inference_cpu_fp16(
-                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
-                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
-                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
-                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
-                    ct.c_int(data_type_num),
-                )
-
-        if dtype != torch.bfloat16:
-            out = out.to(dtype)
-
-        return out.reshape(final_out_shape)
-
-    # ---- gemm_4bit: CPU 注册（2026-09-22 新增）----
-    # 此前这个 op 【没有 cpu 注册】，于是退到 default 后端做"反量化 + 稠密线性"。
-    # autograd/_functions.py:324/334/462/472 调的就是它 ⇒ 真实模型里
-    # 【每一次批量 4-bit 线性前向】都在走慢路径。实测（K=N=4096）：
-    #     M=1   gemm_4bit 14.483 ms  vs 融合内核  1.448 ms  ⇒ 10.00x
-    #     M=8   gemm_4bit 32.494 ms  vs 融合内核  4.124 ms  ⇒  7.88x
-    #     M=64  gemm_4bit 38.576 ms  vs 融合内核 30.382 ms  ⇒  1.27x
-    #     M=256 gemm_4bit 75.532 ms  vs 融合内核 164.840 ms ⇒  0.46x  ← 融合反而慢
-    # 融合内核的 m 循环每 4 行重读一遍权重，所以大 M 时它输给 oneDNN 稠密。
-    # ⇒ 按 M 分档：小 M 走融合内核，大 M 交回 default。
-    # ---- 2026-09-22 实测三方交叉点（K=N=4096，OMP=6）----
-    #        M     融合 ms    default ms   融合优势
-    #        1      1.20       ~14.5       12x
-    #        8      4.71        25.14       5.3x
-    #       16      7.04        25.75       3.7x
-    #       17      7.79        26.43       3.4x
-    #       32     16.20        28.41       1.75x
-    #       64     28.79        31.37       1.09x
-    #       96     43.45        37.01       default 反超
-    #      128     62.11        41.77       default 快 1.49x
-    #      192     85.78        54.67       default 快 1.57x
-    #      256    166.11        64.94       default 快 2.56x
-    # ⇒ 融合内核一直赢到 M=64，default 要到 M≈96 才反超。
-    # 阈值原先写 16 ⇒ M=17~64 被错误地分给 default，白慢最多 3.4x。
-    #
-    # ⚠️ 一个被实测否决的尝试：给 C 层加"权重行解码一次、对全部 M 行复用"的路径，
-    #    结果是【负优化】（M=17 慢 1.52x、M=192 慢 1.44x）——
-    #    它把权重解码从 O(M/4) 降到 O(1)，却让 A 被每个输出列重读一次，
-    #    A 流量变成 O(N*M*K)（M=17 时 1.14 GB），远比重新解码权重贵。
-    #    已删除，只在 cpu_ops.cpp 里留了说明。
-    _GEMM_4BIT_FUSED_MAX_M = 64
-
-    @register_kernel("bitsandbytes::gemm_4bit", "cpu")
-    def _gemm_4bit_cpu(
-        A: torch.Tensor,
-        B: torch.Tensor,
-        shapeB: Sequence[int],
-        absmax: torch.Tensor,
-        blocksize: int,
-        quant_type: str,
-        bias: Optional[torch.Tensor] = None,
-        absmax_8bit: Optional[torch.Tensor] = None,
-        absmax_code: Optional[torch.Tensor] = None,
-        absmax_offset: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        m_rows = A.numel() // A.shape[-1] if A.numel() else 0
-        # 融合路径的限制：无 bias、非 nested absmax、张量是 2D 可 reshape 的、
-        # 且 M 足够小（大 M 时 oneDNN 稠密更快）。
-        can_fuse = (
-            gemm_4bit_forward_kernel is None      # 没有 HF kernel 时才用我们的融合内核
-            and bias is None
-            and absmax_8bit is None
-            and m_rows > 0
-            and m_rows <= _GEMM_4BIT_FUSED_MAX_M
-        )
-        if can_fuse:
-            dtype = A.dtype
-            A_c = A.reshape(-1, A.shape[-1]).contiguous()
-            B_c = B.reshape(shapeB[0], -1)
-            if B_c.dtype != torch.uint8:
-                B_c = B_c.contiguous().view(torch.uint8)
-            else:
-                B_c = B_c.contiguous()
-            absmax_c = absmax.float().contiguous()
-            M, K = A_c.shape[0], A_c.shape[1]
-            N = shapeB[0]
-            data_type_num = 1 if quant_type == "fp4" else 2
-            out = torch.empty((M, N), dtype=dtype, device=A_c.device)
-            fn = {
-                torch.float32: lib.cgemv_4bit_inference_cpu_fp32,
-                torch.bfloat16: lib.cgemv_4bit_inference_cpu_bf16,
-                torch.float16: lib.cgemv_4bit_inference_cpu_fp16,
-            }[dtype]
-            fn(
-                get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
-                ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
-                ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
-                ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
-                ct.c_int(data_type_num),
-            )
-            return out.reshape(*A.shape[:-1], N)
-
-        from ..default.ops import _gemm_4bit_default_impl
-
-        return _gemm_4bit_default_impl(
-            A, B, shapeB, absmax, blocksize, quant_type,
-            bias, absmax_8bit, absmax_code, absmax_offset,
-        )
-
-    # ---- 大 M 分档（2026-09-23 实测）----
-    # 融合内核的内层是 m0 += 4，每 4 行就把整块权重重新读一遍 ⇒ 权重流量 O(M/4)，
-    # 大 M 时输给"反量化 + oneDNN 稠密 matmul"。
-    # 干净基准（gemm8_test.py：每格子进程 + 取 7 次最小值），融合/两步 的比值：
-    #     M=16   0.18x / 0.15x / 0.25x     融合远优
-    #     M=64   0.89x / 0.65x / 0.77x     融合仍优
-    #     M=128  1.34x / 1.14x / 1.32x     ← 交叉点在这一带
-    #     M=256  1.54x / 1.34x / 1.51x
-    #     M=512  2.04x / 1.82x / 2.04x     两步快约一倍
-    # ⇒ 阈值取 64（保守：M=64 时融合还是赢的，且真实推理的 M 是 seq_len）。
-    #
-    # ⚠️ 这个 op 【没有 default 后端回退】（全库只注册了 cpu 与 cuda），
-    #    所以大 M 路径必须在这里自己实现。8-bit 的码表是 create_dynamic_map()，
-    #    与 quantize_blockwise 默认用的同一个 ⇒ 数值一致。
-    _GEMM_8BIT_FUSED_MAX_M = 64
-
-    @register_kernel("bitsandbytes::gemm_8bit", "cpu")
-    def _gemm_8bit_cpu(
-        A: torch.Tensor,
-        B: torch.Tensor,
-        absmax: torch.Tensor,
-        blocksize: int,
-    ) -> torch.Tensor:
-        """融合 8bit 反量化 GEMM：out = A @ dequant8(B)（B 为 uint8 码流，逐行块量化）。
-
-        权重在整个乘法中保持 uint8（DRAM 流量为 fp32 的 1/4），无 fp32 临时张量。
-        小 M（<= 64）走融合内核；大 M 退到"反量化 + 稠密 matmul"（见上面的分档说明）。
-        """
-        if A.numel() == 0:
-            return torch.empty((*A.shape[:-1], B.shape[0]), dtype=torch.float32, device=A.device)
-        A_f = A.reshape(-1, A.shape[-1]).contiguous()
-        B_c = B.contiguous()
-        absmax_c = absmax.contiguous().float()
-        M, K = A_f.shape
-        N = B_c.shape[0]
-
-        if M > _GEMM_8BIT_FUSED_MAX_M:
-            # 大 M：两步更快（实测 M=512 时 1.8~2.0x）
-            from bitsandbytes.functional import create_dynamic_map, dequantize_blockwise
-
-            code = create_dynamic_map().to(B_c.device)
-            W = dequantize_blockwise(B_c, absmax=absmax_c, code=code, blocksize=blocksize)
-            return (A_f @ W.t()).reshape(*A.shape[:-1], N)
-
-        out = torch.empty((M, N), dtype=torch.float32, device=A_f.device)
-        lib.cgemm_8bit_inference_cpu_fp32(
-            get_ptr(A_f),
-            get_ptr(B_c),
-            get_ptr(absmax_c),
-            get_ptr(out),
-            ct.c_longlong(M),
-            ct.c_longlong(N),
-            ct.c_longlong(K),
-            ct.c_longlong(A_f.stride(0)),
-            ct.c_longlong(B_c.stride(0)),
-            ct.c_longlong(out.stride(0)),
-            ct.c_longlong(blocksize),
-        )
-        return out.reshape(*A.shape[:-1], N)
-
-
-# ==================== CPU Optimizer Kernels ====================
-
-
-def _compute_update_norm_and_scale(
-    update: torch.Tensor,
-    unorm_vec: Optional[torch.Tensor],
-    max_unorm: float,
-    param_norm: float,
-) -> float:
-    """Compute trust-ratio scaling factor for LAMB/LARS and store update norm."""
-    if max_unorm <= 0.0:
-        return 1.0
-    unorm = torch.norm(update).item()
-    if unorm_vec is not None:
-        unorm_vec.fill_(unorm)
-    if unorm > max_unorm * param_norm:
-        return (max_unorm * param_norm) / unorm
-    return 1.0
-
-
-@torch.no_grad()
-def _optimizer_update_32bit_cpu(
-    optimizer_name: str,
-    g: torch.Tensor,
-    p: torch.Tensor,
-    state1: torch.Tensor,
-    state2: Optional[torch.Tensor],
-    unorm_vec: Optional[torch.Tensor],
-    max_unorm: float,
-    param_norm: float,
-    beta1: float,
-    beta2: float,
-    beta3: float,
-    alpha: float,
-    eps: float,
-    weight_decay: float,
-    step: int,
-    lr: float,
-    gnorm_scale: float,
-    skip_zeros: bool = False,
-) -> None:
-    g_float = g.float() * gnorm_scale
-    p_float = p.data.float()
-
-    if optimizer_name in ("adam", "lamb"):
-        # Adam / LAMB (2-state): m and v
-        state1.mul_(beta1).add_(g_float, alpha=1.0 - beta1)
-        state2.mul_(beta2).addcmul_(g_float, g_float, value=1.0 - beta2)
-
-        correction1 = 1.0 - beta1**step
-        correction2 = math.sqrt(1.0 - beta2**step)
-        step_size = -lr * correction2 / correction1
-
-        if weight_decay > 0.0:
-            p_float.mul_(1.0 - lr * weight_decay)
-
-        update = state1 / (state2.sqrt() + eps * correction2)
-
-        update_scale = _compute_update_norm_and_scale(update, unorm_vec, max_unorm, param_norm)
-        p_float.add_(update, alpha=step_size * update_scale)
-
-    elif optimizer_name == "ademamix":
-        # AdEMAMix (2-state): state1 shape is (2, *p.shape), state1[0]=m1, state1[1]=m2
-        m1 = state1[0]
-        m2 = state1[1]
-        nu = state2
-
-        m1.mul_(beta1).add_(g_float, alpha=1.0 - beta1)
-        m2.mul_(beta3).add_(g_float, alpha=1.0 - beta3)
-        nu.mul_(beta2).addcmul_(g_float, g_float, value=1.0 - beta2)
-
-        correction1 = 1.0 - beta1**step
-        correction2 = math.sqrt(1.0 - beta2**step)
-
-        if weight_decay > 0.0:
-            p_float.mul_(1.0 - lr * weight_decay)
-
-        mixed_momentum = (m1 / correction1) + (alpha * m2)
-        adaptive_term = (nu.sqrt() / correction2) + eps
-        p_float.add_(mixed_momentum / adaptive_term, alpha=-lr)
-
-    elif optimizer_name in ("momentum", "lars"):
-        # SGD with momentum / LARS (1-state)
-        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
-
-        if step == 1:
-            state1.copy_(g_wd)
-        else:
-            state1.mul_(beta1).add_(g_wd)
-
-        update_scale = _compute_update_norm_and_scale(state1, unorm_vec, max_unorm, param_norm)
-        p_float.add_(state1, alpha=-lr * update_scale)
-
-    elif optimizer_name == "lion":
-        # Lion (2-state sign update)
-        if weight_decay > 0.0:
-            p_float.mul_(1.0 - lr * weight_decay)
-
-        update = state1.mul(beta1).add(g_float, alpha=1.0 - beta1)
-        p_float.add_(update.sign(), alpha=-lr)
-
-        state1.mul_(beta2).add_(g_float, alpha=1.0 - beta2)
-
-    elif optimizer_name == "rmsprop":
-        # RMSprop (1-state)
-        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
-        state1.mul_(beta1).addcmul_(g_wd, g_wd, value=1.0 - beta1)
-
-        update = g_wd / (state1.sqrt() + eps)
-        update_scale = _compute_update_norm_and_scale(update, unorm_vec, max_unorm, param_norm)
-        p_float.add_(update, alpha=-lr * update_scale)
-
-    elif optimizer_name == "adagrad":
-        # Adagrad (1-state)
-        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
-        state1.addcmul_(g_wd, g_wd, value=1.0)
-
-        update = g_wd / (state1.sqrt() + eps)
-        p_float.add_(update, alpha=-lr)
-
-    else:
-        raise ValueError(f"Unsupported optimizer for CPU: {optimizer_name}")
-
-    # Write back to original precision
-    p.data.copy_(p_float)
-
-
-register_kernel("bitsandbytes::optimizer_update_32bit", "cpu")(_optimizer_update_32bit_cpu)
-
-
-@torch.no_grad()
-def _dequant_blockwise_fp32_direct(
-    A_uint8: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int
-) -> torch.Tensor:
-    return torch.ops.bitsandbytes.dequantize_blockwise(A_uint8, absmax, code, blocksize, torch.float32)
-
-
-def _quant_blockwise_fp32_direct(
-    A_fp32: torch.Tensor, code: torch.Tensor, absmax_out: torch.Tensor, out_uint8: torch.Tensor, blocksize: int
-) -> None:
-    out, absmax = torch.ops.bitsandbytes.quantize_blockwise(A_fp32, code, blocksize)
-    out_uint8.copy_(out)
-    absmax_out.copy_(absmax)
-
-
-def _optimizer_update_8bit_blockwise_cpu(
-    optimizer_name: str,
-    g: torch.Tensor,
-    p: torch.Tensor,
-    state1: torch.Tensor,
-    state2: Optional[torch.Tensor],
-    beta1: float,
-    beta2: float,
-    beta3: float,
-    alpha: float,
-    eps: float,
-    step: int,
-    lr: float,
-    qmap1: torch.Tensor,
-    qmap2: Optional[torch.Tensor],
-    absmax1: torch.Tensor,
-    absmax2: Optional[torch.Tensor],
-    weight_decay: float,
-    gnorm_scale: float,
-    skip_zeros: bool = False,
-) -> None:
-    # Fused single-pass C kernel (dequant + update + p update + requant in one
-    # sweep). Falls back to the composite PyTorch path below when the native
-    # library is unavailable or the optimizer is not supported by the kernel.
-    _OPT_IDS = {"adam": 0, "lamb": 0, "momentum": 1, "lars": 1, "lion": 2, "rmsprop": 3, "adagrad": 4, "ademamix": 5}
-    n = p.numel()
-    # The C kernel walks raw pointers and assumes C-contiguous layouts: a
-    # non-contiguous grad (transposed / sliced views are common in training)
-    # would silently read the wrong elements and corrupt the states.
-    contiguous_layout = (
-        g.is_contiguous()
-        and p.is_contiguous()
-        and state1.is_contiguous()
-        and (state2 is None or state2.is_contiguous())
-        and qmap1.is_contiguous()
-        and (qmap2 is None or qmap2.is_contiguous())
-        and absmax1.is_contiguous()
-        and (absmax2 is None or absmax2.is_contiguous())
-    )
-    if (
-        not isinstance(lib, ErrorHandlerMockBNBNativeLibrary)
-        and contiguous_layout
-        and optimizer_name in _OPT_IDS
-        and g.dtype in (torch.float32, torch.bfloat16, torch.float16)
-        # ademamix packs a second momentum into state1[n:2n] with its absmax in
-        # absmax1[blocks:2*blocks]; the flat C layout requires that packing.
-        and (optimizer_name != "ademamix" or state1.numel() >= 2 * n)
-    ):
-        try:
-            dtype_id = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}[g.dtype]
-            # Every scalar needs an explicit ctypes width: a bare Python int is
-            # marshalled as 32-bit c_int, but `n` is a long long. libffi writes
-            # only the low 4 bytes of its 8-byte stack slot; the C side reads
-            # all 8, so stale stack bits leak into n (observed: 0x7FFF00004000
-            # for n=16384) and the kernel walks out of bounds -> SIGSEGV.
-            lib.coptimizer_update_8bit_blockwise_cpu(
-                ct.c_int(_OPT_IDS[optimizer_name]),
-                get_ptr(g),
-                get_ptr(p),
-                get_ptr(state1),
-                get_ptr(state2) if state2 is not None else ct.c_void_p(0),
-                ct.c_float(beta1),
-                ct.c_float(beta2),
-                ct.c_float(beta3),
-                ct.c_float(alpha),
-                ct.c_float(eps),
-                ct.c_int(step),
-                ct.c_float(lr),
-                get_ptr(qmap1),
-                get_ptr(qmap2) if qmap2 is not None else ct.c_void_p(0),
-                get_ptr(absmax1),
-                get_ptr(absmax2) if absmax2 is not None else ct.c_void_p(0),
-                ct.c_float(weight_decay),
-                ct.c_float(gnorm_scale),
-                ct.c_int(1 if skip_zeros else 0),
-                ct.c_longlong(n),
-                ct.c_int(dtype_id),
-            )
-            return
-        except (AttributeError, RuntimeError):
-            pass  # old native library without the fused entry point
-
-    blocksize = 256
-
-    # Dequantize states
-    if optimizer_name == "ademamix" and absmax1.ndim == 2:
-        s1_1 = _dequant_blockwise_fp32_direct(state1[0], absmax1[0], qmap1, blocksize)
-        s1_2 = _dequant_blockwise_fp32_direct(state1[1], absmax1[1], qmap1, blocksize)
-        state1_fp32 = torch.stack([s1_1, s1_2])
-    else:
-        state1_fp32 = _dequant_blockwise_fp32_direct(state1, absmax1, qmap1, blocksize)
-
-    state2_fp32 = None
-    if state2 is not None and qmap2 is not None and absmax2 is not None:
-        state2_fp32 = _dequant_blockwise_fp32_direct(state2, absmax2, qmap2, blocksize)
-
-    grad = g.float() * gnorm_scale
-    p_fp32 = p.data.float()
-
-    if optimizer_name in ("adam", "lamb"):
-        state1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1)
-        state2_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
-
-        correction1 = 1.0 - beta1**step
-        correction2 = math.sqrt(1.0 - beta2**step)
-
-        denom = (state2_fp32.sqrt() / correction2).add_(eps)
-        if weight_decay > 0.0:
-            p_fp32.mul_(1.0 - lr * weight_decay)
-        p_fp32.addcdiv_(state1_fp32, denom, value=-lr / correction1)
-
-    elif optimizer_name == "ademamix":
-        m1_fp32, m2_fp32 = state1_fp32[0], state1_fp32[1]
-        nu_fp32 = state2_fp32
-
-        m1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1)
-        m2_fp32.mul_(beta3).add_(grad, alpha=1.0 - beta3)
-        nu_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
-
-        correction1 = 1.0 - beta1**step
-        correction2 = math.sqrt(1.0 - beta2**step)
-
-        update = (m1_fp32 / correction1 + alpha * m2_fp32) / (nu_fp32.sqrt() / correction2 + eps)
-        if weight_decay > 0.0:
-            p_fp32.mul_(1.0 - lr * weight_decay)
-        p_fp32.add_(update, alpha=-lr)
-
-        state1_fp32 = torch.stack([m1_fp32, m2_fp32])
-
-    elif optimizer_name in ("momentum", "lars"):
-        grad.add_(p_fp32, alpha=weight_decay)
-        if step == 1:
-            state1_fp32.copy_(grad)
-        else:
-            state1_fp32.mul_(beta1).add_(grad)
-        p_fp32.add_(state1_fp32, alpha=-lr)
-
-    elif optimizer_name == "lion":
-        if weight_decay > 0.0:
-            p_fp32.mul_(1.0 - lr * weight_decay)
-
-        update_dir = torch.sign(state1_fp32.mul(beta1) + grad.mul(1.0 - beta1))
-        p_fp32.add_(update_dir, alpha=-lr)
-
-        state1_fp32.mul_(beta2).add_(grad, alpha=1.0 - beta2)
-
-    elif optimizer_name == "rmsprop":
-        grad.add_(p_fp32, alpha=weight_decay)
-        state1_fp32.mul_(beta1).addcmul_(grad, grad, value=1.0 - beta1)
-        p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr)
-
-    elif optimizer_name == "adagrad":
-        grad.add_(p_fp32, alpha=weight_decay)
-        state1_fp32.addcmul_(grad, grad, value=1.0)
-        p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr)
-
-    else:
-        raise ValueError(f"Unsupported optimizer for CPU 8-bit: {optimizer_name}")
-
-    p.data.copy_(p_fp32)
-
-    # Re-quantize states
-    if optimizer_name == "ademamix":
-        _quant_blockwise_fp32_direct(state1_fp32[0], qmap1, absmax1[0], state1[0], blocksize)
-        _quant_blockwise_fp32_direct(state1_fp32[1], qmap1, absmax1[1], state1[1], blocksize)
-        _quant_blockwise_fp32_direct(state2_fp32, qmap2, absmax2, state2, blocksize)
-    else:
-        _quant_blockwise_fp32_direct(state1_fp32, qmap1, absmax1, state1, blocksize)
-        if state2_fp32 is not None:
-            _quant_blockwise_fp32_direct(state2_fp32, qmap2, absmax2, state2, blocksize)
-
-
-register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "cpu")(_optimizer_update_8bit_blockwise_cpu)
+from collections.abc import Sequence
+import ctypes as ct
+import logging
+import math
+from math import prod
+from typing import Optional
+
+import torch
+
+from bitsandbytes.functional import get_ptr, has_avx512bf16
+
+from ..._ops import register_kernel
+from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib
+
+logger = logging.getLogger(__name__)
+
+# 256 项 fp32 8-bit 码表缓存（见 _gemm_8bit_cpu 的说明）
+_GEMM_8BIT_CODE = None
+
+_has_avx512 = torch.backends.cpu.get_cpu_capability() == "AVX512"
+
+# torch._int_mm for s8@s8->s32 is supported on CPU from torch 2.4+.
+# However, we can overflow if we use this without AVX512_VNNI support.
+# This is fixed in torch 2.6+, so we set this as the minimum to be safe.
+# For more information: https://github.com/pytorch/pytorch/pull/136942
+#
+# Without AVX-512 (including aarch64), torch._int_mm uses a scalar fallback
+# that is much slower than fp32 matmul. Only use it when AVX-512 is available.
+if torch.__version__ >= (2, 6) and _has_avx512:
+
+    @register_kernel("bitsandbytes::int8_linear_matmul", "cpu")
+    def _(A: torch.Tensor, B: torch.Tensor):
+        return torch._int_mm(
+            A.reshape(-1, A.shape[-1]),
+            B.t(),
+        ).reshape(*A.shape[:-1], B.shape[0])
+
+
+if not isinstance(lib, ErrorHandlerMockBNBNativeLibrary):
+
+    @register_kernel("bitsandbytes::quantize_blockwise", "cpu")
+    def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
+        A = A.contiguous()
+        n = A.numel()
+        blocks = -(n // -blocksize)
+
+        absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
+        out = torch.empty(A.shape, device=A.device, dtype=torch.uint8)
+
+        if A.dtype == torch.float32:
+            lib.cquantize_blockwise_cpu_fp32(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(n),
+            )
+        elif A.dtype == torch.bfloat16:
+            lib.cquantize_blockwise_cpu_bf16(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(n),
+            )
+        elif A.dtype == torch.float16:
+            lib.cquantize_blockwise_cpu_fp16(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(n),
+            )
+        else:
+            # Generic fallback for other dtypes
+            A_flat = A.reshape(n).float()
+            rem = n % blocksize
+            has_rem = rem > 0
+            A_com = A_flat[: n - rem]
+            A_com_reshaped = A_com.reshape(n // blocksize, blocksize)
+            absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0]
+            scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1)
+            scaled_A = scaled_A.reshape(-1)
+            if has_rem:
+                absmax[-1] = torch.abs(A_flat[n - rem :]).max()
+                scaled_A_rem = torch.clamp(A_flat[n - rem :] * (1 / absmax[-1]), -1, 1)
+                scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0)
+
+            diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device))
+            out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape)
+
+        return out, absmax
+
+    @register_kernel("bitsandbytes::dequantize_blockwise", "cpu")
+    def _(
+        A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype
+    ) -> torch.Tensor:
+        A = A.contiguous()
+        out = torch.empty_like(A, dtype=dtype)
+        if dtype == torch.float32:
+            lib.cdequantize_blockwise_cpu_fp32(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(A.numel()),
+            )
+        elif dtype == torch.bfloat16:
+            lib.cdequantize_blockwise_cpu_bf16(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(A.numel()),
+            )
+        elif dtype == torch.float16:
+            lib.cdequantize_blockwise_cpu_fp16(
+                get_ptr(code),
+                get_ptr(A),
+                get_ptr(absmax),
+                get_ptr(out),
+                ct.c_longlong(blocksize),
+                ct.c_longlong(A.numel()),
+            )
+        else:
+            out = code[A.reshape(-1).int()]
+            blocks = out.shape[-1] // blocksize
+            res = out.shape[-1] % blocksize
+            if res != 0:
+                out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0)
+            out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1)
+            out = out[: blocks * blocksize + res]
+            out = out.reshape(A.shape)
+
+        return out
+
+    @register_kernel("bitsandbytes::dequantize_4bit", "cpu")
+    def _(
+        A: torch.Tensor,
+        absmax: torch.Tensor,
+        blocksize: int,
+        quant_type: str,
+        shape: Sequence[int],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        # Fallback as AVX512 implementation has accuracy issues with blocksize >= 2048.
+        # Note: this is not a common use case.
+        avx512_fallback = _has_avx512 and blocksize >= 2048
+
+        # Odd shape is not supported by this kernel; fallback to generic implementation
+        shape_fallback = shape[-1] % 2 != 0
+
+        if avx512_fallback or shape_fallback:
+            from ..default.ops import _dequantize_4bit_compute
+            from ..utils import _get_4bit_code
+
+            if A.dtype != torch.uint8:
+                A = A.view(torch.uint8)
+            code = _get_4bit_code(quant_type, A.device)
+            return _dequantize_4bit_compute(A.reshape(-1), absmax, code, blocksize, shape, dtype)
+
+        # Enable non uint8 dtype
+        if A.dtype != torch.uint8:
+            A = A.view(torch.uint8)
+
+        # TODO: support half precision absmax
+        if absmax.dtype != torch.float32:
+            absmax = absmax.float()
+
+        if len(shape) == 1:
+            shape = (1, shape[0])
+
+        m = prod(shape[:-1])
+        n = shape[-1]
+
+        A = A.reshape(m, n // 2)
+        out = torch.empty(shape, dtype=dtype, device=A.device)
+
+        if quant_type == "fp4":
+            if dtype == torch.float32:
+                lib.cdequantize_blockwise_cpu_fp4_fp32(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+            elif dtype == torch.bfloat16:
+                lib.cdequantize_blockwise_cpu_fp4_bf16(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+            elif dtype == torch.float16:
+                lib.cdequantize_blockwise_cpu_fp4_fp16(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+        elif quant_type == "nf4":
+            if dtype == torch.float32:
+                lib.cdequantize_blockwise_cpu_nf4_fp32(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+            elif dtype == torch.bfloat16:
+                lib.cdequantize_blockwise_cpu_nf4_bf16(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+            elif dtype == torch.float16:
+                lib.cdequantize_blockwise_cpu_nf4_fp16(
+                    get_ptr(A),
+                    get_ptr(absmax),
+                    get_ptr(out),
+                    ct.c_longlong(blocksize),
+                    ct.c_longlong(m),
+                    ct.c_longlong(n),
+                )
+        else:
+            raise ValueError
+
+        return out
+
+    # NOTE 2026-09-22: the registration below used to live INSIDE `if has_avx512bf16():`,
+    # which meant AVX2-only CPUs (Zen2 4500U / Comet Lake i5 / EPYC Zen2-3) never
+    # registered a CPU kernel for gemv_4bit at all. They silently fell through to the
+    # default backend, which dequantizes the whole weight matrix and then runs a dense
+    # linear -- measured 12.3 ms per call vs 1.16 ms for the fused kernel below (10.6x).
+    # The guard is therefore narrowed to cover only the AVX512-BF16-only HuggingFace
+    # kernel load; the fused registration is unconditional.
+    gemm_4bit_forward_kernel = None
+    if has_avx512bf16():
+        gemm_4bit_forward_kernel = None
+        try:
+            from kernels import get_kernel
+
+            gemm_4bit_forward_kernel = get_kernel(
+                "kernels-community/quantization-bitsandbytes", version=1
+            ).gemm_4bit_forward
+        except Exception as exc:  # pragma: no cover - best effort fallback
+            gemm_4bit_forward_kernel = None
+            logger.warning(
+                "Failed to load CPU gemm_4bit_forward from kernels-community: %s. Please make sure you already `pip install kernels` and the kernels >= 0.11.1",
+                exc,
+            )
+
+    @register_kernel("bitsandbytes::gemv_4bit", "cpu")
+    def _(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        shapeB: Sequence[int],
+        absmax: torch.Tensor,
+        code: torch.Tensor,
+        blocksize: int,
+    ) -> torch.Tensor:
+        if B.dtype != torch.uint8:
+            B = B.contiguous().view(torch.uint8)
+        dtype = A.dtype
+        A_orig = A
+        absmax_orig = absmax
+        quant_type = "fp4" if code[1] > 0 else "nf4"
+        # cpu fused op only support bf16 for now.
+        if dtype != torch.bfloat16:
+            A = A.to(torch.bfloat16)
+        if absmax.dtype != torch.bfloat16:
+            absmax = absmax.to(torch.bfloat16)
+
+        final_out_shape = (*A.shape[:-1], shapeB[0])
+        A = A.reshape(-1, A.shape[-1])
+        out_shape = (*A.shape[:-1], shapeB[0])
+        if gemm_4bit_forward_kernel is not None:
+            quant_type_num = 1 if quant_type == "fp4" else 0
+            # C++ kernel expects weight shape (N, K_packed), ensure 2D contiguous
+            B_2d = B.reshape(shapeB[0], -1).contiguous()
+            out = gemm_4bit_forward_kernel(A, B_2d, absmax, blocksize, quant_type_num)
+        else:
+            # AVX2 回退：用所有 CPU 构建都导出的 cgemv_4bit_inference_cpu_* 融合内核
+            # （lib.gemv_4bit_inference_cpu_fp4/nf4_bf16 只存在于 AVX512+BF16 构建，
+            #  AVX2 机器走到会 AttributeError；data_type: FP4=1 / NF4=2）
+            A_c = A_orig.reshape(-1, A_orig.shape[-1]).contiguous()
+            B_c = B.reshape(shapeB[0], -1).contiguous()
+            absmax_c = absmax_orig.float().contiguous()
+            data_type_num = 1 if quant_type == "fp4" else 2
+            M, K = A_c.shape[0], A_c.shape[1]
+            N = shapeB[0]
+            out = torch.empty(out_shape, dtype=A_c.dtype, device=A_c.device)
+            if dtype == torch.float32:
+                lib.cgemv_4bit_inference_cpu_fp32(
+                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
+                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
+                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
+                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
+                    ct.c_int(data_type_num),
+                )
+            elif dtype == torch.bfloat16:
+                lib.cgemv_4bit_inference_cpu_bf16(
+                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
+                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
+                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
+                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
+                    ct.c_int(data_type_num),
+                )
+            else:
+                lib.cgemv_4bit_inference_cpu_fp16(
+                    get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
+                    ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
+                    ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
+                    ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
+                    ct.c_int(data_type_num),
+                )
+
+        if dtype != torch.bfloat16:
+            out = out.to(dtype)
+
+        return out.reshape(final_out_shape)
+
+    # ---- gemm_4bit: CPU 注册（2026-09-22 新增）----
+    # 此前这个 op 【没有 cpu 注册】，于是退到 default 后端做"反量化 + 稠密线性"。
+    # autograd/_functions.py:324/334/462/472 调的就是它 ⇒ 真实模型里
+    # 【每一次批量 4-bit 线性前向】都在走慢路径。实测（K=N=4096）：
+    #     M=1   gemm_4bit 14.483 ms  vs 融合内核  1.448 ms  ⇒ 10.00x
+    #     M=8   gemm_4bit 32.494 ms  vs 融合内核  4.124 ms  ⇒  7.88x
+    #     M=64  gemm_4bit 38.576 ms  vs 融合内核 30.382 ms  ⇒  1.27x
+    #     M=256 gemm_4bit 75.532 ms  vs 融合内核 164.840 ms ⇒  0.46x  ← 融合反而慢
+    # 融合内核的 m 循环每 4 行重读一遍权重，所以大 M 时它输给 oneDNN 稠密。
+    # ⇒ 按 M 分档：小 M 走融合内核，大 M 交回 default。
+    # ---- 2026-09-22 实测三方交叉点（K=N=4096，OMP=6）----
+    #        M     融合 ms    default ms   融合优势
+    #        1      1.20       ~14.5       12x
+    #        8      4.71        25.14       5.3x
+    #       16      7.04        25.75       3.7x
+    #       17      7.79        26.43       3.4x
+    #       32     16.20        28.41       1.75x
+    #       64     28.79        31.37       1.09x
+    #       96     43.45        37.01       default 反超
+    #      128     62.11        41.77       default 快 1.49x
+    #      192     85.78        54.67       default 快 1.57x
+    #      256    166.11        64.94       default 快 2.56x
+    # ⇒ 融合内核一直赢到 M=64，default 要到 M≈96 才反超。
+    # 阈值原先写 16 ⇒ M=17~64 被错误地分给 default，白慢最多 3.4x。
+    #
+    # ⚠️ 一个被实测否决的尝试：给 C 层加"权重行解码一次、对全部 M 行复用"的路径，
+    #    结果是【负优化】（M=17 慢 1.52x、M=192 慢 1.44x）——
+    #    它把权重解码从 O(M/4) 降到 O(1)，却让 A 被每个输出列重读一次，
+    #    A 流量变成 O(N*M*K)（M=17 时 1.14 GB），远比重新解码权重贵。
+    #    已删除，只在 cpu_ops.cpp 里留了说明。
+    _GEMM_4BIT_FUSED_MAX_M = 64
+
+    @register_kernel("bitsandbytes::gemm_4bit", "cpu")
+    def _gemm_4bit_cpu(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        shapeB: Sequence[int],
+        absmax: torch.Tensor,
+        blocksize: int,
+        quant_type: str,
+        bias: Optional[torch.Tensor] = None,
+        absmax_8bit: Optional[torch.Tensor] = None,
+        absmax_code: Optional[torch.Tensor] = None,
+        absmax_offset: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        m_rows = A.numel() // A.shape[-1] if A.numel() else 0
+        # 融合路径的限制：无 bias、非 nested absmax、张量是 2D 可 reshape 的、
+        # 且 M 足够小（大 M 时 oneDNN 稠密更快）。
+        can_fuse = (
+            gemm_4bit_forward_kernel is None      # 没有 HF kernel 时才用我们的融合内核
+            and bias is None
+            and absmax_8bit is None
+            and m_rows > 0
+            and m_rows <= _GEMM_4BIT_FUSED_MAX_M
+        )
+        if can_fuse:
+            dtype = A.dtype
+            A_c = A.reshape(-1, A.shape[-1]).contiguous()
+            B_c = B.reshape(shapeB[0], -1)
+            if B_c.dtype != torch.uint8:
+                B_c = B_c.contiguous().view(torch.uint8)
+            else:
+                B_c = B_c.contiguous()
+            absmax_c = absmax.float().contiguous()
+            M, K = A_c.shape[0], A_c.shape[1]
+            N = shapeB[0]
+            data_type_num = 1 if quant_type == "fp4" else 2
+            out = torch.empty((M, N), dtype=dtype, device=A_c.device)
+            fn = {
+                torch.float32: lib.cgemv_4bit_inference_cpu_fp32,
+                torch.bfloat16: lib.cgemv_4bit_inference_cpu_bf16,
+                torch.float16: lib.cgemv_4bit_inference_cpu_fp16,
+            }[dtype]
+            fn(
+                get_ptr(A_c), get_ptr(B_c), get_ptr(absmax_c), get_ptr(out),
+                ct.c_longlong(M), ct.c_longlong(N), ct.c_longlong(K),
+                ct.c_longlong(A_c.stride(0)), ct.c_longlong(K // 2),
+                ct.c_longlong(out.stride(0)), ct.c_longlong(blocksize),
+                ct.c_int(data_type_num),
+            )
+            return out.reshape(*A.shape[:-1], N)
+
+        from ..default.ops import _gemm_4bit_default_impl
+
+        return _gemm_4bit_default_impl(
+            A, B, shapeB, absmax, blocksize, quant_type,
+            bias, absmax_8bit, absmax_code, absmax_offset,
+        )
+
+    # ---- 大 M 分档（2026-09-23 实测）----
+    # 融合内核的内层是 m0 += 4，每 4 行就把整块权重重新读一遍 ⇒ 权重流量 O(M/4)，
+    # 大 M 时输给"反量化 + oneDNN 稠密 matmul"。
+    # 干净基准（gemm8_test.py：每格子进程 + 取 7 次最小值），融合/两步 的比值：
+    #     M=16   0.18x / 0.15x / 0.25x     融合远优
+    #     M=64   0.89x / 0.65x / 0.77x     融合仍优
+    #     M=128  1.34x / 1.14x / 1.32x     ← 交叉点在这一带
+    #     M=256  1.54x / 1.34x / 1.51x
+    #     M=512  2.04x / 1.82x / 2.04x     两步快约一倍
+    # ⇒ 阈值取 64（保守：M=64 时融合还是赢的，且真实推理的 M 是 seq_len）。
+    #
+    # ⚠️ 这个 op 【没有 default 后端回退】（全库只注册了 cpu 与 cuda），
+    #    所以大 M 路径必须在这里自己实现。8-bit 的码表是 create_dynamic_map()，
+    #    与 quantize_blockwise 默认用的同一个 ⇒ 数值一致。
+    _GEMM_8BIT_FUSED_MAX_M = 64
+
+    @register_kernel("bitsandbytes::gemm_8bit", "cpu")
+    def _gemm_8bit_cpu(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        absmax: torch.Tensor,
+        blocksize: int,
+    ) -> torch.Tensor:
+        """融合 8bit 反量化 GEMM：out = A @ dequant8(B)（B 为 uint8 码流，逐行块量化）。
+
+        权重在整个乘法中保持 uint8（DRAM 流量为 fp32 的 1/4），无 fp32 临时张量。
+        小 M（<= 64）走融合内核；大 M 退到"反量化 + 稠密 matmul"（见上面的分档说明）。
+        """
+        if A.numel() == 0:
+            return torch.empty((*A.shape[:-1], B.shape[0]), dtype=torch.float32, device=A.device)
+        A_f = A.reshape(-1, A.shape[-1]).contiguous()
+        B_c = B.contiguous()
+        absmax_c = absmax.contiguous().float()
+        M, K = A_f.shape
+        N = B_c.shape[0]
+
+        if M > _GEMM_8BIT_FUSED_MAX_M:
+            # 大 M：两步更快（实测 M=512 时 1.8~2.0x）
+            from bitsandbytes.functional import create_dynamic_map, dequantize_blockwise
+
+            code = create_dynamic_map().to(B_c.device)
+            W = dequantize_blockwise(B_c, absmax=absmax_c, code=code, blocksize=blocksize)
+            return (A_f @ W.t()).reshape(*A.shape[:-1], N)
+
+        out = torch.empty((M, N), dtype=torch.float32, device=A_f.device)
+        lib.cgemm_8bit_inference_cpu_fp32(
+            get_ptr(A_f),
+            get_ptr(B_c),
+            get_ptr(absmax_c),
+            get_ptr(out),
+            ct.c_longlong(M),
+            ct.c_longlong(N),
+            ct.c_longlong(K),
+            ct.c_longlong(A_f.stride(0)),
+            ct.c_longlong(B_c.stride(0)),
+            ct.c_longlong(out.stride(0)),
+            ct.c_longlong(blocksize),
+        )
+        return out.reshape(*A.shape[:-1], N)
+
+
+# ==================== CPU Optimizer Kernels ====================
+
+
+def _compute_update_norm_and_scale(
+    update: torch.Tensor,
+    unorm_vec: Optional[torch.Tensor],
+    max_unorm: float,
+    param_norm: float,
+) -> float:
+    """Compute trust-ratio scaling factor for LAMB/LARS and store update norm."""
+    if max_unorm <= 0.0:
+        return 1.0
+    unorm = torch.norm(update).item()
+    if unorm_vec is not None:
+        unorm_vec.fill_(unorm)
+    if unorm > max_unorm * param_norm:
+        return (max_unorm * param_norm) / unorm
+    return 1.0
+
+
+@torch.no_grad()
+def _optimizer_update_32bit_cpu(
+    optimizer_name: str,
+    g: torch.Tensor,
+    p: torch.Tensor,
+    state1: torch.Tensor,
+    state2: Optional[torch.Tensor],
+    unorm_vec: Optional[torch.Tensor],
+    max_unorm: float,
+    param_norm: float,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    alpha: float,
+    eps: float,
+    weight_decay: float,
+    step: int,
+    lr: float,
+    gnorm_scale: float,
+    skip_zeros: bool = False,
+) -> None:
+    g_float = g.float() * gnorm_scale
+    p_float = p.data.float()
+
+    if optimizer_name in ("adam", "lamb"):
+        # Adam / LAMB (2-state): m and v
+        state1.mul_(beta1).add_(g_float, alpha=1.0 - beta1)
+        state2.mul_(beta2).addcmul_(g_float, g_float, value=1.0 - beta2)
+
+        correction1 = 1.0 - beta1**step
+        correction2 = math.sqrt(1.0 - beta2**step)
+        step_size = -lr * correction2 / correction1
+
+        if weight_decay > 0.0:
+            p_float.mul_(1.0 - lr * weight_decay)
+
+        update = state1 / (state2.sqrt() + eps * correction2)
+
+        update_scale = _compute_update_norm_and_scale(update, unorm_vec, max_unorm, param_norm)
+        p_float.add_(update, alpha=step_size * update_scale)
+
+    elif optimizer_name == "ademamix":
+        # AdEMAMix (2-state): state1 shape is (2, *p.shape), state1[0]=m1, state1[1]=m2
+        m1 = state1[0]
+        m2 = state1[1]
+        nu = state2
+
+        m1.mul_(beta1).add_(g_float, alpha=1.0 - beta1)
+        m2.mul_(beta3).add_(g_float, alpha=1.0 - beta3)
+        nu.mul_(beta2).addcmul_(g_float, g_float, value=1.0 - beta2)
+
+        correction1 = 1.0 - beta1**step
+        correction2 = math.sqrt(1.0 - beta2**step)
+
+        if weight_decay > 0.0:
+            p_float.mul_(1.0 - lr * weight_decay)
+
+        mixed_momentum = (m1 / correction1) + (alpha * m2)
+        adaptive_term = (nu.sqrt() / correction2) + eps
+        p_float.add_(mixed_momentum / adaptive_term, alpha=-lr)
+
+    elif optimizer_name in ("momentum", "lars"):
+        # SGD with momentum / LARS (1-state)
+        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
+
+        if step == 1:
+            state1.copy_(g_wd)
+        else:
+            state1.mul_(beta1).add_(g_wd)
+
+        update_scale = _compute_update_norm_and_scale(state1, unorm_vec, max_unorm, param_norm)
+        p_float.add_(state1, alpha=-lr * update_scale)
+
+    elif optimizer_name == "lion":
+        # Lion (2-state sign update)
+        if weight_decay > 0.0:
+            p_float.mul_(1.0 - lr * weight_decay)
+
+        update = state1.mul(beta1).add(g_float, alpha=1.0 - beta1)
+        p_float.add_(update.sign(), alpha=-lr)
+
+        state1.mul_(beta2).add_(g_float, alpha=1.0 - beta2)
+
+    elif optimizer_name == "rmsprop":
+        # RMSprop (1-state)
+        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
+        state1.mul_(beta1).addcmul_(g_wd, g_wd, value=1.0 - beta1)
+
+        update = g_wd / (state1.sqrt() + eps)
+        update_scale = _compute_update_norm_and_scale(update, unorm_vec, max_unorm, param_norm)
+        p_float.add_(update, alpha=-lr * update_scale)
+
+    elif optimizer_name == "adagrad":
+        # Adagrad (1-state)
+        g_wd = g_float.add(p_float, alpha=weight_decay) if weight_decay > 0.0 else g_float
+        state1.addcmul_(g_wd, g_wd, value=1.0)
+
+        update = g_wd / (state1.sqrt() + eps)
+        p_float.add_(update, alpha=-lr)
+
+    else:
+        raise ValueError(f"Unsupported optimizer for CPU: {optimizer_name}")
+
+    # Write back to original precision
+    p.data.copy_(p_float)
+
+
+register_kernel("bitsandbytes::optimizer_update_32bit", "cpu")(_optimizer_update_32bit_cpu)
+
+
+@torch.no_grad()
+def _dequant_blockwise_fp32_direct(
+    A_uint8: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int
+) -> torch.Tensor:
+    return torch.ops.bitsandbytes.dequantize_blockwise(A_uint8, absmax, code, blocksize, torch.float32)
+
+
+def _quant_blockwise_fp32_direct(
+    A_fp32: torch.Tensor, code: torch.Tensor, absmax_out: torch.Tensor, out_uint8: torch.Tensor, blocksize: int
+) -> None:
+    out, absmax = torch.ops.bitsandbytes.quantize_blockwise(A_fp32, code, blocksize)
+    out_uint8.copy_(out)
+    absmax_out.copy_(absmax)
+
+
+def _optimizer_update_8bit_blockwise_cpu(
+    optimizer_name: str,
+    g: torch.Tensor,
+    p: torch.Tensor,
+    state1: torch.Tensor,
+    state2: Optional[torch.Tensor],
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    alpha: float,
+    eps: float,
+    step: int,
+    lr: float,
+    qmap1: torch.Tensor,
+    qmap2: Optional[torch.Tensor],
+    absmax1: torch.Tensor,
+    absmax2: Optional[torch.Tensor],
+    weight_decay: float,
+    gnorm_scale: float,
+    skip_zeros: bool = False,
+) -> None:
+    # Fused single-pass C kernel (dequant + update + p update + requant in one
+    # sweep). Falls back to the composite PyTorch path below when the native
+    # library is unavailable or the optimizer is not supported by the kernel.
+    _OPT_IDS = {"adam": 0, "lamb": 0, "momentum": 1, "lars": 1, "lion": 2, "rmsprop": 3, "adagrad": 4, "ademamix": 5}
+    n = p.numel()
+    # The C kernel walks raw pointers and assumes C-contiguous layouts: a
+    # non-contiguous grad (transposed / sliced views are common in training)
+    # would silently read the wrong elements and corrupt the states.
+    contiguous_layout = (
+        g.is_contiguous()
+        and p.is_contiguous()
+        and state1.is_contiguous()
+        and (state2 is None or state2.is_contiguous())
+        and qmap1.is_contiguous()
+        and (qmap2 is None or qmap2.is_contiguous())
+        and absmax1.is_contiguous()
+        and (absmax2 is None or absmax2.is_contiguous())
+    )
+    if (
+        not isinstance(lib, ErrorHandlerMockBNBNativeLibrary)
+        and contiguous_layout
+        and optimizer_name in _OPT_IDS
+        and g.dtype in (torch.float32, torch.bfloat16, torch.float16)
+        # ademamix packs a second momentum into state1[n:2n] with its absmax in
+        # absmax1[blocks:2*blocks]; the flat C layout requires that packing.
+        and (optimizer_name != "ademamix" or state1.numel() >= 2 * n)
+    ):
+        try:
+            dtype_id = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}[g.dtype]
+            # Every scalar needs an explicit ctypes width: a bare Python int is
+            # marshalled as 32-bit c_int, but `n` is a long long. libffi writes
+            # only the low 4 bytes of its 8-byte stack slot; the C side reads
+            # all 8, so stale stack bits leak into n (observed: 0x7FFF00004000
+            # for n=16384) and the kernel walks out of bounds -> SIGSEGV.
+            lib.coptimizer_update_8bit_blockwise_cpu(
+                ct.c_int(_OPT_IDS[optimizer_name]),
+                get_ptr(g),
+                get_ptr(p),
+                get_ptr(state1),
+                get_ptr(state2) if state2 is not None else ct.c_void_p(0),
+                ct.c_float(beta1),
+                ct.c_float(beta2),
+                ct.c_float(beta3),
+                ct.c_float(alpha),
+                ct.c_float(eps),
+                ct.c_int(step),
+                ct.c_float(lr),
+                get_ptr(qmap1),
+                get_ptr(qmap2) if qmap2 is not None else ct.c_void_p(0),
+                get_ptr(absmax1),
+                get_ptr(absmax2) if absmax2 is not None else ct.c_void_p(0),
+                ct.c_float(weight_decay),
+                ct.c_float(gnorm_scale),
+                ct.c_int(1 if skip_zeros else 0),
+                ct.c_longlong(n),
+                ct.c_int(dtype_id),
+            )
+            return
+        except (AttributeError, RuntimeError):
+            pass  # old native library without the fused entry point
+
+    blocksize = 256
+
+    # Dequantize states
+    if optimizer_name == "ademamix" and absmax1.ndim == 2:
+        s1_1 = _dequant_blockwise_fp32_direct(state1[0], absmax1[0], qmap1, blocksize)
+        s1_2 = _dequant_blockwise_fp32_direct(state1[1], absmax1[1], qmap1, blocksize)
+        state1_fp32 = torch.stack([s1_1, s1_2])
+    else:
+        state1_fp32 = _dequant_blockwise_fp32_direct(state1, absmax1, qmap1, blocksize)
+
+    state2_fp32 = None
+    if state2 is not None and qmap2 is not None and absmax2 is not None:
+        state2_fp32 = _dequant_blockwise_fp32_direct(state2, absmax2, qmap2, blocksize)
+
+    grad = g.float() * gnorm_scale
+    p_fp32 = p.data.float()
+
+    if optimizer_name in ("adam", "lamb"):
+        state1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+        state2_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+        correction1 = 1.0 - beta1**step
+        correction2 = math.sqrt(1.0 - beta2**step)
+
+        denom = (state2_fp32.sqrt() / correction2).add_(eps)
+        if weight_decay > 0.0:
+            p_fp32.mul_(1.0 - lr * weight_decay)
+        p_fp32.addcdiv_(state1_fp32, denom, value=-lr / correction1)
+
+    elif optimizer_name == "ademamix":
+        m1_fp32, m2_fp32 = state1_fp32[0], state1_fp32[1]
+        nu_fp32 = state2_fp32
+
+        m1_fp32.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+        m2_fp32.mul_(beta3).add_(grad, alpha=1.0 - beta3)
+        nu_fp32.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+        correction1 = 1.0 - beta1**step
+        correction2 = math.sqrt(1.0 - beta2**step)
+
+        update = (m1_fp32 / correction1 + alpha * m2_fp32) / (nu_fp32.sqrt() / correction2 + eps)
+        if weight_decay > 0.0:
+            p_fp32.mul_(1.0 - lr * weight_decay)
+        p_fp32.add_(update, alpha=-lr)
+
+        state1_fp32 = torch.stack([m1_fp32, m2_fp32])
+
+    elif optimizer_name in ("momentum", "lars"):
+        grad.add_(p_fp32, alpha=weight_decay)
+        if step == 1:
+            state1_fp32.copy_(grad)
+        else:
+            state1_fp32.mul_(beta1).add_(grad)
+        p_fp32.add_(state1_fp32, alpha=-lr)
+
+    elif optimizer_name == "lion":
+        if weight_decay > 0.0:
+            p_fp32.mul_(1.0 - lr * weight_decay)
+
+        update_dir = torch.sign(state1_fp32.mul(beta1) + grad.mul(1.0 - beta1))
+        p_fp32.add_(update_dir, alpha=-lr)
+
+        state1_fp32.mul_(beta2).add_(grad, alpha=1.0 - beta2)
+
+    elif optimizer_name == "rmsprop":
+        grad.add_(p_fp32, alpha=weight_decay)
+        state1_fp32.mul_(beta1).addcmul_(grad, grad, value=1.0 - beta1)
+        p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr)
+
+    elif optimizer_name == "adagrad":
+        grad.add_(p_fp32, alpha=weight_decay)
+        state1_fp32.addcmul_(grad, grad, value=1.0)
+        p_fp32.addcdiv_(grad, state1_fp32.sqrt().add_(eps), value=-lr)
+
+    else:
+        raise ValueError(f"Unsupported optimizer for CPU 8-bit: {optimizer_name}")
+
+    p.data.copy_(p_fp32)
+
+    # Re-quantize states
+    if optimizer_name == "ademamix":
+        _quant_blockwise_fp32_direct(state1_fp32[0], qmap1, absmax1[0], state1[0], blocksize)
+        _quant_blockwise_fp32_direct(state1_fp32[1], qmap1, absmax1[1], state1[1], blocksize)
+        _quant_blockwise_fp32_direct(state2_fp32, qmap2, absmax2, state2, blocksize)
+    else:
+        _quant_blockwise_fp32_direct(state1_fp32, qmap1, absmax1, state1, blocksize)
+        if state2_fp32 is not None:
+            _quant_blockwise_fp32_direct(state2_fp32, qmap2, absmax2, state2, blocksize)
+
+
+register_kernel("bitsandbytes::optimizer_update_8bit_blockwise", "cpu")(_optimizer_update_8bit_blockwise_cpu)