Fix absolute imports

flake.lock

@@ -73,11 +73,11 @@
         "nixpkgs": "nixpkgs"
       },
       "locked": {
-        "lastModified": 1750234878,
-        "narHash": "sha256-q9DRC9zdpzUf88qqg1qbhP1qgJbE2cMtn8oUmosuyT8=",
+        "lastModified": 1751968576,
+        "narHash": "sha256-cmKrlWpNTG/hq1bCaHXfbdm9T+Y6V+5//EHAVc1TLBE=",
         "owner": "huggingface",
         "repo": "hf-nix",
-        "rev": "c7132f90763d756da3e77da62e01be0a4546dc57",
+        "rev": "3fcd1e1b46da91b6691261640ffd6b7123d0cb9e",
         "type": "github"
       },
       "original": {
@@ -98,15 +98,16 @@
         ]
       },
       "locked": {
-        "lastModified": 1751630801,
-        "narHash": "sha256-wSVDQZejEifZQfcQF4R7tUd1QJo4wTwu2PMnF3lX0Z0=",
+        "lastModified": 1751968677,
+        "narHash": "sha256-5gtVPN6uk+H3yq2gJRDjSTcaVSgGJZjbMALlO6TBcT8=",
         "owner": "huggingface",
         "repo": "kernel-builder",
-        "rev": "99306a90b243bf603f4d41906bd51af92266fcbd",
+        "rev": "54eea2ce49889202e7018792f407046e36f89bc5",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
+        "ref": "get-kernel-check",
         "repo": "kernel-builder",
         "type": "github"
       }

flake.nix

@@ -2,7 +2,7 @@
   description = "Flake for quantization kernels";
 
   inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder";
+    kernel-builder.url = "github:huggingface/kernel-builder/get-kernel-check";
   };
 
   outputs =

torch-ext/quantization/utils/marlin_utils.py

@@ -6,8 +6,7 @@ from typing import Optional
 import numpy
 import torch
 
-import quantization as ops
-from quantization.scalar_type import ScalarType, scalar_types
+from .. import ScalarType, gptq_marlin_gemm, scalar_types
 
 from .quant_utils import pack_cols, unpack_cols
 
@@ -383,7 +382,7 @@ def apply_gptq_marlin_linear(
                                                   device=input.device,
                                                   dtype=input.dtype)
 
-    output = ops.gptq_marlin_gemm(reshaped_x,
+    output = gptq_marlin_gemm(reshaped_x,
                                   None,
                                   weight,
                                   weight_scale,
@@ -429,7 +428,7 @@ def apply_awq_marlin_linear(
                                                   device=input.device,
                                                   dtype=input.dtype)
 
-    output = ops.gptq_marlin_gemm(reshaped_x,
+    output = gptq_marlin_gemm(reshaped_x,
                                   None,
                                   weight,
                                   weight_scale,

torch-ext/quantization/utils/marlin_utils_fp4.py

@@ -5,12 +5,11 @@ from typing import Optional
 
 import torch
 
-import quantization as ops
-
+from .. import gptq_marlin_gemm, gptq_marlin_repack
 from .marlin_utils import (
     USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_scales,
     should_use_atomic_add_reduce)
-from quantization.scalar_type import scalar_types
+from ..scalar_type import scalar_types
 
 FP4_MARLIN_SUPPORTED_GROUP_SIZES = [16]
 
@@ -90,7 +89,7 @@ def apply_fp4_marlin_linear(
                                                   device=input.device,
                                                   dtype=input.dtype)
 
-    output = ops.gptq_marlin_gemm(a=reshaped_x,
+    output = gptq_marlin_gemm(a=reshaped_x,
                                   c=None,
                                   b_q_weight=weight,
                                   b_scales=weight_scale,
@@ -135,7 +134,7 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
     perm = torch.empty(0, dtype=torch.int, device=device)
     qweight = layer.weight.view(torch.int32).T.contiguous()
 
-    marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
+    marlin_qweight = gptq_marlin_repack(b_q_weight=qweight,
                                             perm=perm,
                                             size_k=part_size_k,
                                             size_n=part_size_n,
@@ -192,7 +191,7 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
         for i in range(e):
             qweight = weight[i].view(torch.int32).T.contiguous()
 
-            marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
+            marlin_qweight = gptq_marlin_repack(b_q_weight=qweight,
                                                     perm=perm,
                                                     size_k=size_k,
                                                     size_n=size_n,
@@ -263,7 +262,7 @@ def rand_marlin_weight_fp4_like(weight, group_size):
     weight_ref = weight_ref * global_scale.to(weight.dtype) * \
         scales.repeat_interleave(group_size, 1).to(weight.dtype)
 
-    marlin_qweight = ops.gptq_marlin_repack(
+    marlin_qweight = gptq_marlin_repack(
         b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
         perm=torch.empty(0, dtype=torch.int, device=device),
         size_k=size_k,

torch-ext/quantization/utils/marlin_utils_fp8.py

@@ -5,7 +5,7 @@ from typing import Optional
 
 import torch
 
-import quantization as ops
+from .. import gptq_marlin_gemm, gptq_marlin_repack
 
 from .marlin_utils import USE_FP32_REDUCE_DEFAULT, marlin_make_workspace, marlin_permute_scales
 
@@ -51,7 +51,7 @@ def apply_fp8_marlin_linear(
                                                   device=input.device,
                                                   dtype=input.dtype)
 
-    output = ops.gptq_marlin_gemm(a=reshaped_x,
+    output = gptq_marlin_gemm(a=reshaped_x,
                                   c=None,
                                   b_q_weight=weight,
                                   b_scales=weight_scale,
@@ -104,7 +104,7 @@ def marlin_quant_fp8_torch(weight, group_size):
         weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
 
     packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous()
-    marlin_qweight = ops.gptq_marlin_repack(
+    marlin_qweight = gptq_marlin_repack(
         b_q_weight=packed_weight,
         perm=torch.empty(0, dtype=torch.int, device=device),
         size_k=size_k,

torch-ext/quantization/utils/marlin_utils_test.py

@@ -5,8 +5,7 @@ from typing import List, Optional
 import numpy as np
 import torch
 
-from quantization.scalar_type import ScalarType
-
+from ..scalar_type import ScalarType
 from .marlin_utils import GPTQ_MARLIN_TILE, marlin_permute_scales, marlin_zero_points
 from .quant_utils import (
     get_pack_factor,

torch-ext/quantization/utils/marlin_utils_test_24.py

@@ -6,8 +6,7 @@ from typing import List
 import numpy
 import torch
 
-from quantization.scalar_type import ScalarType
-
+from ..scalar_type import ScalarType
 from .marlin_utils_test import marlin_weights
 from .quant_utils import gptq_quantize_weights
 

torch-ext/quantization/utils/quant_utils.py

@@ -5,7 +5,7 @@ from typing import List, Optional
 import numpy
 import torch
 
-from quantization.scalar_type import ScalarType, scalar_types
+from ..scalar_type import ScalarType, scalar_types
 
 SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
 SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]