Marlin moe integration

[ElizaWszola]

Unit testing:

pytest tests/kernels/test_moe.py -k test_fused_marlin_moe

(requires to uncomment @pytest.mark.skip in test_moe.py).

End-to-end testing:
Run offline_inference.py with

// quantized moe with act order
llm = LLM(model="TheBloke/Mixtral-8x7B-v0.1-GPTQ",
          revision="gptq-4bit-128g-actorder_True",
          enforce_eager=True)

and

// quantized moe without act order
llm = LLM(model="TheBloke/Mixtral-8x7B-v0.1-GPTQ", enforce_eager=True)

[robertgshaw2-neuralmagic]

This function does not need to adhere to the exact same interface as fused_moe

This function will be called on the hotpath. It should receive INT4 weights and scales and just call the marlin moe kernel directly

[robertgshaw2-neuralmagic]

This will not be called on the hotpath.

Rather, the quantized weights should be an input to this function

[robertgshaw2-neuralmagic]

see comments in code

[robertgshaw2-neuralmagic]

So a couple things. vLLM is layed out in the following way

Models --> llama.py, which uses linear_layers like ColumnParallelLinear

Each layer has a LinearMethod which handles the representation of the weights and the forward pass
For instance, we have an Fp16 linear method, Marlin linear method, etc., etc.

Each LinearMethod exposes the following interface:
create_weights --> defines what the weights look like (e.g. dtype, whether there are scales, etc)
apply --> calls the kernels during the forward pass

So, we will eventually want to create a LinearMethod for MarlinMoE
create_weights --> will load the int4 weights from disk
apply --> passes the weights to your fused_moe_marlin kernel

As a result, the fused_moe_marlin kernel should recieve already quantized weights and just execute the computation. We should not be quantizing the weights inside of that function

For this PR, we should land the kernel + testing code for the kernel. we can work on adding the LinearMethod afterwards

[robertgshaw2-neuralmagic]

Per my comment below, we will load the already compressed weights via create_weights.

So none of this will need to be called on the hotpath

As a result, all of this should be moved into testing utilities

[tlrmchlsmth]

I think it would be good to add a test to make sure this works on other GPUs as well (we do this in the cutlass unit tests, if you want to replicate that here)

[ElizaWszola]

Do you mean testing on different devices? (@pytest.mark.parametrize("device", CUDA_DEVICES))

[tlrmchlsmth]

yes exactly

[ElizaWszola]

Doing anything on cuda:1 results in memory erros (illegal access) in moe_align_block_size_kernel which I rely on, but didn't modify - should I look into it or is it ok to leave it for now?

[tlrmchlsmth]

is 16 a hardcoded block size?

[ElizaWszola]

This is related to Marlin format which is hardcoded

[tlrmchlsmth]

are these guaranteed to be fp16?

[ElizaWszola]

From Pytorch documentation: self.half() is equivalent to self.to(torch.float16)
Scales are not necessarily fp16 when loaded.

[tlrmchlsmth]

how much of this file is copy-pasted from the original marlin code? Could we factor out common functions? It will make it much easier to review if we can see what the new code is

[ElizaWszola]

There is quite a bit of overlap, and many of changes boil down to adding one variable or an extra condition here and there. I don't really want to refactor into common functions until act_order is done, because there might be more of these tiny modifications (or is it better to do the refactor now?).

In any case, running a comparison of this file against csrc/quantization/gptq_marlin/gptq_marlin.cu helps seeing what changed.

Edit: fixed file name

[tlrmchlsmth]

That’s fair for things that may be changed by act_reorder but any functions that are copied over unmodified should be factored out IMO

[ElizaWszola]

Moved to the other repo