Promote latest kernel artifacts to main

README.md

@@ -1,9 +1,116 @@
-# flashrt/fp4-gemm
+# fp4-gemm
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+FlashRT native Blackwell NVFP4 A4W4 GEMM kernels.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fp4-gemm
+This package consumes packed FP4 E2M1 tensors plus CUTLASS Sm1xx SFA/SFB scale
+buffers and produces BF16 output. It is designed to pair with
+`flashrt/fp4-fused-ops` and other static low-bit transformer/diffuser runtime
+paths.
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+## Available Functions
+
+- `sfa_size_bytes(rows, dim)`
+- `quantize_fp4_sfa_fp16(x, packed=None, sfa=None, is_sfb=False)`
+- `quantize_fp4_sfa_bf16(x, packed=None, sfa=None, is_sfb=False)`
+- `dequantize_fp4_sfa_fp16(packed, sfa, out=None, is_sfb=False)`
+- `nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None, variant=-1)`
+- `nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out=None)`
+- `nvfp4_gemm_bias_residual_bf16(a_packed, b_packed, sfa, sfb, bias, residual, out=None)`
+- `nvfp4_gemm_residual_bf16(a_packed, b_packed, sfa, sfb, residual, alpha=1.0, out=None)`
+- `nvfp4_gemm_bias_gelu_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
+- `nvfp4_gemm_bias_gelu_nvfp4(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out_packed=None, out_sfa=None)`
+- `nvfp4_gemm_streamk_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None)`
+- `nvfp4_gemm_streamk_bias_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
+- `fp4_w4a16_linear_bf16(...)` is retained as a compatibility alias
+
+## Tensor Contract
+
+- `a_packed`: `torch.uint8`, shape `(M, K / 2)`.
+- `b_packed`: `torch.uint8`, shape `(N, K / 2)`.
+- `sfa`: `torch.uint8`, CUTLASS SFA layout for `(M, K)`.
+- `sfb`: `torch.uint8`, CUTLASS SFB layout for `(N, K)`.
+- output: `torch.bfloat16`, shape `(M, N)`.
+- `K` must be divisible by 16.
+- Targets: Blackwell `sm_110a` (Jetson AGX Thor, CUDA 13+) and `sm_120a`
+  (RTX Blackwell, CUDA 12.8+).
+
+`variant` selects the CUTLASS schedule:
+
+- `-1`: architecture-aware auto-dispatch (public default).
+- `0`: default `<128,128,256>` cooperative schedule.
+- `1`: widen `<128,256,128>` schedule, intended for very large `N`.
+- `2`: pingpong schedule for A/B testing shape-specific wins.
+
+The canonical linear API and FP4/SFA quantize/dequantize helpers are available
+on both SM110 and SM120. SM110 additionally provides the GROOT N1.7 production
+epilogues `nvfp4_gemm_bias_bf16`, `nvfp4_gemm_bias_residual_bf16`, and
+`nvfp4_gemm_bias_gelu_nvfp4`. The latter emits packed FP4 plus CUTLASS SFA so
+the following projection can consume it without a BF16 materialization and a
+standalone quantization launch. Stream-K and the older BF16 GELU epilogue keep
+their existing SM120 dispatch and reject unsupported architectures explicitly.
+
+The SM110 release gate includes the production `(M,N,K)` shapes
+`(41,4608,1536)`, `(41,6144,1536)`, and `(41,1536,6144)`, plus the legacy
+`M=51` compatibility row. The kernels are the native sources used by FlashRT's
+GROOT N1.7 Thor NVFP4 pipeline.
+
+## Minimal Usage
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
+
+x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
+w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
+
+a_packed, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
+b_packed, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
+
+y = ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0)
+```
+
+For BF16 model activations, use the direct producer so the hot path does not
+materialize an intermediate FP16 tensor:
+
+```python
+x_bf16 = torch.randn((1, 5120), device="cuda", dtype=torch.bfloat16)
+a_packed, sfa = ops.quantize_fp4_sfa_bf16(x_bf16)
+```
+
+The BF16 entry writes the same E2M1 bytes and CUTLASS SFA/SFB layout as
+`quantize_fp4_sfa_fp16(x_bf16.to(torch.float16))` for finite FP16-range
+inputs. It is an additive API; the existing FP16 producer remains unchanged.
+
+The quantize/dequantize helpers are included for examples and validation. A
+production runtime should keep weights prepacked and should avoid quantizing in
+the hot path unless that producer kernel is part of the intended low-bit block.
+
+Use the bias/GELU and residual variants to avoid returning to BF16
+elementwise code between low-bit GEMMs. Stream-K variants are selected only
+for the validated large down-projection shapes; unsupported shapes reject
+rather than silently selecting a losing schedule.
+
+## Validation
+
+```bash
+python fp4-gemm/tests/test_fp4_gemm.py --backend source --mode full
+python fp4-gemm/tests/test_fp4_gemm.py --backend installed --mode full \
+  --artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
+python fp4-gemm/benchmarks/benchmark.py --backend installed --mode headline \
+  --artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
+
+# Thor model-shape gate
+python fp4-gemm/tests/test_fp4_gemm.py --backend installed \
+  --mode thor-models \
+  --artifact fp4-gemm/build/torch211-cxx11-cu130-aarch64-linux
+```
+
+The correctness reference dequantizes the same FP4/SFA and FP4/SFB inputs used
+by the kernel, then computes the PyTorch GEMM reference from those dequantized
+low-bit values.
+
+The producer gate also checks the BF16 direct entry byte-for-byte against the
+established FP16 compatibility chain at decode widths 5120, 6144 and 17408,
+plus multi-row activation and SFB layouts.

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

@@ -19,7 +19,8 @@ def sfa_size_bytes(rows: int, dim: int) -> int:
 def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     return (
         torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
-        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+        # Tile-layout padding entries are not written by every quantizer.
+        torch.zeros((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
     )
 
 
@@ -36,6 +37,36 @@ def _linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_fp16"))
+def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, variant: int = -1) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
+def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
+def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
+def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4_fp16"))
+def _bias_gelu_fp4_fp16_fake(a, b, sfa, sfb, bias, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_fp16"))
+def _bias_residual_fp16_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
 def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
     if a_packed.shape[0] != 1:
@@ -58,11 +89,46 @@ def _legacy_linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_bf16"))
+def _bias_fake(a, b, sfa, sfb, bias, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_bf16"))
+def _bias_residual_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
 def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
+def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
+def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
+def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
+def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
 def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
@@ -112,6 +178,160 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_mse_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack FP16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+
+
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+
+
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
+
+
+def quantize_e0m3_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def e0m3_weight_gemm_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    a_format: int = 1,
+    out: torch.Tensor | None = None,
+):
+    """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.e0m3_weight_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
+    )
+    return out
+
+
+def nvfp4_gemm_relu2_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_relu2_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
 def quantize_fp4_sfa_bf16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -152,6 +372,169 @@ def nvfp4_gemm_bf16(
     return out
 
 
+def nvfp4_gemm_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with FP16 output."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.nvfp4_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+
+
+def nvfp4_gemm_variant_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 7,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with BF16 output.
+
+    ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
+    the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
+    """
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_variant_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+
+
+def nvfp4_gemm_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_geglu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    skinny: bool = False,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GEMM with fused GeGLU and compact NVFP4 output on SM110.
+
+    ``b_interleaved_packed`` stores gate/up rows pairwise, so its first
+    dimension is twice the logical hidden width.
+    """
+    m, n_twice = a_packed.shape[0], b_interleaved_packed.shape[0]
+    hidden = n_twice // 2
+    if scratch is None:
+        scratch = torch.empty((m, hidden), device=a_packed.device, dtype=torch.uint8)
+    if out_packed is None:
+        out_packed = torch.empty((m, hidden // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.zeros((sfa_size_bytes(m, hidden),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed, b_interleaved_packed, sfa, sfb, scratch,
+        out_packed, out_sfa, bool(skinny)
+    )
+    return out_packed, out_sfa
+
+
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
+
+
+def nvfp4_gemm_bias_gelu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP16-bias GEMM with fused GELU and NVFP4 output on SM110."""
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.zeros((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4_fp16(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_bias_residual_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP16-output GEMM with fused FP16 bias and residual on SM110."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_fp16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -191,6 +574,45 @@ def fp4_w4a16_linear_bf16(
     )
 
 
+def nvfp4_gemm_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with a fused per-column BF16 bias."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out)
+    return out
+
+
+def nvfp4_gemm_bias_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with fused BF16 bias and residual add."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
 def nvfp4_gemm_residual_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -243,7 +665,7 @@ def nvfp4_gemm_bias_gelu_nvfp4(
     if out_packed is None:
         out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
     if out_sfa is None:
-        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+        out_sfa = torch.zeros((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
     ops.nvfp4_gemm_bias_gelu_nvfp4(
         a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
     )
@@ -292,16 +714,33 @@ def nvfp4_gemm_streamk_bias_bf16(
 
 
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
+    "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
+    "nvfp4_gemm_fp16",
+    "nvfp4_gemm_variant_bf16",
+    "nvfp4_gemm_nvfp4",
+    "nvfp4_gemm_geglu_nvfp4_fp16",
+    "nvfp4_gemm_bias_gelu_nvfp4_fp16",
+    "nvfp4_gemm_bias_residual_fp16",
+    "nvfp4_gemm_bias_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_bias_residual_bf16",
     "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_relu2_nvfp4",
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
+    "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_83a7cd0.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ac9c00df31abab395c44ea7f019effc0a40fe5e71a6461ca9004540a24f440c2
-size 1825840
+oid sha256:63a46e5ec6142207d1f9f881d2d01238e219823e3e514b75ac232cab42d6c0be
+size 2207552

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

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

build/torch211-cxx11-cu128-x86_64-linux/fp4_gemm/__init__.py

@@ -2,25 +2,13 @@ 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}")
+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)
-    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
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
     return module
 
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_83a7cd0",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -9,5 +9,24 @@
     "archs": [
       "12.0a"
     ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "Y6RuXsYUIgfR+fiB0tASOOIZgj4+UUt1rCMsq0LWwL4=",
+      "_ops.py": "twu0zeqTyaITBKAdr9FudRbkbXpvBQdmAHcVf/OXw2I=",
+      "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu128",
+      "cuda": "12.8"
+    }
   }
 }

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

@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
+def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
+def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
 def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
     return None
@@ -94,6 +104,16 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
+def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -158,6 +178,104 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_mse_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack FP16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+
+
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+
+
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
+
+
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -276,6 +394,52 @@ def nvfp4_gemm_fp16(
     return out
 
 
+def nvfp4_gemm_variant_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 7,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with BF16 output.
+
+    ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
+    the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
+    """
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_variant_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+
+
+def nvfp4_gemm_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
 def nvfp4_gemm_geglu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_interleaved_packed: torch.Tensor,
@@ -307,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 
 
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
+
+
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -527,12 +714,16 @@ def nvfp4_gemm_streamk_bias_bf16(
 
 
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
     "nvfp4_gemm_fp16",
+    "nvfp4_gemm_variant_bf16",
+    "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",
@@ -545,6 +736,10 @@ __all__ = [
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch211-cxx11-cu130-aarch64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f92ea3a945832ebd737cf074eac4983674fd00b160afa1481cd24148e8e32297
-size 5528088
+oid sha256:fa36d417fac09d8ed49fbb867c35562ed06ca80b67747c09624406e21afcd38e
+size 6002944

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

@@ -1,6 +1,6 @@
 import torch
-from . import _fp4_gemm_cuda_230e467
-ops = torch.ops._fp4_gemm_cuda_230e467
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 
 def add_op_namespace_prefix(op_name: str):
-    return f"_fp4_gemm_cuda_230e467::{op_name}"
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_230e467",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,16 +13,20 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
-      "_fp4_gemm_cuda_230e467.abi3.so": "+S6jqUWDLr1zfPB06sSYNnT9ALFgr6FIHNJBSOjjIpc=",
-      "_ops.py": "IfU2iMVrWYRnoe6q8HI3hU8TPnAlLrXDerlv117PcXI=",
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "+jbUF/rAnY7Un7uGfDVWLtBsqAtndHwJYkQG4hr8044=",
+      "_ops.py": "twu0zeqTyaITBKAdr9FudRbkbXpvBQdmAHcVf/OXw2I=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "230e467",
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu130",
+      "cuda": "13.0"
     }
   }
 }

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

@@ -19,7 +19,8 @@ def sfa_size_bytes(rows: int, dim: int) -> int:
 def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
     return (
         torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
-        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+        # Tile-layout padding entries are not written by every quantizer.
+        torch.zeros((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
     )
 
 
@@ -36,6 +37,36 @@ def _linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_fp16"))
+def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, variant: int = -1) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
+def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
+def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
+def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4_fp16"))
+def _bias_gelu_fp4_fp16_fake(a, b, sfa, sfb, bias, out_packed, out_sfa) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_fp16"))
+def _bias_residual_fp16_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
 def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
     if a_packed.shape[0] != 1:
@@ -58,11 +89,46 @@ def _legacy_linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_bf16"))
+def _bias_fake(a, b, sfa, sfb, bias, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_bf16"))
+def _bias_residual_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
 def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
+def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
+def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
+def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
+def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
 def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
@@ -112,6 +178,160 @@ def quantize_fp4_sfa_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_mse_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack FP16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+
+
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+
+
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
+
+
+def quantize_e0m3_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def e0m3_weight_gemm_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    a_format: int = 1,
+    out: torch.Tensor | None = None,
+):
+    """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.e0m3_weight_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
+    )
+    return out
+
+
+def nvfp4_gemm_relu2_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+):
+    """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_relu2_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
 def quantize_fp4_sfa_bf16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -152,6 +372,169 @@ def nvfp4_gemm_bf16(
     return out
 
 
+def nvfp4_gemm_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with FP16 output."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.float16,
+        )
+    ops.nvfp4_gemm_fp16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+
+
+def nvfp4_gemm_variant_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = 7,
+) -> torch.Tensor:
+    """SM110 native NVFP4 GEMM with BF16 output.
+
+    ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
+    the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
+    """
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_variant_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
+    )
+    return out
+
+
+def nvfp4_gemm_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
+    if out_packed is None or out_sfa is None:
+        out_packed, out_sfa = _alloc_fp4(
+            a_packed.shape[0], b_packed.shape[0], a_packed.device
+        )
+    ops.nvfp4_gemm_nvfp4(
+        a_packed, b_packed, sfa, sfb, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_geglu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    skinny: bool = False,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GEMM with fused GeGLU and compact NVFP4 output on SM110.
+
+    ``b_interleaved_packed`` stores gate/up rows pairwise, so its first
+    dimension is twice the logical hidden width.
+    """
+    m, n_twice = a_packed.shape[0], b_interleaved_packed.shape[0]
+    hidden = n_twice // 2
+    if scratch is None:
+        scratch = torch.empty((m, hidden), device=a_packed.device, dtype=torch.uint8)
+    if out_packed is None:
+        out_packed = torch.empty((m, hidden // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.zeros((sfa_size_bytes(m, hidden),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed, b_interleaved_packed, sfa, sfb, scratch,
+        out_packed, out_sfa, bool(skinny)
+    )
+    return out_packed, out_sfa
+
+
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
+
+
+def nvfp4_gemm_bias_gelu_nvfp4_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FP16-bias GEMM with fused GELU and NVFP4 output on SM110."""
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.zeros((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4_fp16(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_bias_residual_fp16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """FP16-output GEMM with fused FP16 bias and residual on SM110."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_fp16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
 def fp4_w4a4_gemv_warpsplit_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -191,6 +574,45 @@ def fp4_w4a16_linear_bf16(
     )
 
 
+def nvfp4_gemm_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with a fused per-column BF16 bias."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out)
+    return out
+
+
+def nvfp4_gemm_bias_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with fused BF16 bias and residual add."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
 def nvfp4_gemm_residual_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -243,7 +665,7 @@ def nvfp4_gemm_bias_gelu_nvfp4(
     if out_packed is None:
         out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
     if out_sfa is None:
-        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+        out_sfa = torch.zeros((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
     ops.nvfp4_gemm_bias_gelu_nvfp4(
         a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
     )
@@ -292,16 +714,33 @@ def nvfp4_gemm_streamk_bias_bf16(
 
 
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
+    "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
+    "nvfp4_gemm_fp16",
+    "nvfp4_gemm_variant_bf16",
+    "nvfp4_gemm_nvfp4",
+    "nvfp4_gemm_geglu_nvfp4_fp16",
+    "nvfp4_gemm_bias_gelu_nvfp4_fp16",
+    "nvfp4_gemm_bias_residual_fp16",
+    "nvfp4_gemm_bias_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_bias_residual_bf16",
     "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_relu2_nvfp4",
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
+    "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",
 ]

build/torch211-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_83a7cd0.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd0e48fd6d7fd9be6c613a6bca6c7bbd6a8e713c8743aca800d28bd0e0a8237b
-size 2422232
+oid sha256:199bf055035b1b7d840c3b83094f0b95fef7be344a8e4ee5e180798a26b74b10
+size 2330536

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

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

build/torch211-cxx11-cu130-x86_64-linux/fp4_gemm/__init__.py

@@ -2,25 +2,13 @@ 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}")
+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)
-    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
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
     return module
 
-
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

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

@@ -1,14 +1,32 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_83a7cd0",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0a",
       "12.0a"
     ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "GZvwVQNbG32EDDuDCU8Llf73vjRKjk7l4YB5iia3SxA=",
+      "_ops.py": "twu0zeqTyaITBKAdr9FudRbkbXpvBQdmAHcVf/OXw2I=",
+      "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu130",
+      "cuda": "13.0"
+    }
   }
 }

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

@@ -109,6 +109,11 @@ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_bf16"))
+def _quant_mse_bf16_fake(x, packed, sfa, is_sfb: bool = False) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
 def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
     return None
@@ -186,6 +191,91 @@ def quantize_fp4_sfa_mse_fp16(
     return packed, sfa
 
 
+def quantize_fp4_sfa_mse_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Pack BF16 weights using per-block reconstruction-MSE scale search."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_mse_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def aligned_fp4_dim(dim: int, alignment: int = 32) -> int:
+    """Return the physical FP4 GEMM dimension for a logical dimension."""
+    if dim <= 0 or alignment <= 0 or alignment % 16 != 0:
+        raise ValueError("dim must be positive and alignment divisible by 16")
+    return ((dim + alignment - 1) // alignment) * alignment
+
+
+def quantize_fp4_sfa_padded_bf16(
+    x: torch.Tensor,
+    *,
+    alignment: int = 32,
+    is_sfb: bool = False,
+):
+    """Bind-time BF16 pack with zero padding on the final dimension.
+
+    The returned packed tensor uses the physical aligned width. Keep the
+    logical width returned as the third item for slicing final model outputs.
+    This helper allocates and is intentionally not a runtime hot-path op.
+    """
+    if x.ndim != 2 or x.dtype != torch.bfloat16:
+        raise ValueError("x must be a 2-D BF16 tensor")
+    logical_dim = int(x.shape[1])
+    physical_dim = aligned_fp4_dim(logical_dim, alignment)
+    if physical_dim == logical_dim:
+        padded = x.contiguous()
+    else:
+        padded = torch.zeros(
+            (x.shape[0], physical_dim), device=x.device, dtype=x.dtype
+        )
+        padded[:, :logical_dim].copy_(x)
+    packed, sfa = quantize_fp4_sfa_bf16(padded, is_sfb=is_sfb)
+    return packed, sfa, logical_dim
+
+
+def pack_nvfp4_weight_bf16(
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    *,
+    alignment: int = 32,
+    mse: bool = False,
+):
+    """Zero-pad and pack a logical ``(N, K)`` BF16 weight at bind time.
+
+    Both physical ``N`` and ``K`` are aligned. This is the model-neutral
+    contract used by widths such as SigLIP's logical 4304, which becomes 4320
+    for Blackwell NVFP4 TMA. The hot path consumes the returned static tensors
+    directly and performs no padding or allocation.
+    """
+    if weight.ndim != 2 or weight.dtype != torch.bfloat16:
+        raise ValueError("weight must be a 2-D BF16 tensor")
+    logical_n, logical_k = map(int, weight.shape)
+    physical_n = aligned_fp4_dim(logical_n, alignment)
+    physical_k = aligned_fp4_dim(logical_k, alignment)
+    padded_weight = torch.zeros(
+        (physical_n, physical_k), device=weight.device, dtype=weight.dtype
+    )
+    padded_weight[:logical_n, :logical_k].copy_(weight)
+    quantize = quantize_fp4_sfa_mse_bf16 if mse else quantize_fp4_sfa_bf16
+    packed, sfb = quantize(padded_weight, is_sfb=True)
+    padded_bias = None
+    if bias is not None:
+        if bias.ndim != 1 or bias.numel() != logical_n:
+            raise ValueError("bias must have shape (N,)")
+        if bias.dtype != torch.bfloat16 or bias.device != weight.device:
+            raise ValueError("bias must match the weight dtype and device")
+        padded_bias = torch.zeros(
+            (physical_n,), device=bias.device, dtype=bias.dtype
+        )
+        padded_bias[:logical_n].copy_(bias)
+    return packed, sfb, padded_bias, (logical_n, logical_k)
+
+
 def quantize_e0m3_sfa_fp16(
     x: torch.Tensor,
     packed: torch.Tensor | None = None,
@@ -381,6 +471,29 @@ def nvfp4_gemm_geglu_nvfp4_fp16(
     return out_packed, out_sfa
 
 
+def cutlass_fp4_gemm_geglu_il_hw_v10(
+    a_packed: torch.Tensor,
+    b_interleaved_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    scratch: torch.Tensor | None = None,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Explicit public alias for the native SM110 v10 GeGLU schedule."""
+    return nvfp4_gemm_geglu_nvfp4_fp16(
+        a_packed,
+        b_interleaved_packed,
+        sfa,
+        sfb,
+        skinny=True,
+        scratch=scratch,
+        out_packed=out_packed,
+        out_sfa=out_sfa,
+    )
+
+
 def nvfp4_gemm_bias_gelu_nvfp4_fp16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -601,6 +714,8 @@ def nvfp4_gemm_streamk_bias_bf16(
 
 
 __all__ = [
+    "aligned_fp4_dim",
+    "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
     "e0m3_weight_gemm_fp16",
     "fp4_w4a16_linear_bf16",
@@ -622,6 +737,9 @@ __all__ = [
     "nvfp4_gemm_streamk_bias_bf16",
     "quantize_fp4_sfa_fp16",
     "quantize_fp4_sfa_mse_fp16",
+    "quantize_fp4_sfa_mse_bf16",
+    "quantize_fp4_sfa_padded_bf16",
+    "pack_nvfp4_weight_bf16",
     "quantize_e0m3_sfa_fp16",
     "quantize_fp4_sfa_bf16",
     "sfa_size_bytes",

build/torch213-cxx11-cu130-aarch64-linux/{fp4_gemm_source_test.abi3.so → _fp4_gemm_cuda_92c4b99.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:73b4cbcb15d13cc6452007e59961f3b217d21771f0c0937cc121cc8e0e6b1e12
-size 6529016
+oid sha256:b40a60dd7cfb9dfd94077016907dae286cd1222905ae26c740a993f73ff0d9d2
+size 6070984

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

@@ -1,6 +1,6 @@
 import torch
-from . import fp4_gemm_source_test
-ops = torch.ops.fp4_gemm_source_test
+from . import _fp4_gemm_cuda_92c4b99
+ops = torch.ops._fp4_gemm_cuda_92c4b99
 
 def add_op_namespace_prefix(op_name: str):
-    return f"fp4_gemm_source_test::{op_name}"
+    return f"_fp4_gemm_cuda_92c4b99::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/metadata.json

@@ -1,27 +1,27 @@
 {
   "name": "fp4-gemm",
-  "id": "fp4_gemm_source_test",
+  "id": "_fp4_gemm_cuda_92c4b99",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0"
+      "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
-      "fp4_gemm_source_test.abi3.so": "c7TLyxXRPMZFIAflmWHzshfSF3HwwJN8wSHMjg5rHhI=",
-      "_ops.py": "ZPjRxrYk7lCA0ZHoCjg8BaJyFY7a5i4NC3jIL6VMRJ4=",
+      "__init__.py": "Mr1/bAUzs+af68D+bRXX6iONDof//dzKgn6QU5pKdu8=",
+      "_fp4_gemm_cuda_92c4b99.abi3.so": "tApg3Xz7nf2UB3AWkH2uKGzRIikFribHQKmT9z/w2dI=",
+      "_ops.py": "twu0zeqTyaITBKAdr9FudRbkbXpvBQdmAHcVf/OXw2I=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
+      "sha": "92c4b991b6f1fbd89744aa6f46f62b1bfac7ca7e",
       "dirty": false
     },
     "validation": {