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.

benchmarks/benchmark.py

@@ -185,12 +185,79 @@ def bench_bf16_producer(ops, native, k: int, warmup: int, iters: int):
     }
 
 
+def bench_w4a16_vocab(ops, warmup: int, iters: int):
+    n, k = 152064, 5120
+    gen = torch.Generator(device="cuda").manual_seed(20260818)
+    weight = torch.randint(
+        0, 256, (n, k // 2), device="cuda", dtype=torch.uint8,
+        generator=gen,
+    )
+    weight_scale = (
+        torch.rand((n, k // 16), device="cuda", generator=gen) * 0.1 + 0.01
+    ).to(torch.float8_e4m3fn)
+    global_scale = torch.tensor([0.00390625], device="cuda", dtype=torch.float32)
+    weight_marlin, scales_marlin, global_marlin, workspace = (
+        ops.adopt_w4a16_marlin(weight, weight_scale, global_scale)
+    )
+
+    vllm_ops = None
+    scalar_types = None
+    try:
+        from vllm import _custom_ops as vllm_ops
+        from vllm.scalar_type import scalar_types
+    except ImportError:
+        pass
+
+    rows = []
+    for m in (1, 7, 8, 16):
+        x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
+        out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+
+        def flashrt_call():
+            ops.w4a16_marlin(
+                x, weight_marlin, scales_marlin, global_marlin, workspace, out
+            )
+
+        flashrt_call()
+        flashrt_us = measure(flashrt_call, warmup, iters)
+        row = {
+            "M": m,
+            "N": n,
+            "K": k,
+            "flashrt_us": flashrt_us,
+            "target_us": 400.0 if m == 8 else None,
+            "target_passed": bool(m != 8 or flashrt_us <= 400.0),
+        }
+        if vllm_ops is not None and scalar_types is not None:
+            native_out = torch.empty_like(out)
+
+            def native_call():
+                vllm_ops.marlin_gemm(
+                    x, native_out, weight_marlin, None, scales_marlin, None,
+                    global_marlin, None, None, None, workspace,
+                    scalar_types.float4_e2m1f, m, n, k,
+                    True, False, False, False,
+                )
+
+            native_call()
+            native_us = measure(native_call, warmup, iters)
+            row.update({
+                "vllm_native_us": native_us,
+                "wrapper_over_native": flashrt_us / native_us,
+                "exact_vs_vllm": bool(torch.equal(out, native_out)),
+            })
+        rows.append(row)
+    return rows
+
+
 def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--backend", choices=["source", "installed"], default="source")
     parser.add_argument("--artifact", default=None)
     parser.add_argument(
-        "--mode", choices=["smoke", "headline", "thor-models"], default="headline"
+        "--mode",
+        choices=["smoke", "headline", "thor-models", "w4a16-vocab"],
+        default="headline",
     )
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
@@ -198,6 +265,31 @@ def main() -> int:
     args = parser.parse_args()
 
     helpers = load_helpers()
+    ops = (
+        helpers.load_source_ops()
+        if args.backend == "source"
+        else helpers.load_installed_ops(args.artifact)
+    )
+    if args.mode == "w4a16-vocab":
+        rows = bench_w4a16_vocab(ops, args.warmup, args.iterations)
+        payload = {
+            "mode": args.mode,
+            "backend": args.backend,
+            "device": torch.cuda.get_device_name(),
+            "torch": torch.__version__,
+            "results": rows,
+            "passed": all(
+                row["target_passed"] and row.get("exact_vs_vllm", True)
+                for row in rows
+            ),
+        }
+        print(json.dumps(payload, indent=2))
+        if args.json_out:
+            out = Path(args.json_out)
+            out.parent.mkdir(parents=True, exist_ok=True)
+            out.write_text(json.dumps(payload, indent=2) + "\n")
+        return 0 if payload["passed"] else 1
+
     native_root = Path(
         os.environ.get("FLASHRT_NATIVE_ROOT", str(ROOT.parent / "official" / "FlashRT"))
     )
@@ -206,11 +298,6 @@ def main() -> int:
         import flash_rt.flash_rt_kernels as native
     finally:
         sys.path.pop(0)
-    ops = (
-        helpers.load_source_ops()
-        if args.backend == "source"
-        else helpers.load_installed_ops(args.artifact)
-    )
     shapes = {
         "small_m16_n128_k128": (16, 128, 128),
         "small_m32_n256_k256": (32, 256, 256),

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_dc57d15.abi3.so → _fp4_gemm_cuda_6971f7a.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d1da1c9c5de85d62d2a39e29178eb3fb6e1fc4d772b04195fbcfd9683b63fcb5
-size 2721416
+oid sha256:f01d9ffadff9a04439f0a83b6095305d1e20d3903fd275624cd4b2c0001ceaaf
+size 4915160

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,9 +13,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "0docnF3oXWLSo54pF46z+24fxNdysEGV+8/ZaDtj/LU=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "8B2f+t/5oEQ58Kg7YJUwXR4g05A/0nViTNSywAAc6q8=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -26,7 +26,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

build/torch211-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_dc57d15.abi3.so → _fp4_gemm_cuda_6971f7a.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:687ec5c797f1519fed83a51f82bc7ecf53059fbe35843981d86fefe4143c3b87
-size 8564792
+oid sha256:555ecae796a12b3fed4ebc261381ef859509f47967e7fc84901fbbca05b31923
+size 10756144

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "aH7Fx5fxUZ/tg6Ufgrx+z1MFn741hDmB2G/v5BQ8O4c=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "VV7K55ahKz/tTrwmE4HvhZUJ9Hln5/yEkB+7ygWzGSM=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_dc57d15.abi3.so → _fp4_gemm_cuda_6971f7a.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0f26894cd8000a3c464811b255f07732758d8d53521c2b84c305e9cc07d96619
-size 8566680
+oid sha256:4ecd0fd14beafd0648b7353999f47e14168ca7a7e051d51527729b836e648548
+size 10758024

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "DyaJTNgACjxGSBGyVfB3MnWNjVNSHCuEwwXpzAfZZhk=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "Ts0P0Uvq/QZItzU5mfR+FBaMp6fgUdUVJ3Kbg25khUg=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_dc57d15.abi3.so → _fp4_gemm_cuda_6971f7a.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3b1fc2155f02132014a4ed86a37eb5c945d705651eeda4cf4fc294c4e4d5a442
-size 8550344
+oid sha256:8d2c4336526e99aec6d7644e1f9bbc3cc0175e0d9b63689e50d8639b58dd0daa
+size 10663864

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "Ox/CFV8CEyAUpO2Go361yUXXBWUe7aTPT8KUxOTVpEI=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "jSxDNlJuma7G12ROH5u8PMAXXg2bY2ieUNhjm1jdDao=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

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

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6075d51ae508759aacba4d16f7bf01c2f27c57de03afcb0745ffb073ee1c3320
+size 10757864

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "oulzVg5pRJ7nlweYS3cD0Z6Pjvm7o5AlE4h8V7ZGgsA=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "YHXVGuUIdZqsuk0W978BwvJ8V94Dr8sHRf+wc+4cMyA=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }

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

@@ -37,6 +37,11 @@ def capabilities() -> dict[str, object]:
         "sm120_warpsplit_mrows": (1, 16),
         "sm120_warpsplit_mrows_n_alignment": 8,
         "sm120_warpsplit_mrows_k_alignment": "64 * warps",
+        "sm120_w4a16_marlin_m": (1, 16),
+        "sm120_w4a16_marlin_n_alignment": 64,
+        "sm120_w4a16_marlin_k_alignment": 128,
+        "sm120_w4a16_marlin_weight_layout": "marlin-repacked-nvfp4",
+        "sm120_w4a16_marlin_scale_layout": "marlin-permuted-fp8-e4m3-block16",
         "errors": "exceptions",
     }
 
@@ -172,6 +177,28 @@ def _m256_fake(a_packed, b_packed, sfa, sfb, workspace, out, alpha: float = 1.0)
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_bf16"))
+def _w4a16_marlin_fake(x, weight, weight_scale, global_scale, workspace, out) -> None:
+    del weight, weight_scale, global_scale, workspace
+    if x.ndim != 2 or out.ndim != 2 or out.shape[0] != x.shape[0]:
+        raise RuntimeError("x/out must have shapes (M, K)/(M, N)")
+    if not 1 <= x.shape[0] <= 16:
+        raise RuntimeError("the Marlin W4A16 tier serves M in [1, 16]")
+    if x.shape[1] % 128 or out.shape[1] % 64:
+        raise RuntimeError("K/N must be divisible by 128/64")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_w4a16_marlin_repack"))
+def _w4a16_marlin_repack_fake(qweight_kn, weight_marlin) -> None:
+    if qweight_kn.ndim != 2:
+        raise RuntimeError("qweight_kn must have shape (K / 8, N)")
+    k, n = qweight_kn.shape[0] * 8, qweight_kn.shape[1]
+    if weight_marlin.shape != (k // 16, n * 2):
+        raise RuntimeError("weight_marlin must have shape (K / 16, 2 * N)")
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
 def _legacy_linear_fake(
     a_packed: torch.Tensor,
@@ -771,6 +798,98 @@ def nvfp4_gemm_m256_bf16(
     return out
 
 
+def allocate_w4a16_marlin_workspace(device: torch.device | str) -> torch.Tensor:
+    """Allocate the persistent Marlin lock workspace before graph capture."""
+    props = torch.cuda.get_device_properties(device)
+    return torch.zeros(props.multi_processor_count, device=device, dtype=torch.int32)
+
+
+def adopt_nvfp4_w4a16_marlin(
+    weight_packed: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_scale_2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Convert standard ModelOpt NVFP4 tensors once at model bind time.
+
+    Inputs use ``weight[N,K/2]``, ``weight_scale[N,K/16]`` and scalar
+    ``weight_scale_2``. Returned tensors are the three runtime operands plus a
+    zeroed lock workspace. This function is intentionally outside the graph
+    hot path.
+    """
+    if weight_packed.dtype != torch.uint8 or weight_packed.ndim != 2:
+        raise ValueError("weight_packed must be uint8 with shape (N, K / 2)")
+    n, k = weight_packed.shape[0], weight_packed.shape[1] * 2
+    if k % 128 or n % 64:
+        raise ValueError("K/N must be divisible by 128/64")
+    if weight_scale.shape != (n, k // NVFP4_BLOCK_SIZE):
+        raise ValueError("weight_scale must have shape (N, K / 16)")
+    if weight_scale_2.numel() != 1:
+        raise ValueError("weight_scale_2 must contain one global scale")
+
+    qweight_kn = weight_packed.view(torch.int32).T.contiguous()
+    weight_marlin = torch.empty(
+        (k // 16, n * 2), device=weight_packed.device, dtype=torch.int32
+    )
+    ops.nvfp4_w4a16_marlin_repack(qweight_kn, weight_marlin)
+
+    scales = weight_scale.T.contiguous().to(torch.bfloat16)
+    scale_perm = [i + 8 * j for i in range(8) for j in range(8)]
+    scales = scales.reshape(-1, 64)[:, scale_perm].reshape(-1, n).contiguous()
+
+    # Match vLLM's NVFP4 Marlin contract exactly: scale permutation uses the
+    # model parameter dtype, then S0E5M3 encoding starts from FP16.
+    scales = scales.to(torch.float16)
+    scales = scales.view(-1, 4)[:, [0, 2, 1, 3]].reshape(scales.shape)
+    scaled = scales.float() * (2**7)
+    nonzero = scaled > 0
+    scale_factor = 1.0
+    if bool(nonzero.any()):
+        max_value = scaled[nonzero].max()
+        if bool(max_value < 448 * (2**7)):
+            scale_factor = float((448 * (2**7) / max_value).log2().floor().exp2())
+    if scale_factor > 1.0:
+        scales = (scales.float() * scale_factor).to(torch.float16)
+    scales = scales * (2**7)
+    scales[scales < 2] = 0
+    scales = (scales.view(torch.int16) << 1).view(torch.float8_e4m3fn)
+    scales = scales[:, 1::2].contiguous()
+
+    # Marlin's BF16 dequantizer uses exponent bias 126 and removes the seven
+    # bits introduced while encoding the special unsigned FP8 scale format.
+    global_scale = (
+        weight_scale_2.to(device=weight_packed.device, dtype=torch.float32)
+        * (2.0 ** 119)
+        / scale_factor
+    ).contiguous()
+    workspace = allocate_w4a16_marlin_workspace(weight_packed.device)
+    return weight_marlin, scales, global_scale, workspace
+
+
+def nvfp4_w4a16_marlin_bf16(
+    x: torch.Tensor,
+    weight_marlin: torch.Tensor,
+    weight_scale_marlin: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    *,
+    workspace: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Small-M BF16 x NVFP4 linear using bind-time Marlin layouts.
+
+    ``weight_marlin``, ``weight_scale_marlin`` and
+    ``weight_global_scale`` are the tensors produced by a standard NVFP4
+    Marlin adoption step. The operation supports M in [1, 16] and never
+    allocates internally, making it suitable for CUDA Graph replay.
+    """
+    if out is None:
+        n = weight_scale_marlin.numel() // (x.shape[1] // NVFP4_BLOCK_SIZE)
+        out = torch.empty((x.shape[0], n), device=x.device, dtype=torch.bfloat16)
+    ops.nvfp4_w4a16_marlin_bf16(
+        x, weight_marlin, weight_scale_marlin, weight_global_scale, workspace, out
+    )
+    return out
+
+
 def fp4_w4a16_linear_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -932,6 +1051,8 @@ __all__ = [
     "SUPPORTED_CUDA_CAPABILITIES",
     "SUPPORTED_LAYOUTS",
     "capabilities",
+    "allocate_w4a16_marlin_workspace",
+    "adopt_nvfp4_w4a16_marlin",
     "aligned_fp4_dim",
     "cutlass_fp4_gemm_geglu_il_hw_v10",
     "dequantize_fp4_sfa_fp16",
@@ -947,6 +1068,7 @@ __all__ = [
     "nvfp4_gemm_nvfp4",
     "nvfp4_gemm_m256_bf16",
     "nvfp4_gemm_m256_workspace_size",
+    "nvfp4_w4a16_marlin_bf16",
     "nvfp4_gemm_geglu_nvfp4_fp16",
     "nvfp4_gemm_bias_gelu_nvfp4_fp16",
     "nvfp4_gemm_bias_residual_fp16",

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

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

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

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

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

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

build/torch213-cxx11-cu132-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_dc57d15",
+  "id": "_fp4_gemm_cuda_6971f7a",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "XycjAPC24qtK6IH+GtF/jrPmWnB6DyNgSny26ep3w1Y=",
-      "_fp4_gemm_cuda_dc57d15.abi3.so": "PFxGYBXp/aNqzA2UJvLJKBu56K8XWvsWhu3Tc0ik+XU=",
-      "_ops.py": "TM5HQsh2RfinhrEZ4U0XyWF2hFyrPWpxR79FQ41D5jk=",
+      "__init__.py": "0hDN2W6X3IMoWIi5BBufrnqrq2yj5N1/q9wBhCwBRxA=",
+      "_fp4_gemm_cuda_6971f7a.abi3.so": "2J36GdJDZxpBXw/a7rimJhxBe/fU9F5dZhdFOH4u8fc=",
+      "_ops.py": "D7InzLughnNL5JQowMPNYpSpNPRFeeD8mlgs38Xw03g=",
       "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   },
@@ -27,7 +27,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "dc57d15b8bb1d4b52d82012f7f789ae5ae6db17f",
+      "sha": "6971f7a37410ba0d6edae93b389dbcbe3ba924d5",
       "dirty": false
     }
   }