fix: router A operand major mode MN (not K) — fixes CuTeDSL local_tile coord error

2026-06-01 00:54:19 +00:00
parent b7a8c44d26
commit c339fe7ad9
1 changed files with 3 additions and 1 deletions
--- a/dsv4/kernels/router/dense_router_decode_kernel.py
+++ b/dsv4/kernels/router/dense_router_decode_kernel.py
@@ -101,7 +101,9 @@ class DenseRouterDecodeKernel:
        self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake)

    def run(self, X, W_gate, e_bias, out_w, out_ids, M, E, K, scaling, top_k, stream=None):
-        self.a_major_mode = OperandMajorMode.K
+        # A operand (hidden_states) is (M, K) — MN-major layout
+        # B operand (W_gate) is (K, E) — K-major layout
+        self.a_major_mode = OperandMajorMode.MN
        self.b_major_mode = OperandMajorMode.K

        if stream is None: