{"record":{"id":"62ba2ebbe3c16be5","repo":"sgl-project/sglang","slug":"mixed-qkv-must-be-2d-got-ndim-mixed-qkv-ndim","errorCode":null,"errorMessage":"`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).","messagePattern":"`mixed_qkv` must be 2D \\(got ndim=(.+?)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py","lineNumber":478,"sourceCode":"    One kernel for both gate granularities, selected by ``is_kda``:\n      * ``is_kda=False`` (GDN): per-head SCALAR gate.  ``a``=[B, HV],\n        ``dt_bias``=[HV], ``g_cache``=[num_slots, HV, L].\n      * ``is_kda=True`` (KDA): per-K-channel gate.  ``a``=[B, HV, K],\n        ``dt_bias``=[HV, K], ``g_cache``=[num_slots, HV, L, K].\n    ``A_log`` is [HV] (per-head scalar) for both.\n\n    Same call surface as the packed decode plus the three ring caches\n    (``d_cache`` / ``k_cache`` / ``g_cache``) and the per-decode-row\n    ``write_pos`` cursor.  ``initial_state`` is both the checkpoint read (h0)\n    and the (flush-only) checkpoint write (ht), in place.\n\n    Allocates nothing persistent: the caller owns the ring tensors and is\n    responsible for advancing / resetting ``write_pos`` (e.g. ``(write_pos+1) %\n    L`` after each step).  This is a STANDALONE kernel; the memory-pool / cache\n    integration is a later phase.\n    \"\"\"\n    if mixed_qkv.ndim != 2:\n        raise ValueError(f\"`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).\")\n    if mixed_qkv.stride(-1) != 1:\n        raise ValueError(\"`mixed_qkv` must be contiguous in the last dim.\")\n    if b.ndim != 2:\n        raise ValueError(f\"`b` must be 2D (got b.ndim={b.ndim}).\")\n    if A_log.ndim != 1:\n        raise ValueError(\"`A_log` must be a 1D tensor.\")\n    if initial_state.ndim != 4:\n        raise ValueError(f\"`initial_state` must be 4D (got ndim={initial_state.ndim}).\")\n    if not out.is_contiguous():\n        raise ValueError(\"`out` must be contiguous.\")\n    if write_pos.ndim != 1 or write_pos.dtype != torch.int32:\n        raise ValueError(\"`write_pos` must be a 1D int32 tensor.\")\n    if force_flush is not None and (\n        force_flush.ndim != 1 or force_flush.dtype != torch.int32\n    ):\n        raise ValueError(\"`force_flush` must be a 1D int32 tensor or None.\")\n\n    B = mixed_qkv.shape[0]","sourceCodeStart":460,"sourceCodeEnd":496,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py#L460-L496","documentation":"fused_recurrent_linear_replayssm_decode validates its inputs before launching: mixed_qkv must be a 2D [num_tokens, qkv_dim] tensor (one row per decode token). Passing 3D/4D tensors (e.g. [B, T, D] or [B, 1, T, D]) trips this check.","triggerScenarios":"Feeding an un-squeezed decode projection of shape [B, 1, D] or a prefill-shaped [B, T, D] tensor into the standalone replaySSM decode kernel.","commonSituations":"Adapting prefill batch code to the decode path; forgetting to reshape after a fused QKV projection in a new model integration; CUDA-graph capture tests passing batched dims.","solutions":["Reshape mixed_qkv to 2D: mixed_qkv.reshape(-1, qkv_dim) before calling","Keep all per-token tensors 2D [num_tokens, dim] on this path","Check b is 2D and initial_state is 4D to satisfy the sibling checks"],"exampleFix":"// before\nout = fused_recurrent_linear_replayssm_decode(mixed_qkv[B,1,D], ...)\n// after\nout = fused_recurrent_linear_replayssm_decode(mixed_qkv.reshape(-1, D), ...)","handlingStrategy":"validation","validationCode":"mixed_qkv = mixed_qkv.reshape(-1, mixed_qkv.shape[-1])\nassert mixed_qkv.ndim == 2","typeGuard":"def is_2d_tokens(t: torch.Tensor) -> bool:\n    return t.ndim == 2","tryCatchPattern":null,"preventionTips":["Normalize all decode-path tensors to [num_tokens, dim] in one helper","Write contract tests asserting ndim for every arg"],"tags":["pytorch","tensor-shape","replayssm","decode","validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}