Hi, it appears that I am encountering a GCC error when attempting to install your project. The error message reads: gcc: error trying to exec 'cc1plus': execvp: No such file or directory. It seems that the issue may be caused by pycuda. Do you have any recommended containers that I could use to quickly run the code?

Can you provide some example code about testing SparTA on float16 Linear or Matmul operators?

I keep getting '[root] All trials failed.' error for every example I try to run. It has been installed properly and everything works fine until it reaches 'best_config = sparta.nn.tune', then I get this error.

updated

SparTA specializer generates customized codes for the sparse operators.

Highlight difference with existing implementation

1. use jinja2 to handle templates
2. use PyTorch JIT Cpp extension to compile generated kernel codes, load kernel function, and test via numpy (torch) tensor
3. decouple TeSA transform (e.g., block coverage) and computing code emitter
4. definition of TeSA class, make it extensible for quantization formats and diverse storage layout (e.g., CSR, BCSR, etc.)
5. tune the performance in op level, other than in kernel function level
6. make the sparse operator class clean for example, let external tuner hold tuning history

Known limits

1. only support single transform and emitter in existing base class; we may need a nested search space and decorate to define a pipeline with multiple transform or emitter choices
2. not support use sparse operator as submodule yet, but I believe current interface is compatible (that's why I use op-level tuning)
3. simplify the burden of tuner, we could support value choice (selecting from a list or set) only

Sparse operators based on single TeSA transformer and emitter

SparTA offers some basic sparse operators e.g., sparse linear, sparse multi-head self-attention. These operators are based on the efficient sparse computing library (template).

from torch.utils.cpp_extension import load

__kernel_binding__ = None # must run op.specialize() first

class DynamicSparseLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(...): 
        return __kernel_binding__.forward_func(...)
    @staticmethod
    def backward(...):
        return __kernel_binding__.backward_func(...)

class DynamicSparseLinear(nn.Linear, SparseOpBase):

    def __init_specialization__(self, ...):
        self.tesa_transform = BlockCoverage
        self.emitter = TemplateEmitterBase(templates=['operators/dynamic_sparse_linear.cu', 'kernels/sparta_sparse_matmul.cu'])
        self.search_space = SearchSpace({ # we could support only value choice first
            'transform_params': ValueChoice((32,32), (32,64), (64,32), (64,64)), # value choice is a list / set of 
            'emitter_params': SearchSpace({
                'BLOCK_SIZE_M_VALUE': ValueChoice(8,16,32,64,128),
                #...
            })
        })

    def forward(self, ...):
        return DynamicSparseLinearFunction.apply(...)
    
    def specialize(self, tesa: dict, sample: TemplateParam):
        transform_params, emitter_params = ..(sample)..
        self.weight_tesa = self.tesa_transform(tesa['weight'], **transform_params)
        # meta information such as BCSR format (block size), number of non-zero block, indexes (even bidirectional) are embedded
        emitter_params.update(...) # update according to tesa
        files = self.emitter.code_gen(emitter_params)
        __kernel_binding__ = load(name, sources = files)
        return self

    def performance_optimization(self, tesa: dict, tuner_cls: type(Tuner), evaluator: Evaluator):
        tuner = tuner_cls(self.search_space)
        while not tuner.completed():
            samp = tuner.get_param()
            self.specialize(tesa, samp)
            metric = evaluator(self) # forward or backward step latency
            tuner.receive_metric(metric)
        return self.specialize(tesa, tuner.best_param()) 

class SparseWeightLinearEvaluator:

    def __init__(self, in_features, out_features):
        self.ref_op = nn.Linear(in_features, out_features)
        self.inputs = ...

    def evaluate(self, spop: SparseWeightLinear):
        pass

# usage example
weight_mask = ..Tensor..
weight_tesa = TeSAUtils.from_mask(weight_mask)
splinear = SparseWeightLinear(in_dim, out_dim)
splinear.tune_specialize({'weight': weight_tesa}, RandomSearchTuner, SparseWeightLinearEvaluator().evaluate)

TeSA and transform

@dataclass
class TeSA:
    mask: torch.Tensor # int tensor
    formats: List[str|dict]
    layout: dict = None # or name it as storage

class TeSAUtils
    @staticmethod
    def from_mask(mask: torch.Tensor, formats: list = None):
        return TeSA(mask=mask, layout={'layout': 'dense', 'meta': None},
                    formats=['float32', 'pruned'] if formats is None else formats)

The mask is integer tensor indicate each element's format, its value is between 0 and #formats - 1.
The formats is the registered formats corresponding to the integer values in mask, for example, the common setting for pruning and quantization could be like

pruning_tesa = TeSA(..mask{0,1}.., ['float32', 'pruned'])
pruning_quant_tesa = TeSA(..mask{0,1}.., [{'format': 'int8', 'scale':.., 'zero-point': 0}, 'pruned'])

The layout is {'layout': 'dense', 'meta': None} by default, which means the tensor is stored in dense format. After transformation (or convert), it could be changed to a detailed sparse format like

{
    'layout': 'BCSR',
    'meta': {
        'block': (32,32),
        'index': ..numpy.array.., # optional
        'transpose_index': ..numpy.array.., # optional
    } 
}

Emitter

The emitter class generate kernel codes for specific functions, e.g., sparse matrix multiplication, from a given template or scheduling primitives.

Here is an example of template based emitter (via jinja2).

from jinja2 import Template

class TemplateEmitterBase:

    def __init_template__(self, template_names: str, compile_opts: List[str]):
        self.templates = [t:Template(..path/to/template_name..) for t in template_names]
        self.compile_opts = compile_opts

    def code_gen(self, params: dict):
        src_code = [k:v.render(*params) for k,v in self.templates]
        files = .. write to disk ..
        return files

In fact we have two kinds of templates, the computing kernels and the operators interface. We could separate them by folders like:

templates
    |- kernels
        |- sparta_sparse_matmul.cu
        |- ...
    |- operators
        |- sparse_weight_linear.cu
        |- sparse_msa.cu

We could leverage the template engine to write complicated templates, for example, matmul with or without bias input as below

__global__ void BLOCK_SPARSE_MATMUL_BIAS(
    float* A, float* W_val, int* W_row, int* W_col,
    float* C, 
{% if USE_BIAS %}
    float *bias, 
{% endif %}
    int M, int K, int N){
        ...
    }

or we could let the template selecting different computing kernels according to its block size

void dynamic_forward_function(float* activation, int* row_ptr, int* col_idx,
                    float * val, float* bias, int M, int K, int N, int block_h, int block_w, float* output)
{
    dim3 gridDim(N/{{ BLOCK_SIZE_N }}, M/{{ BLOCK_SIZE_M }});
    dim3 blockDim({{ BLOCK_SIZE_N }} / {{ THREAD_SIZE_N }}, {{ BLOCK_SIZE_M }}/{{ THREAD_SIZE_M }});

    {% if (BLOCK_SIZE_N == 32) and (BLOCK_SIZE_M == 32) %}
    ...efficient 32x32 kernel...
    {% else %}
    ...general kernel...
    {% endif%}
}

I followed the README to install SparTa, but it failed even though there were no errors during the installation.
My CUDA, pytorch and pycuda are all adapted to 11.8：

But when I use the 'Tune a sparse operator' in the README, I still get an error: ModuleNotFoundError: No module named 'sparse_moe_cpp'. I don't know what's wrong with that.

• Estimated release date:
  • public preview (alpha): 9/1
  • public preview (beta): 9/30
  • refactor kernels & TeSA: 10/15

P0

• Tuning more steps to show the speedup gain of the pytorch sparse modules
• Support the openai kernel/template
• code review
• Usage Interface(8.19) (update one version on 8.26)
• Fix triton speed(8.19)
• Sparse Softmax Kernel
• Biased OpenAI MatMul Kernel
• finegrained 99% + block size 8x8 95% + block size 32 x 32
• Documentation (test)
• package data (test)
• sparta.tune(): hook, set search space
• Fix sparse softmax
• Integration test/example: Linear, Softmax
• Fix JIT latency
• Read the docs
• SparTA DDS MatMul kernel
• Batch MatMul & Softmax
• Sparse Attention
• Add sparse matmul kernel: transpose_A
• Functional
• Support backward
• Add perfermance test: Compare with Triton 1.1.2 (Upload test scripts)
• Test current tuner
• Test Sparse Attention
• Update kernel pycuda interface
• Profile Layout converting
• Construct sparse attention op with linear & softmax ops
• Beta version: docs, docstrings & examples
• Test on V100; backward
• Fix kernel output
• Module tuner: get combined search space of connected ops automatically
• Connect to NNI's new tuner

P1

• Apply roller's rules
• Support multi-process tuning
• BCSR kernel: convert(), inverse(), swapaxes(), sum(), rebuild TeSA Converter when set_mask()
• Auto converter: support value mask in matmul kernels
• PyCUDA device context register & operator.to() (multiple cards)
• Support the multiple sparse formats: sdd dsd, dds for linear
• Support the block quantization kernel/fp16/bf16
• Compare Sparse Softmax with Triton's Sparse Softmax and keep improving.
• unit tests
• Model tuning interface / documents / examples
• Common mask patterns
• Refactor TeSA (Meta, linter)
• Fuse layout converting into kernels

P2

• Support the offline LUT or the kernel cache/DB

The second way to build SparTA in Get Start part in README is wrong.
The correct way to build it from source code may be this:

git clone https://github.com/microsoft/SparTA.git
cd SparTA
pip install -r requirements.txt
python setup.py install

and there is a bug in requirements.txt, it is missing a package ninjia.
If ninjia is not installed, it will cause an undefined symbol: xxxxxx.so xxxxx error

Does SparTa or pit support operator inference on the CPU? If so, are there any test cases?

Work Items

Propagation

• Support propagation for SparTA(import from NNI or move the speedup into SparTA?). 6.1 @zheng-ningxin need discuss about the limitation of current implementation

Codegen

• Block Sparse Matmul Code generation
  • Template(with/without fusion) (formalize the yaml format)
  • Emitter for different frameworks/libraries
  • Transformation policy()
  • Environment & measurement
  • Parallel process pool
• Finegrained
• Block quantization
• Block+finegrained

Test

• Matmul block sparse end2end code generation