pd_addr = amdgpu_gmc_emit_flush_gpu_tlb(ring, vmid, pd_addr);
 
        /* wait for register write */
-       data0 = hub->ctx0_ptb_addr_lo32 + vmid * 2;
+       data0 = hub->ctx0_ptb_addr_lo32 + vmid * hub->ctx_addr_distance;
        data1 = lower_32_bits(pd_addr);
        mask = 0xffffffff;
        jpeg_v2_0_dec_ring_emit_reg_wait(ring, data0, data1, mask);
 
        pd_addr = amdgpu_gmc_emit_flush_gpu_tlb(ring, vmid, pd_addr);
 
        /* wait for register write */
-       data0 = hub->ctx0_ptb_addr_lo32 + vmid * 2;
+       data0 = hub->ctx0_ptb_addr_lo32 + vmid * hub->ctx_addr_distance;
        data1 = lower_32_bits(pd_addr);
        mask = 0xffffffff;
        vcn_v2_0_dec_ring_emit_reg_wait(ring, data0, data1, mask);
        pd_addr = amdgpu_gmc_emit_flush_gpu_tlb(ring, vmid, pd_addr);
 
        /* wait for reg writes */
-       vcn_v2_0_enc_ring_emit_reg_wait(ring, hub->ctx0_ptb_addr_lo32 + vmid * 2,
+       vcn_v2_0_enc_ring_emit_reg_wait(ring, hub->ctx0_ptb_addr_lo32 +
+                                       vmid * hub->ctx_addr_distance,
                                        lower_32_bits(pd_addr), 0xffffffff);
 }