addon for cuda support (pytorch#20822)

Differential Revision: D110282210

Pull Request resolved: pytorch#20822

Author: Gasoonjia

backends/cuda/TARGETS

@@ -3,20 +3,64 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 oncall("executorch")
 
 runtime.python_library(
-    name = "quantize_op_dispatch",
+    name = "coalesced_int4_tensor",
     srcs = [
         "coalesced_int4_tensor.py",
+    ],
+    visibility = [
+        "//executorch/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/ao:torchao",
+    ],
+)
+
+runtime.python_library(
+    name = "dp4a_planar_int5_tensor",
+    srcs = [
+        "dp4a_planar_int5_tensor.py",
+    ],
+    visibility = [
+        "//executorch/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/ao:torchao",
+    ],
+)
+
+runtime.python_library(
+    name = "dp4a_planar_int6_tensor",
+    srcs = [
         "dp4a_planar_int6_tensor.py",
+    ],
+    visibility = [
+        "//executorch/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/ao:torchao",
+    ],
+)
+
+runtime.python_library(
+    name = "quantize_op_dispatch",
+    srcs = [
         "quantize_op_dispatch/__init__.py",
         "quantize_op_dispatch/_library.py",
         "quantize_op_dispatch/int4_dispatch.py",
+        "quantize_op_dispatch/int5_dispatch.py",
         "quantize_op_dispatch/int6_dispatch.py",
         "quantize_op_dispatch/int8_dispatch.py",
     ],
     visibility = [
         "//executorch/...",
     ],
     deps = [
+        ":coalesced_int4_tensor",
+        ":dp4a_planar_int5_tensor",
+        ":dp4a_planar_int6_tensor",
         "//caffe2:torch",
         "//pytorch/ao:torchao",
     ],
@@ -80,6 +124,7 @@ runtime.python_library(
         "triton/kernels/int4_matmul.py",
         "triton/kernels/sdpa.py",
         "triton/kernels/topk.py",
+        "triton/kernels/tq4_sdpa.py",
     ],
     visibility = [
         "//executorch/backends/cuda/...",

backends/cuda/cuda_backend.py

@@ -385,6 +385,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]:
             "aoti_torch_cuda_randint_low_out": None,
             "executorch_cuda::int4_plain_mm": None,
             "aoti_torch_cuda_int4_plain_mm": None,
+            "executorch_cuda::int5_plain_mm": None,
+            "aoti_torch_cuda_int5_plain_mm": None,
             "executorch_cuda::int6_plain_mm": None,
             "aoti_torch_cuda_int6_plain_mm": None,
             "executorch_cuda::int8_plain_mm": None,

extension/llm/export/BUCK

@@ -59,6 +59,26 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "gguf",
+    srcs = [
+        "gguf.py",
+    ],
+    _is_external_target = True,
+    base_module = "executorch.extension.llm.export",
+    visibility = [
+        "//executorch/backends/...",
+        "//executorch/examples/...",
+        "//executorch/extension/llm/...",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/ao:torchao",
+        "fbsource//third-party/pypi/gguf:gguf",
+        "fbsource//third-party/pypi/numpy:numpy",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_binary,
     name = "export_llm",
     srcs = [

extension/llm/modules/BUCK

@@ -37,6 +37,18 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "turboquant",
+    srcs = [
+        "turboquant/__init__.py",
+        "turboquant/codebook.py",
+        "turboquant/kv_cache.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "module_lib",
     srcs = [