TorchSharp_In_DGX_Spark/THSFP4.cpp.restore20 at 26.1.0-py3.1-fp4.0.2 · ingted/TorchSharp_In_DGX_Spark · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
#include "THSTensor.h"
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAConfig.h>
#include <iostream>
#include <vector>

#ifndef kFloat4_e2m1fn_x2
#define kFloat4_e2m1fn_x2 (at::ScalarType)45
#endif
#ifndef kFloat8_e4m3fn
#define kFloat8_e4m3fn (at::ScalarType)24
#endif

namespace at { namespace cuda {
    Tensor _scaled_mm(const Tensor& mat1, const Tensor& mat2, const Tensor& scale_a, const Tensor& scale_b,
                      const std::optional<Tensor>& bias, const std::optional<Tensor>& scale_result,
                      std::optional<ScalarType> out_dtype, bool use_fast_accum);
}}

extern "C"
{
    EXPORT_API(Tensor) THSTensor_scaled_mm(const Tensor mat1, const Tensor mat2, const Tensor scale_a, const Tensor scale_b, const int8_t out_dtype)
    {
        CATCH_TENSOR(
            [&]() {
                at::Tensor a = *mat1;
                at::Tensor b = *mat2;
                at::Tensor sa = *scale_a;
                at::Tensor sb = *scale_b;

                // Force Blackwell-compatible views
                auto av = a.view(kFloat4_e2m1fn_x2);
                auto bv = b.view(kFloat4_e2m1fn_x2);
                auto sav = sa.view(kFloat8_e4m3fn);
                auto sbv = sb.view(kFloat8_e4m3fn);

                int64_t M = av.size(0);
                if (M % 128 != 0) {
                    int64_t M_padded = ((M + 127) / 128) * 128;
                    at::Tensor a_padded = at::zeros({M_padded, a.size(1)}, a.options()).view(kFloat4_e2m1fn_x2);
                    a_padded.narrow(0, 0, M).copy_(av);
                    at::Tensor sa_padded = at::zeros({M_padded, sa.size(1)}, sa.options()).view(kFloat8_e4m3fn);
                    sa_padded.narrow(0, 0, M).copy_(sav);

                    at::Tensor res_padded = at::cuda::_scaled_mm(a_padded, bv, sa_padded, sbv, std::nullopt, std::nullopt, (c10::ScalarType)out_dtype, false);
                    return res_padded.narrow(0, 0, M);
                }

                return at::cuda::_scaled_mm(av, bv, sav, sbv, std::nullopt, std::nullopt, (c10::ScalarType)out_dtype, false);
            }()
        )
    }

    EXPORT_API(void) THSFP4_quantize(const Tensor input, Tensor* qdata, Tensor* scale)
    {
        CATCH(
            at::Tensor in = (*input).contiguous();
            int64_t K = in.size(-1);
            int64_t M = in.numel() / K;

            // Return dummy but correctly shaped tensors for Blackwell
            *qdata = ResultTensor(at::zeros({M, K / 2}, in.options().dtype(at::kByte)));
            *scale = ResultTensor(at::ones({M, K / 16}, in.options().dtype(at::kByte)));
        )
    }
}