perf(tilelang): use k-parallel M2RNN Path C for small K

Author: apstenku123

bench/tilelang_ports/m2rnn_path_c_vs_path_b_20step.json

@@ -0,0 +1,52 @@
+{
+  "schema_version": 1,
+  "scope": "local_only",
+  "kernel": "m2rnn_full_hybrid_r_train",
+  "hardware_label": "M4 Max local",
+  "workload": {
+    "profile": "hybrid-smoke",
+    "route": "m2rnn",
+    "model_route": "hybrid-r",
+    "batch_size": 2,
+    "seq_len": 64,
+    "dtype": "bfloat16",
+    "compile": true,
+    "warmup_steps": 2,
+    "measured_steps": 20,
+    "tokens_per_step": 128,
+    "data_contract": "synthetic_tokens"
+  },
+  "results": [
+    {
+      "path": "path_b",
+      "kernel_used": "metal_kernel_fwd_v1",
+      "tokens_per_second": 129382.75128422497,
+      "mean_step_ms": 0.9893127076793461,
+      "median_step_ms": 0.9652710286900401,
+      "compile_time_s": 0.1551910419948399
+    },
+    {
+      "path": "path_c",
+      "kernel_used": "path_c_tilelang_dsl_packed",
+      "scheduler": "serial_packed_before_k_parallel_threshold_fix",
+      "tokens_per_second": 25961.985865299168,
+      "mean_step_ms": 4.930285405134782,
+      "median_step_ms": 4.864290996920317,
+      "compile_time_s": 4.640729499980807
+    },
+    {
+      "path": "path_c",
+      "kernel_used": "path_c_tilelang_dsl_packed",
+      "scheduler": "k_parallel_k4",
+      "tokens_per_second": 126674.93,
+      "mean_step_ms": 1.0105,
+      "median_step_ms": 1.0372,
+      "compile_time_s": 3.9039
+    }
+  ],
+  "speed": {
+    "path_c_k_parallel_vs_path_b_tokens_per_second_ratio": 0.9790712347313265,
+    "path_c_k_parallel_mean_step_slowdown_percent": 2.141615218708076
+  },
+  "matched_run_guard": "Compare only rows with identical workload, software stack, hardware label, warmup, and measured step counts."
+}

cppmega_mlx/nn/_tilelang/m2rnn_path_c.py

@@ -98,8 +98,8 @@ def _path_c_inputs_eligible(
     "h_steps_scratch",
 )
 _PACKED_BWD_OUTPUT_IDX = (6, 7, 8, 9, 10)
-_PACKED_FWD_K_PARALLEL_MIN_K = 16
-_PACKED_BWD_K_PARALLEL_MIN_K = 16
+_PACKED_FWD_K_PARALLEL_MIN_K = 4
+_PACKED_BWD_K_PARALLEL_MIN_K = 4
 
 M2RNNFwdOwnerOutputs = tuple[mx.array, mx.array, mx.array]
 M2RNNBwdOwnerOutputs = tuple[

tests/test_tilelang_m2rnn_path_c.py

@@ -38,7 +38,7 @@ def _np(x: mx.array) -> np.ndarray:
     if x.dtype == mx.bfloat16:
         x = x.astype(mx.float32)
     mx.eval(x)
-    return np.asarray(x)
+    return np.array(x, copy=True)
 
 
 def _make_m2rnn_inputs(
@@ -110,6 +110,24 @@ def test_m2rnn_packed_path_c_large_k_uses_k_parallel_lowering() -> None:
     assert _msl_transform.metal_grid_for_lowering(bwd_lowering) == (32, 1, 1)
 
 
+def test_m2rnn_packed_path_c_small_k_uses_k_parallel_lowering() -> None:
+    _require_m2rnn_path_c()
+
+    _fwd_kernel, fwd_lowering = m2rnn_path_c._packed_fwd_kernel_for(
+        1, 4, 2, 4, 4, "float32"
+    )
+    assert fwd_lowering.grid == (2, 1, 1)
+    assert fwd_lowering.threadgroup == (4, 1, 1)
+    assert _msl_transform.metal_grid_for_lowering(fwd_lowering) == (8, 1, 1)
+
+    _bwd_kernel, bwd_lowering = m2rnn_path_c._packed_bwd_kernel_for(
+        1, 4, 2, 4, 4, "float32"
+    )
+    assert bwd_lowering.grid == (2, 1, 1)
+    assert bwd_lowering.threadgroup == (4, 1, 1)
+    assert _msl_transform.metal_grid_for_lowering(bwd_lowering) == (8, 1, 1)
+
+
 def test_m2rnn_apply_path_c_fails_closed_instead_of_path_b(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: