diff --git a/vmvx_ukernel/mmt4d_tile.c b/vmvx_ukernel/mmt4d_tile.c
index 05730fb..18ad137 100644
--- a/vmvx_ukernel/mmt4d_tile.c
+++ b/vmvx_ukernel/mmt4d_tile.c
@@ -52,7 +52,7 @@
   iree_uk_int16_t N0 = params->N0;
   iree_uk_int16_t K0 = params->K0;
   // Initialize the accumulator tile.
-  if (!(flags & IREE_UK_FLAG_ACCUMULATE)) {
+  if (!(flags & IREE_UK_FLAG_MMT4D_ACCUMULATE)) {
     memset(out_tile, 0, M0 * N0 * sizeof(iree_uk_int32_t));
   }
   // Accumulation loop.
@@ -81,7 +81,7 @@
   iree_uk_int16_t K0 = params->K0;
   // Initialize the local accumulator tile.
   float acc[iree_uk_mmt4d_tile_generic_max_bytes / sizeof(*out_tile)];
-  if (flags & IREE_UK_FLAG_ACCUMULATE) {
+  if (flags & IREE_UK_FLAG_MMT4D_ACCUMULATE) {
     for (int i = 0; i < M0 * N0; ++i) acc[i] = out_tile[i];
   } else {
     for (int i = 0; i < M0 * N0; ++i) acc[i] = 0;
@@ -107,7 +107,7 @@
 iree_uk_mmt4d_tile_func_t iree_uk_mmt4d_select_tile_func(
     const iree_uk_mmt4d_params_t* params) {
   // TODO(lundong): to be replaced with Kelvin
-  switch (params->type) {
+  switch (iree_uk_mmt4d_type(params->flags)) {
     case iree_uk_mmt4d_type_f32f32f32:
       return iree_uk_mmt4d_tile_f32f32f32_generic;
     case iree_uk_mmt4d_type_i8i8i32: