File size: 316 Bytes
c41750d
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
#pragma once

#include <torch/extension.h>

#include <optional>

torch::Tensor smallm_nvfp4_linear_cuda(
    const torch::Tensor& input,
    const torch::Tensor& packed_weight,
    const torch::Tensor& weight_block_scales,
    const torch::Tensor& weight_tensor_scale,
    const std::optional<torch::Tensor>& bias);