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 15, 13:30 UTC
[BugFix] Accuracy fix for llama4 int4 - improperly casted scales (#16801)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
Apr 18, 2025 at 01:13 UTC
7eb42556281d30436a3a988f2c9184ec63c59338
3 files changed
+6
-9
csrc/moe/moe_wna16.cu
+3
-7
@@ -13,7 +13,6 @@
13
template <typename scalar_t, int bit, int GROUPS>
14
__global__ void moe_wna16_gemm_kernel(
15
const scalar_t* __restrict__ input, scalar_t* __restrict__ output,
16
-
16
const uint32_t* __restrict__ qweight, const scalar_t* __restrict__ scales,
17
const uint32_t* __restrict__ qzeros,
18
@@ -54,8 +53,6 @@ __global__ void moe_wna16_gemm_kernel(
53
if (token_index / top_k >= size_m) break;
54
55
num_valid_tokens = m + 1;
57
- if (blockIdx.z == 0 && offset_n < size_n)
58
- output[token_index * size_n + offset_n] = Dtype::int2num(0);
56
57
if (expert_id != -1) {
58
int k_per_thread = DIVIDE(BLOCK_SIZE_K, BLOCK_SIZE_N);
@@ -284,8 +281,7 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
281
int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
282
int64_t BLOCK_SIZE_K, int64_t bit) {
283
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
287
- auto options =
288
- torch::TensorOptions().dtype(input.dtype()).device(input.device());
284
+ output.zero_();
285
286
const int num_experts = b_qweight.size(0);
287
const int size_m = input.size(0);
@@ -302,9 +298,9 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
298
const uint32_t* b_qzeros_ptr;
299
if (b_qzeros.has_value())
300
b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
305
- const float* topk_weights_ptr;
301
+ const float* topk_weights_ptr = nullptr;
302
if (topk_weights.has_value())
307
- topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
303
+ topk_weights_ptr = (const float*)topk_weights.value().data_ptr<float>();
304
305
int groups_per_block_row = BLOCK_SIZE_K / group_size;
306
TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
vllm/model_executor/layers/fused_moe/layer.py
+1
@@ -422,6 +422,7 @@ class FusedMoE(torch.nn.Module):
422
423
if params_dtype is None:
424
params_dtype = torch.get_default_dtype()
425
+ self.params_dtype = params_dtype
426
427
# Note: here we guard against accessing the TP and DP groups when
428
# uninitialized (this happens when testing)
vllm/model_executor/models/llama4.py
+2
-2
@@ -51,8 +51,8 @@ class Llama4MoE(nn.Module):
51
renormalize: bool,
52
) -> Tuple[torch.Tensor, torch.Tensor]:
53
router_scores, router_indices = fast_topk(gating_output, topk, dim=-1)
54
- router_scores = torch.sigmoid(router_scores.float()).to(
55
- hidden_states.dtype)
54
+ # psuedo-standard is that the router scores are floats
55
+ router_scores = torch.sigmoid(router_scores.float())
56
return (router_scores, router_indices.to(torch.int32))
57
58
def __init__(self,