diff --git a/core/parallel/expert_dispatcher.cpp b/core/parallel/expert_dispatcher.cpp index 84bce7ad..be39ce06 100644 --- a/core/parallel/expert_dispatcher.cpp +++ b/core/parallel/expert_dispatcher.cpp @@ -163,6 +163,9 @@ ExpertDispatcher::ExpertDispatcher(int num_experts, int num_layers, int dtype, case DEEPSEEK_MOE_DENSE_ACT_DENSE: experts_[i][j]->module = new DeepSeekMoEDenseActDense(dtype); break; + case GPT_OSS_MOE_DENSE_ACT_DENSE: + experts_[i][j]->module = new GptOssMoeDenseActDense(); + break; default: DLOG_FATAL("ExpertDispatcher::ExpertDispatcher: unknown expert type ", expert_type); @@ -597,6 +600,10 @@ void ExpertDispatcher::GPUExecFunc(int gpu_id, int thread_idx) { c10::cuda::getStreamFromExternal(stream, gpu_id); c10::cuda::CUDAStreamGuard guard(torch_stream); + if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { + modules_[thread_idx]->DequantMxfp4Params(stream); + } + torch::Tensor output; { #ifndef NVTX_DISABLE diff --git a/core/parallel/expert_module.cpp b/core/parallel/expert_module.cpp index e5a9a2f2..dd336985 100644 --- a/core/parallel/expert_module.cpp +++ b/core/parallel/expert_module.cpp @@ -8,6 +8,9 @@ #include "utils/logger.h" #include "kernel/fused_moe_mlp.h" +void mxfp4_dequant_cuda(const void* packed, const void* scales, void* output, + int rows, int packed_cols, int scale_cols, + int block_size, cudaStream_t stream); extern void fp8_dequant_blockwise_cuda(const void* weight, const void* scale, void* out, int N, int K, cudaStream_t stream); @@ -33,6 +36,8 @@ void ExpertNode::SetTensorsFromBlob(const torch::Device& device) { reinterpret_cast(module)->SetTensorsFromBlob( node->device_memory_ptr, node->tensor_ids, device); break; + case ExpertType::GptOssMoeDenseActDense: + break; default: assert(false); } @@ -48,12 +53,18 @@ MoEMLP::MoEMLP(int dtype, int expert_type) { for (int i = 0; i < 8; i++) { buffer_.push_back(torch::zeros({1}, options)); } - for (int i = 0; i < 4; i++) { + int num_params = expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE ? 6 : 4; + for (int i = 0; i < num_params; i++) { param_.push_back(torch::zeros({1}, options)); } } void MoEMLP::SetTensorsFromIds(const std::vector& tensor_ids) { + if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { + DLOG_FATAL_IF(tensor_ids.size() != 6, + "GPT-OSS expert requires blocks/scales/bias for two " + "projections"); + } int device = at::cuda::current_device(); auto options = torch::TensorOptions() .dtype(dtype_to_torch(dtype_)) @@ -71,8 +82,15 @@ void MoEMLP::SetTensorsFromIds(const std::vector& tensor_ids) { if (!param_init_) { auto allocator = c10::DeviceCachingAllocator::get(device); - int64_t hdim = tensor_shapes[0][1]; - int64_t idim = tensor_shapes[0][0]; + int64_t hdim; + int64_t idim; + if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { + hdim = tensor_shapes[0][1] * 2; + idim = tensor_shapes[0][0] / 2; + } else { + hdim = tensor_shapes[0][1]; + idim = tensor_shapes[0][0]; + } std::vector> data_shapes; data_shapes.push_back({kMaxTokens, hdim}); @@ -102,8 +120,15 @@ void MoEMLP::SetTensorsFromIds(const std::vector& tensor_ids) { // MLP tensor shape: weight is [intermediate, hidden], so // hdim = tensor_shapes[0][1], idim = tensor_shapes[0][0] - int64_t hdim = tensor_shapes[0][1]; - int64_t idim = tensor_shapes[0][0]; + int64_t hdim; + int64_t idim; + if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { + hdim = tensor_shapes[0][1] * 2; + idim = tensor_shapes[0][0] / 2; + } else { + hdim = tensor_shapes[0][1]; + idim = tensor_shapes[0][0]; + } std::vector> data_shapes; data_shapes.push_back({kMaxTokens, hdim}); // input buffer @@ -131,6 +156,32 @@ void MoEMLP::SetTensorsFromIds(const std::vector& tensor_ids) { param_set_ = true; } +void MoEMLP::DequantMxfp4Params(cudaStream_t stream) { + if (expert_type_ != GPT_OSS_MOE_DENSE_ACT_DENSE) return; + int device = at::cuda::current_device(); + gpt_oss_param_.clear(); + for (auto pair : {std::pair{0, 1}, {3, 4}}) { + auto packed = param_[pair.first].contiguous(); + auto scales = param_[pair.second].contiguous(); + DLOG_FATAL_IF(packed.scalar_type() != torch::kUInt8 || + scales.scalar_type() != torch::kUInt8, + "GPT-OSS MXFP4 blocks/scales must be uint8"); + int rows = packed.size(0); + int packed_cols = packed.size(1); + int scale_cols = scales.size(1); + int block_size = packed_cols * 2 / scale_cols; + auto output = + torch::empty({rows, packed_cols * 2}, torch::TensorOptions() + .dtype(torch::kBFloat16) + .device(CUDA_DEVICE(device))); + mxfp4_dequant_cuda(packed.data_ptr(), scales.data_ptr(), output.data_ptr(), + rows, packed_cols, scale_cols, block_size, stream); + gpt_oss_param_.push_back(output); + } + gpt_oss_param_.insert(gpt_oss_param_.begin() + 1, param_[2]); + gpt_oss_param_.push_back(param_[5]); +} + torch::Tensor MoEMLP::forward(torch::Tensor hidden_states, cudaStream_t stream) { DLOG_FATAL_IF(param_set_ == false, "param_set_ should be true"); @@ -199,6 +250,27 @@ void MoEMLP::ForwardHelper(cudaStream_t stream) { auto& input = buffer_[0]; auto& output = buffer_[1]; + if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { + auto& gate_up_weight = gpt_oss_param_[0]; + auto& gate_up_bias = gpt_oss_param_[1]; + auto& down_weight = gpt_oss_param_[2]; + auto& down_bias = gpt_oss_param_[3]; + auto gate_up = + torch::matmul(input, gate_up_weight.transpose(0, 1)) + gate_up_bias; + auto gate = + gate_up.index({torch::indexing::Slice(), + torch::indexing::Slice(0, torch::indexing::None, 2)}); + auto up = + gate_up.index({torch::indexing::Slice(), + torch::indexing::Slice(1, torch::indexing::None, 2)}); + gate = torch::clamp_max(gate, 7.0); + up = torch::clamp(up, -7.0, 7.0); + auto activated = (up + 1.0) * (gate * torch::sigmoid(gate * 1.702)); + output.copy_(torch::matmul(activated, down_weight.transpose(0, 1)) + + down_bias); + return; + } + if (expert_type_ == NLLB_MOE_DENSE_ACT_DENSE) { auto& fc1 = param_[0]; auto& fc2 = param_[1]; diff --git a/core/parallel/expert_module.h b/core/parallel/expert_module.h index da4cee63..3af2ce8e 100644 --- a/core/parallel/expert_module.h +++ b/core/parallel/expert_module.h @@ -14,7 +14,8 @@ enum class ExpertType { NllbMoeDenseActDense = 2, FSGPTMoeDenseActDense = 3, MixtralMoeDenseActDense = 4, - DeepSeekMoeDenseActDense = 5 + DeepSeekMoeDenseActDense = 5, + GptOssMoeDenseActDense = 6 }; // Activation functions enum @@ -172,6 +173,7 @@ using NllbMoeDenseActDense = Expert; using FSGPTMoEDenseActDense = Expert; using MixtralMoEDenseActDense = Expert; using DeepSeekMoEDenseActDense = Expert; +struct GptOssMoeDenseActDense : public torch::nn::Module {}; #ifndef EXPERT_TYPE #define EXPERT_TYPE 0 @@ -181,6 +183,7 @@ using DeepSeekMoEDenseActDense = Expert; #define FSGPT_MOE_DENSE_ACT_DENSE 3 #define MIXTRAL_MOE_DENSE_ACT_DENSE 4 #define DEEPSEEK_MOE_DENSE_ACT_DENSE 5 +#define GPT_OSS_MOE_DENSE_ACT_DENSE 6 // forward declarations torch::Tensor launch_fused_moe_ffn(torch::Tensor hidden, // [M, K] @@ -194,6 +197,7 @@ struct MoEMLP : public torch::nn::Module { torch::Tensor forward(torch::Tensor hidden_states, cudaStream_t stream); void SetTensorsFromIds(const std::vector& tensor_ids); + void DequantMxfp4Params(cudaStream_t stream); void SetFp8Scales(const std::vector& scales); void DequantFp8Params(cudaStream_t stream); @@ -203,6 +207,7 @@ struct MoEMLP : public torch::nn::Module { private: std::vector buffer_; std::vector param_; + std::vector gpt_oss_param_; at::cuda::CUDAGraph graph_; int warmup_count_ = 5; diff --git a/extensions/kernel/v4_fp4/mxfp4_dequant.cu b/extensions/kernel/v4_fp4/mxfp4_dequant.cu new file mode 100644 index 00000000..68f67ee0 --- /dev/null +++ b/extensions/kernel/v4_fp4/mxfp4_dequant.cu @@ -0,0 +1,46 @@ +#include +#include + +#include + +namespace { + +__device__ __constant__ float kMxfp4Values[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + +__global__ void mxfp4_dequant_kernel(const uint8_t* packed, + const uint8_t* scales, + __nv_bfloat16* output, int rows, + int packed_cols, int scale_cols, + int block_size) { + long output_col = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + long output_cols = static_cast(packed_cols) * 2; + long total = static_cast(rows) * output_cols; + if (output_col >= total) return; + + int row = static_cast(output_col / output_cols); + int col = static_cast(output_col % output_cols); + uint8_t byte = packed[static_cast(row) * packed_cols + col / 2]; + uint8_t code = (col & 1) == 0 ? (byte & 0x0F) : (byte >> 4); + int exponent = + static_cast( + scales[static_cast(row) * scale_cols + col / block_size]) - + 127; + exponent = max(-126, min(127, exponent)); + output[output_col] = __float2bfloat16(ldexpf(kMxfp4Values[code], exponent)); +} + +} // namespace + +void mxfp4_dequant_cuda(const void* packed, const void* scales, void* output, + int rows, int packed_cols, int scale_cols, + int block_size, cudaStream_t stream) { + long total = static_cast(rows) * packed_cols * 2; + constexpr int threads = 256; + int blocks = static_cast((total + threads - 1) / threads); + mxfp4_dequant_kernel<<>>( + static_cast(packed), static_cast(scales), + static_cast<__nv_bfloat16*>(output), rows, packed_cols, scale_cols, + block_size); +} diff --git a/extensions/kernel/v4_fp4/v4_fp4_binding.cpp b/extensions/kernel/v4_fp4/v4_fp4_binding.cpp index a88ef556..79401815 100644 --- a/extensions/kernel/v4_fp4/v4_fp4_binding.cpp +++ b/extensions/kernel/v4_fp4/v4_fp4_binding.cpp @@ -11,6 +11,28 @@ void fp4_dequant_to_bf16(const void* packed, const void* scale_e8m0, void* out, int N, int K, cudaStream_t stream); +void mxfp4_dequant_cuda(const void* packed, const void* scales, void* output, + int rows, int packed_cols, int scale_cols, + int block_size, cudaStream_t stream); + +torch::Tensor mxfp4_dequant(torch::Tensor packed, torch::Tensor scales) { + TORCH_CHECK(packed.is_cuda() && scales.is_cuda(), "CUDA tensors required"); + TORCH_CHECK(packed.scalar_type() == torch::kUInt8, "packed must be uint8"); + TORCH_CHECK(scales.scalar_type() == torch::kUInt8, "scales must be uint8"); + TORCH_CHECK(packed.dim() == 2 && scales.dim() == 2, "2D tensors required"); + int rows = packed.size(0); + int packed_cols = packed.size(1); + int scale_cols = scales.size(1); + int block_size = packed_cols * 2 / scale_cols; + packed = packed.contiguous(); + scales = scales.contiguous(); + auto output = torch::empty({rows, packed_cols * 2}, + packed.options().dtype(torch::kBFloat16)); + auto stream = at::cuda::getCurrentCUDAStream(packed.device().index()); + mxfp4_dequant_cuda(packed.data_ptr(), scales.data_ptr(), output.data_ptr(), + rows, packed_cols, scale_cols, block_size, stream); + return output; +} void fp8_dequant_blockwise_cuda(const void* weight, const void* scale, void* out, int N, int K, cudaStream_t stream); @@ -92,6 +114,7 @@ void set_scales(const std::map& /*scales*/) { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fp4_dequant", &fp4_dequant, "FP4 E2M1 packed -> BF16 dequant"); + m.def("mxfp4_dequant", &mxfp4_dequant, "MXFP4 uint8 blocks/scales to BF16"); m.def("v4_expert_forward", &v4_expert_forward, "V4 FP4 routed-expert SwiGLU forward"); m.def("fp8_dequant_blockwise", &fp8_dequant_blockwise, diff --git a/moe_infinity/common/constants.py b/moe_infinity/common/constants.py index 0a54419f..64b5943b 100644 --- a/moe_infinity/common/constants.py +++ b/moe_infinity/common/constants.py @@ -46,7 +46,7 @@ "opt": 3, "deepseek_v3": 5, "deepseek": 5, - "gptoss": 4, + "gptoss": 6, "qwen3": 5, "dbrx": 4, "olmoe": 4, diff --git a/setup.py b/setup.py index 250ac569..ea0e4818 100644 --- a/setup.py +++ b/setup.py @@ -221,6 +221,7 @@ def _find_nvtx_include_dir() -> Optional[str]: "extensions/kernel/fused_moe_mlp.cu", "extensions/kernel/activation_kernels.cu", "extensions/kernel/topk_softmax_kernels.cu", + "extensions/kernel/v4_fp4/mxfp4_dequant.cu", "extensions/kernel/v4_fp4/fp8_dequant.cu", # Python binding "core/python/py_archer_prefetch.cpp", @@ -366,6 +367,7 @@ def _find_nvtx_include_dir() -> Optional[str]: sources=[ "extensions/kernel/v4_fp4/v4_fp4_binding.cpp", "extensions/kernel/v4_fp4/v4_fp4_dequant.cu", + "extensions/kernel/v4_fp4/mxfp4_dequant.cu", "extensions/kernel/v4_fp4/fp8_dequant.cu", ], extra_compile_args={ diff --git a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py new file mode 100644 index 00000000..bbb13cfc --- /dev/null +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -0,0 +1,131 @@ +import pytest +import torch + + +@pytest.mark.gpu +def test_native_mxfp4_gate_up_dequant_is_exact(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + from moe_infinity._v4_fp4 import mxfp4_dequant + from moe_infinity.kernel.mxfp4_gemm import mxfp4_dequantize + + torch.manual_seed(137) + blocks = torch.randint( + 0, 256, (5760, 1440), dtype=torch.uint8, device="cuda" + ) + scales = torch.randint( + 120, 135, (5760, 90), dtype=torch.uint8, device="cuda" + ) + expected = mxfp4_dequantize( + blocks, scales, dtype=torch.bfloat16, block_size=32 + ) + actual = mxfp4_dequant(blocks, scales) + + relative_error = ( + (actual.float() - expected.float()).abs() + / expected.float().abs().clamp_min(1e-12) + ).max() + assert actual.shape == (5760, 2880) + assert actual.dtype == torch.bfloat16 + assert relative_error.item() == 0.0 + + +@pytest.mark.gpu +def test_dequantized_option_a_matches_resident_expert_forward(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + try: + from moe_infinity._v4_fp4 import mxfp4_dequant + except Exception: + pytest.skip("native MXFP4 dequant extension not built") + + from moe_infinity.kernel.mxfp4_gemm import fused_mxfp4_gemm + + torch.manual_seed(137) + tokens, hidden, intermediate = 3, 64, 32 + x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") + gate_blocks = torch.randint( + 0, + 256, + (2 * intermediate, hidden // 2), + dtype=torch.uint8, + device="cuda", + ) + gate_scales = torch.randint( + 120, + 135, + (2 * intermediate, hidden // 32), + dtype=torch.uint8, + device="cuda", + ) + down_blocks = torch.randint( + 0, + 256, + (hidden, intermediate // 2), + dtype=torch.uint8, + device="cuda", + ) + down_scales = torch.randint( + 120, + 135, + (hidden, intermediate // 32), + dtype=torch.uint8, + device="cuda", + ) + gate_bias = torch.randn( + 2 * intermediate, dtype=torch.bfloat16, device="cuda" + ) + down_bias = torch.randn(hidden, dtype=torch.bfloat16, device="cuda") + + resident_gate_up = fused_mxfp4_gemm(x, gate_blocks, gate_scales, gate_bias) + resident_gate, resident_up = ( + resident_gate_up[:, ::2], + resident_gate_up[:, 1::2], + ) + resident_activated = (resident_up.clamp(-7, 7) + 1) * ( + resident_gate.clamp(max=7) + * torch.sigmoid(resident_gate.clamp(max=7) * 1.702) + ) + resident = fused_mxfp4_gemm( + resident_activated.to(torch.bfloat16), + down_blocks, + down_scales, + down_bias, + ) + + gate_weight = mxfp4_dequant(gate_blocks, gate_scales) + down_weight = mxfp4_dequant(down_blocks, down_scales) + option_a_gate_up = x @ gate_weight.t() + gate_bias + option_a_gate, option_a_up = ( + option_a_gate_up[:, ::2], + option_a_gate_up[:, 1::2], + ) + option_a_activated = (option_a_up.clamp(-7, 7) + 1) * ( + option_a_gate.clamp(max=7) + * torch.sigmoid(option_a_gate.clamp(max=7) * 1.702) + ) + option_a = option_a_activated @ down_weight.t() + down_bias + + gate_weight_f = gate_weight.float() + down_weight_f = down_weight.float() + golden_gate_up = x.float() @ gate_weight_f.t() + gate_bias.float() + golden_gate, golden_up = golden_gate_up[:, ::2], golden_gate_up[:, 1::2] + golden_activated = (golden_up.clamp(-7, 7) + 1) * ( + golden_gate.clamp(max=7) + * torch.sigmoid(golden_gate.clamp(max=7) * 1.702) + ) + golden = golden_activated @ down_weight_f.t() + down_bias.float() + + # Bound bf16 rounding by the down-GEMM magnitude instead of comparing two + # cancellation-sensitive bf16 paths directly. + envelope = ( + 8 + * (2**-8) + * ( + golden_activated.abs() @ down_weight_f.abs().t() + + down_bias.float().abs() + ) + + 1e-2 + ) + assert ((option_a.float() - golden).abs() <= envelope).all() + assert ((resident.float() - golden).abs() <= envelope).all() diff --git a/tests/test_gpt_oss_config.py b/tests/test_gpt_oss_config.py index 9ed5710e..4fc24777 100644 --- a/tests/test_gpt_oss_config.py +++ b/tests/test_gpt_oss_config.py @@ -94,13 +94,10 @@ def test_gpt_oss_model_mapping(): assert MODEL_MAPPING_NAMES["gptoss"] is GptOssForCausalLM -def test_gpt_oss_model_type(): - MODEL_MAPPING_TYPES = import_constants_module().MODEL_MAPPING_TYPES +def test_gpt_oss_has_dedicated_dispatcher_expert_type(): + parse_expert_type = import_constants_module().parse_expert_type - assert ( - "gptoss" in MODEL_MAPPING_TYPES - ), "gptoss key missing from MODEL_MAPPING_TYPES" - assert MODEL_MAPPING_TYPES["gptoss"] == 4 + assert parse_expert_type(make_gpt_oss_config()) == 6 def test_gpt_oss_arch_string_matching():