{"record":{"id":"92b6adade425faee","repo":"sgl-project/sglang","slug":"o-partial-tensor-must-match-dtype-partial","errorCode":null,"errorMessage":"O partial tensor must match dtype_partial","messagePattern":"O partial tensor must match dtype_partial","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py","lineNumber":221,"sourceCode":"\n    @cute.jit\n    def __call__(\n        self,\n        mO_partial: cute.Tensor,\n        mLSE_partial: cute.Tensor,\n        mO: cute.Tensor,\n        mLSE: Optional[cute.Tensor] = None,\n        cu_seqlens: Optional[cute.Tensor] = None,\n        seqused: Optional[cute.Tensor] = None,\n        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,\n        varlen_batch_idx: Optional[cute.Tensor] = None,\n        semaphore_to_reset: Optional[cute.Tensor] = None,\n        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).\n        stream: cuda.CUstream = None,\n    ):\n        # Type checking\n        if const_expr(not (mO_partial.element_type == self.dtype_partial)):\n            raise TypeError(\"O partial tensor must match dtype_partial\")\n        if const_expr(not (mO.element_type == self.dtype)):\n            raise TypeError(\"O tensor must match dtype\")\n        if const_expr(mLSE_partial.element_type not in [Float32]):\n            raise TypeError(\"LSE partial tensor must be Float32\")\n        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):\n            raise TypeError(\"LSE tensor must be Float32\")\n\n        # Shape validation - input tensors are in user format, need to be converted to kernel format\n        if const_expr(len(mO_partial.shape) not in [4, 5]):\n            raise ValueError(\n                \"O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)\"\n            )\n        if const_expr(len(mLSE_partial.shape) not in [3, 4]):\n            raise ValueError(\n                \"LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)\"\n            )\n        if const_expr(len(mO.shape) not in [3, 4]):\n            raise ValueError(","sourceCodeStart":203,"sourceCodeEnd":239,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py#L203-L239","documentation":"The flash-attention SplitKV combine kernel requires the partial output tensor mO_partial to match the dtype the operation was constructed with (dtype_partial, normally Float32, since split-KV partials accumulate in fp32). This is a static const_expr type check in __call__.","triggerScenarios":"Calling the combine operation with an mO_partial tensor whose element type differs from the op's dtype_partial configuration, e.g. fp16 partials fed to an op built for fp32 partials.","commonSituations":"Producing partials with a different kernel/configuration than the combine op expects; changing the split path to fp16 partials without reconfiguring the combine op's dtype_partial; buffer reuse across dtypes.","solutions":["Allocate mO_partial as float32 to match the default dtype_partial","Or construct the combine operation with dtype_partial matching your actual partial tensor's element type","Keep partial allocation and combine-op construction driven by one shared dtype variable"],"exampleFix":"// before\nO_partial = torch.empty(shape, dtype=torch.float16, device='cuda')\ncombine(O_partial, O)\n// after\nO_partial = torch.empty(shape, dtype=torch.float32, device='cuda')\ncombine(O_partial, O)","handlingStrategy":"validation","validationCode":"assert O_partial.dtype == torch.float32, 'combine expects fp32 partials'\nassert combine_op.dtype_partial == O_partial.dtype  # if op exposes dtype_partial","typeGuard":"def partial_dtype_ok(O_partial, dtype_partial=torch.float32) -> bool:\n    return O_partial.dtype == dtype_partial","tryCatchPattern":null,"preventionTips":["Derive both the partial allocation and combine-op config from one PARTIAL_DTYPE constant","Assert dtype compatibility in tests covering the split+combine round trip"],"tags":["cuda","dtype","flash-attention","split-kv","combine"],"backgroundTag":null,"analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T11:17:15.048Z"}