A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 13:15 UTC
[Bugfix] Fix persistent_masked_m_silu_mul_quant tests (#28366)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Nov 10, 2025 at 12:21 UTC
b039bfda8f72b442d42dbdac40f51572bf045ad1
3 files changed
+16
-7
csrc/quantization/activation_kernels.cu
+10
-5
@@ -578,11 +578,13 @@ void persistent_masked_m_silu_mul_quant(
578
579
// This kernel currently only supports H % 128 == 0 and assumes a
580
// fixed GROUP_SIZE of 128.
581
+ static constexpr int GROUP_SIZE = 128;
582
+
583
TORCH_CHECK(input.dtype() == torch::kBFloat16);
584
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
585
y_q.dtype() == torch::kFloat8_e4m3fnuz);
586
TORCH_CHECK(y_s.dtype() == torch::kFloat32);
585
- TORCH_CHECK(input.size(-1) % 256 == 0);
587
+ TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
588
589
using Idx_t = int64_t;
590
@@ -601,8 +603,6 @@ void persistent_masked_m_silu_mul_quant(
603
604
Idx_t stride_counts_e = tokens_per_expert.stride(0);
605
604
- static constexpr int GROUP_SIZE = 128;
605
-
606
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
607
608
#define KERNEL(BLOCK_COUNT, USE_UE8M0, THREAD_COUNT, STAGES) \
@@ -628,21 +628,26 @@ void persistent_masked_m_silu_mul_quant(
628
629
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
630
631
+ int const NUM_GROUPS = H / GROUP_SIZE;
632
if (!use_ue8m0) {
632
- if (H >= 4096) {
633
+ if (H >= 4096 && (NUM_GROUPS % 8 == 0)) {
634
+ /* 8 warps config */
635
static constexpr int NUM_STAGES = 4;
636
static constexpr int THREAD_COUNT = 256;
637
KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, NUM_STAGES);
638
} else {
639
+ /* 1 warp config */
640
static constexpr int THREAD_COUNT = 32;
641
KERNEL(SILU_V2_BLOCK_COUNT, false, THREAD_COUNT, 2);
642
}
643
} else {
641
- if (H >= 4096) {
644
+ if (H >= 4096 && (NUM_GROUPS % 8 == 0)) {
645
+ /* 8 warps config */
646
static constexpr int NUM_STAGES = 4;
647
static constexpr int THREAD_COUNT = 256;
648
KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, NUM_STAGES);
649
} else {
650
+ /* 1 warp config */
651
static constexpr int THREAD_COUNT = 32;
652
KERNEL(SILU_V2_BLOCK_COUNT, true, THREAD_COUNT, 2);
653
}
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+4
-1
@@ -25,6 +25,7 @@ CASES = [
25
(8, 16, 128 * 2, fp8_dtype),
26
(8, 16, 128 * 3, fp8_dtype),
27
(8, 64, 7168, fp8_dtype),
28
+ (8, 128, 128 * 33, fp8_dtype),
29
(8, 128, 7168, fp8_dtype),
30
(8, 512, 7168, fp8_dtype),
31
(8, 1024, 7168, fp8_dtype),
@@ -54,8 +55,10 @@ def test_silu_mul_fp8_quant_deep_gemm(E, T, H, fp8_type):
55
)
56
57
# Run the SiLU V2 kernel
58
+ # TODO (varun): use_e8m0 is set to false as the reference impl does
59
+ # not handle that case.
60
y_q, y_s = persistent_masked_m_silu_mul_quant(
58
- y, tokens_per_expert, group_size=group_size
61
+ y, tokens_per_expert, group_size=group_size, use_ue8m0=False
62
)
63
64
torch.cuda.synchronize()
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+2
-1
@@ -100,6 +100,7 @@ def persistent_masked_m_silu_mul_quant(
100
tokens_per_expert: torch.Tensor, # (E,) number of valid tokens per expert
101
num_parallel_tokens=16,
102
group_size: int = 128,
103
+ use_ue8m0: bool | None = None,
104
) -> tuple[torch.Tensor, torch.Tensor]:
105
"""Quantize silu(y[..., :H]) * y[..., H:] to FP8 with group per-token scales
106
y has shape (E, T, 2*H). The first half of the last dimension is
@@ -164,7 +165,7 @@ def persistent_masked_m_silu_mul_quant(
165
device=y.device,
166
)
167
167
- use_ue8m0 = is_deep_gemm_e8m0_used()
168
+ use_ue8m0 = use_ue8m0 if use_ue8m0 is not None else is_deep_gemm_e8m0_used()
169
170
cuda_arch = current_platform.get_device_capability(
171
device_id=y.device.index