sgl-project/sglang · error · ValueError

Helion KDA decode requires power-of-two key and value head d

Error message

Helion KDA decode requires power-of-two key and value head dimensions (got K={K}, V={V}).

What it means

The Helion KDA decode kernel is specialized for power-of-two head dimensions: it uses the fast power-of-two code path (tower-of-power decays) for key dim K and value dim V taken from initial_state.shape[-3:]. validate_packed_decode_inputs rejects any non-power-of-two K or V so the Triton/Helion kernel never receives a shape it cannot compile a cache entry for.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_decode.py:288

            "Mismatched batch sizes: "
            f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, "
            f"b.shape[0]={b.shape[0]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] "
            f"(got {tuple(ssm_state_indices.shape)}; expected ({B},))."
        )

    if initial_state.ndim != 4:
        raise ValueError(
            f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
        )
    if initial_state.stride(-1) != 1:
        raise ValueError("`initial_state` must be contiguous in the last dim.")
    HV, V, K = initial_state.shape[-3:]
    if not _is_power_of_two(K) or not _is_power_of_two(V):
        raise ValueError(
            "Helion KDA decode requires power-of-two key and value head "
            f"dimensions (got K={K}, V={V})."
        )
    if a.shape[1] != HV * K:
        raise ValueError(
            f"`a` must have shape [B, HV*K] with HV={HV}, K={K} "
            f"(got a.shape={tuple(a.shape)})."
        )
    if b.shape[1] != HV:
        raise ValueError(
            f"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)})."
        )
    if A_log.numel() != HV:
        raise ValueError(f"`A_log` must have {HV} elements (got {A_log.numel()}).")
    if dt_bias.numel() != HV * K:
        raise ValueError(
            f"`dt_bias` must have {HV * K} elements (got {dt_bias.numel()})."
        )

View on GitHub (pinned to 0132848349)

Solutions

Fall back to the reference/Triton KDA decode path when K or V is not a power of two
Change the model config so key_dim and value_dim are powers of two (64/128/256) if you control the checkpoint
Gate model registration: assert _is_power_of_two(head_k_dim) and _is_power_of_two(head_v_dim) before selecting the Helion backend

Example fix

// before
use_helion = True  # model has head_v_dim=96
// after
use_helion = _is_power_of_two(cfg.head_k_dim) and _is_power_of_two(cfg.head_v_dim)
if not use_helion:
    backend = "triton"

Defensive patterns

Strategy: fallback

Validate before calling

def helion_ok(cfg):
    return cfg.head_k_dim & (cfg.head_k_dim - 1) == 0 and cfg.head_v_dim & (cfg.head_v_dim - 1) == 0
backend = "helion" if helion_ok(cfg) else "triton"

Prevention

Check power-of-two head dims at model registration time
Keep a reference decode path for non-power-of-two configs

When it happens

Trigger: Running a KDA/GDN model whose value_dim or key_dim is e.g. 96, 192, or any non-power-of-two, with the Helion decode path enabled; passing a hand-built initial_state with odd V/K in a unit test.

Common situations: Loading an experimental checkpoint with non-standard head dims; a config where head_v_dim differs from head_k_dim and neither is a power of two; upgrading a model config without re-checking the Helion kernel constraints.

Related errors

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/3f96fb67332cbea0. Report an issue: GitHub.