[KERNEL] int8 quantization kernel refactoring & optimization WIP

[ZelboK]

This refactors the int8 quantization kernel so that it will optimize f32 -> i8. However, please note that fp16 will require use of half2 though. That can be done with some template specializations but for brevity I will leave that outside of this PR for now.

1. Reorders instructions to improve ILP with unrolling(manual unrolling can also be done)
2. Use vectorized loads and stores for fp32, this will significantly increase memory throughput.
3. Refactors to use the reciprocal of the scale so that we can avoid doing division on the GPU, which is categorically slower because of how many more instructions it leads to.
4. From the profiling I've done, the problem is that there is a dependency chain of instructions going on. Warps get stalled at the multiplication becasue they are dependent on the memory instruction to complete before they get to do anything. So we want to interleave independent instructions(this is handled by the unrolling) in order to minimize the stalling.
5. Memory accesses are already coalesced from the original kernel. Since this streams into global memory, we aren't reusing data. Follows the same pattern as the original.

This was profiled on my 3080 and cuda 12.4.

PR Checklist (Click to Expand)

Thank you for your contribution to vLLM! Before submitting the pull request, please ensure the PR meets the following criteria. This helps vLLM maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Only specific types of PRs will be reviewed. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Model] for adding a new model or improving an existing model. Model name should appear in the title.
• [Frontend] For changes on the vLLM frontend (e.g., OpenAI API server, LLM class, etc.)
• [Kernel] for changes affecting CUDA kernels or other compute kernels.
• [Core] for changes in the core vLLM logic (e.g., LLMEngine, AsyncLLMEngine, Scheduler, etc.)
• [Hardware][Vendor] for hardware-specific changes. Vendor name should appear in the prefix (e.g., [Hardware][AMD]).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• We adhere to Google Python style guide and Google C++ style guide.
• Pass all linter checks. Please use format.sh to format your code.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Include sufficient tests to ensure the project to stay correct and robust. This includes both unit tests and integration tests.
• Please add documentation to docs/source/ if the PR modifies the user-facing behaviors of vLLM. It helps vLLM user understand and utilize the new features or changes.

Notes for Large Changes

Please keep the changes as concise as possible. For major architectural changes (>500 LOC excluding kernel/data/config/test), we would expect a GitHub issue (RFC) discussing the technical design and justification. Otherwise, we will tag it with rfc-required and might not go through the PR.

What to Expect for the Reviews

The goal of the vLLM team is to be a transparent reviewing machine. We would like to make the review process transparent and efficient and make sure no contributor feel confused or frustrated. However, the vLLM team is small, so we need to prioritize some PRs over others. Here is what you can expect from the review process:

• After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability.
• After the PR is assigned, the reviewer will provide status update every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team.
• After the review, the reviewer will put an action-required label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR.
• Please respond to all comments within a reasonable time frame. If a comment isn't clear or you disagree with a suggestion, feel free to ask for clarification or discuss the suggestion.

Thank You

Finally, thank you for taking the time to read these guidelines and for your interest in contributing to vLLM. Your contributions make vLLM a great tool for everyone!

[ZelboK]

just divides this by 4(since we're using float4 we need to do this as we're handling 4x as much work per iteration). Just avoiding the division assembly here.

[ZelboK]

This will improve ILP by reordering the instructions a bit to reduce the impact of the warp stalling from the memory load operation.

It seems that manual unrolling might give a very slight advantage over the macro though. However for maintainability it might not be worth doing that.

[ZelboK]

divide this on host so that we can use multiplication instructions instead of division in the kernel itself. Value is constant so works perfectly fine and simply.

[ZelboK]

So for fp16 we would need to do something like...

    struct half4
    {
        half2 x, y; // 2 half2s since cuda doesn't support half4
    };
    
     template <>
    __global__ void static_scaled_int8_quant_kernel<half4, float>(
        const half4 *__restrict__ input, int8_t *__restrict__ out,
        float inverted_scale, const int hidden_size)
    {
    ... 
     unroll 4
        for (int i = tid; i < traverse_space; i += blockDim.x)
        {
            int index = token_idx * traverse_space + i;

            half4 data_half4 = vectorized[index];

            // Convert half2 to float
            float2 data0 = __half22float2(data_half4.x);
            float2 data1 = __half22float2(data_half4.y);
            ...

etc which can just be a specialization. Can experiment with overloading as well to see what impacts this has on compile times(if important) since we are introducing another kernel

[robertgshaw2-redhat]

Could you focus on the fp16 / bf16 case? We won’t run the fp32 activations very often

[ZelboK]

Could you focus on the fp16 / bf16 case? We won’t run the fp32 activations very often

Yep that's fine with me. The code should mostly follow the same logic, will push when I get the chance!

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ZelboK.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork