submission 567194

  """
- v5: Safe approach - only env vars, no custom function params.
- Shape branching with env var tuning only.
+ submission__v4: runtime hybrid + selective ksplit.
+  
+ Use ksplit=2 only on the small 33-expert / 512-intermediate shapes where it
+ materially helps; keep submission__v1 behavior everywhere else.
  """
  import os
+  
  from task import input_t, output_t
   
- # Set env vars BEFORE importing aiter (some are read at import time)
  os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
  os.environ["AITER_USE_NT"] = "0"
   
- import aiter
  from aiter import ActivationType, QuantType
- from aiter.fused_moe import fused_moe
+ import aiter.fused_moe as fused_moe_mod
   
   
- def _call_default(data):
+ _LAST_RUNTIME_STATE = None
+  
+  
+ def _set_runtime_state(use_opus: bool, use_nt: bool, ksplit: int | None = None) -> None:
+     global _LAST_RUNTIME_STATE
+  
+     state = (use_opus, use_nt, ksplit)
+     if state == _LAST_RUNTIME_STATE:
+         return
+  
+     fused_moe_mod._USE_OPUS_MOE_SORTING = use_opus
+     os.environ["AITER_USE_NT"] = "1" if use_nt else "0"
+  
+     if ksplit in (None, 0):
+         os.environ.pop("AITER_KSPLIT", None)
+     else:
+         os.environ["AITER_KSPLIT"] = str(ksplit)
+  
+     fused_moe_mod.use_nt.cache_clear()
+     fused_moe_mod.get_ksplit.cache_clear()
+     fused_moe_mod.get_2stage_cfgs.cache_clear()
+     _LAST_RUNTIME_STATE = state
+  
+  
+ def _policy(config: dict) -> tuple[bool, bool, int | None]:
+     total_experts = config["n_routed_experts"] + config["n_shared_experts"]
+  
+     if total_experts == 257:
+         return True, False, None
+  
+     if config["d_expert"] == 512 and config["bs"] <= 128:
+         return False, True, 2
+  
+     return False, True, None
+  
+  
+ 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_shuffled, down_weight_shuffled,
-         gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
-         topk_weights, topk_ids, config,
+         hidden_states,
+         gate_up_weight,
+         down_weight,
+         gate_up_weight_scale,
+         down_weight_scale,
+         gate_up_weight_shuffled,
+         down_weight_shuffled,
+         gate_up_weight_scale_shuffled,
+         down_weight_scale_shuffled,
+         topk_weights,
+         topk_ids,
+         config,
      ) = data
   
+     use_opus, use_nt, ksplit = _policy(config)
+     _set_runtime_state(use_opus=use_opus, use_nt=use_nt, ksplit=ksplit)
+  
      hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
      intermediate_pad = config["d_expert_pad"] - config["d_expert"]
   
-     return fused_moe(
+     return fused_moe_mod.fused_moe(
          hidden_states,
          gate_up_weight_shuffled,
          down_weight_shuffled,
  ⋯ 10 unchanged lines
          hidden_pad=hidden_pad,
          intermediate_pad=intermediate_pad,
      )
-  
-  
- def custom_kernel(data: input_t) -> output_t:
-     return _call_default(data)