sgl-project/sglang · error · NotImplementedError

kda_prefill is the inference forward path: cp_context, and d

Error message

kda_prefill is the inference forward path: cp_context, and disable_recompute are training-side knobs it does not implement

What it means

kda_prefill (kda_ptx_prefill) is the inference-only forward. The public entry asserts that training-side context-parallelism and recompute knobs (cp_context, disable_recompute) are not passed. Passing either one raises NotImplementedError immediately.

Source

Thrown at python/sglang/kernels/ops/attention/linear/kda_ptx_prefill/__init__.py:114

      beta [B,T,H] bf16 (fp32 also accepted; widened to fp32 in the ext).
      use_qk_l2norm_in_kernel=True accepts raw q/k and applies FLA-compatible
        L2 normalization (eps=1e-6, bf16 rounding) in the CUDA tile loads.
      use_beta_sigmoid_in_kernel=True accepts beta logits and fuses sigmoid.
      cu_seqlens: host values are needed for the kernel's per-sequence piece
        table -- pass cu_seqlens_cpu to avoid the D2H sync; chunk_indices is
        accepted and ignored (the kernel derives its own piece table).
      initial_state [N,H,128,128] fp32 or None (zeros).
      return_intermediate_states=True returns dense fp32 chunk-boundary states
        [1, NT, H, 128, 128] at tuple index 10.

    Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state
    [N,H,128,128] fp32 or None, then Nones, ..., h, initial_state).
    """
    assert (
        chunk_size == CHUNK
    ), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}"
    if cp_context is not None or disable_recompute:
        raise NotImplementedError(
            "kda_prefill is the inference forward path: cp_context, "
            "and disable_recompute are training-side knobs it does not implement"
        )
    if allow_neg_eigval and use_beta_sigmoid_in_kernel:
        raise NotImplementedError(
            "allow_neg_eigval=True requires 2*sigmoid(beta), which is not "
            "implemented by the fused beta path; pass pre-activated beta with "
            "use_beta_sigmoid_in_kernel=False"
        )
    if state_v_first and initial_state is not None:
        # [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call
        initial_state = initial_state.transpose(-1, -2).contiguous()
    assert (
        q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K
    ), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}"
    B, T, H, _ = q.shape

    cu_cpu = None

View on GitHub (pinned to 0132848349)

Solutions

Pass cp_context=None and disable_recompute=False (or omit them) when calling the inference prefill
Use the dedicated training-side kernel entry (fla/training path) for cp / recompute semantics

Example fix

// before
kda_prefill(q, k, v, ..., cp_context=cp_ctx, disable_recompute=True)
// after
kda_prefill(q, k, v, ..., cp_context=None, disable_recompute=False)

Defensive patterns

Strategy: validation

Validate before calling

assert cp_context is None and not disable_recompute, \
    'kda_prefill is inference-only: drop training knobs'

Prevention

Keep separate call sites for training and inference KDA paths
Omit training kwargs entirely on the inference path

When it happens

Trigger: Calling the kda_prefill forward with cp_context not None or disable_recompute=True (values copied from a training code path / fla-style chunk_kda API).

Common situations: Porting a training implementation (e.g. FLA chunk_kda with context parallelism) into SGLang serving and reusing the same kwargs; a model definition shared between train and inference passing all knobs unconditionally.

Related errors

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