[Feature][Humming] Humming feature integration (#56685)
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>
Jinzhen Lin committed
Sep 20, 2026 at 23:11 UTC
7d99c2c4fd61cdeebd29cc7b60584bad10d1ef99
62 files changed
+2038
-1084
.buildkite/test_areas/lm_eval.yaml
+63
-23
@@ -347,47 +347,67 @@ steps:
347
commands:
348
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
349
350
-- label: ":nvidia: (A100) LM Eval Humming FP16 TEMPORARY Shard %N"
351
- key: lm-eval-humming-f16-a100
350
+- label: ":nvidia: (L4) LM Eval Humming TEMPORARY"
351
+ key: lm-eval-humming-l4
352
+ optional: true
353
+ timeout_in_minutes: 90
354
+ device: l4
355
+ num_devices: 1
356
+ source_file_dependencies:
357
+ - .buildkite/test_areas/lm_eval.yaml
358
+ - tests/evals/gsm8k/configs/humming/
359
+ - requirements/cuda.txt
360
+ - vllm/utils/humming.py
361
+ - vllm/model_executor/layers/quantization/humming.py
362
+ - vllm/model_executor/layers/quantization/utils/humming/
363
+ - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
364
+ - vllm/model_executor/layers/fused_moe/oracle/
365
+ - vllm/model_executor/kernels/linear/
366
+ commands:
367
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt
368
+
369
+- label: ":nvidia: (A100) LM Eval Humming Default TEMPORARY Shard %N"
370
+ key: lm-eval-humming-a100-default
371
+ optional: true
372
timeout_in_minutes: 75
373
device: a100
354
- optional: true
374
num_devices: 1
375
parallelism: 3
376
source_file_dependencies:
377
- vllm/model_executor/layers/quantization/humming.py
359
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
378
+ - vllm/model_executor/layers/quantization/utils/humming/
379
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
380
- vllm/model_executor/layers/fused_moe/oracle/
381
- vllm/model_executor/kernels/linear/
382
commands:
383
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt
384
366
-- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY"
367
- key: lm-eval-humming-act-a100
385
+- label: ":nvidia: (A100) LM Eval Humming Activation INT8 TEMPORARY Shard %N"
386
+ key: lm-eval-humming-a100-act
387
+ optional: true
388
timeout_in_minutes: 45
389
device: a100
370
- optional: true
390
num_devices: 1
391
+ parallelism: 2
392
source_file_dependencies:
393
- vllm/model_executor/layers/quantization/humming.py
374
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
394
+ - vllm/model_executor/layers/quantization/utils/humming/
395
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
396
- vllm/model_executor/layers/fused_moe/oracle/
397
- vllm/model_executor/kernels/linear/
398
commands:
379
- - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
399
+ - 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
400
381
-- label: ":nvidia: (H100) LM Eval Humming FP16 TEMPORARY Shard %N"
382
- key: lm-eval-humming-f16-h100
401
+- label: ":nvidia: (H100) LM Eval Humming Default TEMPORARY Shard %N"
402
+ key: lm-eval-humming-h100-default
403
+ optional: true
404
timeout_in_minutes: 70
405
device: h100
385
- optional: true
406
num_devices: 1
407
parallelism: 3
408
source_file_dependencies:
409
- vllm/model_executor/layers/quantization/humming.py
390
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
410
+ - vllm/model_executor/layers/quantization/utils/humming/
411
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
412
- vllm/model_executor/layers/fused_moe/oracle/
413
- vllm/model_executor/kernels/linear/
@@ -395,14 +415,14 @@ steps:
415
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt
416
417
- label: ":nvidia: (H100) LM Eval Humming Activation FP8/INT8 TEMPORARY"
398
- key: lm-eval-humming-act-h100
418
+ key: lm-eval-humming-h100-act
419
+ optional: true
420
timeout_in_minutes: 70
421
device: h100
401
- optional: true
422
num_devices: 1
423
source_file_dependencies:
424
- vllm/model_executor/layers/quantization/humming.py
405
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
425
+ - vllm/model_executor/layers/quantization/utils/humming/
426
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
427
- vllm/model_executor/layers/fused_moe/oracle/
428
- vllm/model_executor/kernels/linear/
@@ -410,15 +430,35 @@ steps:
430
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
431
# - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
432
413
-- label: ":nvidia: (B200) LM Eval Humming FP16 TEMPORARY"
414
- key: lm-eval-humming-f16-b200
433
+- label: ":nvidia: (H100) LM Eval Humming DeepEP TEMPORARY"
434
+ key: lm-eval-humming-h100-deepep
435
+ optional: true
436
+ timeout_in_minutes: 60
437
+ device: h100
438
+ num_devices: 2
439
+ source_file_dependencies:
440
+ - .buildkite/test_areas/lm_eval.yaml
441
+ - tests/evals/gsm8k/configs/humming/
442
+ - requirements/cuda.txt
443
+ - vllm/utils/humming.py
444
+ - vllm/model_executor/layers/quantization/humming.py
445
+ - vllm/model_executor/layers/quantization/utils/humming/
446
+ - vllm/model_executor/layers/fused_moe/
447
+ - vllm/model_executor/kernels/linear/
448
+ - vllm/distributed/
449
+ - csrc/moe/
450
+ commands:
451
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-deepep.txt
452
+
453
+- label: ":nvidia: (B200) LM Eval Humming Default TEMPORARY"
454
+ key: lm-eval-humming-b200-default
455
+ optional: true
456
timeout_in_minutes: 50
457
device: b200-k8s
417
- optional: true
458
num_devices: 1
459
source_file_dependencies:
460
- vllm/model_executor/layers/quantization/humming.py
421
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
461
+ - vllm/model_executor/layers/quantization/utils/humming/
462
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
463
- vllm/model_executor/layers/fused_moe/oracle/
464
- vllm/model_executor/kernels/linear/
@@ -426,14 +466,14 @@ steps:
466
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
467
468
- label: ":nvidia: (B200) LM Eval Humming Activation FP8/INT8 TEMPORARY"
429
- key: lm-eval-humming-act-b200
469
+ key: lm-eval-humming-b200-act
470
+ optional: true
471
timeout_in_minutes: 50
472
device: b200-k8s
432
- optional: true
473
num_devices: 1
474
source_file_dependencies:
475
- vllm/model_executor/layers/quantization/humming.py
436
- - vllm/model_executor/layers/quantization/utils/humming_utils.py
476
+ - vllm/model_executor/layers/quantization/utils/humming/
477
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
478
- vllm/model_executor/layers/fused_moe/oracle/
479
- vllm/model_executor/kernels/linear/
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+47
-19
@@ -236,7 +236,7 @@ __device__ void _moe_align_block_size_small_batch_expert(
236
int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
237
size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks,
238
int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num,
239
- int32_t* token_mask, bool has_expert_map) {
239
+ int32_t* token_mask, bool has_expert_map, int32_t* scatter_idx = nullptr) {
240
// Compute input buffer offsets. Typically these will all be 0, except when
241
// using Multi LoRA.
242
int sorted_token_ids_offset = max_num_tokens_padded * model_offset;
@@ -318,9 +318,14 @@ __device__ void _moe_align_block_size_small_batch_expert(
318
}
319
320
for (size_t i = tid; i < numel; i += stride) {
321
- if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
322
- num_experts, has_expert_map);
323
- expert_id != -1) {
321
+ int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
322
+ has_expert_map);
323
+ bool valid =
324
+ expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
325
+ if (scatter_idx != nullptr) {
326
+ scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
327
+ }
328
+ if (expert_id != -1) {
329
int32_t rank_post_pad =
330
tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
331
@@ -338,14 +343,20 @@ __device__ void _count_and_sort_expert_tokens(
343
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
344
int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
345
int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask,
341
- int32_t model_offset, int32_t topk_num, bool has_expert_map) {
346
+ int32_t model_offset, int32_t topk_num, bool has_expert_map,
347
+ int32_t* scatter_idx = nullptr) {
348
const size_t tid = blockIdx.y * blockDim.x + threadIdx.x;
349
const size_t stride = blockDim.x * gridDim.y;
350
351
for (size_t i = tid; i < numel; i += stride) {
346
- if (int expert_id = get_local_expert_id(i, topk_ids, expert_map,
347
- num_experts, has_expert_map);
348
- expert_id != -1) {
352
+ int expert_id = get_local_expert_id(i, topk_ids, expert_map, num_experts,
353
+ has_expert_map);
354
+ bool valid =
355
+ expert_id != -1 && (token_mask == nullptr || token_mask[i / topk_num]);
356
+ if (scatter_idx != nullptr) {
357
+ scatter_idx[i] = valid ? static_cast<int32_t>(i) : -1;
358
+ }
359
+ if (expert_id != -1) {
360
if (token_mask == nullptr || token_mask[i / topk_num]) {
361
int32_t rank_post_pad = atomicAdd(
362
&cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
@@ -377,10 +388,11 @@ __global__ void count_and_sort_expert_tokens_kernel(
388
const scalar_t* __restrict__ topk_ids,
389
int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer,
390
int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts,
380
- int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) {
391
+ int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map,
392
+ int32_t* scatter_idx) {
393
_count_and_sort_expert_tokens(
394
topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts,
383
- max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
395
+ max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map, scatter_idx);
396
}
397
398
// 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(
545
int32_t* __restrict__ total_tokens_post_pad,
546
int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size,
547
size_t numel, int32_t max_num_tokens_padded, int32_t topk_num,
536
- bool has_expert_map) {
548
+ bool has_expert_map, int32_t* scatter_idx) {
549
_moe_align_block_size_small_batch_expert<scalar_t, fill_threads>(
550
topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map,
551
num_experts, block_size, numel, max_num_tokens_padded,
552
CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr,
541
- has_expert_map);
553
+ has_expert_map, scatter_idx);
554
}
555
556
template <typename scalar_t>
@@ -658,16 +670,31 @@ __global__ void moe_lora_align_block_size_small_batch_expert_kernel(
670
671
// taken from
672
// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc
661
-void moe_align_block_size(
662
- torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
663
- torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
664
- torch::stable::Tensor num_tokens_post_pad,
665
- std::optional<torch::stable::Tensor> maybe_expert_map) {
673
+void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
674
+ int64_t block_size,
675
+ torch::stable::Tensor sorted_token_ids,
676
+ torch::stable::Tensor experts_ids,
677
+ torch::stable::Tensor num_tokens_post_pad,
678
+ std::optional<torch::stable::Tensor> maybe_expert_map,
679
+ std::optional<torch::stable::Tensor> scatter_idx) {
680
const torch::stable::accelerator::DeviceGuard device_guard(
681
topk_ids.get_device_index());
682
const cudaStream_t stream =
683
get_current_cuda_stream(topk_ids.get_device_index());
684
685
+ int32_t* scatter_ptr = nullptr;
686
+ if (scatter_idx.has_value()) {
687
+ STD_TORCH_CHECK(
688
+ scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int,
689
+ "scatter_idx must be int32");
690
+ STD_TORCH_CHECK(scatter_idx->device() == topk_ids.device() &&
691
+ scatter_idx->is_contiguous() &&
692
+ scatter_idx->numel() == topk_ids.numel(),
693
+ "scatter_idx must be contiguous, on the same device, and "
694
+ "match topk_ids.numel()");
695
+ scatter_ptr = reinterpret_cast<int32_t*>(scatter_idx->mutable_data_ptr());
696
+ }
697
+
698
int64_t padded_num_experts =
699
((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
700
int experts_per_warp = WARP_SIZE;
@@ -713,7 +740,8 @@ void moe_align_block_size(
740
num_tokens_post_pad.mutable_data_ptr()),
741
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
742
num_experts, block_size, topk_ids.numel(),
716
- sorted_token_ids.size(0), topk_ids.size(1), has_expert_map);
743
+ sorted_token_ids.size(0), topk_ids.size(1), has_expert_map,
744
+ scatter_ptr);
745
} else {
746
torch::stable::Tensor cumsum_buffer = torch::stable::new_empty(
747
topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
@@ -753,7 +781,7 @@ void moe_align_block_size(
781
reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
782
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
783
topk_ids.numel(), num_experts, sorted_token_ids.size(0),
756
- topk_ids.size(1), has_expert_map);
784
+ topk_ids.size(1), has_expert_map, scatter_ptr);
785
}
786
});
787
}
csrc/libtorch_stable/moe/moe_ops.h
+7
-5
@@ -36,11 +36,13 @@ void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
36
std::optional<torch::stable::Tensor> topk_ids,
37
std::optional<torch::stable::Tensor> expert_map);
38
39
-void moe_align_block_size(
40
- torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
41
- torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids,
42
- torch::stable::Tensor num_tokens_post_pad,
43
- std::optional<torch::stable::Tensor> maybe_expert_map);
39
+void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
40
+ int64_t block_size,
41
+ torch::stable::Tensor sorted_token_ids,
42
+ torch::stable::Tensor experts_ids,
43
+ torch::stable::Tensor num_tokens_post_pad,
44
+ std::optional<torch::stable::Tensor> maybe_expert_map,
45
+ std::optional<torch::stable::Tensor> scatter_idx);
46
47
void batched_moe_align_block_size(
48
int64_t max_tokens_per_batch, int64_t block_size,
csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
+75
-18
@@ -56,20 +56,18 @@ int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
56
CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
57
}
58
59
-void moe_permute_impl(
60
- const torch::stable::Tensor& input, // [n_token, hidden]
61
- const torch::stable::Tensor& topk_ids, // [n_token, topk]
62
- const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
59
+torch::stable::Tensor moe_sort_routing(
60
+ const torch::stable::Tensor& topk_ids, // [n_token, topk]
61
+ const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
62
const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
63
int64_t n_expert, int64_t n_local_expert, int64_t topk,
65
- torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
64
torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
65
torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
68
- torch::stable::Tensor& permuted_idx, // [permute_size]
66
const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
67
const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
68
const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
72
- const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
69
+ const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort,
70
+ bool build_inverse) {
71
STD_TORCH_CHECK(expert_first_token_offset.scalar_type() ==
72
torch::headeronly::ScalarType::Long,
73
"expert_first_token_offset must be int64");
@@ -88,12 +86,11 @@ void moe_permute_impl(
86
"token_expert_indices shape must be same as inv_permuted_idx");
87
88
const torch::stable::accelerator::DeviceGuard device_guard(
91
- input.get_device_index());
92
- auto device = input.device();
93
- auto n_token = input.sizes()[0];
94
- auto n_hidden = input.sizes()[1];
89
+ topk_ids.get_device_index());
90
+ auto device = topk_ids.device();
91
+ auto n_token = topk_ids.sizes()[0];
92
auto expanded_rows = n_token * topk;
96
- auto stream = get_current_cuda_stream(input.get_device_index());
93
+ auto stream = get_current_cuda_stream(topk_ids.get_device_index());
94
95
auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
96
auto sort_workspace = maybe_allocate_tensor(
@@ -107,13 +104,11 @@ void moe_permute_impl(
104
torch::headeronly::ScalarType::Int, device, "sorted_row_idx");
105
106
CubKeyValueSorter sorter{};
110
- int64_t* valid_num_ptr = nullptr;
107
torch::stable::Tensor topk_ids_for_sort = topk_ids;
108
113
- if (expert_map.has_value()) {
114
- const int* expert_map_ptr = get_ptr<int>(expert_map.value());
115
- valid_num_ptr =
116
- get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
109
+ if (expert_map.has_value() || build_inverse) {
110
+ const int* expert_map_ptr =
111
+ expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr;
112
topk_ids_for_sort = maybe_allocate_tensor(
113
maybe_topk_ids_for_sort, topk_ids.sizes(),
114
torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort");
@@ -126,7 +121,39 @@ void moe_permute_impl(
121
get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
122
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
123
get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
129
- n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
124
+ n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream,
125
+ build_inverse ? get_ptr<int>(inv_permuted_idx) : nullptr);
126
+ return sorted_row_idx;
127
+}
128
+
129
+void moe_permute_impl(
130
+ const torch::stable::Tensor& input, // [n_token, hidden]
131
+ const torch::stable::Tensor& topk_ids, // [n_token, topk]
132
+ const torch::stable::Tensor& token_expert_indices, // [n_token, topk]
133
+ const std::optional<torch::stable::Tensor>& expert_map, // [n_expert]
134
+ int64_t n_expert, int64_t n_local_expert, int64_t topk,
135
+ torch::stable::Tensor& permuted_input, // [permuted_size, hidden]
136
+ torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1]
137
+ torch::stable::Tensor& inv_permuted_idx, // [n_token, topk]
138
+ torch::stable::Tensor& permuted_idx, // [permute_size]
139
+ const std::optional<torch::stable::Tensor>& maybe_sort_workspace,
140
+ const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id,
141
+ const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx,
142
+ const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) {
143
+ const torch::stable::accelerator::DeviceGuard device_guard(
144
+ input.get_device_index());
145
+ auto sorted_row_idx = moe_sort_routing(
146
+ topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert,
147
+ topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace,
148
+ maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort,
149
+ false);
150
+ auto stream = get_current_cuda_stream(input.get_device_index());
151
+ auto n_token = input.size(0);
152
+ auto n_hidden = input.size(1);
153
+ auto valid_num_ptr =
154
+ expert_map.has_value()
155
+ ? get_ptr<int64_t>(expert_first_token_offset) + n_local_expert
156
+ : nullptr;
157
158
MOE_DISPATCH(input.scalar_type(), [&] {
159
expandInputRowsKernelLauncher<scalar_t>(
@@ -171,6 +198,22 @@ void moe_permute_with_scratch(
198
topk_ids_for_sort);
199
}
200
201
+void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
202
+ const torch::stable::Tensor& token_expert_indices,
203
+ const std::optional<torch::stable::Tensor>& expert_map,
204
+ int64_t n_expert, int64_t n_local_expert,
205
+ torch::stable::Tensor& expert_first_token_offset,
206
+ torch::stable::Tensor& scatter_idx,
207
+ torch::stable::Tensor& sort_workspace,
208
+ torch::stable::Tensor& sorted_experts,
209
+ torch::stable::Tensor& sorted_rows,
210
+ torch::stable::Tensor& topk_ids_for_sort) {
211
+ moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert,
212
+ n_local_expert, topk_ids.size(1), expert_first_token_offset,
213
+ scatter_idx, sort_workspace, sorted_experts, sorted_rows,
214
+ topk_ids_for_sort, true);
215
+}
216
+
217
void moe_unpermute(
218
const torch::stable::Tensor&
219
permuted_hidden_states, // [n_token * topk, hidden]
@@ -315,6 +358,19 @@ void moe_unpermute(
358
STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0");
359
}
360
361
+void moe_prepare_scatter(const torch::stable::Tensor& topk_ids,
362
+ const torch::stable::Tensor& token_expert_indices,
363
+ const std::optional<torch::stable::Tensor>& expert_map,
364
+ int64_t n_expert, int64_t n_local_expert,
365
+ torch::stable::Tensor& expert_first_token_offset,
366
+ torch::stable::Tensor& scatter_idx,
367
+ torch::stable::Tensor& sort_workspace,
368
+ torch::stable::Tensor& sorted_experts,
369
+ torch::stable::Tensor& sorted_rows,
370
+ torch::stable::Tensor& topk_ids_for_sort) {
371
+ STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0");
372
+}
373
+
374
#endif
375
376
bool moe_permute_unpermute_supported() {
@@ -328,5 +384,6 @@ bool moe_permute_unpermute_supported() {
384
STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
385
m.impl("moe_permute", TORCH_BOX(&moe_permute));
386
m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
387
+ m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter));
388
m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
389
}
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+30
-24
@@ -82,34 +82,40 @@ __device__ inline int64_t findTotalEltsLessThanTarget(T const* sorted_indices,
82
return target_location + 1;
83
}
84
85
-// Calculates the start offset of the tokens for a given expert. The last
86
-// element is the total number of valid tokens
87
-__global__ void computeExpertFirstTokenOffsetKernel(
85
+// Computes expert offsets (ending with the valid row count) and optional
86
+// inverse indices.
87
+__global__ void computeExpertOffsetsAndInverseKernel(
88
int const* sorted_experts, int64_t const sorted_experts_len,
89
- int const num_experts, int64_t* expert_first_token_offset) {
90
- // First, compute the global tid. We only need 1 thread per expert.
91
- int const expert = blockIdx.x * blockDim.x + threadIdx.x;
89
+ int const num_experts, int64_t* expert_first_token_offset,
90
+ int const* sorted_rows, int* inverse) {
91
+ int const index = blockIdx.x * blockDim.x + threadIdx.x;
92
93
// Note that expert goes [0, num_experts] (inclusive) because we want a count
94
// for the total number of active tokens at the end of the scan.
95
- if (expert >= num_experts + 1) {
96
- return;
95
+ if (index < num_experts + 1) {
96
+ expert_first_token_offset[index] =
97
+ findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, index);
98
+ }
99
+ if (inverse != nullptr && index < sorted_experts_len) {
100
+ inverse[sorted_rows[index]] = index;
101
}
98
- expert_first_token_offset[expert] =
99
- findTotalEltsLessThanTarget(sorted_experts, sorted_experts_len, expert);
102
}
103
102
-void computeExpertFirstTokenOffset(int const* sorted_indices,
103
- int const total_indices,
104
- int const num_experts,
105
- int64_t* expert_first_token_offset,
106
- cudaStream_t stream) {
107
- int const num_entries = num_experts + 1;
108
- int const threads = std::min(1024, num_entries);
104
+void computeExpertOffsetsAndInverse(int const* sorted_indices,
105
+ int const total_indices,
106
+ int const num_experts,
107
+ int64_t* expert_first_token_offset,
108
+ int const* sorted_rows, int* inverse,
109
+ cudaStream_t stream) {
110
+ int const num_entries = inverse != nullptr
111
+ ? std::max(num_experts + 1, total_indices)
112
+ : num_experts + 1;
113
+ int const threads = inverse != nullptr ? 256 : std::min(1024, num_entries);
114
int const blocks = (num_entries + threads - 1) / threads;
115
111
- computeExpertFirstTokenOffsetKernel<<<blocks, threads, 0, stream>>>(
112
- sorted_indices, total_indices, num_experts, expert_first_token_offset);
116
+ computeExpertOffsetsAndInverseKernel<<<blocks, threads, 0, stream>>>(
117
+ sorted_indices, total_indices, num_experts, expert_first_token_offset,
118
+ sorted_rows, inverse);
119
}
120
121
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,
123
int64_t* expert_first_token_offset, int num_rows,
124
int num_experts, int num_experts_per_node, int k,
125
CubKeyValueSorter& sorter, void* sorter_ws,
120
- cudaStream_t stream) {
126
+ cudaStream_t stream, int* inverse) {
127
int64_t const expanded_num_rows = static_cast<int64_t>(k) * num_rows;
128
// We need to use the full num_experts because that is the sentinel value used
129
// by topk for disabled experts
@@ -127,9 +133,9 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
133
sorter.run((void*)sorter_ws, sorter_ws_size_bytes, expert_for_source_row,
134
permuted_experts, source_rows, permuted_rows, expanded_num_rows,
135
stream);
130
- computeExpertFirstTokenOffset(permuted_experts, expanded_num_rows,
131
- num_experts_per_node, expert_first_token_offset,
132
- stream);
136
+ computeExpertOffsetsAndInverse(
137
+ permuted_experts, expanded_num_rows, num_experts_per_node,
138
+ expert_first_token_offset, permuted_rows, inverse, stream);
139
}
140
141
__global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
@@ -142,7 +148,7 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
148
extern __shared__ int smem_expert_map[];
149
// store expert_map in smem
150
for (int i = tidx; i < num_experts; i += blockDim.x) {
145
- smem_expert_map[i] = expert_map_ptr[i];
151
+ smem_expert_map[i] = expert_map_ptr ? expert_map_ptr[i] : i;
152
}
153
__syncthreads();
154
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+7
-6
@@ -43,18 +43,19 @@ class CubKeyValueSorter {
43
int num_bits_;
44
};
45
46
-void computeExpertFirstTokenOffset(int const* sorted_indices,
47
- int const total_indices,
48
- int const num_experts,
49
- int64_t* expert_first_token_offset,
50
- cudaStream_t stream);
46
+void computeExpertOffsetsAndInverse(int const* sorted_indices,
47
+ int const total_indices,
48
+ int const num_experts,
49
+ int64_t* expert_first_token_offset,
50
+ int const* sorted_rows, int* inverse,
51
+ cudaStream_t stream);
52
53
void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
54
int* permuted_experts, int* permuted_rows,
55
int64_t* expert_first_token_offset, int num_rows,
56
int num_experts, int num_experts_per_node, int k,
57
CubKeyValueSorter& sorter, void* sorter_ws,
57
- cudaStream_t stream);
58
+ cudaStream_t stream, int* inverse = nullptr);
59
60
template <typename T>
61
void expandInputRowsKernelLauncher(
csrc/libtorch_stable/moe/torch_bindings.cpp
+9
-1
@@ -39,7 +39,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
39
" int block_size, Tensor! sorted_token_ids,"
40
" Tensor! experts_ids,"
41
" Tensor! num_tokens_post_pad,"
42
- " Tensor? maybe_expert_map) -> ()");
42
+ " Tensor? maybe_expert_map,"
43
+ " Tensor(a!)? scatter_idx=None) -> ()");
44
45
// Aligning the number of tokens to be processed by each expert such
46
// that it is divisible by the block size, but for the batched case.
@@ -105,6 +106,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
106
"permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
107
"Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
108
109
+ m.def(
110
+ "moe_prepare_scatter(Tensor topk_ids, Tensor token_expert_indices, "
111
+ "Tensor? expert_map, int n_expert, int n_local_expert, "
112
+ "Tensor! expert_first_token_offset, Tensor! scatter_idx, "
113
+ "Tensor! sort_workspace, Tensor! sorted_experts, Tensor! sorted_rows, "
114
+ "Tensor! topk_ids_for_sort)->()");
115
+
116
m.def(
117
"moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
118
"Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
requirements/cuda.txt
+1
-1
@@ -33,4 +33,4 @@ quack-kernels==0.6.5 # Required by MSA
33
tokenspeed-mla==0.1.8; platform_system == "Linux"
34
35
# Humming kernels for quantization gemm
36
-humming-kernels[cu13]==0.1.12
36
+humming-kernels[cu13]==0.1.15
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-AWQ-humming.yaml
+2
@@ -9,3 +9,5 @@ server_args: >-
9
--moe-backend humming
10
--linear-backend humming
11
--dtype bfloat16
12
+env:
13
+ VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
new
+14
@@ -0,0 +1,14 @@
1
+model_name: "nm-testing/Qwen3-30B-A3B-FP8-block"
2
+accuracy_threshold: 0.86
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: >-
6
+ --enforce-eager
7
+ --max-model-len 8192
8
+ --moe-backend humming
9
+ --linear-backend humming
10
+ --data-parallel-size 2
11
+ --enable-expert-parallel
12
+ --all2all-backend deepep_high_throughput
13
+env:
14
+ VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
new
+14
@@ -0,0 +1,14 @@
1
+model_name: "nm-testing/Qwen3-30B-A3B-FP8-block"
2
+accuracy_threshold: 0.86
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: >-
6
+ --enforce-eager
7
+ --max-model-len 8192
8
+ --moe-backend humming
9
+ --linear-backend humming
10
+ --data-parallel-size 2
11
+ --enable-expert-parallel
12
+ --all2all-backend deepep_high_throughput
13
+env:
14
+ VLLM_HUMMING_MOE_GEMM_TYPE: "indexed"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml
+2
@@ -7,3 +7,5 @@ server_args: >-
7
--max-model-len 8192
8
--quantization humming
9
--linear-backend humming
10
+env:
11
+ VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-NVFP4-humming.yaml
+2
@@ -7,3 +7,5 @@ server_args: >-
7
--max-model-len 8192
8
--moe-backend humming
9
--linear-backend humming
10
+env:
11
+ VLLM_HUMMING_MOE_GEMM_TYPE: "grouped_contiguous"
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
deleted
-13
@@ -1,13 +0,0 @@
1
-model_name: "Qwen/Qwen3-30B-A3B"
2
-accuracy_threshold: 0.80
3
-num_questions: 1319
4
-num_fewshot: 5
5
-server_args: >-
6
- --enforce-eager
7
- --max-model-len 8192
8
- --tensor-parallel-size 1
9
- --quantization humming
10
- --kernel-config.enable_flashinfer_autotune=False
11
-env:
12
- VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
13
- VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
renamed
+4
-3
@@ -1,7 +1,8 @@
1
-model_name: "Qwen/Qwen3-30B-A3B"
2
-accuracy_threshold: 0.80
1
+model_name: "Qwen/Qwen3.5-35B-A3B"
2
+accuracy_threshold: 0.90
3
num_questions: 1319
4
num_fewshot: 5
5
+gen_prefix: " <think>\n\n</think>\n"
6
server_args: >-
7
--enforce-eager
8
--max-model-len 8192
@@ -9,5 +10,5 @@ server_args: >-
10
--quantization humming
11
--kernel-config.enable_flashinfer_autotune=False
12
env:
12
- VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
13
+ VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
14
VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}'
tests/evals/gsm8k/configs/humming/Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
renamed
+4
-3
@@ -1,7 +1,8 @@
1
-model_name: "Qwen/Qwen3-30B-A3B"
2
-accuracy_threshold: 0.80
1
+model_name: "Qwen/Qwen3.5-35B-A3B"
2
+accuracy_threshold: 0.90
3
num_questions: 1319
4
num_fewshot: 5
5
+gen_prefix: " <think>\n\n</think>\n"
6
server_args: >-
7
--enforce-eager
8
--max-model-len 8192
@@ -9,4 +10,4 @@ server_args: >-
10
--quantization humming
11
--kernel-config.enable_flashinfer_autotune=False
12
env:
12
- VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int5","hadamard_block_size":-1}'
13
+ 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
new
+13
@@ -0,0 +1,13 @@
1
+model_name: "Qwen/Qwen3.5-9B"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+gen_prefix: " <think>\n\n</think>\n"
6
+server_args: >-
7
+ --enforce-eager
8
+ --max-model-len 8192
9
+ --quantization humming
10
+ --kernel-config.enable_flashinfer_autotune=False
11
+env:
12
+ VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
13
+ VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"int8"}'
tests/evals/gsm8k/configs/humming/Qwen3.5-9B-int6wc-hadamard-humming.yaml
new
+12
@@ -0,0 +1,12 @@
1
+model_name: "Qwen/Qwen3.5-9B"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+gen_prefix: " <think>\n\n</think>\n"
6
+server_args: >-
7
+ --enforce-eager
8
+ --max-model-len 8192
9
+ --quantization humming
10
+ --kernel-config.enable_flashinfer_autotune=False
11
+env:
12
+ VLLM_HUMMING_ONLINE_QUANT_CONFIG: '{"dtype":"int6","hadamard_block_size":-1}'
tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-0.txt
new
+3
@@ -0,0 +1,3 @@
1
+gpt-oss-20b-humming.yaml
2
+Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml
3
+Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
tests/evals/gsm8k/configs/humming/config-a100-act-int8-shard-1.txt
new
+4
@@ -0,0 +1,4 @@
1
+Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml
2
+Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml
3
+Qwen3-30B-A3B-AWQ-humming.yaml
4
+Qwen3-4B-mixed-quant-RTN-humming.yaml
tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt
+1
@@ -1,4 +1,5 @@
1
Qwen3.5-35B-A3B-experts-int8-humming.yaml
2
+Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
3
gpt-oss-20b-humming.yaml
4
Qwen3-30B-A3B-GPTQ-Int4-humming.yaml
5
Qwen3-30B-A3B-FP8-block-humming.yaml
tests/evals/gsm8k/configs/humming/config-act-int8.txt
+1
@@ -3,4 +3,5 @@ Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming-act-int8.yaml
3
Qwen3-30B-A3B-GPTQ-Int4-humming-act-int8.yaml
4
Qwen3-30B-A3B-AWQ-humming.yaml
5
Qwen3.5-35B-A3B-experts-int8-humming-act-int8.yaml
6
+Qwen3.5-35B-A3B-int6wc-hadamard-humming-act-int8.yaml
7
Qwen3-4B-mixed-quant-RTN-humming.yaml
tests/evals/gsm8k/configs/humming/config-deepep.txt
new
+2
@@ -0,0 +1,2 @@
1
+Qwen3-30B-A3B-FP8-block-humming-deepep-ht-grouped.yaml
2
+Qwen3-30B-A3B-FP8-block-humming-deepep-ht-indexed.yaml
tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt
+1
@@ -1,4 +1,5 @@
1
Qwen3.5-35B-A3B-experts-int8-humming.yaml
2
+Qwen3.5-35B-A3B-int6wc-hadamard-humming.yaml
3
Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml
4
Qwen3-30B-A3B-Fp8-v1-humming.yaml
5
Qwen3-30B-A3B-NVFP4-humming.yaml
tests/evals/gsm8k/configs/humming/config-int5wc-hadamard.txt
deleted
-3
@@ -1,3 +0,0 @@
1
-Qwen3-30B-A3B-int5wc-hadamard-humming.yaml
2
-Qwen3-30B-A3B-int5wc-hadamard-humming-act-fp8.yaml
3
-Qwen3-30B-A3B-int5wc-hadamard-humming-act-int8.yaml
tests/evals/gsm8k/configs/humming/config-l4.txt
new
+9
@@ -0,0 +1,9 @@
1
+Qwen2-1.5B-Instruct-FP8W8-humming.yaml
2
+Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml
3
+Qwen3-0.6B-MXFP8-humming.yaml
4
+Qwen3-0.6B-MXFP8-humming-act-fp8.yaml
5
+Qwen3-30B-A3B-AWQ-humming.yaml
6
+Qwen3-30B-A3B-GPTQ-Int4-humming-act-fp8.yaml
7
+gpt-oss-20b-humming.yaml
8
+Qwen3.5-9B-int6wc-hadamard-humming.yaml
9
+Qwen3.5-9B-int6wc-hadamard-humming-act-int8.yaml
tests/kernels/core/test_activation.py
+109
-42
@@ -352,49 +352,46 @@ def test_masked_situ_and_mul(
352
)
353
354
355
-@pytest.mark.parametrize(
356
- ("activation", "activation_config"),
357
- [
358
- (MoEActivation.SILU, ApplyMoEActivationConfig()),
359
- (
360
- MoEActivation.SILU,
361
- ApplyMoEActivationConfig(clamp_limit=3.0),
355
+MOE_ACTIVATION_CASES = [
356
+ pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"),
357
+ pytest.param(
358
+ MoEActivation.SILU, ApplyMoEActivationConfig(clamp_limit=3.0), id="silu_clamp"
359
+ ),
360
+ pytest.param(MoEActivation.GELU, ApplyMoEActivationConfig(), id="gelu"),
361
+ pytest.param(MoEActivation.GELU_TANH, ApplyMoEActivationConfig(), id="gelu_tanh"),
362
+ pytest.param(
363
+ MoEActivation.SITU,
364
+ ApplyMoEActivationConfig(
365
+ activation_situ_beta=1.5,
366
+ activation_situ_linear_beta=2.0,
367
),
363
- (MoEActivation.GELU, ApplyMoEActivationConfig()),
364
- (MoEActivation.GELU_TANH, ApplyMoEActivationConfig()),
365
- (
366
- MoEActivation.SITU,
367
- ApplyMoEActivationConfig(
368
- activation_situ_beta=1.5,
369
- activation_situ_linear_beta=2.0,
370
- ),
371
- ),
372
- (MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig()),
373
- (
374
- MoEActivation.SWIGLUOAI_UNINTERLEAVE,
375
- ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
376
- ),
377
- (MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig()),
378
- (MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig()),
379
- (MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig()),
380
- (MoEActivation.GELU_TANH_NO_MUL, ApplyMoEActivationConfig()),
381
- (MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig()),
382
- ],
383
- ids=[
384
- "silu",
385
- "silu_clamp",
386
- "gelu",
387
- "gelu_tanh",
388
- "situ",
389
- "swigluoai",
390
- "swigluoai_uninterleave",
391
- "swiglustep",
392
- "silu_no_mul",
393
- "gelu_no_mul",
394
- "gelu_tanh_no_mul",
395
- "relu2_no_mul",
396
- ],
397
-)
368
+ id="situ",
369
+ ),
370
+ pytest.param(MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig(), id="swigluoai"),
371
+ pytest.param(
372
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
373
+ ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
374
+ id="swigluoai_uninterleave",
375
+ ),
376
+ pytest.param(MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig(), id="swiglustep"),
377
+ pytest.param(
378
+ MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig(), id="silu_no_mul"
379
+ ),
380
+ pytest.param(
381
+ MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig(), id="gelu_no_mul"
382
+ ),
383
+ pytest.param(
384
+ MoEActivation.GELU_TANH_NO_MUL,
385
+ ApplyMoEActivationConfig(),
386
+ id="gelu_tanh_no_mul",
387
+ ),
388
+ pytest.param(
389
+ MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig(), id="relu2_no_mul"
390
+ ),
391
+]
392
+
393
+
394
+@pytest.mark.parametrize(("activation", "activation_config"), MOE_ACTIVATION_CASES)
395
@torch.inference_mode()
396
def test_masked_moe_activation_dispatch(
397
default_vllm_config,
@@ -497,3 +494,73 @@ def test_activation(
494
495
out = torch.empty_like(x)
496
opcheck(fn, (out, x))
497
+
498
+
499
+HUMMING_ACTIVATION_CASES = MOE_ACTIVATION_CASES + [
500
+ pytest.param(MoEActivation.RELU2, ApplyMoEActivationConfig(), id="relu2"),
501
+]
502
+HUMMING_ACTIVATION_CASES += [
503
+ pytest.param(
504
+ MoEActivation.SITU,
505
+ ApplyMoEActivationConfig(
506
+ activation_situ_beta=1.5,
507
+ activation_situ_linear_beta=linear_beta,
508
+ ),
509
+ id=f"situ-linear-beta-{linear_beta}",
510
+ )
511
+ for linear_beta in (None, 0.0, -1.0)
512
+]
513
+
514
+
515
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Humming requires CUDA")
516
+@pytest.mark.parametrize(("activation", "activation_config"), HUMMING_ACTIVATION_CASES)
517
+@pytest.mark.parametrize("dtype", DTYPES)
518
+@pytest.mark.parametrize(("num_tokens", "d"), [(1, 512), (7, 768), (83, 512)])
519
+@torch.inference_mode()
520
+def test_humming_activation_matches_framework(
521
+ activation: MoEActivation,
522
+ activation_config: ApplyMoEActivationConfig,
523
+ dtype: torch.dtype,
524
+ num_tokens: int,
525
+ d: int,
526
+) -> None:
527
+ """Compare activation math/layouts without quantization or Hadamard error."""
528
+ pytest.importorskip("humming")
529
+ from humming.ops import process_input
530
+
531
+ from vllm.model_executor.layers.quantization.utils.humming.activation import (
532
+ get_humming_activation,
533
+ )
534
+
535
+ set_random_seed(0)
536
+ width = d * 2 if activation.is_gated else d
537
+ x = 4 * torch.randn(num_tokens, width, dtype=dtype, device=CUDA_DEVICES[0])
538
+ # Exercise zero, saturation, and both sides of the clamp limits (3 and 7).
539
+ edges = x.new_tensor([0, 0.001, 1, 2.99, 3, 3.01, 6.99, 7, 7.01, 8, 16])
540
+ edges = torch.cat((-edges[1:].flip(0), edges))
541
+ x[0] = edges[torch.arange(width, device=x.device) % edges.numel()]
542
+
543
+ actual, _, _ = process_input(
544
+ x,
545
+ quant_mode="none",
546
+ hadamard_block_size=0,
547
+ **get_humming_activation(activation, activation_config),
548
+ )
549
+ expected = torch.empty(num_tokens, d, dtype=dtype, device=x.device)
550
+ if activation == MoEActivation.RELU2:
551
+ # apply_moe_activation has only the non-gated ReLU2 variant.
552
+ gate, up = x.float().chunk(2, dim=-1)
553
+ activated_gate = torch.empty_like(gate)
554
+ apply_moe_activation(MoEActivation.RELU2_NO_MUL, activated_gate, gate.clone())
555
+ expected.copy_(activated_gate * up)
556
+ else:
557
+ # The framework's non-gated ReLU2 path modifies its input in place.
558
+ apply_moe_activation(
559
+ activation, expected, x.clone(), activation_config=activation_config
560
+ )
561
+ torch.testing.assert_close(
562
+ actual,
563
+ expected,
564
+ atol=get_default_atol(expected),
565
+ rtol=get_default_rtol(expected),
566
+ )
tests/kernels/moe/test_moe.py
+75
-26
@@ -1245,7 +1245,7 @@ def _make_humming_indexed_experts(activation: MoEActivation):
1245
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1246
HummingIndexedExperts,
1247
)
1248
- from vllm.model_executor.layers.quantization.utils import humming_utils
1248
+ from vllm.model_executor.layers.quantization.utils import humming as humming_utils
1249
from vllm.utils import humming
1250
1251
top_k, num_experts = 6, 12
@@ -1306,6 +1306,56 @@ def _make_humming_indexed_experts(activation: MoEActivation):
1306
return experts, layer
1307
1308
1309
+@pytest.mark.parametrize("valid_shape_m", [8192, 8193])
1310
+@pytest.mark.parametrize("w2_block_size", [64, 96, 128])
1311
+def test_humming_indexed_routing_matches_each_projection_tile(
1312
+ valid_shape_m: int, w2_block_size: int
1313
+):
1314
+ """Independently tuned projections preserve routed rows and expert identity."""
1315
+ from types import SimpleNamespace
1316
+
1317
+ from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1318
+ HummingIndexedExperts,
1319
+ )
1320
+
1321
+ topk_ids = torch.tensor([[0, 1], [1, 2], [2, 0]], device="cuda")
1322
+ expert_map = torch.tensor([1, -1, 0], dtype=torch.int32, device="cuda")
1323
+ experts = SimpleNamespace(
1324
+ global_num_experts=3,
1325
+ estimate_local_valid_shape_m=lambda _: valid_shape_m,
1326
+ w13_tuning_config=[
1327
+ [0, 8192, {"block_shape": [64, 128, 64]}],
1328
+ [8192, 16384, {"block_shape": [128, 128, 64]}],
1329
+ ],
1330
+ w2_tuning_config=[
1331
+ [0, 8192, {"block_shape": [64, 128, 64]}],
1332
+ [8192, 16384, {"block_shape": [w2_block_size, 128, 64]}],
1333
+ ],
1334
+ compute_config_str="compute",
1335
+ w13_tuning_config_str="w13",
1336
+ w2_tuning_config_str="w2",
1337
+ )
1338
+ first, second, _ = HummingIndexedExperts.prepare_humming_moe_kwargs(
1339
+ experts, topk_ids, expert_map, None
1340
+ )
1341
+ block_sizes = (64, 64) if valid_shape_m == 8192 else (128, w2_block_size)
1342
+ expected_rows = torch.where(expert_map[topk_ids.flatten()] >= 0)[0]
1343
+ for kwargs, block_size in zip((first, second), block_sizes):
1344
+ padded = kwargs["num_tokens_padded"].item()
1345
+ assert padded % block_size == 0
1346
+ rows = kwargs["sorted_ids"][:padded].long()
1347
+ valid = rows < topk_ids.numel()
1348
+ torch.testing.assert_close(rows[valid].sort().values, expected_rows)
1349
+ routed_experts = kwargs["expert_ids"][: padded // block_size]
1350
+ torch.testing.assert_close(
1351
+ routed_experts.repeat_interleave(block_size)[valid],
1352
+ expert_map[topk_ids.flatten()[rows[valid]]],
1353
+ )
1354
+ assert (first["top_k"], second["top_k"]) == (2, 1)
1355
+ for key in ("sorted_ids", "expert_ids", "num_tokens_padded"):
1356
+ assert (first[key] is second[key]) == (block_sizes[0] == block_sizes[1])
1357
+
1358
+
1359
@pytest.mark.parametrize("activation", list(MoEActivation))
1360
def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation):
1361
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
@@ -1464,54 +1514,46 @@ def test_humming_delegates_to_instance_activation():
1514
def test_humming_grouped_apply_forwards_valid_prefix(
1515
monkeypatch: pytest.MonkeyPatch,
1516
):
1467
- from types import MethodType, SimpleNamespace
1517
+ from types import SimpleNamespace
1518
from unittest.mock import Mock
1519
1520
import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
1521
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1472
- HummingExpertsBase,
1522
HummingGroupedExperts,
1523
)
1524
1476
- activation_func = Mock()
1525
buffers = {
1526
"gate_up_output": torch.empty(6, 4),
1479
- "activation_output": torch.empty(6, 2),
1527
+ "quanted_down_input": torch.empty(6, 2),
1528
"down_output": torch.empty(6, 4),
1529
}
1530
experts = SimpleNamespace(
1483
- activation=activation_func,
1531
num_experts=2,
1532
estimate_local_valid_shape_m=lambda _: 6,
1533
prepare_buffers=lambda *_: buffers,
1487
- _get_permute_scratch=lambda _: None,
1488
- quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: (
1489
- inputs,
1490
- input_scale,
1534
+ _get_permute_scratch=lambda _, *, indices_only=False: object(),
1535
+ process_input=Mock(
1536
+ side_effect=lambda _, **kwargs: (
1537
+ kwargs["inputs"],
1538
+ kwargs.get("input_scale"),
1539
+ None,
1540
+ )
1541
),
1542
humming_forward=Mock(),
1543
compute_config_str="",
1544
w13_tuning_config_str="",
1545
w2_tuning_config_str="",
1546
)
1497
- experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts)
1547
1548
hidden_states = torch.empty(3, 4)
1549
output = torch.empty_like(hidden_states)
1550
topk_ids = torch.zeros((3, 2), dtype=torch.int64)
1551
topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
1503
- expected_counts = torch.tensor([4], dtype=torch.int32)
1552
expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64)
1553
monkeypatch.setattr(
1554
humming,
1507
- "moe_permute",
1508
- lambda **_: (
1509
- torch.empty(6, 4),
1510
- None,
1511
- expert_offsets,
1512
- torch.arange(6),
1513
- None,
1514
- ),
1555
+ "moe_prepare_scatter",
1556
+ lambda *_: (expert_offsets, torch.arange(6).view(3, 2)),
1557
)
1558
monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None)
1559
@@ -1534,12 +1576,11 @@ def test_humming_grouped_apply_forwards_valid_prefix(
1576
apply_router_weight_on_input=False,
1577
)
1578
1537
- activation_func.assert_called_once()
1538
- call_kwargs = activation_func.call_args.kwargs
1579
+ call_kwargs = experts.process_input.call_args.kwargs
1580
assert call_kwargs["activation"] == MoEActivation.SITU
1540
- assert call_kwargs["input"].shape == (6, 4)
1541
- assert call_kwargs["output"].shape == (6, 2)
1542
- torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts)
1581
+ assert call_kwargs["inputs"].shape == (6, 4)
1582
+ assert call_kwargs["quanted_input"].shape == (6, 2)
1583
+ torch.testing.assert_close(call_kwargs["num_valid_tokens"], expert_offsets[-1:])
1584
1585
1586
def test_batched_marlin_activation_uses_expert_token_counts(
@@ -1653,7 +1694,8 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1694
activation=moe_config.activation,
1695
)
1696
assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size
1656
- assert buffer_metas["activation_output"]["shape"][-1] == intermediate_size
1697
+ assert "activation_output" not in buffer_metas
1698
+ assert buffer_metas["quanted_down_input"]["shape"][-1] == intermediate_size
1699
assert experts.moe_problem_size(
1700
a1=torch.empty(1, hidden_size),
1701
w1=torch.empty(num_experts, 1),
@@ -1661,6 +1703,13 @@ def test_humming_gated_non_gated_shape_contract(activation: MoEActivation):
1703
topk_ids=torch.empty(1, top_k, dtype=torch.long),
1704
) == (num_experts, 1, intermediate_size, hidden_size, top_k)
1705
1706
+ # Quantized batched input storage must cover every dispatcher's token slots.
1707
+ experts.is_batched = lambda: True
1708
+ experts.max_num_tokens = 17
1709
+ experts.num_dispatchers = 2
1710
+ batched_metas, _ = experts.get_buffer_metas(17, top_k, activation)
1711
+ assert batched_metas["quanted_gate_up_input"]["shape"][0] == num_experts * 17 * 2
1712
+
1713
1714
def test_humming_indexed_writes_supplied_output_buffer():
1715
from vllm.forward_context import set_forward_context
tests/kernels/moe/test_moe_align_block_size.py
+40
@@ -22,6 +22,46 @@ BLOCK_SIZES = [32, 128]
22
set_random_seed(0)
23
24
25
+@pytest.mark.parametrize("num_experts,num_tokens", [(32, 8), (256, 8), (256, 1024)])
26
+@pytest.mark.parametrize("use_expert_map", [False, True])
27
+def test_moe_align_scatter_masks_invalid_routes(
28
+ num_experts, num_tokens, use_expert_map
29
+):
30
+ topk_ids = torch.arange(num_tokens * 4, device="cuda").reshape(-1, 4)
31
+ topk_ids %= num_experts
32
+ topk_ids[0] = torch.tensor([-1, num_experts, 0, 1], device="cuda")
33
+ expert_map = None
34
+ valid = (topk_ids >= 0) & (topk_ids < num_experts)
35
+ if use_expert_map:
36
+ expert_map = torch.full((num_experts,), -1, dtype=torch.int32, device="cuda")
37
+ expert_map[::2] = torch.arange(
38
+ num_experts // 2, device="cuda", dtype=torch.int32
39
+ )
40
+ valid &= topk_ids % 2 == 0
41
+
42
+ sorted_ids, expert_ids, num_padded, scatter_idx = moe_align_block_size(
43
+ topk_ids,
44
+ 32,
45
+ num_experts,
46
+ expert_map,
47
+ ignore_invalid_experts=True,
48
+ return_scatter_idx=True,
49
+ )
50
+ expected = torch.arange(topk_ids.numel(), device="cuda", dtype=torch.int32)
51
+ expected[~valid.flatten()] = -1
52
+ torch.testing.assert_close(scatter_idx, expected[:, None])
53
+ active_sorted = sorted_ids[: num_padded.item()]
54
+ active_sorted = active_sorted[active_sorted < topk_ids.numel()]
55
+ torch.testing.assert_close(active_sorted.sort().values, expected[expected >= 0])
56
+
57
+ baseline = moe_align_block_size(
58
+ topk_ids, 32, num_experts, expert_map, ignore_invalid_experts=True
59
+ )
60
+ assert len(baseline) == 3
61
+ torch.testing.assert_close(num_padded, baseline[2])
62
+ torch.testing.assert_close(expert_ids, baseline[1])
63
+
64
+
65
def _group_tokens_by_expert(
66
sorted_ids: torch.Tensor,
67
expert_ids: torch.Tensor,
tests/kernels/moe/test_moe_permute_unpermute.py
+26
-1
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
18
get_moe_permute_scratch,
19
moe_permute,
20
moe_permute_unpermute_supported,
21
+ moe_prepare_scatter,
22
moe_unpermute,
23
)
24
from vllm.platforms import current_platform
@@ -458,7 +459,7 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
459
hidden_dtype=hidden_states.dtype,
460
)
461
461
- permuted, _, expert_offsets, _, _ = moe_permute(
462
+ permuted, _, expert_offsets, inverse, _ = moe_permute(
463
hidden_states=hidden_states,
464
a1q_scale=None,
465
topk_ids=topk_ids,
@@ -473,3 +474,27 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
474
torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"),
475
)
476
torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]])
477
+ assert torch.all(inverse[[1, 3, 4]] >= expert_offsets[-1])
478
+ expected = torch.zeros_like(hidden_states)
479
+ expected[[0, 2]] = hidden_states[[0, 2]]
480
+ output = torch.empty_like(hidden_states)
481
+ moe_unpermute(
482
+ output, permuted, torch.ones(5, 1, device="cuda"), inverse, expert_offsets
483
+ )
484
+ torch.testing.assert_close(output, expected)
485
+
486
+ expected_inverse = inverse.clone()
487
+ expected_offsets = expert_offsets.clone()
488
+ expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch)
489
+ torch.testing.assert_close(indices.flatten(), expected_inverse)
490
+ torch.testing.assert_close(expert_offsets, expected_offsets)
491
+ graph = torch.cuda.CUDAGraph()
492
+ with torch.cuda.graph(graph):
493
+ expert_offsets, indices = moe_prepare_scatter(topk_ids, expert_map, scratch)
494
+ topk_ids.zero_()
495
+ graph.replay()
496
+
497
+ expected_inverse = torch.arange(5, dtype=torch.int32, device="cuda")
498
+ expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda")
499
+ torch.testing.assert_close(indices.flatten(), expected_inverse)
500
+ torch.testing.assert_close(expert_offsets, expected_offsets)
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+1
-2
@@ -14,7 +14,6 @@ from vllm.model_executor.kernels.linear import (
14
AiterMxfp4LinearKernel,
15
EmulationMxfp4LinearKernel,
16
FlashInferMxFp4LinearKernel,
17
- HummingMxFp4LinearKernel,
17
MarlinMxFp4LinearKernel,
18
MxFp4LinearKernel,
19
MxFp4LinearLayerConfig,
@@ -44,7 +43,7 @@ _TRUE_W4A4_KERNELS = [
43
44
# Weight-only (A16) kernels: they never quantize activations. They still accept
45
# MXFP4 activation keys as an intentional compatibility fallback.
47
-_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel]
46
+_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel]
47
48
49
def test_can_implement_is_abstract():
tests/quantization/test_humming_mxfp4_block_fp8.py
+12
-8
@@ -15,6 +15,7 @@ it in vLLM comes down to two things, both exercised here:
15
"""
16
17
import pytest
18
+import torch
19
20
from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
21
HummingExpertsBase,
@@ -56,8 +57,8 @@ def test_mxfp4_weight_with_per_token_fp8_activation_still_supported():
57
],
58
)
59
def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key):
59
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
60
- _humming_input_schema_to_quant_key,
60
+ from vllm.model_executor.layers.quantization.utils.humming import (
61
+ input_schema_to_quant_key,
62
)
63
from vllm.utils.humming import HummingInputSchema
64
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):
67
a_dtype=humming_dtypes.float8e4m3,
68
input_scale_group_size=group_size,
69
)
69
- assert _humming_input_schema_to_quant_key(schema) == expected_key
70
+ assert input_schema_to_quant_key(schema, param_dtype=torch.bfloat16) == expected_key
71
72
73
@pytest.mark.skipif(not has_humming(), reason="humming is not installed")
74
def test_humming_bf16_input_schema_is_unquantized():
74
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
75
- _humming_input_schema_to_quant_key,
75
+ from vllm.model_executor.layers.quantization.utils.humming import (
76
+ input_schema_to_quant_key,
77
)
78
from vllm.utils.humming import HummingInputSchema
79
80
# No a_dtype -> unquantized (bf16/fp16) inputs -> None.
80
- assert _humming_input_schema_to_quant_key(HummingInputSchema()) is None
81
+ assert (
82
+ input_schema_to_quant_key(HummingInputSchema(), param_dtype=torch.bfloat16)
83
+ is None
84
+ )
85
86
87
def test_block_fp8_activation_quant_config_is_block_quantized():
@@ -86,7 +90,7 @@ def test_block_fp8_activation_quant_config_is_block_quantized():
90
DeepEP prepare/finalize step quantize activations to block FP8 *before* the
91
all-to-all dispatch instead of deferring to Humming -- and what makes
92
HummingExpertsBase.expects_unquantized_inputs return False."""
89
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
93
+ from vllm.model_executor.layers.quantization.utils.humming import (
94
make_humming_moe_quant_config,
95
)
96
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -108,7 +112,7 @@ def test_default_activation_quant_config_defers_to_humming():
112
"""Without a block activation shape the config stays per-token (the deferred
113
path): Humming quantizes internally, preserving pre-existing behavior for
114
every non-block-FP8 scheme."""
111
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
115
+ from vllm.model_executor.layers.quantization.utils.humming import (
116
make_humming_moe_quant_config,
117
)
118
from vllm.platforms import current_platform
tests/quantization/test_moe_wna16.py
+1
-1
@@ -229,7 +229,7 @@ def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
229
230
231
def test_moe_wna16_uses_humming_quant_config(monkeypatch):
232
- from vllm.model_executor.layers.quantization.utils import humming_utils
232
+ from vllm.model_executor.layers.quantization.utils import humming as humming_utils
233
234
method = object.__new__(MoeWNA16Method)
235
method.wna16_backend = WNA16MoEBackend.HUMMING
vllm/_custom_ops.py
+2
@@ -2272,6 +2272,7 @@ def moe_align_block_size(
2272
experts_ids: torch.Tensor,
2273
num_tokens_post_pad: torch.Tensor,
2274
expert_map: torch.Tensor | None = None,
2275
+ scatter_idx: torch.Tensor | None = None,
2276
) -> None:
2277
torch.ops._moe_C.moe_align_block_size(
2278
topk_ids,
@@ -2281,6 +2282,7 @@ def moe_align_block_size(
2282
experts_ids,
2283
num_tokens_post_pad,
2284
expert_map,
2285
+ scatter_idx,
2286
)
2287
2288
vllm/envs.py
+2
-1
@@ -1524,7 +1524,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
1524
"VLLM_HUMMING_ONLINE_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
1525
os.environ.get("VLLM_HUMMING_ONLINE_QUANT_CONFIG", None)
1526
),
1527
- # The activation dtype config for humming kernel
1527
+ # The activation dtype config for humming kernel. Explicit schemas disable
1528
+ # fallback unless the config includes "allow_fallback": true.
1529
"VLLM_HUMMING_INPUT_QUANT_CONFIG": lambda: maybe_convert_json_str_or_file(
1530
os.environ.get("VLLM_HUMMING_INPUT_QUANT_CONFIG", None)
1531
),
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -101,6 +101,7 @@ from vllm.model_executor.kernels.linear.mxfp6 import (
101
from vllm.model_executor.kernels.linear.mxfp6.emulation import (
102
EmulationMxfp6LinearKernel,
103
)
104
+from vllm.model_executor.kernels.linear.mxfp6.humming import HummingMxFp6LinearKernel
105
from vllm.model_executor.kernels.linear.mxfp8 import (
106
Mxfp8LinearKernel,
107
Mxfp8LinearLayerConfig,
@@ -295,6 +296,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
296
HummingLinearKernel,
297
HummingMxfp8LinearKernel,
298
HummingMxFp4LinearKernel,
299
+ HummingMxFp6LinearKernel,
300
HummingNvFp4LinearKernel,
301
},
302
"marlin": {
@@ -574,6 +576,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
576
577
_POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = {
578
PlatformEnum.CUDA: [
579
+ HummingMxFp6LinearKernel,
580
EmulationMxfp6LinearKernel,
581
],
582
PlatformEnum.ROCM: [
@@ -1294,6 +1297,7 @@ __all__ = [
1297
"MxFp6LinearLayerConfig",
1298
"init_mxfp6_linear_kernel",
1299
"EmulationMxfp6LinearKernel",
1300
+ "HummingMxFp6LinearKernel",
1301
"AiterMxfp4LinearKernel",
1302
"EmulationMxfp4LinearKernel",
1303
"FlashInferMxFp4LinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/humming.py
+3
-2
@@ -24,7 +24,7 @@ class HummingLinearKernel(MPLinearKernel):
24
return True, None
25
26
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
27
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
27
+ from vllm.model_executor.layers.quantization.utils.humming import (
28
convert_linear_layer_to_humming_standard,
29
get_humming_linear_compute_config,
30
prepare_humming_linear_layer_config,
@@ -57,13 +57,14 @@ class HummingLinearKernel(MPLinearKernel):
57
x: torch.Tensor,
58
bias: torch.Tensor | None = None,
59
) -> torch.Tensor:
60
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
60
+ from vllm.model_executor.layers.quantization.utils.humming import (
61
apply_humming_linear,
62
)
63
64
return apply_humming_linear(
65
layer,
66
x,
67
+ skip_bias_add=bias is None,
68
layer_config=self.layer_config,
69
compute_config=self.compute_config,
70
locks=self.locks,
vllm/model_executor/kernels/linear/mxfp4/humming.py
+11
-14
@@ -3,21 +3,18 @@
3
4
import torch
5
6
-from vllm.logger import init_logger
7
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
6
+from vllm.model_executor.layers.quantization.utils.humming import (
7
apply_humming_linear,
8
convert_linear_layer_to_humming_standard,
9
get_humming_linear_compute_config,
10
prepare_humming_linear_layer_config,
11
+ quant_key_to_input_schema,
12
)
13
-from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
13
from vllm.platforms import current_platform
14
from vllm.utils.import_utils import has_humming
15
16
from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
17
19
-logger = init_logger(__name__)
20
-
18
19
class HummingMxFp4LinearKernel(MxFp4LinearKernel):
20
"""Humming GEMM Kernel for MXFP4."""
@@ -39,14 +36,10 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel):
36
37
@classmethod
38
def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
42
- if config.activation_quant_key not in (None, kMxfp4Dynamic):
43
- return False, "only supports MXFP4 dynamic or unquantized activations"
44
- if config.activation_quant_key is not None:
45
- logger.warning_once(
46
- "HummingMxFp4LinearKernel is a weight-only (A16) kernel; "
47
- "the requested activation quantization (%s) is ignored.",
48
- config.activation_quant_key,
49
- )
39
+ try:
40
+ quant_key_to_input_schema(config.activation_quant_key)
41
+ except ValueError as error:
42
+ return False, str(error)
43
return True, None
44
45
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -62,7 +55,10 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel):
55
}
56
57
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
65
- self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
58
+ input_schema = quant_key_to_input_schema(self.config.activation_quant_key)
59
+ self.layer_config = prepare_humming_linear_layer_config(
60
+ layer, quant_config, input_schema=input_schema
61
+ )
62
self.compute_config = get_humming_linear_compute_config()
63
self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
64
@@ -75,6 +71,7 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel):
71
return apply_humming_linear(
72
layer,
73
x,
74
+ skip_bias_add=bias is None,
75
layer_config=self.layer_config,
76
compute_config=self.compute_config,
77
locks=self.locks,
vllm/model_executor/kernels/linear/mxfp6/base.py
+2
-2
@@ -13,8 +13,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
13
class MxFp6LinearLayerConfig:
14
"""Configuration for an MXFP6 linear layer.
15
16
- All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per
17
- byte) and per-block weight scales (group size 32).
16
+ All MXFP6 layers share the same structure: packed uint8 weights (4 FP6 values
17
+ per 3 bytes) and per-block weight scales (group size 32).
18
19
Attributes:
20
weight_quant_key: Identifies the weight quantization format. Can be
vllm/model_executor/kernels/linear/mxfp6/humming.py
new
+84
@@ -0,0 +1,84 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.model_executor.layers.quantization.utils.humming import (
7
+ apply_humming_linear,
8
+ convert_linear_layer_to_humming_standard,
9
+ get_humming_linear_compute_config,
10
+ prepare_humming_linear_layer_config,
11
+ quant_key_to_input_schema,
12
+)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp6E2M3Static
14
+from vllm.platforms import current_platform
15
+from vllm.utils.import_utils import has_humming
16
+
17
+from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig
18
+
19
+
20
+class HummingMxFp6LinearKernel(MxFp6LinearKernel):
21
+ """Humming GEMM for packed MXFP6 E2M3 and E3M2 weights."""
22
+
23
+ @classmethod
24
+ def is_supported(
25
+ cls, compute_capability: int | None = None
26
+ ) -> tuple[bool, str | None]:
27
+ if not current_platform.is_cuda():
28
+ return False, "Humming only supported on CUDA"
29
+
30
+ if not has_humming():
31
+ return False, "Humming is not installed"
32
+
33
+ if not current_platform.has_device_capability(75):
34
+ return False, "Humming only supported on SM75+"
35
+
36
+ return True, None
37
+
38
+ @classmethod
39
+ def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
40
+ try:
41
+ quant_key_to_input_schema(config.activation_quant_key)
42
+ except ValueError as error:
43
+ return False, str(error)
44
+ return True, None
45
+
46
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
47
+ weight_dtype = "float6e3m2"
48
+ if self.config.weight_quant_key == kMxfp6E2M3Static:
49
+ weight_dtype = "float6e2m3"
50
+
51
+ quant_config = {
52
+ "quant_method": "humming",
53
+ "dtype": weight_dtype,
54
+ "scale_dtype": "float8e8m0",
55
+ "group_size": 32,
56
+ "weight_scale_type": "group",
57
+ }
58
+ layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
59
+
60
+ convert_linear_layer_to_humming_standard(
61
+ layer=layer,
62
+ name_map={"weight": "weight", "weight_scale": "weight_scale"},
63
+ )
64
+ input_schema = quant_key_to_input_schema(self.config.activation_quant_key)
65
+ self.layer_config = prepare_humming_linear_layer_config(
66
+ layer, quant_config, input_schema=input_schema
67
+ )
68
+ self.compute_config = get_humming_linear_compute_config()
69
+ self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
70
+
71
+ def apply_weights(
72
+ self,
73
+ layer: torch.nn.Module,
74
+ x: torch.Tensor,
75
+ bias: torch.Tensor | None = None,
76
+ ) -> torch.Tensor:
77
+ return apply_humming_linear(
78
+ layer,
79
+ x,
80
+ skip_bias_add=bias is None,
81
+ layer_config=self.layer_config,
82
+ compute_config=self.compute_config,
83
+ locks=self.locks,
84
+ )
vllm/model_executor/kernels/linear/mxfp8/humming.py
+8
-2
@@ -3,12 +3,14 @@
3
4
import torch
5
6
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
6
+from vllm.model_executor.layers.quantization.utils.humming import (
7
apply_humming_linear,
8
convert_linear_layer_to_humming_standard,
9
get_humming_linear_compute_config,
10
prepare_humming_linear_layer_config,
11
+ quant_key_to_input_schema,
12
)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
14
from vllm.platforms import current_platform
15
from vllm.utils.import_utils import has_humming
16
@@ -50,7 +52,10 @@ class HummingMxfp8LinearKernel(Mxfp8LinearKernel):
52
}
53
54
convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
53
- self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
55
+ input_schema = quant_key_to_input_schema(kMxfp8Dynamic)
56
+ self.layer_config = prepare_humming_linear_layer_config(
57
+ layer, quant_config, input_schema=input_schema
58
+ )
59
self.compute_config = get_humming_linear_compute_config()
60
self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
61
@@ -63,6 +68,7 @@ class HummingMxfp8LinearKernel(Mxfp8LinearKernel):
68
return apply_humming_linear(
69
layer,
70
x,
71
+ skip_bias_add=bias is None,
72
layer_config=self.layer_config,
73
compute_config=self.compute_config,
74
locks=self.locks,
vllm/model_executor/kernels/linear/nvfp4/humming.py
+15
-5
@@ -4,11 +4,13 @@
4
import torch
5
6
from vllm.logger import init_logger
7
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
7
+from vllm.model_executor.layers.quantization.utils.humming import (
8
apply_humming_linear,
9
get_humming_linear_compute_config,
10
prepare_humming_linear_layer_config,
11
+ quant_key_to_input_schema,
12
)
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import kNvfp4Dynamic
14
from vllm.platforms import current_platform
15
from vllm.utils.import_utils import has_humming
16
@@ -40,9 +42,7 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel):
42
return True, None
43
44
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
43
- # Route through humming's compressed-tensors nvfp4 loader (same path as
44
- # the MoE oracle); the native group_tensor schema mishandles a scalar
45
- # global scale.
45
+ # Reuse the CT loader for packed FP4 weights and tensor scales.
46
quant_config = {
47
"quant_method": "compressed-tensors",
48
"format": "nvfp4-pack-quantized",
@@ -60,7 +60,16 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel):
60
layer.weight_global_scale = torch.nn.Parameter(
61
1.0 / layer.weight_global_scale, requires_grad=False
62
)
63
- self.layer_config = prepare_humming_linear_layer_config(layer, quant_config)
63
+ input_quant_key = None
64
+ if hasattr(layer, "input_global_scale"):
65
+ layer.input_scale_2 = torch.nn.Parameter(
66
+ layer.input_global_scale.reshape(-1), requires_grad=False
67
+ )
68
+ input_quant_key = kNvfp4Dynamic
69
+ input_schema = quant_key_to_input_schema(input_quant_key)
70
+ self.layer_config = prepare_humming_linear_layer_config(
71
+ layer, quant_config, input_schema=input_schema
72
+ )
73
self.compute_config = get_humming_linear_compute_config()
74
self.locks = torch.zeros(1024, dtype=torch.int32, device=layer.weight.device)
75
@@ -73,6 +82,7 @@ class HummingNvFp4LinearKernel(NvFp4LinearKernel):
82
return apply_humming_linear(
83
layer,
84
x,
85
+ skip_bias_add=bias is None,
86
layer_config=self.layer_config,
87
compute_config=self.compute_config,
88
locks=self.locks,
vllm/model_executor/kernels/linear/scaled_mm/humming.py
+3
-1
@@ -4,7 +4,7 @@
4
import torch
5
6
from vllm.logger import init_logger
7
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
7
+from vllm.model_executor.layers.quantization.utils.humming import (
8
apply_humming_linear,
9
convert_linear_layer_to_humming_standard,
10
get_humming_linear_compute_config,
@@ -92,6 +92,7 @@ class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
92
return apply_humming_linear(
93
layer,
94
x,
95
+ skip_bias_add=bias is None,
96
layer_config=self.layer_config,
97
compute_config=self.compute_config,
98
locks=self.locks,
@@ -158,6 +159,7 @@ class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
159
return apply_humming_linear(
160
layer,
161
x,
162
+ skip_bias_add=bias is None,
163
layer_config=self.layer_config,
164
compute_config=self.compute_config,
165
locks=self.locks,
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+181
-249
@@ -4,6 +4,7 @@
4
5
import json
6
import math
7
+from dataclasses import replace
8
from typing import TYPE_CHECKING, Any, cast
9
10
import torch
@@ -30,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
31
get_moe_permute_scratch,
32
moe_permute,
33
moe_permute_unpermute_supported,
34
+ moe_prepare_scatter,
35
moe_unpermute,
36
)
37
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 (
59
kMxfp8Static,
60
kNvfp4Dynamic,
61
kNvfp4Static,
62
+ kStaticTensorScale,
63
)
64
from vllm.platforms import current_platform
65
from vllm.scalar_type import ScalarType
66
from vllm.utils.import_utils import has_humming
67
+from vllm.utils.platform_utils import num_compute_units
68
from vllm.v1.worker.workspace import current_workspace_manager
69
70
if TYPE_CHECKING:
67
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
71
+ from vllm.model_executor.layers.quantization.utils.humming import (
72
HummingMoEQuantConfig,
73
)
74
from vllm.utils.humming import (
@@ -118,38 +122,9 @@ def get_humming_moe_gemm_type(moe_config: FusedMoEConfig | None = None) -> str:
122
return gemm_type
123
124
121
-def _fixup_moe_tuning_config(tuning_config: list, max_k_block: int = 128) -> None:
122
- """Fix up each MoE tile in place: cap the K-block and widen warp-N.
123
-
124
- - block_shape[2] (K-block) > ``max_k_block``: the driver rejects the TMA
125
- descriptor at launch (CUDA_ERROR_MISALIGNED_ADDRESS). Cap at 128, which
126
- Humming already uses for larger M.
127
- - warp_shape[1] (warp-N) < 32: block-FP8 (group-128) activations route the
128
- w13 (gate/up) GEMM to a tuning table that pins warp-N to 16, under-filling
129
- the Hopper WGMMA N dimension and corrupting the GEMM output (gsm8k
130
- 0.94 -> 0.89). Widen to 32 whenever block_n % 32 == 0 -- on every tile,
131
- not just the K-capped ones. w2 (down) already uses warp_n=32.
132
- """
133
- logger.info_once("Attempting to override humming GEMM config")
134
- for entry in tuning_config:
135
- config = entry[2]
136
- block_shape = config.get("block_shape")
137
- if not (block_shape and len(block_shape) == 3):
138
- continue
139
- block_m, block_n, block_k = block_shape
140
-
141
- logger.info_once(f"Overriding humming GEMM config. Previous config\n: {config}")
142
- if block_k > max_k_block:
143
- config["block_shape"] = [block_m, block_n, max_k_block]
144
-
145
- warp_shape = config.get("warp_shape")
146
- if warp_shape and warp_shape[1] < 32 and block_n % 32 == 0:
147
- config["warp_shape"] = [warp_shape[0], 32, warp_shape[2]]
148
-
149
- logger.info_once(f"Overridden humming GEMM config. Current config\n: {config}")
150
-
151
-
125
class HummingExpertsBase(mk.FusedMoEExpertsModular):
126
+ quant_config: "HummingMoEQuantConfig"
127
+
128
def __init__(
129
self,
130
moe_config: FusedMoEConfig,
@@ -165,7 +140,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
140
self.locks = torch.zeros(1024, dtype=torch.int32, device=moe_config.device)
141
self.num_experts = moe_config.num_local_experts
142
self.global_num_experts = moe_config.num_experts
168
- self.quant_config = quant_config
143
+ self.quant_config = humming_quant_config
144
self.init_humming_moe()
145
146
if self.is_batched():
@@ -199,31 +174,61 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
174
gemm_type=self.humming_gemm_type(),
175
)
176
self.compute_config_str = json.dumps(self.compute_config)
202
- # Fix up the heuristic-chosen tiles before freezing to JSON -- see
203
- # _fixup_moe_tuning_config.
204
- _fixup_moe_tuning_config(self.w13_tuning_config)
205
- _fixup_moe_tuning_config(self.w2_tuning_config)
177
self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
178
self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)
179
209
- def quantize_input(
180
+ def process_input(
181
self,
182
sublayer_name: str,
183
inputs: torch.Tensor,
184
quanted_input: torch.Tensor | None,
185
input_scale: torch.Tensor | None = None,
215
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
216
- from vllm.utils.humming import may_quant_input
186
+ expert_tokens: torch.Tensor | None = None,
187
+ activation: MoEActivation | None = None,
188
+ scatter_idx: torch.Tensor | None = None,
189
+ num_valid_tokens: torch.Tensor | None = None,
190
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
191
+ from vllm.model_executor.layers.quantization.utils.humming.activation import (
192
+ get_humming_activation,
193
+ )
194
+ from vllm.utils.humming import may_process_input
195
218
- # input_scale is set for block-FP8 (group-128) activations that were
219
- # quantized before the EP dispatch: may_quant_input then skips the
220
- # redundant w13 quantization and forwards the pre-computed scale.
221
- return may_quant_input(
222
- self.humming_configs[sublayer_name],
196
+ quant_config = self.quant_config
197
+ config = self.humming_configs[sublayer_name]
198
+ if input_scale is not None:
199
+ assert scatter_idx is None
200
+ if activation is not None:
201
+ raise ValueError("Cannot apply activation to prequantized input")
202
+ return inputs, input_scale, None
203
+
204
+ activation_kwargs = {}
205
+ if activation is not None:
206
+ activation_config = self.activation_config
207
+ activation_kwargs = get_humming_activation(activation, activation_config)
208
+
209
+ prefix = "w1" if sublayer_name == "w13" else "w2"
210
+ scale = getattr(quant_config, f"{prefix}_input_scale")
211
+ scale2 = getattr(quant_config, f"{prefix}_input_scale_2")
212
+ mode = config.input_quant_mode
213
+ layout = "grouped_mask" if self.is_batched() else "normal"
214
+ if scatter_idx is not None:
215
+ layout = "scatter"
216
+
217
+ inputs, group_scales, token_scales = may_process_input(
218
+ config,
219
inputs=inputs,
224
- input_scale=input_scale,
225
- quanted_input=quanted_input,
220
+ outputs=quanted_input,
221
+ token_scales=scale2 if mode.has_secondary_scale else scale,
222
+ hadamard_block_size=getattr(quant_config, f"{prefix}_hadamard_block_size"),
223
+ layout=layout,
224
+ expert_tokens=expert_tokens if layout == "grouped_mask" else None,
225
+ scatter_idx=scatter_idx,
226
+ num_valid_tokens=num_valid_tokens,
227
+ **activation_kwargs,
228
)
229
+ input_scale = group_scales if mode.has_group_scale else token_scales
230
+ input_scale_2 = token_scales if mode.has_secondary_scale else None
231
+ return inputs, input_scale, input_scale_2
232
233
def humming_forward(
234
self,
@@ -232,30 +237,30 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
237
weight: torch.Tensor,
238
input_scale: torch.Tensor | None,
239
outputs: torch.Tensor,
240
+ input_scale_2: torch.Tensor | None = None,
241
**kwargs: Any,
242
) -> torch.Tensor:
243
from vllm.utils.humming import humming_forward
244
239
- is_w13 = sublayer_name == "w13"
245
+ index = 1 if sublayer_name == "w13" else 2
246
return humming_forward(
247
self.humming_configs[sublayer_name],
248
inputs=inputs,
249
weight=weight,
244
- weight_scale=(
245
- self.quant_config.w1_scale if is_w13 else self.quant_config.w2_scale
246
- ),
247
- zero_point=(self.quant_config.w1_zp if is_w13 else self.quant_config.w2_zp),
248
- bias=(self.quant_config.w1_bias if is_w13 else self.quant_config.w2_bias),
249
- weight_scale_2=(
250
- self.quant_config.g1_alphas if is_w13 else self.quant_config.g2_alphas
251
- ),
250
+ weight_scale=getattr(self.quant_config, f"w{index}_scale"),
251
+ zero_point=getattr(self.quant_config, f"w{index}_zp"),
252
+ bias=getattr(self.quant_config, f"w{index}_bias"),
253
+ weight_scale_2=getattr(self.quant_config, f"g{index}_alphas"),
254
input_scale=input_scale,
255
+ input_scale_2=input_scale_2,
256
outputs=outputs,
257
locks=self.locks,
258
**kwargs,
259
)
260
258
- def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None:
261
+ def _get_permute_scratch(
262
+ self, topk: int, indices_only: bool = False
263
+ ) -> MoEPermuteScratch | None:
264
if not moe_permute_unpermute_supported():
265
return None
266
@@ -270,8 +275,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
275
num_experts=self.moe_config.num_experts,
276
num_local_experts=self.moe_config.num_local_experts,
277
device=torch.device(self.moe_config.device),
273
- hidden_size=self.moe_config.hidden_dim,
274
- hidden_dtype=self.moe_config.in_dtype,
278
+ hidden_size=None if indices_only else self.moe_config.hidden_dim,
279
+ hidden_dtype=None if indices_only else self.moe_config.in_dtype,
280
)
281
282
def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
@@ -302,6 +307,16 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
307
weight_key: QuantKey | None,
308
activation_key: QuantKey | None,
309
) -> bool:
310
+ if (
311
+ activation_key is not None
312
+ and activation_key.scale == kNvfp4Dynamic.scale
313
+ and activation_key.dtype == kNvfp4Dynamic.dtype
314
+ and activation_key.scale2 is not None
315
+ and not activation_key.scale2.static
316
+ and activation_key.scale2.group_shape == GroupShape.PER_TOKEN
317
+ and activation_key.scale2.dtype == torch.float32
318
+ ):
319
+ activation_key = replace(activation_key, scale2=kStaticTensorScale)
320
SUPPORTED_W_A = [
321
(kMxfp4Static, None),
322
(kMxfp4Static, kMxfp4Dynamic),
@@ -343,7 +358,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
358
(kMxfp8Static, kMxfp8Dynamic),
359
]
360
return (weight_key, activation_key) in SUPPORTED_W_A or (
346
- activation_key in (None, kFp8DynamicTokenSym)
361
+ activation_key in (None, kFp8DynamicTokenSym, kInt8DynamicTokenSym)
362
and _is_supported_wna16_weight_key(weight_key)
363
)
364
@@ -355,7 +370,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
370
to FP8 before dispatch sends FP8 rather than BF16 over the interconnect,
371
and Humming then consumes the pre-quantized FP8 + scale directly (see
372
the apply() methods, which forward the dispatch scale into
358
- HummingMethod.may_quant_input, and may_quant_input itself, which is a
373
+ HummingExpertsBase.process_input, which is a
374
no-op when an input scale is already supplied). The scale layout
375
produced by vLLM's block-FP8 quantization ([M, K // 128] float32,
376
row-major) matches what the Humming WGMMA grouped GEMM expects.
@@ -373,7 +388,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
388
inputs).
389
390
Humming normally quantizes inputs internally via
376
- HummingMethod.may_quant_input() in apply(), so we defer quantization
391
+ HummingExpertsBase.process_input() in apply(), so we defer quantization
392
(return True) to avoid quantizing twice -- once in prepare and once in
393
Humming's apply().
394
@@ -398,8 +413,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
413
414
@staticmethod
415
def _supports_activation(activation: MoEActivation) -> bool:
401
- # Humming uses apply_moe_activation() callback for activation,
402
- # so any activation supported there can be used here.
416
return apply_moe_activation_supported(activation)
417
418
@staticmethod
@@ -456,10 +469,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
469
assert intermediate_dim == self.adjust_N_for_activation(gate_up_dim, activation)
470
471
# hidden_states
459
- # (-> quanted_gate_up_input) (if not BF16/FP16 activation)
472
+ # -> quanted_gate_up_input (Hadamard + optional quantization)
473
# -> gate_up_output
461
- # -> activation_output
462
- # (-> quanted_down_input) (if not BF16/FP16 activation)
474
+ # -> quanted_down_input (activation + Hadamard + optional quantization)
475
# -> down_output
476
# (-> output) (if not is_batched)
477
# Neighboring nodes are required to utilize distinct workspaces.
@@ -471,8 +483,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
483
max_num_tokens = self.max_num_tokens
484
num_dispatchers = self.num_dispatchers
485
assert max_num_tokens is not None and num_dispatchers is not None
474
- input_shape_m = num_experts * max_num_tokens
486
real_shape_m = num_experts * max_num_tokens * num_dispatchers
487
+ input_shape_m = real_shape_m
488
output_shape = (num_experts, max_num_tokens * num_dispatchers, K)
489
else:
490
input_shape_m = M
@@ -483,15 +495,18 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
495
496
a_dtype = self.humming_configs["w13"].a_dtype
497
c_dtype = self.humming_configs["w13"].c_dtype
486
- num_bits = a_dtype.num_bits
498
+ down_a_dtype = w2_config.a_dtype
499
torch_dtype_map = {
500
dtypes.float16: torch.float16,
501
dtypes.bfloat16: torch.bfloat16,
502
dtypes.float32: torch.float32,
503
+ dtypes.float8e3m4: torch.uint8,
504
dtypes.float8e4m3: torch.float8_e4m3fn,
505
dtypes.float8e5m2: torch.float8_e5m2,
506
dtypes.int8: torch.int8,
507
dtypes.int4: torch.uint8,
508
+ dtypes.float4e0m3: torch.uint8,
509
+ dtypes.float4e2m1: torch.uint8,
510
}
511
512
buffer_metas = {
@@ -503,13 +518,9 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
518
"shape": (real_shape_m, gate_up_dim),
519
"dtype": torch_dtype_map[c_dtype],
520
},
506
- "activation_output": {
507
- "shape": (real_shape_m, intermediate_dim),
508
- "dtype": torch_dtype_map[c_dtype],
509
- },
521
"quanted_down_input": {
522
"shape": (real_shape_m, intermediate_dim),
512
- "dtype": torch_dtype_map[a_dtype],
523
+ "dtype": torch_dtype_map[down_a_dtype],
524
},
525
"down_output": {
526
"shape": output_shape if self.is_batched() else (real_shape_m, K),
@@ -523,7 +534,8 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
534
535
for key in buffer_metas:
536
meta = buffer_metas[key]
526
- if "quanted" in key and a_dtype.num_bits == 4:
537
+ input_dtype = down_a_dtype if key == "quanted_down_input" else a_dtype
538
+ if "quanted" in key and input_dtype.num_bits == 4:
539
last_dim = meta["shape"][-1]
540
if last_dim % 2 != 0:
541
raise ValueError(
@@ -532,25 +544,12 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
544
)
545
meta["shape"] = meta["shape"][:-1] + (last_dim // 2,)
546
535
- if num_bits == 16:
536
- required_buffers = ["gate_up_output", "activation_output", "down_output"]
537
- else:
538
- required_buffers = [
539
- "quanted_gate_up_input",
540
- "gate_up_output",
541
- "activation_output",
542
- "quanted_down_input",
543
- "down_output",
544
- ]
545
-
546
- # Note: The fused SITU+FP8 quant goes gate_up_output ->
547
- # quanted_down_input directly, never materializing
548
- # activation_output. Dropping activation_output from the chain
549
- # flips the even/odd workspace 2-coloring below so
550
- # gate_up_output and quanted_down_input land
551
- # on DIFFERENT workspaces.
552
- if self.fused_situ_quant_enabled(activation):
553
- required_buffers.remove("activation_output")
547
+ required_buffers = [
548
+ "quanted_gate_up_input",
549
+ "gate_up_output",
550
+ "quanted_down_input",
551
+ "down_output",
552
+ ]
553
554
# batched moe use down_output as output
555
if not self.is_batched():
@@ -673,81 +672,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
672
activation_kwargs["valid_token_counts"] = valid_token_counts
673
self.activation(**activation_kwargs)
674
676
- def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool:
677
- """Whether the SITU activation + w2 quant can be fused into one kernel.
678
-
679
- Fused only for k-major block-FP8 group-128 e4m3 with float32 scales --
680
- the sole layout situ_and_mul_quant supports. A float32 as_dtype rules out
681
- MXMMA (which uses e8m0 m-major scales), so a group-128 scale is
682
- guaranteed k-major. Everything else (m-major MXFP8, other group sizes,
683
- 16-bit passthrough, non-SITU) falls back to the separate
684
- situ_and_mul + quantize_input path.
685
- """
686
- if activation != MoEActivation.SITU:
687
- return False
688
- w2cfg = self.humming_configs["w2"]
689
- # Report the specific blocker once so a silent fallback to the unfused
690
- # situ_and_mul + quant pair is never a mystery in a trace.
691
- reason: str | None = None
692
- if not (w2cfg.a_dtype.num_bits == 8 and str(w2cfg.a_dtype) == "float8e4m3"):
693
- reason = f"w2 a_dtype is {w2cfg.a_dtype} (need float8e4m3)"
694
- elif w2cfg.input_scale_group_size != 128:
695
- reason = (
696
- f"w2 input_scale_group_size is {w2cfg.input_scale_group_size} "
697
- "(need 128)"
698
- )
699
- elif not (w2cfg.as_dtype is not None and str(w2cfg.as_dtype) == "float32"):
700
- reason = f"w2 as_dtype is {w2cfg.as_dtype} (need float32, k-major)"
701
- if reason is not None:
702
- logger.warning_once(
703
- "Humming fused SITU+FP8 quant disabled, using unfused "
704
- "situ_and_mul + quant: %s",
705
- reason,
706
- )
707
- return False
708
- return True
709
-
710
- def fused_situ_quant(
711
- self,
712
- gate_up_output: torch.Tensor,
713
- quanted_down_input: torch.Tensor,
714
- num_valid_tokens: torch.Tensor | None,
715
- topk: int = 1,
716
- ) -> tuple[torch.Tensor, torch.Tensor]:
717
- """Run the fused SITU activation + FP8 quant for the w2 input.
718
-
719
- Returns ``(quanted_down_input, input_scale)`` in the same layout the
720
- unfused ``apply_activation`` + ``quantize_input("w2")`` pair produces: an
721
- fp8 [M, d] tensor plus a k-major block-FP8 float32 [M, d // 128] group
722
- scale. Only reached when ``fused_situ_quant_enabled`` (group_size 128).
723
- """
724
- from vllm.model_executor.layers.fused_moe.activation import (
725
- situ_and_mul_quant,
726
- )
727
-
728
- cfg = self.activation_config
729
- assert cfg.activation_situ_beta is not None, (
730
- "SITU requires activation_situ_beta from FusedMoEConfig"
731
- )
732
- group_size = self.humming_configs["w2"].input_scale_group_size
733
- m, d = quanted_down_input.size(0), quanted_down_input.size(1)
734
- input_scale = torch.empty(
735
- (m, d // group_size),
736
- dtype=torch.float32,
737
- device=quanted_down_input.device,
738
- )
739
- situ_and_mul_quant(
740
- quanted_down_input,
741
- input_scale,
742
- gate_up_output,
743
- beta=cfg.activation_situ_beta,
744
- linear_beta=cfg.activation_situ_linear_beta,
745
- group_size=group_size,
746
- num_valid_tokens=num_valid_tokens,
747
- topk=topk,
748
- )
749
- return quanted_down_input, input_scale
750
-
675
676
class HummingIndexedExperts(HummingExpertsBase):
677
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
@@ -768,7 +692,7 @@ class HummingIndexedExperts(HummingExpertsBase):
692
topk_ids: torch.Tensor,
693
expert_map: torch.Tensor | None,
694
expert_tokens_meta: mk.ExpertTokensMetadata | None,
771
- ) -> tuple[dict[str, Any], dict[str, Any]]:
695
+ ) -> tuple[dict[str, Any], dict[str, Any], torch.Tensor | None]:
696
valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
697
698
moe_block_size = None
@@ -784,13 +708,21 @@ class HummingIndexedExperts(HummingExpertsBase):
708
)
709
moe_block_size = 64
710
787
- sorted_ids, expert_ids, num_tokens_padded = moe_align_block_size(
711
+ use_scatter = False
712
+ if expert_map is not None:
713
+ num_sms = num_compute_units(topk_ids.get_device())
714
+ use_scatter = topk_ids.numel() > 2 * num_sms
715
+
716
+ alignment = moe_align_block_size(
717
topk_ids=topk_ids,
718
block_size=moe_block_size,
719
num_experts=self.global_num_experts,
720
expert_map=expert_map,
721
ignore_invalid_experts=True,
722
+ return_scatter_idx=use_scatter,
723
)
724
+ sorted_ids, expert_ids, num_tokens_padded = alignment[:3]
725
+ scatter_idx = alignment[3] if len(alignment) == 4 else None
726
727
moe_common_kwargs = {
728
"sorted_ids": sorted_ids,
@@ -806,7 +738,26 @@ class HummingIndexedExperts(HummingExpertsBase):
738
moe_kwargs1.update(moe_common_kwargs)
739
moe_kwargs2.update(moe_common_kwargs)
740
809
- return moe_kwargs1, moe_kwargs2
741
+ w2_block_size = 64
742
+ for lower, upper, config in self.w2_tuning_config:
743
+ if lower < valid_shape_m <= upper:
744
+ w2_block_size = config["block_shape"][0]
745
+ break
746
+ if w2_block_size != moe_block_size:
747
+ sorted_ids2, expert_ids2, num_tokens_padded2 = moe_align_block_size(
748
+ topk_ids=topk_ids,
749
+ block_size=w2_block_size,
750
+ num_experts=self.global_num_experts,
751
+ expert_map=expert_map,
752
+ ignore_invalid_experts=True,
753
+ )
754
+ moe_kwargs2.update(
755
+ sorted_ids=sorted_ids2,
756
+ expert_ids=expert_ids2,
757
+ num_tokens_padded=num_tokens_padded2,
758
+ )
759
+
760
+ return moe_kwargs1, moe_kwargs2, scatter_idx
761
762
def apply(
763
self,
@@ -832,7 +783,7 @@ class HummingIndexedExperts(HummingExpertsBase):
783
object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
784
path (Humming quantizes the w13 input itself); for block-FP8 activations
785
it carries the scale computed before dispatch, which is forwarded to
835
- may_quant_input so Humming skips the redundant w13 quantization. The
786
+ process_input so Humming skips the redundant w13 quantization. The
787
output is written into workspace13 via the buffer management.
788
"""
789
assert not apply_router_weight_on_input
@@ -847,13 +798,13 @@ class HummingIndexedExperts(HummingExpertsBase):
798
)
799
buffers["output"] = output
800
850
- moe_kwargs1, moe_kwargs2 = self.prepare_humming_moe_kwargs(
801
+ moe_kwargs1, moe_kwargs2, scatter_idx = self.prepare_humming_moe_kwargs(
802
topk_ids=topk_ids,
803
expert_map=expert_map,
804
expert_tokens_meta=expert_tokens_meta,
805
)
806
856
- inputs, input_scale = self.quantize_input(
807
+ inputs, input_scale, input_scale_2 = self.process_input(
808
"w13",
809
inputs=hidden_states,
810
input_scale=a1q_scale,
@@ -865,52 +816,25 @@ class HummingIndexedExperts(HummingExpertsBase):
816
inputs=inputs,
817
weight=w1,
818
input_scale=input_scale,
819
+ input_scale_2=input_scale_2,
820
outputs=buffers["gate_up_output"],
821
**moe_kwargs1,
822
)
823
872
- # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view.
873
- # The fused path and mul_sum consume tokens directly. The unfused
874
- # activation consumes the expanded row count (tokens * topk).
875
- valid_tokens = None
876
- topk = topk_ids.size(1)
877
- if (
878
- expert_tokens_meta is not None
879
- and expert_tokens_meta.psum_recv_per_rank is not None
880
- ):
881
- valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]
882
-
883
- if self.fused_situ_quant_enabled(activation):
884
- # Fused SITU + FP8 quant (per-token or block-FP8 group-128) straight
885
- # into the w2 input, skipping the bf16 activation_output round-trip.
886
- inputs, input_scale = self.fused_situ_quant(
887
- gate_up_output=buffers["gate_up_output"],
888
- quanted_down_input=buffers["quanted_down_input"],
889
- num_valid_tokens=valid_tokens,
890
- topk=topk,
891
- )
892
- else:
893
- valid_token_counts = (
894
- valid_tokens * topk if valid_tokens is not None else None
895
- )
896
- self.apply_activation(
897
- activation=activation,
898
- input=buffers["gate_up_output"],
899
- output=buffers["activation_output"],
900
- valid_token_counts=valid_token_counts,
901
- )
902
-
903
- inputs, input_scale = self.quantize_input(
904
- "w2",
905
- inputs=buffers["activation_output"],
906
- quanted_input=buffers.get("quanted_down_input", None),
907
- )
824
+ inputs, input_scale, input_scale_2 = self.process_input(
825
+ "w2",
826
+ inputs=buffers["gate_up_output"],
827
+ quanted_input=buffers["quanted_down_input"],
828
+ activation=activation,
829
+ scatter_idx=scatter_idx,
830
+ )
831
832
self.humming_forward(
833
"w2",
834
inputs=inputs,
835
weight=w2,
836
input_scale=input_scale,
837
+ input_scale_2=input_scale_2,
838
outputs=buffers["down_output"].view(-1, hidden_states.size(-1)),
839
**moe_kwargs2,
840
)
@@ -918,6 +842,10 @@ class HummingIndexedExperts(HummingExpertsBase):
842
# expert_map masks any non-local id; num_valid_tokens bounds the
843
# persistent kernel to the real token rows [0, num_recv) so the padding
844
# tail is never iterated (CUDA-graph-safe device scalar).
845
+ valid_tokens = None
846
+ if expert_tokens_meta and expert_tokens_meta.psum_recv_per_rank is not None:
847
+ valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]
848
+
849
moe_fused_mul_sum(
850
inputs=buffers["down_output"].view(*topk_ids.shape, -1),
851
topk_weights=topk_weights,
@@ -983,21 +911,33 @@ class HummingGroupedExperts(HummingExpertsBase):
911
)
912
buffers["output"] = output
913
986
- hidden_states, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute(
987
- hidden_states=hidden_states,
988
- a1q_scale=a1q_scale,
989
- topk_ids=topk_ids,
990
- n_expert=global_num_experts,
991
- n_local_expert=self.num_experts,
992
- expert_map=expert_map,
993
- scratch=self._get_permute_scratch(topk_ids.size(1)),
994
- )
914
+ if a1q_scale is None:
915
+ scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True)
916
+ assert scratch is not None
917
+ scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch)
918
+ expert_offsets, scatter_idx = scatter_metadata
919
+ inv_perm = scatter_idx.flatten()
920
+ num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None
921
+ else:
922
+ hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
923
+ hidden_states=hidden_states,
924
+ a1q_scale=a1q_scale,
925
+ topk_ids=topk_ids,
926
+ n_expert=global_num_experts,
927
+ n_local_expert=self.num_experts,
928
+ expert_map=expert_map,
929
+ scratch=self._get_permute_scratch(topk_ids.size(1)),
930
+ )
931
+ scatter_idx = None
932
+ num_valid_tokens = None
933
996
- inputs, input_scale = self.quantize_input(
934
+ inputs, input_scale, input_scale_2 = self.process_input(
935
"w13",
936
inputs=hidden_states,
937
input_scale=a1q_scale,
938
quanted_input=buffers.get("quanted_gate_up_input", None),
939
+ scatter_idx=scatter_idx,
940
+ num_valid_tokens=num_valid_tokens,
941
)
942
943
self.humming_forward(
@@ -1005,28 +945,20 @@ class HummingGroupedExperts(HummingExpertsBase):
945
inputs=inputs,
946
weight=w1,
947
input_scale=input_scale,
948
+ input_scale_2=input_scale_2,
949
outputs=buffers["gate_up_output"],
950
valid_shape_m=valid_shape_m,
1010
- expert_layout=expert_first_token_offset,
951
+ expert_layout=expert_offsets,
952
compute_config=self.compute_config_str,
953
tuning_config=self.w13_tuning_config_str,
954
)
955
1015
- self.apply_activation(
1016
- activation=activation,
1017
- input=buffers["gate_up_output"],
1018
- output=buffers["activation_output"],
1019
- valid_token_counts=(
1020
- expert_first_token_offset[-1:].to(torch.int32)
1021
- if expert_tokens_meta is None
1022
- else None
1023
- ),
1024
- )
1025
-
1026
- inputs, input_scale = self.quantize_input(
956
+ inputs, input_scale, input_scale_2 = self.process_input(
957
"w2",
1028
- inputs=buffers["activation_output"],
1029
- quanted_input=buffers.get("quanted_down_input", None),
958
+ inputs=buffers["gate_up_output"],
959
+ quanted_input=buffers["quanted_down_input"],
960
+ activation=activation,
961
+ num_valid_tokens=expert_offsets[-1:],
962
)
963
964
self.humming_forward(
@@ -1034,9 +966,10 @@ class HummingGroupedExperts(HummingExpertsBase):
966
inputs=inputs,
967
weight=w2,
968
input_scale=input_scale,
969
+ input_scale_2=input_scale_2,
970
outputs=buffers["down_output"],
971
valid_shape_m=valid_shape_m,
1039
- expert_layout=expert_first_token_offset,
972
+ expert_layout=expert_offsets,
973
compute_config=self.compute_config_str,
974
tuning_config=self.w2_tuning_config_str,
975
)
@@ -1046,7 +979,7 @@ class HummingGroupedExperts(HummingExpertsBase):
979
permuted_hidden_states=buffers["down_output"].view(*topk_ids.shape, -1),
980
topk_weights=topk_weights,
981
inv_permuted_idx=inv_perm,
1049
- expert_first_token_offset=expert_first_token_offset,
982
+ expert_first_token_offset=expert_offsets,
983
)
984
985
@@ -1088,7 +1021,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1021
object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
1022
path (Humming quantizes the w13 input itself); for block-FP8 activations
1023
it carries the scale computed before dispatch, which is forwarded to
1091
- may_quant_input so Humming skips the redundant w13 quantization. The
1024
+ process_input so Humming skips the redundant w13 quantization. The
1025
output is written into workspace13 via the buffer management.
1026
"""
1027
assert not apply_router_weight_on_input
@@ -1111,11 +1044,12 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1044
)
1045
buffers["down_output"] = output
1046
1114
- inputs, input_scale = self.quantize_input(
1047
+ inputs, input_scale, input_scale_2 = self.process_input(
1048
"w13",
1049
inputs=hidden_states,
1050
input_scale=a1q_scale,
1051
quanted_input=buffers.get("quanted_gate_up_input", None),
1052
+ expert_tokens=expert_num_tokens,
1053
)
1054
1055
self.humming_forward(
@@ -1123,6 +1057,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1057
inputs=inputs,
1058
weight=w1,
1059
input_scale=input_scale,
1060
+ input_scale_2=input_scale_2,
1061
outputs=buffers["gate_up_output"],
1062
valid_shape_m=valid_shape_m,
1063
expert_layout=expert_num_tokens,
@@ -1130,16 +1065,12 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1065
tuning_config=self.w13_tuning_config_str,
1066
)
1067
1133
- self.apply_activation(
1134
- activation=activation,
1135
- input=buffers["gate_up_output"],
1136
- output=buffers["activation_output"],
1137
- )
1138
-
1139
- inputs, input_scale = self.quantize_input(
1068
+ inputs, input_scale, input_scale_2 = self.process_input(
1069
"w2",
1141
- inputs=buffers["activation_output"],
1142
- quanted_input=buffers.get("quanted_down_input", None),
1070
+ inputs=buffers["gate_up_output"],
1071
+ quanted_input=buffers["quanted_down_input"],
1072
+ expert_tokens=expert_num_tokens,
1073
+ activation=activation,
1074
)
1075
1076
self.humming_forward(
@@ -1147,6 +1078,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1078
inputs=inputs,
1079
weight=w2,
1080
input_scale=input_scale,
1081
+ input_scale_2=input_scale_2,
1082
outputs=output.view(-1, hidden_states.size(-1)),
1083
valid_shape_m=valid_shape_m,
1084
expert_layout=expert_num_tokens,
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+67
-1
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Literal, overload
5
+
6
import torch
7
8
from vllm import _custom_ops as ops
@@ -8,6 +10,7 @@ from vllm.triton_utils import triton
10
from vllm.utils.math_utils import round_up
11
12
13
+@overload
14
def moe_align_block_size(
15
topk_ids: torch.Tensor,
16
block_size: int,
@@ -15,7 +18,52 @@ def moe_align_block_size(
18
expert_map: torch.Tensor | None = None,
19
pad_sorted_ids: bool = False,
20
ignore_invalid_experts: bool = False,
18
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
21
+ *,
22
+ return_scatter_idx: Literal[False] = False,
23
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
24
+
25
+
26
+@overload
27
+def moe_align_block_size(
28
+ topk_ids: torch.Tensor,
29
+ block_size: int,
30
+ num_experts: int,
31
+ expert_map: torch.Tensor | None = None,
32
+ pad_sorted_ids: bool = False,
33
+ ignore_invalid_experts: bool = False,
34
+ *,
35
+ return_scatter_idx: Literal[True],
36
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ...
37
+
38
+
39
+@overload
40
+def moe_align_block_size(
41
+ topk_ids: torch.Tensor,
42
+ block_size: int,
43
+ num_experts: int,
44
+ expert_map: torch.Tensor | None = None,
45
+ pad_sorted_ids: bool = False,
46
+ ignore_invalid_experts: bool = False,
47
+ *,
48
+ return_scatter_idx: bool,
49
+) -> (
50
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]
51
+ | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
52
+): ...
53
+
54
+
55
+def moe_align_block_size(
56
+ topk_ids: torch.Tensor,
57
+ block_size: int,
58
+ num_experts: int,
59
+ expert_map: torch.Tensor | None = None,
60
+ pad_sorted_ids: bool = False,
61
+ ignore_invalid_experts: bool = False,
62
+ return_scatter_idx: bool = False,
63
+) -> (
64
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor]
65
+ | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
66
+):
67
"""Aligns the token distribution across experts to be compatible with block
68
size for matrix multiplication.
69
@@ -42,6 +90,8 @@ def moe_align_block_size(
90
as -1. When True, all invalid expert_ids in topk_ids will be ignored
91
and will not participate in counting or ranking, and there will be no
92
-1 in expert_ids.
93
+ return_scatter_idx: Whether to additionally return a masked identity
94
+ mapping for the original routed rows.
95
96
Returns:
97
- sorted_token_ids: A tensor containing the sorted token indices according
@@ -49,6 +99,10 @@ def moe_align_block_size(
99
- expert_ids: A tensor indicating the assigned expert index for each block.
100
- num_tokens_post_padded: The total number of tokens after padding,
101
ensuring divisibility by block_size.
102
+ - scatter_idx: Only returned when return_scatter_idx=True. An int32 tensor
103
+ shaped [topk_ids.numel(), 1], containing the original routed row index
104
+ for valid routes and -1 otherwise. With expert_map, this requires
105
+ ignore_invalid_experts=True.
106
107
This function pads the number of tokens that each expert needs to process
108
so that it is divisible by block_size.
@@ -71,6 +125,15 @@ def moe_align_block_size(
125
by block_size for proper block matrix operations.
126
127
"""
128
+ if return_scatter_idx and expert_map is not None and not ignore_invalid_experts:
129
+ raise ValueError("scatter_idx with expert_map requires ignore_invalid_experts")
130
+ scatter_idx = None
131
+ if return_scatter_idx:
132
+ scatter_idx = torch.empty(
133
+ (topk_ids.numel(), 1),
134
+ dtype=torch.int32,
135
+ device=topk_ids.device,
136
+ )
137
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
138
if pad_sorted_ids:
139
max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
@@ -95,11 +158,14 @@ def moe_align_block_size(
158
expert_ids,
159
num_tokens_post_pad,
160
expert_map if ignore_invalid_experts else None,
161
+ scatter_idx,
162
)
163
164
if expert_map is not None and not ignore_invalid_experts:
165
expert_ids = expert_map[expert_ids]
166
167
+ if scatter_idx is not None:
168
+ return sorted_ids, expert_ids, num_tokens_post_pad, scatter_idx
169
return sorted_ids, expert_ids, num_tokens_post_pad
170
171
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+31
@@ -107,6 +107,37 @@ class MoEPermuteScratch:
107
return topk_ids_int32
108
109
110
+def moe_prepare_scatter(
111
+ topk_ids: torch.Tensor,
112
+ expert_map: torch.Tensor | None,
113
+ scratch: MoEPermuteScratch,
114
+) -> tuple[torch.Tensor, torch.Tensor]:
115
+ """Generate expert offsets and shared scatter/unpermute destination indices."""
116
+ assert topk_ids.device == scratch.device
117
+ assert topk_ids.size(0) <= scratch.max_num_tokens
118
+ assert topk_ids.size(1) == scratch.topk
119
+ n_token, topk = topk_ids.shape
120
+ expanded_rows = topk_ids.numel()
121
+ inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk)
122
+ if expanded_rows == 0:
123
+ scratch.expert_first_token_offset.zero_()
124
+ return scratch.expert_first_token_offset, inverse
125
+ torch.ops._moe_C.moe_prepare_scatter(
126
+ scratch.prepare_topk_ids(topk_ids),
127
+ scratch.token_expert_indices_view(n_token),
128
+ expert_map,
129
+ scratch.num_experts,
130
+ scratch.num_local_experts,
131
+ scratch.expert_first_token_offset,
132
+ inverse,
133
+ scratch.sort_workspace,
134
+ scratch.permuted_experts_id[:expanded_rows].view(n_token, topk),
135
+ scratch.sorted_row_idx[:expanded_rows].view(n_token, topk),
136
+ scratch.topk_ids_for_sort[:expanded_rows].view(n_token, topk),
137
+ )
138
+ return scratch.expert_first_token_offset, inverse
139
+
140
+
141
def get_moe_permute_scratch(
142
*,
143
max_num_tokens: int,
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+2
-2
@@ -603,7 +603,7 @@ def convert_to_fp8_moe_kernel_format(
603
w13.is_shuffled = True
604
w2.is_shuffled = True
605
elif fp8_backend == Fp8MoeBackend.HUMMING:
606
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
606
+ from vllm.model_executor.layers.quantization.utils.humming import (
607
convert_to_humming_moe_kernel_format,
608
)
609
@@ -739,7 +739,7 @@ def make_fp8_moe_quant_config(
739
)
740
elif fp8_backend == Fp8MoeBackend.HUMMING:
741
from vllm.model_executor.layers.fused_moe import RoutedExperts
742
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
742
+ from vllm.model_executor.layers.quantization.utils.humming import (
743
get_humming_moe_quant_config,
744
)
745
vllm/model_executor/layers/fused_moe/oracle/int8.py
+2
-2
@@ -193,7 +193,7 @@ def make_int8_moe_quant_config(
193
194
if int8_backend == Int8MoeBackend.HUMMING:
195
from vllm.model_executor.layers.fused_moe import RoutedExperts
196
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
196
+ from vllm.model_executor.layers.quantization.utils.humming import (
197
get_humming_moe_quant_config,
198
)
199
@@ -254,7 +254,7 @@ def convert_to_int8_moe_kernel_format(
254
) -> tuple[torch.Tensor, torch.Tensor]:
255
"""Convert INT8 MoE weights to backend-specific kernel format."""
256
if int8_backend == Int8MoeBackend.HUMMING:
257
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
257
+ from vllm.model_executor.layers.quantization.utils.humming import (
258
convert_to_humming_moe_kernel_format,
259
)
260
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+1
-1
@@ -1456,7 +1456,7 @@ def convert_to_wna16_moe_kernel_format(
1456
"""
1457
if backend == WNA16MoEBackend.HUMMING:
1458
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
1459
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
1459
+ from vllm.model_executor.layers.quantization.utils.humming import (
1460
convert_to_humming_moe_kernel_format,
1461
)
1462
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+3
-3
@@ -844,7 +844,7 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
844
sf_block_size = 32 # mxfp4 block size
845
846
if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
847
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
847
+ from vllm.model_executor.layers.quantization.utils.humming import (
848
convert_to_humming_moe_kernel_format,
849
)
850
@@ -1428,7 +1428,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1428
)
1429
1430
if mxfp4_backend == Mxfp4MoeBackend.HUMMING:
1431
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
1431
+ from vllm.model_executor.layers.quantization.utils.humming import (
1432
convert_to_humming_moe_kernel_format,
1433
)
1434
@@ -1966,7 +1966,7 @@ def make_mxfp4_moe_quant_config(
1966
gemm1_clamp_limit=swiglu_limit,
1967
)
1968
elif mxfp4_backend == Mxfp4MoeBackend.HUMMING:
1969
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
1969
+ from vllm.model_executor.layers.quantization.utils.humming import (
1970
get_humming_moe_quant_config,
1971
)
1972
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+2
-2
@@ -409,7 +409,7 @@ def convert_to_nvfp4_moe_kernel_format(
409
is_act_and_mul=is_act_and_mul,
410
)
411
elif nvfp4_backend == NvFp4MoeBackend.HUMMING:
412
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
412
+ from vllm.model_executor.layers.quantization.utils.humming import (
413
convert_to_humming_moe_kernel_format,
414
)
415
@@ -523,7 +523,7 @@ def make_nvfp4_moe_quant_config(
523
use_a16 = _use_a16(backend, use_a16)
524
if backend == NvFp4MoeBackend.HUMMING:
525
from vllm.model_executor.layers.fused_moe import RoutedExperts
526
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
526
+ from vllm.model_executor.layers.quantization.utils.humming import (
527
get_humming_moe_quant_config,
528
)
529
vllm/model_executor/layers/quantization/auto_awq.py
+1
-1
@@ -713,7 +713,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):
713
714
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
715
if self.wna16_moe_backend == WNA16MoEBackend.HUMMING:
716
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
716
+ from vllm.model_executor.layers.quantization.utils.humming import (
717
get_humming_moe_quant_config,
718
)
719
vllm/model_executor/layers/quantization/auto_gptq.py
+1
-1
@@ -689,7 +689,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
689
690
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
691
if self.wna16_moe_backend == WNA16MoEBackend.HUMMING:
692
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
692
+ from vllm.model_executor.layers.quantization.utils.humming import (
693
get_humming_moe_quant_config,
694
)
695
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+1
-1
@@ -474,7 +474,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
474
self, layer: torch.nn.Module
475
) -> FusedMoEQuantConfig | None:
476
if self.wna16_backend == WNA16MoEBackend.HUMMING:
477
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
477
+ from vllm.model_executor.layers.quantization.utils.humming import (
478
get_humming_moe_quant_config,
479
)
480
vllm/model_executor/layers/quantization/humming.py
+106
-52
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Callable
4
+from math import gcd
5
from typing import TYPE_CHECKING, Any
6
7
import regex as re
@@ -28,12 +29,15 @@ from vllm.model_executor.layers.quantization.base_config import (
29
QuantizationConfig,
30
QuantizeMethodBase,
31
)
31
-from vllm.model_executor.layers.quantization.utils.humming_utils import (
32
+from vllm.model_executor.layers.quantization.utils.humming import (
33
+ check_and_fallback_input_schema,
34
convert_to_humming_moe_kernel_format,
35
get_humming_linear_compute_config,
36
get_humming_moe_quant_config,
37
+ humming_update_schema_hadamard_block_size,
38
input_schema_to_quant_key,
39
make_humming_moe_kernel,
40
+ resolve_humming_layer_config,
41
select_humming_moe_experts,
42
weight_schema_to_quant_key,
43
)
@@ -227,23 +231,13 @@ class HummingConfig(QuantizationConfig):
231
return None
232
config = group_config
233
230
- layer_config = config
231
- layer_dynamic = config.get("dynamic", {})
232
- if not isinstance(layer_dynamic, dict):
233
- layer_dynamic = {}
234
- for regex, override_config in layer_dynamic.items():
235
- if regex[:1] != "+":
236
- continue
237
- if re.match(regex[2:], prefix):
238
- layer_config = config.copy()
239
- layer_config.update(override_config)
240
- break
234
+ layer_config = resolve_humming_layer_config(config, prefix)
235
236
if "quant_method" in layer_config:
237
return _hm.BaseWeightSchema.from_config(layer_config)
238
return None
239
246
- def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
240
+ def get_layer_input_config(self, config: dict[str, Any], prefix: str):
241
if self.is_layer_skipped(config, prefix):
242
return None
243
if config["quant_method"] in ["compressed-tensors", "modelopt"]:
@@ -252,10 +246,15 @@ class HummingConfig(QuantizationConfig):
246
return None
247
config = group_config
248
249
+ config = resolve_humming_layer_config(config, prefix)
250
if config.get("quant_method", None) in _hm.BaseInputSchema.INPUT_SCHEMA_MAP:
256
- return _hm.BaseInputSchema.from_config(config)
251
+ return config
252
return None
253
254
+ def get_layer_input_schema(self, config: dict[str, Any], prefix: str):
255
+ config = self.get_layer_input_config(config, prefix)
256
+ return None if config is None else _hm.BaseInputSchema.from_config(config)
257
+
258
def get_quant_config_for_layer(
259
self, prefix: str, layer_type: str
260
) -> "HummingLayerQuantizationConfig | None":
@@ -279,14 +278,25 @@ class HummingConfig(QuantizationConfig):
278
if weight_schema is not None:
279
input_schema = None
280
force_input_schema = None
281
+ allow_input_schema_fallback = True
282
283
if self.full_config:
284
- input_schema = self.get_layer_input_schema(self.full_config, prefix)
284
+ input_config = self.get_layer_input_config(self.full_config, prefix)
285
+ if input_config is not None:
286
+ allow_input_schema_fallback = input_config.pop(
287
+ "allow_fallback", True
288
+ )
289
+ input_schema = _hm.BaseInputSchema.from_config(input_config)
290
291
if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
292
quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG.copy()
293
quant_config["quant_method"] = "humming"
289
- force_input_schema = self.get_layer_input_schema(quant_config, prefix)
294
+ input_config = self.get_layer_input_config(quant_config, prefix)
295
+ if input_config is not None:
296
+ allow_input_schema_fallback = input_config.pop(
297
+ "allow_fallback", False
298
+ )
299
+ force_input_schema = _hm.BaseInputSchema.from_config(input_config)
300
if input_schema is None:
301
input_schema = force_input_schema
302
@@ -299,6 +309,7 @@ class HummingConfig(QuantizationConfig):
309
force_weight_schema=force_weight_schema,
310
force_input_schema=force_input_schema,
311
is_online_quant=is_online_quant,
312
+ allow_input_schema_fallback=allow_input_schema_fallback,
313
)
314
return None
315
@@ -318,6 +329,11 @@ class HummingConfig(QuantizationConfig):
329
elif isinstance(layer, LinearBase):
330
return UnquantizedLinearMethod()
331
elif isinstance(layer, LinearBase):
332
+ input_size = layer.input_size
333
+ partition_size = getattr(layer, "input_size_per_partition", input_size)
334
+ is_online_quant = quant_config.is_online_quant
335
+ if is_online_quant and (input_size % 32 != 0 or partition_size % 32 != 0):
336
+ return UnquantizedLinearMethod()
337
return HummingLinearMethod(quant_config)
338
elif isinstance(layer, RoutedExperts):
339
return HummingMoEMethod(quant_config, layer.moe_config)
@@ -332,6 +348,7 @@ class HummingLayerQuantizationConfig(HummingConfig):
348
force_weight_schema: "HummingWeightSchema | None" = None,
349
force_input_schema: "HummingInputSchema | None" = None,
350
is_online_quant: bool = False,
351
+ allow_input_schema_fallback: bool = True,
352
):
353
self.weight_schema = weight_schema
354
if input_schema is None:
@@ -340,6 +357,7 @@ class HummingLayerQuantizationConfig(HummingConfig):
357
self.force_weight_schema = force_weight_schema
358
self.force_input_schema = force_input_schema
359
self.is_online_quant = is_online_quant
360
+ self.allow_input_schema_fallback = allow_input_schema_fallback
361
362
@classmethod
363
def from_config(cls, config):
@@ -373,21 +391,13 @@ class HummingLinearMethod(LinearMethodBase):
391
if is_unquantized and self.is_online_quant:
392
# online quant (fp16/bf16 -> quant_type)
393
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
376
- f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
377
- has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
378
- tensor_list = _hm.quantize_weight(
379
- weight=loaded_weight,
380
- dtype=self.weight_schema.b_dtype,
381
- scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
382
- group_size=self.weight_schema.weight_scale_group_size,
383
- has_zero_point=self.weight_schema.has_zero_point,
384
- has_global_scale=has_global_scale,
385
- is_fp_zero_point=self.weight_schema.is_fp_zero_point,
386
- pack=True,
394
+ tensors = self.weight_schema.quant_tensor(
395
+ loaded_weight.to(param.device),
396
+ self.weight_schema,
397
+ layer.param_dtype,
398
)
399
389
- key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
390
- for key, tensor in zip(key_list, tensor_list):
400
+ for key, tensor in tensors.items():
401
if tensor is None or tensor.nelement() == 0:
402
continue
403
param = getattr(layer, key)
@@ -459,6 +469,18 @@ class HummingLinearMethod(LinearMethodBase):
469
layer.output_partition_sizes = output_partition_sizes
470
layer.extra_weight_attrs = extra_weight_attrs.copy()
471
472
+ input_schema = self.force_input_schema or self.input_schema
473
+ input_schema = input_schema.to_humming_schema(params_dtype)
474
+ for name in ("weight_schema", "force_weight_schema"):
475
+ schema = getattr(self, name)
476
+ if getattr(schema, "hadamard_block_size", None) == -1:
477
+ schema = humming_update_schema_hadamard_block_size(
478
+ weight_schema=schema,
479
+ input_schema=input_schema,
480
+ shape_k=input_size_per_partition,
481
+ )
482
+ setattr(self, name, schema)
483
+
484
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
485
new_weight_loader = self.prepare_weight_loader(layer, weight_loader)
486
extra_weight_attrs["weight_loader"] = new_weight_loader
@@ -502,7 +524,10 @@ class HummingLinearMethod(LinearMethodBase):
524
return None
525
526
# convert from checkpoint format to humming format
505
- if not isinstance(self.weight_schema, _hm.HummingWeightSchema):
527
+ is_humming_weight = isinstance(self.weight_schema, _hm.HummingWeightSchema)
528
+ is_humming_input = isinstance(self.input_schema, _hm.HummingInputSchema)
529
+
530
+ if not is_humming_weight or not is_humming_input:
531
self.weight_schema, tensors = self.weight_schema.convert_humming(
532
tensors=layer.state_dict(),
533
shape_n_stacks=layer.output_partition_sizes,
@@ -510,12 +535,13 @@ class HummingLinearMethod(LinearMethodBase):
535
param_dtype=layer.param_dtype,
536
)
537
513
- self.input_schema, _ = self.input_schema.convert_humming(
538
+ self.input_schema, input_tensors = self.input_schema.convert_humming(
539
tensors=layer.state_dict(),
540
shape_n_stacks=layer.output_partition_sizes,
541
shape_k_stacks=[layer.input_size_per_partition],
542
param_dtype=layer.param_dtype,
543
)
544
+ tensors.update(input_tensors)
545
546
for name, _ in list(layer.named_parameters()):
547
delattr(layer, name)
@@ -530,11 +556,15 @@ class HummingLinearMethod(LinearMethodBase):
556
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
557
force_requant = self.force_weight_schema is not None
558
if force_requant and self.weight_schema != self.force_weight_schema:
559
+ source_tensors = layer.state_dict()
560
tensors = self.weight_schema.requant_tensors(
534
- tensors=layer.state_dict(),
561
+ tensors=source_tensors,
562
target_weight_schema=self.force_weight_schema,
563
param_dtype=layer.param_dtype,
564
)
565
+ name = self.input_schema.static_tensor_scale_name
566
+ if name is not None:
567
+ tensors[name] = source_tensors[name]
568
569
self.weight_schema = self.force_weight_schema
570
@@ -548,6 +578,12 @@ class HummingLinearMethod(LinearMethodBase):
578
579
del tensors
580
581
+ self.input_schema = check_and_fallback_input_schema(
582
+ weight_schema=self.weight_schema,
583
+ input_schema=self.input_schema,
584
+ param_dtype=layer.param_dtype,
585
+ allow_fallback=self.quant_config.allow_input_schema_fallback,
586
+ )
587
self.layer_config = _hm.prepare_layer_config(
588
shape_n=layer.output_partition_sizes_sum,
589
shape_k=layer.input_size_per_partition,
@@ -557,10 +593,16 @@ class HummingLinearMethod(LinearMethodBase):
593
pad_k_to_multiple=128,
594
has_bias=layer.has_bias,
595
torch_dtype=layer.param_dtype,
596
+ device=layer.weight.device,
597
)
561
- tensors = _hm.transform_humming_tensors(
562
- self.layer_config, dict(layer.named_parameters())
563
- )
598
+ source_tensors = dict(layer.named_parameters())
599
+ tensors = _hm.transform_humming_tensors(self.layer_config, source_tensors)
600
+ name = self.input_schema.static_tensor_scale_name
601
+ if name is not None:
602
+ tensors[name] = source_tensors[name]
603
+ if "bias" in tensors:
604
+ zero_bias = torch.zeros_like(tensors["bias"])
605
+ layer.register_buffer("zero_bias", zero_bias)
606
for name, _ in list(layer.named_parameters()):
607
delattr(layer, name)
608
for name, tensor in tensors.items():
@@ -583,8 +625,11 @@ class HummingLinearMethod(LinearMethodBase):
625
weight=layer.weight,
626
weight_scale=getattr(layer, "weight_scale", None),
627
zero_point=getattr(layer, "zero_point", None),
586
- bias=getattr(layer, "bias", None),
628
+ bias=getattr(layer, "bias" if bias is not None else "zero_bias", None),
629
weight_scale_2=getattr(layer, "weight_scale_2", None),
630
+ input_scale=getattr(layer, "input_scale", None),
631
+ input_scale_2=getattr(layer, "input_scale_2", None),
632
+ hadamard_block_size=self.weight_schema.hadamard_block_size,
633
locks=self.locks,
634
compute_config=self.compute_config,
635
)
@@ -606,11 +651,15 @@ class HummingMoEMethod(FusedMoEMethodBase):
651
# Derive QuantKeys from humming schemas.
652
# Prefer force schemas (the final format after requant) over base.
653
weight_key = weight_schema_to_quant_key(
609
- self.force_weight_schema or self.weight_schema
654
+ self.force_weight_schema or self.weight_schema, moe.in_dtype
655
)
611
- activation_key = input_schema_to_quant_key(
612
- self.force_input_schema or self.input_schema
656
+ runtime_input_schema = check_and_fallback_input_schema(
657
+ weight_schema=self.force_weight_schema or self.weight_schema,
658
+ input_schema=self.force_input_schema or self.input_schema,
659
+ param_dtype=moe.in_dtype,
660
+ allow_fallback=quant_config.allow_input_schema_fallback,
661
)
662
+ activation_key = input_schema_to_quant_key(runtime_input_schema, moe.in_dtype)
663
664
# Select Humming MoE experts
665
self.experts_cls = select_humming_moe_experts(
@@ -634,22 +683,14 @@ class HummingMoEMethod(FusedMoEMethodBase):
683
# online quant (fp16/bf16 -> quant_type)
684
if is_unquantized:
685
assert isinstance(self.weight_schema, _hm.HummingWeightSchema)
637
- f16_dtype = _hm.DataType.from_torch_dtype(layer.param_dtype)
638
- has_global_scale = "TENSOR" in str(self.weight_schema.weight_scale_type)
639
- tensor_list = _hm.quantize_weight(
640
- weight=loaded_weight,
641
- dtype=self.weight_schema.b_dtype,
642
- scale_dtype=self.weight_schema.bs_dtype or f16_dtype,
643
- group_size=self.weight_schema.weight_scale_group_size,
644
- has_zero_point=self.weight_schema.has_zero_point,
645
- has_global_scale=has_global_scale,
646
- is_fp_zero_point=self.weight_schema.is_fp_zero_point,
647
- pack=True,
686
+ tensors = self.weight_schema.quant_tensor(
687
+ loaded_weight.to(param.device),
688
+ self.weight_schema,
689
+ layer.param_dtype,
690
)
691
650
- key_list = ["weight", "weight_scale", "zero_point", "global_scale"]
692
success = True
652
- for key, tensor in zip(key_list, tensor_list):
693
+ for key, tensor in tensors.items():
694
if tensor is None or tensor.nelement() == 0:
695
continue
696
sublayer_name = "w2" if shard_id == "w2" else "w13"
@@ -695,6 +736,18 @@ class HummingMoEMethod(FusedMoEMethodBase):
736
layer.num_experts = num_experts
737
layer.param_dtype = params_dtype
738
layer.intermediate_size = intermediate_size_per_partition
739
+ input_schema = self.force_input_schema or self.input_schema
740
+ input_schema = input_schema.to_humming_schema(params_dtype)
741
+ shape_k = gcd(hidden_size, intermediate_size_per_partition)
742
+ for name in ("weight_schema", "force_weight_schema"):
743
+ schema = getattr(self, name)
744
+ if getattr(schema, "hadamard_block_size", None) == -1:
745
+ schema = humming_update_schema_hadamard_block_size(
746
+ weight_schema=schema,
747
+ input_schema=input_schema,
748
+ shape_k=shape_k,
749
+ )
750
+ setattr(self, name, schema)
751
weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader)
752
weight_loader = self.prepare_weight_loader(layer, weight_loader)
753
extra_weight_attrs["weight_loader"] = weight_loader
@@ -761,6 +814,7 @@ class HummingMoEMethod(FusedMoEMethodBase):
814
weight_schema=self.weight_schema,
815
input_schema=self.input_schema,
816
force_weight_schema=self.force_weight_schema,
817
+ allow_input_schema_fallback=self.quant_config.allow_input_schema_fallback,
818
)
819
820
# Build the MoE kernel
vllm/model_executor/layers/quantization/moe_wna16.py
+1
-1
@@ -389,7 +389,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
389
self, layer: RoutedExperts
390
) -> FusedMoEQuantConfig | None:
391
if self.wna16_backend == WNA16MoEBackend.HUMMING:
392
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
392
+ from vllm.model_executor.layers.quantization.utils.humming import (
393
get_humming_moe_quant_config,
394
)
395
vllm/model_executor/layers/quantization/utils/humming/__init__.py
new
+51
@@ -0,0 +1,51 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Humming quantization integration."""
4
+
5
+from vllm.model_executor.layers.quantization.utils.humming.activation import (
6
+ get_humming_activation,
7
+)
8
+from vllm.model_executor.layers.quantization.utils.humming.linear import (
9
+ apply_humming_linear,
10
+ convert_linear_layer_to_humming_standard,
11
+ get_humming_linear_compute_config,
12
+ prepare_humming_linear_layer_config,
13
+)
14
+from vllm.model_executor.layers.quantization.utils.humming.moe import (
15
+ HummingMoEQuantConfig,
16
+ convert_to_humming_moe_kernel_format,
17
+ get_humming_moe_quant_config,
18
+ make_humming_moe_kernel,
19
+ make_humming_moe_quant_config,
20
+ select_humming_moe_experts,
21
+)
22
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
23
+ check_and_fallback_input_schema,
24
+ humming_is_layer_skipped,
25
+ humming_update_schema_hadamard_block_size,
26
+ input_schema_to_quant_key,
27
+ quant_key_to_input_schema,
28
+ resolve_humming_layer_config,
29
+ weight_schema_to_quant_key,
30
+)
31
+
32
+__all__ = [
33
+ "check_and_fallback_input_schema",
34
+ "get_humming_activation",
35
+ "HummingMoEQuantConfig",
36
+ "get_humming_moe_quant_config",
37
+ "humming_is_layer_skipped",
38
+ "humming_update_schema_hadamard_block_size",
39
+ "make_humming_moe_quant_config",
40
+ "apply_humming_linear",
41
+ "convert_linear_layer_to_humming_standard",
42
+ "get_humming_linear_compute_config",
43
+ "prepare_humming_linear_layer_config",
44
+ "make_humming_moe_kernel",
45
+ "select_humming_moe_experts",
46
+ "input_schema_to_quant_key",
47
+ "quant_key_to_input_schema",
48
+ "resolve_humming_layer_config",
49
+ "weight_schema_to_quant_key",
50
+ "convert_to_humming_moe_kernel_format",
51
+]
vllm/model_executor/layers/quantization/utils/humming/activation.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MoE activation expressions for Humming input processing."""
4
+
5
+from dataclasses import replace
6
+
7
+from vllm.model_executor.layers.fused_moe.activation import (
8
+ ApplyMoEActivationConfig,
9
+ MoEActivation,
10
+)
11
+
12
+_SILU_IMPL = "a / (1.f + expf(-a))"
13
+_GELU_IMPL = "0.5f * a * (1.f + erff(a * 0.7071067811865476f))"
14
+_RELU2_IMPL = "fmaxf(a, 0.f) * fmaxf(a, 0.f)"
15
+_GELU_TANH_IMPL = """[](float a) {
16
+ const float cubic = 0.044715f * a * a * a;
17
+ const float inner = 0.7978845608028654f * (a + cubic);
18
+ return 0.5f * a * (1.f + tanhf(inner));
19
+}(a)"""
20
+_CLAMPED_SILU_IMPL = """[](float a, float b) {{
21
+ const float gate = fminf(a, {limit});
22
+ const float up = fmaxf(fminf(b, {limit}), -{limit});
23
+ const float activated_gate = gate / (1.f + expf(-{alpha} * gate));
24
+ return activated_gate * (up + {beta});
25
+}}(a, b)"""
26
+
27
+# Parameterized expressions are resolved by get_humming_activation().
28
+_HUMMING_ACTIVATIONS: dict[MoEActivation, dict[str, str]] = {
29
+ MoEActivation.SILU: {
30
+ "activation_type": "binary_split",
31
+ "activation_impl": f"({_SILU_IMPL}) * b",
32
+ },
33
+ MoEActivation.GELU: {
34
+ "activation_type": "binary_split",
35
+ "activation_impl": f"({_GELU_IMPL}) * b",
36
+ },
37
+ MoEActivation.GELU_TANH: {
38
+ "activation_type": "binary_split",
39
+ "activation_impl": f"({_GELU_TANH_IMPL}) * b",
40
+ },
41
+ MoEActivation.RELU2: {
42
+ "activation_type": "binary_split",
43
+ "activation_impl": f"({_RELU2_IMPL}) * b",
44
+ },
45
+ MoEActivation.SWIGLUOAI: {
46
+ "activation_type": "binary_interleaved",
47
+ "activation_impl": _CLAMPED_SILU_IMPL,
48
+ },
49
+ MoEActivation.SITU: {
50
+ "activation_type": "binary_split",
51
+ "activation_impl": """[](float a, float b) {{
52
+ const float gate = {situ_beta} * tanhf(a / {situ_beta});
53
+ const float activated_gate = gate / (1.f + expf(-a));
54
+ const float up = {situ_up};
55
+ return activated_gate * up;
56
+ }}(a, b)""",
57
+ },
58
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE: {
59
+ "activation_type": "binary_split",
60
+ "activation_impl": _CLAMPED_SILU_IMPL,
61
+ },
62
+ MoEActivation.SWIGLUSTEP: {
63
+ "activation_type": "binary_split",
64
+ "activation_impl": """[](float a, float b) {
65
+ const float activated_gate = a / (1.f + expf(-a));
66
+ const float gate = fminf(activated_gate, 7.f);
67
+ const float up = fmaxf(fminf(b, 7.f), -7.f);
68
+ return gate * up;
69
+ }(a, b)""",
70
+ },
71
+ MoEActivation.SILU_NO_MUL: {
72
+ "activation_type": "unary",
73
+ "activation_impl": _SILU_IMPL,
74
+ },
75
+ MoEActivation.GELU_NO_MUL: {
76
+ "activation_type": "unary",
77
+ "activation_impl": _GELU_IMPL,
78
+ },
79
+ MoEActivation.GELU_TANH_NO_MUL: {
80
+ "activation_type": "unary",
81
+ "activation_impl": _GELU_TANH_IMPL,
82
+ },
83
+ MoEActivation.RELU2_NO_MUL: {
84
+ "activation_type": "unary",
85
+ "activation_impl": _RELU2_IMPL,
86
+ },
87
+}
88
+
89
+
90
+def get_humming_activation(
91
+ activation: MoEActivation,
92
+ config: ApplyMoEActivationConfig | None = None,
93
+) -> dict[str, str]:
94
+ config = config or ApplyMoEActivationConfig()
95
+ use_clamped_silu = False
96
+ if activation == MoEActivation.SWIGLUOAI:
97
+ config = ApplyMoEActivationConfig(clamp_limit=7.0, alpha=1.702, beta=1.0)
98
+ use_clamped_silu = True
99
+ elif activation == MoEActivation.SILU and config.clamp_limit is not None:
100
+ config = replace(config, alpha=1.0, beta=0.0)
101
+ use_clamped_silu = True
102
+ elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
103
+ if config.clamp_limit is None:
104
+ raise ValueError("SWIGLUOAI_UNINTERLEAVE requires clamp_limit")
105
+ use_clamped_silu = True
106
+
107
+ definition = _HUMMING_ACTIVATIONS[activation].copy()
108
+ if activation == MoEActivation.SITU:
109
+ if config.activation_situ_beta is None:
110
+ raise ValueError("SITU requires activation_situ_beta")
111
+ linear_beta = config.activation_situ_linear_beta
112
+ situ_up = "b"
113
+ if linear_beta is not None and linear_beta > 0:
114
+ situ_up = f"({float(linear_beta)}f) * tanhf(b / ({float(linear_beta)}f))"
115
+ situ_beta = f"({float(config.activation_situ_beta)}f)"
116
+ activation_impl = definition["activation_impl"]
117
+ activation_impl = activation_impl.format(situ_beta=situ_beta, situ_up=situ_up)
118
+ definition["activation_impl"] = activation_impl
119
+ elif use_clamped_silu:
120
+ assert config.clamp_limit is not None
121
+ limit = f"({float(config.clamp_limit)}f)"
122
+ alpha = f"({float(config.alpha)}f)"
123
+ beta = f"({float(config.beta)}f)"
124
+ activation_impl = _CLAMPED_SILU_IMPL.format(limit=limit, alpha=alpha, beta=beta)
125
+ definition["activation_impl"] = activation_impl
126
+ return definition
vllm/model_executor/layers/quantization/utils/humming/linear.py
new
+183
@@ -0,0 +1,183 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Prepare and execute Humming linear layers."""
4
+
5
+import json
6
+from typing import TYPE_CHECKING
7
+
8
+import torch
9
+
10
+from vllm import envs
11
+from vllm.model_executor.layers.linear import LinearBase
12
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
13
+ check_and_fallback_input_schema,
14
+)
15
+
16
+if TYPE_CHECKING:
17
+ from vllm.utils.humming import BaseInputSchema, LayerConfig
18
+
19
+
20
+def convert_linear_layer_to_humming_standard(
21
+ layer: LinearBase, name_map: dict[str, str]
22
+):
23
+ """Rename/reshape a linear layer's quantized params (the canonical MPLinear
24
+ layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names
25
+ and layout humming's weight schema expects (``weight`` / ``weight_scale``)."""
26
+ for name, checkpoint_name in name_map.items():
27
+ tensor = getattr(layer, checkpoint_name)
28
+ delattr(layer, checkpoint_name)
29
+
30
+ if name == "weight":
31
+ input_dim = getattr(tensor, "input_dim", 1)
32
+ output_dim = getattr(tensor, "output_dim", 0)
33
+
34
+ if input_dim == 0 and output_dim == 1:
35
+ tensor = tensor.transpose(1, 0).contiguous()
36
+ else:
37
+ assert output_dim == 0 and input_dim == 1
38
+
39
+ tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
40
+ elif name in ["weight_scale", "zero_point"]:
41
+ if getattr(tensor, "output_dim", 0) == 1:
42
+ tensor = tensor.transpose(0, 1).contiguous()
43
+ if tensor.ndim == 1:
44
+ tensor = tensor.unsqueeze(1)
45
+
46
+ tensor = tensor.view(torch.int32) if name == "zero_point" else tensor
47
+
48
+ if isinstance(tensor, torch.nn.Parameter):
49
+ param = tensor
50
+ else:
51
+ param = torch.nn.Parameter(tensor, requires_grad=False)
52
+
53
+ setattr(layer, name, param)
54
+
55
+
56
+def prepare_humming_linear_layer_config(
57
+ layer: LinearBase,
58
+ quant_config: dict,
59
+ input_quant_config: dict | None = None,
60
+ *,
61
+ input_schema: "BaseInputSchema | None" = None,
62
+) -> "LayerConfig":
63
+ from vllm.utils.humming import (
64
+ BaseInputSchema,
65
+ BaseWeightSchema,
66
+ HummingInputSchema,
67
+ prepare_layer_config,
68
+ transform_humming_tensors,
69
+ )
70
+
71
+ weight_schema = BaseWeightSchema.from_config(quant_config)
72
+ if input_schema is None:
73
+ if input_quant_config is not None:
74
+ input_schema = BaseInputSchema.from_config(input_quant_config)
75
+ else:
76
+ input_schema = HummingInputSchema()
77
+
78
+ # ReplicatedLinear has no TP partitioning and so does not set
79
+ # input_size_per_partition; for it that is just input_size. Use hasattr
80
+ # rather than getattr's default arg, which is evaluated eagerly and would
81
+ # raise on layers lacking input_size (e.g. ParallelLMHead).
82
+ if hasattr(layer, "input_size_per_partition"):
83
+ input_size_per_partition = layer.input_size_per_partition
84
+ else:
85
+ input_size_per_partition = layer.input_size
86
+ shape_k_stacks = [input_size_per_partition]
87
+ shape_n_stacks = layer.output_partition_sizes
88
+
89
+ # Step 1: convert weight and input schemas to humming standard format
90
+ source_tensors = dict(layer.named_parameters())
91
+ weight_schema, tensors = weight_schema.convert_humming(
92
+ tensors=source_tensors,
93
+ shape_n_stacks=shape_n_stacks,
94
+ shape_k_stacks=shape_k_stacks,
95
+ param_dtype=layer.params_dtype,
96
+ )
97
+ input_schema, input_tensors = input_schema.convert_humming(
98
+ tensors=source_tensors,
99
+ shape_n_stacks=shape_n_stacks,
100
+ shape_k_stacks=shape_k_stacks,
101
+ param_dtype=layer.params_dtype,
102
+ )
103
+
104
+ tensors.update(input_tensors)
105
+ layer.weight_schema = weight_schema
106
+ input_schema = check_and_fallback_input_schema(
107
+ weight_schema=weight_schema,
108
+ input_schema=input_schema,
109
+ param_dtype=layer.params_dtype,
110
+ )
111
+ for name in ("input_scale", "input_scale_2"):
112
+ if name != input_schema.static_tensor_scale_name:
113
+ tensors.pop(name, None)
114
+ input_tensors.pop(name, None)
115
+
116
+ # Step 2: transform weight (humming standard format) for forwarding.
117
+ config = prepare_layer_config(
118
+ shape_n=sum(layer.output_partition_sizes),
119
+ shape_k=input_size_per_partition,
120
+ weight_schema=weight_schema,
121
+ input_schema=input_schema,
122
+ pad_n_to_multiple=256,
123
+ pad_k_to_multiple=128,
124
+ has_bias=layer.has_bias,
125
+ torch_dtype=layer.params_dtype,
126
+ device=tensors["weight"].device,
127
+ )
128
+ tensors = transform_humming_tensors(config, tensors)
129
+ tensors.update(input_tensors)
130
+ if "bias" in tensors:
131
+ zero_bias = torch.zeros_like(tensors["bias"])
132
+ layer.register_buffer("zero_bias", zero_bias)
133
+ for name, _ in list(layer.named_parameters()):
134
+ delattr(layer, name)
135
+ for name, tensor in tensors.items():
136
+ param = torch.nn.Parameter(tensor, requires_grad=False)
137
+ setattr(layer, name, param)
138
+
139
+ return config
140
+
141
+
142
+def get_humming_linear_compute_config() -> str:
143
+ return json.dumps(
144
+ {
145
+ "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
146
+ "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
147
+ "gemm_type": "dense",
148
+ }
149
+ )
150
+
151
+
152
+def apply_humming_linear(
153
+ layer: LinearBase,
154
+ x: torch.Tensor,
155
+ *,
156
+ skip_bias_add: bool = False,
157
+ layer_config: "LayerConfig",
158
+ compute_config: str,
159
+ locks: torch.Tensor,
160
+) -> torch.Tensor:
161
+ from vllm.utils.humming import humming_forward
162
+
163
+ mode = layer_config.input_quant_mode
164
+ input_scale = getattr(layer, "input_scale", None) if mode.should_quantize else None
165
+ input_scale_2 = (
166
+ getattr(layer, "input_scale_2", None) if mode.has_secondary_scale else None
167
+ )
168
+ flatten_inputs = x.reshape(-1, x.size(-1))
169
+ output = humming_forward(
170
+ layer_config,
171
+ inputs=flatten_inputs,
172
+ weight=layer.weight,
173
+ weight_scale=getattr(layer, "weight_scale", None),
174
+ zero_point=getattr(layer, "zero_point", None),
175
+ bias=getattr(layer, "zero_bias" if skip_bias_add else "bias", None),
176
+ weight_scale_2=getattr(layer, "weight_scale_2", None),
177
+ input_scale=input_scale,
178
+ input_scale_2=input_scale_2,
179
+ hadamard_block_size=layer.weight_schema.hadamard_block_size,
180
+ locks=locks,
181
+ compute_config=compute_config,
182
+ )
183
+ return output.view(*x.shape[:-1], output.size(-1))
vllm/model_executor/layers/quantization/utils/humming/moe.py
renamed
+65
-539
@@ -1,13 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import json
3
+"""Configure, prepare weights for, and assemble Humming MoE kernels."""
4
+
5
from dataclasses import dataclass
6
from typing import TYPE_CHECKING, Any
7
7
-import regex as re
8
import torch
9
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+import vllm.model_executor.layers.quantization.utils.humming.schema as humming_schema
12
from vllm import envs
13
from vllm.logger import init_logger
14
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 (
24
HummingGroupedExperts,
25
HummingIndexedExperts,
26
)
26
-from vllm.model_executor.layers.linear import LinearBase
27
+from vllm.model_executor.layers.quantization.utils.humming.schema import (
28
+ humming_is_layer_skipped,
29
+)
30
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
- FP4_DTYPE,
31
FP8_DTYPE,
30
- INT4_DTYPE,
31
- INT8_DTYPE,
32
- MXFP_SCALE_DTYPE,
32
GroupShape,
33
QuantKey,
35
- ScaleDesc,
34
)
35
from vllm.utils.import_utils import has_humming
36
37
if TYPE_CHECKING:
38
from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
41
- from vllm.utils.humming import (
42
- AWQWeightSchema,
43
- BaseInputSchema,
44
- BaseWeightSchema,
45
- CompressedTensorsInputSchema,
46
- CompressedTensorsWeightSchema,
47
- Fp8WeightSchema,
48
- GPTQWeightSchema,
49
- HummingInputSchema,
50
- HummingWeightSchema,
51
- LayerConfig,
52
- )
53
- from vllm.utils.humming import dtypes as humming_dtypes
39
+ from vllm.utils.humming import LayerConfig
40
41
logger = init_logger(__name__)
42
@@ -59,504 +45,12 @@ logger = init_logger(__name__)
45
class HummingMoEQuantConfig(FusedMoEQuantConfig):
46
w1_humming_config: "LayerConfig"
47
w2_humming_config: "LayerConfig"
62
-
63
-
64
-if has_humming():
65
- from vllm.utils.humming import dtypes as humming_dtypes
66
-
67
- _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {
68
- humming_dtypes.float4e2m1: FP4_DTYPE,
69
- humming_dtypes.float8e4m3: FP8_DTYPE,
70
- humming_dtypes.float8e5m2: torch.float8_e5m2,
71
- humming_dtypes.int8: torch.int8,
72
- humming_dtypes.uint4: INT4_DTYPE,
73
- humming_dtypes.uint8: INT8_DTYPE,
74
- humming_dtypes.uint2: torch.uint8,
75
- humming_dtypes.uint3: torch.uint8,
76
- humming_dtypes.uint5: torch.uint8,
77
- humming_dtypes.uint6: torch.uint8,
78
- humming_dtypes.uint7: torch.uint8,
79
- }
80
-
81
- _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
82
- humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,
83
- humming_dtypes.float8e4m3: FP8_DTYPE,
84
- humming_dtypes.float16: torch.float16,
85
- humming_dtypes.bfloat16: torch.bfloat16,
86
- humming_dtypes.float32: torch.float32,
87
- }
88
-
89
-
90
-def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:
91
- """Map humming group sizes to QuantKey GroupShape.
92
-
93
- group_size: elements per group along K (col); 0 means full dimension.
94
- group_size_n: elements per group along N (row); 0 means 1 (per-row).
95
-
96
- GroupShape convention: row = N dim, col = K dim.
97
- """
98
- if group_size == 0 and group_size_n == 0:
99
- return GroupShape.PER_CHANNEL
100
-
101
- row = group_size_n if group_size_n > 0 else 1
102
- col = group_size if group_size > 0 else -1
103
- return GroupShape(row=row, col=col)
104
-
105
-
106
-# ---- HummingWeightSchema (post-conversion) --------------------------------
107
-
108
-
109
-def _humming_weight_schema_to_quant_key(
110
- schema: "HummingWeightSchema",
111
-) -> QuantKey:
112
- from vllm.utils.humming import WeightScaleType
113
-
114
- """Convert a HummingWeightSchema to a QuantKey."""
115
- dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]
116
-
117
- if schema.bs_dtype is not None:
118
- scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]
119
- else:
120
- scale_dtype = torch.float32
121
-
122
- group_shape = _group_shape(
123
- schema.weight_scale_group_size,
124
- schema.weight_scale_group_size_n,
125
- )
126
-
127
- scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
128
-
129
- scale2 = None
130
- if schema.weight_scale_type == WeightScaleType.GROUP_TENSOR:
131
- scale2 = ScaleDesc(
132
- dtype=torch.float32,
133
- static=True,
134
- group_shape=GroupShape.PER_TENSOR,
135
- )
136
-
137
- return QuantKey(
138
- dtype=dtype,
139
- scale=scale,
140
- scale2=scale2,
141
- symmetric=not schema.has_zero_point,
142
- )
143
-
144
-
145
-# ---- Checkpoint-format weight schemas (pre-conversion) --------------------
146
-
147
-
148
-def _fp8_weight_schema_to_quant_key(schema: "Fp8WeightSchema") -> QuantKey:
149
- if schema.weight_block_size is not None:
150
- gs_n, gs_k = schema.weight_block_size
151
- group_shape = GroupShape(row=gs_n, col=gs_k)
152
- else:
153
- group_shape = GroupShape.PER_CHANNEL
154
-
155
- scale = ScaleDesc(dtype=torch.float32, static=True, group_shape=group_shape)
156
- return QuantKey(dtype=FP8_DTYPE, scale=scale, symmetric=True)
157
-
158
-
159
-def _awq_weight_schema_to_quant_key(schema: "AWQWeightSchema") -> QuantKey:
160
- group_shape = _group_shape(schema.group_size)
161
- scale = ScaleDesc(
162
- dtype=torch.float16,
163
- static=True,
164
- group_shape=group_shape,
165
- )
166
- return QuantKey(
167
- dtype=INT4_DTYPE,
168
- scale=scale,
169
- symmetric=not schema.zero_point,
170
- )
171
-
172
-
173
-def _gptq_weight_schema_to_quant_key(schema: "GPTQWeightSchema") -> QuantKey:
174
- group_shape = _group_shape(schema.group_size)
175
- scale = ScaleDesc(
176
- dtype=torch.float16,
177
- static=True,
178
- group_shape=group_shape,
179
- )
180
- return QuantKey(dtype=INT4_DTYPE, scale=scale, symmetric=schema.sym)
181
-
182
-
183
-def _compressed_tensors_weight_schema_to_quant_key(
184
- schema: "CompressedTensorsWeightSchema",
185
-) -> QuantKey:
186
- # Determine dtype from format/type/num_bits
187
- fmt = schema.format
188
- if fmt in ("int-quantized", "float-quantized", "naive-quantized"):
189
- dtype = INT8_DTYPE if schema.type == "int" else FP8_DTYPE
190
- elif "nvfp4" in fmt or "mxfp4" in fmt:
191
- dtype = FP4_DTYPE
192
- else:
193
- dtype = _HUMMING_TO_QUANT_DTYPE[
194
- humming_dtypes.DataType.from_str(f"uint{schema.num_bits}")
195
- ]
196
-
197
- # Determine group shape from strategy
198
- if schema.strategy in ("group", "tensor_group"):
199
- group_shape = _group_shape(schema.group_size or 0)
200
- elif schema.strategy == "block" and schema.block_structure is not None:
201
- group_shape = GroupShape(
202
- row=schema.block_structure[0],
203
- col=schema.block_structure[1],
204
- )
205
- else:
206
- group_shape = GroupShape.PER_CHANNEL
207
-
208
- # Determine scale dtype
209
- if "mxfp" in fmt:
210
- scale_dtype = MXFP_SCALE_DTYPE
211
- elif "nvfp4" in fmt:
212
- scale_dtype = FP8_DTYPE
213
- else:
214
- scale_dtype = torch.float32
215
-
216
- scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
217
-
218
- scale2 = None
219
- if "nvfp4" in fmt or schema.strategy == "tensor_group":
220
- scale2 = ScaleDesc(
221
- dtype=torch.float32,
222
- static=True,
223
- group_shape=GroupShape.PER_TENSOR,
224
- )
225
-
226
- return QuantKey(
227
- dtype=dtype,
228
- scale=scale,
229
- scale2=scale2,
230
- symmetric=schema.symmetric,
231
- )
232
-
233
-
234
-# ---- Dispatch for any BaseWeightSchema ------------------------------------
235
-
236
-
237
-def weight_schema_to_quant_key(
238
- schema: "BaseWeightSchema",
239
-) -> QuantKey:
240
- from vllm.utils.humming import (
241
- AWQWeightSchema,
242
- BitnetWeightSchema,
243
- CompressedTensorsWeightSchema,
244
- Fp8WeightSchema,
245
- GptOssMxfp4WeightSchema,
246
- GPTQWeightSchema,
247
- HummingWeightSchema,
248
- ModeloptMxfp8WeightSchema,
249
- ModeloptNvfp4WeightSchema,
250
- Mxfp4WeightSchema,
251
- )
252
-
253
- """Convert any BaseWeightSchema to a QuantKey."""
254
- if isinstance(schema, HummingWeightSchema):
255
- return _humming_weight_schema_to_quant_key(schema)
256
-
257
- # Schemas with fixed QuantKeys
258
- if isinstance(schema, (Mxfp4WeightSchema, GptOssMxfp4WeightSchema)):
259
- return QuantKey(
260
- dtype=FP4_DTYPE,
261
- scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),
262
- )
263
- if isinstance(schema, ModeloptMxfp8WeightSchema):
264
- return QuantKey(
265
- dtype=FP8_DTYPE,
266
- scale=ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)),
267
- )
268
- if isinstance(schema, ModeloptNvfp4WeightSchema):
269
- return QuantKey(
270
- dtype=FP4_DTYPE,
271
- scale=ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16)),
272
- scale2=ScaleDesc(torch.float32, True, GroupShape.PER_TENSOR),
273
- )
274
- if isinstance(schema, BitnetWeightSchema):
275
- return QuantKey(
276
- dtype=torch.uint8,
277
- scale=ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL),
278
- )
279
-
280
- # Schemas requiring config inspection
281
- if isinstance(schema, Fp8WeightSchema):
282
- return _fp8_weight_schema_to_quant_key(schema)
283
- if isinstance(schema, AWQWeightSchema):
284
- return _awq_weight_schema_to_quant_key(schema)
285
- if isinstance(schema, GPTQWeightSchema):
286
- return _gptq_weight_schema_to_quant_key(schema)
287
- if isinstance(schema, CompressedTensorsWeightSchema):
288
- return _compressed_tensors_weight_schema_to_quant_key(schema)
289
-
290
- raise TypeError(f"Unsupported weight schema type: {type(schema)}")
291
-
292
-
293
-# ---- HummingInputSchema (post-conversion) ----------------------------------
294
-
295
-
296
-def _humming_input_schema_to_quant_key(
297
- schema: "HummingInputSchema",
298
-) -> QuantKey | None:
299
- """Convert a HummingInputSchema to a QuantKey. Returns None if
300
- the schema represents unquantized (bf16/fp16) inputs."""
301
- if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:
302
- return None
303
-
304
- dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]
305
-
306
- gs = schema.input_scale_group_size
307
- group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
308
-
309
- # Pick the scale dtype the Humming kernel actually consumes. An explicit
310
- # input_scale_dtype always wins. Otherwise infer from the grouping: MX
311
- # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
312
- # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
313
- # float32 scale. Getting this right lets a grouped FP8 activation match
314
- # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
315
- if schema.input_scale_dtype is not None:
316
- scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
317
- elif gs == 32:
318
- scale_dtype = MXFP_SCALE_DTYPE
319
- else:
320
- scale_dtype = torch.float32
321
-
322
- scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)
323
-
324
- return QuantKey(dtype=dtype, scale=scale, symmetric=True)
325
-
326
-
327
-# ---- Checkpoint-format input schemas (pre-conversion) ----------------------
328
-
329
-
330
-def _resolve_input_quant_key(
331
- origin_a_dtype: "humming_dtypes.DataType",
332
- group_size: int,
333
-) -> QuantKey | None:
334
- from vllm.utils.humming import HummingInputSchema
335
-
336
- """Resolve the actual activation QuantKey after platform fallback."""
337
- a_dtype = HummingInputSchema().get_fallback_input_dtype(origin_a_dtype)
338
- if a_dtype is None or a_dtype.num_bits >= 16:
339
- return None
340
-
341
- dtype = _HUMMING_TO_QUANT_DTYPE[a_dtype]
342
- gs = group_size if a_dtype == humming_dtypes.float4e2m1 else 0
343
- group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
344
- scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32
345
-
346
- scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)
347
- return QuantKey(dtype=dtype, scale=scale, symmetric=True)
348
-
349
-
350
-def _compressed_tensors_input_schema_to_quant_key(
351
- schema: "CompressedTensorsInputSchema",
352
-) -> QuantKey | None:
353
- type_bits_to_dtype = {
354
- ("float", 8): humming_dtypes.float8e4m3,
355
- ("float", 4): humming_dtypes.float4e2m1,
356
- ("int", 8): humming_dtypes.int8,
357
- ("int", 4): humming_dtypes.int4,
358
- }
359
- origin = type_bits_to_dtype.get((schema.type, schema.num_bits))
360
- if origin is None:
361
- return None
362
- return _resolve_input_quant_key(origin, schema.group_size)
363
-
364
-
365
-# ---- Dispatch for any BaseInputSchema -------------------------------------
366
-
367
-
368
-def input_schema_to_quant_key(
369
- schema: "BaseInputSchema",
370
-) -> QuantKey | None:
371
- from vllm.utils.humming import (
372
- CompressedTensorsInputSchema,
373
- Fp8InputSchema,
374
- HummingInputSchema,
375
- ModeloptNvfp4InputSchema,
376
- )
377
-
378
- """Convert any BaseInputSchema to a QuantKey. Returns None if
379
- the schema represents unquantized (bf16/fp16) inputs."""
380
- if isinstance(schema, HummingInputSchema):
381
- return _humming_input_schema_to_quant_key(schema)
382
-
383
- if isinstance(schema, Fp8InputSchema):
384
- return _resolve_input_quant_key(humming_dtypes.float8e4m3, 0)
385
-
386
- if isinstance(schema, ModeloptNvfp4InputSchema):
387
- return _resolve_input_quant_key(
388
- humming_dtypes.float8e4m3,
389
- schema.group_size,
390
- )
391
-
392
- if isinstance(schema, CompressedTensorsInputSchema):
393
- return _compressed_tensors_input_schema_to_quant_key(schema)
394
-
395
- raise TypeError(f"Unsupported input schema type: {type(schema)}")
396
-
397
-
398
-def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
399
- if not config:
400
- return True
401
-
402
- keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
403
- ignored_layers: list[str] = []
404
- for key in keys:
405
- candidate = config.get(key, []) or []
406
- if candidate:
407
- ignored_layers = candidate
408
- break
409
-
410
- if any(module_name in prefix for module_name in ignored_layers):
411
- return True
412
- if "lm_head" in prefix:
413
- return True
414
-
415
- for regex in config.get("dynamic", {}):
416
- if regex[:1] != "-":
417
- continue
418
- if re.match(regex[2:], prefix):
419
- return True
420
-
421
- return False
422
-
423
-
424
-def convert_linear_layer_to_humming_standard(
425
- layer: LinearBase, name_map: dict[str, str]
426
-):
427
- """Rename/reshape a linear layer's quantized params (the canonical MPLinear
428
- layout: ``weight_packed`` int32 + ``weight_scale``) into the parameter names
429
- and layout humming's weight schema expects (``weight`` / ``weight_scale``)."""
430
- for name, checkpoint_name in name_map.items():
431
- tensor = getattr(layer, checkpoint_name)
432
- delattr(layer, checkpoint_name)
433
-
434
- if name == "weight":
435
- input_dim = getattr(tensor, "input_dim", 1)
436
- output_dim = getattr(tensor, "output_dim", 0)
437
-
438
- if input_dim == 0 and output_dim == 1:
439
- tensor = tensor.transpose(1, 0).contiguous()
440
- else:
441
- assert output_dim == 0 and input_dim == 1
442
-
443
- tensor = tensor.view(tensor.size(0), -1).view(torch.int32)
444
- elif name in ["weight_scale", "zero_point"]:
445
- if getattr(tensor, "output_dim", 0) == 1:
446
- tensor = tensor.transpose(0, 1).contiguous()
447
- if tensor.ndim == 1:
448
- tensor = tensor.unsqueeze(1)
449
-
450
- tensor = tensor.view(torch.int32) if name == "zero_point" else tensor
451
-
452
- if isinstance(tensor, torch.nn.Parameter):
453
- param = tensor
454
- else:
455
- param = torch.nn.Parameter(tensor, requires_grad=False)
456
-
457
- setattr(layer, name, param)
458
-
459
-
460
-def prepare_humming_linear_layer_config(
461
- layer: LinearBase,
462
- quant_config: dict,
463
- input_quant_config: dict | None = None,
464
-) -> "LayerConfig":
465
- from vllm.utils.humming import (
466
- BaseInputSchema,
467
- BaseWeightSchema,
468
- HummingInputSchema,
469
- prepare_layer_config,
470
- transform_humming_tensors,
471
- )
472
-
473
- weight_schema = BaseWeightSchema.from_config(quant_config)
474
- if input_quant_config is not None:
475
- input_schema = BaseInputSchema.from_config(input_quant_config)
476
- else:
477
- input_schema = HummingInputSchema()
478
-
479
- # ReplicatedLinear has no TP partitioning and so does not set
480
- # input_size_per_partition; for it that is just input_size. Use hasattr
481
- # rather than getattr's default arg, which is evaluated eagerly and would
482
- # raise on layers lacking input_size (e.g. ParallelLMHead).
483
- if hasattr(layer, "input_size_per_partition"):
484
- input_size_per_partition = layer.input_size_per_partition
485
- else:
486
- input_size_per_partition = layer.input_size
487
- shape_k_stacks = [input_size_per_partition]
488
- shape_n_stacks = layer.output_partition_sizes
489
-
490
- # Step 1: convert weight and input schemas to humming standard format
491
- weight_schema, tensors = weight_schema.convert_humming(
492
- tensors=dict(layer.named_parameters()),
493
- shape_n_stacks=shape_n_stacks,
494
- shape_k_stacks=shape_k_stacks,
495
- param_dtype=layer.params_dtype,
496
- )
497
- input_schema, _ = input_schema.convert_humming(
498
- tensors={},
499
- shape_n_stacks=shape_n_stacks,
500
- shape_k_stacks=shape_k_stacks,
501
- param_dtype=layer.params_dtype,
502
- )
503
-
504
- layer.weight_schema = weight_schema
505
-
506
- # Step 2: transform weight (humming standard format) for forwarding.
507
- config = prepare_layer_config(
508
- shape_n=sum(layer.output_partition_sizes),
509
- shape_k=input_size_per_partition,
510
- weight_schema=weight_schema,
511
- input_schema=input_schema,
512
- pad_n_to_multiple=256,
513
- pad_k_to_multiple=128,
514
- has_bias=layer.has_bias,
515
- torch_dtype=layer.params_dtype,
516
- )
517
- tensors = transform_humming_tensors(config, tensors)
518
- for name, _ in list(layer.named_parameters()):
519
- delattr(layer, name)
520
- for name, tensor in tensors.items():
521
- param = torch.nn.Parameter(tensor, requires_grad=False)
522
- setattr(layer, name, param)
523
-
524
- return config
525
-
526
-
527
-def get_humming_linear_compute_config() -> str:
528
- return json.dumps(
529
- {
530
- "use_batch_invariant": envs.VLLM_BATCH_INVARIANT,
531
- "use_f16_accum": envs.VLLM_HUMMING_USE_F16_ACCUM,
532
- "gemm_type": "dense",
533
- }
534
- )
535
-
536
-
537
-def apply_humming_linear(
538
- layer: LinearBase,
539
- x: torch.Tensor,
540
- *,
541
- layer_config: "LayerConfig",
542
- compute_config: str,
543
- locks: torch.Tensor,
544
-) -> torch.Tensor:
545
- from vllm.utils.humming import humming_forward
546
-
547
- flatten_inputs = x.reshape(-1, x.size(-1))
548
- output = humming_forward(
549
- layer_config,
550
- inputs=flatten_inputs,
551
- weight=layer.weight,
552
- weight_scale=getattr(layer, "weight_scale", None),
553
- zero_point=getattr(layer, "zero_point", None),
554
- bias=getattr(layer, "bias", None),
555
- weight_scale_2=getattr(layer, "weight_scale_2", None),
556
- locks=locks,
557
- compute_config=compute_config,
558
- )
559
- return output.view(*x.shape[:-1], output.size(-1))
48
+ w1_hadamard_block_size: int = 1
49
+ w2_hadamard_block_size: int = 1
50
+ w1_input_scale: torch.Tensor | None = None
51
+ w2_input_scale: torch.Tensor | None = None
52
+ w1_input_scale_2: torch.Tensor | None = None
53
+ w2_input_scale_2: torch.Tensor | None = None
54
55
56
def make_humming_moe_quant_config(
@@ -648,8 +142,12 @@ def get_humming_moe_quant_config(
142
q_dtype is not None
143
and q_dtype.startswith("float8")
144
and input_scale_group_size == 128
145
+ and weight_schema.hadamard_block_size <= 1
146
+ and input_schema.input_quant_mode in (None, "dynamic_group")
147
):
652
- q_dtype = _HUMMING_TO_QUANT_DTYPE.get(input_schema.a_dtype, FP8_DTYPE)
148
+ q_dtype = humming_schema._HUMMING_TO_QUANT_DTYPE.get(
149
+ input_schema.a_dtype, FP8_DTYPE
150
+ )
151
activation_group_shape = GroupShape(row=1, col=input_scale_group_size)
152
153
weight_scale_group_size = weight_schema.weight_scale_group_size
@@ -665,7 +163,7 @@ def get_humming_moe_quant_config(
163
else:
164
weight_group_shape = GroupShape(row=weight_scale_group_size, col=1)
165
668
- return make_humming_moe_quant_config(
166
+ config = make_humming_moe_quant_config(
167
quant_dtype=q_dtype,
168
weight_dtype=str(weight_schema.b_dtype),
169
weight_group_shape=weight_group_shape,
@@ -683,6 +181,17 @@ def get_humming_moe_quant_config(
181
gemm1_clamp_limit=gemm1_clamp_limit,
182
humming_configs=humming_configs,
183
)
184
+ for prefix, sublayer in (("w1", "w13"), ("w2", "w2")):
185
+ setattr(
186
+ config,
187
+ f"{prefix}_hadamard_block_size",
188
+ layer.weight_schemas[sublayer].hadamard_block_size,
189
+ )
190
+ for name in ("input_scale", "input_scale_2"):
191
+ key = f"{prefix}_{name}"
192
+ value = getattr(layer, f"{sublayer}_{name}", None)
193
+ setattr(config, key, value)
194
+ return config
195
196
197
def select_humming_moe_experts(
@@ -846,12 +355,6 @@ def _convert_sublayer_to_humming(
355
Tuple of (converted_weight_schema, converted_input_schema)
356
357
"""
849
- from vllm.utils.humming import HummingWeightSchema
850
-
851
- if isinstance(weight_schema, HummingWeightSchema):
852
- # Already in Humming format
853
- return weight_schema, input_schema
854
-
358
tensors = _extract_sublayer_tensors(layer, sublayer_name)
359
360
shape_k_stacks = [shape_k]
@@ -867,13 +370,14 @@ def _convert_sublayer_to_humming(
370
num_experts=num_experts,
371
)
372
870
- converted_input_schema, _ = input_schema.convert_humming(
871
- tensors=converted_tensors,
373
+ converted_input_schema, input_tensors = input_schema.convert_humming(
374
+ tensors=tensors,
375
shape_n_stacks=shape_n_stacks,
376
shape_k_stacks=shape_k_stacks,
377
param_dtype=param_dtype,
378
num_experts=num_experts,
379
)
380
+ converted_tensors.update(input_tensors)
381
382
_replace_layer_parameters(layer, sublayer_name, converted_tensors)
383
@@ -907,11 +411,16 @@ def _prepare_and_transform_sublayer(
411
has_bias=has_bias,
412
num_experts=num_experts,
413
torch_dtype=param_dtype,
414
+ device=getattr(layer, f"{sublayer_name}_weight").device,
415
)
911
- tensors = transform_humming_tensors(
912
- config,
913
- _extract_sublayer_tensors(layer, sublayer_name),
914
- )
416
+ source_tensors = _extract_sublayer_tensors(layer, sublayer_name)
417
+ for name in ("input_scale", "input_scale_2"):
418
+ if name != input_schema.static_tensor_scale_name:
419
+ source_tensors.pop(name, None)
420
+ tensors = transform_humming_tensors(config, source_tensors)
421
+ for name in ("input_scale", "input_scale_2"):
422
+ if name in source_tensors:
423
+ tensors[name] = source_tensors[name]
424
_replace_layer_parameters(layer, sublayer_name, tensors)
425
return config
426
@@ -927,6 +436,7 @@ def _process_single_sublayer(
436
num_experts: int,
437
param_dtype: torch.dtype,
438
force_weight_schema: Any | None = None,
439
+ allow_input_schema_fallback: bool = True,
440
) -> tuple[Any, Any, "LayerConfig"]:
441
"""Process a single sublayer: convert, optionally requant, prepare, and transform.
442
@@ -944,6 +454,7 @@ def _process_single_sublayer(
454
num_experts: Number of experts
455
param_dtype: Parameter data type
456
force_weight_schema: Optional schema to force requantization to
457
+ allow_input_schema_fallback: Whether incompatible input schemas may be replaced.
458
459
Returns:
460
Tuple of the final weight schema, input schema, and Humming layer config.
@@ -966,19 +477,28 @@ def _process_single_sublayer(
477
# Step 2: Force requant if needed
478
assert isinstance(current_weight_schema, HummingWeightSchema)
479
if force_weight_schema is not None and current_weight_schema != force_weight_schema:
969
- tensors = _extract_sublayer_tensors(layer, sublayer_name)
480
+ source_tensors = _extract_sublayer_tensors(layer, sublayer_name)
481
482
tensors = current_weight_schema.requant_tensors(
972
- tensors=tensors,
483
+ tensors=source_tensors,
484
target_weight_schema=force_weight_schema,
485
param_dtype=param_dtype,
486
)
487
+ name = current_input_schema.static_tensor_scale_name
488
+ if name is not None:
489
+ tensors[name] = source_tensors[name]
490
491
current_weight_schema = force_weight_schema
492
_replace_layer_parameters(layer, sublayer_name, tensors, preserve_bias=True)
493
del tensors
494
495
# Step 3: Prepare layer metadata and transform weights
496
+ current_input_schema = humming_schema.check_and_fallback_input_schema(
497
+ weight_schema=current_weight_schema,
498
+ input_schema=current_input_schema,
499
+ param_dtype=param_dtype,
500
+ allow_fallback=allow_input_schema_fallback,
501
+ )
502
config = _prepare_and_transform_sublayer(
503
layer=layer,
504
sublayer_name=sublayer_name,
@@ -1001,6 +521,7 @@ def convert_to_humming_moe_kernel_format(
521
weight_schema: Any | None = None,
522
input_schema: Any | None = None,
523
force_weight_schema: Any | None = None,
524
+ allow_input_schema_fallback: bool = True,
525
) -> dict[str, "LayerConfig"]:
526
"""Convert MoE weights from checkpoint format to Humming kernel format.
527
@@ -1023,6 +544,7 @@ def convert_to_humming_moe_kernel_format(
544
input_schema: Optional initial input quantization schema.
545
If None, built from quant_config or env vars.
546
force_weight_schema: Optional schema to force requantization to
547
+ allow_input_schema_fallback: Whether incompatible input schemas may be replaced.
548
549
Side effects:
550
- Modifies layer parameters in place
@@ -1041,20 +563,23 @@ def convert_to_humming_moe_kernel_format(
563
"Must provide either weight_schema/input_schema or quant_config"
564
)
565
1044
- from vllm.model_executor.layers.quantization.utils.humming_utils import (
1045
- humming_is_layer_skipped,
1046
- )
566
from vllm.utils.humming import BaseWeightSchema, HummingInputSchema
567
568
if weight_schema is None:
569
weight_schema = BaseWeightSchema.from_config(quant_config)
570
571
if input_schema is None:
1053
- input_quant_config = envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}
572
+ input_quant_config = (envs.VLLM_HUMMING_INPUT_QUANT_CONFIG or {}).copy()
573
if humming_is_layer_skipped(input_quant_config, layer.layer_name):
574
input_schema = HummingInputSchema()
575
else:
576
# TODO: read input_quant_config from quant_config
577
+ input_quant_config = humming_schema.resolve_humming_layer_config(
578
+ input_quant_config, layer.layer_name
579
+ )
580
+ allow_input_schema_fallback = input_quant_config.pop(
581
+ "allow_fallback", False
582
+ )
583
input_schema = HummingInputSchema.from_config(input_quant_config)
584
585
# Build sublayer configs from layer properties if not provided
@@ -1089,6 +614,7 @@ def convert_to_humming_moe_kernel_format(
614
num_experts=num_experts,
615
param_dtype=param_dtype,
616
force_weight_schema=force_weight_schema,
617
+ allow_input_schema_fallback=allow_input_schema_fallback,
618
)
619
)
620
vllm/model_executor/layers/quantization/utils/humming/schema.py
new
+437
@@ -0,0 +1,437 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Map Humming schemas and handle shared checkpoint quantization settings."""
4
+
5
+import dataclasses
6
+from typing import TYPE_CHECKING, Any
7
+
8
+import regex as re
9
+import torch
10
+
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ FP4_DTYPE,
14
+ FP8_DTYPE,
15
+ INT4_DTYPE,
16
+ INT8_DTYPE,
17
+ MXFP_SCALE_DTYPE,
18
+ GroupShape,
19
+ QuantKey,
20
+ ScaleDesc,
21
+)
22
+from vllm.platforms import current_platform
23
+from vllm.scalar_type import scalar_types
24
+from vllm.utils.import_utils import has_humming
25
+
26
+if TYPE_CHECKING:
27
+ from vllm.utils.humming import (
28
+ BaseInputSchema,
29
+ BaseWeightSchema,
30
+ HummingInputSchema,
31
+ HummingWeightSchema,
32
+ )
33
+ from vllm.utils.humming import dtypes as humming_dtypes
34
+
35
+
36
+def humming_update_schema_hadamard_block_size(
37
+ weight_schema: "HummingWeightSchema",
38
+ input_schema: "HummingInputSchema",
39
+ shape_k: int,
40
+) -> "HummingWeightSchema":
41
+ assert shape_k > 0
42
+ block_size = 32
43
+ weight_group_size = weight_schema.weight_scale_group_size
44
+ input_group_size = input_schema.input_scale_group_size
45
+ for group_size in (weight_group_size, input_group_size):
46
+ if group_size is not None and group_size > 0:
47
+ block_size = min(block_size, group_size)
48
+ block_size = 1 << (block_size.bit_length() - 1)
49
+ while shape_k % block_size:
50
+ block_size //= 2
51
+ return dataclasses.replace(weight_schema, hadamard_block_size=block_size)
52
+
53
+
54
+if has_humming():
55
+ from vllm.utils.humming import dtypes as humming_dtypes
56
+
57
+ _HUMMING_TO_QUANT_DTYPE: dict[humming_dtypes.DataType, Any] = {
58
+ humming_dtypes.float4e0m3: FP4_DTYPE,
59
+ humming_dtypes.float4e2m1: FP4_DTYPE,
60
+ humming_dtypes.float6e2m3: scalar_types.float6_e2m3f,
61
+ humming_dtypes.float6e3m2: scalar_types.float6_e3m2f,
62
+ humming_dtypes.float8e3m4: FP8_DTYPE,
63
+ humming_dtypes.float8e4m3: FP8_DTYPE,
64
+ humming_dtypes.float8e5m2: torch.float8_e5m2,
65
+ humming_dtypes.int8: torch.int8,
66
+ humming_dtypes.uint4: INT4_DTYPE,
67
+ humming_dtypes.uint8: INT8_DTYPE,
68
+ humming_dtypes.uint2: torch.uint8,
69
+ humming_dtypes.uint3: torch.uint8,
70
+ humming_dtypes.uint5: torch.uint8,
71
+ humming_dtypes.uint6: torch.uint8,
72
+ humming_dtypes.uint7: torch.uint8,
73
+ }
74
+
75
+ _HUMMING_TO_SCALE_DTYPE: dict[humming_dtypes.DataType, torch.dtype] = {
76
+ humming_dtypes.float8e8m0: MXFP_SCALE_DTYPE,
77
+ humming_dtypes.float8e4m3: FP8_DTYPE,
78
+ humming_dtypes.float16: torch.float16,
79
+ humming_dtypes.bfloat16: torch.bfloat16,
80
+ humming_dtypes.float32: torch.float32,
81
+ }
82
+
83
+
84
+logger = init_logger(__name__)
85
+
86
+
87
+def _group_shape(group_size: int, group_size_n: int = 0) -> GroupShape:
88
+ """Map humming group sizes to QuantKey GroupShape.
89
+
90
+ group_size: elements per group along K (col); 0 means full dimension.
91
+ group_size_n: elements per group along N (row); 0 means 1 (per-row).
92
+
93
+ GroupShape convention: row = N dim, col = K dim.
94
+ """
95
+ if group_size == 0 and group_size_n == 0:
96
+ return GroupShape.PER_CHANNEL
97
+
98
+ row = group_size_n if group_size_n > 0 else 1
99
+ col = group_size if group_size > 0 else -1
100
+ return GroupShape(row=row, col=col)
101
+
102
+
103
+def weight_schema_to_quant_key(
104
+ schema: "BaseWeightSchema",
105
+ param_dtype: torch.dtype | None = None,
106
+) -> QuantKey:
107
+ from vllm.utils.humming import WeightScale2Type, WeightScaleType
108
+
109
+ if param_dtype is None:
110
+ param_dtype = torch.get_default_dtype()
111
+ if param_dtype not in [torch.float16, torch.bfloat16]:
112
+ param_dtype = torch.bfloat16
113
+ if not current_platform.has_device_capability(80):
114
+ param_dtype = torch.float16
115
+
116
+ schema = schema.to_humming_schema(param_dtype)
117
+ dtype = _HUMMING_TO_QUANT_DTYPE[schema.b_dtype]
118
+
119
+ if schema.bs_dtype is not None:
120
+ scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.bs_dtype]
121
+ else:
122
+ scale_dtype = torch.float32
123
+
124
+ group_shape = _group_shape(
125
+ schema.weight_scale_group_size,
126
+ schema.weight_scale_group_size_n,
127
+ )
128
+ if schema.weight_scale_type == WeightScaleType.TENSOR:
129
+ group_shape = GroupShape.PER_TENSOR
130
+
131
+ scale = ScaleDesc(dtype=scale_dtype, static=True, group_shape=group_shape)
132
+
133
+ scale2 = None
134
+ if schema.weight_scale_2_type == WeightScale2Type.TENSOR:
135
+ scale2 = ScaleDesc(
136
+ dtype=torch.float32,
137
+ static=True,
138
+ group_shape=GroupShape.PER_TENSOR,
139
+ )
140
+
141
+ return QuantKey(
142
+ dtype=dtype,
143
+ scale=scale,
144
+ scale2=scale2,
145
+ symmetric=not schema.has_zero_point,
146
+ )
147
+
148
+
149
+def input_schema_to_quant_key(
150
+ schema: "BaseInputSchema",
151
+ param_dtype: torch.dtype | None = None,
152
+) -> QuantKey | None:
153
+ from vllm.utils.humming import InputQuantizationMode
154
+
155
+ if param_dtype is None:
156
+ param_dtype = torch.get_default_dtype()
157
+ if param_dtype not in [torch.float16, torch.bfloat16]:
158
+ param_dtype = torch.bfloat16
159
+ if not current_platform.has_device_capability(80):
160
+ param_dtype = torch.float16
161
+
162
+ schema = schema.to_humming_schema(param_dtype)
163
+ if schema.a_dtype is None or schema.a_dtype.num_bits >= 16:
164
+ return None
165
+
166
+ mode = schema.input_quant_mode
167
+ if mode == InputQuantizationMode.Disabled:
168
+ return None
169
+ dtype = _HUMMING_TO_QUANT_DTYPE[schema.a_dtype]
170
+
171
+ gs = schema.input_scale_group_size
172
+ group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
173
+
174
+ # Pick the scale dtype the Humming kernel actually consumes. An explicit
175
+ # input_scale_dtype always wins. Otherwise infer from the grouping: MX
176
+ # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
177
+ # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
178
+ # float32 scale. Getting this right lets a grouped FP8 activation match
179
+ # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
180
+ if schema.input_scale_dtype is not None:
181
+ scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
182
+ elif gs == 16:
183
+ scale_dtype = FP8_DTYPE
184
+ elif gs == 32:
185
+ scale_dtype = MXFP_SCALE_DTYPE
186
+ else:
187
+ scale_dtype = torch.float32
188
+
189
+ if mode == InputQuantizationMode.StaticTensor:
190
+ group_shape = GroupShape.PER_TENSOR
191
+ scale = ScaleDesc(
192
+ dtype=scale_dtype,
193
+ static=mode == InputQuantizationMode.StaticTensor,
194
+ group_shape=group_shape,
195
+ )
196
+ scale2 = None
197
+ has_scale2 = (
198
+ InputQuantizationMode.StaticTensorDynamicGroup,
199
+ InputQuantizationMode.DynamicGroupToken,
200
+ )
201
+ if mode in has_scale2:
202
+ static = mode == InputQuantizationMode.StaticTensorDynamicGroup
203
+ scale2 = ScaleDesc(
204
+ torch.float32,
205
+ static,
206
+ GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN,
207
+ )
208
+
209
+ return QuantKey(dtype=dtype, scale=scale, scale2=scale2, symmetric=True)
210
+
211
+
212
+def quant_key_to_input_schema(key: QuantKey | None) -> "HummingInputSchema":
213
+ from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
214
+
215
+ if key is None:
216
+ return HummingInputSchema(input_quant_mode=InputQuantizationMode.Disabled)
217
+
218
+ if not key.symmetric:
219
+ raise ValueError("Humming input quantization must be symmetric")
220
+
221
+ quant_dtypes = {value: dtype for dtype, value in _HUMMING_TO_QUANT_DTYPE.items()}
222
+ quant_dtypes[FP4_DTYPE] = humming_dtypes.float4e2m1
223
+ quant_dtypes[FP8_DTYPE] = humming_dtypes.float8e4m3
224
+ scale_dtypes = {value: dtype for dtype, value in _HUMMING_TO_SCALE_DTYPE.items()}
225
+ if key.dtype not in quant_dtypes or key.scale.dtype not in scale_dtypes:
226
+ raise ValueError(f"Unsupported Humming input or scale dtype: {key}")
227
+
228
+ scale = key.scale
229
+ group_size = 0
230
+ if scale.static and scale.group_shape.is_per_tensor() and key.scale2 is None:
231
+ mode = InputQuantizationMode.StaticTensor
232
+ elif not scale.static and scale.group_shape.is_per_token() and key.scale2 is None:
233
+ mode = InputQuantizationMode.DynamicToken
234
+ elif not scale.static and scale.group_shape.is_per_group():
235
+ group_size = scale.group_shape.col
236
+ scale2 = key.scale2
237
+ if scale2 is None:
238
+ mode = InputQuantizationMode.DynamicGroup
239
+ elif scale2.dtype != torch.float32:
240
+ raise ValueError("Humming secondary input scales must be float32")
241
+ elif scale2.static and scale2.group_shape.is_per_tensor():
242
+ mode = InputQuantizationMode.StaticTensorDynamicGroup
243
+ elif not scale2.static and scale2.group_shape.is_per_token():
244
+ mode = InputQuantizationMode.DynamicGroupToken
245
+ else:
246
+ raise ValueError(f"Unsupported Humming secondary input scale: {scale2}")
247
+ else:
248
+ raise ValueError(f"Unsupported Humming input scale: {key}")
249
+
250
+ if group_size == 0 and scale.dtype != torch.float32:
251
+ raise ValueError("Humming tensor and token input scales must be float32")
252
+
253
+ return HummingInputSchema(
254
+ a_dtype=quant_dtypes[key.dtype],
255
+ input_scale_group_size=group_size,
256
+ input_scale_dtype=scale_dtypes[scale.dtype],
257
+ input_quant_mode=mode,
258
+ )
259
+
260
+
261
+def check_and_fallback_input_schema(
262
+ weight_schema: "BaseWeightSchema",
263
+ input_schema: "BaseInputSchema",
264
+ param_dtype: torch.dtype | None = None,
265
+ allow_fallback: bool = True,
266
+) -> "HummingInputSchema":
267
+ from vllm.utils.humming import HummingInputSchema, InputQuantizationMode
268
+
269
+ capability = current_platform.get_device_capability()
270
+ assert capability is not None
271
+ sm_version = capability.to_int()
272
+ if param_dtype is None:
273
+ param_dtype = torch.get_default_dtype()
274
+ if param_dtype not in (torch.float16, torch.bfloat16):
275
+ param_dtype = torch.bfloat16 if sm_version >= 80 else torch.float16
276
+ assert param_dtype in (torch.float16, torch.bfloat16)
277
+
278
+ dtype_fallback_order_map = {
279
+ humming_dtypes.float8e5m2: (
280
+ humming_dtypes.float8e4m3,
281
+ humming_dtypes.float8e3m4,
282
+ humming_dtypes.int8,
283
+ ),
284
+ humming_dtypes.float8e4m3: (
285
+ humming_dtypes.float8e3m4,
286
+ humming_dtypes.int8,
287
+ ),
288
+ humming_dtypes.float8e3m4: (
289
+ humming_dtypes.float8e4m3,
290
+ humming_dtypes.int8,
291
+ ),
292
+ humming_dtypes.float4e2m1: (
293
+ humming_dtypes.float8e4m3,
294
+ humming_dtypes.float8e3m4,
295
+ humming_dtypes.int8,
296
+ ),
297
+ humming_dtypes.float4e0m3: (
298
+ humming_dtypes.int4,
299
+ humming_dtypes.float8e4m3,
300
+ humming_dtypes.float8e3m4,
301
+ humming_dtypes.int8,
302
+ ),
303
+ humming_dtypes.int8: (
304
+ humming_dtypes.float8e4m3,
305
+ humming_dtypes.float8e3m4,
306
+ ),
307
+ humming_dtypes.int4: (
308
+ humming_dtypes.int8,
309
+ humming_dtypes.float8e4m3,
310
+ humming_dtypes.float8e3m4,
311
+ ),
312
+ }
313
+
314
+ weight_schema = weight_schema.to_humming_schema(param_dtype)
315
+ input_schema = input_schema.to_humming_schema(param_dtype)
316
+ fp6_dtypes = (humming_dtypes.float6e2m3, humming_dtypes.float6e3m2)
317
+ if input_schema.a_dtype in fp6_dtypes:
318
+ input_schema.a_dtype = humming_dtypes.float8e4m3
319
+ a_dtype = input_schema.a_dtype
320
+ input_bits = a_dtype.num_bits if a_dtype is not None else 16
321
+ weight_group_size = weight_schema.weight_scale_group_size
322
+ input_group_size = input_schema.input_scale_group_size
323
+ input_scale_dtype = input_schema.input_scale_dtype
324
+
325
+ def is_deprecated(dtype: "humming_dtypes.DataType | None") -> bool:
326
+ if dtype is None or dtype.num_bits == 16:
327
+ return False
328
+ is_int4_deprecated = dtype == humming_dtypes.int4 and sm_version >= 90
329
+ is_int8_deprecated = dtype == humming_dtypes.int8 and 103 <= sm_version < 110
330
+ # Humming has not yet implemented high-performance FP8 computation
331
+ # on SM10X and SM11X. Therefore, it is temporarily considered deprecated.
332
+ # TODO: removed this after Humming has high-performance FP8 computation
333
+ is_fp8_deprecated = dtype.num_bits < 16 and 100 <= sm_version < 120
334
+ return is_int4_deprecated or is_int8_deprecated or is_fp8_deprecated
335
+
336
+ if input_schema.is_compatible_with(weight_schema, param_dtype):
337
+ if not allow_fallback:
338
+ if is_deprecated(a_dtype):
339
+ msg = f"Humming {a_dtype} activation is slow on SM{sm_version} now."
340
+ logger.warning_once(msg)
341
+ return input_schema
342
+
343
+ if a_dtype is None or a_dtype.num_bits == 16:
344
+ return input_schema
345
+
346
+ if not is_deprecated(a_dtype):
347
+ is_mxfp8 = (
348
+ a_dtype.is_floating_point_type
349
+ and a_dtype.num_bits == 8
350
+ and input_group_size == 32
351
+ and input_scale_dtype == humming_dtypes.float8e8m0
352
+ and sm_version >= 120
353
+ )
354
+ is_groupwise = input_group_size > 0 or weight_group_size > 0
355
+ if input_bits == 8 and is_groupwise and not is_mxfp8:
356
+ # prefer wna16 for groupwise weight or groupwise input
357
+ return HummingInputSchema()
358
+
359
+ return input_schema
360
+
361
+ if allow_fallback:
362
+ fallback_dtypes = (
363
+ a_dtype,
364
+ *dtype_fallback_order_map.get(a_dtype, ()),
365
+ humming_dtypes.DataType.from_any(param_dtype),
366
+ )
367
+ for dtype in fallback_dtypes:
368
+ if dtype is None or is_deprecated(dtype):
369
+ continue
370
+
371
+ group_size = 0
372
+ scale_dtype = humming_dtypes.float32
373
+ quant_mode = InputQuantizationMode.DynamicToken
374
+ if dtype.num_bits == 16:
375
+ scale_dtype = None
376
+ quant_mode = InputQuantizationMode.Disabled
377
+ elif dtype in (humming_dtypes.float8e4m3, humming_dtypes.float8e3m4):
378
+ is_channelwise_weight = weight_schema.weight_scale_group_size == 0
379
+ weight_group_size = weight_schema.weight_scale_group_size
380
+ is_e8m0_scale = weight_schema.bs_dtype == humming_dtypes.float8e8m0
381
+ is_mx_weight = weight_group_size == 32 and is_e8m0_scale
382
+ if sm_version >= 120 and (is_channelwise_weight or is_mx_weight):
383
+ group_size = 32
384
+ scale_dtype = humming_dtypes.float8e8m0
385
+ quant_mode = InputQuantizationMode.DynamicGroup
386
+
387
+ candidate = HummingInputSchema(
388
+ a_dtype=dtype,
389
+ input_scale_group_size=group_size,
390
+ input_scale_dtype=scale_dtype,
391
+ input_quant_mode=quant_mode,
392
+ )
393
+ if candidate.is_compatible_with(weight_schema, param_dtype):
394
+ return candidate
395
+
396
+ raise ValueError(
397
+ f"No compatible Humming input schema for {input_schema} with "
398
+ f"weight schema {weight_schema} on SM{sm_version} "
399
+ f"({allow_fallback=})"
400
+ )
401
+
402
+
403
+def resolve_humming_layer_config(config: dict[str, Any], prefix: str) -> dict[str, Any]:
404
+ layer_config = config.copy()
405
+ dynamic = config.get("dynamic", {})
406
+ if isinstance(dynamic, dict):
407
+ for pattern, overrides in dynamic.items():
408
+ if pattern.startswith("+") and re.match(pattern[2:], prefix):
409
+ layer_config.update(overrides)
410
+ break
411
+ return layer_config
412
+
413
+
414
+def humming_is_layer_skipped(config: dict[str, Any], prefix: str):
415
+ if not config:
416
+ return True
417
+
418
+ keys = ["ignored_layers", "ignore", "modules_to_not_convert"]
419
+ ignored_layers: list[str] = []
420
+ for key in keys:
421
+ candidate = config.get(key, []) or []
422
+ if candidate:
423
+ ignored_layers = candidate
424
+ break
425
+
426
+ if any(module_name in prefix for module_name in ignored_layers):
427
+ return True
428
+ if "lm_head" in prefix:
429
+ return True
430
+
431
+ for regex in config.get("dynamic", {}):
432
+ if regex[:1] != "-":
433
+ continue
434
+ if re.match(regex[2:], prefix):
435
+ return True
436
+
437
+ return False
vllm/utils/humming.py
+52
-3
@@ -4,12 +4,56 @@
4
5
vLLM code should import humming symbols from here so that ``import humming``
6
(which has import-time side effects) is deferred until first use. Add new
7
-symbols by appending one entry to ``_EXPORTS`` as ``"module.path:attr"``,
8
-or ``"module.path"`` for a whole-module re-export.
7
+symbols to both the ``TYPE_CHECKING`` imports and ``_EXPORTS``, using
8
+``"module.path:attr"`` or ``"module.path"`` for a whole-module re-export.
9
"""
10
11
import importlib
12
-from typing import Any
12
+from typing import TYPE_CHECKING, Any
13
+
14
+if TYPE_CHECKING:
15
+ from humming import dtypes as dtypes
16
+ from humming.config import GemmType as GemmType
17
+ from humming.config import InputQuantizationMode as InputQuantizationMode
18
+ from humming.config import LayerConfig as LayerConfig
19
+ from humming.config import MmaType as MmaType
20
+ from humming.config import WeightScale2Type as WeightScale2Type
21
+ from humming.config import WeightScaleType as WeightScaleType
22
+ from humming.dtypes import DataType as DataType
23
+ from humming.forward import humming_forward as humming_forward
24
+ from humming.forward import may_process_input as may_process_input
25
+ from humming.forward import may_quant_input as may_quant_input
26
+ from humming.ops import process_input as process_input
27
+ from humming.schema import AWQWeightSchema as AWQWeightSchema
28
+ from humming.schema import BaseInputSchema as BaseInputSchema
29
+ from humming.schema import BaseWeightSchema as BaseWeightSchema
30
+ from humming.schema import BitnetWeightSchema as BitnetWeightSchema
31
+ from humming.schema import (
32
+ CompressedTensorsInputSchema as CompressedTensorsInputSchema,
33
+ )
34
+ from humming.schema import (
35
+ CompressedTensorsWeightSchema as CompressedTensorsWeightSchema,
36
+ )
37
+ from humming.schema import Fp8InputSchema as Fp8InputSchema
38
+ from humming.schema import GptOssMxfp4WeightSchema as GptOssMxfp4WeightSchema
39
+ from humming.schema import GPTQWeightSchema as GPTQWeightSchema
40
+ from humming.schema import HummingInputSchema as HummingInputSchema
41
+ from humming.schema import HummingWeightSchema as HummingWeightSchema
42
+ from humming.schema import Mxfp4WeightSchema as Mxfp4WeightSchema
43
+ from humming.schema.fp8 import Fp8WeightSchema as Fp8WeightSchema
44
+ from humming.schema.modelopt import (
45
+ ModeloptMxfp8WeightSchema as ModeloptMxfp8WeightSchema,
46
+ )
47
+ from humming.schema.modelopt import (
48
+ ModeloptNvfp4InputSchema as ModeloptNvfp4InputSchema,
49
+ )
50
+ from humming.schema.modelopt import (
51
+ ModeloptNvfp4WeightSchema as ModeloptNvfp4WeightSchema,
52
+ )
53
+ from humming.transform import prepare_layer_config as prepare_layer_config
54
+ from humming.transform import transform_humming_tensors as transform_humming_tensors
55
+ from humming.tune import get_heuristics_config as get_heuristics_config
56
+ from humming.utils.weight import quantize_weight as quantize_weight
57
58
_EXPORTS: dict[str, str] = {
59
"dtypes": "humming.dtypes",
@@ -17,7 +61,12 @@ _EXPORTS: dict[str, str] = {
61
"GemmType": "humming.config:GemmType",
62
"LayerConfig": "humming.config:LayerConfig",
63
"WeightScaleType": "humming.config:WeightScaleType",
64
+ "WeightScale2Type": "humming.config:WeightScale2Type",
65
+ "InputQuantizationMode": "humming.config:InputQuantizationMode",
66
+ "MmaType": "humming.config:MmaType",
67
"humming_forward": "humming.forward:humming_forward",
68
+ "may_process_input": "humming.forward:may_process_input",
69
+ "process_input": "humming.ops:process_input",
70
"may_quant_input": "humming.forward:may_quant_input",
71
"prepare_layer_config": "humming.transform:prepare_layer_config",
72
"transform_humming_tensors": "humming.transform:transform_humming_tensors",