hypertensor 1.0.0__tar.gz

hypertensor-1.0.0/HyperTensor/__init__.py

@@ -0,0 +1,17 @@
+from .formats import PrecisionFormat
+from .weight_options import WeightOption
+from .quantize import scan_model
+from .modules import HyperLinear, wrap_model_with_hypertensor
+from .hyper_attention import HyperAttention
+from .serialize import save_hypertensors, load_hypertensors
+
+__all__ = [
+    "PrecisionFormat",
+    "WeightOption",
+    "scan_model",
+    "HyperLinear",
+    "HyperAttention",
+    "wrap_model_with_hypertensor",
+    "save_hypertensors",
+    "load_hypertensors",
+]

hypertensor-1.0.0/HyperTensor/block_utils.py

@@ -0,0 +1,8 @@
+from typing import Tuple
+
+
+def compute_block_grid(shape: Tuple[int, int], block_size: int) -> Tuple[int, int]:
+    out_features, in_features = shape
+    n_blocks_out = (out_features + block_size - 1) // block_size
+    n_blocks_in = (in_features + block_size - 1) // block_size
+    return n_blocks_out, n_blocks_in

hypertensor-1.0.0/HyperTensor/cuda_backend.py

@@ -0,0 +1,147 @@
+import os
+import torch
+from torch.utils.cpp_extension import load
+
+_src_dir = os.path.dirname(__file__)
+_ext_name = "hypertensor_kernels"
+
+_hypertensor_ext = load(
+    name=_ext_name,
+    sources=[
+        os.path.join(_src_dir, "kernels.cu"),
+        os.path.join(_src_dir, "binding.cpp"),
+    ],
+    verbose=False,
+    extra_cflags=[
+        "-O3",
+        "-DTORCH_DISABLE_DYNAMO",
+    ],
+    extra_cuda_cflags=[
+        "-O3",
+        "--use_fast_math",
+        "-DTORCH_DISABLE_DYNAMO",
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+        "-gencode=arch=compute_89,code=sm_89",
+    ],
+)
+
+
+def _ensure_same_device(a: torch.Tensor, b: torch.Tensor):
+    if a.device != b.device:
+        b = b.to(a.device)
+    return a, b
+
+
+# ---------- FP32 / FP16 / BF16: let cuBLAS do its job ----------
+
+def matmul_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    a, b = _ensure_same_device(a, b)
+    return torch.matmul(a, b)
+
+
+def matmul_fp16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    a, b = _ensure_same_device(a, b)
+    return torch.matmul(a, b)
+
+
+def matmul_bf16(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    a, b = _ensure_same_device(a, b)
+    return torch.matmul(a, b)
+
+
+# ---------- INT8 / INT4 weight GEMM (for HyperLinear) ----------
+
+def matmul_int8(a: torch.Tensor, w_int8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    a, w_int8 = _ensure_same_device(a, w_int8)
+    scale = scale.to(a.device)
+
+    if a.dtype != torch.float16:
+        a = a.to(torch.float16)
+
+    if (
+        a.is_cuda
+        and w_int8.is_cuda
+        and scale.is_cuda
+        and w_int8.dtype == torch.int8
+        and scale.dtype == torch.float32
+    ):
+        return _hypertensor_ext.gemm_int8_fp16_tc(a, w_int8, scale)
+
+    w_f = w_int8.to(torch.float32)
+    if scale.numel() == 1:
+        w_f = w_f * scale.view(1, 1)
+    else:
+        w_f = w_f * scale.view(1, -1)
+
+    return torch.matmul(a.to(torch.float32), w_f)
+
+
+def matmul_int4(a: torch.Tensor,
+                w_int4_packed: torch.Tensor,
+                scale: torch.Tensor,
+                out_features: int) -> torch.Tensor:
+
+    a, w_int4_packed = _ensure_same_device(a, w_int4_packed)
+    scale = scale.to(a.device)
+
+    if a.dtype != torch.float16:
+        a = a.to(torch.float16)
+
+    if (
+        a.is_cuda
+        and w_int4_packed.is_cuda
+        and scale.is_cuda
+        and w_int4_packed.dtype == torch.int8
+        and scale.dtype == torch.float32
+    ):
+        return _hypertensor_ext.gemm_int4_fp16_tc(a, w_int4_packed, scale)
+
+    K = w_int4_packed.size(0)
+    N_packed = w_int4_packed.size(1)
+    N = out_features
+
+    w_full = torch.empty((K, N), dtype=torch.float32, device=a.device)
+
+    for k in range(K):
+        row = w_int4_packed[k]
+        for n in range(N):
+            col_packed = n // 2
+            high = (n % 2) == 1
+            packed = int(row[col_packed].item())
+            if high:
+                v = (packed >> 4) & 0x0F
+            else:
+                v = packed & 0x0F
+            if v & 0x08:
+                v |= 0xF0
+            w_full[k, n] = float(torch.tensor(v, dtype=torch.int8))
+
+    if scale.numel() == 1:
+        w_full = w_full * scale.view(1, 1)
+    else:
+        w_full = w_full * scale.view(1, -1)
+
+    return torch.matmul(a.to(torch.float32), w_full)
+
+
+
+def matmul_bw_dx_fp16(grad_out: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    grad_out, w = _ensure_same_device(grad_out, w)
+    return _hypertensor_ext.gemm_bw_dx_fp16(grad_out, w)
+
+
+def matmul_bw_dw_fp16(grad_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    grad_out, x = _ensure_same_device(grad_out, x)
+    return _hypertensor_ext.gemm_bw_dw_fp16(grad_out, x)
+
+def matmul_bw_dx_bf16(grad_out: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    grad_out, w = _ensure_same_device(grad_out, w)
+    return _hypertensor_ext.gemm_bw_dx_bf16(grad_out, w)
+
+
+def matmul_bw_dw_bf16(grad_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    grad_out, x = _ensure_same_device(grad_out, x)
+    return _hypertensor_ext.gemm_bw_dw_bf16(grad_out, x)

hypertensor-1.0.0/HyperTensor/formats.py

@@ -0,0 +1,118 @@
+from enum import IntEnum
+import torch
+
+
+class PrecisionFormat(IntEnum):
+    INT4 = 1
+    INT8 = 2
+    FP16 = 5
+    BF16 = 6
+    FP32 = 7
+
+
+PRECISION_ORDER = [
+    PrecisionFormat.INT4,
+    PrecisionFormat.INT8,
+    PrecisionFormat.FP16,
+    PrecisionFormat.BF16,
+    PrecisionFormat.FP32,
+]
+
+
+def _can_be_int4(x: torch.Tensor) -> torch.Tensor:
+    return (x == x.round()) & (x >= -8) & (x <= 7)
+
+
+def _can_be_int8(x: torch.Tensor) -> torch.Tensor:
+    return (x == x.round()) & (x >= -128) & (x <= 127)
+
+
+def _roundtrip_error(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return (x - y).abs()
+
+
+def _can_be_fp16(x: torch.Tensor, eps: float = 0.0) -> torch.Tensor:
+    x16 = x.to(torch.float16)
+    x32 = x16.to(torch.float32)
+    if eps == 0.0:
+        return x32 == x
+    else:
+        return _roundtrip_error(x, x32) <= eps
+
+
+def _can_be_bf16(x: torch.Tensor, eps: float = 0.0) -> torch.Tensor:
+    x16 = x.to(torch.bfloat16)
+    x32 = x16.to(torch.float32)
+    if eps == 0.0:
+        return x32 == x
+    else:
+        return _roundtrip_error(x, x32) <= eps
+
+
+def choose_format_for_tensor(
+    t: torch.Tensor,
+    eps_int4: float = 1e-1,
+    eps_int8: float = 5e-2,
+    eps_fp16: float = 1e-3,
+    eps_bf16: float = 1e-3,
+    prefer_bf16: bool = False,
+) -> torch.Tensor:
+    """
+    Per-weight format choice based on approximate fit.
+    - INT4: allowed if quantization error <= eps_int4
+    - INT8: if INT4 too inaccurate, but error <= eps_int8
+    - FP16/BF16: if INT8 too inaccurate, but error <= eps_fp16/eps_bf16
+    - FP32: fallback if nothing else fits well
+    """
+
+    w32 = t.to(torch.float32)
+    fmt = torch.full_like(w32, fill_value=int(PrecisionFormat.FP32), dtype=torch.int8)
+
+    # ----- INT4 -----
+    q4 = w32.round().clamp(-8, 7)
+    err4 = (w32 - q4).abs()
+    mask_int4 = err4 <= eps_int4
+    fmt[mask_int4] = int(PrecisionFormat.INT4)
+
+    # ----- INT8 (only where not INT4) -----
+    q8 = w32.round().clamp(-128, 127)
+    err8 = (w32 - q8).abs()
+    mask_int8 = (err8 <= eps_int8) & ~mask_int4
+    fmt[mask_int8] = int(PrecisionFormat.INT8)
+
+    # ----- FP16 / BF16 (only where not INT4/INT8) -----
+    remaining = ~mask_int4 & ~mask_int8
+
+    if prefer_bf16:
+        # BF16 first
+        x_bf16 = w32.to(torch.bfloat16)
+        x_bf32 = x_bf16.to(torch.float32)
+        err_bf16 = (w32 - x_bf32).abs()
+        mask_bf16 = (err_bf16 <= eps_bf16) & remaining
+        fmt[mask_bf16] = int(PrecisionFormat.BF16)
+
+        remaining = remaining & ~mask_bf16
+
+        x_f16 = w32.to(torch.float16)
+        x_f32 = x_f16.to(torch.float32)
+        err_f16 = (w32 - x_f32).abs()
+        mask_fp16 = (err_f16 <= eps_fp16) & remaining
+        fmt[mask_fp16] = int(PrecisionFormat.FP16)
+    else:
+        # FP16 first
+        x_f16 = w32.to(torch.float16)
+        x_f32 = x_f16.to(torch.float32)
+        err_f16 = (w32 - x_f32).abs()
+        mask_fp16 = (err_f16 <= eps_fp16) & remaining
+        fmt[mask_fp16] = int(PrecisionFormat.FP16)
+
+        remaining = remaining & ~mask_fp16
+
+        x_bf16 = w32.to(torch.bfloat16)
+        x_bf32 = x_bf16.to(torch.float32)
+        err_bf16 = (w32 - x_bf32).abs()
+        mask_bf16 = (err_bf16 <= eps_bf16) & remaining
+        fmt[mask_bf16] = int(PrecisionFormat.BF16)
+
+    # Anything that didn't fit any of the above stays FP32
+    return fmt.to(torch.int8)

hypertensor-1.0.0/HyperTensor/hyper_attention.py

@@ -0,0 +1,115 @@
+# HyperTensor/hyper_attention.py
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from .modules import HyperLinear
+
+
+class HyperAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: Optional[int] = None,
+        bias: bool = False,
+        rope_theta: float = 10000.0,
+        prefer_bf16: bool = False,
+        inference_only: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.head_dim = hidden_size // num_heads
+        self.rope_theta = rope_theta
+
+        self.q_proj = HyperLinear(
+            hidden_size,
+            num_heads * self.head_dim,
+            bias=bias,
+            prefer_bf16=prefer_bf16,
+            inference_only=inference_only,
+        )
+        self.k_proj = HyperLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=bias,
+            prefer_bf16=prefer_bf16,
+            inference_only=inference_only,
+        )
+        self.v_proj = HyperLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=bias,
+            prefer_bf16=prefer_bf16,
+            inference_only=inference_only,
+        )
+        self.o_proj = HyperLinear(
+            num_heads * self.head_dim,
+            hidden_size,
+            bias=bias,
+            prefer_bf16=prefer_bf16,
+            inference_only=inference_only,
+        )
+
+        self.prefer_bf16 = prefer_bf16
+        self.inference_only = inference_only
+
+    def _shape(self, x: torch.Tensor, bsz: int, n_heads: int) -> torch.Tensor:
+        return x.view(bsz, -1, n_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        **kwargs,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        q = self._shape(q, bsz, self.num_heads)          # [B,H,T,D]
+        k = self._shape(k, bsz, self.num_kv_heads)
+        v = self._shape(v, bsz, self.num_kv_heads)
+
+        if past_key_value is not None:
+            k = torch.cat([past_key_value[0], k], dim=2)
+            v = torch.cat([past_key_value[1], v], dim=2)
+
+        present = (k, v) if use_cache else None
+
+        if self.num_kv_heads != self.num_heads:
+            repeat_factor = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(repeat_factor, dim=1)
+            v = v.repeat_interleave(repeat_factor, dim=1)
+
+        dtype = torch.bfloat16 if self.prefer_bf16 else torch.float16
+        q = q.to(dtype)
+        k = k.to(dtype)
+        v = v.to(dtype)
+
+        # Fast FP16 attention using cuBLAS/FlashAttention
+        attn_scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:
+            attn_scores = attn_scores + attention_mask
+
+        attn_weights = torch.softmax(attn_scores, dim=-1)
+        attn_output = torch.matmul(attn_weights, v)
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, -1)
+        attn_output = self.o_proj(attn_output)
+
+        if output_attentions:
+            return attn_output, present, attn_weights
+        return attn_output, present