sgl-project/sglang · error · ValueError

Unexpected initial_state_source shape: {initial_state_source

Error message

Unexpected initial_state_source shape: {initial_state_source.shape}

What it means

The fused sigmoid-gating GDN delta-rule update accepts an initial recurrent-state tensor that is either a flat pool (dim 2, interpreted as (pool_size, HV, K, V) via numel arithmetic) or an explicit 4D (pool_size, HV, K, V) batch. Any other dimensionality raises ValueError('Unexpected initial_state_source shape: ...').

Source

Thrown at python/sglang/kernels/ops/attention/cutedsl_gdn.py:1409

    B_q, T_q, H, K = q.shape
    HV = v.shape[2]
    V = v.shape[3]
    N = initial_state_indices.shape[0]

    is_varlen_decode = B_q == 1 and T_q == N and N > 1
    if scale is None:
        scale = K**-0.5

    use_small_batch = N < SMALL_BATCH_THRESHOLD

    if initial_state_source.dim() == 1:
        pool_size = initial_state_source.numel() // (HV * K * V)
        h0_source = initial_state_source.view(pool_size, HV, K, V)
    elif initial_state_source.dim() == 4:
        pool_size = initial_state_source.shape[0]
        h0_source = initial_state_source
    else:
        raise ValueError(
            f"Unexpected initial_state_source shape: {initial_state_source.shape}"
        )

    if is_varlen_decode:
        if a.dim() == 3:
            a = a.squeeze(0)
        if b.dim() == 3:
            b = b.squeeze(0)
        o = q.new_empty(1, N, HV, V, dtype=torch.bfloat16)
    else:
        if a.dim() == 2:
            a = a.unsqueeze(1)
        if b.dim() == 2:
            b = b.unsqueeze(1)
        o = q.new_empty(N, 1, HV, V, dtype=torch.bfloat16)

    q, k, v = [t.contiguous() for t in (q, k, v)]

View on GitHub (pinned to 0132848349)

Solutions

  1. Reshape to 4D (pool_size, HV, K, V) before the call
  2. Or pass the flat pool tensor and let the kernel view it: tensor.view(numel // (HV*K*V), HV, K, V)
  3. Print initial_state_source.shape at the call site and compare with the HV, K, V arguments to find the mismatched dim

Example fix

# before
state = state.squeeze(0)          # dim 3 -> ValueError
update = cutedsl_fused_sigmoid_gating_delta_rule_update(..., initial_state_source=state)
# after
state = state.reshape(pool_size, HV, K, V)  # dim 4
update = cutedsl_fused_sigmoid_gating_delta_rule_update(..., initial_state_source=state)
Defensive patterns

Strategy: validation

Validate before calling

HV_K_V = HV * K * V
assert initial_state_source.dim() in (2, 4), initial_state_source.shape
assert initial_state_source.numel() % HV_K_V == 0, 'state numel must be a multiple of HV*K*V'

Type guard

def is_valid_initial_state(t) -> bool:
    return t.dim() == 4 or (t.dim() == 2 and t.numel() % (HV * K * V) == 0)

Try / catch

try:
    out = cutedsl_fused_sigmoid_gating_delta_rule_update(..., initial_state_source=state)
except ValueError as e:
    if 'initial_state_source' in str(e):
        state = state.reshape(-1, HV, K, V)
        out = cutedsl_fused_sigmoid_gating_delta_rule_update(..., initial_state_source=state)
    else:
        raise

Prevention

When it happens

Trigger: Calling cutedsl_fused_sigmoid_gating_delta_rule_update with an initial_state_source of dim 3 or dim >= 5 — e.g. passing (1, HV, K, V) with an unwanted leading batch dim of size 1 that is actually meant to be 5D, or a per-request (HV,K,V) tensor without the pool dimension.

Common situations: Wiring a custom GDN/Gated-DeltaNet model's recurrent state cache directly into the fused kernel; migration from another attention backend whose state layout is (batch, heads, d_k, d_v) but squeezed incorrectly; passing cu_seqlen varlen states without the pool dimension.

Related errors


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