test new partition scheme

[yiakwy-xpu-ml-framework-team]

Description

• eliminate unnecessary dtype transform (aka, skip with template dtype checking if is_same<T, TCompute>()::value is true)
• row_id_map : the previous version map "source" row id to "dest " row id, which creates uncoalesce access of global store, since "source" id is randomly located:

// "source_topK_id * num_rows + source_token_id" is actually a random number for each thread
row_id_map[source_topK_id * num_rows + source_token_id] = idx;

instead,

    remapped_key_slot[tid] = source_topK_id * num_rows + source_token_id;
    
    cg::sync(cta);

    // Store : store the row to
    row_id_map[out_idx] = remapped_key_slot[tid];

is more friendly in continuously accessing global memory

• add row_id_map non-inverted version in permute function (demo)