submission 692250

+ import functools
+ import importlib
  import os
+  
  import torch
+  
  from task import input_t, output_t
   
- os.environ["VLLM_MOE_WTYPE"] = "fp8"
- os.environ["VLLM_QUANT_OVERRIDE"] = "0"
+ os.environ["VLLM_MOE_CHUNK_SIZE"] = "512"
   
  from aiter import ActivationType, QuantType
  from aiter.fused_moe import fused_moe
   
   
- @torch.inference_mode()
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Submission template for DeepSeek-R1 MXFP4 MoE kernel.
+ _PATCH_DONE = False
   
-     Input data tuple:
-         hidden_states:                [M, d_hidden]                           bf16
-         gate_up_weight:               [E, 2*d_expert_pad, d_hidden_pad//2]    fp4x2  (raw)
-         down_weight:                  [E, d_hidden_pad, d_expert_pad//2]      fp4x2  (raw)
-         gate_up_weight_scale:         [E, 2*d_expert_pad, scale_K]            e8m0   (raw)
-         down_weight_scale:            [E, d_hidden_pad, scale_K]              e8m0   (raw)
-         gate_up_weight_shuffled:      [E, 2*d_expert_pad, d_hidden_pad//2]    fp4x2  (shuffled)
-         down_weight_shuffled:         [E, d_hidden_pad, d_expert_pad//2]      fp4x2  (shuffled)
-         gate_up_weight_scale_shuffled:[padded, flat]                          e8m0   (shuffled)
-         down_weight_scale_shuffled:   [padded, flat]                          e8m0   (shuffled)
-         topk_weights:                 [M, total_top_k]                        float32
-         topk_ids:                     [M, total_top_k]                        int32
-         config:                       dict
   
-     Returns:
-         output: [M, d_hidden] bf16
-     """
+ def _install_patch() -> None:
+     global _PATCH_DONE
+     if _PATCH_DONE:
+         return
+     _PATCH_DONE = True
+  
+     try:
+         fused_moe_module = importlib.import_module("aiter.fused_moe")
+         original = getattr(fused_moe_module, "get_2stage_cfgs", None)
+         if original is None:
+             return
+  
+         def patched_get_2stage_cfgs(*args, **kwargs):
+             meta = original(*args, **kwargs)
+  
+             token = kwargs.get("token", args[0] if len(args) > 0 else None)
+             model_dim = kwargs.get("model_dim", args[1] if len(args) > 1 else None)
+             inter_dim = kwargs.get("inter_dim", args[2] if len(args) > 2 else None)
+             expert = kwargs.get("expert", args[3] if len(args) > 3 else None)
+             topk = kwargs.get("topk", args[4] if len(args) > 4 else None)
+  
+             if token == 512 and model_dim == 7168 and inter_dim == 512 and expert == 33 and topk == 9:
+                 try:
+                     meta.block_m = 128
+                 except Exception:
+                     pass
+                 for attr in ("stage1", "stage2"):
+                     fn = getattr(meta, attr, None)
+                     if isinstance(fn, functools.partial):
+                         keywords = dict(fn.keywords or {})
+                         keywords["block_m"] = 128
+                         setattr(meta, attr, functools.partial(fn.func, *(fn.args or ()), **keywords))
+             return meta
+  
+         fused_moe_module.get_2stage_cfgs = patched_get_2stage_cfgs
+  
+         fused_moe_globals = getattr(fused_moe, "__globals__", None)
+         if isinstance(fused_moe_globals, dict) and fused_moe_globals.get("get_2stage_cfgs") is original:
+             fused_moe_globals["get_2stage_cfgs"] = patched_get_2stage_cfgs
+     except Exception:
+         pass
+  
+  
+ @torch.inference_mode()
+ def custom_kernel(data: input_t) -> output_t:
      (
          hidden_states,
-         gate_up_weight,
-         down_weight,
-         gate_up_weight_scale,
-         down_weight_scale,
+         _gate_up_weight,
+         _down_weight,
+         _gate_up_weight_scale,
+         _down_weight_scale,
          gate_up_weight_shuffled,
          down_weight_shuffled,
          gate_up_weight_scale_shuffled,
  ⋯ 3 unchanged lines
          config,
      ) = data
   
-     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
-     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
+     _install_patch()
   
-     output = fused_moe(
+     hidden_pad = int(config["d_hidden_pad"]) - int(config["d_hidden"])
+     intermediate_pad = int(config["d_expert_pad"]) - int(config["d_expert"])
+  
+     return fused_moe(
          hidden_states,
          gate_up_weight_shuffled,
          down_weight_shuffled,
  ⋯ 10 unchanged lines
          hidden_pad=hidden_pad,
          intermediate_pad=intermediate_pad,
      )
-  
-     return output