Vmap over random keys is significantly slower with RBG

[imoneoi]

Why is random generation significantly (~50x) slower when vmapped over random keys? See the code below for an example.

UPDATE: It seems that only rbg and unsafe_rbg PRNG is slow, the default PRNG have similar speed when vmapped.

import jax
import jax.numpy as jnp

import time


jax.config.update("jax_default_prng_impl", "rbg")


def test_random_slow(
    seed: int = 0,

    batch_size: int = 10240,
    ep_length: int = 1000,
    dims: int = 10
):
    def _step(carry, input):
        key, state = carry
        new_key, cur_key = jax.random.split(key)

        new_state = state + jax.random.uniform(cur_key, (dims,))

        return (new_key, new_state), None

    # Initial state
    key = jax.random.split(jax.random.PRNGKey(seed), batch_size)
    state = jnp.zeros((batch_size, dims))

    # Vmap & Scan
    (new_key, new_state), _ = jax.lax.scan(jax.vmap(_step, in_axes=(0, None)), (key, state), None, ep_length)
    return new_state


def test_random_fast(
    seed: int = 0,

    batch_size: int = 10240,
    ep_length: int = 1000,
    dims: int = 10
):
    def _step(carry, input):
        key, state = carry
        new_key, cur_key = jax.random.split(key)

        new_state = state + jax.random.uniform(cur_key, (batch_size, dims))

        return (new_key, new_state), None

    # Initial state
    key = jax.random.PRNGKey(seed)
    state = jnp.zeros((batch_size, dims))

    # Vmap & Scan
    (new_key, new_state), _ = jax.lax.scan(_step, (key, state), None, ep_length)
    return new_state


if __name__ == "__main__":
    start_time = time.time()
    jax.jit(test_random_fast)().block_until_ready()

    print(f"{time.time() - start_time:.2f}")

    start_time = time.time()
    jax.jit(test_random_slow)().block_until_ready()

    print(f"{time.time() - start_time:.2f}")

Test result on a RTX 3090 GPU (rbg):

1.27
61.66

Test result on a RTX 3090 GPU (default):

0.97
0.22