submission 537658

  """
- MXFP4 MoE submission v1 - AITER fused_moe baseline.
- Full pipeline: quant -> gate+up GEMM -> SwiGLU -> down GEMM -> weighted reduce
+ MoE MXFP4 v29 — OPUS sorting + doweight_stage1 exploration.
+  
+ From probe 16:
+ - AITER_USE_OPUS_MOE_SORTING=1 → uses moe_sorting_opus_fwd (potentially faster)
+ - doweight_stage1=True → fuses weight multiply into stage1 GEMM
+ - splitk parameter → parallelize K dimension
+  
+ Strategy: Test OPUS sorting first (safest change), then doweight_stage1.
+ NO CSV changes — avoid the benchmark timeout from CSV injection.
  """
+ import os
+  
+ # Environment vars BEFORE importing aiter
+ os.environ["AITER_USE_NT"] = "1"
+ os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
+  
+ import sys
  import torch
- from typing import Dict
  from task import input_t, output_t
-  
  from aiter import ActivationType, QuantType
  from aiter.fused_moe import fused_moe
   
+ _warmed = False
   
+ def _warmup():
+     global _warmed
+     if _warmed:
+         return
+     _warmed = True
+  
+     configs = [
+         (2, 256, 1, 7168, 256, 8),
+         (2, 32, 1, 7168, 512, 8),
+         (2, 32, 1, 7168, 2048, 8),
+     ]
+  
+     for bs, n_routed, n_shared, d_hidden, d_expert, n_experts_per_token in configs:
+         E = n_routed + n_shared
+         total_topk = n_experts_per_token + n_shared
+         d_hidden_pad = ((d_hidden + 255) // 256) * 256
+         d_expert_pad = ((d_expert + 255) // 256) * 256
+         hidden_pad = d_hidden_pad - d_hidden
+         intermediate_pad = d_expert_pad - d_expert
+  
+         h = torch.randn(bs, d_hidden, dtype=torch.bfloat16, device="cuda")
+         w1 = torch.empty(E, 2 * d_expert_pad, d_hidden_pad // 2,
+                          dtype=torch.float4_e2m1fn_x2, device="cuda")
+         w2 = torch.empty(E, d_hidden_pad, d_expert_pad // 2,
+                          dtype=torch.float4_e2m1fn_x2, device="cuda")
+         w1_s = torch.empty(E, 2 * d_expert_pad, d_hidden_pad // 32,
+                            dtype=torch.float8_e8m0fnu, device="cuda")
+         w2_s = torch.empty(E, d_hidden_pad, d_expert_pad // 32,
+                            dtype=torch.float8_e8m0fnu, device="cuda")
+         topk_w = torch.ones(bs, total_topk, dtype=torch.float32, device="cuda")
+         topk_i = torch.zeros(bs, total_topk, dtype=torch.int32, device="cuda")
+         for t in range(bs):
+             for k in range(n_experts_per_token):
+                 topk_i[t, k] = k % n_routed
+             for k in range(n_shared):
+                 topk_i[t, n_experts_per_token + k] = n_routed + k
+  
+         try:
+             fused_moe(
+                 h, w1, w2, topk_w, topk_i,
+                 activation=ActivationType.Silu,
+                 quant_type=QuantType.per_1x32,
+                 w1_scale=w1_s, w2_scale=w2_s,
+                 hidden_pad=hidden_pad,
+                 intermediate_pad=intermediate_pad,
+             )
+             torch.cuda.synchronize()
+             print(f"  Warmup OK: E={E} d_h={d_hidden} d_e={d_expert}", file=sys.stderr)
+         except Exception as e:
+             print(f"  Warmup FAIL: E={E} d_h={d_hidden} d_e={d_expert}: {e}", file=sys.stderr)
+  
+     print("Warmup complete (v29 OPUS sorting)", file=sys.stderr)
+  
+ _warmup()
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      (
          hidden_states,
  ⋯ 10 unchanged lines
          config,
      ) = data
   
-     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
-     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
-  
-     output = fused_moe(
+     return fused_moe(
          hidden_states,
          gate_up_weight_shuffled,
          down_weight_shuffled,
          topk_weights,
          topk_ids,
-         expert_mask=None,
          activation=ActivationType.Silu,
          quant_type=QuantType.per_1x32,
-         doweight_stage1=False,
          w1_scale=gate_up_weight_scale_shuffled,
          w2_scale=down_weight_scale_shuffled,
-         a1_scale=None,
-         a2_scale=None,
-         hidden_pad=hidden_pad,
-         intermediate_pad=intermediate_pad,
+         hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
+         intermediate_pad=config["d_expert_pad"] - config["d_expert"],
      )
-  
-     return output