sgl-project/sglang · error · NotImplementedError

FlashKDAKernel only supports prefill (extend)

Error message

FlashKDAKernel only supports prefill (extend)

What it means

FlashKDAKernel wraps a chunked prefill kernel only; it intentionally implements extend() and raises NotImplementedError from decode(). Decode must be routed to a different backend.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py:99

    Requires an SM90+ GPU with the ``flash_kda`` package.
    """

    def decode(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        a: torch.Tensor,
        b: torch.Tensor,
        *,
        A_log: torch.Tensor,
        dt_bias: torch.Tensor,
        ssm_states: torch.Tensor,
        cache_indices: torch.Tensor,
        query_start_loc: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        raise NotImplementedError("FlashKDAKernel only supports prefill (extend)")

    def extend(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        g: torch.Tensor,
        beta: torch.Tensor,
        *,
        ssm_states: torch.Tensor,
        cache_indices: torch.Tensor,
        query_start_loc: torch.Tensor,
        A_log: Optional[torch.Tensor] = None,
        dt_bias: Optional[torch.Tensor] = None,
        lower_bound: Optional[float] = None,
        extend_seq_lens_cpu: Optional[list] = None,
        is_spec_decode: bool = False,
        return_intermediate_states: bool = False,

View on GitHub (pinned to 0132848349)

Solutions

1. Route decode to a decode-capable kernel (flashinfer recurrent_kda, triton, or helion)
2. Keep flashkda as the prefill backend only
3. Check the KDA dispatcher config for a per-phase backend split

Example fix

# before
kernel = FlashKDAKernel(); out = kernel.decode(q,k,v,...)
# after
decode_kernel = TritonKDAKernel(); out = decode_kernel.decode(q,k,v,...)

if phase == 'decode':
    assert not isinstance(kernel, FlashKDAKernel), 'flashkda is prefill-only'

def kernel_supports_decode(kernel) -> bool:
    return type(kernel).decode.__code__ is not FlashKDAKernel.decode.__code__

When it happens

Trigger: Calling decode() on FlashKDAKernel — e.g. a dispatcher bug or manual use that routes the decode phase to the flashkda backend.

Common situations: Misconfigured backend selection that applies flashkda to all phases instead of prefill-only.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• NvidiaKDAKernel is prefill-only
• PtxKDAKernel is prefill-only
• CuteDSLKDAKernel does not support target_verify
• FlashInferKDAKernel has no prefill kernel; keep prefill on T
• The 'flashkda' KDA prefill backend requires the flash_kda mo

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