sgl-project/sglang · error · ValueError

f"recurrent_kda state pool breaks the compiled stride contra

Error message

f"recurrent_kda state pool breaks the compiled stride contract: slot stride {ssm_states.stride(0)} elements must be a multiple of 16 and the base byte offset {base_bytes} a multiple of 32 (sym_int64(divisibility=16) / assumed_align=32)"

What it means

The FlashInfer recurrent_kda kernel is compiled with sym_int64(divisibility=16) and assumed_align=32 assumptions: the slot stride (elements) must be divisible by 16 and the pool's base byte offset by 32. Violating this trips an alignment assert in the compiled kernel, so the Python side pre-checks and raises.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py:118

        """
        key = id(ssm_states)
        if key in self._state_contract_ok:
            return
        if ssm_states.dim() != 4:
            raise ValueError(
                f"recurrent_kda needs a [N, HV, V, K] state pool; got "
                f"shape {tuple(ssm_states.shape)}"
            )
        _, hv, v, k = ssm_states.shape
        if ssm_states.stride()[1:] != (v * k, k, 1):
            raise ValueError(
                "recurrent_kda state inner strides must be compact "
                f"(V*K, K, 1)=({v * k}, {k}, 1); got {ssm_states.stride()[1:]} "
                "(only the slot stride may be non-compact)"
            )
        base_bytes = ssm_states.storage_offset() * ssm_states.element_size()
        if ssm_states.stride(0) % 16 != 0 or base_bytes % 32 != 0:
            raise ValueError(
                "recurrent_kda state pool breaks the compiled stride contract: "
                f"slot stride {ssm_states.stride(0)} elements must be a multiple "
                f"of 16 and the base byte offset {base_bytes} a multiple of 32 "
                "(sym_int64(divisibility=16) / assumed_align=32)"
            )
        self._state_contract_ok.add(key)

    # ---- gate / beta normalization (shared by decode + verify) ----

    def _prep_gate_params(self, A_log: torch.Tensor, dt_bias: torch.Tensor):
        # A_log: [1, 1, H, 1] -> [H] fp32; dt_bias: [H*K] (1D) -> fp32. Cached per
        # layer (constant weights) so this is a dict lookup on the hot path.
        key = (id(A_log), id(dt_bias))
        cached = self._gate_cache.get(key)
        if cached is not None:
            return cached
        A_log_fi = A_log.reshape(-1).float().contiguous()
        dt_bias_fi = (

View on GitHub (pinned to 0132848349)

Solutions

  1. Pad the slot stride (dim 0 stride) up to a multiple of 16 elements
  2. Ensure the pool tensor starts at a 32-byte-aligned storage offset (allocate fresh, avoid unaligned slices)
  3. Use the standard SGLang MambaCache allocation which already satisfies the contract

Example fix

# before
slot_stride = hv * v * k  # e.g. 8*64*63 not divisible by 16
# after
slot_stride = ((hv * v * k + 15) // 16) * 16
pool = torch.empty(num_slots * slot_stride, dtype=torch.float32).view(num_slots, slot_stride)[:, :hv*v*k].view(num_slots, hv, v, k).contiguous()
Defensive patterns

Strategy: validation

Validate before calling

elem = ssm_states.element_size()
base_bytes = ssm_states.storage_offset() * elem
assert ssm_states.stride(0) % 16 == 0, 'slot stride must be 16-element aligned'
assert base_bytes % 32 == 0, 'base offset must be 32-byte aligned'

Type guard

def meets_kda_alignment(t: torch.Tensor) -> bool:
    return t.stride(0) % 16 == 0 and (t.storage_offset() * t.element_size()) % 32 == 0

Prevention

When it happens

Trigger: Allocating the state pool with a slot stride not a multiple of 16 elements (e.g. odd V*K product) or passing a sliced view whose storage_offset yields a non-32-byte-aligned base.

Common situations: Custom cache allocators that pad slots to arbitrary sizes; slicing a shared pool at an unaligned offset; models with head dims whose product V*K is not divisible by 16.

Related errors


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