xai-org/x-algorithm · error · ValueError

the fused rowwise Adagrad update needs a row shard that divi

Error message

the fused rowwise Adagrad update needs a row shard that divides a warp, got shard_width={context.shard_width} (emb_width={context.emb_width})

What it means

The fused rowwise Adagrad CUDA kernel processes one embedding row across a warp (32 threads), so each row shard must evenly divide 32. This error is raised when shard_width (derived from emb_width/sharding context) does not divide 32, making the warp-cooperative update impossible.

Source

Thrown at phoenix/xrex/optimizers/recsys/rowwise_adagrad.py:106

        return embeddings, None

    @property
    def decay_factor(self) -> float:
        return math.exp(-self._decay_rate) if self._decay_rate is not None else 1.0

    def gradient_update_start(
        self,
        context: async_emb.AsyncEmbContextHandle,
        update: AsyncEmbGradientUpdate,
        table: jax.Array,
        state: RecsysRowwiseAdagradState,
        gate: jax.Array,
    ) -> tuple[tuple[jax.Array, ...], jax.Array, RecsysRowwiseAdagradState]:
        if self._lazy_decay and (state.step is None or state.last_step is None):
            raise ValueError("fused lazy decay needs timestamped state (step/last_step)")

        if 32 % context.shard_width != 0:
            raise ValueError(
                f"the fused rowwise Adagrad update needs a row shard that divides a warp, "
                f"got shard_width={context.shard_width} (emb_width={context.emb_width})"
            )

        from xrex.cuda.async_emb import async_emb

        metrics: dict[str, jax.Array] = {}

        if self._lazy_decay:

            @shard_map(
                mesh=context.mesh,
                in_specs=(
                    P(context.data_axis, None),
                    P(context.data_axis),
                    P(),
                    P(None, context.table_axis),
                    P(),

View on GitHub (pinned to 24c60942c5)

Solutions

Change emb_width/sharding so shard_width divides 32 (1, 2, 4, 8, 16, or 32) — e.g. pad emb_width to 64/128/256
Check context.shard_width vs context.emb_width to confirm how many shards the row is split into
If the width cannot change, use a non-fused optimizer path instead of the fused rowwise Adagrad update

Example fix

# before
emb = Embedding(num_embeddings=N, emb_width=96)  # shard_width=3 -> 32 % 3 != 0

# after
emb = Embedding(num_embeddings=N, emb_width=128)  # shard_width=4 -> ok

Defensive patterns

Strategy: validation

Validate before calling

assert 32 % shard_width == 0, f"shard_width={shard_width} must divide 32 (warp size)"

Type guard

def warp_compatible(emb_width: int, shards: int) -> bool:
    return 32 % (emb_width // shards) == 0

Prevention

Restrict emb_width to multiples of 32
Validate shard geometry in config parsing before building the optimizer

When it happens

Trigger: Calling gradient_update_start with an embedding width (or shard width) like 48, 96, or any value where 32 % shard_width != 0, e.g. emb_width producing shard_width=3 or 6.

Common situations: Switching an embedding table to a non-power-of-two width (e.g. 96 or 768/7=... odd widths), changing sharding so each shard is narrower, or porting a config from a different optimizer without the warp-divisibility constraint.

Related errors

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/1b038d0c04980f80. Report an issue: GitHub.