Promote latest kernel artifacts to main

README.md

@@ -1,9 +1,98 @@
-# flashrt/grouped-moe-gemv
+# grouped-moe-gemv
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+FlashRT native CUDA grouped expert projection kernels for Blackwell decode and
+small verify batches. Version 2 adds W4A4 with device-side top-k routing while
+preserving the version 1 W4A16 APIs.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/grouped-moe-gemv
+## Hardware Backends
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+- SM110 (Jetson AGX Thor): W4A16 decode and grouped expert GEMV use the
+  FlashRT edge backend validated by FlashRT PR #169. This target is compiled
+  independently with `FLASHRT_W4A16_EDGE_UNROLL=2`; the SM120 value remains 4.
+- SM120/SM121: W4A16 and block-scaled-MMA W4A4 paths are available.
+- W4A4 is intentionally rejected on SM110 because that implementation uses
+  the SM120 block-scaled MMA path. It never silently falls back or launches an
+  incompatible cubin.
+
+## Functions
+
+- `w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, alpha=1.0, out=None)`
+- `grouped_w4a16_gemv_bf16(activations, weight_stack, sfb_stack, alpha_stack, expert_idx, n, w_stride=None, sfb_stride=None, out=None)`
+- `quantize_activations_nvfp4_bf16(activations, packed=None, sfa=None)`
+- `quantize_weights_nvfp4_bf16(weights, packed=None, sfb=None)`
+- `grouped_w4a4_gemv_bf16(activations_packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out=None)`
+- `grouped_w4a4_gemv_from_bf16(activations, weight_stack, sfb_stack, alpha_stack, expert_idx, packed=None, sfa=None, out=None)`
+
+The grouped API runs one BF16-activation x NVFP4-weight GEMV per routed slot.
+It is intended for static routed expert batches where the caller already owns
+packed weights and swizzled scale-factor buffers.
+
+On SM120/SM121, the W4A4 API accepts packed activations `[M,K/2]`, expert weights
+`[E,N,K/2]`, and a contiguous device routing tensor `[M,top_k]`. It emits
+`[M,top_k,N]` in one grouped compute launch. For down projections with a
+different activation per routed pair, flatten to `M=routed_pairs, top_k=1`.
+
+`K` must be divisible by 16 and `N` by 8. Target `K%64==0` shapes use tuned
+SM120 paths; the remaining `K%16` shapes use a fixed-order SIMT contract path.
+No atomics, host synchronization, or dynamic workspace are used by the native
+ops. Pass `packed`, `sfa`, and `out` buffers to the composed helper for
+allocation-free CUDA Graph capture.
+
+## Example
+
+```python
+from kernels import get_kernel
+import torch
+
+try:
+    moe = get_kernel(
+        "flashrt/grouped-moe-gemv", version=2, trust_remote_code=True
+    )
+except TypeError:  # kernels==0.12.x compatibility
+    moe = get_kernel("flashrt/grouped-moe-gemv", version=2)
+
+M, TOP_K, E, N, K = 7, 8, 8, 1024, 2048
+x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16)
+expert_idx = torch.randint(E, (M, TOP_K), device="cuda", dtype=torch.int32)
+
+def sf_bytes(rows, dim):
+    return ((rows + 127) // 128) * (((dim // 16) + 3) // 4) * 512
+
+# Do this once while loading the checkpoint, not in the inference hot path.
+weights_bf16 = torch.randn(E, N, K, device="cuda", dtype=torch.bfloat16)
+weights_packed = torch.empty(E, N, K // 2, device="cuda", dtype=torch.uint8)
+weight_sfs = torch.empty(E, sf_bytes(N, K), device="cuda", dtype=torch.uint8)
+for expert in range(E):
+    moe.quantize_weights_nvfp4_bf16(
+        weights_bf16[expert],
+        packed=weights_packed[expert],
+        sfb=weight_sfs[expert],
+    )
+weight_alpha = torch.ones(E, device="cuda", dtype=torch.float32)
+
+packed = torch.empty(M, K // 2, device="cuda", dtype=torch.uint8)
+sfa = torch.empty(sf_bytes(M, K), device="cuda", dtype=torch.uint8)
+out = torch.empty(M, TOP_K, N, device="cuda", dtype=torch.bfloat16)
+y = moe.grouped_w4a4_gemv_from_bf16(
+    x, weights_packed, weight_sfs, weight_alpha, expert_idx,
+    packed=packed, sfa=sfa, out=out,
+)
+```
+
+The example buffers may be larger than the minimum; wrappers validate storage.
+For production code, derive SF sizes from the checkpoint packer metadata.
+
+## Dispatch guidance
+
+Use the built-artifact benchmark to dispatch rather than selecting only by
+dtype. On the tested cu128 artifact W4A16 wins gate-up, W4A4 wins down verify,
+and down decode is effectively tied kernel-only. A fused/upstream FP4 producer
+removes the standalone quantization charge, but callers still should not assume
+lower precision is automatically faster.
+
+## Validation
+
+```bash
+python grouped-moe-gemv/tests/test_grouped_moe_gemv.py --backend source --mode full
+python grouped-moe-gemv/benchmarks/benchmark.py --backend source
+```

build/torch213-cxx11-cu130-x86_64-linux/__init__.py

@@ -0,0 +1,240 @@
+"""FlashRT grouped MoE GEMV 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("w4a16_decode_gemv_bf16"))
+def _w4a16_decode_gemv_fake(
+    x_bf16: torch.Tensor,
+    weight_packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
+    if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
+        raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
+def _grouped_w4a16_gemv_fake(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    w_stride: int,
+    sfb_stride: int,
+    out: torch.Tensor,
+) -> None:
+    if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
+        raise RuntimeError("expected activations (slots,K), out (slots,N)")
+    if expert_idx.shape != (activations.shape[0],):
+        raise RuntimeError("expert_idx must have shape (slots,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_activations_nvfp4_bf16"))
+def _quantize_activations_nvfp4_fake(
+    activations: torch.Tensor,
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+) -> None:
+    if activations.dim() != 2:
+        raise RuntimeError("activations must have shape (M,K)")
+    if packed.shape != (activations.shape[0], activations.shape[1] // 2):
+        raise RuntimeError("packed must have shape (M,K/2)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_weights_nvfp4_bf16"))
+def _quantize_weights_nvfp4_fake(
+    weights: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+) -> None:
+    if weights.dim() != 2 or packed.shape != (weights.shape[0], weights.shape[1] // 2):
+        raise RuntimeError("expected weights (N,K), packed (N,K/2)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("grouped_w4a4_gemv_bf16"))
+def _grouped_w4a4_gemv_fake(
+    activations_packed: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if activations_packed.dim() != 2 or weight_stack.dim() != 3:
+        raise RuntimeError("expected activations_packed (M,K/2), weight_stack (E,N,K/2)")
+    if expert_idx.dim() != 2:
+        raise RuntimeError("expert_idx must have shape (M,top_k)")
+    expected = (activations_packed.shape[0], expert_idx.shape[1], weight_stack.shape[1])
+    if out.shape != expected:
+        raise RuntimeError(f"out must have shape {expected}")
+    return None
+
+
+def _swizzled_sf_bytes(rows: int, dim: int) -> int:
+    return ((int(rows) + 127) // 128) * (((int(dim) // 16) + 3) // 4) * 512
+
+
+def w4a16_decode_gemv_bf16(
+    x_bf16: torch.Tensor,
+    weight_packed: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
+    ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
+    return out
+
+
+def grouped_w4a16_gemv_bf16(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    n: int,
+    w_stride: Optional[int] = None,
+    sfb_stride: Optional[int] = None,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Run one W4A16 GEMV per routed slot.
+
+    `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
+    strides between experts; by default `w_stride = n * K / 2`.
+    """
+
+    k = activations.shape[1]
+    if out is None:
+        out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
+    if w_stride is None:
+        w_stride = int(n) * k // 2
+    if sfb_stride is None:
+        raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
+    ops.grouped_w4a16_gemv_bf16(
+        activations,
+        weight_stack,
+        sfb_stack,
+        alpha_stack,
+        expert_idx,
+        int(w_stride),
+        int(sfb_stride),
+        out,
+    )
+    return out
+
+
+def quantize_activations_nvfp4_bf16(
+    activations: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfa: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a BF16 ``[M,K]`` activation once for routed W4A4 GEMV.
+
+    Pass preallocated ``packed`` and ``sfa`` buffers on CUDA Graph hot paths.
+    """
+
+    m, k = activations.shape
+    if packed is None:
+        packed = torch.empty((m, k // 2), device=activations.device, dtype=torch.uint8)
+    if sfa is None:
+        sfa = torch.empty((_swizzled_sf_bytes(m, k),), device=activations.device, dtype=torch.uint8)
+    ops.quantize_activations_nvfp4_bf16(activations, packed, sfa)
+    return packed, sfa
+
+
+def quantize_weights_nvfp4_bf16(
+    weights: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfb: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Offline/helper quantization for one expert's BF16 ``[N,K]`` weight."""
+
+    n, k = weights.shape
+    if packed is None:
+        packed = torch.empty((n, k // 2), device=weights.device, dtype=torch.uint8)
+    if sfb is None:
+        sfb = torch.empty((_swizzled_sf_bytes(n, k),), device=weights.device, dtype=torch.uint8)
+    ops.quantize_weights_nvfp4_bf16(weights, packed, sfb)
+    return packed, sfb
+
+
+def grouped_w4a4_gemv_bf16(
+    activations_packed: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute all token/top-k W4A4 expert projections in one launch.
+
+    Inputs use token-major routing: ``expert_idx[M,top_k]`` and output
+    ``[M,top_k,N]``. The device index tensor is read on every graph replay.
+    """
+
+    m = activations_packed.shape[0]
+    top_k = expert_idx.shape[1]
+    n = weight_stack.shape[1]
+    if out is None:
+        out = torch.empty((m, top_k, n), device=activations_packed.device, dtype=torch.bfloat16)
+    ops.grouped_w4a4_gemv_bf16(
+        activations_packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out
+    )
+    return out
+
+
+def grouped_w4a4_gemv_from_bf16(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfa: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Quantize ``[M,K]`` once, then launch all ``M*top_k`` projections.
+
+    Supplying all three buffers makes this two-launch composition allocation
+    free and CUDA Graph replay safe.
+    """
+
+    packed, sfa = quantize_activations_nvfp4_bf16(
+        activations, packed=packed, sfa=sfa
+    )
+    return grouped_w4a4_gemv_bf16(
+        packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out=out
+    )
+
+
+__all__ = [
+    "grouped_w4a4_gemv_bf16",
+    "grouped_w4a4_gemv_from_bf16",
+    "grouped_w4a16_gemv_bf16",
+    "quantize_activations_nvfp4_bf16",
+    "quantize_weights_nvfp4_bf16",
+    "w4a16_decode_gemv_bf16",
+]

build/torch213-cxx11-cu130-x86_64-linux/_grouped_moe_gemv_cuda_1c03930.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71893fb7929f7bcdbfd877c98c877949d26c86349f5c498edc8912274e2ab2bc
+size 543200

build/torch213-cxx11-cu130-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemv_cuda_1c03930
+ops = torch.ops._grouped_moe_gemv_cuda_1c03930
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemv_cuda_1c03930::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/grouped_moe_gemv/__init__.py

@@ -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/torch213-cxx11-cu130-x86_64-linux/metadata.json

@@ -0,0 +1,34 @@
+{
+  "name": "grouped-moe-gemv",
+  "id": "_grouped_moe_gemv_cuda_1c03930",
+  "version": 2,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "6/7oa/DPHaXhsOCeDDw3KAkskrPXDsVGh0ll4lrN3DU=",
+      "_grouped_moe_gemv_cuda_1c03930.abi3.so": "cYk/t5Kfe82/2HfJjId5SdJshjSfXEmO3IkSJ04qsrw=",
+      "_ops.py": "dpEMiXfiniPw/mtHu4nyuKUeQDYLyaqFcpmnqiniKIA=",
+      "grouped_moe_gemv/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "b39ca23f1b36383df00b27b3ffe1276cd5dbea85",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "1c0393023e02142361a2de8ae1b4ae8ececb9730",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/__init__.py

@@ -0,0 +1,240 @@
+"""FlashRT grouped MoE GEMV 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("w4a16_decode_gemv_bf16"))
+def _w4a16_decode_gemv_fake(
+    x_bf16: torch.Tensor,
+    weight_packed: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
+    if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
+        raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
+def _grouped_w4a16_gemv_fake(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    w_stride: int,
+    sfb_stride: int,
+    out: torch.Tensor,
+) -> None:
+    if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
+        raise RuntimeError("expected activations (slots,K), out (slots,N)")
+    if expert_idx.shape != (activations.shape[0],):
+        raise RuntimeError("expert_idx must have shape (slots,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_activations_nvfp4_bf16"))
+def _quantize_activations_nvfp4_fake(
+    activations: torch.Tensor,
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+) -> None:
+    if activations.dim() != 2:
+        raise RuntimeError("activations must have shape (M,K)")
+    if packed.shape != (activations.shape[0], activations.shape[1] // 2):
+        raise RuntimeError("packed must have shape (M,K/2)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_weights_nvfp4_bf16"))
+def _quantize_weights_nvfp4_fake(
+    weights: torch.Tensor,
+    packed: torch.Tensor,
+    sfb: torch.Tensor,
+) -> None:
+    if weights.dim() != 2 or packed.shape != (weights.shape[0], weights.shape[1] // 2):
+        raise RuntimeError("expected weights (N,K), packed (N,K/2)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("grouped_w4a4_gemv_bf16"))
+def _grouped_w4a4_gemv_fake(
+    activations_packed: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if activations_packed.dim() != 2 or weight_stack.dim() != 3:
+        raise RuntimeError("expected activations_packed (M,K/2), weight_stack (E,N,K/2)")
+    if expert_idx.dim() != 2:
+        raise RuntimeError("expert_idx must have shape (M,top_k)")
+    expected = (activations_packed.shape[0], expert_idx.shape[1], weight_stack.shape[1])
+    if out.shape != expected:
+        raise RuntimeError(f"out must have shape {expected}")
+    return None
+
+
+def _swizzled_sf_bytes(rows: int, dim: int) -> int:
+    return ((int(rows) + 127) // 128) * (((int(dim) // 16) + 3) // 4) * 512
+
+
+def w4a16_decode_gemv_bf16(
+    x_bf16: torch.Tensor,
+    weight_packed: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
+    ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
+    return out
+
+
+def grouped_w4a16_gemv_bf16(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    n: int,
+    w_stride: Optional[int] = None,
+    sfb_stride: Optional[int] = None,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Run one W4A16 GEMV per routed slot.
+
+    `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
+    strides between experts; by default `w_stride = n * K / 2`.
+    """
+
+    k = activations.shape[1]
+    if out is None:
+        out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
+    if w_stride is None:
+        w_stride = int(n) * k // 2
+    if sfb_stride is None:
+        raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
+    ops.grouped_w4a16_gemv_bf16(
+        activations,
+        weight_stack,
+        sfb_stack,
+        alpha_stack,
+        expert_idx,
+        int(w_stride),
+        int(sfb_stride),
+        out,
+    )
+    return out
+
+
+def quantize_activations_nvfp4_bf16(
+    activations: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfa: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a BF16 ``[M,K]`` activation once for routed W4A4 GEMV.
+
+    Pass preallocated ``packed`` and ``sfa`` buffers on CUDA Graph hot paths.
+    """
+
+    m, k = activations.shape
+    if packed is None:
+        packed = torch.empty((m, k // 2), device=activations.device, dtype=torch.uint8)
+    if sfa is None:
+        sfa = torch.empty((_swizzled_sf_bytes(m, k),), device=activations.device, dtype=torch.uint8)
+    ops.quantize_activations_nvfp4_bf16(activations, packed, sfa)
+    return packed, sfa
+
+
+def quantize_weights_nvfp4_bf16(
+    weights: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfb: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Offline/helper quantization for one expert's BF16 ``[N,K]`` weight."""
+
+    n, k = weights.shape
+    if packed is None:
+        packed = torch.empty((n, k // 2), device=weights.device, dtype=torch.uint8)
+    if sfb is None:
+        sfb = torch.empty((_swizzled_sf_bytes(n, k),), device=weights.device, dtype=torch.uint8)
+    ops.quantize_weights_nvfp4_bf16(weights, packed, sfb)
+    return packed, sfb
+
+
+def grouped_w4a4_gemv_bf16(
+    activations_packed: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Compute all token/top-k W4A4 expert projections in one launch.
+
+    Inputs use token-major routing: ``expert_idx[M,top_k]`` and output
+    ``[M,top_k,N]``. The device index tensor is read on every graph replay.
+    """
+
+    m = activations_packed.shape[0]
+    top_k = expert_idx.shape[1]
+    n = weight_stack.shape[1]
+    if out is None:
+        out = torch.empty((m, top_k, n), device=activations_packed.device, dtype=torch.bfloat16)
+    ops.grouped_w4a4_gemv_bf16(
+        activations_packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out
+    )
+    return out
+
+
+def grouped_w4a4_gemv_from_bf16(
+    activations: torch.Tensor,
+    weight_stack: torch.Tensor,
+    sfb_stack: torch.Tensor,
+    alpha_stack: torch.Tensor,
+    expert_idx: torch.Tensor,
+    *,
+    packed: Optional[torch.Tensor] = None,
+    sfa: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Quantize ``[M,K]`` once, then launch all ``M*top_k`` projections.
+
+    Supplying all three buffers makes this two-launch composition allocation
+    free and CUDA Graph replay safe.
+    """
+
+    packed, sfa = quantize_activations_nvfp4_bf16(
+        activations, packed=packed, sfa=sfa
+    )
+    return grouped_w4a4_gemv_bf16(
+        packed, weight_stack, sfa, sfb_stack, alpha_stack, expert_idx, out=out
+    )
+
+
+__all__ = [
+    "grouped_w4a4_gemv_bf16",
+    "grouped_w4a4_gemv_from_bf16",
+    "grouped_w4a16_gemv_bf16",
+    "quantize_activations_nvfp4_bf16",
+    "quantize_weights_nvfp4_bf16",
+    "w4a16_decode_gemv_bf16",
+]

build/torch213-cxx11-cu132-x86_64-linux/_grouped_moe_gemv_cuda_1c03930.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4486ccc3b61668cb3790d36afc14bfdef065982bbc0a95f417fe29c8470d2bd3
+size 555536

build/torch213-cxx11-cu132-x86_64-linux/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _grouped_moe_gemv_cuda_1c03930
+ops = torch.ops._grouped_moe_gemv_cuda_1c03930
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_grouped_moe_gemv_cuda_1c03930::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/grouped_moe_gemv/__init__.py

@@ -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/torch213-cxx11-cu132-x86_64-linux/metadata.json

@@ -0,0 +1,34 @@
+{
+  "name": "grouped-moe-gemv",
+  "id": "_grouped_moe_gemv_cuda_1c03930",
+  "version": 2,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0a",
+      "12.0a"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "6/7oa/DPHaXhsOCeDDw3KAkskrPXDsVGh0ll4lrN3DU=",
+      "_grouped_moe_gemv_cuda_1c03930.abi3.so": "RIbMw7YWaMs3kNNq/BS/3vBlmCu8CpX0F/4pyEcNK9M=",
+      "_ops.py": "dpEMiXfiniPw/mtHu4nyuKUeQDYLyaqFcpmnqiniKIA=",
+      "grouped_moe_gemv/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "b39ca23f1b36383df00b27b3ffe1276cd5dbea85",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "1c0393023e02142361a2de8ae1b4ae8ececb9730",
+      "dirty": false
+    }
+  }
+}