sgl-project/sglang · error · NotImplementedError

PtxKDAKernel is prefill-only

Error message

PtxKDAKernel is prefill-only

What it means

PtxKDAKernel is a GB300/sm_103a PTX chunked prefill kernel; it has no decode kernel and raises NotImplementedError from decode() to keep decode on another backend.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_ptx.py:79

        self._param_flat = {}
        self._unsupported_logged = False
        # (bucket, H, K, V, device) -> staging dict for ragged token counts.
        self._staging = {}

    def _ensure_loaded(self):
        if self._fwd is None:
            from sglang.kernels.ops.attention.linear.kda_ptx_prefill import (
                chunk_kda_fwd,
                load_ext,
            )

            logger.info("Building the PTX KDA prefill extension (first use, ~1-2 min)")
            load_ext()
            self._fwd = chunk_kda_fwd
            logger.info("Using PTX KDA chunked prefill (GB300 / sm_103a)")

    def decode(self, *args, **kwargs):
        raise NotImplementedError("PtxKDAKernel is prefill-only")

    def target_verify(self, *args, **kwargs):
        raise NotImplementedError("PtxKDAKernel does not support target_verify")

    def _flat_param(self, t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
        if t is None:
            return None
        key = (t.data_ptr(), t.dtype, tuple(t.shape))
        flat = self._param_flat.get(key)
        if flat is None:
            flat = t.detach().reshape(-1).float().contiguous()
            self._param_flat[key] = flat
        return flat

    def _get_staging(self, bucket, num_heads, head_k_dim, head_v_dim, dev):
        key = (bucket, num_heads, head_k_dim, head_v_dim, dev.index)
        st = self._staging.get(key)
        if st is None:

View on GitHub (pinned to 0132848349)

Solutions

  1. Route decode to a decode-capable KDA kernel (triton/flashinfer/helion)
  2. Use PTX kernel for prefill only
  3. Verify per-phase backend routing config

Example fix

# before
kernel = PtxKDAKernel(); out = kernel.decode(...)
# after
out = triton_decode_kernel.decode(...)
Defensive patterns

Strategy: type-guard

Validate before calling

if phase == 'decode':
    assert not isinstance(kernel, PtxKDAKernel), 'PTX KDA kernel is prefill-only'

Type guard

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

Prevention

When it happens

Trigger: Dispatching decode to PtxKDAKernel — direct call or misrouted backend configuration.

Common situations: Running on GB300 with the PTX backend enabled for all phases; dispatcher change that removed the decode fallback.

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


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