feat(gpt-oss): expert offloading Phase 2 — MXFP4 dequant-on-copy execution

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

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<DeepSeekMoEDenseActDense*>(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<std::uint32_t>& 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<std::uint32_t>& 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<std::vector<int64_t>> data_shapes;
     data_shapes.push_back({kMaxTokens, hdim});
@@ -102,8 +120,15 @@ void MoEMLP::SetTensorsFromIds(const std::vector<std::uint32_t>& 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<std::vector<int64_t>> data_shapes;
     data_shapes.push_back({kMaxTokens, hdim});  // input buffer
@@ -131,6 +156,32 @@ void MoEMLP::SetTensorsFromIds(const std::vector<std::uint32_t>& 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<int, int>{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];

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<ExpertType::NllbMoeDenseActDense>;
 using FSGPTMoEDenseActDense = Expert<ExpertType::FSGPTMoeDenseActDense>;
 using MixtralMoEDenseActDense = Expert<ExpertType::MixtralMoeDenseActDense>;
 using DeepSeekMoEDenseActDense = Expert<ExpertType::DeepSeekMoeDenseActDense>;
+struct GptOssMoeDenseActDense : public torch::nn::Module {};
 
 #ifndef EXPERT_TYPE
   #define EXPERT_TYPE 0
@@ -181,6 +183,7 @@ using DeepSeekMoEDenseActDense = Expert<ExpertType::DeepSeekMoeDenseActDense>;
 #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<std::uint32_t>& tensor_ids);
+  void DequantMxfp4Params(cudaStream_t stream);
   void SetFp8Scales(const std::vector<torch::Tensor>& scales);
   void DequantFp8Params(cudaStream_t stream);
 
@@ -203,6 +207,7 @@ struct MoEMLP : public torch::nn::Module {
  private:
   std::vector<torch::Tensor> buffer_;
   std::vector<torch::Tensor> param_;
+  std::vector<torch::Tensor> gpt_oss_param_;
 
   at::cuda::CUDAGraph graph_;
   int warmup_count_ = 5;

extensions/kernel/v4_fp4/mxfp4_dequant.cu

@@ -0,0 +1,46 @@
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+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<long>(blockIdx.x) * blockDim.x + threadIdx.x;
+  long output_cols = static_cast<long>(packed_cols) * 2;
+  long total = static_cast<long>(rows) * output_cols;
+  if (output_col >= total) return;
+
+  int row = static_cast<int>(output_col / output_cols);
+  int col = static_cast<int>(output_col % output_cols);
+  uint8_t byte = packed[static_cast<long>(row) * packed_cols + col / 2];
+  uint8_t code = (col & 1) == 0 ? (byte & 0x0F) : (byte >> 4);
+  int exponent =
+      static_cast<int>(
+          scales[static_cast<long>(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<long>(rows) * packed_cols * 2;
+  constexpr int threads = 256;
+  int blocks = static_cast<int>((total + threads - 1) / threads);
+  mxfp4_dequant_kernel<<<blocks, threads, 0, stream>>>(
+      static_cast<const uint8_t*>(packed), static_cast<const uint8_t*>(scales),
+      static_cast<__nv_bfloat16*>(output), rows, packed_cols, scale_cols,
+      block_size);
+}

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<std::string, torch::Tensor>& /*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,

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,

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={