From 98913c9b1a17389770f2d604c5b3158ca0066b1f Mon Sep 17 00:00:00 2001 From: biondizzle Date: Thu, 14 May 2026 09:38:50 +0000 Subject: [PATCH] Fix stage_activation: use Triton staging kernel from vLLM patch instead of fp4_utils --- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 8eb158fa..62e96311 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -161,13 +161,13 @@ def nvfp4_mega_moe_l2( def stage_activation(x_bf16): """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales. - - This replaces the Triton staging kernel from patches/staging_kernel.py. + + Uses the Triton staging kernel from the vLLM deepseek_v4 patch. """ - from vllm.model_executor.layers.quantization.utils.fp4_utils import ( - per_token_cast_to_fp4, + from vllm.model_executor.models.staging_kernel import ( + _stage_deepseek_v4_mega_moe_inputs, ) - return per_token_cast_to_fp4(x_bf16) + return _stage_deepseek_v4_mega_moe_inputs(x_bf16) def nvfp4_mega_moe_full(