Add torch211-cxx11-cu130-aarch64-linux SM110 artifact

Browse files

Files changed (5) hide show

build/torch211-cxx11-cu130-aarch64-linux/__init__.py +122 -0
build/torch211-cxx11-cu130-aarch64-linux/_flashrt_fused_quant_cuda_c4d802d.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch211-cxx11-cu130-aarch64-linux/flashrt_fused_quant/__init__.py +14 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +22 -0

build/torch211-cxx11-cu130-aarch64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,122 @@

+"""FlashRT fused activation and quantization kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_quant_nvfp4_swizzled_bf16"))
+def _silu_mul_quant_nvfp4_swizzled_bf16_fake(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    packed: torch.Tensor,
+    scales: torch.Tensor,
+) -> None:
+    if packed.shape != (gate.shape[0], gate.shape[1] // 2):
+        raise RuntimeError("packed shape must be (rows, cols / 2)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("silu_mul_merged_quant_nvfp4_swizzled_bf16"))
+def _silu_mul_merged_quant_nvfp4_swizzled_bf16_fake(
+    merged_gate_up: torch.Tensor,
+    packed: torch.Tensor,
+    scales: torch.Tensor,
+) -> None:
+    rows, merged_cols = merged_gate_up.shape
+    if packed.shape != (rows, merged_cols // 4):
+        raise RuntimeError("packed shape must be (rows, merged_cols / 4)")
+    return None
+def nvfp4_swizzled_scale_bytes(rows: int, cols: int) -> int:
+    """Return byte count for a CUTLASS Sm1xx NVFP4 swizzled scale buffer."""
+    if rows <= 0:
+        raise ValueError("rows must be positive")
+    if cols <= 0 or cols % 16 != 0:
+        raise ValueError("cols must be positive and divisible by 16")
+    n_blocks = cols // 16
+    return ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 512
+def _allocate_outputs(
+    rows: int,
+    cols: int,
+    device: torch.device,
+    packed: Optional[torch.Tensor],
+    scales: Optional[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if packed is None:
+        packed = torch.empty((rows, cols // 2), device=device, dtype=torch.uint8)
+    if scales is None:
+        scales = torch.zeros(
+            (nvfp4_swizzled_scale_bytes(rows, cols),),
+            device=device,
+            dtype=torch.uint8,
+        )
+    return packed, scales
+def silu_mul_quant_nvfp4_swizzled_bf16(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    scales: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute ``SiLU(gate) * up`` and quantize to NVFP4 swizzled layout.
+    ``gate`` and ``up`` must be contiguous BF16 tensors with shape
+    ``(rows, cols)``. ``cols`` must be divisible by 16. The returned ``packed``
+    tensor has shape ``(rows, cols / 2)`` and dtype ``torch.uint8``. The
+    returned ``scales`` tensor is a flat CUTLASS Sm1xx swizzled UE4M3
+    scale-factor byte buffer.
+    """
+    if gate.dim() != 2:
+        raise ValueError("gate must have shape (rows, cols)")
+    rows, cols = gate.shape
+    packed, scales = _allocate_outputs(rows, cols, gate.device, packed, scales)
+    ops.silu_mul_quant_nvfp4_swizzled_bf16(gate, up, packed, scales)
+    return packed, scales
+def silu_mul_merged_quant_nvfp4_swizzled_bf16(
+    merged_gate_up: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    scales: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute merged ``[gate | up]`` SiLU product and NVFP4 quantization.
+    ``merged_gate_up`` must be contiguous BF16 with shape ``(rows, 2 * cols)``
+    and row layout ``[gate | up]``.
+    """
+    if merged_gate_up.dim() != 2:
+        raise ValueError("merged_gate_up must have shape (rows, 2 * cols)")
+    rows, merged_cols = merged_gate_up.shape
+    if merged_cols % 2 != 0:
+        raise ValueError("merged_gate_up.shape[1] must be even")
+    cols = merged_cols // 2
+    packed, scales = _allocate_outputs(
+        rows, cols, merged_gate_up.device, packed, scales
+    )
+    ops.silu_mul_merged_quant_nvfp4_swizzled_bf16(
+        merged_gate_up,
+        packed,
+        scales,
+    )
+    return packed, scales
+__all__ = [
+    "nvfp4_swizzled_scale_bytes",
+    "silu_mul_quant_nvfp4_swizzled_bf16",
+    "silu_mul_merged_quant_nvfp4_swizzled_bf16",
+]

build/torch211-cxx11-cu130-aarch64-linux/_flashrt_fused_quant_cuda_c4d802d.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:f6d95803f9bd5045b0c93677efa93f216d4dc726a006a0302c0c07af917d3c88
+size 238224

build/torch211-cxx11-cu130-aarch64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,6 @@

+import torch
+from . import _flashrt_fused_quant_cuda_c4d802d
+ops = torch.ops._flashrt_fused_quant_cuda_c4d802d
+def add_op_namespace_prefix(op_name: str):
+    return f"_flashrt_fused_quant_cuda_c4d802d::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/flashrt_fused_quant/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

build/torch211-cxx11-cu130-aarch64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,22 @@

+{
+  "name": "flashrt-fused-quant",
+  "id": "_flashrt_fused_quant_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "O7W74GfPuHsHXV2HuZEIUUcyhmlCL+2rOXdYH4eW4VE=",
+      "_flashrt_fused_quant_cuda_c4d802d.abi3.so": "9tlYA/m9UEWwyTZ376k/IW1NxyagBqAwLAwHr5F9PIg=",
+      "_ops.py": "568RL653nmkbqVCzLNOwZsqHvzbRpUvL8pCVpSTGlYQ=",
+      "flashrt_fused_quant/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}