liangsu9988 commited on Jul 18

Commit

af4f1aa

verified ·

1 Parent(s): b3874be

Uploaded using `kernel-builder`.

Browse files

Files changed (37) hide show

benchmarks/benchmark_linear_bias.py +338 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +112 -0
build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so → _flashrt_fp8_ffn_cuda_255584f.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +112 -0
build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so → _flashrt_fp8_ffn_cuda_255584f.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-rocm71-x86_64-linux/__init__.py +112 -0
build/torch211-cxx11-rocm71-x86_64-linux/{_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so → _flashrt_fp8_ffn_rocm_255584f.abi3.so} +2 -2
build/torch211-cxx11-rocm71-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-rocm71-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-rocm72-x86_64-linux/__init__.py +112 -0
build/torch211-cxx11-rocm72-x86_64-linux/{_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so → _flashrt_fp8_ffn_rocm_255584f.abi3.so} +2 -2
build/torch211-cxx11-rocm72-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-rocm72-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +112 -0
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so +0 -3
build/torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_255584f.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +112 -0
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so +0 -3
build/torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_255584f.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-rocm71-x86_64-linux/__init__.py +112 -0
build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so +0 -3
build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_rocm_255584f.abi3.so +3 -0
build/torch212-cxx11-rocm71-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-rocm71-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-rocm72-x86_64-linux/__init__.py +112 -0
build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so +0 -3
build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_255584f.abi3.so +3 -0
build/torch212-cxx11-rocm72-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-rocm72-x86_64-linux/metadata.json +1 -1

benchmarks/benchmark_linear_bias.py ADDED Viewed

	@@ -0,0 +1,338 @@

+#!/usr/bin/env python3
+"""Benchmark generic FP8 linear+bias projection paths."""
+from __future__ import annotations
+import argparse
+import importlib
+import importlib.util
+import json
+import statistics
+import sys
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+TEST_FILE = ROOT / "flashrt-fp8-ffn" / "tests" / "test_fp8_ffn.py"
+SHAPES = {
+    "decode_m1_d2048_o": (1, 2048, 2048),
+    "small_m8_d1536_o": (8, 1536, 1536),
+    "groot_dit_m51_qkv": (51, 1536, 4608),
+    "groot_dit_m51_o": (51, 1536, 1536),
+    "mid_m64_d2048_o": (64, 2048, 2048),
+    "groot_backbone_m105_qkv": (105, 2048, 4096),
+    "groot_backbone_m105_o": (105, 2048, 2048),
+    "mid_m128_siglip_qkv": (128, 1152, 3456),
+    "prefill_m256_d1536_o": (256, 1536, 1536),
+    "prefill_m512_d2048_o": (512, 2048, 2048),
+}
+def load_test_helpers():
+    spec = importlib.util.spec_from_file_location("flashrt_fp8_ffn_test_helpers", TEST_FILE)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load {TEST_FILE}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+def load_ops(backend: str, artifact: str | None):
+    helpers = load_test_helpers()
+    if backend == "source":
+        return helpers.load_source_ops(), helpers
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("flashrt_fp8_ffn"), helpers
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def cuda_time_us(fn, warmup: int, iters: int, rounds: int) -> float:
+    samples = []
+    for _ in range(rounds):
+        for _ in range(warmup):
+            fn()
+        torch.cuda.synchronize()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(iters):
+            fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0 / iters)
+    return statistics.median(samples)
+def graph_runner(fn):
+    fn()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    graph.replay()
+    torch.cuda.synchronize()
+    return graph.replay
+def maybe_fvk_runner(x_fp8, w_fp8, bias, out, x_scale, w_scale, M, N, K):
+    try:
+        from flash_rt import flash_rt_kernels as fvk
+    except ImportError:
+        return None, "flash_rt.flash_rt_kernels is unavailable"
+    gemm = fvk.GemmRunner()
+    alpha = float((x_scale * w_scale).item())
+    stream = torch.cuda.current_stream().cuda_stream
+    def run():
+        gemm.fp8_nn_bias_bf16(
+            x_fp8.data_ptr(), w_fp8.data_ptr(), out.data_ptr(), bias.data_ptr(),
+            M, N, K, alpha, stream,
+        )
+    try:
+        run()
+    except RuntimeError as error:
+        return None, str(error)
+    return run, None
+def maybe_fvk_fp16_runners(
+    x_fp8, w_fp8, bias_fp16, out_fp16, x_scale, w_scale, M, N, K
+):
+    try:
+        from flash_rt import flash_rt_kernels as fvk
+    except ImportError:
+        return None, None, "flash_rt.flash_rt_kernels is unavailable"
+    gemm = fvk.GemmRunner()
+    alpha = float((x_scale * w_scale).item())
+    stream = torch.cuda.current_stream().cuda_stream
+    def fused():
+        gemm.fp8_nn_bias(
+            x_fp8.data_ptr(), w_fp8.data_ptr(), out_fp16.data_ptr(),
+            bias_fp16.data_ptr(), M, N, K, alpha, stream,
+        )
+    def decomposed():
+        gemm.fp8_descale_fp16(
+            x_fp8.data_ptr(), w_fp8.data_ptr(), out_fp16.data_ptr(), M, N, K,
+            x_scale.data_ptr(), w_scale.data_ptr(), stream,
+        )
+        fvk.add_bias_fp16(
+            out_fp16.data_ptr(), bias_fp16.data_ptr(), M, N, stream
+        )
+    try:
+        fused()
+        fused_runner = fused
+        fused_error = None
+    except RuntimeError as error:
+        fused_runner = None
+        fused_error = str(error)
+    decomposed()
+    return fused_runner, decomposed, fused_error
+def metrics(got: torch.Tensor, expected: torch.Tensor) -> dict[str, float]:
+    got_f = got.float().flatten()
+    exp_f = expected.float().flatten()
+    diff = (got_f - exp_f).abs()
+    return {
+        "max_abs": float(diff.max().item()),
+        "mean_abs": float(diff.mean().item()),
+        "p99_abs": float(torch.quantile(diff, 0.99).item()),
+        "cosine": float(torch.nn.functional.cosine_similarity(got_f, exp_f, dim=0).item()),
+    }
+def run(args) -> dict:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    torch.manual_seed(20260718)
+    ops, helpers = load_ops(args.backend, args.artifact)
+    labels = list(SHAPES) if args.shapes == "all" else args.shapes.split(",")
+    results = []
+    for label in labels:
+        if label not in SHAPES:
+            raise ValueError(f"unknown shape {label!r}")
+        M, K, N = SHAPES[label]
+        x_bf16 = torch.randn((M, K), device="cuda", dtype=torch.bfloat16) * 0.25
+        w_bf16 = torch.randn((N, K), device="cuda", dtype=torch.bfloat16) * (K**-0.5)
+        bias = torch.randn((N,), device="cuda", dtype=torch.bfloat16) * 0.01
+        x_scale = (
+            x_bf16.float().abs().max() / (0.9 * helpers.fp8_max())
+        ).clamp_min(1e-6).reshape(1)
+        w_scale = (
+            w_bf16.float().abs().max() / (0.9 * helpers.fp8_max())
+        ).clamp_min(1e-6).reshape(1)
+        x_fp8 = helpers.quantize_fp8_reciprocal(x_bf16, x_scale)
+        w_fp8 = helpers.quantize_fp8(w_bf16, w_scale)
+        out = torch.empty((M, N), device="cuda", dtype=torch.bfloat16)
+        region_out = torch.empty_like(out)
+        input_fp8 = torch.empty_like(x_fp8)
+        fvk_out = torch.empty_like(out)
+        fvk_fp16_out = torch.empty((M, N), device="cuda", dtype=torch.float16)
+        bias_fp16 = bias.to(torch.float16)
+        def package_fp8():
+            ops.fp8_linear_bias_bf16(
+                x_fp8, w_fp8, bias, x_scale, w_scale, out=out
+            )
+        def package_region():
+            ops.bf16_fp8_linear_bias_bf16(
+                x_bf16, w_fp8, bias, x_scale, w_scale,
+                input_fp8=input_fp8, out=region_out, pad_to=M,
+            )
+        reference = helpers.ref_linear_bias(
+            x_fp8, w_fp8, bias, x_scale, w_scale
+        )
+        package_fp8()
+        package_metrics = metrics(out, reference)
+        if package_metrics["p99_abs"] > 0.015625 or package_metrics["cosine"] < 0.9999:
+            raise AssertionError(f"{label} package correctness failed: {package_metrics}")
+        package_region()
+        if not torch.equal(region_out, out) or not torch.equal(input_fp8, x_fp8):
+            raise AssertionError(f"{label} BF16 region does not match FP8 entry")
+        graph = graph_runner(package_region)
+        if args.compare_fvk:
+            fvk, fvk_error = maybe_fvk_runner(
+                x_fp8, w_fp8, bias, fvk_out, x_scale, w_scale, M, N, K
+            )
+        else:
+            fvk, fvk_error = None, "not requested"
+        if args.compare_fvk:
+            fvk_fp16_fused, fvk_fp16_decomposed, fvk_fp16_error = (
+                maybe_fvk_fp16_runners(
+                    x_fp8, w_fp8, bias_fp16, fvk_fp16_out, x_scale, w_scale,
+                    M, N, K,
+                )
+            )
+        else:
+            fvk_fp16_fused = None
+            fvk_fp16_decomposed = None
+            fvk_fp16_error = "not requested"
+        if fvk is not None:
+            fvk()
+            fvk_metrics = metrics(fvk_out, reference)
+            if fvk_metrics["p99_abs"] > 0.015625 or fvk_metrics["cosine"] < 0.9999:
+                raise AssertionError(f"{label} FVK correctness failed: {fvk_metrics}")
+        else:
+            fvk_metrics = None
+        def eager(value):
+            return torch.addmm(bias, value, w_bf16.T)
+        compiled = torch.compile(eager, fullgraph=True) if args.compile_baseline else None
+        if compiled is not None:
+            compiled_out = compiled(x_bf16)
+            changed_out = compiled(x_bf16 + torch.ones_like(x_bf16) * 0.125)
+            torch.cuda.synchronize()
+            if torch.equal(compiled_out, changed_out):
+                raise AssertionError(
+                    f"{label} compiled baseline did not respond to changed input"
+                )
+        package_us = cuda_time_us(package_fp8, args.warmup, args.iters, args.rounds)
+        region_us = cuda_time_us(package_region, args.warmup, args.iters, args.rounds)
+        graph_us = cuda_time_us(graph, args.warmup, args.iters, args.rounds)
+        eager_call = lambda: eager(x_bf16)
+        compiled_call = (lambda: compiled(x_bf16)) if compiled is not None else None
+        eager_us = cuda_time_us(eager_call, args.warmup, args.iters, args.rounds)
+        compiled_us = (
+            cuda_time_us(compiled_call, args.warmup, args.iters, args.rounds)
+            if compiled_call is not None else None
+        )
+        fvk_us = (
+            cuda_time_us(fvk, args.warmup, args.iters, args.rounds)
+            if fvk is not None else None
+        )
+        fvk_fp16_fused_us = (
+            cuda_time_us(
+                fvk_fp16_fused, args.warmup, args.iters, args.rounds
+            )
+            if fvk_fp16_fused is not None else None
+        )
+        fvk_fp16_decomposed_us = (
+            cuda_time_us(
+                fvk_fp16_decomposed, args.warmup, args.iters, args.rounds
+            )
+            if fvk_fp16_decomposed is not None else None
+        )
+        row = {
+            "shape": label,
+            "M": M,
+            "K": K,
+            "N": N,
+            "fp8_linear_bias_us": package_us,
+            "bf16_region_us": region_us,
+            "bf16_region_graph_us": graph_us,
+            "torch_bf16_eager_us": eager_us,
+            "torch_bf16_compiled_us": compiled_us,
+            "fvk_fp8_nn_bias_bf16_us": fvk_us,
+            "fvk_status": "ok" if fvk is not None else "unsupported",
+            "fvk_error": fvk_error,
+            "fvk_fp16_fused_us": fvk_fp16_fused_us,
+            "fvk_fp16_decomposed_us": fvk_fp16_decomposed_us,
+            "fvk_fp16_fused_error": fvk_fp16_error,
+            "package_vs_eager": eager_us / package_us,
+            "region_vs_eager": eager_us / region_us,
+            "graph_vs_eager": eager_us / graph_us,
+            "package_vs_fvk": fvk_us / package_us if fvk_us is not None else None,
+            "package_vs_fvk_fp16_fused": (
+                fvk_fp16_fused_us / package_us
+                if fvk_fp16_fused_us is not None else None
+            ),
+            "package_vs_fvk_fp16_decomposed": (
+                fvk_fp16_decomposed_us / package_us
+                if fvk_fp16_decomposed_us is not None else None
+            ),
+            "package_correctness": package_metrics,
+            "fvk_correctness": fvk_metrics,
+        }
+        results.append(row)
+        print(json.dumps(row, sort_keys=True))
+    return {
+        "backend": args.backend,
+        "device": torch.cuda.get_device_name(0),
+        "capability": list(torch.cuda.get_device_capability(0)),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "warmup": args.warmup,
+        "iters": args.iters,
+        "rounds": args.rounds,
+        "results": results,
+    }
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--shapes", default="all")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--rounds", type=int, default=5)
+    parser.add_argument("--compile-baseline", action="store_true")
+    parser.add_argument("--compare-fvk", action="store_true")
+    parser.add_argument("--output", type=Path, default=None)
+    args = parser.parse_args()
+    payload = run(args)
+    if args.output is not None:
+        args.output.parent.mkdir(parents=True, exist_ok=True)
+        args.output.write_text(json.dumps(payload, indent=2) + "\n")
+if __name__ == "__main__":
+    main()

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

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so → _flashrt_fp8_ffn_cuda_255584f.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:3d274681fba749480e852c3da2670c29041914b6edfb234ebce6e46bec1afa8c
+size 698640

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so → _flashrt_fp8_ffn_cuda_255584f.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:2636112b16bd21bec7d249cf99d86206b873df2f8672fe47a560e57171b55cd2
+size 668320

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-rocm71-x86_64-linux/__init__.py CHANGED Viewed

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch211-cxx11-rocm71-x86_64-linux/{_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so → _flashrt_fp8_ffn_rocm_255584f.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7e9b5e1e2b4d7de01d7da70219952c7587879ddd254f15f9f60060c2e2f9838c
-size 167328

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

build/torch211-cxx11-rocm71-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch211-cxx11-rocm71-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_rocm_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-rocm72-x86_64-linux/__init__.py CHANGED Viewed

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch211-cxx11-rocm72-x86_64-linux/{_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so → _flashrt_fp8_ffn_rocm_255584f.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:01f36636ddf259f85c94157b09cb5267b9cd4b371b148acd5150f4f3c13f3b78
+size 178992

build/torch211-cxx11-rocm72-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch211-cxx11-rocm72-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_rocm_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch212-cxx11-cu130-x86_64-linux/_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so DELETED Viewed

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:9cf974f0a6585860d33d00d4269205d2046cbaa8e74442d4e289ec8ea7cddc6e
+size 668544

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch212-cxx11-cu132-x86_64-linux/_flashrt_fp8_ffn_cuda_1b7b2db.abi3.so DELETED Viewed

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

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

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:92f9e9ba898c08f78f8ea086de134a0f1dcfb1201ce9d04ca932d4ed30844919
+size 668544

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_cuda_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_cuda_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-rocm71-x86_64-linux/__init__.py CHANGED Viewed

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so DELETED Viewed

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

build/torch212-cxx11-rocm71-x86_64-linux/_flashrt_fp8_ffn_rocm_255584f.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:6203acdac17af6b44944b57bb831a4584a39b8ad401f0e8cd58a7c2b5760862e
+size 180000

build/torch212-cxx11-rocm71-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch212-cxx11-rocm71-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_rocm_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-rocm72-x86_64-linux/__init__.py CHANGED Viewed

@@ -55,6 +55,50 @@ def _fp8_gemm_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
@@ -157,6 +201,33 @@ def fp8_gemm_bf16(
     return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -260,6 +331,45 @@ def _midm_padded_rows(input: torch.Tensor) -> int:
     return rows
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
@@ -324,8 +434,10 @@ def bf16_fp8_gelu_mlp_bf16(
 __all__ = [
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

     return None
+@torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16"))
+def _fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (input.shape[0], weight.shape[0]):
+        raise RuntimeError("out shape must be (input.shape[0], weight.shape[0])")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("bf16_fp8_linear_bias_bf16"))
+def _bf16_fp8_linear_bias_bf16_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 2 or weight.dim() != 2:
+        raise RuntimeError("input and weight must be rank-2 tensors")
+    if input.shape[1] != weight.shape[1]:
+        raise RuntimeError("input.shape[1] must match weight.shape[1]")
+    padded_m = input_fp8.shape[0]
+    if padded_m < input.shape[0]:
+        raise RuntimeError("input_fp8 padded rows must cover input rows")
+    if input_fp8.shape != (padded_m, input.shape[1]):
+        raise RuntimeError("input_fp8 must have shape (padded_m, input.shape[1])")
+    if bias.shape != (weight.shape[0],):
+        raise RuntimeError("bias shape must be (weight.shape[0],)")
+    if out.shape != (padded_m, weight.shape[0]):
+        raise RuntimeError("out must have shape (padded_m, weight.shape[0])")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_quant_bf16"))
 def _fp8_linear_bias_gelu_quant_bf16_fake(
     input: torch.Tensor,
     return out
+def fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run an FP8-weight linear projection with BF16 bias and output.
+    ``input`` and ``weight`` use FP8 E4M3 with shapes ``(M, K)`` and
+    ``(N, K)``. Scales are device float32 scalar tensors. Pass ``out`` to keep
+    the call allocation-free for CUDA Graph capture.
+    """
+    if out is None:
+        out = torch.empty(
+            (input.shape[0], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, out
+    )
+    return out
 def fp8_linear_bias_gelu_quant_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
     return rows
+def bf16_fp8_linear_bias_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_fp8: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    *,
+    pad_to: int | None = None,
+) -> torch.Tensor:
+    """Run a static-scale FP8 linear projection from a BF16 boundary.
+    This traceable region entry quantizes ``input`` and computes the FP8
+    projection with BF16 bias/output. Thor SM110 pads M=9..128 to a multiple of
+    64 by default. Preallocate ``input_fp8`` and ``out`` for allocation-free
+    CUDA Graph replay.
+    """
+    logical_m = input.shape[0]
+    padded_m = _midm_padded_rows(input) if pad_to is None else pad_to
+    if padded_m < logical_m:
+        raise ValueError("pad_to must be >= input.shape[0]")
+    if input_fp8 is None:
+        input_fp8 = torch.empty(
+            (padded_m, input.shape[1]), device=input.device, dtype=_fp8_dtype()
+        )
+    if out is None:
+        out = torch.empty(
+            (padded_m, weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.bf16_fp8_linear_bias_bf16(
+        input, weight, bias, input_scale, weight_scale, input_fp8, out
+    )
+    return out[:logical_m]
 def bf16_fp8_gelu_mlp_bf16(
     input: torch.Tensor,
     up_weight: torch.Tensor,
 __all__ = [
+    "bf16_fp8_linear_bias_bf16",
     "bf16_fp8_gelu_mlp_bf16",
     "fp8_gemm_bf16",
     "fp8_gelu_mlp_bf16",
+    "fp8_linear_bias_bf16",
     "fp8_linear_bias_gelu_quant_bf16",
 ]

build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_1b7b2db.abi3.so DELETED Viewed

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

build/torch212-cxx11-rocm72-x86_64-linux/_flashrt_fp8_ffn_rocm_255584f.abi3.so ADDED Viewed

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

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

build/torch212-cxx11-rocm72-x86_64-linux/_ops.py CHANGED Viewed

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

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

build/torch212-cxx11-rocm72-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-fp8-ffn",
-  "id": "_flashrt_fp8_ffn_rocm_1b7b2db",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-fp8-ffn",
+  "id": "_flashrt_fp8_ffn_rocm_255584f",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],