From 51936faa0c8ae8e1c11e3514f9487732b6c7c53b Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:16:43 +0000 Subject: [PATCH 1/7] feat(gpt-oss): add MXFP4 CUDA dequant primitive Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- extensions/kernel/v4_fp4/mxfp4_dequant.cu | 45 +++++++++++++++++++++++ setup.py | 2 + 2 files changed, 47 insertions(+) create mode 100644 extensions/kernel/v4_fp4/mxfp4_dequant.cu diff --git a/extensions/kernel/v4_fp4/mxfp4_dequant.cu b/extensions/kernel/v4_fp4/mxfp4_dequant.cu new file mode 100644 index 00000000..e15a3139 --- /dev/null +++ b/extensions/kernel/v4_fp4/mxfp4_dequant.cu @@ -0,0 +1,45 @@ +#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/setup.py b/setup.py index a2d0614c..6e6697b5 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", # Python binding "core/python/py_archer_prefetch.cpp", ] @@ -354,6 +355,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", ], extra_compile_args={ "cxx": ["-O3", "-std=c++17", "-fPIC"], From e2c2b7698c8b5d9a754342b926d9c2e90d404b23 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:16:47 +0000 Subject: [PATCH 2/7] test(gpt-oss): expose exact MXFP4 dequant parity Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- extensions/kernel/v4_fp4/v4_fp4_binding.cpp | 24 ++++++++++++++ .../unit/test_gpt_oss_mxfp4_dispatch.py | 31 +++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 tests/python/unit/test_gpt_oss_mxfp4_dispatch.py diff --git a/extensions/kernel/v4_fp4/v4_fp4_binding.cpp b/extensions/kernel/v4_fp4/v4_fp4_binding.cpp index 2e87e36f..d72667a5 100644 --- a/extensions/kernel/v4_fp4/v4_fp4_binding.cpp +++ b/extensions/kernel/v4_fp4/v4_fp4_binding.cpp @@ -8,6 +8,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; +} // packed: [N, K/2] uint8 (view of float4_e2m1fn_x2); scale: [N, K/32] e8m0. // Returns dequantized BF16 weight [N, K]. @@ -51,6 +73,8 @@ torch::Tensor v4_expert_forward(torch::Tensor x, torch::Tensor w1, 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"); } 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..72067ae9 --- /dev/null +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -0,0 +1,31 @@ +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 From acbe510372e203955eab4ea9713ea1b3f1e14744 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:23:34 +0000 Subject: [PATCH 3/7] feat(gpt-oss): assign dedicated dispatcher expert type Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- moe_infinity/common/constants.py | 2 +- tests/test_gpt_oss_config.py | 9 +++------ 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/moe_infinity/common/constants.py b/moe_infinity/common/constants.py index 0ced40ad..ccb1cd45 100644 --- a/moe_infinity/common/constants.py +++ b/moe_infinity/common/constants.py @@ -41,7 +41,7 @@ "opt": 3, "deepseek_v3": 5, "deepseek": 5, - "gptoss": 4, + "gptoss": 6, "qwen3": 5, "dbrx": 4, "olmoe": 4, 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(): From ac143a6055852ebf934843466495a7de1784cf6c Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:23:38 +0000 Subject: [PATCH 4/7] feat(gpt-oss): execute fetched MXFP4 experts in Archer Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- core/parallel/expert_dispatcher.cpp | 7 ++ core/parallel/expert_module.cpp | 84 +++++++++++++++++-- core/parallel/expert_module.h | 7 +- .../unit/test_gpt_oss_mxfp4_dispatch.py | 81 ++++++++++++++++++ 4 files changed, 173 insertions(+), 6 deletions(-) diff --git a/core/parallel/expert_dispatcher.cpp b/core/parallel/expert_dispatcher.cpp index d78799c8..fb4f7fd8 100644 --- a/core/parallel/expert_dispatcher.cpp +++ b/core/parallel/expert_dispatcher.cpp @@ -96,6 +96,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); @@ -519,6 +522,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 7543088f..0273997f 100644 --- a/core/parallel/expert_module.cpp +++ b/core/parallel/expert_module.cpp @@ -9,6 +9,10 @@ #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); + static const int64_t kMaxTokens = 256; void ExpertNode::SetTensorsFromBlob(const torch::Device& device) { @@ -30,6 +34,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); } @@ -45,12 +51,19 @@ 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_)) @@ -68,8 +81,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}); @@ -99,8 +119,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 @@ -128,6 +155,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"); @@ -196,6 +249,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 b09f8c52..9a9dcd33 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); private: void ForwardHelper(cudaStream_t stream); @@ -201,6 +205,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/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py index 72067ae9..4c9f74e7 100644 --- a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -29,3 +29,84 @@ def test_native_mxfp4_gate_up_dequant_is_exact(): 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", + ).zero_() + 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 + + torch.testing.assert_close(option_a, resident, rtol=1e-2, atol=5e-2) From 367ab338f160a2e564e8a11ffe8338f4deab4e3e Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:28:34 +0000 Subject: [PATCH 5/7] style(gpt-oss): sort MXFP4 test imports Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- tests/python/unit/test_gpt_oss_mxfp4_dispatch.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py index 4c9f74e7..fb623872 100644 --- a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -7,7 +7,6 @@ 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) From 544ff7052cb30972257bbb39f446b8a52d04358f Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 10:48:20 +0000 Subject: [PATCH 6/7] test(gpt-oss): fp32-golden envelope parity Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../unit/test_gpt_oss_mxfp4_dispatch.py | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py index fb623872..81ca9dc2 100644 --- a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -50,7 +50,7 @@ def test_dequantized_option_a_matches_resident_expert_forward(): (2 * intermediate, hidden // 2), dtype=torch.uint8, device="cuda", - ).zero_() + ) gate_scales = torch.randint( 120, 135, @@ -108,4 +108,21 @@ def test_dequantized_option_a_matches_resident_expert_forward(): ) option_a = option_a_activated @ down_weight.t() + down_bias - torch.testing.assert_close(option_a, resident, rtol=1e-2, atol=5e-2) + 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() From 97ae90798cb33065b86a5a7adddf5f11e16a848e Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 19:44:06 +0000 Subject: [PATCH 7/7] style: apply ruff-format and clang-format --- core/parallel/expert_module.cpp | 33 +++++++++---------- extensions/kernel/v4_fp4/mxfp4_dequant.cu | 23 ++++++------- extensions/kernel/v4_fp4/v4_fp4_binding.cpp | 13 ++++---- moe_infinity/runtime/model_offload.py | 4 +-- .../unit/test_gpt_oss_mxfp4_dispatch.py | 17 ++++++---- 5 files changed, 45 insertions(+), 45 deletions(-) diff --git a/core/parallel/expert_module.cpp b/core/parallel/expert_module.cpp index 0273997f..ca7ff9d9 100644 --- a/core/parallel/expert_module.cpp +++ b/core/parallel/expert_module.cpp @@ -10,8 +10,8 @@ #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); + int rows, int packed_cols, int scale_cols, + int block_size, cudaStream_t stream); static const int64_t kMaxTokens = 256; @@ -51,8 +51,7 @@ MoEMLP::MoEMLP(int dtype, int expert_type) { for (int i = 0; i < 8; i++) { buffer_.push_back(torch::zeros({1}, options)); } - int num_params = - expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE ? 6 : 4; + 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)); } @@ -169,12 +168,12 @@ void MoEMLP::DequantMxfp4Params(cudaStream_t stream) { 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))); + 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); + 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]); @@ -254,14 +253,14 @@ void MoEMLP::ForwardHelper(cudaStream_t stream) { 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)}); + 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)); diff --git a/extensions/kernel/v4_fp4/mxfp4_dequant.cu b/extensions/kernel/v4_fp4/mxfp4_dequant.cu index e15a3139..68f67ee0 100644 --- a/extensions/kernel/v4_fp4/mxfp4_dequant.cu +++ b/extensions/kernel/v4_fp4/mxfp4_dequant.cu @@ -9,9 +9,11 @@ __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) { +__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; @@ -21,10 +23,10 @@ __global__ void mxfp4_dequant_kernel( 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; + 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)); } @@ -32,14 +34,13 @@ __global__ void mxfp4_dequant_kernel( } // 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) { + 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(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 d72667a5..d42e6b30 100644 --- a/extensions/kernel/v4_fp4/v4_fp4_binding.cpp +++ b/extensions/kernel/v4_fp4/v4_fp4_binding.cpp @@ -9,8 +9,8 @@ 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); + 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"); @@ -23,11 +23,11 @@ torch::Tensor mxfp4_dequant(torch::Tensor packed, torch::Tensor scales) { 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 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); + rows, packed_cols, scale_cols, block_size, stream); return output; } @@ -73,8 +73,7 @@ torch::Tensor v4_expert_forward(torch::Tensor x, torch::Tensor w1, 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("mxfp4_dequant", &mxfp4_dequant, "MXFP4 uint8 blocks/scales to BF16"); m.def("v4_expert_forward", &v4_expert_forward, "V4 FP4 routed-expert SwiGLU forward"); } diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index a5f88c5d..3db35d70 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -738,9 +738,7 @@ def archer_from_pretrained(cls, *args, **kwargs): state_dict = torch.load(ckpt) _remap_v5_batched_experts(state_dict, self.config) - _expand_gpt_oss_packed_experts( - state_dict, self.config - ) + _expand_gpt_oss_packed_experts(state_dict, self.config) is_gptq_ckpt = is_gptq_quantized(self.config) self._cast_state_dict_tensors( diff --git a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py index 81ca9dc2..bbb13cfc 100644 --- a/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py +++ b/tests/python/unit/test_gpt_oss_mxfp4_dispatch.py @@ -77,9 +77,7 @@ def test_dequantized_option_a_matches_resident_expert_forward(): ) 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_up = fused_mxfp4_gemm(x, gate_blocks, gate_scales, gate_bias) resident_gate, resident_up = ( resident_gate_up[:, ::2], resident_gate_up[:, 1::2], @@ -120,9 +118,14 @@ def test_dequantized_option_a_matches_resident_expert_forward(): # 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 + 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()