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

bitsandbytes/backends/triton/kernels_4bit.py

@@ -1,577 +1,583 @@
-import torch
-
-import triton
-import triton.language as tl
-
-
-# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeFP4
-# @triton.autotune(
-#     configs=[
-#         triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 1}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 4}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 8}),
-#     ],
-#     key=["n_elements"],
-# )
-@triton.jit
-def quantize_fp4_blockwise_kernel(
-    A_ptr,
-    absmax_ptr,
-    out_ptr,
-    n_elements,
-    BLOCK_SIZE: tl.constexpr,
-    SPLIT_NUM_BLOCKS: tl.constexpr,
-):
-    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
-    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
-    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-
-    offsets = block_start_idx * BLOCK_SIZE + thread_idx
-    mask = offsets < n_elements
-
-    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
-
-    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
-    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
-
-    # Calculating absamax for each block
-    absmax = tl.max(tl.abs(A_reshaped), axis=1)
-    tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax)
-
-    A_normalized = A_reshaped / absmax[:, None]
-    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
-
-    sign = tl.where(A_normalized < 0, 0b1000, 0b0000)
-    A_absf = tl.abs(A_normalized)
-
-    result = tl.where(
-        A_absf > 0.29166667,
-        tl.where(
-            A_absf > 0.583333, tl.where(A_absf > 0.8333333, 0b011, 0b010), tl.where(A_absf > 0.4166667, 0b101, 0b100)
-        ),
-        tl.where(
-            A_absf > 0.0859375,
-            tl.where(A_absf > 0.20833333, 0b0111, 0b0110),
-            tl.where(A_absf > 0.00260417, 0b0001, 0b0000),
-        ),
-    )
-    quantized = (result ^ sign).to(tl.uint8)
-
-    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
-    left, right = quantized.split()
-    packed = left << 4 | (right & 0xF)
-
-    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
-    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-    # Use n - n//2 instead of (n+1)//2 to avoid integer overflow for large n
-    out_mask = out_offsets < (n_elements - n_elements // 2)
-    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)
-
-
-# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeNF4
-# @triton.autotune(
-#     configs=[
-#         triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 1}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 4}),
-#         triton.Config({"SPLIT_NUM_BLOCKS": 8}),
-#     ],
-#     key=["n_elements"],
-# )
-@triton.jit
-def quantize_nf4_blockwise_kernel(
-    A_ptr,
-    absmax_ptr,
-    out_ptr,
-    n_elements,
-    BLOCK_SIZE: tl.constexpr,
-    SPLIT_NUM_BLOCKS: tl.constexpr,
-):
-    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
-    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
-    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-
-    offsets = block_start_idx * BLOCK_SIZE + thread_idx
-    mask = offsets < n_elements
-
-    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
-
-    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
-    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
-
-    # Calculating absamax for each block
-    absmax = tl.max(tl.abs(A_reshaped), axis=1)
-    tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax)
-
-    A_normalized = A_reshaped / absmax[:, None]
-    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
-
-    result = tl.where(
-        A_normalized > 0.03979014977812767,
-        tl.where(
-            A_normalized > 0.3893125355243683,
-            tl.where(
-                A_normalized > 0.6427869200706482,
-                tl.where(A_normalized > 0.8614784181118011, 0b1111, 0b1110),
-                tl.where(A_normalized > 0.5016634166240692, 0b1101, 0b1100),
-            ),
-            tl.where(
-                A_normalized > 0.2035212516784668,
-                tl.where(A_normalized > 0.2920137718319893, 0b1011, 0b1010),
-                tl.where(A_normalized > 0.1202552504837513, 0b1001, 0b1000),
-            ),
-        ),
-        tl.where(
-            A_normalized > -0.33967943489551544,
-            tl.where(
-                A_normalized > -0.13791173323988914,
-                tl.where(A_normalized > -0.045525018125772476, 0b0111, 0b0110),
-                tl.where(A_normalized > -0.23460740596055984, 0b0101, 0b0100),
-            ),
-            tl.where(
-                A_normalized > -0.6106329262256622,
-                tl.where(A_normalized > -0.4599952697753906, 0b0011, 0b0010),
-                tl.where(A_normalized > -0.8480964004993439, 0b0001, 0b0000),
-            ),
-        ),
-    )
-    quantized = result.to(tl.uint8)
-
-    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
-
-    left, right = quantized.split()
-    packed = left << 4 | (right & 0xF)
-
-    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
-    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-    # Use n - n//2 instead of (n+1)//2 to avoid integer overflow for large n
-    out_mask = out_offsets < (n_elements - n_elements // 2)
-    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)
-
-
-def quantize_4bit_blockwise_triton(A, blocksize, quant_type, blocks, absmax, num_elements, quantized_out):
-    # grid = lambda META: (triton.cdiv(blocks, META["SPLIT_NUM_BLOCKS"]),)
-    split_num_blocks = 4
-    grid = (triton.cdiv(blocks, split_num_blocks),)
-    if quant_type == "fp4":
-        quantize_fp4_blockwise_kernel[grid](
-            A_ptr=A,
-            absmax_ptr=absmax,
-            out_ptr=quantized_out,
-            n_elements=num_elements,
-            BLOCK_SIZE=blocksize,
-            SPLIT_NUM_BLOCKS=split_num_blocks,
-        )
-    else:
-        quantize_nf4_blockwise_kernel[grid](
-            A_ptr=A,
-            absmax_ptr=absmax,
-            out_ptr=quantized_out,
-            n_elements=num_elements,
-            BLOCK_SIZE=blocksize,
-            SPLIT_NUM_BLOCKS=split_num_blocks,
-        )
-    return quantized_out, absmax
-
-
-@triton.jit
-def dequant_4bit_body_util(a, offsets, quant_ptr, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
-    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
-    mask = offsets < n_elems
-    higher = a & 0xF
-    # lower 4bits
-    lower = a >> 4
-
-    abs_offsets = offsets // PAIRED_QUANT_BLOCK
-    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
-
-    # apply conversion
-    lower_4 = tl.load(quant_ptr + lower, eviction_policy="evict_last")
-    higher_4 = tl.load(quant_ptr + higher, eviction_policy="evict_last")
-
-    mul_high = higher_4 * absmax
-    mul_low = lower_4 * absmax
-    out_dq = tl.interleave(mul_low, mul_high)
-    return out_dq
-
-
-# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeFP4Tree
-@triton.jit
-def dequantize_fp4_tree(val, absmax):
-    # val: tl.tensor (uint8)
-    # absmax: tl.tensor (float32/float16)
-    #  00001100  00001011  00001001  00001111
-    sign = tl.where((val & 0b1000) == 0b1000, -1.0, 1.0)  # -1
-    third_bit = (val & 0b0100) == 0b0100  # True
-    second_bit = (val & 0b0010) == 0b0010  # False
-    first_bit = (val & 0b0001) == 0b0001  # False
-
-    branch1 = tl.where(
-        second_bit,
-        tl.where(first_bit, 0.25, 0.16666667),  # 1111, 1110
-        tl.where(first_bit, 0.5, 0.33333333),  # 1101, 1100
-    )
-    branch2 = tl.where(
-        second_bit,
-        tl.where(first_bit, 1.0, 0.66666667),  # 1011, 1010
-        tl.where(first_bit, 0.005208333333, 0.0),  # 1001, 1000
-    )
-    out = tl.where(third_bit, branch1, branch2)
-    return out * sign * absmax
-
-
-@triton.jit
-def dequant_fp4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
-    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
-    mask = offsets < n_elems
-    higher = a & 0xF
-    lower = a >> 4
-
-    abs_offsets = offsets // PAIRED_QUANT_BLOCK
-    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
-    mul_high = dequantize_fp4_tree(higher, absmax)
-    mul_low = dequantize_fp4_tree(lower, absmax)
-    out_dq = tl.interleave(mul_low, mul_high)
-    return out_dq
-
-
-# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeNF4
-@triton.jit
-def dequantize_nf4_tree(val):
-    # val: tl.tensor (uint8)
-    cond0 = (val & 0b1000) == 0b1000
-    cond1 = (val & 0b0100) == 0b0100
-    cond2 = (val & 0b0010) == 0b0010
-    cond3 = (val & 0b0001) == 0b0001
-
-    # Positive branch (val & 0b1000) == 8
-    branch_pos = tl.where(
-        cond1,
-        tl.where(
-            cond2,
-            tl.where(cond3, 1.0, 0.7229568362236023),  # 1111, 1110
-            tl.where(cond3, 0.5626170039176941, 0.44070982933044434),  # 1101, 1100
-        ),
-        tl.where(
-            cond2,
-            tl.where(cond3, 0.33791524171829224, 0.24611230194568634),  # 1011, 1010
-            tl.where(cond3, 0.16093020141124725, 0.07958029955625534),  # 1001, 1000
-        ),
-    )
-
-    # Negative branch (val & 0b1000) == 0
-    branch_neg = tl.where(
-        cond1,
-        tl.where(
-            cond2,
-            tl.where(cond3, 0.0, -0.09105003625154495),  # 0111, 0110
-            tl.where(cond3, -0.18477343022823334, -0.28444138169288635),  # 0101, 0100
-        ),
-        tl.where(
-            cond2,
-            tl.where(cond3, -0.39491748809814453, -0.5250730514526367),  # 0011, 0010
-            tl.where(cond3, -0.6961928009986877, -1.0),  # 0001, 0000
-        ),
-    )
-    return tl.where(cond0, branch_pos, branch_neg)
-
-
-@triton.jit
-def dequant_nf4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
-    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
-    mask = offsets < n_elems
-    higher = a & 0xF
-    # lower 4bits
-    lower = a >> 4
-
-    abs_offsets = offsets // PAIRED_QUANT_BLOCK
-    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
-    mul_high = dequantize_nf4_tree(higher) * absmax
-    mul_low = dequantize_nf4_tree(lower) * absmax
-    out_dq = tl.interleave(mul_low, mul_high)
-    return out_dq
-
-
-# All such kernels are similar, so maybe code can be generalised.
-# @triton.autotune(
-#     configs=[
-# #         # triton.Config({'SPLIT_SIZE': 64}),
-# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
-#         triton.Config({'SPLIT_SIZE': 128}),
-#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
-# #         # triton.Config({'SPLIT_SIZE': 128}, num_warps = 4, num_stages = 4),
-# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
-#         triton.Config({'SPLIT_SIZE': 256}),
-#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
-#         # triton.Config({'SPLIT_SIZE': 256}, num_warps = 4, num_stages = 4),
-#         triton.Config({'SPLIT_SIZE': 512}),
-#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
-#         # triton.Config({'SPLIT_SIZE': 512}, num_warps = 4, num_stages = 4),
-# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
-# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
-# #         # triton.Config({'SPLIT_SIZE': 1024}),
-# #         # # triton.Config({'SPLIT_SIZE': 2048}),
-# #         # # triton.Config({'SPLIT_SIZE': 4096}),
-# #         # # triton.Config({'SPLIT_SIZE': 8192}),
-# #         # # triton.Config({'SPLIT_SIZE': 16384}),
-#     ],
-#     key=['num_paired_elements'],
-# )
-@triton.jit
-def dequant_4bit_kernel(
-    a_ptr,
-    c_ptr,
-    quant_ptr,
-    absmax_ptr,
-    num_paired_elements,
-    num_output_elements,
-    QUANT_BLOCK: tl.constexpr,
-    SPLIT_SIZE: tl.constexpr,
-):
-    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    block_start = pid * SPLIT_SIZE
-    offsets = block_start + tl.arange(0, SPLIT_SIZE)
-    mask = offsets < num_paired_elements
-
-    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
-
-    out_dq = dequant_4bit_body_util(
-        a=a,
-        offsets=offsets,
-        quant_ptr=quant_ptr,
-        absmax_ptr=absmax_ptr,
-        n_elems=num_paired_elements,
-        QUANT_BLOCK=QUANT_BLOCK,
-    )
-
-    out_block_start = pid * SPLIT_SIZE * 2
-    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
-    mask = offs < num_output_elements
-    tl.store(c_ptr + offs, out_dq, mask)
-
-
-# @triton.autotune(
-#     configs=[
-#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 256}),
-#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 512}),
-#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2),
-#     ],
-#     key=['num_paired_elements'],
-# )
-@triton.jit
-def dequant_fp4_kernel(
-    a_ptr,
-    c_ptr,
-    absmax_ptr,
-    num_paired_elements,
-    num_output_elements,
-    QUANT_BLOCK: tl.constexpr,
-    SPLIT_SIZE: tl.constexpr,
-):
-    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    block_start = pid * SPLIT_SIZE
-    offsets = block_start + tl.arange(0, SPLIT_SIZE)
-    mask = offsets < num_paired_elements
-
-    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
-
-    out_dq = dequant_fp4_body_util(
-        a=a,
-        offsets=offsets,
-        absmax_ptr=absmax_ptr,
-        n_elems=num_paired_elements,
-        QUANT_BLOCK=QUANT_BLOCK,
-    )
-
-    out_block_start = pid * SPLIT_SIZE * 2
-    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
-    mask = offs < num_output_elements
-    tl.store(c_ptr + offs, out_dq, mask)
-
-
-# @triton.autotune(
-#     configs=[
-#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 256}),
-#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 512}),
-#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
-#         triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2),
-#     ],
-#     key=['num_paired_elements'],
-# )
-@triton.jit
-def dequant_nf4_kernel(
-    a_ptr,
-    c_ptr,
-    absmax_ptr,
-    num_paired_elements,
-    num_output_elements,
-    QUANT_BLOCK: tl.constexpr,
-    SPLIT_SIZE: tl.constexpr,
-):
-    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
-    block_start = pid * SPLIT_SIZE
-    offsets = block_start + tl.arange(0, SPLIT_SIZE)
-    mask = offsets < num_paired_elements
-
-    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
-
-    out_dq = dequant_nf4_body_util(
-        a=a,
-        offsets=offsets,
-        absmax_ptr=absmax_ptr,
-        n_elems=num_paired_elements,
-        QUANT_BLOCK=QUANT_BLOCK,
-    )
-
-    out_block_start = pid * SPLIT_SIZE * 2
-    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
-    mask = offs < num_output_elements
-    tl.store(c_ptr + offs, out_dq, mask)
-
-
-def dequantize_4bit_impl(
-    A: torch.Tensor,
-    absmax: torch.Tensor,
-    blocksize: int,
-    quant_type: str,
-    dtype: torch.dtype,
-    out: torch.Tensor,
-) -> None:
-    # It's will be processed as an array, so
-    # actual length is row * col
-    # Elements are in uint8 format, so interleaved
-    # so total amount of data is 2 * elem_count
-    number_of_paired_elements = A.numel()
-    num_output_elements = out.numel()
-    # we assume that split_size > quant_blocksize
-
-    SPLIT_SIZE = 256
-    # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), )
-    grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),)
-    if quant_type == "fp4":
-        dequant_fp4_kernel[grid](A, out, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE)
-    else:
-        dequant_nf4_kernel[grid](A, out, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE)
-
-
-def dequantize_4bit_impl_passing_code(
-    A: torch.Tensor,
-    absmax: torch.Tensor,
-    blocksize: int,
-    code: torch.Tensor,
-    dtype: torch.dtype,
-    out: torch.Tensor,
-) -> None:
-    number_of_paired_elements = A.numel()
-    num_output_elements = out.numel()
-    # we assume that split_size > quant_blocksize
-
-    SPLIT_SIZE = 256
-    # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), )
-    grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),)
-    dequant_4bit_kernel[grid](
-        A, out, code, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE
-    )
-
-
-######################### Fallback dequantization functions #########################
-## for debug ##
-
-
-# @triton.autotune(
-#     configs=[
-#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
-#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
-#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
-#         # #
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         #
-#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "large"}, num_stages=2, num_warps=32),
-#         # # triton.Config({'SPLIT_NUM_BLOCKS': 2, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=2, num_warps=32),
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=2, num_warps=32),
-#         # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=4, num_warps=32),
-#         # triton.Config({'SPLIT_NUM_BLOCKS': 8, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
-#     ],
-#     key=["n_elements", "BLOCK_SIZE"],
-# )
-@triton.jit
-def quantize_4bit_blockwise_kernel(
-    A_ptr,
-    code_ptr,
-    absmax_ptr,
-    out_ptr,
-    n_elements,
-    BLOCK_SIZE: tl.constexpr,
-    CODE_SIZE: tl.constexpr,
-    SPLIT_NUM_BLOCKS: tl.constexpr,
-):
-    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
-    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
-    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-
-    offsets = block_start_idx * BLOCK_SIZE + thread_idx
-    mask = offsets < n_elements
-
-    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
-
-    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
-    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
-
-    # Calculating absamax for each block
-    absmax = tl.max(tl.abs(A_reshaped), axis=1)
-    tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax)
-
-    A_normalized = A_reshaped / absmax[:, None]
-    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
-
-    lower_pivot = tl.zeros((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), dtype=tl.int32)
-    upper_pivot = tl.full((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), CODE_SIZE - 1, dtype=tl.int32)
-
-    for _ in range(4):  # ceil(log2(code_size)) = 4, actually, in general case should be input parameter
-        pivot = (lower_pivot + upper_pivot) // 2
-        val = tl.load(code_ptr + pivot)
-        is_higher = A_normalized > val  # code[pivot]
-        lower_pivot = tl.where(is_higher, pivot, lower_pivot)
-        upper_pivot = tl.where(is_higher, upper_pivot, pivot)
-
-    # Choose closest level
-    lower_val = tl.load(code_ptr + lower_pivot)
-    upper_val = tl.load(code_ptr + upper_pivot)
-    lower_dist = tl.abs(A_normalized - lower_val)
-    upper_dist = tl.abs(A_normalized - upper_val)
-    quantized = tl.where(lower_dist <= upper_dist, lower_pivot, upper_pivot).to(tl.uint8)
-
-    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
-    quantized = quantized.to(tl.uint8, bitcast=True)
-    left, right = quantized.split()
-    packed = left << 4 | (right & 0xF)
-
-    # Reduce don't guarantee the order of the elements passed to unite_2_int4
-    # packed = tl.reduce(quantized, axis=2, combine_fn=unite_2_int4)
-    # packed = packed.to(tl.uint8, bitcast=True)
-
-    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
-    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
-    out_mask = out_offsets < n_elements // 2
-    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)
+import torch
+
+import triton
+import triton.language as tl
+
+
+# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeFP4
+# @triton.autotune(
+#     configs=[
+#         triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 1}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 4}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 8}),
+#     ],
+#     key=["n_elements"],
+# )
+@triton.jit
+def quantize_fp4_blockwise_kernel(
+    A_ptr,
+    absmax_ptr,
+    out_ptr,
+    n_elements,
+    absmax_elements,
+    BLOCK_SIZE: tl.constexpr,
+    SPLIT_NUM_BLOCKS: tl.constexpr,
+):
+    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
+    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
+    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+
+    offsets = block_start_idx * BLOCK_SIZE + thread_idx
+    mask = offsets < n_elements
+
+    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
+
+    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
+    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
+
+    # Calculating absamax for each block
+    absmax = tl.max(tl.abs(A_reshaped), axis=1)
+    absmax_offsets = block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS)
+    tl.store(absmax_ptr + absmax_offsets, absmax, mask=absmax_offsets < absmax_elements)
+
+    A_normalized = A_reshaped / absmax[:, None]
+    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
+
+    sign = tl.where(A_normalized < 0, 0b1000, 0b0000)
+    A_absf = tl.abs(A_normalized)
+
+    result = tl.where(
+        A_absf > 0.29166667,
+        tl.where(
+            A_absf > 0.583333, tl.where(A_absf > 0.8333333, 0b011, 0b010), tl.where(A_absf > 0.4166667, 0b101, 0b100)
+        ),
+        tl.where(
+            A_absf > 0.0859375,
+            tl.where(A_absf > 0.20833333, 0b0111, 0b0110),
+            tl.where(A_absf > 0.00260417, 0b0001, 0b0000),
+        ),
+    )
+    quantized = (result ^ sign).to(tl.uint8)
+
+    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
+    left, right = quantized.split()
+    packed = left << 4 | (right & 0xF)
+
+    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
+    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+    # Use n - n//2 instead of (n+1)//2 to avoid integer overflow for large n
+    out_mask = out_offsets < (n_elements - n_elements // 2)
+    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)
+
+
+# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dQuantizeNF4
+# @triton.autotune(
+#     configs=[
+#         triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 1}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 4}),
+#         triton.Config({"SPLIT_NUM_BLOCKS": 8}),
+#     ],
+#     key=["n_elements"],
+# )
+@triton.jit
+def quantize_nf4_blockwise_kernel(
+    A_ptr,
+    absmax_ptr,
+    out_ptr,
+    n_elements,
+    absmax_elements,
+    BLOCK_SIZE: tl.constexpr,
+    SPLIT_NUM_BLOCKS: tl.constexpr,
+):
+    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
+    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
+    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+
+    offsets = block_start_idx * BLOCK_SIZE + thread_idx
+    mask = offsets < n_elements
+
+    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
+
+    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
+    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
+
+    # Calculating absamax for each block
+    absmax = tl.max(tl.abs(A_reshaped), axis=1)
+    absmax_offsets = block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS)
+    tl.store(absmax_ptr + absmax_offsets, absmax, mask=absmax_offsets < absmax_elements)
+
+    A_normalized = A_reshaped / absmax[:, None]
+    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
+
+    result = tl.where(
+        A_normalized > 0.03979014977812767,
+        tl.where(
+            A_normalized > 0.3893125355243683,
+            tl.where(
+                A_normalized > 0.6427869200706482,
+                tl.where(A_normalized > 0.8614784181118011, 0b1111, 0b1110),
+                tl.where(A_normalized > 0.5016634166240692, 0b1101, 0b1100),
+            ),
+            tl.where(
+                A_normalized > 0.2035212516784668,
+                tl.where(A_normalized > 0.2920137718319893, 0b1011, 0b1010),
+                tl.where(A_normalized > 0.1202552504837513, 0b1001, 0b1000),
+            ),
+        ),
+        tl.where(
+            A_normalized > -0.33967943489551544,
+            tl.where(
+                A_normalized > -0.13791173323988914,
+                tl.where(A_normalized > -0.045525018125772476, 0b0111, 0b0110),
+                tl.where(A_normalized > -0.23460740596055984, 0b0101, 0b0100),
+            ),
+            tl.where(
+                A_normalized > -0.6106329262256622,
+                tl.where(A_normalized > -0.4599952697753906, 0b0011, 0b0010),
+                tl.where(A_normalized > -0.8480964004993439, 0b0001, 0b0000),
+            ),
+        ),
+    )
+    quantized = result.to(tl.uint8)
+
+    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
+
+    left, right = quantized.split()
+    packed = left << 4 | (right & 0xF)
+
+    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
+    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+    # Use n - n//2 instead of (n+1)//2 to avoid integer overflow for large n
+    out_mask = out_offsets < (n_elements - n_elements // 2)
+    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)
+
+
+def quantize_4bit_blockwise_triton(A, blocksize, quant_type, blocks, absmax, num_elements, quantized_out):
+    # grid = lambda META: (triton.cdiv(blocks, META["SPLIT_NUM_BLOCKS"]),)
+    split_num_blocks = 4
+    grid = (triton.cdiv(blocks, split_num_blocks),)
+    if quant_type == "fp4":
+        quantize_fp4_blockwise_kernel[grid](
+            A_ptr=A,
+            absmax_ptr=absmax,
+            out_ptr=quantized_out,
+            n_elements=num_elements,
+            absmax_elements=absmax.numel(),
+            BLOCK_SIZE=blocksize,
+            SPLIT_NUM_BLOCKS=split_num_blocks,
+        )
+    else:
+        quantize_nf4_blockwise_kernel[grid](
+            A_ptr=A,
+            absmax_ptr=absmax,
+            out_ptr=quantized_out,
+            n_elements=num_elements,
+            absmax_elements=absmax.numel(),
+            BLOCK_SIZE=blocksize,
+            SPLIT_NUM_BLOCKS=split_num_blocks,
+        )
+    return quantized_out, absmax
+
+
+@triton.jit
+def dequant_4bit_body_util(a, offsets, quant_ptr, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
+    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
+    mask = offsets < n_elems
+    higher = a & 0xF
+    # lower 4bits
+    lower = a >> 4
+
+    abs_offsets = offsets // PAIRED_QUANT_BLOCK
+    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
+
+    # apply conversion
+    lower_4 = tl.load(quant_ptr + lower, eviction_policy="evict_last")
+    higher_4 = tl.load(quant_ptr + higher, eviction_policy="evict_last")
+
+    mul_high = higher_4 * absmax
+    mul_low = lower_4 * absmax
+    out_dq = tl.interleave(mul_low, mul_high)
+    return out_dq
+
+
+# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeFP4Tree
+@triton.jit
+def dequantize_fp4_tree(val, absmax):
+    # val: tl.tensor (uint8)
+    # absmax: tl.tensor (float32/float16)
+    #  00001100  00001011  00001001  00001111
+    sign = tl.where((val & 0b1000) == 0b1000, -1.0, 1.0)  # -1
+    third_bit = (val & 0b0100) == 0b0100  # True
+    second_bit = (val & 0b0010) == 0b0010  # False
+    first_bit = (val & 0b0001) == 0b0001  # False
+
+    branch1 = tl.where(
+        second_bit,
+        tl.where(first_bit, 0.25, 0.16666667),  # 1111, 1110
+        tl.where(first_bit, 0.5, 0.33333333),  # 1101, 1100
+    )
+    branch2 = tl.where(
+        second_bit,
+        tl.where(first_bit, 1.0, 0.66666667),  # 1011, 1010
+        tl.where(first_bit, 0.005208333333, 0.0),  # 1001, 1000
+    )
+    out = tl.where(third_bit, branch1, branch2)
+    return out * sign * absmax
+
+
+@triton.jit
+def dequant_fp4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
+    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
+    mask = offsets < n_elems
+    higher = a & 0xF
+    lower = a >> 4
+
+    abs_offsets = offsets // PAIRED_QUANT_BLOCK
+    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
+    mul_high = dequantize_fp4_tree(higher, absmax)
+    mul_low = dequantize_fp4_tree(lower, absmax)
+    out_dq = tl.interleave(mul_low, mul_high)
+    return out_dq
+
+
+# Triton implementation of similar CUDA kernel to avoid loading code from csrc/kernels.cu::dDequantizeNF4
+@triton.jit
+def dequantize_nf4_tree(val):
+    # val: tl.tensor (uint8)
+    cond0 = (val & 0b1000) == 0b1000
+    cond1 = (val & 0b0100) == 0b0100
+    cond2 = (val & 0b0010) == 0b0010
+    cond3 = (val & 0b0001) == 0b0001
+
+    # Positive branch (val & 0b1000) == 8
+    branch_pos = tl.where(
+        cond1,
+        tl.where(
+            cond2,
+            tl.where(cond3, 1.0, 0.7229568362236023),  # 1111, 1110
+            tl.where(cond3, 0.5626170039176941, 0.44070982933044434),  # 1101, 1100
+        ),
+        tl.where(
+            cond2,
+            tl.where(cond3, 0.33791524171829224, 0.24611230194568634),  # 1011, 1010
+            tl.where(cond3, 0.16093020141124725, 0.07958029955625534),  # 1001, 1000
+        ),
+    )
+
+    # Negative branch (val & 0b1000) == 0
+    branch_neg = tl.where(
+        cond1,
+        tl.where(
+            cond2,
+            tl.where(cond3, 0.0, -0.09105003625154495),  # 0111, 0110
+            tl.where(cond3, -0.18477343022823334, -0.28444138169288635),  # 0101, 0100
+        ),
+        tl.where(
+            cond2,
+            tl.where(cond3, -0.39491748809814453, -0.5250730514526367),  # 0011, 0010
+            tl.where(cond3, -0.6961928009986877, -1.0),  # 0001, 0000
+        ),
+    )
+    return tl.where(cond0, branch_pos, branch_neg)
+
+
+@triton.jit
+def dequant_nf4_body_util(a, offsets, absmax_ptr, n_elems, QUANT_BLOCK: tl.constexpr):
+    PAIRED_QUANT_BLOCK: tl.constexpr = QUANT_BLOCK // 2
+    mask = offsets < n_elems
+    higher = a & 0xF
+    # lower 4bits
+    lower = a >> 4
+
+    abs_offsets = offsets // PAIRED_QUANT_BLOCK
+    absmax = tl.load(absmax_ptr + abs_offsets, mask=mask, other=1.0, eviction_policy="evict_last")
+    mul_high = dequantize_nf4_tree(higher) * absmax
+    mul_low = dequantize_nf4_tree(lower) * absmax
+    out_dq = tl.interleave(mul_low, mul_high)
+    return out_dq
+
+
+# All such kernels are similar, so maybe code can be generalised.
+# @triton.autotune(
+#     configs=[
+# #         # triton.Config({'SPLIT_SIZE': 64}),
+# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 64, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
+#         triton.Config({'SPLIT_SIZE': 128}),
+#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
+# #         # triton.Config({'SPLIT_SIZE': 128}, num_warps = 4, num_stages = 4),
+# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 128, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
+#         triton.Config({'SPLIT_SIZE': 256}),
+#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
+#         # triton.Config({'SPLIT_SIZE': 256}, num_warps = 4, num_stages = 4),
+#         triton.Config({'SPLIT_SIZE': 512}),
+#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
+#         # triton.Config({'SPLIT_SIZE': 512}, num_warps = 4, num_stages = 4),
+# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
+# #         # # triton.Config({'SPLIT_SIZE': 512, 'grf_mode': 'auto'}, num_stages=4, num_warps=32),
+# #         # triton.Config({'SPLIT_SIZE': 1024}),
+# #         # # triton.Config({'SPLIT_SIZE': 2048}),
+# #         # # triton.Config({'SPLIT_SIZE': 4096}),
+# #         # # triton.Config({'SPLIT_SIZE': 8192}),
+# #         # # triton.Config({'SPLIT_SIZE': 16384}),
+#     ],
+#     key=['num_paired_elements'],
+# )
+@triton.jit
+def dequant_4bit_kernel(
+    a_ptr,
+    c_ptr,
+    quant_ptr,
+    absmax_ptr,
+    num_paired_elements,
+    num_output_elements,
+    QUANT_BLOCK: tl.constexpr,
+    SPLIT_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+    block_start = pid * SPLIT_SIZE
+    offsets = block_start + tl.arange(0, SPLIT_SIZE)
+    mask = offsets < num_paired_elements
+
+    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
+
+    out_dq = dequant_4bit_body_util(
+        a=a,
+        offsets=offsets,
+        quant_ptr=quant_ptr,
+        absmax_ptr=absmax_ptr,
+        n_elems=num_paired_elements,
+        QUANT_BLOCK=QUANT_BLOCK,
+    )
+
+    out_block_start = pid * SPLIT_SIZE * 2
+    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
+    mask = offs < num_output_elements
+    tl.store(c_ptr + offs, out_dq, mask)
+
+
+# @triton.autotune(
+#     configs=[
+#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 256}),
+#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 512}),
+#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2),
+#     ],
+#     key=['num_paired_elements'],
+# )
+@triton.jit
+def dequant_fp4_kernel(
+    a_ptr,
+    c_ptr,
+    absmax_ptr,
+    num_paired_elements,
+    num_output_elements,
+    QUANT_BLOCK: tl.constexpr,
+    SPLIT_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+    block_start = pid * SPLIT_SIZE
+    offsets = block_start + tl.arange(0, SPLIT_SIZE)
+    mask = offsets < num_paired_elements
+
+    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
+
+    out_dq = dequant_fp4_body_util(
+        a=a,
+        offsets=offsets,
+        absmax_ptr=absmax_ptr,
+        n_elems=num_paired_elements,
+        QUANT_BLOCK=QUANT_BLOCK,
+    )
+
+    out_block_start = pid * SPLIT_SIZE * 2
+    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
+    mask = offs < num_output_elements
+    tl.store(c_ptr + offs, out_dq, mask)
+
+
+# @triton.autotune(
+#     configs=[
+#         triton.Config({'SPLIT_SIZE': 128}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 256}),
+#         triton.Config({'SPLIT_SIZE': 256}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 512}),
+#         triton.Config({'SPLIT_SIZE': 512}, num_warps = 32, num_stages = 2),
+#         triton.Config({'SPLIT_SIZE': 1024}, num_warps = 32, num_stages = 2),
+#     ],
+#     key=['num_paired_elements'],
+# )
+@triton.jit
+def dequant_nf4_kernel(
+    a_ptr,
+    c_ptr,
+    absmax_ptr,
+    num_paired_elements,
+    num_output_elements,
+    QUANT_BLOCK: tl.constexpr,
+    SPLIT_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
+    block_start = pid * SPLIT_SIZE
+    offsets = block_start + tl.arange(0, SPLIT_SIZE)
+    mask = offsets < num_paired_elements
+
+    a = tl.load(a_ptr + offsets, mask, eviction_policy="evict_first")
+
+    out_dq = dequant_nf4_body_util(
+        a=a,
+        offsets=offsets,
+        absmax_ptr=absmax_ptr,
+        n_elems=num_paired_elements,
+        QUANT_BLOCK=QUANT_BLOCK,
+    )
+
+    out_block_start = pid * SPLIT_SIZE * 2
+    offs = out_block_start + tl.arange(0, SPLIT_SIZE * 2)
+    mask = offs < num_output_elements
+    tl.store(c_ptr + offs, out_dq, mask)
+
+
+def dequantize_4bit_impl(
+    A: torch.Tensor,
+    absmax: torch.Tensor,
+    blocksize: int,
+    quant_type: str,
+    dtype: torch.dtype,
+    out: torch.Tensor,
+) -> None:
+    # It's will be processed as an array, so
+    # actual length is row * col
+    # Elements are in uint8 format, so interleaved
+    # so total amount of data is 2 * elem_count
+    number_of_paired_elements = A.numel()
+    num_output_elements = out.numel()
+    # we assume that split_size > quant_blocksize
+
+    SPLIT_SIZE = 256
+    # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), )
+    grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),)
+    if quant_type == "fp4":
+        dequant_fp4_kernel[grid](A, out, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE)
+    else:
+        dequant_nf4_kernel[grid](A, out, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE)
+
+
+def dequantize_4bit_impl_passing_code(
+    A: torch.Tensor,
+    absmax: torch.Tensor,
+    blocksize: int,
+    code: torch.Tensor,
+    dtype: torch.dtype,
+    out: torch.Tensor,
+) -> None:
+    number_of_paired_elements = A.numel()
+    num_output_elements = out.numel()
+    # we assume that split_size > quant_blocksize
+
+    SPLIT_SIZE = 256
+    # grid = lambda META: (triton.cdiv(number_of_paired_elements, META['SPLIT_SIZE']), )
+    grid = (triton.cdiv(number_of_paired_elements, SPLIT_SIZE),)
+    dequant_4bit_kernel[grid](
+        A, out, code, absmax, number_of_paired_elements, num_output_elements, blocksize, SPLIT_SIZE
+    )
+
+
+######################### Fallback dequantization functions #########################
+## for debug ##
+
+
+# @triton.autotune(
+#     configs=[
+#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
+#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'auto'}, num_stages=2, num_warps=32),
+#         # triton.Config({'SPLIT_NUM_BLOCKS': 1, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
+#         # #
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 1, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         #
+#         triton.Config({"SPLIT_NUM_BLOCKS": 2}),
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "large"}, num_stages=2, num_warps=32),
+#         # # triton.Config({'SPLIT_NUM_BLOCKS': 2, 'grf_mode': 'large'}, num_stages=4, num_warps=32),
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=2, num_warps=32),
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 2, "grf_mode": "auto"}, num_stages=4, num_warps=32),
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=2, num_warps=32),
+#         # triton.Config({"SPLIT_NUM_BLOCKS": 4, "grf_mode": "large"}, num_stages=4, num_warps=32),
+#         # triton.Config({'SPLIT_NUM_BLOCKS': 8, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
+#     ],
+#     key=["n_elements", "BLOCK_SIZE"],
+# )
+@triton.jit
+def quantize_4bit_blockwise_kernel(
+    A_ptr,
+    code_ptr,
+    absmax_ptr,
+    out_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    CODE_SIZE: tl.constexpr,
+    SPLIT_NUM_BLOCKS: tl.constexpr,
+):
+    PAIRED_SPLIT_NUM_BLOCKS: tl.constexpr = SPLIT_NUM_BLOCKS * 2
+    block_start_idx = tl.program_id(0) * PAIRED_SPLIT_NUM_BLOCKS
+    thread_idx = tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+
+    offsets = block_start_idx * BLOCK_SIZE + thread_idx
+    mask = offsets < n_elements
+
+    A = tl.load(A_ptr + offsets, mask=mask, other=0.0)
+
+    # To be able process several blocks -> (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE)
+    A_reshaped = tl.reshape(A, (PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE))
+
+    # Calculating absamax for each block
+    absmax = tl.max(tl.abs(A_reshaped), axis=1)
+    tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax)
+
+    A_normalized = A_reshaped / absmax[:, None]
+    A_normalized = tl.clamp(A_normalized, -1.0, 1.0)
+
+    lower_pivot = tl.zeros((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), dtype=tl.int32)
+    upper_pivot = tl.full((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE), CODE_SIZE - 1, dtype=tl.int32)
+
+    for _ in range(4):  # ceil(log2(code_size)) = 4, actually, in general case should be input parameter
+        pivot = (lower_pivot + upper_pivot) // 2
+        val = tl.load(code_ptr + pivot)
+        is_higher = A_normalized > val  # code[pivot]
+        lower_pivot = tl.where(is_higher, pivot, lower_pivot)
+        upper_pivot = tl.where(is_higher, upper_pivot, pivot)
+
+    # Choose closest level
+    lower_val = tl.load(code_ptr + lower_pivot)
+    upper_val = tl.load(code_ptr + upper_pivot)
+    lower_dist = tl.abs(A_normalized - lower_val)
+    upper_dist = tl.abs(A_normalized - upper_val)
+    quantized = tl.where(lower_dist <= upper_dist, lower_pivot, upper_pivot).to(tl.uint8)
+
+    quantized = quantized.reshape((PAIRED_SPLIT_NUM_BLOCKS, BLOCK_SIZE // 2, 2))
+    quantized = quantized.to(tl.uint8, bitcast=True)
+    left, right = quantized.split()
+    packed = left << 4 | (right & 0xF)
+
+    # Reduce don't guarantee the order of the elements passed to unite_2_int4
+    # packed = tl.reduce(quantized, axis=2, combine_fn=unite_2_int4)
+    # packed = packed.to(tl.uint8, bitcast=True)
+
+    packed_flat = tl.reshape(packed, (BLOCK_SIZE * SPLIT_NUM_BLOCKS,))
+    out_offsets = block_start_idx * BLOCK_SIZE // 2 + tl.arange(0, SPLIT_NUM_BLOCKS * BLOCK_SIZE)
+    out_mask = out_offsets < n_elements // 2
+    tl.store(out_ptr + out_offsets, packed_flat, mask=out_mask)