Support MoE sigmoid scoring

[catswe]

Description

This enables running GLM models #1104. I referenced this commit from Kyuyeun 652318d

Tests

Updated tpu-inference/tests/kernels/fused_moe_v1_test.py

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[kyuyeunk]

looks good. but requires approval from @bythew3i on changes for fused_moe/v1/kernel.py.

[bvrockwell]

This is sweet @catswe @kyuyeunk - great job!

[catswe]

The newly written unit test failed since bf16 input causes a jax.nn.sigmoid Pallas Mosaic lowering error. I already filed a fix PR here jax-ml/jax#34674. One workaround that works properly and passes the test is this

def apply_scoring_fn(scoring_fn: str, x):
    match scoring_fn:
        case "softmax":
            return jax.nn.softmax(x, axis=-1)
        case "sigmoid":
            # TODO(catswe): switch to jax.nn.sigmoid once bf16 Mosaic lowering bug is fixed
            ones = jnp.ones_like(x)
            return ones / (ones + jnp.exp(-x))
        case _:
            raise NotImplementedError(
                f"Unsupported scoring function: {scoring_fn}")

Let me know what direction we want to move towards (e.g., move forward with the workaround, or wait for the bug fix to land and upgrade Jax version)

[kyuyeunk]

The newly written unit test failed since bf16 input causes jax.nn.sigmoid Pallas Mosaic lowering error. I already filed a fix here jax-ml/jax#34674. One workaround that works properly and passes the test is this

def apply_scoring_fn(scoring_fn: str, x):
    match scoring_fn:
        case "softmax":
            return jax.nn.softmax(x, axis=-1)
        case "sigmoid":
            # TODO(catswe): switch to jax.nn.sigmoid once bf16 Mosaic lowering bug is fixed
            ones = jnp.ones_like(x)
            return ones / (ones + jnp.exp(-x))
        case _:
            raise NotImplementedError(
                f"Unsupported scoring function: {scoring_fn}")

Let me know what direction we want to move towards

thanks for identifying this issue! I'll run some tests myself as well and get back to you if i find anything.

[catswe]

I pushed up a small change to manually implement sigmoid. jax.nn.sigmoid Pallas Mosaic lowering uses the same 1 / (1 + exp(-x)) formula, although there's a bug with bf16 input. I added a TODO to use jax.nn.sigmoid once the lowering bug is fixed and Jax version is updated

Edit: Apologies for the mass reviewers ping. Messed up my force push 😅

[kyuyeunk]

I pushed up a small change to manually implement sigmoid. jax.nn.sigmoid Pallas Mosaic lowering uses the same 1 / (1 + exp(-x)) formula, although there's a bug with bf16 input. I added a TODO to use jax.nn.sigmoid once the lowering bug is fixed and Jax version is updated

Edit: Apologies for the mass reviewers ping. Messed up my force push 😅

just ran a quick test with latest jax version and it still triggered error. so current approach seems good to me. current approach of manually implementing sigmoid lgtm. thank you for the contribution!