[Bugfix][Kernel] Fix moe align block issue for mixtral

[ElizaWszola]

Fix an issue with shared arrays in moe_align_block_size_kernel that was causing Mixtral inference to crash.

Testing: run inference with
llm = LLM(model="TheBloke/Mixtral-8x7B-v0.1-GPTQ")

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[mgoin]

Could you please take a look @jinzhen-lin

[tlrmchlsmth]

Thanks for the fix!

[tlrmchlsmth]

Basically we were only asking for num_experts + 1 elements in shared_memory for cumsum. But the offset for tokens_cnts in the kernel was as if cumsum needed max(num_experts, warp_size) + 1.

vllm/csrc/moe/moe_align_sum_kernels.cu

Lines 35 to 36 in 3132a93

	int32_t* cumsum = shared_mem; // 1d tensor with shape (num_experts + 1)
	token_cnts_t* tokens_cnts = (token_cnts_t*)(shared_mem + blockDim.x + 1);

vllm/csrc/moe/moe_align_sum_kernels.cu

Lines 229 to 230 in 3132a93

	const int32_t shared_mem_i32 =
	((num_thread + 1) * num_experts + (num_experts + 1)) * sizeof(int32_t);

So we were trying to use more shared memory than we asked for. Since we only need num_experts + 1 elements, we have this fix

[mgoin]

Thank you for the careful fix