liangsu9988 commited on Jun 20

Commit

2adb8f7

verified ·

1 Parent(s): a9a6160

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark.py +107 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +223 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +23 -0
build/torch211-cxx11-cu128-x86_64-linux/vl_transformer_primitives/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +223 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch211-cxx11-cu130-x86_64-linux/vl_transformer_primitives/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +223 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu130-x86_64-linux/vl_transformer_primitives/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +223 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +22 -0
build/torch212-cxx11-cu132-x86_64-linux/vl_transformer_primitives/__init__.py +26 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,107 @@

+#!/usr/bin/env python3
+"""Benchmark vl-transformer-primitives against PyTorch eager references."""
+from __future__ import annotations
+import argparse
+import sys
+from pathlib import Path
+import torch
+PACKAGE = Path(__file__).resolve().parents[1]
+TESTS = PACKAGE / "tests"
+sys.path.insert(0, str(TESTS))
+from test_vl_transformer_primitives import (  # noqa: E402
+    load_installed_ops,
+    load_source_ops,
+    make_decode_case,
+    ref_avg_pool,
+    ref_norm_rope,
+)
+def bench(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / iters
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--warmup", type=int, default=50)
+    parser.add_argument("--iters", type=int, default=500)
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    torch.manual_seed(1234)
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    print("| Workload | Shape | FlashRT us | PyTorch eager us | Speedup |")
+    print("|---|---:|---:|---:|---:|")
+    for heads in [1, 4, 8, 16, 32, 40]:
+        q, k, v, q_w, k_w, cos, sin = make_decode_case(heads)
+        fused = bench(
+            lambda: ops.qwen3_q_norm_rope_qstage_bf16(q, q_w, cos, sin),
+            args.warmup,
+            args.iters,
+        )
+        eager = bench(
+            lambda: ref_norm_rope(q, q_w, cos, sin, 1e-6),
+            args.warmup,
+            args.iters,
+        )
+        print(f"| q_norm_rope | heads={heads}, d=128 | {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |")
+        fused = bench(
+            lambda: ops.qwen3_k_norm_rope_kvwrite_bf16(k, v, k_w, cos, sin),
+            args.warmup,
+            args.iters,
+        )
+        eager = bench(
+            lambda: (ref_norm_rope(k, k_w, cos, sin, 1e-6), v.clone()),
+            args.warmup,
+            args.iters,
+        )
+        print(f"| k_norm_rope_vwrite | heads={heads}, d=128 | {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |")
+    for nv, h, w, dim, pool in [
+        (1, 16, 16, 1024, 2),
+        (2, 16, 16, 1152, 2),
+        (4, 16, 16, 2048, 4),
+        (2, 32, 32, 1024, 4),
+    ]:
+        x = torch.randn((nv * h * w, dim), device="cuda", dtype=torch.bfloat16)
+        fused = bench(
+            lambda: ops.avg_pool_vision_tokens_bf16(x, nv, h, w, pool),
+            args.warmup,
+            args.iters,
+        )
+        eager = bench(
+            lambda: ref_avg_pool(x, nv, h, w, pool),
+            args.warmup,
+            args.iters,
+        )
+        print(
+            f"| avg_pool_vision | nv={nv}, h={h}, w={w}, dim={dim}, pool={pool} "
+            f"| {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |"
+        )
+    return 0
+if __name__ == "__main__":
+    raise SystemExit(main())

build/torch211-cxx11-cu128-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,223 @@

+"""FlashRT VL transformer primitive kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_decode_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    name: str,
+) -> None:
+    if x.dim() != 2 or x.shape[1] != 128:
+        raise RuntimeError(f"{name} must have shape (heads, 128)")
+    if weight.shape != (128,):
+        raise RuntimeError("norm weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as input")
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_q_norm_rope_qstage_bf16"))
+def _qwen3_q_norm_rope_qstage_bf16_fake(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+) -> None:
+    _check_decode_rope(q_pre, q_norm_weight, cos, sin, q_out, "q_pre")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_bf16"))
+def _qwen3_k_norm_rope_kvwrite_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    k_cache_dst: torch.Tensor,
+    v_cache_dst: torch.Tensor,
+) -> None:
+    _check_decode_rope(k_pre, k_norm_weight, cos, sin, k_cache_dst, "k_pre")
+    if v_pre.shape != k_pre.shape or v_cache_dst.shape != k_pre.shape:
+        raise RuntimeError("v_pre and v_cache_dst must have shape (n_kv_heads, 128)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_devpos_bf16"))
+def _qwen3_k_norm_rope_kvwrite_devpos_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    eps: float,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+) -> None:
+    if k_pre.dim() != 2 or k_pre.shape[1] != 128:
+        raise RuntimeError("k_pre must have shape (n_kv_heads, 128)")
+    n_kv = k_pre.shape[0]
+    if v_pre.shape != k_pre.shape:
+        raise RuntimeError("v_pre must have shape (n_kv_heads, 128)")
+    if k_norm_weight.shape != (128,):
+        raise RuntimeError("k_norm_weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if cur_pos.numel() != 1:
+        raise RuntimeError("cur_pos must have one int32 element")
+    if k_cache.dim() != 3 or k_cache.shape[1:] != (n_kv, 128):
+        raise RuntimeError("k_cache must have shape (max_seq_len, n_kv_heads, 128)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache must have the same shape as k_cache")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool_vision_tokens_bf16"))
+def _avg_pool_vision_tokens_bf16_fake(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (nv * h * w, dim)")
+    dim = x.shape[1]
+    if x.shape[0] != nv * h * w:
+        raise RuntimeError("x.shape[0] must equal nv * h * w")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise RuntimeError("h and w must be divisible by pool_factor")
+    expected = (nv * (h // pool_factor) * (w // pool_factor), dim)
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled shape")
+    return None
+def qwen3_q_norm_rope_qstage_bf16(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Apply Q RMSNorm + full RoPE and write decode staging output.
+    Inputs are contiguous BF16 tensors. ``q_pre`` has shape
+    ``(n_q_heads, 128)``. ``q_norm_weight`` has shape ``(128,)``.
+    ``cos`` and ``sin`` have shape ``(64,)``.
+    """
+    if out is None:
+        out = torch.empty_like(q_pre)
+    ops.qwen3_q_norm_rope_qstage_bf16(q_pre, q_norm_weight, cos, sin, float(eps), out)
+    return out
+def qwen3_k_norm_rope_kvwrite_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    k_cache_dst: Optional[torch.Tensor] = None,
+    v_cache_dst: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply K RMSNorm + full RoPE and copy V into the current KV slot."""
+    if k_cache_dst is None:
+        k_cache_dst = torch.empty_like(k_pre)
+    if v_cache_dst is None:
+        v_cache_dst = torch.empty_like(v_pre)
+    ops.qwen3_k_norm_rope_kvwrite_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        float(eps),
+        k_cache_dst,
+        v_cache_dst,
+    )
+    return k_cache_dst, v_cache_dst
+def qwen3_k_norm_rope_kvwrite_devpos_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Device-position KV write for CUDA Graph decode loops."""
+    ops.qwen3_k_norm_rope_kvwrite_devpos_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        cur_pos,
+        float(eps),
+        k_cache,
+        v_cache,
+    )
+    return k_cache, v_cache
+def avg_pool_vision_tokens_bf16(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Average-pool flattened vision tokens in the spatial grid.
+    ``x`` is ``(nv * h * w, dim)`` BF16. The result is
+    ``(nv * h / pool_factor * w / pool_factor, dim)`` BF16.
+    """
+    if x.dim() != 2:
+        raise ValueError("x must have shape (nv * h * w, dim)")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise ValueError("h and w must be divisible by pool_factor")
+    rows = nv * (h // pool_factor) * (w // pool_factor)
+    if out is None:
+        out = torch.empty((rows, x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.avg_pool_vision_tokens_bf16(x, int(nv), int(h), int(w), int(pool_factor), out)
+    return out
+__all__ = [
+    "avg_pool_vision_tokens_bf16",
+    "qwen3_k_norm_rope_kvwrite_bf16",
+    "qwen3_k_norm_rope_kvwrite_devpos_bf16",
+    "qwen3_q_norm_rope_qstage_bf16",
+]

build/torch211-cxx11-cu128-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _vl_transformer_primitives_cuda_5596053
+ops = torch.ops._vl_transformer_primitives_cuda_5596053
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_vl_transformer_primitives_cuda_5596053::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:1a659ff89026ca179c7d220e9b761d6a31ce5159182f9b9956ba259b89f26e30
+size 535240

build/torch211-cxx11-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,23 @@

+{
+  "name": "vl-transformer-primitives",
+  "id": "_vl_transformer_primitives_cuda_5596053",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/vl_transformer_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,223 @@

+"""FlashRT VL transformer primitive kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_decode_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    name: str,
+) -> None:
+    if x.dim() != 2 or x.shape[1] != 128:
+        raise RuntimeError(f"{name} must have shape (heads, 128)")
+    if weight.shape != (128,):
+        raise RuntimeError("norm weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as input")
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_q_norm_rope_qstage_bf16"))
+def _qwen3_q_norm_rope_qstage_bf16_fake(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+) -> None:
+    _check_decode_rope(q_pre, q_norm_weight, cos, sin, q_out, "q_pre")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_bf16"))
+def _qwen3_k_norm_rope_kvwrite_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    k_cache_dst: torch.Tensor,
+    v_cache_dst: torch.Tensor,
+) -> None:
+    _check_decode_rope(k_pre, k_norm_weight, cos, sin, k_cache_dst, "k_pre")
+    if v_pre.shape != k_pre.shape or v_cache_dst.shape != k_pre.shape:
+        raise RuntimeError("v_pre and v_cache_dst must have shape (n_kv_heads, 128)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_devpos_bf16"))
+def _qwen3_k_norm_rope_kvwrite_devpos_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    eps: float,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+) -> None:
+    if k_pre.dim() != 2 or k_pre.shape[1] != 128:
+        raise RuntimeError("k_pre must have shape (n_kv_heads, 128)")
+    n_kv = k_pre.shape[0]
+    if v_pre.shape != k_pre.shape:
+        raise RuntimeError("v_pre must have shape (n_kv_heads, 128)")
+    if k_norm_weight.shape != (128,):
+        raise RuntimeError("k_norm_weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if cur_pos.numel() != 1:
+        raise RuntimeError("cur_pos must have one int32 element")
+    if k_cache.dim() != 3 or k_cache.shape[1:] != (n_kv, 128):
+        raise RuntimeError("k_cache must have shape (max_seq_len, n_kv_heads, 128)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache must have the same shape as k_cache")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool_vision_tokens_bf16"))
+def _avg_pool_vision_tokens_bf16_fake(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (nv * h * w, dim)")
+    dim = x.shape[1]
+    if x.shape[0] != nv * h * w:
+        raise RuntimeError("x.shape[0] must equal nv * h * w")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise RuntimeError("h and w must be divisible by pool_factor")
+    expected = (nv * (h // pool_factor) * (w // pool_factor), dim)
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled shape")
+    return None
+def qwen3_q_norm_rope_qstage_bf16(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Apply Q RMSNorm + full RoPE and write decode staging output.
+    Inputs are contiguous BF16 tensors. ``q_pre`` has shape
+    ``(n_q_heads, 128)``. ``q_norm_weight`` has shape ``(128,)``.
+    ``cos`` and ``sin`` have shape ``(64,)``.
+    """
+    if out is None:
+        out = torch.empty_like(q_pre)
+    ops.qwen3_q_norm_rope_qstage_bf16(q_pre, q_norm_weight, cos, sin, float(eps), out)
+    return out
+def qwen3_k_norm_rope_kvwrite_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    k_cache_dst: Optional[torch.Tensor] = None,
+    v_cache_dst: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply K RMSNorm + full RoPE and copy V into the current KV slot."""
+    if k_cache_dst is None:
+        k_cache_dst = torch.empty_like(k_pre)
+    if v_cache_dst is None:
+        v_cache_dst = torch.empty_like(v_pre)
+    ops.qwen3_k_norm_rope_kvwrite_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        float(eps),
+        k_cache_dst,
+        v_cache_dst,
+    )
+    return k_cache_dst, v_cache_dst
+def qwen3_k_norm_rope_kvwrite_devpos_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Device-position KV write for CUDA Graph decode loops."""
+    ops.qwen3_k_norm_rope_kvwrite_devpos_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        cur_pos,
+        float(eps),
+        k_cache,
+        v_cache,
+    )
+    return k_cache, v_cache
+def avg_pool_vision_tokens_bf16(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Average-pool flattened vision tokens in the spatial grid.
+    ``x`` is ``(nv * h * w, dim)`` BF16. The result is
+    ``(nv * h / pool_factor * w / pool_factor, dim)`` BF16.
+    """
+    if x.dim() != 2:
+        raise ValueError("x must have shape (nv * h * w, dim)")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise ValueError("h and w must be divisible by pool_factor")
+    rows = nv * (h // pool_factor) * (w // pool_factor)
+    if out is None:
+        out = torch.empty((rows, x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.avg_pool_vision_tokens_bf16(x, int(nv), int(h), int(w), int(pool_factor), out)
+    return out
+__all__ = [
+    "avg_pool_vision_tokens_bf16",
+    "qwen3_k_norm_rope_kvwrite_bf16",
+    "qwen3_k_norm_rope_kvwrite_devpos_bf16",
+    "qwen3_q_norm_rope_qstage_bf16",
+]

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

	@@ -0,0 +1,9 @@

+import torch
+from . import _vl_transformer_primitives_cuda_5596053
+ops = torch.ops._vl_transformer_primitives_cuda_5596053
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_vl_transformer_primitives_cuda_5596053::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:54ffd2375eea20a8bf09f2832eba1eede35711e7cf9c370fcc3929b5d04038f4
+size 519384

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

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

+{
+  "name": "vl-transformer-primitives",
+  "id": "_vl_transformer_primitives_cuda_5596053",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,223 @@

+"""FlashRT VL transformer primitive kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_decode_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    name: str,
+) -> None:
+    if x.dim() != 2 or x.shape[1] != 128:
+        raise RuntimeError(f"{name} must have shape (heads, 128)")
+    if weight.shape != (128,):
+        raise RuntimeError("norm weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as input")
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_q_norm_rope_qstage_bf16"))
+def _qwen3_q_norm_rope_qstage_bf16_fake(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+) -> None:
+    _check_decode_rope(q_pre, q_norm_weight, cos, sin, q_out, "q_pre")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_bf16"))
+def _qwen3_k_norm_rope_kvwrite_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    k_cache_dst: torch.Tensor,
+    v_cache_dst: torch.Tensor,
+) -> None:
+    _check_decode_rope(k_pre, k_norm_weight, cos, sin, k_cache_dst, "k_pre")
+    if v_pre.shape != k_pre.shape or v_cache_dst.shape != k_pre.shape:
+        raise RuntimeError("v_pre and v_cache_dst must have shape (n_kv_heads, 128)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_devpos_bf16"))
+def _qwen3_k_norm_rope_kvwrite_devpos_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    eps: float,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+) -> None:
+    if k_pre.dim() != 2 or k_pre.shape[1] != 128:
+        raise RuntimeError("k_pre must have shape (n_kv_heads, 128)")
+    n_kv = k_pre.shape[0]
+    if v_pre.shape != k_pre.shape:
+        raise RuntimeError("v_pre must have shape (n_kv_heads, 128)")
+    if k_norm_weight.shape != (128,):
+        raise RuntimeError("k_norm_weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if cur_pos.numel() != 1:
+        raise RuntimeError("cur_pos must have one int32 element")
+    if k_cache.dim() != 3 or k_cache.shape[1:] != (n_kv, 128):
+        raise RuntimeError("k_cache must have shape (max_seq_len, n_kv_heads, 128)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache must have the same shape as k_cache")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool_vision_tokens_bf16"))
+def _avg_pool_vision_tokens_bf16_fake(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (nv * h * w, dim)")
+    dim = x.shape[1]
+    if x.shape[0] != nv * h * w:
+        raise RuntimeError("x.shape[0] must equal nv * h * w")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise RuntimeError("h and w must be divisible by pool_factor")
+    expected = (nv * (h // pool_factor) * (w // pool_factor), dim)
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled shape")
+    return None
+def qwen3_q_norm_rope_qstage_bf16(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Apply Q RMSNorm + full RoPE and write decode staging output.
+    Inputs are contiguous BF16 tensors. ``q_pre`` has shape
+    ``(n_q_heads, 128)``. ``q_norm_weight`` has shape ``(128,)``.
+    ``cos`` and ``sin`` have shape ``(64,)``.
+    """
+    if out is None:
+        out = torch.empty_like(q_pre)
+    ops.qwen3_q_norm_rope_qstage_bf16(q_pre, q_norm_weight, cos, sin, float(eps), out)
+    return out
+def qwen3_k_norm_rope_kvwrite_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    k_cache_dst: Optional[torch.Tensor] = None,
+    v_cache_dst: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply K RMSNorm + full RoPE and copy V into the current KV slot."""
+    if k_cache_dst is None:
+        k_cache_dst = torch.empty_like(k_pre)
+    if v_cache_dst is None:
+        v_cache_dst = torch.empty_like(v_pre)
+    ops.qwen3_k_norm_rope_kvwrite_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        float(eps),
+        k_cache_dst,
+        v_cache_dst,
+    )
+    return k_cache_dst, v_cache_dst
+def qwen3_k_norm_rope_kvwrite_devpos_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Device-position KV write for CUDA Graph decode loops."""
+    ops.qwen3_k_norm_rope_kvwrite_devpos_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        cur_pos,
+        float(eps),
+        k_cache,
+        v_cache,
+    )
+    return k_cache, v_cache
+def avg_pool_vision_tokens_bf16(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Average-pool flattened vision tokens in the spatial grid.
+    ``x`` is ``(nv * h * w, dim)`` BF16. The result is
+    ``(nv * h / pool_factor * w / pool_factor, dim)`` BF16.
+    """
+    if x.dim() != 2:
+        raise ValueError("x must have shape (nv * h * w, dim)")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise ValueError("h and w must be divisible by pool_factor")
+    rows = nv * (h // pool_factor) * (w // pool_factor)
+    if out is None:
+        out = torch.empty((rows, x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.avg_pool_vision_tokens_bf16(x, int(nv), int(h), int(w), int(pool_factor), out)
+    return out
+__all__ = [
+    "avg_pool_vision_tokens_bf16",
+    "qwen3_k_norm_rope_kvwrite_bf16",
+    "qwen3_k_norm_rope_kvwrite_devpos_bf16",
+    "qwen3_q_norm_rope_qstage_bf16",
+]

build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _vl_transformer_primitives_cuda_5596053
+ops = torch.ops._vl_transformer_primitives_cuda_5596053
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_vl_transformer_primitives_cuda_5596053::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so ADDED Viewed

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

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

build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "vl-transformer-primitives",
+  "id": "_vl_transformer_primitives_cuda_5596053",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/vl_transformer_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,223 @@

+"""FlashRT VL transformer primitive kernels."""
+from __future__ import annotations
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+def _check_decode_rope(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    name: str,
+) -> None:
+    if x.dim() != 2 or x.shape[1] != 128:
+        raise RuntimeError(f"{name} must have shape (heads, 128)")
+    if weight.shape != (128,):
+        raise RuntimeError("norm weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if out.shape != x.shape:
+        raise RuntimeError("out must have the same shape as input")
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_q_norm_rope_qstage_bf16"))
+def _qwen3_q_norm_rope_qstage_bf16_fake(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    q_out: torch.Tensor,
+) -> None:
+    _check_decode_rope(q_pre, q_norm_weight, cos, sin, q_out, "q_pre")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_bf16"))
+def _qwen3_k_norm_rope_kvwrite_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float,
+    k_cache_dst: torch.Tensor,
+    v_cache_dst: torch.Tensor,
+) -> None:
+    _check_decode_rope(k_pre, k_norm_weight, cos, sin, k_cache_dst, "k_pre")
+    if v_pre.shape != k_pre.shape or v_cache_dst.shape != k_pre.shape:
+        raise RuntimeError("v_pre and v_cache_dst must have shape (n_kv_heads, 128)")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qwen3_k_norm_rope_kvwrite_devpos_bf16"))
+def _qwen3_k_norm_rope_kvwrite_devpos_bf16_fake(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    eps: float,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+) -> None:
+    if k_pre.dim() != 2 or k_pre.shape[1] != 128:
+        raise RuntimeError("k_pre must have shape (n_kv_heads, 128)")
+    n_kv = k_pre.shape[0]
+    if v_pre.shape != k_pre.shape:
+        raise RuntimeError("v_pre must have shape (n_kv_heads, 128)")
+    if k_norm_weight.shape != (128,):
+        raise RuntimeError("k_norm_weight must have shape (128,)")
+    if cos.shape != (64,) or sin.shape != (64,):
+        raise RuntimeError("cos and sin must have shape (64,)")
+    if cur_pos.numel() != 1:
+        raise RuntimeError("cur_pos must have one int32 element")
+    if k_cache.dim() != 3 or k_cache.shape[1:] != (n_kv, 128):
+        raise RuntimeError("k_cache must have shape (max_seq_len, n_kv_heads, 128)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache must have the same shape as k_cache")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool_vision_tokens_bf16"))
+def _avg_pool_vision_tokens_bf16_fake(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    out: torch.Tensor,
+) -> None:
+    if x.dim() != 2:
+        raise RuntimeError("x must have shape (nv * h * w, dim)")
+    dim = x.shape[1]
+    if x.shape[0] != nv * h * w:
+        raise RuntimeError("x.shape[0] must equal nv * h * w")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise RuntimeError("h and w must be divisible by pool_factor")
+    expected = (nv * (h // pool_factor) * (w // pool_factor), dim)
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled shape")
+    return None
+def qwen3_q_norm_rope_qstage_bf16(
+    q_pre: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Apply Q RMSNorm + full RoPE and write decode staging output.
+    Inputs are contiguous BF16 tensors. ``q_pre`` has shape
+    ``(n_q_heads, 128)``. ``q_norm_weight`` has shape ``(128,)``.
+    ``cos`` and ``sin`` have shape ``(64,)``.
+    """
+    if out is None:
+        out = torch.empty_like(q_pre)
+    ops.qwen3_q_norm_rope_qstage_bf16(q_pre, q_norm_weight, cos, sin, float(eps), out)
+    return out
+def qwen3_k_norm_rope_kvwrite_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+    k_cache_dst: Optional[torch.Tensor] = None,
+    v_cache_dst: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply K RMSNorm + full RoPE and copy V into the current KV slot."""
+    if k_cache_dst is None:
+        k_cache_dst = torch.empty_like(k_pre)
+    if v_cache_dst is None:
+        v_cache_dst = torch.empty_like(v_pre)
+    ops.qwen3_k_norm_rope_kvwrite_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        float(eps),
+        k_cache_dst,
+        v_cache_dst,
+    )
+    return k_cache_dst, v_cache_dst
+def qwen3_k_norm_rope_kvwrite_devpos_bf16(
+    k_pre: torch.Tensor,
+    v_pre: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    cur_pos: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Device-position KV write for CUDA Graph decode loops."""
+    ops.qwen3_k_norm_rope_kvwrite_devpos_bf16(
+        k_pre,
+        v_pre,
+        k_norm_weight,
+        cos,
+        sin,
+        cur_pos,
+        float(eps),
+        k_cache,
+        v_cache,
+    )
+    return k_cache, v_cache
+def avg_pool_vision_tokens_bf16(
+    x: torch.Tensor,
+    nv: int,
+    h: int,
+    w: int,
+    pool_factor: int,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Average-pool flattened vision tokens in the spatial grid.
+    ``x`` is ``(nv * h * w, dim)`` BF16. The result is
+    ``(nv * h / pool_factor * w / pool_factor, dim)`` BF16.
+    """
+    if x.dim() != 2:
+        raise ValueError("x must have shape (nv * h * w, dim)")
+    if h % pool_factor != 0 or w % pool_factor != 0:
+        raise ValueError("h and w must be divisible by pool_factor")
+    rows = nv * (h // pool_factor) * (w // pool_factor)
+    if out is None:
+        out = torch.empty((rows, x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.avg_pool_vision_tokens_bf16(x, int(nv), int(h), int(w), int(pool_factor), out)
+    return out
+__all__ = [
+    "avg_pool_vision_tokens_bf16",
+    "qwen3_k_norm_rope_kvwrite_bf16",
+    "qwen3_k_norm_rope_kvwrite_devpos_bf16",
+    "qwen3_q_norm_rope_qstage_bf16",
+]

build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _vl_transformer_primitives_cuda_5596053
+ops = torch.ops._vl_transformer_primitives_cuda_5596053
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_vl_transformer_primitives_cuda_5596053::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/_vl_transformer_primitives_cuda_5596053.abi3.so ADDED Viewed

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

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

build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

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

+{
+  "name": "vl-transformer-primitives",
+  "id": "_vl_transformer_primitives_cuda_5596053",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/vl_transformer_primitives/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))