[BENCH] [NFC] Update routing.py (#7503)

Minor refactor

Author: apgoucher

python/triton_kernels/triton_kernels/routing.py

@@ -276,6 +276,19 @@ def compute_expt_data(expt_hist, n_expts_tot, n_gates):
 # --------------------------
 
 
+def routing_from_bitmatrix(bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act):
+    hist, topk_indx, gate_indx, gate_scal, token_offs_raw, token_offs_pad, block_pid_map = sort_tokens(
+        expt_scal, expt_indx, n_expts_tot, bitmatrix)
+    token_offs_pad = _unpack_into_dict(token_offs_pad)
+    block_pid_map = _unpack_into_dict(block_pid_map)
+    expt_data = ExptData(hist, token_offs_raw, token_offs_pad, block_pid_map)
+
+    # pack the matmul data structure
+    gather_indx = GatherIndx(src_indx=topk_indx, dst_indx=gate_indx)
+    scatter_indx = ScatterIndx(src_indx=gate_indx, dst_indx=topk_indx)
+    return RoutingData(gate_scal, hist, n_expts_tot, n_expts_act, expt_data), gather_indx, scatter_indx
+
+
 def routing(logits, n_expts_act, sm_first=False, expt_indx=None, simulated_ep=1, n_rows=None):
     from .topk import topk
     if sm_first:
@@ -286,16 +299,8 @@ def routing(logits, n_expts_act, sm_first=False, expt_indx=None, simulated_ep=1,
     # mutate bitmatrix
     if simulated_ep > 1:
         expt_scal, expt_indx, bitmatrix = prune_routing(expt_scal, expt_indx, bitmatrix, logits.shape[-1], simulated_ep)
-    hist, topk_indx, gate_indx, gate_scal, token_offs_raw, token_offs_pad, block_pid_map = sort_tokens(
-        expt_scal, expt_indx, n_expts_tot, bitmatrix)
-    token_offs_pad = _unpack_into_dict(token_offs_pad)
-    block_pid_map = _unpack_into_dict(block_pid_map)
-    expt_data = ExptData(hist, token_offs_raw, token_offs_pad, block_pid_map)
 
-    # pack the matmul data structure
-    gather_indx = GatherIndx(src_indx=topk_indx, dst_indx=gate_indx)
-    scatter_indx = ScatterIndx(src_indx=gate_indx, dst_indx=topk_indx)
-    return RoutingData(gate_scal, hist, n_expts_tot, n_expts_act, expt_data), gather_indx, scatter_indx
+    return routing_from_bitmatrix(bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act)
 
 
 # --------------------------