Skip to content

GitLab

Unverified Commit 9101d832 authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[Bugfix] Make moe_align_block_size AMD-compatible (#3470)

parent 93348d94
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
csrc/moe_align_block_size_kernels.cu
View file @ 9101d832
......@@ -111,7 +111,8 @@ void moe_align_block_size(
// set dynamic shared mem
auto kernel = vllm::moe_align_block_size_kernel<scalar_t>;
AT_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_mem));
AT_CUDA_CHECK(
VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize((void *)kernel, shared_mem));
kernel<<<1, num_experts, shared_mem, stream>>>(
topk_ids.data_ptr<scalar_t>(),
sorted_token_ids.data_ptr<int32_t>(),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment