iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Feature][Humming] Humming feature integration (#56685)

Browse Source 
Signed-off-by: jinzhen.ljz <jinzhen.ljz@antgroup.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Jinzhen Lin
2026-09-20 08:11:01 -07:00
committed by GitHub
co-authored by Codex mgoin
parent 27757dde02
commit 7d99c2c4fd
62 changed files with 2040 additions and 1086 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/lm_eval.yaml
+63 -23
View File
@@ -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/
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+47 -19
View File
@@ -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<int32_t>(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<int32_t>(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<scalar_t, fill_threads>(
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 <typename scalar_t>
@@ -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<torch::stable::Tensor> 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<torch::stable::Tensor> maybe_expert_map,
std::optional<torch::stable::Tensor> 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<int32_t*>(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<int32_t*>(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<int32_t*>(cumsum_buffer.mutable_data_ptr()),
reinterpret_cast<int32_t*>(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);
}
});
}
csrc/libtorch_stable/moe/moe_ops.h
+7 -5
View File
@@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> maybe_expert_map,
std::optional<torch::stable::Tensor> scatter_idx);
void batched_moe_align_block_size(
int64_t max_tokens_per_batch, int64_t block_size,
csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
+75 -18
View File
@@ -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<torch::stable::Tensor>& 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<torch::stable::Tensor>& maybe_sort_workspace,
const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
const std::optional<torch::stable::Tensor>& 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<int>(expert_map.value());
valid_num_ptr =
get_ptr<int64_t>(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<int>(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<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream,
build_inverse ? get_ptr<int>(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<torch::stable::Tensor>& 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<torch::stable::Tensor>& maybe_sort_workspace,
const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
const std::optional<torch::stable::Tensor>& 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<int64_t>(expert_first_token_offset) + n_local_expert
: nullptr;
MOE_DISPATCH(input.scalar_type(), [&] {
expandInputRowsKernelLauncher<scalar_t>(
@@ -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<torch::stable::Tensor>& 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<torch::stable::Tensor>& 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));
}
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+30 -24
View File
@@ -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<<<blocks, threads, 0, stream>>>(
sorted_indices, total_indices, num_experts, expert_first_token_offset);
computeExpertOffsetsAndInverseKernel<<<blocks, threads, 0, stream>>>(
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<int64_t>(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();
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+7 -6
View File
@@ -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 <typename T>
void expandInputRowsKernelLauncher(
csrc/libtorch_stable/moe/torch_bindings.cpp
+9 -1
View File
@@ -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, "
requirements/cuda.txt
+1 -1
View File
@@ -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
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml
+2
View File
@@ -9,3 +9,5 @@ server_args: >-
--moe-backend humming
--linear-backend humming
--dtype bfloat16
env:
VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
+14
View File
@@ -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"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
+14
View File
@@ -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"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml
+2
View File
@@ -7,3 +7,5 @@ server_args: >-
--max-model-len 8192
--quantization humming
--linear-backend humming
env:
VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml
+2
View File
@@ -7,3 +7,5 @@ server_args: >-
--max-model-len 8192
--moe-backend humming
--linear-backend humming
env:
VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
-13
View File
@@ -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"}'
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml → tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
+4 -3
View File
@@ -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: " <think>\n\n</think>\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"}'
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming.yaml → tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
+4 -3
View File
@@ -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: " <think>\n\n</think>\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}'
tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml
+13
View File
@@ -0,0 +1,13 @@
model_name: "Qwen/Qwen3.5-9B"
accuracy_threshold: 0.88
num_questions: 1319
num_fewshot: 5
gen_prefix: " <think>\n\n</think>\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"}'
tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml
+12
View File
@@ -0,0 +1,12 @@
model_name: "Qwen/Qwen3.5-9B"
accuracy_threshold: 0.88
num_questions: 1319
num_fewshot: 5
gen_prefix: " <think>\n\n</think>\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}'
tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt
+3
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt
+4
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
+1
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-act-int8.txt
+1
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-deepep.txt
+2
View File
@@ -0,0 +1,2 @@
Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
+1
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt
-3
View File
@@ -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
tests/evals/gsm8k/configs/humming/config-l4.txt
+9
View File
@@ -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
tests/kernels/core/test_activation.py
+109 -42
View File
@@ -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),
)
tests/kernels/moe/test_moe.py
+75 -26
View File
@@ -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
tests/kernels/moe/test_moe_align_block_size.py
+40
View File
@@ -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,
tests/kernels/moe/test_moe_permute_unpermute.py
+26 -1
View File
@@ -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)
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+1 -2
View File
@@ -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():
tests/quantization/test_humming_mxfp4_block_fp8.py
+12 -8
View File
@@ -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
tests/quantization/test_moe_wna16.py
+1 -1
View File
@@ -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
vllm/_custom_ops.py
+2
View File
@@ -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,
)
vllm/envs.py
+2 -1
View File
@@ -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)
),
vllm/model_executor/kernels/linear/__init__.py
+4
View File
@@ -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",
vllm/model_executor/kernels/linear/mixed_precision/humming.py
+3 -2
View File
@@ -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,
vllm/model_executor/kernels/linear/mxfp4/humming.py
+11 -14
View File
@@ -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,
vllm/model_executor/kernels/linear/mxfp6/base.py
+2 -2
View File
@@ -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
vllm/model_executor/kernels/linear/mxfp6/humming.py
+84
View File
@@ -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,
)
vllm/model_executor/kernels/linear/mxfp8/humming.py
+8 -2
View File
@@ -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,
vllm/model_executor/kernels/linear/nvfp4/humming.py
+15 -5
View File
@@ -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,
vllm/model_executor/kernels/linear/scaled_mm/humming.py
+3 -1
View File
@@ -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,
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+183 -251
View File
@@ -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,
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+67 -1
View File
@@ -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
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+31
View File
@@ -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,
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+2 -2
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/int8.py
+2 -2
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+1 -1
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+3 -3
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+2 -2
View File
@@ -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,
)
vllm/model_executor/layers/quantization/auto_awq.py
+1 -1
View File
@@ -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,
)
vllm/model_executor/layers/quantization/auto_gptq.py
+1 -1
View File
@@ -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,
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+1 -1
View File
@@ -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,
)
vllm/model_executor/layers/quantization/humming.py
+106 -52
View File
@@ -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
vllm/model_executor/layers/quantization/moe_wna16.py
+1 -1
View File
@@ -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,
)
vllm/model_executor/layers/quantization/utils/humming/__init__.py
+51
View File
@@ -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",
]
vllm/model_executor/layers/quantization/utils/humming/activation.py
+126
View File
@@ -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
vllm/model_executor/layers/quantization/utils/humming/linear.py
+183
View File
@@ -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))
vllm/model_executor/layers/quantization/utils/humming_utils.py → vllm/model_executor/layers/quantization/utils/humming/moe.py
+65 -539
View File
@@ -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,
)
)
vllm/model_executor/layers/quantization/utils/humming/schema.py
+437
View File
@@ -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
vllm/utils/humming.py
+52 -3
View File
@@ -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",