diff --git a/dsv4/kernels/router/dense_router_decode_kernel.py b/dsv4/kernels/router/dense_router_decode_kernel.py index 43e9d651..e537406e 100644 --- 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: