diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 3eb7f0a444..56aeca2a8e 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -347,47 +347,67 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N" - key: lm-eval-humming-f16-a100 +- label: ":nvidia: (L4) LM Eval Humming TEMPORARY" + key: lm-eval-humming-l4 + optional: true + timeout_in_minutes: 90 + device: l4 + num_devices: 1 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/configs/humming/ + - requirements/cuda.txt + - vllm/utils/humming.py + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming/ + - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py + - vllm/model_executor/layers/fused_moe/oracle/ + - vllm/model_executor/kernels/linear/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt + +- label: ":nvidia: (A100) LM Eval Humming Default TEMPORARY Shard %N" + key: lm-eval-humming-a100-default + optional: true timeout_in_minutes: 75 device: a100 - optional: true num_devices: 1 parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY" - key: lm-eval-humming-act-a100 +- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY Shard %N" + key: lm-eval-humming-a100-act + optional: true timeout_in_minutes: 45 device: a100 - optional: true num_devices: 1 + parallelism: 2 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt -- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N" - key: lm-eval-humming-f16-h100 +- label: ":nvidia: (H100) LM Eval Humming Default TEMPORARY Shard %N" + key: lm-eval-humming-h100-default + optional: true timeout_in_minutes: 70 device: h100 - optional: true num_devices: 1 parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -395,14 +415,14 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY" - key: lm-eval-humming-act-h100 + key: lm-eval-humming-h100-act + optional: true timeout_in_minutes: 70 device: h100 - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -410,15 +430,35 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt # - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY" - key: lm-eval-humming-f16-b200 +- label: ":nvidia: (H100) LM Eval Humming DeepEP TEMPORARY" + key: lm-eval-humming-h100-deepep + optional: true + timeout_in_minutes: 60 + device: h100 + num_devices: 2 + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/configs/humming/ + - requirements/cuda.txt + - vllm/utils/humming.py + - vllm/model_executor/layers/quantization/humming.py + - vllm/model_executor/layers/quantization/utils/humming/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/kernels/linear/ + - vllm/distributed/ + - csrc/moe/ + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-deepep.txt + +- label: ":nvidia: (B200) LM Eval Humming Default TEMPORARY" + key: lm-eval-humming-b200-default + optional: true timeout_in_minutes: 50 device: b200-k8s - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ @@ -426,14 +466,14 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt - label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY" - key: lm-eval-humming-act-b200 + key: lm-eval-humming-b200-act + optional: true timeout_in_minutes: 50 device: b200-k8s - optional: true num_devices: 1 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming_utils.py + - vllm/model_executor/layers/quantization/utils/humming/ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ diff --git a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu index d66b78fe83..5c9a522feb 100644 --- a/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu +++ b/csrc/libtorch_stable/moe/moe_align_sum_kernels.cu @@ -236,7 +236,7 @@ __device__ void _moe_align_block_size_small_batch_expert( int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num, - int32_t* token_mask, bool has_expert_map) { + int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) { // Compute input buffer offsets. Typically these will all be 0, except when // using Multi LoRA. int sorted_token_ids_offset = max_num_tokens_padded * model_offset; @@ -318,9 +318,14 @@ __device__ void _moe_align_block_size_small_batch_expert( } for (size_t i = tid; i < numel; i += stride) { - if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, - num_experts, has_expert_map); - expert_id != -1) { + int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, + has_expert_map); + bool valid = + expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); + if (scatter_idx != nullptr) { + scatter_idx[i] = valid ? static_cast(i) : -1; + } + if (expert_id != -1) { int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; @@ -338,14 +343,20 @@ __device__ void _count_and_sort_expert_tokens( int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask, - int32_t model_offset, int32_t topk_num, bool has_expert_map) { + int32_t model_offset, int32_t topk_num, bool has_expert_map, + int32_t* scatter_idx = nullptr) { const size_t tid = blockIdx.y * blockDim.x + threadIdx.x; const size_t stride = blockDim.x * gridDim.y; for (size_t i = tid; i < numel; i += stride) { - if (int expert_id = get_local_expert_id(i, topk_ids, expert_map, - num_experts, has_expert_map); - expert_id != -1) { + int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts, + has_expert_map); + bool valid = + expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]); + if (scatter_idx != nullptr) { + scatter_idx[i] = valid ? static_cast(i) : -1; + } + if (expert_id != -1) { if (token_mask == nullptr || token_mask[i / topk_num]) { int32_t rank_post_pad = atomicAdd( &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); @@ -377,10 +388,11 @@ __global__ void count_and_sort_expert_tokens_kernel( const scalar_t* __restrict__ topk_ids, int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, - int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) { + int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map, + int32_t* scatter_idx) { _count_and_sort_expert_tokens( topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, - max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map); + max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx); } // Reduce the topk expert outputs per token (summed in fp32). The output is @@ -533,12 +545,12 @@ __global__ void moe_align_block_size_small_batch_expert_kernel( int32_t* __restrict__ total_tokens_post_pad, int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, size_t numel, int32_t max_num_tokens_padded, int32_t topk_num, - bool has_expert_map) { + bool has_expert_map, int32_t* scatter_idx) { _moe_align_block_size_small_batch_expert( topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, num_experts, block_size, numel, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr, - has_expert_map); + has_expert_map, scatter_idx); } template @@ -658,16 +670,31 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel( // taken from // https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc -void moe_align_block_size( - torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, - torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, - torch::stable::Tensor num_tokens_post_pad, - std::optional maybe_expert_map) { +void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, + int64_t block_size, + torch::stable::Tensor sorted_token_ids, + torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map, + std::optional scatter_idx) { const torch::stable::accelerator::DeviceGuard device_guard( topk_ids.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(topk_ids.get_device_index()); + int32_t* scatter_ptr = nullptr; + if (scatter_idx.has_value()) { + STD_TORCH_CHECK( + scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int, + "scatter_idx must be int32"); + STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() && + scatter_idx->is_contiguous() && + scatter_idx->numel() == topk_ids.numel(), + "scatter_idx must be contiguous, on the same device, and " + "match topk_ids.numel()"); + scatter_ptr = reinterpret_cast(scatter_idx->mutable_data_ptr()); + } + int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; int experts_per_warp = WARP_SIZE; @@ -713,7 +740,8 @@ void moe_align_block_size( num_tokens_post_pad.mutable_data_ptr()), reinterpret_cast(expert_map.mutable_data_ptr()), num_experts, block_size, topk_ids.numel(), - sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map, + scatter_ptr); } else { torch::stable::Tensor cumsum_buffer = torch::stable::new_empty( topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); @@ -753,7 +781,7 @@ void moe_align_block_size( reinterpret_cast(cumsum_buffer.mutable_data_ptr()), reinterpret_cast(expert_map.mutable_data_ptr()), topk_ids.numel(), num_experts, sorted_token_ids.size(0), - topk_ids.size(1), has_expert_map); + topk_ids.size(1), has_expert_map, scatter_ptr); } }); } diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 03057c3b14..549a19e981 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, std::optional expert_map); -void moe_align_block_size( - torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, - torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, - torch::stable::Tensor num_tokens_post_pad, - std::optional maybe_expert_map); +void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts, + int64_t block_size, + torch::stable::Tensor sorted_token_ids, + torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map, + std::optional scatter_idx); void batched_moe_align_block_size( int64_t max_tokens_per_batch, int64_t block_size, diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu index 251097456f..957cd60e49 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu @@ -56,20 +56,18 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert)); } -void moe_permute_impl( - const torch::stable::Tensor& input, // [n_token, hidden] - const torch::stable::Tensor& topk_ids, // [n_token, topk] - const torch::stable::Tensor& token_expert_indices, // [n_token, topk] +torch::stable::Tensor moe_sort_routing( + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] const std::optional& expert_map, // [n_expert] int64_t n_expert, int64_t n_local_expert, int64_t topk, - torch::stable::Tensor& permuted_input, // [permuted_size, hidden] torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] - torch::stable::Tensor& permuted_idx, // [permute_size] const std::optional& maybe_sort_workspace, const std::optional& maybe_permuted_experts_id, const std::optional& maybe_sorted_row_idx, - const std::optional& maybe_topk_ids_for_sort) { + const std::optional& maybe_topk_ids_for_sort, + bool build_inverse) { STD_TORCH_CHECK(expert_first_token_offset.scalar_type() == torch::headeronly::ScalarType::Long, "expert_first_token_offset must be int64"); @@ -88,12 +86,11 @@ void moe_permute_impl( "token_expert_indices shape must be same as inv_permuted_idx"); const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - auto device = input.device(); - auto n_token = input.sizes()[0]; - auto n_hidden = input.sizes()[1]; + topk_ids.get_device_index()); + auto device = topk_ids.device(); + auto n_token = topk_ids.sizes()[0]; auto expanded_rows = n_token * topk; - auto stream = get_current_cuda_stream(input.get_device_index()); + auto stream = get_current_cuda_stream(topk_ids.get_device_index()); auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert); auto sort_workspace = maybe_allocate_tensor( @@ -107,13 +104,11 @@ void moe_permute_impl( torch::headeronly::ScalarType::Int, device, "sorted_row_idx"); CubKeyValueSorter sorter{}; - int64_t* valid_num_ptr = nullptr; torch::stable::Tensor topk_ids_for_sort = topk_ids; - if (expert_map.has_value()) { - const int* expert_map_ptr = get_ptr(expert_map.value()); - valid_num_ptr = - get_ptr(expert_first_token_offset) + n_local_expert; + if (expert_map.has_value() || build_inverse) { + const int* expert_map_ptr = + expert_map.has_value() ? get_ptr(expert_map.value()) : nullptr; topk_ids_for_sort = maybe_allocate_tensor( maybe_topk_ids_for_sort, topk_ids.sizes(), torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort"); @@ -126,7 +121,39 @@ void moe_permute_impl( get_ptr(topk_ids_for_sort), get_ptr(token_expert_indices), get_ptr(permuted_experts_id), get_ptr(sorted_row_idx), get_ptr(expert_first_token_offset), n_token, n_expert, - n_local_expert, topk, sorter, get_ptr(sort_workspace), stream); + n_local_expert, topk, sorter, get_ptr(sort_workspace), stream, + build_inverse ? get_ptr(inv_permuted_idx) : nullptr); + return sorted_row_idx; +} + +void moe_permute_impl( + const torch::stable::Tensor& input, // [n_token, hidden] + const torch::stable::Tensor& topk_ids, // [n_token, topk] + const torch::stable::Tensor& token_expert_indices, // [n_token, topk] + const std::optional& expert_map, // [n_expert] + int64_t n_expert, int64_t n_local_expert, int64_t topk, + torch::stable::Tensor& permuted_input, // [permuted_size, hidden] + torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] + torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] + torch::stable::Tensor& permuted_idx, // [permute_size] + const std::optional& maybe_sort_workspace, + const std::optional& maybe_permuted_experts_id, + const std::optional& maybe_sorted_row_idx, + const std::optional& maybe_topk_ids_for_sort) { + const torch::stable::accelerator::DeviceGuard device_guard( + input.get_device_index()); + auto sorted_row_idx = moe_sort_routing( + topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert, + topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace, + maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort, + false); + auto stream = get_current_cuda_stream(input.get_device_index()); + auto n_token = input.size(0); + auto n_hidden = input.size(1); + auto valid_num_ptr = + expert_map.has_value() + ? get_ptr(expert_first_token_offset) + n_local_expert + : nullptr; MOE_DISPATCH(input.scalar_type(), [&] { expandInputRowsKernelLauncher( @@ -171,6 +198,22 @@ void moe_permute_with_scratch( topk_ids_for_sort); } +void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, + int64_t n_expert, int64_t n_local_expert, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& scatter_idx, + torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& sorted_experts, + torch::stable::Tensor& sorted_rows, + torch::stable::Tensor& topk_ids_for_sort) { + moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert, + n_local_expert, topk_ids.size(1), expert_first_token_offset, + scatter_idx, sort_workspace, sorted_experts, sorted_rows, + topk_ids_for_sort, true); +} + void moe_unpermute( const torch::stable::Tensor& permuted_hidden_states, // [n_token * topk, hidden] @@ -315,6 +358,19 @@ void moe_unpermute( STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0"); } +void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& token_expert_indices, + const std::optional& expert_map, + int64_t n_expert, int64_t n_local_expert, + torch::stable::Tensor& expert_first_token_offset, + torch::stable::Tensor& scatter_idx, + torch::stable::Tensor& sort_workspace, + torch::stable::Tensor& sorted_experts, + torch::stable::Tensor& sorted_rows, + torch::stable::Tensor& topk_ids_for_sort) { + STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0"); +} + #endif bool moe_permute_unpermute_supported() { @@ -328,5 +384,6 @@ bool moe_permute_unpermute_supported() { STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { m.impl("moe_permute", TORCH_BOX(&moe_permute)); m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch)); + m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter)); m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute)); } diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu index ec7851da3d..8d8bdf8baa 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu @@ -82,34 +82,40 @@ __device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices, return target_location + 1; } -// Calculates the start offset of the tokens for a given expert. The last -// element is the total number of valid tokens -__global__ void computeExpertFirstTokenOffsetKernel( +// Computes expert offsets (ending with the valid row count) and optional +// inverse indices. +__global__ void computeExpertOffsetsAndInverseKernel( int const* sorted_experts, int64_t const sorted_experts_len, - int const num_experts, int64_t* expert_first_token_offset) { - // First, compute the global tid. We only need 1 thread per expert. - int const expert = blockIdx.x * blockDim.x + threadIdx.x; + int const num_experts, int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse) { + int const index = blockIdx.x * blockDim.x + threadIdx.x; // Note that expert goes [0, num_experts] (inclusive) because we want a count // for the total number of active tokens at the end of the scan. - if (expert >= num_experts + 1) { - return; + if (index < num_experts + 1) { + expert_first_token_offset[index] = + findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, index); + } + if (inverse != nullptr && index < sorted_experts_len) { + inverse[sorted_rows[index]] = index; } - expert_first_token_offset[expert] = - findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert); } -void computeExpertFirstTokenOffset(int const* sorted_indices, - int const total_indices, - int const num_experts, - int64_t* expert_first_token_offset, - cudaStream_t stream) { - int const num_entries = num_experts + 1; - int const threads = std::min(1024, num_entries); +void computeExpertOffsetsAndInverse(int const* sorted_indices, + int const total_indices, + int const num_experts, + int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse, + cudaStream_t stream) { + int const num_entries = inverse != nullptr + ? std::max(num_experts + 1, total_indices) + : num_experts + 1; + int const threads = inverse != nullptr ? 256 : std::min(1024, num_entries); int const blocks = (num_entries + threads - 1) / threads; - computeExpertFirstTokenOffsetKernel<<>>( - sorted_indices, total_indices, num_experts, expert_first_token_offset); + computeExpertOffsetsAndInverseKernel<<>>( + sorted_indices, total_indices, num_experts, expert_first_token_offset, + sorted_rows, inverse); } void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, @@ -117,7 +123,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, int64_t* expert_first_token_offset, int num_rows, int num_experts, int num_experts_per_node, int k, CubKeyValueSorter& sorter, void* sorter_ws, - cudaStream_t stream) { + cudaStream_t stream, int* inverse) { int64_t const expanded_num_rows = static_cast(k) * num_rows; // We need to use the full num_experts because that is the sentinel value used // by topk for disabled experts @@ -127,9 +133,9 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row, permuted_experts, source_rows, permuted_rows, expanded_num_rows, stream); - computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows, - num_experts_per_node, expert_first_token_offset, - stream); + computeExpertOffsetsAndInverse( + permuted_experts, expanded_num_rows, num_experts_per_node, + expert_first_token_offset, permuted_rows, inverse, stream); } __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size, @@ -142,7 +148,7 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size, extern __shared__ int smem_expert_map[]; // store expert_map in smem for (int i = tidx; i < num_experts; i += blockDim.x) { - smem_expert_map[i] = expert_map_ptr[i]; + smem_expert_map[i] = expert_map_ptr ? expert_map_ptr[i] : i; } __syncthreads(); diff --git a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h index 89c278a4ed..53afad2f5d 100644 --- a/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h +++ b/csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h @@ -43,18 +43,19 @@ class CubKeyValueSorter { int num_bits_; }; -void computeExpertFirstTokenOffset(int const* sorted_indices, - int const total_indices, - int const num_experts, - int64_t* expert_first_token_offset, - cudaStream_t stream); +void computeExpertOffsetsAndInverse(int const* sorted_indices, + int const total_indices, + int const num_experts, + int64_t* expert_first_token_offset, + int const* sorted_rows, int* inverse, + cudaStream_t stream); void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows, int* permuted_experts, int* permuted_rows, int64_t* expert_first_token_offset, int num_rows, int num_experts, int num_experts_per_node, int k, CubKeyValueSorter& sorter, void* sorter_ws, - cudaStream_t stream); + cudaStream_t stream, int* inverse = nullptr); template void expandInputRowsKernelLauncher( diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index c0e966b2be..2feec38389 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -39,7 +39,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { " int block_size, Tensor! sorted_token_ids," " Tensor! experts_ids," " Tensor! num_tokens_post_pad," - " Tensor? maybe_expert_map) -> ()"); + " Tensor? maybe_expert_map," + " Tensor(a!)? scatter_idx=None) -> ()"); // Aligning the number of tokens to be processed by each expert such // that it is divisible by the block size, but for the batched case. @@ -105,6 +106,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, " "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()"); + m.def( + "moe_prepare_scatter(Tensor topk_ids, Tensor token_expert_indices, " + "Tensor? expert_map, int n_expert, int n_local_expert, " + "Tensor! expert_first_token_offset, Tensor! scatter_idx, " + "Tensor! sort_workspace, Tensor! sorted_experts, Tensor! sorted_rows, " + "Tensor! topk_ids_for_sort)->()"); + m.def( "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights," "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, " diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 00a56ac3bd..72609d8983 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -33,4 +33,4 @@ quack-kernels==0.6.5 # Required by MSA tokenspeed-mla==0.1.8; platform_system == "Linux" # Humming kernels for quantization gemm -humming-kernels[cu13]==0.1.12 +humming-kernels[cu13]==0.1.15 diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml index 310255f6bf..2ee157359a 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml @@ -9,3 +9,5 @@ server_args: >- --moe-backend humming --linear-backend humming --dtype bfloat16 +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml new file mode 100644 index 0000000000..a6bf4f6525 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml @@ -0,0 +1,14 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --data-parallel-size 2 + --enable-expert-parallel + --all2all-backend deepep_high_throughput +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml new file mode 100644 index 0000000000..cf8e01d70f --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml @@ -0,0 +1,14 @@ +model_name: "nm-testing/Qwen3-30B-A3B-FP8-block" +accuracy_threshold: 0.86 +num_questions: 1319 +num_fewshot: 5 +server_args: >- + --enforce-eager + --max-model-len 8192 + --moe-backend humming + --linear-backend humming + --data-parallel-size 2 + --enable-expert-parallel + --all2all-backend deepep_high_throughput +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "indexed" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml index 67725fce64..a8b289aaac 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml @@ -7,3 +7,5 @@ server_args: >- --max-model-len 8192 --quantization humming --linear-backend humming +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml index a20c143310..57923c63b0 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml @@ -7,3 +7,5 @@ server_args: >- --max-model-len 8192 --moe-backend humming --linear-backend humming +env: + VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous" diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml deleted file mode 100644 index b58c0d710e..0000000000 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml +++ /dev/null @@ -1,13 +0,0 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 -num_questions: 1319 -num_fewshot: 5 -server_args: >- - --enforce-eager - --max-model-len 8192 - --tensor-parallel-size 1 - --quantization humming - --kernel-config.enable_flashinfer_autotune=False -env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' - VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml similarity index 62% rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml index c932091dbe..d2f04348b9 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml @@ -1,7 +1,8 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 +gen_prefix: " \n\n\n" server_args: >- --enforce-eager --max-model-len 8192 @@ -9,5 +10,5 @@ server_args: >- --quantization humming --kernel-config.enable_flashinfer_autotune=False env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml similarity index 57% rename from tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml rename to tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml index fa5095cc88..08687aa6a3 100644 --- a/tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml @@ -1,7 +1,8 @@ -model_name: "Qwen/Qwen3-30B-A3B" -accuracy_threshold: 0.80 +model_name: "Qwen/Qwen3.5-35B-A3B" +accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 +gen_prefix: " \n\n\n" server_args: >- --enforce-eager --max-model-len 8192 @@ -9,4 +10,4 @@ server_args: >- --quantization humming --kernel-config.enable_flashinfer_autotune=False env: - VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}' + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml new file mode 100644 index 0000000000..067b96924d --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml @@ -0,0 +1,13 @@ +model_name: "Qwen/Qwen3.5-9B" +accuracy_threshold: 0.88 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}' diff --git a/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml new file mode 100644 index 0000000000..24915abbc2 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml @@ -0,0 +1,12 @@ +model_name: "Qwen/Qwen3.5-9B" +accuracy_threshold: 0.88 +num_questions: 1319 +num_fewshot: 5 +gen_prefix: " \n\n\n" +server_args: >- + --enforce-eager + --max-model-len 8192 + --quantization humming + --kernel-config.enable_flashinfer_autotune=False +env: + VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}' diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt new file mode 100644 index 0000000000..0537ed974a --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt @@ -0,0 +1,3 @@ +gpt-oss-20b-humming.yaml +Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt new file mode 100644 index 0000000000..512cbe99ff --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt @@ -0,0 +1,4 @@ +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-4B-mixed-quant-RTN-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt index 63b428ac0e..bdfa9033df 100644 --- a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt @@ -1,4 +1,5 @@ Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml gpt-oss-20b-humming.yaml Qwen3-30B-A3B-GPTQ-Int4-humming.yaml Qwen3-30B-A3B-FP8-block-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-act-int8.txt b/tests/evals/gsm8k/configs/humming/config-act-int8.txt index a7efbac3d0..68f5cd1c15 100644 --- a/tests/evals/gsm8k/configs/humming/config-act-int8.txt +++ b/tests/evals/gsm8k/configs/humming/config-act-int8.txt @@ -3,4 +3,5 @@ Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml Qwen3-30B-A3B-AWQ-humming.yaml Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml Qwen3-4B-mixed-quant-RTN-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-deepep.txt b/tests/evals/gsm8k/configs/humming/config-deepep.txt new file mode 100644 index 0000000000..2c97531d75 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-deepep.txt @@ -0,0 +1,2 @@ +Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml +Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt index 60179f1159..737a92cf48 100644 --- a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt +++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt @@ -1,4 +1,5 @@ Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml Qwen3-30B-A3B-Fp8-v1-humming.yaml Qwen3-30B-A3B-NVFP4-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt b/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt deleted file mode 100644 index 2c10777a09..0000000000 --- a/tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt +++ /dev/null @@ -1,3 +0,0 @@ -Qwen3-30B-A3B-int5wc-hadamard-humming.yaml -Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml -Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-l4.txt b/tests/evals/gsm8k/configs/humming/config-l4.txt new file mode 100644 index 0000000000..d1434b1cab --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-l4.txt @@ -0,0 +1,9 @@ +Qwen2-1.5B-Instruct-FP8W8-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml +Qwen3-0.6B-MXFP8-humming.yaml +Qwen3-0.6B-MXFP8-humming-act-fp8.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml +gpt-oss-20b-humming.yaml +Qwen3.5-9B-int6wc-hadamard-humming.yaml +Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index b1a52a5d55..c1ba96a792 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -352,49 +352,46 @@ def test_masked_situ_and_mul( ) -@pytest.mark.parametrize( - ("activation", "activation_config"), - [ - (MoEActivation.SILU, ApplyMoEActivationConfig()), - ( - MoEActivation.SILU, - ApplyMoEActivationConfig(clamp_limit=3.0), +MOE_ACTIVATION_CASES = [ + pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"), + pytest.param( + MoEActivation.SILU, ApplyMoEActivationConfig(clamp_limit=3.0), id="silu_clamp" + ), + pytest.param(MoEActivation.GELU, ApplyMoEActivationConfig(), id="gelu"), + pytest.param(MoEActivation.GELU_TANH, ApplyMoEActivationConfig(), id="gelu_tanh"), + pytest.param( + MoEActivation.SITU, + ApplyMoEActivationConfig( + activation_situ_beta=1.5, + activation_situ_linear_beta=2.0, ), - (MoEActivation.GELU, ApplyMoEActivationConfig()), - (MoEActivation.GELU_TANH, ApplyMoEActivationConfig()), - ( - MoEActivation.SITU, - ApplyMoEActivationConfig( - activation_situ_beta=1.5, - activation_situ_linear_beta=2.0, - ), - ), - (MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig()), - ( - MoEActivation.SWIGLUOAI_UNINTERLEAVE, - ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5), - ), - (MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig()), - (MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.GELU_TANH_NO_MUL, ApplyMoEActivationConfig()), - (MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig()), - ], - ids=[ - "silu", - "silu_clamp", - "gelu", - "gelu_tanh", - "situ", - "swigluoai", - "swigluoai_uninterleave", - "swiglustep", - "silu_no_mul", - "gelu_no_mul", - "gelu_tanh_no_mul", - "relu2_no_mul", - ], -) + id="situ", + ), + pytest.param(MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig(), id="swigluoai"), + pytest.param( + MoEActivation.SWIGLUOAI_UNINTERLEAVE, + ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5), + id="swigluoai_uninterleave", + ), + pytest.param(MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig(), id="swiglustep"), + pytest.param( + MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig(), id="silu_no_mul" + ), + pytest.param( + MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig(), id="gelu_no_mul" + ), + pytest.param( + MoEActivation.GELU_TANH_NO_MUL, + ApplyMoEActivationConfig(), + id="gelu_tanh_no_mul", + ), + pytest.param( + MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig(), id="relu2_no_mul" + ), +] + + +@pytest.mark.parametrize(("activation", "activation_config"), MOE_ACTIVATION_CASES) @torch.inference_mode() def test_masked_moe_activation_dispatch( default_vllm_config, @@ -497,3 +494,73 @@ def test_activation( out = torch.empty_like(x) opcheck(fn, (out, x)) + + +HUMMING_ACTIVATION_CASES = MOE_ACTIVATION_CASES + [ + pytest.param(MoEActivation.RELU2, ApplyMoEActivationConfig(), id="relu2"), +] +HUMMING_ACTIVATION_CASES += [ + pytest.param( + MoEActivation.SITU, + ApplyMoEActivationConfig( + activation_situ_beta=1.5, + activation_situ_linear_beta=linear_beta, + ), + id=f"situ-linear-beta-{linear_beta}", + ) + for linear_beta in (None, 0.0, -1.0) +] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Humming requires CUDA") +@pytest.mark.parametrize(("activation", "activation_config"), HUMMING_ACTIVATION_CASES) +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize(("num_tokens", "d"), [(1, 512), (7, 768), (83, 512)]) +@torch.inference_mode() +def test_humming_activation_matches_framework( + activation: MoEActivation, + activation_config: ApplyMoEActivationConfig, + dtype: torch.dtype, + num_tokens: int, + d: int, +) -> None: + """Compare activation math/layouts without quantization or Hadamard error.""" + pytest.importorskip("humming") + from humming.ops import process_input + + from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, + ) + + set_random_seed(0) + width = d * 2 if activation.is_gated else d + x = 4 * torch.randn(num_tokens, width, dtype=dtype, device=CUDA_DEVICES[0]) + # Exercise zero, saturation, and both sides of the clamp limits (3 and 7). + edges = x.new_tensor([0, 0.001, 1, 2.99, 3, 3.01, 6.99, 7, 7.01, 8, 16]) + edges = torch.cat((-edges[1:].flip(0), edges)) + x[0] = edges[torch.arange(width, device=x.device) % edges.numel()] + + actual, _, _ = process_input( + x, + quant_mode="none", + hadamard_block_size=0, + **get_humming_activation(activation, activation_config), + ) + expected = torch.empty(num_tokens, d, dtype=dtype, device=x.device) + if activation == MoEActivation.RELU2: + # apply_moe_activation has only the non-gated ReLU2 variant. + gate, up = x.float().chunk(2, dim=-1) + activated_gate = torch.empty_like(gate) + apply_moe_activation(MoEActivation.RELU2_NO_MUL, activated_gate, gate.clone()) + expected.copy_(activated_gate * up) + else: + # The framework's non-gated ReLU2 path modifies its input in place. + apply_moe_activation( + activation, expected, x.clone(), activation_config=activation_config + ) + torch.testing.assert_close( + actual, + expected, + atol=get_default_atol(expected), + rtol=get_default_rtol(expected), + ) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 1aca8593ed..8d6deac4f0 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -1245,7 +1245,7 @@ def _make_humming_indexed_experts(activation: MoEActivation): from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( HummingIndexedExperts, ) - from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.model_executor.layers.quantization.utils import humming as humming_utils from vllm.utils import humming top_k, num_experts = 6, 12 @@ -1306,6 +1306,56 @@ def _make_humming_indexed_experts(activation: MoEActivation): return experts, layer +@pytest.mark.parametrize("valid_shape_m", [8192, 8193]) +@pytest.mark.parametrize("w2_block_size", [64, 96, 128]) +def test_humming_indexed_routing_matches_each_projection_tile( + valid_shape_m: int, w2_block_size: int +): + """Independently tuned projections preserve routed rows and expert identity.""" + from types import SimpleNamespace + + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( + HummingIndexedExperts, + ) + + topk_ids = torch.tensor([[0, 1], [1, 2], [2, 0]], device="cuda") + expert_map = torch.tensor([1, -1, 0], dtype=torch.int32, device="cuda") + experts = SimpleNamespace( + global_num_experts=3, + estimate_local_valid_shape_m=lambda _: valid_shape_m, + w13_tuning_config=[ + [0, 8192, {"block_shape": [64, 128, 64]}], + [8192, 16384, {"block_shape": [128, 128, 64]}], + ], + w2_tuning_config=[ + [0, 8192, {"block_shape": [64, 128, 64]}], + [8192, 16384, {"block_shape": [w2_block_size, 128, 64]}], + ], + compute_config_str="compute", + w13_tuning_config_str="w13", + w2_tuning_config_str="w2", + ) + first, second, _ = HummingIndexedExperts.prepare_humming_moe_kwargs( + experts, topk_ids, expert_map, None + ) + block_sizes = (64, 64) if valid_shape_m == 8192 else (128, w2_block_size) + expected_rows = torch.where(expert_map[topk_ids.flatten()] >= 0)[0] + for kwargs, block_size in zip((first, second), block_sizes): + padded = kwargs["num_tokens_padded"].item() + assert padded % block_size == 0 + rows = kwargs["sorted_ids"][:padded].long() + valid = rows < topk_ids.numel() + torch.testing.assert_close(rows[valid].sort().values, expected_rows) + routed_experts = kwargs["expert_ids"][: padded // block_size] + torch.testing.assert_close( + routed_experts.repeat_interleave(block_size)[valid], + expert_map[topk_ids.flatten()[rows[valid]]], + ) + assert (first["top_k"], second["top_k"]) == (2, 1) + for key in ("sorted_ids", "expert_ids", "num_tokens_padded"): + assert (first[key] is second[key]) == (block_sizes[0] == block_sizes[1]) + + @pytest.mark.parametrize("activation", list(MoEActivation)) def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation): from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( @@ -1464,54 +1514,46 @@ def test_humming_delegates_to_instance_activation(): def test_humming_grouped_apply_forwards_valid_prefix( monkeypatch: pytest.MonkeyPatch, ): - from types import MethodType, SimpleNamespace + from types import SimpleNamespace from unittest.mock import Mock import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( - HummingExpertsBase, HummingGroupedExperts, ) - activation_func = Mock() buffers = { "gate_up_output": torch.empty(6, 4), - "activation_output": torch.empty(6, 2), + "quanted_down_input": torch.empty(6, 2), "down_output": torch.empty(6, 4), } experts = SimpleNamespace( - activation=activation_func, num_experts=2, estimate_local_valid_shape_m=lambda _: 6, prepare_buffers=lambda *_: buffers, - _get_permute_scratch=lambda _: None, - quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: ( - inputs, - input_scale, + _get_permute_scratch=lambda _, *, indices_only=False: object(), + process_input=Mock( + side_effect=lambda _, **kwargs: ( + kwargs["inputs"], + kwargs.get("input_scale"), + None, + ) ), humming_forward=Mock(), compute_config_str="", w13_tuning_config_str="", w2_tuning_config_str="", ) - experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts) hidden_states = torch.empty(3, 4) output = torch.empty_like(hidden_states) topk_ids = torch.zeros((3, 2), dtype=torch.int64) topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) - expected_counts = torch.tensor([4], dtype=torch.int32) expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64) monkeypatch.setattr( humming, - "moe_permute", - lambda **_: ( - torch.empty(6, 4), - None, - expert_offsets, - torch.arange(6), - None, - ), + "moe_prepare_scatter", + lambda *_: (expert_offsets, torch.arange(6).view(3, 2)), ) monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None) @@ -1534,12 +1576,11 @@ def test_humming_grouped_apply_forwards_valid_prefix( apply_router_weight_on_input=False, ) - activation_func.assert_called_once() - call_kwargs = activation_func.call_args.kwargs + call_kwargs = experts.process_input.call_args.kwargs assert call_kwargs["activation"] == MoEActivation.SITU - assert call_kwargs["input"].shape == (6, 4) - assert call_kwargs["output"].shape == (6, 2) - torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts) + assert call_kwargs["inputs"].shape == (6, 4) + assert call_kwargs["quanted_input"].shape == (6, 2) + torch.testing.assert_close(call_kwargs["num_valid_tokens"], expert_offsets[-1:]) def test_batched_marlin_activation_uses_expert_token_counts( @@ -1653,7 +1694,8 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): activation=moe_config.activation, ) assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size - assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size + assert "activation_output" not in buffer_metas + assert buffer_metas["quanted_down_input"]["shape"][-1] == intermediate_size assert experts.moe_problem_size( a1=torch.empty(1, hidden_size), w1=torch.empty(num_experts, 1), @@ -1661,6 +1703,13 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): topk_ids=torch.empty(1, top_k, dtype=torch.long), ) == (num_experts, 1, intermediate_size, hidden_size, top_k) + # Quantized batched input storage must cover every dispatcher's token slots. + experts.is_batched = lambda: True + experts.max_num_tokens = 17 + experts.num_dispatchers = 2 + batched_metas, _ = experts.get_buffer_metas(17, top_k, activation) + assert batched_metas["quanted_gate_up_input"]["shape"][0] == num_experts * 17 * 2 + def test_humming_indexed_writes_supplied_output_buffer(): from vllm.forward_context import set_forward_context diff --git a/tests/kernels/moe/test_moe_align_block_size.py b/tests/kernels/moe/test_moe_align_block_size.py index a1599f5ad3..ca339eeed8 100644 --- a/tests/kernels/moe/test_moe_align_block_size.py +++ b/tests/kernels/moe/test_moe_align_block_size.py @@ -22,6 +22,46 @@ BLOCK_SIZES = [32, 128] set_random_seed(0) +@pytest.mark.parametrize("num_experts,num_tokens", [(32, 8), (256, 8), (256, 1024)]) +@pytest.mark.parametrize("use_expert_map", [False, True]) +def test_moe_align_scatter_masks_invalid_routes( + num_experts, num_tokens, use_expert_map +): + topk_ids = torch.arange(num_tokens * 4, device="cuda").reshape(-1, 4) + topk_ids %= num_experts + topk_ids[0] = torch.tensor([-1, num_experts, 0, 1], device="cuda") + expert_map = None + valid = (topk_ids >= 0) & (topk_ids < num_experts) + if use_expert_map: + expert_map = torch.full((num_experts,), -1, dtype=torch.int32, device="cuda") + expert_map[::2] = torch.arange( + num_experts // 2, device="cuda", dtype=torch.int32 + ) + valid &= topk_ids % 2 == 0 + + sorted_ids, expert_ids, num_padded, scatter_idx = moe_align_block_size( + topk_ids, + 32, + num_experts, + expert_map, + ignore_invalid_experts=True, + return_scatter_idx=True, + ) + expected = torch.arange(topk_ids.numel(), device="cuda", dtype=torch.int32) + expected[~valid.flatten()] = -1 + torch.testing.assert_close(scatter_idx, expected[:, None]) + active_sorted = sorted_ids[: num_padded.item()] + active_sorted = active_sorted[active_sorted < topk_ids.numel()] + torch.testing.assert_close(active_sorted.sort().values, expected[expected >= 0]) + + baseline = moe_align_block_size( + topk_ids, 32, num_experts, expert_map, ignore_invalid_experts=True + ) + assert len(baseline) == 3 + torch.testing.assert_close(num_padded, baseline[2]) + torch.testing.assert_close(expert_ids, baseline[1]) + + def _group_tokens_by_expert( sorted_ids: torch.Tensor, expert_ids: torch.Tensor, diff --git a/tests/kernels/moe/test_moe_permute_unpermute.py b/tests/kernels/moe/test_moe_permute_unpermute.py index 52fc8963ec..76f4b57cf0 100644 --- a/tests/kernels/moe/test_moe_permute_unpermute.py +++ b/tests/kernels/moe/test_moe_permute_unpermute.py @@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, + moe_prepare_scatter, moe_unpermute, ) from vllm.platforms import current_platform @@ -458,7 +459,7 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: hidden_dtype=hidden_states.dtype, ) - permuted, _, expert_offsets, _, _ = moe_permute( + permuted, _, expert_offsets, inverse, _ = moe_permute( hidden_states=hidden_states, a1q_scale=None, topk_ids=topk_ids, @@ -473,3 +474,27 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None: torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), ) torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]]) + assert torch.all(inverse[[1, 3, 4]] >= expert_offsets[-1]) + expected = torch.zeros_like(hidden_states) + expected[[0, 2]] = hidden_states[[0, 2]] + output = torch.empty_like(hidden_states) + moe_unpermute( + output, permuted, torch.ones(5, 1, device="cuda"), inverse, expert_offsets + ) + torch.testing.assert_close(output, expected) + + expected_inverse = inverse.clone() + expected_offsets = expert_offsets.clone() + expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) + torch.testing.assert_close(indices.flatten(), expected_inverse) + torch.testing.assert_close(expert_offsets, expected_offsets) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch) + topk_ids.zero_() + graph.replay() + + expected_inverse = torch.arange(5, dtype=torch.int32, device="cuda") + expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda") + torch.testing.assert_close(indices.flatten(), expected_inverse) + torch.testing.assert_close(expert_offsets, expected_offsets) diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 4adfffa7c8..46bef106e6 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -14,7 +14,6 @@ from vllm.model_executor.kernels.linear import ( AiterMxfp4LinearKernel, EmulationMxfp4LinearKernel, FlashInferMxFp4LinearKernel, - HummingMxFp4LinearKernel, MarlinMxFp4LinearKernel, MxFp4LinearKernel, MxFp4LinearLayerConfig, @@ -44,7 +43,7 @@ _TRUE_W4A4_KERNELS = [ # Weight-only (A16) kernels: they never quantize activations. They still accept # MXFP4 activation keys as an intentional compatibility fallback. -_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel] +_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel] def test_can_implement_is_abstract(): diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py index dbec989df3..74c266e0bf 100644 --- a/tests/quantization/test_humming_mxfp4_block_fp8.py +++ b/tests/quantization/test_humming_mxfp4_block_fp8.py @@ -15,6 +15,7 @@ it in vLLM comes down to two things, both exercised here: """ import pytest +import torch from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( HummingExpertsBase, @@ -56,8 +57,8 @@ def test_mxfp4_weight_with_per_token_fp8_activation_still_supported(): ], ) def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key): - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - _humming_input_schema_to_quant_key, + from vllm.model_executor.layers.quantization.utils.humming import ( + input_schema_to_quant_key, ) from vllm.utils.humming import HummingInputSchema from vllm.utils.humming import dtypes as humming_dtypes @@ -66,18 +67,21 @@ def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key): a_dtype=humming_dtypes.float8e4m3, input_scale_group_size=group_size, ) - assert _humming_input_schema_to_quant_key(schema) == expected_key + assert input_schema_to_quant_key(schema, param_dtype=torch.bfloat16) == expected_key @pytest.mark.skipif(not has_humming(), reason="humming is not installed") def test_humming_bf16_input_schema_is_unquantized(): - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - _humming_input_schema_to_quant_key, + from vllm.model_executor.layers.quantization.utils.humming import ( + input_schema_to_quant_key, ) from vllm.utils.humming import HummingInputSchema # No a_dtype -> unquantized (bf16/fp16) inputs -> None. - assert _humming_input_schema_to_quant_key(HummingInputSchema()) is None + assert ( + input_schema_to_quant_key(HummingInputSchema(), param_dtype=torch.bfloat16) + is None + ) def test_block_fp8_activation_quant_config_is_block_quantized(): @@ -86,7 +90,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized(): DeepEP prepare/finalize step quantize activations to block FP8 *before* the all-to-all dispatch instead of deferring to Humming -- and what makes HummingExpertsBase.expects_unquantized_inputs return False.""" - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( make_humming_moe_quant_config, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape @@ -108,7 +112,7 @@ def test_default_activation_quant_config_defers_to_humming(): """Without a block activation shape the config stays per-token (the deferred path): Humming quantizes internally, preserving pre-existing behavior for every non-block-FP8 scheme.""" - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( make_humming_moe_quant_config, ) from vllm.platforms import current_platform diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 0fd5a10120..989bf967fa 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -229,7 +229,7 @@ def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): def test_moe_wna16_uses_humming_quant_config(monkeypatch): - from vllm.model_executor.layers.quantization.utils import humming_utils + from vllm.model_executor.layers.quantization.utils import humming as humming_utils method = object.__new__(MoeWNA16Method) method.wna16_backend = WNA16MoEBackend.HUMMING diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2a1654231f..1d38bcade0 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2272,6 +2272,7 @@ def moe_align_block_size( experts_ids: torch.Tensor, num_tokens_post_pad: torch.Tensor, expert_map: torch.Tensor | None = None, + scatter_idx: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.moe_align_block_size( topk_ids, @@ -2281,6 +2282,7 @@ def moe_align_block_size( experts_ids, num_tokens_post_pad, expert_map, + scatter_idx, ) diff --git a/vllm/envs.py b/vllm/envs.py index 26c00c6da6..259c08e7e7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1524,7 +1524,8 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None) ), - # The activation dtype config for humming kernel + # The activation dtype config for humming kernel. Explicit schemas disable + # fallback unless the config includes "allow_fallback": true. "VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file( os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None) ), diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index fb7a170b82..764bb73e88 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -101,6 +101,7 @@ from vllm.model_executor.kernels.linear.mxfp6 import ( from vllm.model_executor.kernels.linear.mxfp6.emulation import ( EmulationMxfp6LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp6.humming import HummingMxFp6LinearKernel from vllm.model_executor.kernels.linear.mxfp8 import ( Mxfp8LinearKernel, Mxfp8LinearLayerConfig, @@ -295,6 +296,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { HummingLinearKernel, HummingMxfp8LinearKernel, HummingMxFp4LinearKernel, + HummingMxFp6LinearKernel, HummingNvFp4LinearKernel, }, "marlin": { @@ -574,6 +576,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { _POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = { PlatformEnum.CUDA: [ + HummingMxFp6LinearKernel, EmulationMxfp6LinearKernel, ], PlatformEnum.ROCM: [ @@ -1294,6 +1297,7 @@ __all__ = [ "MxFp6LinearLayerConfig", "init_mxfp6_linear_kernel", "EmulationMxfp6LinearKernel", + "HummingMxFp6LinearKernel", "AiterMxfp4LinearKernel", "EmulationMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/humming.py b/vllm/model_executor/kernels/linear/mixed_precision/humming.py index 3de8b9ca41..b98dfee18b 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/humming.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/humming.py @@ -24,7 +24,7 @@ class HummingLinearKernel(MPLinearKernel): return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, @@ -57,13 +57,14 @@ class HummingLinearKernel(MPLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, ) return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index 660e31dcf4..0504fc5482 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -3,21 +3,18 @@ import torch -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig -logger = init_logger(__name__) - class HummingMxFp4LinearKernel(MxFp4LinearKernel): """Humming GEMM Kernel for MXFP4.""" @@ -39,14 +36,10 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel): @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: - if config.activation_quant_key not in (None, kMxfp4Dynamic): - return False, "only supports MXFP4 dynamic or unquantized activations" - if config.activation_quant_key is not None: - logger.warning_once( - "HummingMxFp4LinearKernel is a weight-only (A16) kernel; " - "the requested activation quantization (%s) is ignored.", - config.activation_quant_key, - ) + try: + quant_key_to_input_schema(config.activation_quant_key) + except ValueError as error: + return False, str(error) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: @@ -62,7 +55,10 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel): } convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_schema = quant_key_to_input_schema(self.config.activation_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -75,6 +71,7 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel): return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py index 3206936987..947ea60e3f 100644 --- a/vllm/model_executor/kernels/linear/mxfp6/base.py +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -13,8 +13,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey class MxFp6LinearLayerConfig: """Configuration for an MXFP6 linear layer. - All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per - byte) and per-block weight scales (group size 32). + All MXFP6 layers share the same structure: packed uint8 weights (4 FP6 values + per 3 bytes) and per-block weight scales (group size 32). Attributes: weight_quant_key: Identifies the weight quantization format. Can be diff --git a/vllm/model_executor/kernels/linear/mxfp6/humming.py b/vllm/model_executor/kernels/linear/mxfp6/humming.py new file mode 100644 index 0000000000..fff7584d93 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/humming.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.humming import ( + apply_humming_linear, + convert_linear_layer_to_humming_standard, + get_humming_linear_compute_config, + prepare_humming_linear_layer_config, + quant_key_to_input_schema, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp6E2M3Static +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming + +from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig + + +class HummingMxFp6LinearKernel(MxFp6LinearKernel): + """Humming GEMM for packed MXFP6 E2M3 and E3M2 weights.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_cuda(): + return False, "Humming only supported on CUDA" + + if not has_humming(): + return False, "Humming is not installed" + + if not current_platform.has_device_capability(75): + return False, "Humming only supported on SM75+" + + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + try: + quant_key_to_input_schema(config.activation_quant_key) + except ValueError as error: + return False, str(error) + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_dtype = "float6e3m2" + if self.config.weight_quant_key == kMxfp6E2M3Static: + weight_dtype = "float6e2m3" + + quant_config = { + "quant_method": "humming", + "dtype": weight_dtype, + "scale_dtype": "float8e8m0", + "group_size": 32, + "weight_scale_type": "group", + } + layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu) + + convert_linear_layer_to_humming_standard( + layer=layer, + name_map={"weight": "weight", "weight_scale": "weight_scale"}, + ) + input_schema = quant_key_to_input_schema(self.config.activation_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) + self.compute_config = get_humming_linear_compute_config() + self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return apply_humming_linear( + layer, + x, + skip_bias_add=bias is None, + layer_config=self.layer_config, + compute_config=self.compute_config, + locks=self.locks, + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/humming.py b/vllm/model_executor/kernels/linear/mxfp8/humming.py index bde1565d2a..63e9514f4e 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp8/humming.py @@ -3,12 +3,14 @@ import torch -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming @@ -50,7 +52,10 @@ class HummingMxfp8LinearKernel(Mxfp8LinearKernel): } convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_schema = quant_key_to_input_schema(kMxfp8Dynamic) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -63,6 +68,7 @@ class HummingMxfp8LinearKernel(Mxfp8LinearKernel): return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/nvfp4/humming.py b/vllm/model_executor/kernels/linear/nvfp4/humming.py index 54a0cf6c1d..3829c39b74 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/humming.py +++ b/vllm/model_executor/kernels/linear/nvfp4/humming.py @@ -4,11 +4,13 @@ import torch from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, get_humming_linear_compute_config, prepare_humming_linear_layer_config, + quant_key_to_input_schema, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic from vllm.platforms import current_platform from vllm.utils.import_utils import has_humming @@ -40,9 +42,7 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel): return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Route through humming's compressed-tensors nvfp4 loader (same path as - # the MoE oracle); the native group_tensor schema mishandles a scalar - # global scale. + # Reuse the CT loader for packed FP4 weights and tensor scales. quant_config = { "quant_method": "compressed-tensors", "format": "nvfp4-pack-quantized", @@ -60,7 +60,16 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel): layer.weight_global_scale = torch.nn.Parameter( 1.0 / layer.weight_global_scale, requires_grad=False ) - self.layer_config = prepare_humming_linear_layer_config(layer, quant_config) + input_quant_key = None + if hasattr(layer, "input_global_scale"): + layer.input_scale_2 = torch.nn.Parameter( + layer.input_global_scale.reshape(-1), requires_grad=False + ) + input_quant_key = kNvfp4Dynamic + input_schema = quant_key_to_input_schema(input_quant_key) + self.layer_config = prepare_humming_linear_layer_config( + layer, quant_config, input_schema=input_schema + ) self.compute_config = get_humming_linear_compute_config() self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device) @@ -73,6 +82,7 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel): return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/kernels/linear/scaled_mm/humming.py b/vllm/model_executor/kernels/linear/scaled_mm/humming.py index b52e73b5a4..3381a3381f 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/humming.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/humming.py @@ -4,7 +4,7 @@ import torch from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( apply_humming_linear, convert_linear_layer_to_humming_standard, get_humming_linear_compute_config, @@ -92,6 +92,7 @@ class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, @@ -158,6 +159,7 @@ class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): return apply_humming_linear( layer, x, + skip_bias_add=bias is None, layer_config=self.layer_config, compute_config=self.compute_config, locks=self.locks, diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py index 18a2fb5cc6..a6d44c3141 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py @@ -4,6 +4,7 @@ import json import math +from dataclasses import replace from typing import TYPE_CHECKING, Any, cast import torch @@ -30,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( get_moe_permute_scratch, moe_permute, moe_permute_unpermute_supported, + moe_prepare_scatter, moe_unpermute, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( @@ -57,14 +59,16 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp8Static, kNvfp4Dynamic, kNvfp4Static, + kStaticTensorScale, ) from vllm.platforms import current_platform from vllm.scalar_type import ScalarType from vllm.utils.import_utils import has_humming +from vllm.utils.platform_utils import num_compute_units from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( HummingMoEQuantConfig, ) from vllm.utils.humming import ( @@ -118,38 +122,9 @@ def get_humming_moe_gemm_type(moe_config: FusedMoEConfig | None = None) -> str: return gemm_type -def _fixup_moe_tuning_config(tuning_config: list, max_k_block: int = 128) -> None: - """Fix up each MoE tile in place: cap the K-block and widen warp-N. - - - block_shape[2] (K-block) > ``max_k_block``: the driver rejects the TMA - descriptor at launch (CUDA_ERROR_MISALIGNED_ADDRESS). Cap at 128, which - Humming already uses for larger M. - - warp_shape[1] (warp-N) < 32: block-FP8 (group-128) activations route the - w13 (gate/up) GEMM to a tuning table that pins warp-N to 16, under-filling - the Hopper WGMMA N dimension and corrupting the GEMM output (gsm8k - 0.94 -> 0.89). Widen to 32 whenever block_n % 32 == 0 -- on every tile, - not just the K-capped ones. w2 (down) already uses warp_n=32. - """ - logger.info_once("Attempting to override humming GEMM config") - for entry in tuning_config: - config = entry[2] - block_shape = config.get("block_shape") - if not (block_shape and len(block_shape) == 3): - continue - block_m, block_n, block_k = block_shape - - logger.info_once(f"Overriding humming GEMM config. Previous config\n: {config}") - if block_k > max_k_block: - config["block_shape"] = [block_m, block_n, max_k_block] - - warp_shape = config.get("warp_shape") - if warp_shape and warp_shape[1] < 32 and block_n % 32 == 0: - config["warp_shape"] = [warp_shape[0], 32, warp_shape[2]] - - logger.info_once(f"Overridden humming GEMM config. Current config\n: {config}") - - class HummingExpertsBase(mk.FusedMoEExpertsModular): + quant_config: "HummingMoEQuantConfig" + def __init__( self, moe_config: FusedMoEConfig, @@ -165,7 +140,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): self.locks = torch.zeros(1024, dtype=torch.int32, device=moe_config.device) self.num_experts = moe_config.num_local_experts self.global_num_experts = moe_config.num_experts - self.quant_config = quant_config + self.quant_config = humming_quant_config self.init_humming_moe() if self.is_batched(): @@ -199,31 +174,61 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): gemm_type=self.humming_gemm_type(), ) self.compute_config_str = json.dumps(self.compute_config) - # Fix up the heuristic-chosen tiles before freezing to JSON -- see - # _fixup_moe_tuning_config. - _fixup_moe_tuning_config(self.w13_tuning_config) - _fixup_moe_tuning_config(self.w2_tuning_config) self.w13_tuning_config_str = json.dumps(self.w13_tuning_config) self.w2_tuning_config_str = json.dumps(self.w2_tuning_config) - def quantize_input( + def process_input( self, sublayer_name: str, inputs: torch.Tensor, quanted_input: torch.Tensor | None, input_scale: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - from vllm.utils.humming import may_quant_input - - # input_scale is set for block-FP8 (group-128) activations that were - # quantized before the EP dispatch: may_quant_input then skips the - # redundant w13 quantization and forwards the pre-computed scale. - return may_quant_input( - self.humming_configs[sublayer_name], - inputs=inputs, - input_scale=input_scale, - quanted_input=quanted_input, + expert_tokens: torch.Tensor | None = None, + activation: MoEActivation | None = None, + scatter_idx: torch.Tensor | None = None, + num_valid_tokens: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, ) + from vllm.utils.humming import may_process_input + + quant_config = self.quant_config + config = self.humming_configs[sublayer_name] + if input_scale is not None: + assert scatter_idx is None + if activation is not None: + raise ValueError("Cannot apply activation to prequantized input") + return inputs, input_scale, None + + activation_kwargs = {} + if activation is not None: + activation_config = self.activation_config + activation_kwargs = get_humming_activation(activation, activation_config) + + prefix = "w1" if sublayer_name == "w13" else "w2" + scale = getattr(quant_config, f"{prefix}_input_scale") + scale2 = getattr(quant_config, f"{prefix}_input_scale_2") + mode = config.input_quant_mode + layout = "grouped_mask" if self.is_batched() else "normal" + if scatter_idx is not None: + layout = "scatter" + + inputs, group_scales, token_scales = may_process_input( + config, + inputs=inputs, + outputs=quanted_input, + token_scales=scale2 if mode.has_secondary_scale else scale, + hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"), + layout=layout, + expert_tokens=expert_tokens if layout == "grouped_mask" else None, + scatter_idx=scatter_idx, + num_valid_tokens=num_valid_tokens, + **activation_kwargs, + ) + input_scale = group_scales if mode.has_group_scale else token_scales + input_scale_2 = token_scales if mode.has_secondary_scale else None + return inputs, input_scale, input_scale_2 def humming_forward( self, @@ -232,30 +237,30 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): weight: torch.Tensor, input_scale: torch.Tensor | None, outputs: torch.Tensor, + input_scale_2: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: from vllm.utils.humming import humming_forward - is_w13 = sublayer_name == "w13" + index = 1 if sublayer_name == "w13" else 2 return humming_forward( self.humming_configs[sublayer_name], inputs=inputs, weight=weight, - weight_scale=( - self.quant_config.w1_scale if is_w13 else self.quant_config.w2_scale - ), - zero_point=(self.quant_config.w1_zp if is_w13 else self.quant_config.w2_zp), - bias=(self.quant_config.w1_bias if is_w13 else self.quant_config.w2_bias), - weight_scale_2=( - self.quant_config.g1_alphas if is_w13 else self.quant_config.g2_alphas - ), + weight_scale=getattr(self.quant_config, f"w{index}_scale"), + zero_point=getattr(self.quant_config, f"w{index}_zp"), + bias=getattr(self.quant_config, f"w{index}_bias"), + weight_scale_2=getattr(self.quant_config, f"g{index}_alphas"), input_scale=input_scale, + input_scale_2=input_scale_2, outputs=outputs, locks=self.locks, **kwargs, ) - def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None: + def _get_permute_scratch( + self, topk: int, indices_only: bool = False + ) -> MoEPermuteScratch | None: if not moe_permute_unpermute_supported(): return None @@ -270,8 +275,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): num_experts=self.moe_config.num_experts, num_local_experts=self.moe_config.num_local_experts, device=torch.device(self.moe_config.device), - hidden_size=self.moe_config.hidden_dim, - hidden_dtype=self.moe_config.in_dtype, + hidden_size=None if indices_only else self.moe_config.hidden_dim, + hidden_dtype=None if indices_only else self.moe_config.in_dtype, ) def get_global_valid_shape_m(self, topk_ids: torch.Tensor): @@ -302,6 +307,16 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: + if ( + activation_key is not None + and activation_key.scale == kNvfp4Dynamic.scale + and activation_key.dtype == kNvfp4Dynamic.dtype + and activation_key.scale2 is not None + and not activation_key.scale2.static + and activation_key.scale2.group_shape == GroupShape.PER_TOKEN + and activation_key.scale2.dtype == torch.float32 + ): + activation_key = replace(activation_key, scale2=kStaticTensorScale) SUPPORTED_W_A = [ (kMxfp4Static, None), (kMxfp4Static, kMxfp4Dynamic), @@ -343,7 +358,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A or ( - activation_key in (None, kFp8DynamicTokenSym) + activation_key in (None, kFp8DynamicTokenSym, kInt8DynamicTokenSym) and _is_supported_wna16_weight_key(weight_key) ) @@ -355,7 +370,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): to FP8 before dispatch sends FP8 rather than BF16 over the interconnect, and Humming then consumes the pre-quantized FP8 + scale directly (see the apply() methods, which forward the dispatch scale into - HummingMethod.may_quant_input, and may_quant_input itself, which is a + HummingExpertsBase.process_input, which is a no-op when an input scale is already supplied). The scale layout produced by vLLM's block-FP8 quantization ([M, K // 128] float32, row-major) matches what the Humming WGMMA grouped GEMM expects. @@ -373,7 +388,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): inputs). Humming normally quantizes inputs internally via - HummingMethod.may_quant_input() in apply(), so we defer quantization + HummingExpertsBase.process_input() in apply(), so we defer quantization (return True) to avoid quantizing twice -- once in prepare and once in Humming's apply(). @@ -398,8 +413,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - # Humming uses apply_moe_activation() callback for activation, - # so any activation supported there can be used here. return apply_moe_activation_supported(activation) @staticmethod @@ -456,10 +469,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): assert intermediate_dim == self.adjust_N_for_activation(gate_up_dim, activation) # hidden_states - # (-> quanted_gate_up_input) (if not BF16/FP16 activation) + # -> quanted_gate_up_input (Hadamard + optional quantization) # -> gate_up_output - # -> activation_output - # (-> quanted_down_input) (if not BF16/FP16 activation) + # -> quanted_down_input (activation + Hadamard + optional quantization) # -> down_output # (-> output) (if not is_batched) # Neighboring nodes are required to utilize distinct workspaces. @@ -471,8 +483,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): max_num_tokens = self.max_num_tokens num_dispatchers = self.num_dispatchers assert max_num_tokens is not None and num_dispatchers is not None - input_shape_m = num_experts * max_num_tokens real_shape_m = num_experts * max_num_tokens * num_dispatchers + input_shape_m = real_shape_m output_shape = (num_experts, max_num_tokens * num_dispatchers, K) else: input_shape_m = M @@ -483,15 +495,18 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): a_dtype = self.humming_configs["w13"].a_dtype c_dtype = self.humming_configs["w13"].c_dtype - num_bits = a_dtype.num_bits + down_a_dtype = w2_config.a_dtype torch_dtype_map = { dtypes.float16: torch.float16, dtypes.bfloat16: torch.bfloat16, dtypes.float32: torch.float32, + dtypes.float8e3m4: torch.uint8, dtypes.float8e4m3: torch.float8_e4m3fn, dtypes.float8e5m2: torch.float8_e5m2, dtypes.int8: torch.int8, dtypes.int4: torch.uint8, + dtypes.float4e0m3: torch.uint8, + dtypes.float4e2m1: torch.uint8, } buffer_metas = { @@ -503,13 +518,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): "shape": (real_shape_m, gate_up_dim), "dtype": torch_dtype_map[c_dtype], }, - "activation_output": { - "shape": (real_shape_m, intermediate_dim), - "dtype": torch_dtype_map[c_dtype], - }, "quanted_down_input": { "shape": (real_shape_m, intermediate_dim), - "dtype": torch_dtype_map[a_dtype], + "dtype": torch_dtype_map[down_a_dtype], }, "down_output": { "shape": output_shape if self.is_batched() else (real_shape_m, K), @@ -523,7 +534,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): for key in buffer_metas: meta = buffer_metas[key] - if "quanted" in key and a_dtype.num_bits == 4: + input_dtype = down_a_dtype if key == "quanted_down_input" else a_dtype + if "quanted" in key and input_dtype.num_bits == 4: last_dim = meta["shape"][-1] if last_dim % 2 != 0: raise ValueError( @@ -532,25 +544,12 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): ) meta["shape"] = meta["shape"][:-1] + (last_dim // 2,) - if num_bits == 16: - required_buffers = ["gate_up_output", "activation_output", "down_output"] - else: - required_buffers = [ - "quanted_gate_up_input", - "gate_up_output", - "activation_output", - "quanted_down_input", - "down_output", - ] - - # Note: The fused SITU+FP8 quant goes gate_up_output -> - # quanted_down_input directly, never materializing - # activation_output. Dropping activation_output from the chain - # flips the even/odd workspace 2-coloring below so - # gate_up_output and quanted_down_input land - # on DIFFERENT workspaces. - if self.fused_situ_quant_enabled(activation): - required_buffers.remove("activation_output") + required_buffers = [ + "quanted_gate_up_input", + "gate_up_output", + "quanted_down_input", + "down_output", + ] # batched moe use down_output as output if not self.is_batched(): @@ -673,81 +672,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular): activation_kwargs["valid_token_counts"] = valid_token_counts self.activation(**activation_kwargs) - def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool: - """Whether the SITU activation + w2 quant can be fused into one kernel. - - Fused only for k-major block-FP8 group-128 e4m3 with float32 scales -- - the sole layout situ_and_mul_quant supports. A float32 as_dtype rules out - MXMMA (which uses e8m0 m-major scales), so a group-128 scale is - guaranteed k-major. Everything else (m-major MXFP8, other group sizes, - 16-bit passthrough, non-SITU) falls back to the separate - situ_and_mul + quantize_input path. - """ - if activation != MoEActivation.SITU: - return False - w2cfg = self.humming_configs["w2"] - # Report the specific blocker once so a silent fallback to the unfused - # situ_and_mul + quant pair is never a mystery in a trace. - reason: str | None = None - if not (w2cfg.a_dtype.num_bits == 8 and str(w2cfg.a_dtype) == "float8e4m3"): - reason = f"w2 a_dtype is {w2cfg.a_dtype} (need float8e4m3)" - elif w2cfg.input_scale_group_size != 128: - reason = ( - f"w2 input_scale_group_size is {w2cfg.input_scale_group_size} " - "(need 128)" - ) - elif not (w2cfg.as_dtype is not None and str(w2cfg.as_dtype) == "float32"): - reason = f"w2 as_dtype is {w2cfg.as_dtype} (need float32, k-major)" - if reason is not None: - logger.warning_once( - "Humming fused SITU+FP8 quant disabled, using unfused " - "situ_and_mul + quant: %s", - reason, - ) - return False - return True - - def fused_situ_quant( - self, - gate_up_output: torch.Tensor, - quanted_down_input: torch.Tensor, - num_valid_tokens: torch.Tensor | None, - topk: int = 1, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the fused SITU activation + FP8 quant for the w2 input. - - Returns ``(quanted_down_input, input_scale)`` in the same layout the - unfused ``apply_activation`` + ``quantize_input("w2")`` pair produces: an - fp8 [M, d] tensor plus a k-major block-FP8 float32 [M, d // 128] group - scale. Only reached when ``fused_situ_quant_enabled`` (group_size 128). - """ - from vllm.model_executor.layers.fused_moe.activation import ( - situ_and_mul_quant, - ) - - cfg = self.activation_config - assert cfg.activation_situ_beta is not None, ( - "SITU requires activation_situ_beta from FusedMoEConfig" - ) - group_size = self.humming_configs["w2"].input_scale_group_size - m, d = quanted_down_input.size(0), quanted_down_input.size(1) - input_scale = torch.empty( - (m, d // group_size), - dtype=torch.float32, - device=quanted_down_input.device, - ) - situ_and_mul_quant( - quanted_down_input, - input_scale, - gate_up_output, - beta=cfg.activation_situ_beta, - linear_beta=cfg.activation_situ_linear_beta, - group_size=group_size, - num_valid_tokens=num_valid_tokens, - topk=topk, - ) - return quanted_down_input, input_scale - class HummingIndexedExperts(HummingExpertsBase): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: @@ -768,7 +692,7 @@ class HummingIndexedExperts(HummingExpertsBase): topk_ids: torch.Tensor, expert_map: torch.Tensor | None, expert_tokens_meta: mk.ExpertTokensMetadata | None, - ) -> tuple[dict[str, Any], dict[str, Any]]: + ) -> tuple[dict[str, Any], dict[str, Any], torch.Tensor | None]: valid_shape_m = self.estimate_local_valid_shape_m(topk_ids) moe_block_size = None @@ -784,13 +708,21 @@ class HummingIndexedExperts(HummingExpertsBase): ) moe_block_size = 64 - sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size( + use_scatter = False + if expert_map is not None: + num_sms = num_compute_units(topk_ids.get_device()) + use_scatter = topk_ids.numel() > 2 * num_sms + + alignment = moe_align_block_size( topk_ids=topk_ids, block_size=moe_block_size, num_experts=self.global_num_experts, expert_map=expert_map, ignore_invalid_experts=True, + return_scatter_idx=use_scatter, ) + sorted_ids, expert_ids, num_tokens_padded = alignment[:3] + scatter_idx = alignment[3] if len(alignment) == 4 else None moe_common_kwargs = { "sorted_ids": sorted_ids, @@ -806,7 +738,26 @@ class HummingIndexedExperts(HummingExpertsBase): moe_kwargs1.update(moe_common_kwargs) moe_kwargs2.update(moe_common_kwargs) - return moe_kwargs1, moe_kwargs2 + w2_block_size = 64 + for lower, upper, config in self.w2_tuning_config: + if lower < valid_shape_m <= upper: + w2_block_size = config["block_shape"][0] + break + if w2_block_size != moe_block_size: + sorted_ids2, expert_ids2, num_tokens_padded2 = moe_align_block_size( + topk_ids=topk_ids, + block_size=w2_block_size, + num_experts=self.global_num_experts, + expert_map=expert_map, + ignore_invalid_experts=True, + ) + moe_kwargs2.update( + sorted_ids=sorted_ids2, + expert_ids=expert_ids2, + num_tokens_padded=num_tokens_padded2, + ) + + return moe_kwargs1, moe_kwargs2, scatter_idx def apply( self, @@ -832,7 +783,7 @@ class HummingIndexedExperts(HummingExpertsBase): object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to - may_quant_input so Humming skips the redundant w13 quantization. The + process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management. """ assert not apply_router_weight_on_input @@ -847,13 +798,13 @@ class HummingIndexedExperts(HummingExpertsBase): ) buffers["output"] = output - moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs( + moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs( topk_ids=topk_ids, expert_map=expert_map, expert_tokens_meta=expert_tokens_meta, ) - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, @@ -865,52 +816,25 @@ class HummingIndexedExperts(HummingExpertsBase): inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], **moe_kwargs1, ) - # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view. - # The fused path and mul_sum consume tokens directly. The unfused - # activation consumes the expanded row count (tokens * topk). - valid_tokens = None - topk = topk_ids.size(1) - if ( - expert_tokens_meta is not None - and expert_tokens_meta.psum_recv_per_rank is not None - ): - valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:] - - if self.fused_situ_quant_enabled(activation): - # Fused SITU + FP8 quant (per-token or block-FP8 group-128) straight - # into the w2 input, skipping the bf16 activation_output round-trip. - inputs, input_scale = self.fused_situ_quant( - gate_up_output=buffers["gate_up_output"], - quanted_down_input=buffers["quanted_down_input"], - num_valid_tokens=valid_tokens, - topk=topk, - ) - else: - valid_token_counts = ( - valid_tokens * topk if valid_tokens is not None else None - ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - valid_token_counts=valid_token_counts, - ) - - inputs, input_scale = self.quantize_input( - "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), - ) + inputs, input_scale, input_scale_2 = self.process_input( + "w2", + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + activation=activation, + scatter_idx=scatter_idx, + ) self.humming_forward( "w2", inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["down_output"].view(-1, hidden_states.size(-1)), **moe_kwargs2, ) @@ -918,6 +842,10 @@ class HummingIndexedExperts(HummingExpertsBase): # expert_map masks any non-local id; num_valid_tokens bounds the # persistent kernel to the real token rows [0, num_recv) so the padding # tail is never iterated (CUDA-graph-safe device scalar). + valid_tokens = None + if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None: + valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:] + moe_fused_mul_sum( inputs=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, @@ -983,21 +911,33 @@ class HummingGroupedExperts(HummingExpertsBase): ) buffers["output"] = output - hidden_states, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute( - hidden_states=hidden_states, - a1q_scale=a1q_scale, - topk_ids=topk_ids, - n_expert=global_num_experts, - n_local_expert=self.num_experts, - expert_map=expert_map, - scratch=self._get_permute_scratch(topk_ids.size(1)), - ) + if a1q_scale is None: + scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True) + assert scratch is not None + scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch) + expert_offsets, scatter_idx = scatter_metadata + inv_perm = scatter_idx.flatten() + num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None + else: + hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute( + hidden_states=hidden_states, + a1q_scale=a1q_scale, + topk_ids=topk_ids, + n_expert=global_num_experts, + n_local_expert=self.num_experts, + expert_map=expert_map, + scratch=self._get_permute_scratch(topk_ids.size(1)), + ) + scatter_idx = None + num_valid_tokens = None - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, quanted_input=buffers.get("quanted_gate_up_input", None), + scatter_idx=scatter_idx, + num_valid_tokens=num_valid_tokens, ) self.humming_forward( @@ -1005,28 +945,20 @@ class HummingGroupedExperts(HummingExpertsBase): inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], valid_shape_m=valid_shape_m, - expert_layout=expert_first_token_offset, + expert_layout=expert_offsets, compute_config=self.compute_config_str, tuning_config=self.w13_tuning_config_str, ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - valid_token_counts=( - expert_first_token_offset[-1:].to(torch.int32) - if expert_tokens_meta is None - else None - ), - ) - - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + activation=activation, + num_valid_tokens=expert_offsets[-1:], ) self.humming_forward( @@ -1034,9 +966,10 @@ class HummingGroupedExperts(HummingExpertsBase): inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["down_output"], valid_shape_m=valid_shape_m, - expert_layout=expert_first_token_offset, + expert_layout=expert_offsets, compute_config=self.compute_config_str, tuning_config=self.w2_tuning_config_str, ) @@ -1046,7 +979,7 @@ class HummingGroupedExperts(HummingExpertsBase): permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1), topk_weights=topk_weights, inv_permuted_idx=inv_perm, - expert_first_token_offset=expert_first_token_offset, + expert_first_token_offset=expert_offsets, ) @@ -1088,7 +1021,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual path (Humming quantizes the w13 input itself); for block-FP8 activations it carries the scale computed before dispatch, which is forwarded to - may_quant_input so Humming skips the redundant w13 quantization. The + process_input so Humming skips the redundant w13 quantization. The output is written into workspace13 via the buffer management. """ assert not apply_router_weight_on_input @@ -1111,11 +1044,12 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): ) buffers["down_output"] = output - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w13", inputs=hidden_states, input_scale=a1q_scale, quanted_input=buffers.get("quanted_gate_up_input", None), + expert_tokens=expert_num_tokens, ) self.humming_forward( @@ -1123,6 +1057,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): inputs=inputs, weight=w1, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=buffers["gate_up_output"], valid_shape_m=valid_shape_m, expert_layout=expert_num_tokens, @@ -1130,16 +1065,12 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): tuning_config=self.w13_tuning_config_str, ) - self.apply_activation( - activation=activation, - input=buffers["gate_up_output"], - output=buffers["activation_output"], - ) - - inputs, input_scale = self.quantize_input( + inputs, input_scale, input_scale_2 = self.process_input( "w2", - inputs=buffers["activation_output"], - quanted_input=buffers.get("quanted_down_input", None), + inputs=buffers["gate_up_output"], + quanted_input=buffers["quanted_down_input"], + expert_tokens=expert_num_tokens, + activation=activation, ) self.humming_forward( @@ -1147,6 +1078,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase): inputs=inputs, weight=w2, input_scale=input_scale, + input_scale_2=input_scale_2, outputs=output.view(-1, hidden_states.size(-1)), valid_shape_m=valid_shape_m, expert_layout=expert_num_tokens, diff --git a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py index c8587d3b71..663475306c 100644 --- a/vllm/model_executor/layers/fused_moe/moe_align_block_size.py +++ b/vllm/model_executor/layers/fused_moe/moe_align_block_size.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Literal, overload + import torch from vllm import _custom_ops as ops @@ -8,6 +10,48 @@ from vllm.triton_utils import triton from vllm.utils.math_utils import round_up +@overload +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + *, + return_scatter_idx: Literal[False] = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ... + + +@overload +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + *, + return_scatter_idx: Literal[True], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ... + + +@overload +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + expert_map: torch.Tensor | None = None, + pad_sorted_ids: bool = False, + ignore_invalid_experts: bool = False, + *, + return_scatter_idx: bool, +) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +): ... + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, @@ -15,7 +59,11 @@ def moe_align_block_size( expert_map: torch.Tensor | None = None, pad_sorted_ids: bool = False, ignore_invalid_experts: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return_scatter_idx: bool = False, +) -> ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +): """Aligns the token distribution across experts to be compatible with block size for matrix multiplication. @@ -42,6 +90,8 @@ def moe_align_block_size( as -1. When True, all invalid expert_ids in topk_ids will be ignored and will not participate in counting or ranking, and there will be no -1 in expert_ids. + return_scatter_idx: Whether to additionally return a masked identity + mapping for the original routed rows. Returns: - sorted_token_ids: A tensor containing the sorted token indices according @@ -49,6 +99,10 @@ def moe_align_block_size( - expert_ids: A tensor indicating the assigned expert index for each block. - num_tokens_post_padded: The total number of tokens after padding, ensuring divisibility by block_size. + - scatter_idx: Only returned when return_scatter_idx=True. An int32 tensor + shaped [topk_ids.numel(), 1], containing the original routed row index + for valid routes and -1 otherwise. With expert_map, this requires + ignore_invalid_experts=True. This function pads the number of tokens that each expert needs to process so that it is divisible by block_size. @@ -71,6 +125,15 @@ def moe_align_block_size( by block_size for proper block matrix operations. """ + if return_scatter_idx and expert_map is not None and not ignore_invalid_experts: + raise ValueError("scatter_idx with expert_map requires ignore_invalid_experts") + scatter_idx = None + if return_scatter_idx: + scatter_idx = torch.empty( + (topk_ids.numel(), 1), + dtype=torch.int32, + device=topk_ids.device, + ) max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) if pad_sorted_ids: max_num_tokens_padded = round_up(max_num_tokens_padded, block_size) @@ -95,11 +158,14 @@ def moe_align_block_size( expert_ids, num_tokens_post_pad, expert_map if ignore_invalid_experts else None, + scatter_idx, ) if expert_map is not None and not ignore_invalid_experts: expert_ids = expert_map[expert_ids] + if scatter_idx is not None: + return sorted_ids, expert_ids, num_tokens_post_pad, scatter_idx return sorted_ids, expert_ids, num_tokens_post_pad diff --git a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py index 0aa9a38071..d83152fc69 100644 --- a/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py +++ b/vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py @@ -107,6 +107,37 @@ class MoEPermuteScratch: return topk_ids_int32 +def moe_prepare_scatter( + topk_ids: torch.Tensor, + expert_map: torch.Tensor | None, + scratch: MoEPermuteScratch, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate expert offsets and shared scatter/unpermute destination indices.""" + assert topk_ids.device == scratch.device + assert topk_ids.size(0) <= scratch.max_num_tokens + assert topk_ids.size(1) == scratch.topk + n_token, topk = topk_ids.shape + expanded_rows = topk_ids.numel() + inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk) + if expanded_rows == 0: + scratch.expert_first_token_offset.zero_() + return scratch.expert_first_token_offset, inverse + torch.ops._moe_C.moe_prepare_scatter( + scratch.prepare_topk_ids(topk_ids), + scratch.token_expert_indices_view(n_token), + expert_map, + scratch.num_experts, + scratch.num_local_experts, + scratch.expert_first_token_offset, + inverse, + scratch.sort_workspace, + scratch.permuted_experts_id[:expanded_rows].view(n_token, topk), + scratch.sorted_row_idx[:expanded_rows].view(n_token, topk), + scratch.topk_ids_for_sort[:expanded_rows].view(n_token, topk), + ) + return scratch.expert_first_token_offset, inverse + + def get_moe_permute_scratch( *, max_num_tokens: int, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index abf0896204..5219fff617 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -603,7 +603,7 @@ def convert_to_fp8_moe_kernel_format( w13.is_shuffled = True w2.is_shuffled = True elif fp8_backend == Fp8MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -739,7 +739,7 @@ def make_fp8_moe_quant_config( ) elif fp8_backend == Fp8MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 1a23b01736..29e9888098 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -193,7 +193,7 @@ def make_int8_moe_quant_config( if int8_backend == Int8MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) @@ -254,7 +254,7 @@ def convert_to_int8_moe_kernel_format( ) -> tuple[torch.Tensor, torch.Tensor]: """Convert INT8 MoE weights to backend-specific kernel format.""" if int8_backend == Int8MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index c70bea9aa3..ad591d9ea7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1456,7 +1456,7 @@ def convert_to_wna16_moe_kernel_format( """ if backend == WNA16MoEBackend.HUMMING: from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 7fed4124f5..5450f73f85 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -844,7 +844,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( sf_block_size = 32 # mxfp4 block size if mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -1428,7 +1428,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( ) if mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -1966,7 +1966,7 @@ def make_mxfp4_moe_quant_config( gemm1_clamp_limit=swiglu_limit, ) elif mxfp4_backend == Mxfp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index ef6a9695f3..0d794584d7 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -409,7 +409,7 @@ def convert_to_nvfp4_moe_kernel_format( is_act_and_mul=is_act_and_mul, ) elif nvfp4_backend == NvFp4MoeBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( convert_to_humming_moe_kernel_format, ) @@ -523,7 +523,7 @@ def make_nvfp4_moe_quant_config( use_a16 = _use_a16(backend, use_a16) if backend == NvFp4MoeBackend.HUMMING: from vllm.model_executor.layers.fused_moe import RoutedExperts - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 92132f3ffb..b17a17d083 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -713,7 +713,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index ffbc6fba22..960ee95d60 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -689,7 +689,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index feddb77cff..6eb2455d0e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -474,7 +474,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: if self.wna16_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 128a5146f1..6bb0faa7c8 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from math import gcd from typing import TYPE_CHECKING, Any import regex as re @@ -28,12 +29,15 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.layers.quantization.utils.humming_utils import ( +from vllm.model_executor.layers.quantization.utils.humming import ( + check_and_fallback_input_schema, convert_to_humming_moe_kernel_format, get_humming_linear_compute_config, get_humming_moe_quant_config, + humming_update_schema_hadamard_block_size, input_schema_to_quant_key, make_humming_moe_kernel, + resolve_humming_layer_config, select_humming_moe_experts, weight_schema_to_quant_key, ) @@ -227,23 +231,13 @@ class HummingConfig(QuantizationConfig): return None config = group_config - layer_config = config - layer_dynamic = config.get("dynamic", {}) - if not isinstance(layer_dynamic, dict): - layer_dynamic = {} - for regex, override_config in layer_dynamic.items(): - if regex[:1] != "+": - continue - if re.match(regex[2:], prefix): - layer_config = config.copy() - layer_config.update(override_config) - break + layer_config = resolve_humming_layer_config(config, prefix) if "quant_method" in layer_config: return _hm.BaseWeightSchema.from_config(layer_config) return None - def get_layer_input_schema(self, config: dict[str, Any], prefix: str): + def get_layer_input_config(self, config: dict[str, Any], prefix: str): if self.is_layer_skipped(config, prefix): return None if config["quant_method"] in ["compressed-tensors", "modelopt"]: @@ -252,10 +246,15 @@ class HummingConfig(QuantizationConfig): return None config = group_config + config = resolve_humming_layer_config(config, prefix) if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP: - return _hm.BaseInputSchema.from_config(config) + return config return None + def get_layer_input_schema(self, config: dict[str, Any], prefix: str): + config = self.get_layer_input_config(config, prefix) + return None if config is None else _hm.BaseInputSchema.from_config(config) + def get_quant_config_for_layer( self, prefix: str, layer_type: str ) -> "HummingLayerQuantizationConfig | None": @@ -279,14 +278,25 @@ class HummingConfig(QuantizationConfig): if weight_schema is not None: input_schema = None force_input_schema = None + allow_input_schema_fallback = True if self.full_config: - input_schema = self.get_layer_input_schema(self.full_config, prefix) + input_config = self.get_layer_input_config(self.full_config, prefix) + if input_config is not None: + allow_input_schema_fallback = input_config.pop( + "allow_fallback", True + ) + input_schema = _hm.BaseInputSchema.from_config(input_config) if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG: quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy() quant_config["quant_method"] = "humming" - force_input_schema = self.get_layer_input_schema(quant_config, prefix) + input_config = self.get_layer_input_config(quant_config, prefix) + if input_config is not None: + allow_input_schema_fallback = input_config.pop( + "allow_fallback", False + ) + force_input_schema = _hm.BaseInputSchema.from_config(input_config) if input_schema is None: input_schema = force_input_schema @@ -299,6 +309,7 @@ class HummingConfig(QuantizationConfig): force_weight_schema=force_weight_schema, force_input_schema=force_input_schema, is_online_quant=is_online_quant, + allow_input_schema_fallback=allow_input_schema_fallback, ) return None @@ -318,6 +329,11 @@ class HummingConfig(QuantizationConfig): elif isinstance(layer, LinearBase): return UnquantizedLinearMethod() elif isinstance(layer, LinearBase): + input_size = layer.input_size + partition_size = getattr(layer, "input_size_per_partition", input_size) + is_online_quant = quant_config.is_online_quant + if is_online_quant and (input_size % 32 != 0 or partition_size % 32 != 0): + return UnquantizedLinearMethod() return HummingLinearMethod(quant_config) elif isinstance(layer, RoutedExperts): return HummingMoEMethod(quant_config, layer.moe_config) @@ -332,6 +348,7 @@ class HummingLayerQuantizationConfig(HummingConfig): force_weight_schema: "HummingWeightSchema | None" = None, force_input_schema: "HummingInputSchema | None" = None, is_online_quant: bool = False, + allow_input_schema_fallback: bool = True, ): self.weight_schema = weight_schema if input_schema is None: @@ -340,6 +357,7 @@ class HummingLayerQuantizationConfig(HummingConfig): self.force_weight_schema = force_weight_schema self.force_input_schema = force_input_schema self.is_online_quant = is_online_quant + self.allow_input_schema_fallback = allow_input_schema_fallback @classmethod def from_config(cls, config): @@ -373,21 +391,13 @@ class HummingLinearMethod(LinearMethodBase): if is_unquantized and self.is_online_quant: # online quant (fp16/bf16 -> quant_type) assert isinstance(self.weight_schema, _hm.HummingWeightSchema) - f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype) - has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type) - tensor_list = _hm.quantize_weight( - weight=loaded_weight, - dtype=self.weight_schema.b_dtype, - scale_dtype=self.weight_schema.bs_dtype or f16_dtype, - group_size=self.weight_schema.weight_scale_group_size, - has_zero_point=self.weight_schema.has_zero_point, - has_global_scale=has_global_scale, - is_fp_zero_point=self.weight_schema.is_fp_zero_point, - pack=True, + tensors = self.weight_schema.quant_tensor( + loaded_weight.to(param.device), + self.weight_schema, + layer.param_dtype, ) - key_list = ["weight", "weight_scale", "zero_point", "global_scale"] - for key, tensor in zip(key_list, tensor_list): + for key, tensor in tensors.items(): if tensor is None or tensor.nelement() == 0: continue param = getattr(layer, key) @@ -459,6 +469,18 @@ class HummingLinearMethod(LinearMethodBase): layer.output_partition_sizes = output_partition_sizes layer.extra_weight_attrs = extra_weight_attrs.copy() + input_schema = self.force_input_schema or self.input_schema + input_schema = input_schema.to_humming_schema(params_dtype) + for name in ("weight_schema", "force_weight_schema"): + schema = getattr(self, name) + if getattr(schema, "hadamard_block_size", None) == -1: + schema = humming_update_schema_hadamard_block_size( + weight_schema=schema, + input_schema=input_schema, + shape_k=input_size_per_partition, + ) + setattr(self, name, schema) + weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) new_weight_loader = self.prepare_weight_loader(layer, weight_loader) extra_weight_attrs["weight_loader"] = new_weight_loader @@ -502,7 +524,10 @@ class HummingLinearMethod(LinearMethodBase): return None # convert from checkpoint format to humming format - if not isinstance(self.weight_schema, _hm.HummingWeightSchema): + is_humming_weight = isinstance(self.weight_schema, _hm.HummingWeightSchema) + is_humming_input = isinstance(self.input_schema, _hm.HummingInputSchema) + + if not is_humming_weight or not is_humming_input: self.weight_schema, tensors = self.weight_schema.convert_humming( tensors=layer.state_dict(), shape_n_stacks=layer.output_partition_sizes, @@ -510,12 +535,13 @@ class HummingLinearMethod(LinearMethodBase): param_dtype=layer.param_dtype, ) - self.input_schema, _ = self.input_schema.convert_humming( + self.input_schema, input_tensors = self.input_schema.convert_humming( tensors=layer.state_dict(), shape_n_stacks=layer.output_partition_sizes, shape_k_stacks=[layer.input_size_per_partition], param_dtype=layer.param_dtype, ) + tensors.update(input_tensors) for name, _ in list(layer.named_parameters()): delattr(layer, name) @@ -530,11 +556,15 @@ class HummingLinearMethod(LinearMethodBase): assert isinstance(self.weight_schema, _hm.HummingWeightSchema) force_requant = self.force_weight_schema is not None if force_requant and self.weight_schema != self.force_weight_schema: + source_tensors = layer.state_dict() tensors = self.weight_schema.requant_tensors( - tensors=layer.state_dict(), + tensors=source_tensors, target_weight_schema=self.force_weight_schema, param_dtype=layer.param_dtype, ) + name = self.input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] self.weight_schema = self.force_weight_schema @@ -548,6 +578,12 @@ class HummingLinearMethod(LinearMethodBase): del tensors + self.input_schema = check_and_fallback_input_schema( + weight_schema=self.weight_schema, + input_schema=self.input_schema, + param_dtype=layer.param_dtype, + allow_fallback=self.quant_config.allow_input_schema_fallback, + ) self.layer_config = _hm.prepare_layer_config( shape_n=layer.output_partition_sizes_sum, shape_k=layer.input_size_per_partition, @@ -557,10 +593,16 @@ class HummingLinearMethod(LinearMethodBase): pad_k_to_multiple=128, has_bias=layer.has_bias, torch_dtype=layer.param_dtype, + device=layer.weight.device, ) - tensors = _hm.transform_humming_tensors( - self.layer_config, dict(layer.named_parameters()) - ) + source_tensors = dict(layer.named_parameters()) + tensors = _hm.transform_humming_tensors(self.layer_config, source_tensors) + name = self.input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] + if "bias" in tensors: + zero_bias = torch.zeros_like(tensors["bias"]) + layer.register_buffer("zero_bias", zero_bias) for name, _ in list(layer.named_parameters()): delattr(layer, name) for name, tensor in tensors.items(): @@ -583,8 +625,11 @@ class HummingLinearMethod(LinearMethodBase): weight=layer.weight, weight_scale=getattr(layer, "weight_scale", None), zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "bias", None), + bias=getattr(layer, "bias" if bias is not None else "zero_bias", None), weight_scale_2=getattr(layer, "weight_scale_2", None), + input_scale=getattr(layer, "input_scale", None), + input_scale_2=getattr(layer, "input_scale_2", None), + hadamard_block_size=self.weight_schema.hadamard_block_size, locks=self.locks, compute_config=self.compute_config, ) @@ -606,11 +651,15 @@ class HummingMoEMethod(FusedMoEMethodBase): # Derive QuantKeys from humming schemas. # Prefer force schemas (the final format after requant) over base. weight_key = weight_schema_to_quant_key( - self.force_weight_schema or self.weight_schema + self.force_weight_schema or self.weight_schema, moe.in_dtype ) - activation_key = input_schema_to_quant_key( - self.force_input_schema or self.input_schema + runtime_input_schema = check_and_fallback_input_schema( + weight_schema=self.force_weight_schema or self.weight_schema, + input_schema=self.force_input_schema or self.input_schema, + param_dtype=moe.in_dtype, + allow_fallback=quant_config.allow_input_schema_fallback, ) + activation_key = input_schema_to_quant_key(runtime_input_schema, moe.in_dtype) # Select Humming MoE experts self.experts_cls = select_humming_moe_experts( @@ -634,22 +683,14 @@ class HummingMoEMethod(FusedMoEMethodBase): # online quant (fp16/bf16 -> quant_type) if is_unquantized: assert isinstance(self.weight_schema, _hm.HummingWeightSchema) - f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype) - has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type) - tensor_list = _hm.quantize_weight( - weight=loaded_weight, - dtype=self.weight_schema.b_dtype, - scale_dtype=self.weight_schema.bs_dtype or f16_dtype, - group_size=self.weight_schema.weight_scale_group_size, - has_zero_point=self.weight_schema.has_zero_point, - has_global_scale=has_global_scale, - is_fp_zero_point=self.weight_schema.is_fp_zero_point, - pack=True, + tensors = self.weight_schema.quant_tensor( + loaded_weight.to(param.device), + self.weight_schema, + layer.param_dtype, ) - key_list = ["weight", "weight_scale", "zero_point", "global_scale"] success = True - for key, tensor in zip(key_list, tensor_list): + for key, tensor in tensors.items(): if tensor is None or tensor.nelement() == 0: continue sublayer_name = "w2" if shard_id == "w2" else "w13" @@ -695,6 +736,18 @@ class HummingMoEMethod(FusedMoEMethodBase): layer.num_experts = num_experts layer.param_dtype = params_dtype layer.intermediate_size = intermediate_size_per_partition + input_schema = self.force_input_schema or self.input_schema + input_schema = input_schema.to_humming_schema(params_dtype) + shape_k = gcd(hidden_size, intermediate_size_per_partition) + for name in ("weight_schema", "force_weight_schema"): + schema = getattr(self, name) + if getattr(schema, "hadamard_block_size", None) == -1: + schema = humming_update_schema_hadamard_block_size( + weight_schema=schema, + input_schema=input_schema, + shape_k=shape_k, + ) + setattr(self, name, schema) weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) weight_loader = self.prepare_weight_loader(layer, weight_loader) extra_weight_attrs["weight_loader"] = weight_loader @@ -761,6 +814,7 @@ class HummingMoEMethod(FusedMoEMethodBase): weight_schema=self.weight_schema, input_schema=self.input_schema, force_weight_schema=self.force_weight_schema, + allow_input_schema_fallback=self.quant_config.allow_input_schema_fallback, ) # Build the MoE kernel diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index fef2c0c7ee..6a27332b06 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -389,7 +389,7 @@ class MoeWNA16Method(FusedMoEMethodBase): self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: if self.wna16_backend == WNA16MoEBackend.HUMMING: - from vllm.model_executor.layers.quantization.utils.humming_utils import ( + from vllm.model_executor.layers.quantization.utils.humming import ( get_humming_moe_quant_config, ) diff --git a/vllm/model_executor/layers/quantization/utils/humming/__init__.py b/vllm/model_executor/layers/quantization/utils/humming/__init__.py new file mode 100644 index 0000000000..4303f74221 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/__init__.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Humming quantization integration.""" + +from vllm.model_executor.layers.quantization.utils.humming.activation import ( + get_humming_activation, +) +from vllm.model_executor.layers.quantization.utils.humming.linear import ( + apply_humming_linear, + convert_linear_layer_to_humming_standard, + get_humming_linear_compute_config, + prepare_humming_linear_layer_config, +) +from vllm.model_executor.layers.quantization.utils.humming.moe import ( + HummingMoEQuantConfig, + convert_to_humming_moe_kernel_format, + get_humming_moe_quant_config, + make_humming_moe_kernel, + make_humming_moe_quant_config, + select_humming_moe_experts, +) +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + check_and_fallback_input_schema, + humming_is_layer_skipped, + humming_update_schema_hadamard_block_size, + input_schema_to_quant_key, + quant_key_to_input_schema, + resolve_humming_layer_config, + weight_schema_to_quant_key, +) + +__all__ = [ + "check_and_fallback_input_schema", + "get_humming_activation", + "HummingMoEQuantConfig", + "get_humming_moe_quant_config", + "humming_is_layer_skipped", + "humming_update_schema_hadamard_block_size", + "make_humming_moe_quant_config", + "apply_humming_linear", + "convert_linear_layer_to_humming_standard", + "get_humming_linear_compute_config", + "prepare_humming_linear_layer_config", + "make_humming_moe_kernel", + "select_humming_moe_experts", + "input_schema_to_quant_key", + "quant_key_to_input_schema", + "resolve_humming_layer_config", + "weight_schema_to_quant_key", + "convert_to_humming_moe_kernel_format", +] diff --git a/vllm/model_executor/layers/quantization/utils/humming/activation.py b/vllm/model_executor/layers/quantization/utils/humming/activation.py new file mode 100644 index 0000000000..3942944d2d --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/activation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoE activation expressions for Humming input processing.""" + +from dataclasses import replace + +from vllm.model_executor.layers.fused_moe.activation import ( + ApplyMoEActivationConfig, + MoEActivation, +) + +_SILU_IMPL = "a / (1.f + expf(-a))" +_GELU_IMPL = "0.5f * a * (1.f + erff(a * 0.7071067811865476f))" +_RELU2_IMPL = "fmaxf(a, 0.f) * fmaxf(a, 0.f)" +_GELU_TANH_IMPL = """[](float a) { + const float cubic = 0.044715f * a * a * a; + const float inner = 0.7978845608028654f * (a + cubic); + return 0.5f * a * (1.f + tanhf(inner)); +}(a)""" +_CLAMPED_SILU_IMPL = """[](float a, float b) {{ + const float gate = fminf(a, {limit}); + const float up = fmaxf(fminf(b, {limit}), -{limit}); + const float activated_gate = gate / (1.f + expf(-{alpha} * gate)); + return activated_gate * (up + {beta}); +}}(a, b)""" + +# Parameterized expressions are resolved by get_humming_activation(). +_HUMMING_ACTIVATIONS: dict[MoEActivation, dict[str, str]] = { + MoEActivation.SILU: { + "activation_type": "binary_split", + "activation_impl": f"({_SILU_IMPL}) * b", + }, + MoEActivation.GELU: { + "activation_type": "binary_split", + "activation_impl": f"({_GELU_IMPL}) * b", + }, + MoEActivation.GELU_TANH: { + "activation_type": "binary_split", + "activation_impl": f"({_GELU_TANH_IMPL}) * b", + }, + MoEActivation.RELU2: { + "activation_type": "binary_split", + "activation_impl": f"({_RELU2_IMPL}) * b", + }, + MoEActivation.SWIGLUOAI: { + "activation_type": "binary_interleaved", + "activation_impl": _CLAMPED_SILU_IMPL, + }, + MoEActivation.SITU: { + "activation_type": "binary_split", + "activation_impl": """[](float a, float b) {{ + const float gate = {situ_beta} * tanhf(a / {situ_beta}); + const float activated_gate = gate / (1.f + expf(-a)); + const float up = {situ_up}; + return activated_gate * up; + }}(a, b)""", + }, + MoEActivation.SWIGLUOAI_UNINTERLEAVE: { + "activation_type": "binary_split", + "activation_impl": _CLAMPED_SILU_IMPL, + }, + MoEActivation.SWIGLUSTEP: { + "activation_type": "binary_split", + "activation_impl": """[](float a, float b) { + const float activated_gate = a / (1.f + expf(-a)); + const float gate = fminf(activated_gate, 7.f); + const float up = fmaxf(fminf(b, 7.f), -7.f); + return gate * up; + }(a, b)""", + }, + MoEActivation.SILU_NO_MUL: { + "activation_type": "unary", + "activation_impl": _SILU_IMPL, + }, + MoEActivation.GELU_NO_MUL: { + "activation_type": "unary", + "activation_impl": _GELU_IMPL, + }, + MoEActivation.GELU_TANH_NO_MUL: { + "activation_type": "unary", + "activation_impl": _GELU_TANH_IMPL, + }, + MoEActivation.RELU2_NO_MUL: { + "activation_type": "unary", + "activation_impl": _RELU2_IMPL, + }, +} + + +def get_humming_activation( + activation: MoEActivation, + config: ApplyMoEActivationConfig | None = None, +) -> dict[str, str]: + config = config or ApplyMoEActivationConfig() + use_clamped_silu = False + if activation == MoEActivation.SWIGLUOAI: + config = ApplyMoEActivationConfig(clamp_limit=7.0, alpha=1.702, beta=1.0) + use_clamped_silu = True + elif activation == MoEActivation.SILU and config.clamp_limit is not None: + config = replace(config, alpha=1.0, beta=0.0) + use_clamped_silu = True + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE: + if config.clamp_limit is None: + raise ValueError("SWIGLUOAI_UNINTERLEAVE requires clamp_limit") + use_clamped_silu = True + + definition = _HUMMING_ACTIVATIONS[activation].copy() + if activation == MoEActivation.SITU: + if config.activation_situ_beta is None: + raise ValueError("SITU requires activation_situ_beta") + linear_beta = config.activation_situ_linear_beta + situ_up = "b" + if linear_beta is not None and linear_beta > 0: + situ_up = f"({float(linear_beta)}f) * tanhf(b / ({float(linear_beta)}f))" + situ_beta = f"({float(config.activation_situ_beta)}f)" + activation_impl = definition["activation_impl"] + activation_impl = activation_impl.format(situ_beta=situ_beta, situ_up=situ_up) + definition["activation_impl"] = activation_impl + elif use_clamped_silu: + assert config.clamp_limit is not None + limit = f"({float(config.clamp_limit)}f)" + alpha = f"({float(config.alpha)}f)" + beta = f"({float(config.beta)}f)" + activation_impl = _CLAMPED_SILU_IMPL.format(limit=limit, alpha=alpha, beta=beta) + definition["activation_impl"] = activation_impl + return definition diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py new file mode 100644 index 0000000000..339ff640b0 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/linear.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepare and execute Humming linear layers.""" + +import json +from typing import TYPE_CHECKING + +import torch + +from vllm import envs +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + check_and_fallback_input_schema, +) + +if TYPE_CHECKING: + from vllm.utils.humming import BaseInputSchema, LayerConfig + + +def convert_linear_layer_to_humming_standard( + layer: LinearBase, name_map: dict[str, str] +): + """Rename/reshape a linear layer's quantized params (the canonical MPLinear + layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names + and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" + for name, checkpoint_name in name_map.items(): + tensor = getattr(layer, checkpoint_name) + delattr(layer, checkpoint_name) + + if name == "weight": + input_dim = getattr(tensor, "input_dim", 1) + output_dim = getattr(tensor, "output_dim", 0) + + if input_dim == 0 and output_dim == 1: + tensor = tensor.transpose(1, 0).contiguous() + else: + assert output_dim == 0 and input_dim == 1 + + tensor = tensor.view(tensor.size(0), -1).view(torch.int32) + elif name in ["weight_scale", "zero_point"]: + if getattr(tensor, "output_dim", 0) == 1: + tensor = tensor.transpose(0, 1).contiguous() + if tensor.ndim == 1: + tensor = tensor.unsqueeze(1) + + tensor = tensor.view(torch.int32) if name == "zero_point" else tensor + + if isinstance(tensor, torch.nn.Parameter): + param = tensor + else: + param = torch.nn.Parameter(tensor, requires_grad=False) + + setattr(layer, name, param) + + +def prepare_humming_linear_layer_config( + layer: LinearBase, + quant_config: dict, + input_quant_config: dict | None = None, + *, + input_schema: "BaseInputSchema | None" = None, +) -> "LayerConfig": + from vllm.utils.humming import ( + BaseInputSchema, + BaseWeightSchema, + HummingInputSchema, + prepare_layer_config, + transform_humming_tensors, + ) + + weight_schema = BaseWeightSchema.from_config(quant_config) + if input_schema is None: + if input_quant_config is not None: + input_schema = BaseInputSchema.from_config(input_quant_config) + else: + input_schema = HummingInputSchema() + + # ReplicatedLinear has no TP partitioning and so does not set + # input_size_per_partition; for it that is just input_size. Use hasattr + # rather than getattr's default arg, which is evaluated eagerly and would + # raise on layers lacking input_size (e.g. ParallelLMHead). + if hasattr(layer, "input_size_per_partition"): + input_size_per_partition = layer.input_size_per_partition + else: + input_size_per_partition = layer.input_size + shape_k_stacks = [input_size_per_partition] + shape_n_stacks = layer.output_partition_sizes + + # Step 1: convert weight and input schemas to humming standard format + source_tensors = dict(layer.named_parameters()) + weight_schema, tensors = weight_schema.convert_humming( + tensors=source_tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) + input_schema, input_tensors = input_schema.convert_humming( + tensors=source_tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) + + tensors.update(input_tensors) + layer.weight_schema = weight_schema + input_schema = check_and_fallback_input_schema( + weight_schema=weight_schema, + input_schema=input_schema, + param_dtype=layer.params_dtype, + ) + for name in ("input_scale", "input_scale_2"): + if name != input_schema.static_tensor_scale_name: + tensors.pop(name, None) + input_tensors.pop(name, None) + + # Step 2: transform weight (humming standard format) for forwarding. + config = prepare_layer_config( + shape_n=sum(layer.output_partition_sizes), + shape_k=input_size_per_partition, + weight_schema=weight_schema, + input_schema=input_schema, + pad_n_to_multiple=256, + pad_k_to_multiple=128, + has_bias=layer.has_bias, + torch_dtype=layer.params_dtype, + device=tensors["weight"].device, + ) + tensors = transform_humming_tensors(config, tensors) + tensors.update(input_tensors) + if "bias" in tensors: + zero_bias = torch.zeros_like(tensors["bias"]) + layer.register_buffer("zero_bias", zero_bias) + for name, _ in list(layer.named_parameters()): + delattr(layer, name) + for name, tensor in tensors.items(): + param = torch.nn.Parameter(tensor, requires_grad=False) + setattr(layer, name, param) + + return config + + +def get_humming_linear_compute_config() -> str: + return json.dumps( + { + "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, + "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, + "gemm_type": "dense", + } + ) + + +def apply_humming_linear( + layer: LinearBase, + x: torch.Tensor, + *, + skip_bias_add: bool = False, + layer_config: "LayerConfig", + compute_config: str, + locks: torch.Tensor, +) -> torch.Tensor: + from vllm.utils.humming import humming_forward + + mode = layer_config.input_quant_mode + input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None + input_scale_2 = ( + getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None + ) + flatten_inputs = x.reshape(-1, x.size(-1)) + output = humming_forward( + layer_config, + inputs=flatten_inputs, + weight=layer.weight, + weight_scale=getattr(layer, "weight_scale", None), + zero_point=getattr(layer, "zero_point", None), + bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None), + weight_scale_2=getattr(layer, "weight_scale_2", None), + input_scale=input_scale, + input_scale_2=input_scale_2, + hadamard_block_size=layer.weight_schema.hadamard_block_size, + locks=locks, + compute_config=compute_config, + ) + return output.view(*x.shape[:-1], output.size(-1)) diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming/moe.py similarity index 51% rename from vllm/model_executor/layers/quantization/utils/humming_utils.py rename to vllm/model_executor/layers/quantization/utils/humming/moe.py index f48c33d4cd..37eb9c26f0 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming/moe.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import json +"""Configure, prepare weights for, and assemble Humming MoE kernels.""" + from dataclasses import dataclass from typing import TYPE_CHECKING, Any -import regex as re import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +import vllm.model_executor.layers.quantization.utils.humming.schema as humming_schema from vllm import envs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( @@ -23,34 +24,19 @@ from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( HummingGroupedExperts, HummingIndexedExperts, ) -from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.layers.quantization.utils.humming.schema import ( + humming_is_layer_skipped, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( - FP4_DTYPE, FP8_DTYPE, - INT4_DTYPE, - INT8_DTYPE, - MXFP_SCALE_DTYPE, GroupShape, QuantKey, - ScaleDesc, ) from vllm.utils.import_utils import has_humming if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from vllm.utils.humming import ( - AWQWeightSchema, - BaseInputSchema, - BaseWeightSchema, - CompressedTensorsInputSchema, - CompressedTensorsWeightSchema, - Fp8WeightSchema, - GPTQWeightSchema, - HummingInputSchema, - HummingWeightSchema, - LayerConfig, - ) - from vllm.utils.humming import dtypes as humming_dtypes + from vllm.utils.humming import LayerConfig logger = init_logger(__name__) @@ -59,504 +45,12 @@ logger = init_logger(__name__) class HummingMoEQuantConfig(FusedMoEQuantConfig): w1_humming_config: "LayerConfig" w2_humming_config: "LayerConfig" - - -if has_humming(): - from vllm.utils.humming import dtypes as humming_dtypes - - _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = { - humming_dtypes.float4e2m1: FP4_DTYPE, - humming_dtypes.float8e4m3: FP8_DTYPE, - humming_dtypes.float8e5m2: torch.float8_e5m2, - humming_dtypes.int8: torch.int8, - humming_dtypes.uint4: INT4_DTYPE, - humming_dtypes.uint8: INT8_DTYPE, - humming_dtypes.uint2: torch.uint8, - humming_dtypes.uint3: torch.uint8, - humming_dtypes.uint5: torch.uint8, - humming_dtypes.uint6: torch.uint8, - humming_dtypes.uint7: torch.uint8, - } - - _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { - humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE, - humming_dtypes.float8e4m3: FP8_DTYPE, - humming_dtypes.float16: torch.float16, - humming_dtypes.bfloat16: torch.bfloat16, - humming_dtypes.float32: torch.float32, - } - - -def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: - """Map humming group sizes to QuantKey GroupShape. - - group_size: elements per group along K (col); 0 means full dimension. - group_size_n: elements per group along N (row); 0 means 1 (per-row). - - GroupShape convention: row = N dim, col = K dim. - """ - if group_size == 0 and group_size_n == 0: - return GroupShape.PER_CHANNEL - - row = group_size_n if group_size_n > 0 else 1 - col = group_size if group_size > 0 else -1 - return GroupShape(row=row, col=col) - - -# ---- HummingWeightSchema (post-conversion) -------------------------------- - - -def _humming_weight_schema_to_quant_key( - schema: "HummingWeightSchema", -) -> QuantKey: - from vllm.utils.humming import WeightScaleType - - """Convert a HummingWeightSchema to a QuantKey.""" - dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype] - - if schema.bs_dtype is not None: - scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype] - else: - scale_dtype = torch.float32 - - group_shape = _group_shape( - schema.weight_scale_group_size, - schema.weight_scale_group_size_n, - ) - - scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) - - scale2 = None - if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR: - scale2 = ScaleDesc( - dtype=torch.float32, - static=True, - group_shape=GroupShape.PER_TENSOR, - ) - - return QuantKey( - dtype=dtype, - scale=scale, - scale2=scale2, - symmetric=not schema.has_zero_point, - ) - - -# ---- Checkpoint-format weight schemas (pre-conversion) -------------------- - - -def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey: - if schema.weight_block_size is not None: - gs_n, gs_k = schema.weight_block_size - group_shape = GroupShape(row=gs_n, col=gs_k) - else: - group_shape = GroupShape.PER_CHANNEL - - scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape) - return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True) - - -def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey: - group_shape = _group_shape(schema.group_size) - scale = ScaleDesc( - dtype=torch.float16, - static=True, - group_shape=group_shape, - ) - return QuantKey( - dtype=INT4_DTYPE, - scale=scale, - symmetric=not schema.zero_point, - ) - - -def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey: - group_shape = _group_shape(schema.group_size) - scale = ScaleDesc( - dtype=torch.float16, - static=True, - group_shape=group_shape, - ) - return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym) - - -def _compressed_tensors_weight_schema_to_quant_key( - schema: "CompressedTensorsWeightSchema", -) -> QuantKey: - # Determine dtype from format/type/num_bits - fmt = schema.format - if fmt in ("int-quantized", "float-quantized", "naive-quantized"): - dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE - elif "nvfp4" in fmt or "mxfp4" in fmt: - dtype = FP4_DTYPE - else: - dtype = _HUMMING_TO_QUANT_DTYPE[ - humming_dtypes.DataType.from_str(f"uint{schema.num_bits}") - ] - - # Determine group shape from strategy - if schema.strategy in ("group", "tensor_group"): - group_shape = _group_shape(schema.group_size or 0) - elif schema.strategy == "block" and schema.block_structure is not None: - group_shape = GroupShape( - row=schema.block_structure[0], - col=schema.block_structure[1], - ) - else: - group_shape = GroupShape.PER_CHANNEL - - # Determine scale dtype - if "mxfp" in fmt: - scale_dtype = MXFP_SCALE_DTYPE - elif "nvfp4" in fmt: - scale_dtype = FP8_DTYPE - else: - scale_dtype = torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) - - scale2 = None - if "nvfp4" in fmt or schema.strategy == "tensor_group": - scale2 = ScaleDesc( - dtype=torch.float32, - static=True, - group_shape=GroupShape.PER_TENSOR, - ) - - return QuantKey( - dtype=dtype, - scale=scale, - scale2=scale2, - symmetric=schema.symmetric, - ) - - -# ---- Dispatch for any BaseWeightSchema ------------------------------------ - - -def weight_schema_to_quant_key( - schema: "BaseWeightSchema", -) -> QuantKey: - from vllm.utils.humming import ( - AWQWeightSchema, - BitnetWeightSchema, - CompressedTensorsWeightSchema, - Fp8WeightSchema, - GptOssMxfp4WeightSchema, - GPTQWeightSchema, - HummingWeightSchema, - ModeloptMxfp8WeightSchema, - ModeloptNvfp4WeightSchema, - Mxfp4WeightSchema, - ) - - """Convert any BaseWeightSchema to a QuantKey.""" - if isinstance(schema, HummingWeightSchema): - return _humming_weight_schema_to_quant_key(schema) - - # Schemas with fixed QuantKeys - if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)): - return QuantKey( - dtype=FP4_DTYPE, - scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), - ) - if isinstance(schema, ModeloptMxfp8WeightSchema): - return QuantKey( - dtype=FP8_DTYPE, - scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)), - ) - if isinstance(schema, ModeloptNvfp4WeightSchema): - return QuantKey( - dtype=FP4_DTYPE, - scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)), - scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR), - ) - if isinstance(schema, BitnetWeightSchema): - return QuantKey( - dtype=torch.uint8, - scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL), - ) - - # Schemas requiring config inspection - if isinstance(schema, Fp8WeightSchema): - return _fp8_weight_schema_to_quant_key(schema) - if isinstance(schema, AWQWeightSchema): - return _awq_weight_schema_to_quant_key(schema) - if isinstance(schema, GPTQWeightSchema): - return _gptq_weight_schema_to_quant_key(schema) - if isinstance(schema, CompressedTensorsWeightSchema): - return _compressed_tensors_weight_schema_to_quant_key(schema) - - raise TypeError(f"Unsupported weight schema type: {type(schema)}") - - -# ---- HummingInputSchema (post-conversion) ---------------------------------- - - -def _humming_input_schema_to_quant_key( - schema: "HummingInputSchema", -) -> QuantKey | None: - """Convert a HummingInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs.""" - if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: - return None - - dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype] - - gs = schema.input_scale_group_size - group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN - - # Pick the scale dtype the Humming kernel actually consumes. An explicit - # input_scale_dtype always wins. Otherwise infer from the grouping: MX - # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8) - # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a - # float32 scale. Getting this right lets a grouped FP8 activation match - # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key. - if schema.input_scale_dtype is not None: - scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype] - elif gs == 32: - scale_dtype = MXFP_SCALE_DTYPE - else: - scale_dtype = torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) - - return QuantKey(dtype=dtype, scale=scale, symmetric=True) - - -# ---- Checkpoint-format input schemas (pre-conversion) ---------------------- - - -def _resolve_input_quant_key( - origin_a_dtype: "humming_dtypes.DataType", - group_size: int, -) -> QuantKey | None: - from vllm.utils.humming import HummingInputSchema - - """Resolve the actual activation QuantKey after platform fallback.""" - a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype) - if a_dtype is None or a_dtype.num_bits >= 16: - return None - - dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype] - gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0 - group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN - scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32 - - scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape) - return QuantKey(dtype=dtype, scale=scale, symmetric=True) - - -def _compressed_tensors_input_schema_to_quant_key( - schema: "CompressedTensorsInputSchema", -) -> QuantKey | None: - type_bits_to_dtype = { - ("float", 8): humming_dtypes.float8e4m3, - ("float", 4): humming_dtypes.float4e2m1, - ("int", 8): humming_dtypes.int8, - ("int", 4): humming_dtypes.int4, - } - origin = type_bits_to_dtype.get((schema.type, schema.num_bits)) - if origin is None: - return None - return _resolve_input_quant_key(origin, schema.group_size) - - -# ---- Dispatch for any BaseInputSchema ------------------------------------- - - -def input_schema_to_quant_key( - schema: "BaseInputSchema", -) -> QuantKey | None: - from vllm.utils.humming import ( - CompressedTensorsInputSchema, - Fp8InputSchema, - HummingInputSchema, - ModeloptNvfp4InputSchema, - ) - - """Convert any BaseInputSchema to a QuantKey. Returns None if - the schema represents unquantized (bf16/fp16) inputs.""" - if isinstance(schema, HummingInputSchema): - return _humming_input_schema_to_quant_key(schema) - - if isinstance(schema, Fp8InputSchema): - return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0) - - if isinstance(schema, ModeloptNvfp4InputSchema): - return _resolve_input_quant_key( - humming_dtypes.float8e4m3, - schema.group_size, - ) - - if isinstance(schema, CompressedTensorsInputSchema): - return _compressed_tensors_input_schema_to_quant_key(schema) - - raise TypeError(f"Unsupported input schema type: {type(schema)}") - - -def humming_is_layer_skipped(config: dict[str, Any], prefix: str): - if not config: - return True - - keys = ["ignored_layers", "ignore", "modules_to_not_convert"] - ignored_layers: list[str] = [] - for key in keys: - candidate = config.get(key, []) or [] - if candidate: - ignored_layers = candidate - break - - if any(module_name in prefix for module_name in ignored_layers): - return True - if "lm_head" in prefix: - return True - - for regex in config.get("dynamic", {}): - if regex[:1] != "-": - continue - if re.match(regex[2:], prefix): - return True - - return False - - -def convert_linear_layer_to_humming_standard( - layer: LinearBase, name_map: dict[str, str] -): - """Rename/reshape a linear layer's quantized params (the canonical MPLinear - layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names - and layout humming's weight schema expects (``weight`` / ``weight_scale``).""" - for name, checkpoint_name in name_map.items(): - tensor = getattr(layer, checkpoint_name) - delattr(layer, checkpoint_name) - - if name == "weight": - input_dim = getattr(tensor, "input_dim", 1) - output_dim = getattr(tensor, "output_dim", 0) - - if input_dim == 0 and output_dim == 1: - tensor = tensor.transpose(1, 0).contiguous() - else: - assert output_dim == 0 and input_dim == 1 - - tensor = tensor.view(tensor.size(0), -1).view(torch.int32) - elif name in ["weight_scale", "zero_point"]: - if getattr(tensor, "output_dim", 0) == 1: - tensor = tensor.transpose(0, 1).contiguous() - if tensor.ndim == 1: - tensor = tensor.unsqueeze(1) - - tensor = tensor.view(torch.int32) if name == "zero_point" else tensor - - if isinstance(tensor, torch.nn.Parameter): - param = tensor - else: - param = torch.nn.Parameter(tensor, requires_grad=False) - - setattr(layer, name, param) - - -def prepare_humming_linear_layer_config( - layer: LinearBase, - quant_config: dict, - input_quant_config: dict | None = None, -) -> "LayerConfig": - from vllm.utils.humming import ( - BaseInputSchema, - BaseWeightSchema, - HummingInputSchema, - prepare_layer_config, - transform_humming_tensors, - ) - - weight_schema = BaseWeightSchema.from_config(quant_config) - if input_quant_config is not None: - input_schema = BaseInputSchema.from_config(input_quant_config) - else: - input_schema = HummingInputSchema() - - # ReplicatedLinear has no TP partitioning and so does not set - # input_size_per_partition; for it that is just input_size. Use hasattr - # rather than getattr's default arg, which is evaluated eagerly and would - # raise on layers lacking input_size (e.g. ParallelLMHead). - if hasattr(layer, "input_size_per_partition"): - input_size_per_partition = layer.input_size_per_partition - else: - input_size_per_partition = layer.input_size - shape_k_stacks = [input_size_per_partition] - shape_n_stacks = layer.output_partition_sizes - - # Step 1: convert weight and input schemas to humming standard format - weight_schema, tensors = weight_schema.convert_humming( - tensors=dict(layer.named_parameters()), - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.params_dtype, - ) - input_schema, _ = input_schema.convert_humming( - tensors={}, - shape_n_stacks=shape_n_stacks, - shape_k_stacks=shape_k_stacks, - param_dtype=layer.params_dtype, - ) - - layer.weight_schema = weight_schema - - # Step 2: transform weight (humming standard format) for forwarding. - config = prepare_layer_config( - shape_n=sum(layer.output_partition_sizes), - shape_k=input_size_per_partition, - weight_schema=weight_schema, - input_schema=input_schema, - pad_n_to_multiple=256, - pad_k_to_multiple=128, - has_bias=layer.has_bias, - torch_dtype=layer.params_dtype, - ) - tensors = transform_humming_tensors(config, tensors) - for name, _ in list(layer.named_parameters()): - delattr(layer, name) - for name, tensor in tensors.items(): - param = torch.nn.Parameter(tensor, requires_grad=False) - setattr(layer, name, param) - - return config - - -def get_humming_linear_compute_config() -> str: - return json.dumps( - { - "use_batch_invariant": envs.VLLM_BATCH_INVARIANT, - "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM, - "gemm_type": "dense", - } - ) - - -def apply_humming_linear( - layer: LinearBase, - x: torch.Tensor, - *, - layer_config: "LayerConfig", - compute_config: str, - locks: torch.Tensor, -) -> torch.Tensor: - from vllm.utils.humming import humming_forward - - flatten_inputs = x.reshape(-1, x.size(-1)) - output = humming_forward( - layer_config, - inputs=flatten_inputs, - weight=layer.weight, - weight_scale=getattr(layer, "weight_scale", None), - zero_point=getattr(layer, "zero_point", None), - bias=getattr(layer, "bias", None), - weight_scale_2=getattr(layer, "weight_scale_2", None), - locks=locks, - compute_config=compute_config, - ) - return output.view(*x.shape[:-1], output.size(-1)) + w1_hadamard_block_size: int = 1 + w2_hadamard_block_size: int = 1 + w1_input_scale: torch.Tensor | None = None + w2_input_scale: torch.Tensor | None = None + w1_input_scale_2: torch.Tensor | None = None + w2_input_scale_2: torch.Tensor | None = None def make_humming_moe_quant_config( @@ -648,8 +142,12 @@ def get_humming_moe_quant_config( q_dtype is not None and q_dtype.startswith("float8") and input_scale_group_size == 128 + and weight_schema.hadamard_block_size <= 1 + and input_schema.input_quant_mode in (None, "dynamic_group") ): - q_dtype = _HUMMING_TO_QUANT_DTYPE.get(input_schema.a_dtype, FP8_DTYPE) + q_dtype = humming_schema._HUMMING_TO_QUANT_DTYPE.get( + input_schema.a_dtype, FP8_DTYPE + ) activation_group_shape = GroupShape(row=1, col=input_scale_group_size) weight_scale_group_size = weight_schema.weight_scale_group_size @@ -665,7 +163,7 @@ def get_humming_moe_quant_config( else: weight_group_shape = GroupShape(row=weight_scale_group_size, col=1) - return make_humming_moe_quant_config( + config = make_humming_moe_quant_config( quant_dtype=q_dtype, weight_dtype=str(weight_schema.b_dtype), weight_group_shape=weight_group_shape, @@ -683,6 +181,17 @@ def get_humming_moe_quant_config( gemm1_clamp_limit=gemm1_clamp_limit, humming_configs=humming_configs, ) + for prefix, sublayer in (("w1", "w13"), ("w2", "w2")): + setattr( + config, + f"{prefix}_hadamard_block_size", + layer.weight_schemas[sublayer].hadamard_block_size, + ) + for name in ("input_scale", "input_scale_2"): + key = f"{prefix}_{name}" + value = getattr(layer, f"{sublayer}_{name}", None) + setattr(config, key, value) + return config def select_humming_moe_experts( @@ -846,12 +355,6 @@ def _convert_sublayer_to_humming( Tuple of (converted_weight_schema, converted_input_schema) """ - from vllm.utils.humming import HummingWeightSchema - - if isinstance(weight_schema, HummingWeightSchema): - # Already in Humming format - return weight_schema, input_schema - tensors = _extract_sublayer_tensors(layer, sublayer_name) shape_k_stacks = [shape_k] @@ -867,13 +370,14 @@ def _convert_sublayer_to_humming( num_experts=num_experts, ) - converted_input_schema, _ = input_schema.convert_humming( - tensors=converted_tensors, + converted_input_schema, input_tensors = input_schema.convert_humming( + tensors=tensors, shape_n_stacks=shape_n_stacks, shape_k_stacks=shape_k_stacks, param_dtype=param_dtype, num_experts=num_experts, ) + converted_tensors.update(input_tensors) _replace_layer_parameters(layer, sublayer_name, converted_tensors) @@ -907,11 +411,16 @@ def _prepare_and_transform_sublayer( has_bias=has_bias, num_experts=num_experts, torch_dtype=param_dtype, + device=getattr(layer, f"{sublayer_name}_weight").device, ) - tensors = transform_humming_tensors( - config, - _extract_sublayer_tensors(layer, sublayer_name), - ) + source_tensors = _extract_sublayer_tensors(layer, sublayer_name) + for name in ("input_scale", "input_scale_2"): + if name != input_schema.static_tensor_scale_name: + source_tensors.pop(name, None) + tensors = transform_humming_tensors(config, source_tensors) + for name in ("input_scale", "input_scale_2"): + if name in source_tensors: + tensors[name] = source_tensors[name] _replace_layer_parameters(layer, sublayer_name, tensors) return config @@ -927,6 +436,7 @@ def _process_single_sublayer( num_experts: int, param_dtype: torch.dtype, force_weight_schema: Any | None = None, + allow_input_schema_fallback: bool = True, ) -> tuple[Any, Any, "LayerConfig"]: """Process a single sublayer: convert, optionally requant, prepare, and transform. @@ -944,6 +454,7 @@ def _process_single_sublayer( num_experts: Number of experts param_dtype: Parameter data type force_weight_schema: Optional schema to force requantization to + allow_input_schema_fallback: Whether incompatible input schemas may be replaced. Returns: Tuple of the final weight schema, input schema, and Humming layer config. @@ -966,19 +477,28 @@ def _process_single_sublayer( # Step 2: Force requant if needed assert isinstance(current_weight_schema, HummingWeightSchema) if force_weight_schema is not None and current_weight_schema != force_weight_schema: - tensors = _extract_sublayer_tensors(layer, sublayer_name) + source_tensors = _extract_sublayer_tensors(layer, sublayer_name) tensors = current_weight_schema.requant_tensors( - tensors=tensors, + tensors=source_tensors, target_weight_schema=force_weight_schema, param_dtype=param_dtype, ) + name = current_input_schema.static_tensor_scale_name + if name is not None: + tensors[name] = source_tensors[name] current_weight_schema = force_weight_schema _replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True) del tensors # Step 3: Prepare layer metadata and transform weights + current_input_schema = humming_schema.check_and_fallback_input_schema( + weight_schema=current_weight_schema, + input_schema=current_input_schema, + param_dtype=param_dtype, + allow_fallback=allow_input_schema_fallback, + ) config = _prepare_and_transform_sublayer( layer=layer, sublayer_name=sublayer_name, @@ -1001,6 +521,7 @@ def convert_to_humming_moe_kernel_format( weight_schema: Any | None = None, input_schema: Any | None = None, force_weight_schema: Any | None = None, + allow_input_schema_fallback: bool = True, ) -> dict[str, "LayerConfig"]: """Convert MoE weights from checkpoint format to Humming kernel format. @@ -1023,6 +544,7 @@ def convert_to_humming_moe_kernel_format( input_schema: Optional initial input quantization schema. If None, built from quant_config or env vars. force_weight_schema: Optional schema to force requantization to + allow_input_schema_fallback: Whether incompatible input schemas may be replaced. Side effects: - Modifies layer parameters in place @@ -1041,20 +563,23 @@ def convert_to_humming_moe_kernel_format( "Must provide either weight_schema/input_schema or quant_config" ) - from vllm.model_executor.layers.quantization.utils.humming_utils import ( - humming_is_layer_skipped, - ) from vllm.utils.humming import BaseWeightSchema, HummingInputSchema if weight_schema is None: weight_schema = BaseWeightSchema.from_config(quant_config) if input_schema is None: - input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {} + input_quant_config = (envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}).copy() if humming_is_layer_skipped(input_quant_config, layer.layer_name): input_schema = HummingInputSchema() else: # TODO: read input_quant_config from quant_config + input_quant_config = humming_schema.resolve_humming_layer_config( + input_quant_config, layer.layer_name + ) + allow_input_schema_fallback = input_quant_config.pop( + "allow_fallback", False + ) input_schema = HummingInputSchema.from_config(input_quant_config) # Build sublayer configs from layer properties if not provided @@ -1089,6 +614,7 @@ def convert_to_humming_moe_kernel_format( num_experts=num_experts, param_dtype=param_dtype, force_weight_schema=force_weight_schema, + allow_input_schema_fallback=allow_input_schema_fallback, ) ) diff --git a/vllm/model_executor/layers/quantization/utils/humming/schema.py b/vllm/model_executor/layers/quantization/utils/humming/schema.py new file mode 100644 index 0000000000..c84e8a370d --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/humming/schema.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Map Humming schemas and handle shared checkpoint quantization settings.""" + +import dataclasses +from typing import TYPE_CHECKING, Any + +import regex as re +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP4_DTYPE, + FP8_DTYPE, + INT4_DTYPE, + INT8_DTYPE, + MXFP_SCALE_DTYPE, + GroupShape, + QuantKey, + ScaleDesc, +) +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types +from vllm.utils.import_utils import has_humming + +if TYPE_CHECKING: + from vllm.utils.humming import ( + BaseInputSchema, + BaseWeightSchema, + HummingInputSchema, + HummingWeightSchema, + ) + from vllm.utils.humming import dtypes as humming_dtypes + + +def humming_update_schema_hadamard_block_size( + weight_schema: "HummingWeightSchema", + input_schema: "HummingInputSchema", + shape_k: int, +) -> "HummingWeightSchema": + assert shape_k > 0 + block_size = 32 + weight_group_size = weight_schema.weight_scale_group_size + input_group_size = input_schema.input_scale_group_size + for group_size in (weight_group_size, input_group_size): + if group_size is not None and group_size > 0: + block_size = min(block_size, group_size) + block_size = 1 << (block_size.bit_length() - 1) + while shape_k % block_size: + block_size //= 2 + return dataclasses.replace(weight_schema, hadamard_block_size=block_size) + + +if has_humming(): + from vllm.utils.humming import dtypes as humming_dtypes + + _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = { + humming_dtypes.float4e0m3: FP4_DTYPE, + humming_dtypes.float4e2m1: FP4_DTYPE, + humming_dtypes.float6e2m3: scalar_types.float6_e2m3f, + humming_dtypes.float6e3m2: scalar_types.float6_e3m2f, + humming_dtypes.float8e3m4: FP8_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float8e5m2: torch.float8_e5m2, + humming_dtypes.int8: torch.int8, + humming_dtypes.uint4: INT4_DTYPE, + humming_dtypes.uint8: INT8_DTYPE, + humming_dtypes.uint2: torch.uint8, + humming_dtypes.uint3: torch.uint8, + humming_dtypes.uint5: torch.uint8, + humming_dtypes.uint6: torch.uint8, + humming_dtypes.uint7: torch.uint8, + } + + _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = { + humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE, + humming_dtypes.float8e4m3: FP8_DTYPE, + humming_dtypes.float16: torch.float16, + humming_dtypes.bfloat16: torch.bfloat16, + humming_dtypes.float32: torch.float32, + } + + +logger = init_logger(__name__) + + +def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape: + """Map humming group sizes to QuantKey GroupShape. + + group_size: elements per group along K (col); 0 means full dimension. + group_size_n: elements per group along N (row); 0 means 1 (per-row). + + GroupShape convention: row = N dim, col = K dim. + """ + if group_size == 0 and group_size_n == 0: + return GroupShape.PER_CHANNEL + + row = group_size_n if group_size_n > 0 else 1 + col = group_size if group_size > 0 else -1 + return GroupShape(row=row, col=col) + + +def weight_schema_to_quant_key( + schema: "BaseWeightSchema", + param_dtype: torch.dtype | None = None, +) -> QuantKey: + from vllm.utils.humming import WeightScale2Type, WeightScaleType + + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in [torch.float16, torch.bfloat16]: + param_dtype = torch.bfloat16 + if not current_platform.has_device_capability(80): + param_dtype = torch.float16 + + schema = schema.to_humming_schema(param_dtype) + dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype] + + if schema.bs_dtype is not None: + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype] + else: + scale_dtype = torch.float32 + + group_shape = _group_shape( + schema.weight_scale_group_size, + schema.weight_scale_group_size_n, + ) + if schema.weight_scale_type == WeightScaleType.TENSOR: + group_shape = GroupShape.PER_TENSOR + + scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape) + + scale2 = None + if schema.weight_scale_2_type == WeightScale2Type.TENSOR: + scale2 = ScaleDesc( + dtype=torch.float32, + static=True, + group_shape=GroupShape.PER_TENSOR, + ) + + return QuantKey( + dtype=dtype, + scale=scale, + scale2=scale2, + symmetric=not schema.has_zero_point, + ) + + +def input_schema_to_quant_key( + schema: "BaseInputSchema", + param_dtype: torch.dtype | None = None, +) -> QuantKey | None: + from vllm.utils.humming import InputQuantizationMode + + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in [torch.float16, torch.bfloat16]: + param_dtype = torch.bfloat16 + if not current_platform.has_device_capability(80): + param_dtype = torch.float16 + + schema = schema.to_humming_schema(param_dtype) + if schema.a_dtype is None or schema.a_dtype.num_bits >= 16: + return None + + mode = schema.input_quant_mode + if mode == InputQuantizationMode.Disabled: + return None + dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype] + + gs = schema.input_scale_group_size + group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN + + # Pick the scale dtype the Humming kernel actually consumes. An explicit + # input_scale_dtype always wins. Otherwise infer from the grouping: MX + # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8) + # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a + # float32 scale. Getting this right lets a grouped FP8 activation match + # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key. + if schema.input_scale_dtype is not None: + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype] + elif gs == 16: + scale_dtype = FP8_DTYPE + elif gs == 32: + scale_dtype = MXFP_SCALE_DTYPE + else: + scale_dtype = torch.float32 + + if mode == InputQuantizationMode.StaticTensor: + group_shape = GroupShape.PER_TENSOR + scale = ScaleDesc( + dtype=scale_dtype, + static=mode == InputQuantizationMode.StaticTensor, + group_shape=group_shape, + ) + scale2 = None + has_scale2 = ( + InputQuantizationMode.StaticTensorDynamicGroup, + InputQuantizationMode.DynamicGroupToken, + ) + if mode in has_scale2: + static = mode == InputQuantizationMode.StaticTensorDynamicGroup + scale2 = ScaleDesc( + torch.float32, + static, + GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN, + ) + + return QuantKey(dtype=dtype, scale=scale, scale2=scale2, symmetric=True) + + +def quant_key_to_input_schema(key: QuantKey | None) -> "HummingInputSchema": + from vllm.utils.humming import HummingInputSchema, InputQuantizationMode + + if key is None: + return HummingInputSchema(input_quant_mode=InputQuantizationMode.Disabled) + + if not key.symmetric: + raise ValueError("Humming input quantization must be symmetric") + + quant_dtypes = {value: dtype for dtype, value in _HUMMING_TO_QUANT_DTYPE.items()} + quant_dtypes[FP4_DTYPE] = humming_dtypes.float4e2m1 + quant_dtypes[FP8_DTYPE] = humming_dtypes.float8e4m3 + scale_dtypes = {value: dtype for dtype, value in _HUMMING_TO_SCALE_DTYPE.items()} + if key.dtype not in quant_dtypes or key.scale.dtype not in scale_dtypes: + raise ValueError(f"Unsupported Humming input or scale dtype: {key}") + + scale = key.scale + group_size = 0 + if scale.static and scale.group_shape.is_per_tensor() and key.scale2 is None: + mode = InputQuantizationMode.StaticTensor + elif not scale.static and scale.group_shape.is_per_token() and key.scale2 is None: + mode = InputQuantizationMode.DynamicToken + elif not scale.static and scale.group_shape.is_per_group(): + group_size = scale.group_shape.col + scale2 = key.scale2 + if scale2 is None: + mode = InputQuantizationMode.DynamicGroup + elif scale2.dtype != torch.float32: + raise ValueError("Humming secondary input scales must be float32") + elif scale2.static and scale2.group_shape.is_per_tensor(): + mode = InputQuantizationMode.StaticTensorDynamicGroup + elif not scale2.static and scale2.group_shape.is_per_token(): + mode = InputQuantizationMode.DynamicGroupToken + else: + raise ValueError(f"Unsupported Humming secondary input scale: {scale2}") + else: + raise ValueError(f"Unsupported Humming input scale: {key}") + + if group_size == 0 and scale.dtype != torch.float32: + raise ValueError("Humming tensor and token input scales must be float32") + + return HummingInputSchema( + a_dtype=quant_dtypes[key.dtype], + input_scale_group_size=group_size, + input_scale_dtype=scale_dtypes[scale.dtype], + input_quant_mode=mode, + ) + + +def check_and_fallback_input_schema( + weight_schema: "BaseWeightSchema", + input_schema: "BaseInputSchema", + param_dtype: torch.dtype | None = None, + allow_fallback: bool = True, +) -> "HummingInputSchema": + from vllm.utils.humming import HummingInputSchema, InputQuantizationMode + + capability = current_platform.get_device_capability() + assert capability is not None + sm_version = capability.to_int() + if param_dtype is None: + param_dtype = torch.get_default_dtype() + if param_dtype not in (torch.float16, torch.bfloat16): + param_dtype = torch.bfloat16 if sm_version >= 80 else torch.float16 + assert param_dtype in (torch.float16, torch.bfloat16) + + dtype_fallback_order_map = { + humming_dtypes.float8e5m2: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float8e4m3: ( + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float8e3m4: ( + humming_dtypes.float8e4m3, + humming_dtypes.int8, + ), + humming_dtypes.float4e2m1: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.float4e0m3: ( + humming_dtypes.int4, + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + humming_dtypes.int8, + ), + humming_dtypes.int8: ( + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + ), + humming_dtypes.int4: ( + humming_dtypes.int8, + humming_dtypes.float8e4m3, + humming_dtypes.float8e3m4, + ), + } + + weight_schema = weight_schema.to_humming_schema(param_dtype) + input_schema = input_schema.to_humming_schema(param_dtype) + fp6_dtypes = (humming_dtypes.float6e2m3, humming_dtypes.float6e3m2) + if input_schema.a_dtype in fp6_dtypes: + input_schema.a_dtype = humming_dtypes.float8e4m3 + a_dtype = input_schema.a_dtype + input_bits = a_dtype.num_bits if a_dtype is not None else 16 + weight_group_size = weight_schema.weight_scale_group_size + input_group_size = input_schema.input_scale_group_size + input_scale_dtype = input_schema.input_scale_dtype + + def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool: + if dtype is None or dtype.num_bits == 16: + return False + is_int4_deprecated = dtype == humming_dtypes.int4 and sm_version >= 90 + is_int8_deprecated = dtype == humming_dtypes.int8 and 103 <= sm_version < 110 + # Humming has not yet implemented high-performance FP8 computation + # on SM10X and SM11X. Therefore, it is temporarily considered deprecated. + # TODO: removed this after Humming has high-performance FP8 computation + is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120 + return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated + + if input_schema.is_compatible_with(weight_schema, param_dtype): + if not allow_fallback: + if is_deprecated(a_dtype): + msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now." + logger.warning_once(msg) + return input_schema + + if a_dtype is None or a_dtype.num_bits == 16: + return input_schema + + if not is_deprecated(a_dtype): + is_mxfp8 = ( + a_dtype.is_floating_point_type + and a_dtype.num_bits == 8 + and input_group_size == 32 + and input_scale_dtype == humming_dtypes.float8e8m0 + and sm_version >= 120 + ) + is_groupwise = input_group_size > 0 or weight_group_size > 0 + if input_bits == 8 and is_groupwise and not is_mxfp8: + # prefer wna16 for groupwise weight or groupwise input + return HummingInputSchema() + + return input_schema + + if allow_fallback: + fallback_dtypes = ( + a_dtype, + *dtype_fallback_order_map.get(a_dtype, ()), + humming_dtypes.DataType.from_any(param_dtype), + ) + for dtype in fallback_dtypes: + if dtype is None or is_deprecated(dtype): + continue + + group_size = 0 + scale_dtype = humming_dtypes.float32 + quant_mode = InputQuantizationMode.DynamicToken + if dtype.num_bits == 16: + scale_dtype = None + quant_mode = InputQuantizationMode.Disabled + elif dtype in (humming_dtypes.float8e4m3, humming_dtypes.float8e3m4): + is_channelwise_weight = weight_schema.weight_scale_group_size == 0 + weight_group_size = weight_schema.weight_scale_group_size + is_e8m0_scale = weight_schema.bs_dtype == humming_dtypes.float8e8m0 + is_mx_weight = weight_group_size == 32 and is_e8m0_scale + if sm_version >= 120 and (is_channelwise_weight or is_mx_weight): + group_size = 32 + scale_dtype = humming_dtypes.float8e8m0 + quant_mode = InputQuantizationMode.DynamicGroup + + candidate = HummingInputSchema( + a_dtype=dtype, + input_scale_group_size=group_size, + input_scale_dtype=scale_dtype, + input_quant_mode=quant_mode, + ) + if candidate.is_compatible_with(weight_schema, param_dtype): + return candidate + + raise ValueError( + f"No compatible Humming input schema for {input_schema} with " + f"weight schema {weight_schema} on SM{sm_version} " + f"({allow_fallback=})" + ) + + +def resolve_humming_layer_config(config: dict[str, Any], prefix: str) -> dict[str, Any]: + layer_config = config.copy() + dynamic = config.get("dynamic", {}) + if isinstance(dynamic, dict): + for pattern, overrides in dynamic.items(): + if pattern.startswith("+") and re.match(pattern[2:], prefix): + layer_config.update(overrides) + break + return layer_config + + +def humming_is_layer_skipped(config: dict[str, Any], prefix: str): + if not config: + return True + + keys = ["ignored_layers", "ignore", "modules_to_not_convert"] + ignored_layers: list[str] = [] + for key in keys: + candidate = config.get(key, []) or [] + if candidate: + ignored_layers = candidate + break + + if any(module_name in prefix for module_name in ignored_layers): + return True + if "lm_head" in prefix: + return True + + for regex in config.get("dynamic", {}): + if regex[:1] != "-": + continue + if re.match(regex[2:], prefix): + return True + + return False diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py index 014603ff39..574a00f16b 100644 --- a/vllm/utils/humming.py +++ b/vllm/utils/humming.py @@ -4,12 +4,56 @@ vLLM code should import humming symbols from here so that ``import humming`` (which has import-time side effects) is deferred until first use. Add new -symbols by appending one entry to ``_EXPORTS`` as ``"module.path:attr"``, -or ``"module.path"`` for a whole-module re-export. +symbols to both the ``TYPE_CHECKING`` imports and ``_EXPORTS``, using +``"module.path:attr"`` or ``"module.path"`` for a whole-module re-export. """ import importlib -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from humming import dtypes as dtypes + from humming.config import GemmType as GemmType + from humming.config import InputQuantizationMode as InputQuantizationMode + from humming.config import LayerConfig as LayerConfig + from humming.config import MmaType as MmaType + from humming.config import WeightScale2Type as WeightScale2Type + from humming.config import WeightScaleType as WeightScaleType + from humming.dtypes import DataType as DataType + from humming.forward import humming_forward as humming_forward + from humming.forward import may_process_input as may_process_input + from humming.forward import may_quant_input as may_quant_input + from humming.ops import process_input as process_input + from humming.schema import AWQWeightSchema as AWQWeightSchema + from humming.schema import BaseInputSchema as BaseInputSchema + from humming.schema import BaseWeightSchema as BaseWeightSchema + from humming.schema import BitnetWeightSchema as BitnetWeightSchema + from humming.schema import ( + CompressedTensorsInputSchema as CompressedTensorsInputSchema, + ) + from humming.schema import ( + CompressedTensorsWeightSchema as CompressedTensorsWeightSchema, + ) + from humming.schema import Fp8InputSchema as Fp8InputSchema + from humming.schema import GptOssMxfp4WeightSchema as GptOssMxfp4WeightSchema + from humming.schema import GPTQWeightSchema as GPTQWeightSchema + from humming.schema import HummingInputSchema as HummingInputSchema + from humming.schema import HummingWeightSchema as HummingWeightSchema + from humming.schema import Mxfp4WeightSchema as Mxfp4WeightSchema + from humming.schema.fp8 import Fp8WeightSchema as Fp8WeightSchema + from humming.schema.modelopt import ( + ModeloptMxfp8WeightSchema as ModeloptMxfp8WeightSchema, + ) + from humming.schema.modelopt import ( + ModeloptNvfp4InputSchema as ModeloptNvfp4InputSchema, + ) + from humming.schema.modelopt import ( + ModeloptNvfp4WeightSchema as ModeloptNvfp4WeightSchema, + ) + from humming.transform import prepare_layer_config as prepare_layer_config + from humming.transform import transform_humming_tensors as transform_humming_tensors + from humming.tune import get_heuristics_config as get_heuristics_config + from humming.utils.weight import quantize_weight as quantize_weight _EXPORTS: dict[str, str] = { "dtypes": "humming.dtypes", @@ -17,7 +61,12 @@ _EXPORTS: dict[str, str] = { "GemmType": "humming.config:GemmType", "LayerConfig": "humming.config:LayerConfig", "WeightScaleType": "humming.config:WeightScaleType", + "WeightScale2Type": "humming.config:WeightScale2Type", + "InputQuantizationMode": "humming.config:InputQuantizationMode", + "MmaType": "humming.config:MmaType", "humming_forward": "humming.forward:humming_forward", + "may_process_input": "humming.forward:may_process_input", + "process_input": "humming.ops:process_input", "may_quant_input": "humming.forward:may_quant_input", "prepare_layer_config": "humming.transform:prepare_layer_config", "transform_humming_tensors": "humming.transform:transform_humming_tensors",