Вам понадобится тип, который позволяет искать тип по ключу, но я не думаю, что тегированные объединения допускают, к сожалению. Например, вы можете изменить свою функцию следующим образом:
type ActionsLookup = {
A: A,
B: B,
};
type Actions = $Values<ActionsLookup>;
function dispatch<K: $Keys<ActionsLookup>>(
type: K,
payload: $ElementType<$ElementType<ActionsLookup, K>, 'payload'>
) {}