first commit

dev/core/kernels/__init__.py

@@ -0,0 +1,7 @@
+from .tritons.fp8_block_sys import dequant_fp8_block_sym_triton
+from .pytorchs.fp8_block_sys import dequant_fp8_block_sym
+from .pytorchs.int4_group_sys import dequant_int4_group_sym
+from .pytorchs.fp8_channel_sys import dequant_fp8_channel_sym
+
+from .pytorchs.int8 import weight_quant_bf16_to_int8
+from .pytorchs.int4_pack8 import weight_quant_bf16_to_int4pack8

dev/core/kernels/fp8e4m3_to_bf16.py

@@ -0,0 +1,154 @@
+# triton_fp8_dequant.py
+import torch
+import triton
+import triton.language as tl
+from typing import Tuple
+
+
+# ----------------------------
+# Triton kernel: dequant one block tile (BLOCK x BLOCK)
+# ----------------------------
+@triton.jit
+def _dequant_block_kernel(
+    x_ptr,         # pointer to int8/uint8 input padded matrix flattened
+    s_ptr,         # pointer to float32 scales [M1 * N1]
+    out_ptr,       # pointer to float32 output padded matrix flattened
+    target_N,      # total columns in padded matrix (int)
+    M1,            # number of block-rows
+    N1,            # number of block-cols
+    BLOCK: tl.constexpr
+):
+    pid_m = tl.program_id(0)  # block row index (0..M1-1)
+    pid_n = tl.program_id(1)  # block col index (0..N1-1)
+
+    # block row/col start (in element coords)
+    m_start = pid_m * BLOCK
+    n_start = pid_n * BLOCK
+
+    # row/col indices inside global matrix for this tile
+    rows = tl.arange(0, BLOCK)
+    cols = tl.arange(0, BLOCK)
+
+    # create 2D grid of indices (BLOCK x BLOCK)
+    row_idx = m_start + rows[:, None]   # shape [BLOCK, 1]
+    col_idx = n_start + cols[None, :]   # shape [1, BLOCK]
+
+    # flatten indices into 1D offsets for load/store:
+    # idx = row_idx * target_N + col_idx  -> shape [BLOCK, BLOCK]
+    idx = row_idx * target_N + col_idx
+
+    # load block (int8) with mask; safe since padded array supplies zeros for out-of-bound
+    mask = tl.where(idx < (M1 * BLOCK) * target_N, True, False)  # always True for valid padded region
+    # Note: We rely on padded input being at least (M1*BLOCK, N1*BLOCK). Mask still required for correct loads.
+    vals_i8 = tl.load(x_ptr + idx, mask=mask, other=0)  # returns int8 values
+
+    # load scale for this block
+    s_idx = pid_m * N1 + pid_n
+    s_val = tl.load(s_ptr + s_idx)  # float32 scalar
+
+    # convert to float32 and multiply
+    vals_f32 = tl.cast(vals_i8, tl.float32) * s_val
+
+    # store results to output (float32)
+    tl.store(out_ptr + idx, vals_f32, mask=mask)
+
+
+# ----------------------------
+# Python wrapper
+# ----------------------------
+def weight_dequant_fp8e4m3_to_bf16(
+    x: torch.Tensor,        # int8/uint8 tensor shape [M, N]
+    s: torch.Tensor,        # float32 tensor shape [M1, N1]
+    block_size: int = 128,
+    pad_value: float = 0.0,
+) -> torch.Tensor:
+    """
+    Dequantize FP8-block-packed integer matrix `x` with per-block scales `s` to BF16.
+
+    Returns: torch.Tensor of shape [M, N] and dtype torch.bfloat16 (on same device as inputs).
+
+    Implementation notes:
+      - Kernel writes float32 outputs; we convert to bfloat16 in PyTorch afterwards
+      - `x` must be on CUDA device (Triton runs on CUDA)
+      - `s` must be float32 on same device
+      - This function pads `x` to full block grid using `pad_value`
+    """
+    if not x.is_cuda:
+        raise ValueError("triton_fp8_dequant_to_bf16 requires CUDA tensors (x must be on CUDA).")
+    if not s.is_cuda:
+        raise ValueError("s must be on CUDA device too.")
+    if x.dim() != 2 or s.dim() != 2:
+        raise ValueError("x must be 2D, s must be 2D (block-grid).")
+
+    device = x.device
+    M, N = x.shape
+    M1, N1 = s.shape
+    target_M = M1 * block_size
+    target_N = N1 * block_size
+    if M > target_M or N > target_N:
+        raise ValueError(f"Input x shape {(M,N)} is larger than block grid implied by s: {(target_M, target_N)}")
+
+    # ensure x dtype is int8 or uint8. If not, cast (we assume FP8 mantissa stored as int8)
+    if x.dtype not in (torch.int8, torch.uint8):
+        x_int8 = x.to(torch.int8)
+    else:
+        x_int8 = x
+
+    # prepare padded input (int8) - use zeros to avoid uninitialized values
+    if M == target_M and N == target_N and x_int8.is_contiguous():
+        padded_x = x_int8
+    else:
+        padded_x = torch.full((target_M, target_N), fill_value=int(pad_value), dtype=torch.int8, device=device)
+        padded_x[:M, :N] = x_int8
+
+    # allocate output padded float32 buffer
+    out_f32 = torch.empty((target_M, target_N), dtype=torch.float32, device=device)
+
+    # pointers
+    x_ptr = padded_x
+    s_ptr = s.to(device=device, dtype=torch.float32)  # ensure float32 and on device
+    out_ptr = out_f32
+
+    # launch grid: each program handles one block tile
+    grid = (M1, N1)
+
+    # compute num warps heuristics: bigger blocks -> fewer warps; set default to 8
+    num_warps = 8
+
+    _dequant_block_kernel[grid](x_ptr, s_ptr, out_ptr, target_N, M1, N1, BLOCK=block_size, num_warps=num_warps)
+
+    # crop back to original shape and convert to bfloat16
+    res_bf16 = out_f32[:M, :N].to(dtype=torch.bfloat16)
+
+    return res_bf16
+
+
+# ----------------------------
+# Quick sanity test (only run if this file executed directly)
+# ----------------------------
+if __name__ == "__main__":
+    # small example to sanity-check on GPU
+    if not torch.cuda.is_available():
+        print("CUDA not available; cannot run Triton kernel.")
+    else:
+        device = torch.device("cuda")
+        M, N = 300, 500
+        block_size = 128
+        M1 = (M + block_size - 1) // block_size
+        N1 = (N + block_size - 1) // block_size
+
+        # create small int8-like data and small scales
+        x = torch.randint(-10, 11, (M, N), dtype=torch.int8, device=device)
+        s = torch.rand((M1, N1), dtype=torch.float32, device=device) * 0.05 + 0.01
+
+        out = triton_fp8_dequant_to_bf16(x, s, block_size=block_size)
+        print("out.shape=", out.shape, " dtype=", out.dtype, " device=", out.device)
+        # quick numeric check against naive CPU implementation
+        # naive:
+        padded = torch.zeros((M1*block_size, N1*block_size), dtype=torch.int8, device=device)
+        padded[:M, :N] = x
+        blocks = padded.view(M1, block_size, N1, block_size).permute(0,2,1,3)
+        s_expanded = s.unsqueeze(-1).unsqueeze(-1).to(torch.float32)
+        naive = (blocks.to(torch.float32) * s_expanded).permute(0,2,1,3).reshape(M1*block_size, N1*block_size)[:M, :N].to(torch.bfloat16)
+        diff = (naive.to(torch.float32) - out.to(torch.float32)).abs().max()
+        print("max abs diff (float32):", diff.item())

dev/core/kernels/pytorchs/fp8_block_sys.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn.functional as F
+
+
+def dequant_fp8_block_sym(
+    artifacts,
+    rq,
+    output_dtype=torch.bfloat16,
+    device: torch.device = torch.device("cuda")
+) -> torch.Tensor:
+    """
+    Fully vectorized FP8 block-wise symmetric dequantization.
+
+    Supports:
+    - column block scale         [num_blocks]
+    - row-wise block scale       [out_features, num_blocks]
+    - 2D block scale             [row_blocks, col_blocks]
+
+    Works with partial blocks.
+    """
+
+    qweight = artifacts.qweight.to(device, non_blocking=True).contiguous()
+    scale = artifacts.scale.to(device, non_blocking=True).contiguous()
+
+    M, N = qweight.shape
+
+    block_size = (
+        rq.block_structure[1]
+        if rq.block_structure
+        else rq.group_size or N
+    )
+
+    # number of column blocks
+    col_blocks = (N + block_size - 1) // block_size
+    col_pad = col_blocks * block_size - N
+
+    # pad column
+    if col_pad > 0:
+        qweight = F.pad(qweight, (0, col_pad))
+
+    # reshape → block layout
+    # [M, col_blocks, block_size]
+    q_blocks = qweight.view(M, col_blocks, block_size)
+
+    # =========================
+    # build broadcast scale
+    # =========================
+
+    if scale.ndim == 1:
+        # [col_blocks]
+        if scale.shape[0] != col_blocks:
+            raise ValueError(
+                f"Invalid scale shape {scale.shape} for col_blocks={col_blocks}"
+            )
+
+        scale_broadcast = scale.view(1, col_blocks, 1)
+
+    elif scale.ndim == 2:
+
+        if scale.shape == (M, col_blocks):
+            # per-row block scale
+            scale_broadcast = scale.unsqueeze(-1)
+
+        else:
+            # 2D block scale
+            row_blocks, col_blocks_scale = scale.shape
+
+            if col_blocks_scale != col_blocks:
+                raise ValueError(
+                    f"Invalid block scale {scale.shape}, expected col_blocks={col_blocks}"
+                )
+
+            # expand row blocks
+            # [row_blocks] → rows
+            expanded_rows = row_blocks * block_size
+
+            # broadcast to rows
+            scale_rows = scale.repeat_interleave(block_size, dim=0)
+
+            # crop to M
+            scale_rows = scale_rows[:M]
+
+            scale_broadcast = scale_rows.unsqueeze(-1)
+
+    else:
+        raise ValueError(f"Unsupported scale shape {scale.shape}")
+
+    # =========================
+    # dequantization
+    # =========================
+
+    dequant_blocks = q_blocks.float() * scale_broadcast
+
+    # reshape back
+    dequant = dequant_blocks.reshape(M, col_blocks * block_size)
+
+    if col_pad > 0:
+        dequant = dequant[:, :N]
+
+    return dequant.to(output_dtype)
+
+# def dequant_fp8_block_sym(
+#     artifacts,
+#     rq,
+#     output_dtype=torch.bfloat16,
+#     device: torch.device = torch.device("cuda")
+# ) -> torch.Tensor:
+#     """
+#     Fully vectorized FP8 block-wise symmetric dequantization on GPU.
+#     Avoids Python loops for large tensors.
+#     """
+#     qweight = artifacts.qweight.to(device, non_blocking=True)
+#     scale = artifacts.scale.to(device, non_blocking=True)
+#     out_features, in_features = qweight.shape
+#     block_size = rq.block_structure[1] if rq.block_structure else rq.group_size or in_features
+#     num_blocks = (in_features + block_size - 1) // block_size
+#     pad_size = num_blocks * block_size - in_features
+
+#     # pad qweight if necessary
+#     if pad_size > 0:
+#         qweight = F.pad(qweight, (0, pad_size))
+
+#     # reshape为 [out_features, num_blocks, block_size]
+#     qweight_blocks = qweight.view(out_features, num_blocks, block_size)
+
+#     # 处理 scale
+#     if scale.ndim == 1 and scale.shape[0] == num_blocks:
+#         scale_broadcast = scale.view(1, num_blocks, 1)
+#     elif scale.ndim == 2:
+#         if scale.shape == (out_features, num_blocks):
+#             scale_broadcast = scale.unsqueeze(-1)
+#         else:
+#             # block-wise scale [num_row_blocks, num_col_blocks]
+#             num_row_blocks, num_col_blocks = scale.shape
+#             if num_row_blocks * block_size != out_features or num_col_blocks != num_blocks:
+#                 raise ValueError(f"Unsupported block-wise scale {scale.shape}, {qweight.shape}, {block_size}")
+#             scale_broadcast = scale.repeat_interleave(block_size, dim=0).unsqueeze(-1)
+#     else:
+#         raise ValueError(f"Unsupported scale shape {scale.shape}")
+
+#     # 向量化反量化
+#     dequant_blocks = qweight_blocks.float() * scale_broadcast
+
+#     # reshape回原始形状
+#     dequant = dequant_blocks.view(out_features, num_blocks * block_size)
+#     if pad_size > 0:
+#         dequant = dequant[:, :in_features]
+
+#     output = dequant.to(output_dtype)
+
+#     return output

dev/core/kernels/pytorchs/fp8_channel_sys.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn.functional as F
+
+def dequant_fp8_channel_sym(
+    artifacts,
+    output_dtype=torch.bfloat16,
+    device: torch.device = torch.device("cuda"),
+    *args,
+    **kwargs
+) -> torch.Tensor:
+    """
+    Fully vectorized FP8 channel-wise symmetric dequantization on GPU.
+    Channel dimension = out_features
+    """
+    qweight = artifacts.qweight.to(device, non_blocking=True)
+    scale = artifacts.scale.to(device, non_blocking=True)
+
+    # qweight: [out_features, in_features]
+    out_features, in_features = qweight.shape
+
+    # -------- scale 处理（channel-wise）--------
+    if scale.ndim == 1:
+        # [out_features]
+        if scale.shape[0] != out_features:
+            raise ValueError(
+                f"Channel-wise scale must have shape [{out_features}], got {scale.shape}"
+            )
+        scale_broadcast = scale.view(out_features, 1)
+
+    elif scale.ndim == 2:
+        # [out_features, 1]
+        if scale.shape == (out_features, 1):
+            scale_broadcast = scale
+        else:
+            raise ValueError(
+                f"Unsupported channel-wise scale shape {scale.shape}"
+            )
+    else:
+        raise ValueError(f"Unsupported scale shape {scale.shape}")
+
+    # -------- 向量化反量化 --------
+    dequant = qweight.float() * scale_broadcast
+
+    return dequant.to(output_dtype)

dev/core/kernels/pytorchs/int4_group_sys.py

@@ -0,0 +1,147 @@
+import torch
+
+def unpack_int4_from_int32(
+    packed: torch.Tensor,
+    original_shape: tuple,
+    *,
+    out_dtype=torch.float32,
+):
+    """
+    Unpack INT4 weights packed in INT32.
+
+    Layout per int32 (little-endian nibble order):
+    [ 3..0 | 7..4 | 11..8 | ... | 31..28 ]  -> 8 x INT4
+
+    Args:
+        packed: int32 tensor, arbitrary shape
+        original_shape: target shape after unpack
+        out_dtype: output dtype (default float32)
+
+    Returns:
+        Tensor of shape original_shape, dtype = out_dtype
+    """
+    assert packed.dtype == torch.int32
+    device = packed.device
+
+    # [N] -> [N, 1]
+    x = packed.view(-1, 1)
+
+    # shifts = [0,4,8,...,28]
+    shifts = torch.arange(
+        0, 32, 4,
+        device=device,
+        dtype=torch.int32
+    ).view(1, 8)
+
+    # [N, 8] uint4 in [0, 15]
+    vals = (x >> shifts) & 0xF
+
+    # uint4 -> int4 : [0,15] → [-8,7]
+    vals = vals.to(torch.int8).sub_(8)
+
+    # flatten & trim padding
+    numel = torch.prod(
+        torch.tensor(original_shape, device=device)
+    ).item()
+
+    vals = vals.view(-1)[:numel]
+
+    return vals.view(original_shape).to(out_dtype)
+
+
+def apply_group_scaling(unpacked, scale_tensor, group_size=32):
+    """
+    Apply group-wise scaling to unpacked INT4 values.
+
+    Args:
+        unpacked: torch.Tensor, shape [out_features, in_features]
+        scale_tensor: torch.Tensor, shape [out_features, in_features//group_size]
+        group_size: int, number of elements per group (default 32 for Kimi)
+
+    Returns:
+        torch.Tensor with scaled values
+    """
+    if scale_tensor.numel() == 1:
+        # Single scale value
+        return unpacked * scale_tensor.item()
+
+    out_features, in_features = unpacked.shape
+    scale_out, scale_in = scale_tensor.shape
+
+    if scale_out == out_features and scale_in * group_size == in_features:
+        # Standard group-wise scaling
+        # Reshape: [out_features, in_features] -> [out_features, scale_in, group_size]
+        weight_grouped = unpacked.view(out_features, scale_in, group_size)
+        # Expand scale: [out_features, scale_in] -> [out_features, scale_in, 1]
+        scale_expanded = scale_tensor.view(out_features, scale_in, 1)
+        # Apply scaling
+        scaled_grouped = weight_grouped * scale_expanded.float()
+        # Reshape back: [out_features, scale_in, group_size] -> [out_features, in_features]
+        return scaled_grouped.view(out_features, in_features)
+
+    elif scale_out == out_features:
+        # Try to handle irregular group sizes
+        actual_group_size = in_features // scale_in
+        if actual_group_size > 0 and in_features % scale_in == 0:
+            weight_grouped = unpacked.view(out_features, scale_in, actual_group_size)
+            scale_expanded = scale_tensor.view(out_features, scale_in, 1)
+            scaled_grouped = weight_grouped * scale_expanded.float()
+            return scaled_grouped.view(out_features, in_features)
+        else:
+            # Fallback: repeat scales to match dimensions
+            scale_repeated = scale_tensor.repeat_interleave(
+                (in_features + scale_in - 1) // scale_in, dim=1
+            )[:, :in_features]
+            return unpacked * scale_repeated.float()
+
+    else:
+        # Last resort: try direct broadcasting
+        try:
+            return unpacked * scale_tensor.float()
+        except RuntimeError:
+            # Use mean scale as fallback
+            return unpacked * scale_tensor.mean().item()
+
+
+def weight_dequant_int4(packed_tensor, scale_tensor, shape_tensor, group_size):
+    """
+    Dequantize INT4 weights to BF16.
+
+    This is the INT4 equivalent of weight_dequant() in the FP8 version.
+
+    Args:
+        packed_tensor: torch.Tensor, INT32 tensor with packed INT4 values
+        scale_tensor: torch.Tensor, BF16 tensor with group-wise scales
+        shape_tensor: torch.Tensor, INT32 tensor with original shape [H, W]
+
+    Returns:
+        torch.Tensor in BF16 format with original shape
+    """
+    # Get original shape
+    original_shape = tuple(shape_tensor.tolist())
+
+    # Step 1: Unpack INT4 values from INT32 storage
+    unpacked = unpack_int4_from_int32(packed_tensor, original_shape).to(scale_tensor.device)
+
+    # Step 2: Apply group-wise scaling
+    scaled = apply_group_scaling(unpacked, scale_tensor, group_size=group_size)
+
+    # Step 3: Convert to BF16
+    return scaled.to(torch.bfloat16)
+
+
+def dequant_int4_group_sym(
+    artifacts,
+    rq,
+    output_dtype=torch.bfloat16,
+    device: torch.device = torch.device("cuda")
+):
+    qweight = artifacts.qweight.to(device, non_blocking=True)
+    scale = artifacts.scale.to(device, non_blocking=True)
+    shape = artifacts.shape.to(device, non_blocking=True)
+
+    group_size = rq.group_size
+
+    output = weight_dequant_int4(qweight, scale, shape, group_size)
+
+    return output.to(output_dtype)

dev/core/kernels/pytorchs/int4_pack8.py

@@ -0,0 +1,58 @@
+import torch
+import ixformer.inference.functions as ixfop
+
+
+def weight_quant_bf16_to_int4pack8(
+    v: torch.Tensor,            # [B, R, C]
+    block_size: int = 128,
+    group_size: int = -1,
+    format: str = "TN",
+    symmetric: bool = True,
+    version: int = 2,
+):
+    """
+    Batch 版本 INT4 量化 + 打包。
+
+    Args:
+        v: [batch, rows, cols], float Tensor
+
+    Returns:
+        i4_weights: [batch, rows, packed_cols]
+        scale:      [batch, rows, 1]
+        i8scales:   来自 ixfop.quant_repack_int4
+        i8zeros:    来自 ixfop.quant_repack_int4
+    """
+    assert v.dim() == 3, f"expected [batch, rows, cols], got {v.shape}"
+
+    B, R, C = v.shape
+
+    qmax = 127.0
+
+    # abs_max: [B, R, 1]
+    abs_max = torch.abs(v).amax(dim=2, keepdim=True)
+    scale = abs_max / qmax     # [B, R, 1]
+
+    # quantized: [B, R, C]
+    quantized = torch.round(v / scale)
+    quantized = torch.clamp(quantized, -qmax, qmax).to(torch.int8)
+
+    # ixfop.quant_repack_int4 需要 [batch, rows, cols]
+    # 它本来就是 batch-first，可直接送进去
+    # 返回形状一般是:
+    #   i4_weights: [B, R, packed_C]
+    #   i8scales:   [B, R, groups]
+    #   i8zeros:    [B, R, groups]
+    i4_weights, i8scales, i8zeros = ixfop.quant_repack_int4(
+        quantized,          # 不需要 unsqueeze，因为本来就是 [B, R, C]
+        group_size,
+        version,
+        format,
+        not symmetric,  # ixformer parameter is isAsymQuant
+    )
+
+    return (
+        i4_weights,          # [B, R, packed_C]
+        scale.to(torch.float32),  # [B, R, 1]
+        i8scales,            # 来自 repack
+        i8zeros
+    )

dev/core/kernels/pytorchs/int4_pack8.py.pre_glm5_group64.bak

@@ -0,0 +1,58 @@
+import torch
+import ixformer.inference.functions as ixfop
+
+
+def weight_quant_bf16_to_int4pack8(
+    v: torch.Tensor,            # [B, R, C]
+    block_size: int = 128,
+    group_size: int = -1,
+    format: str = "TN",
+    symmetric: bool = True,
+    version: int = 2,
+):
+    """
+    Batch 版本 INT4 量化 + 打包。
+
+    Args:
+        v: [batch, rows, cols], float Tensor
+
+    Returns:
+        i4_weights: [batch, rows, packed_cols]
+        scale:      [batch, rows, 1]
+        i8scales:   来自 ixfop.quant_repack_int4
+        i8zeros:    来自 ixfop.quant_repack_int4
+    """
+    assert v.dim() == 3, f"expected [batch, rows, cols], got {v.shape}"
+
+    B, R, C = v.shape
+
+    qmax = 127.0
+
+    # abs_max: [B, R, 1]
+    abs_max = torch.abs(v).amax(dim=2, keepdim=True)
+    scale = abs_max / qmax     # [B, R, 1]
+
+    # quantized: [B, R, C]
+    quantized = torch.round(v / scale)
+    quantized = torch.clamp(quantized, -qmax, qmax).to(torch.int8)
+
+    # ixfop.quant_repack_int4 需要 [batch, rows, cols]
+    # 它本来就是 batch-first，可直接送进去
+    # 返回形状一般是:
+    #   i4_weights: [B, R, packed_C]
+    #   i8scales:   [B, R, groups]
+    #   i8zeros:    [B, R, groups]
+    i4_weights, i8scales, i8zeros = ixfop.quant_repack_int4(
+        quantized,          # 不需要 unsqueeze，因为本来就是 [B, R, C]
+        group_size,
+        version,
+        format,
+        symmetric,
+    )
+
+    return (
+        i4_weights,          # [B, R, packed_C]
+        scale.to(torch.float32),  # [B, R, 1]
+        i8scales,            # 来自 repack
+        i8zeros
+    )

dev/core/kernels/pytorchs/int8.py

@@ -0,0 +1,16 @@
+import torch
+
+
+def weight_quant_bf16_to_int8(inputs: torch.Tensor):
+    assert inputs.dim() == 3, f"inputs shape is [batch, output_dim, input_dim], but got {inputs.dim()}"
+
+    qmax = 127.0
+    abs_max = torch.abs(inputs).max(dim=2, keepdim=True)[0]
+    scale = abs_max / qmax
+
+    assert scale.shape == (*inputs.shape[:2], 1)
+
+    quantized = torch.round(inputs / scale)
+    quantized = torch.clamp(quantized, -qmax, qmax)
+    return quantized.to(torch.int8), scale.to(torch.float32)
+

dev/core/kernels/tests/test_dequant_fp8_block_sys.py

@@ -0,0 +1,95 @@
+import torch
+import time
+import numpy as np
+
+from pathlib import Path
+
+current_dir = Path(__file__).resolve()
+obj_dir = current_dir.parent.parent
+
+import sys
+sys.path.insert(0, str(obj_dir))
+
+from pytorchs.fp8_block_sys import dequant_fp8_block_sym
+from tritons.fp8_block_sys import dequant_fp8_block_sym_triton
+
+# 假设你已经有以下函数
+# dequant_fp8_block_sym    : PyTorch 实现
+# dequant_fp8_block_sym_triton : Triton 实现
+
+# -------------------------------
+# 测试参数
+M, N = 1024, 1024       # 权重矩阵大小
+block_size = 32
+
+# 随机生成量化权重
+torch.manual_seed(0)
+qweight = torch.randint(-127, 128, (M, N), dtype=torch.int8, device='cuda')
+
+# 随机生成 block-wise scale
+num_blocks = (N + block_size - 1) // block_size
+scale = torch.rand((M, num_blocks), dtype=torch.float32, device='cuda')
+
+# 构造 artifacts 和 RQ
+from dataclasses import dataclass
+
+@dataclass
+class QuantArtifacts:
+    qweight: torch.Tensor
+    scale: torch.Tensor
+    zero_point: torch.Tensor | None = None
+    shape: torch.Tensor | None = None
+
+@dataclass
+class ResolvedQuantConfig:
+    dtype: str
+    num_bits: int
+    packed: bool
+    strategy: str
+    group_size: int | None
+    block_structure: tuple | None
+    symmetric: bool
+    dynamic: bool
+
+artifacts = QuantArtifacts(qweight=qweight, scale=scale)
+rq = ResolvedQuantConfig(
+    dtype="float",
+    num_bits=8,
+    packed=False,
+    strategy="block",
+    group_size=None,
+    block_structure=(M, block_size),
+    symmetric=True,
+    dynamic=False
+)
+
+# -------------------------------
+# PyTorch 反量化
+torch.cuda.synchronize()
+t0 = time.time()
+dequant_pt = dequant_fp8_block_sym(artifacts, rq)
+torch.cuda.synchronize()
+t_pt = time.time() - t0
+print(f"PyTorch dequant time: {t_pt*1000:.2f} ms")
+# -------------------------------
+# Triton 反量化
+torch.cuda.synchronize()
+t0 = time.time()
+dequant_tr = dequant_fp8_block_sym_triton(artifacts, rq)
+torch.cuda.synchronize()
+t_tr = time.time() - t0
+print(f"Triton dequant time: {t_tr*1000:.2f} ms")
+
+# -------------------------------
+# 验证数值一致性
+diff = (dequant_pt.float() - dequant_tr.float()).abs().max().item()
+print(f"Max absolute difference: {diff}")
+if diff < 1e-5:
+    print("✅ PyTorch and Triton results match!")
+else:
+    print("❌ Results mismatch!")
+
+# -------------------------------
+# 可选：统计平均误差
+mean_diff = (dequant_pt.float() - dequant_tr.float()).abs().mean().item()
+print(f"Mean absolute difference: {mean_diff}")

dev/core/kernels/tritons/fp8_block_sys.py

@@ -0,0 +1,84 @@
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def fp8_block_dequant_kernel(
+    qweight_ptr, scale_ptr, dequant_ptr,
+    M, N, block_size, num_blocks,
+    SCALE_ROWWISE: tl.constexpr,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    row_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    col_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    row_mask = row_offsets < M
+    col_mask = col_offsets < N
+
+    offs = row_offsets[:, None] * N + col_offsets[None, :]
+    x = tl.load(qweight_ptr + offs, mask=row_mask[:, None] & col_mask[None, :], other=0)
+    x = x.to(tl.float32)
+
+    block_id = col_offsets // block_size
+
+    if SCALE_ROWWISE:
+        # scale: [M, num_blocks] -> load per row
+        scale_offs = row_offsets[:, None] * num_blocks + block_id[None, :]
+        s = tl.load(scale_ptr + scale_offs, mask=row_mask[:, None] & col_mask[None, :], other=1.0)
+    else:
+        # scale: [num_blocks] -> broadcast row
+        s = tl.load(scale_ptr + block_id[None, :])
+        s = s[None, :]  # broadcast to rows
+
+    dequant = x * s
+    tl.store(dequant_ptr + offs, dequant, mask=row_mask[:, None] & col_mask[None, :])
+
+
+def dequant_fp8_block_sym_triton(artifacts, rq):
+    qweight = artifacts.qweight
+    scale = artifacts.scale
+    M, N = qweight.shape
+
+    # block size
+    if rq.block_structure:
+        block_size = rq.block_structure[1]
+    elif rq.group_size:
+        block_size = rq.group_size
+    else:
+        block_size = N
+
+    num_blocks = (N + block_size - 1) // block_size
+
+    # 判断 scale 是否 row-wise
+    SCALE_ROWWISE = scale.ndim == 2
+    if SCALE_ROWWISE:
+        assert scale.shape[0] == M and scale.shape[1] == num_blocks, \
+            f"Row-wise scale shape must be [M, num_blocks], got {scale.shape}"
+    else:
+        # scale 1D
+        assert scale.shape[0] == num_blocks, \
+            f"Scale shape must be [num_blocks], got {scale.shape}"
+
+    dequant = torch.empty((M, N), dtype=torch.float32, device=qweight.device)
+
+    BLOCK_M = 64
+    BLOCK_N = 128
+    grid_m = (M + BLOCK_M - 1) // BLOCK_M
+    grid_n = (N + BLOCK_N - 1) // BLOCK_N
+    grid = (grid_m, grid_n)
+
+    fp8_block_dequant_kernel[grid](
+        qweight, scale, dequant,
+        M, N, block_size, num_blocks,
+        SCALE_ROWWISE,
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N
+    )
+
+    if rq.dtype.lower() == "bf16":
+        dequant = dequant.to(torch.bfloat16)
+
+    return dequant

dev/core/kernels/tritons/int4_block_sys.py

@@ -0,0 +1,349 @@
+import torch
+import numpy as np
+
+from typing import List
+from torch.multiprocessing import Pool
+
+
+# @csabakecskemeti
+# def unpack_int4_from_int32(packed_tensor, original_shape):
+#     """
+#     Unpack INT4 values from INT32 storage.
+
+#     Each INT32 contains 8 INT4 values:
+#     INT32: [31..28][27..24][23..20][19..16][15..12][11..8][7..4][3..0]
+#            │  val7 ││  val6 ││  val5 ││  val4 ││  val3 ││ val2││val1││val0│
+
+#     Args:
+#         packed_tensor: torch.Tensor with dtype int32 containing packed INT4 values
+#         original_shape: tuple, the original shape of the weight before packing
+
+#     Returns:
+#         torch.Tensor with unpacked INT4 values as float32
+#     """
+#     # Convert to numpy for reliable bit operations
+#     packed_np = packed_tensor.detach().cpu().numpy().astype(np.uint32)
+
+#     # Extract 8 INT4 values from each UINT32
+#     unpacked_list = []
+#     for i in range(8):
+#         shift = i * 4
+#         mask = 0xF  # 4-bit mask (0b1111)
+#         int4_val = (packed_np >> shift) & mask
+#         # Convert from unsigned [0,15] to signed [-8,7]
+#         int4_signed = int4_val.astype(np.int8) - 8
+#         unpacked_list.append(int4_signed)
+
+#     # Stack along new axis: [..., 8]
+#     unpacked_np = np.stack(unpacked_list, axis=-1)
+
+#     # Flatten the last dimension: [..., 8] -> [..., 8*elements]
+#     flat_shape = list(unpacked_np.shape[:-1]) + [-1]
+#     unpacked_flat = unpacked_np.reshape(flat_shape)
+
+#     # Calculate expected total elements
+#     total_elements = 1
+#     for dim in original_shape:
+#         total_elements *= dim
+
+#     # Handle potential padding
+#     unpacked_1d = unpacked_flat.flatten()
+#     if len(unpacked_1d) > total_elements:
+#         unpacked_1d = unpacked_1d[:total_elements]
+#     elif len(unpacked_1d) < total_elements:
+#         # Pad if needed (shouldn't happen with correct packing)
+#         padding = total_elements - len(unpacked_1d)
+#         unpacked_1d = np.concatenate([unpacked_1d, np.zeros(padding, dtype=np.int8)])
+
+#     # Reshape to original shape
+#     result = unpacked_1d.reshape(original_shape)
+#     return torch.from_numpy(result).float()
+
+# def unpack_int4_from_int32(packed_tensor: torch.Tensor, original_shape):
+#     """
+#     High-performance INT4 unpack using pure PyTorch tensor ops.
+#     packed_tensor: int32 tensor containing 8 x INT4 packed
+#     original_shape: target shape (H, W)
+#     return: float32 tensor with unpacked int4 values
+#     """
+
+#     # packed_tensor: [N] int32
+#     device = packed_tensor.device
+#     x = packed_tensor.to(torch.int32)
+
+#     # Prepare shifts for extracting 8 nibbles
+#     # shift = [0, 4, 8, 12, 16, 20, 24, 28]
+#     shifts = torch.arange(0, 32, 4, device=device, dtype=torch.int32)
+
+#     # Extract 8 nibble values: [(x >> s) & 0xF for s in shifts]
+#     # Expand: x[..., None] shape => [N, 1]
+#     # Broadcast: shifts => [1, 8]
+#     # Result => [N, 8]
+#     unpacked = (x.unsqueeze(-1) >> shifts) & 0xF
+
+#     # Convert unsigned [0..15] → signed [-8..7]
+#     unpacked = unpacked.to(torch.int8) - 8
+
+#     # Final reshape
+#     total = 1
+#     for d in original_shape:
+#         total *= d
+
+#     unpacked = unpacked.reshape(-1)[:total]
+#     return unpacked.reshape(original_shape).float()
+
+
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def kernel_unpack_int4(
+    packed_ptr, out_ptr, N, OUT_SIZE: tl.constexpr, BLOCK: tl.constexpr
+):
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK
+
+    offsets = block_start + tl.arange(0, BLOCK)
+    mask = offsets < N  # Triton mask，0/1 即可
+
+    # load int32 packed data
+    x = tl.load(packed_ptr + offsets, mask=mask, other=0)
+    x = x.to(tl.int32)
+
+    # unpack 8 int4 per int32
+    shifts = tl.arange(0, 8) * 4
+    x = (x[:, None] >> shifts[None, :]) & 0xF
+    x = x.to(tl.int8) - 8
+
+    # flatten
+    x_flat = tl.reshape(x, (BLOCK * 8,))
+
+    # write to output
+    out_offsets = block_start * 8 + tl.arange(0, BLOCK * 8)
+    mask_out = out_offsets < OUT_SIZE  # 0/1 mask
+    tl.store(out_ptr + out_offsets, x_flat.to(tl.float32), mask=mask_out)
+
+
+def unpack_int4_from_int32(packed: torch.Tensor, original_shape, out_dtype=torch.float32):
+    device = packed.device
+    packed = packed.flatten().to(torch.int32)
+    N = packed.numel()
+
+    total = int(torch.prod(torch.tensor(original_shape)))
+    out = torch.empty(total, dtype=out_dtype, device=device)
+
+    BLOCK = 1024
+    grid = lambda meta: ((N + BLOCK - 1) // BLOCK,)
+
+    kernel_unpack_int4[grid](
+        packed_ptr=packed,
+        out_ptr=out,
+        N=N,
+        OUT_SIZE=total,
+        BLOCK=BLOCK,
+    )
+
+    return out.reshape(original_shape)
+
+
+def apply_group_scaling(unpacked, scale_tensor, group_size=32):
+    """
+    Apply group-wise scaling to unpacked INT4 values.
+
+    Args:
+        unpacked: torch.Tensor, shape [out_features, in_features]
+        scale_tensor: torch.Tensor, shape [out_features, in_features//group_size]
+        group_size: int, number of elements per group (default 32 for Kimi)
+
+    Returns:
+        torch.Tensor with scaled values
+    """
+    if scale_tensor.numel() == 1:
+        # Single scale value
+        return unpacked * scale_tensor.item()
+
+    out_features, in_features = unpacked.shape
+    scale_out, scale_in = scale_tensor.shape
+
+    if scale_out == out_features and scale_in * group_size == in_features:
+        # Standard group-wise scaling
+        # Reshape: [out_features, in_features] -> [out_features, scale_in, group_size]
+        weight_grouped = unpacked.view(out_features, scale_in, group_size)
+        # Expand scale: [out_features, scale_in] -> [out_features, scale_in, 1]
+        scale_expanded = scale_tensor.view(out_features, scale_in, 1)
+        # Apply scaling
+        scaled_grouped = weight_grouped * scale_expanded.float()
+        # Reshape back: [out_features, scale_in, group_size] -> [out_features, in_features]
+        return scaled_grouped.view(out_features, in_features)
+
+    elif scale_out == out_features:
+        # Try to handle irregular group sizes
+        actual_group_size = in_features // scale_in
+        if actual_group_size > 0 and in_features % scale_in == 0:
+            weight_grouped = unpacked.view(out_features, scale_in, actual_group_size)
+            scale_expanded = scale_tensor.view(out_features, scale_in, 1)
+            scaled_grouped = weight_grouped * scale_expanded.float()
+            return scaled_grouped.view(out_features, in_features)
+        else:
+            # Fallback: repeat scales to match dimensions
+            scale_repeated = scale_tensor.repeat_interleave(
+                (in_features + scale_in - 1) // scale_in, dim=1
+            )[:, :in_features]
+            return unpacked * scale_repeated.float()
+
+    else:
+        # Last resort: try direct broadcasting
+        try:
+            return unpacked * scale_tensor.float()
+        except RuntimeError:
+            # Use mean scale as fallback
+            return unpacked * scale_tensor.mean().item()
+
+
+def weight_dequant_int4(packed_tensor, scale_tensor, shape_tensor):
+    """
+    Dequantize INT4 weights to BF16.
+
+    This is the INT4 equivalent of weight_dequant() in the FP8 version.
+
+    Args:
+        packed_tensor: torch.Tensor, INT32 tensor with packed INT4 values
+        scale_tensor: torch.Tensor, BF16 tensor with group-wise scales
+        shape_tensor: torch.Tensor, INT32 tensor with original shape [H, W]
+
+    Returns:
+        torch.Tensor in BF16 format with original shape
+    """
+    # Get original shape
+    original_shape = tuple(shape_tensor.tolist())
+
+    # Step 1: Unpack INT4 values from INT32 storage
+    unpacked = unpack_int4_from_int32(packed_tensor, original_shape).to(scale_tensor.device)
+
+    # Step 2: Apply group-wise scaling
+    scaled = apply_group_scaling(unpacked, scale_tensor, group_size=32)
+
+    # Step 3: Convert to BF16
+    return scaled.to(torch.bfloat16)
+
+
+def _dequant_one_prefix(args):
+    """
+    子进程函数：对一个 prefix（layer）进行反量化
+    """
+    prefix_name, suffix_map, state_dict, device = args
+
+    packed_tensor = state_dict[suffix_map[".weight_packed"]].to(device, non_blocking=True)
+    scale_tensor  = state_dict[suffix_map[".weight_scale"]].to(device, non_blocking=True)
+    shape_tensor  = state_dict[suffix_map[".weight_shape"]].to(device, non_blocking=True)
+
+    with torch.no_grad():
+        new_weight = weight_dequant_int4(
+            packed_tensor,
+            scale_tensor,
+            shape_tensor
+        ).cpu()
+
+    # 避免显存泄漏
+    del packed_tensor, scale_tensor, shape_tensor
+    torch.cuda.empty_cache()
+
+    return prefix_name, new_weight
+
+
+def weight_dequant_int4_wrapper(state_dict, suffixes: list[str], device: str = "cuda", batch_size: int = 128):
+    """
+    批量堆叠反量化版本：
+    - 将多个 prefix 的 tensor 堆叠到 GPU
+    - 一次性反量化，提高显存和计算利用率
+    """
+    suffix_set = set(suffixes)
+    prefix_map: dict[str, dict[str, str]] = {}
+    new_sd = {}
+    weight_names = list(state_dict.keys())
+
+    for w in weight_names:
+        for suf in suffixes:
+            if w.endswith(suf):
+                prefix = w.removesuffix(suf)
+                prefix_map.setdefault(prefix, {})[suf] = w
+                break
+        else:
+            new_sd[w] = state_dict[w]
+
+    for prefix_name, sufmap in prefix_map.items():
+        missing = suffix_set - set(sufmap.keys())
+        if missing:
+            raise ValueError(f"[反量化失败] prefix={prefix_name} 缺少 {missing}，无法反量化。")
+
+    prefixes = list(prefix_map.items())
+    for i in range(0, len(prefixes), batch_size):
+        batch = prefixes[i:i+batch_size]
+
+        packed_tensors = []
+        scale_tensors = []
+        shape_tensors = []
+        batch_names = []
+
+        for prefix_name, sufmap in batch:
+            packed_tensors.append(state_dict[sufmap[".weight_packed"]])
+            scale_tensors.append(state_dict[sufmap[".weight_scale"]])
+            shape_tensors.append(state_dict[sufmap[".weight_shape"]])
+            batch_names.append(prefix_name)
+
+        # 堆叠 tensor（这里使用 torch.stack，如果 shape 不一致，可保留 list）
+        # 如果 shape 不一致，weight_dequant_int4 需支持 list 输入
+        packed_tensors = [t.to(device, non_blocking=True) for t in packed_tensors]
+        scale_tensors = [t.to(device, non_blocking=True) for t in scale_tensors]
+        shape_tensors = [t.to(device, non_blocking=True) for t in shape_tensors]
+
+        # 批量反量化
+        with torch.no_grad():
+            new_weights = []
+            for pt, st, sht in zip(packed_tensors, scale_tensors, shape_tensors):
+                w = weight_dequant_int4(pt, st, sht).cpu()
+                new_weights.append(w)
+
+        # 保存结果
+        for name, w in zip(batch_names, new_weights):
+            new_sd[name + ".weight"] = w
+
+        # 显存管理
+        del packed_tensors, scale_tensors, shape_tensors, new_weights
+        # torch.cuda.empty_cache()  # 可选
+
+    return new_sd
+
+def tensor_memory_bytes(t: torch.Tensor):
+    return t.numel() * t.element_size()
+
+def state_dict_memory_bytes(sd: dict):
+    return sum(tensor_memory_bytes(v) for v in sd.values())
+
+def bytes_to_human(n):
+    for unit in ["B", "KB", "MB", "GB"]:
+        if n < 1024:
+            return f"{n:.2f}{unit}"
+        n /= 1024
+    return f"{n:.2f}TB"
+
+
+if __name__ == "__main__":
+    from glob import glob
+    from safetensors.torch import load_file
+
+    all_files = glob("/home/mashun/vllm_project/checkpoints/Kimi-K2-Thinking/*.safetensors")
+    f = all_files[-1]
+
+    sd = load_file(f)
+
+    suffixes = ["weight_packed", "weight_scale", "weight_shape"]
+
+    nsd = weight_dequant_int4_wrapper(sd, suffixes, "cuda:0")
+
+    total = state_dict_memory_bytes(nsd)
+    print("显存占用:", bytes_to_human(total))
+
+