xai-org/x-algorithm · error · TypeError

Only Float16 or BFloat16 is supported

Error message

Only Float16 or BFloat16 is supported

What it means

The flash backward postprocess kernel (__call__ in flash_bwd_postprocess.py) only accepts Float16 or BFloat16 for the Q output tensor mdQ. Other dtypes (FP32, FP8, TF32) are rejected at compile time via const_expr.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_postprocess.py:221

                major_mode_size=self.tile_hdim // wg_d_dQ,
            )
        else:
            self.sdQ_layout = sm100_utils_basic.make_smem_layout_epi(
                self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim), 1
            )

    @cute.jit
    def __call__(
        self,
        mdQaccum: cute.Tensor,
        mdQ: cute.Tensor,
        scale: cutlass.Float32,
        mCuSeqlensQ: Optional[cute.Tensor],
        mSeqUsedQ: Optional[cute.Tensor],
        stream: cuda.CUstream = None,
    ):
        if const_expr(mdQ.element_type not in [cutlass.Float16, cutlass.BFloat16]):
            raise TypeError("Only Float16 or BFloat16 is supported")
        if const_expr(mdQaccum is not None):
            if const_expr(mdQaccum.element_type not in [cutlass.Float32]):
                raise TypeError("dQaccum tensor must be Float32")

        mdQaccum, mdQ = [assume_tensor_aligned(t) for t in (mdQaccum, mdQ)]

        self.tiled_mma = self._get_tiled_mma()
        self._setup_attributes()

        smem_size = max(
            cute.size_in_bytes(cutlass.Float32, self.sdQaccum_layout),
            cute.size_in_bytes(self.dtype, self.sdQ_layout),
        )

        if const_expr(mCuSeqlensQ is not None):
            TileScheduler = SingleTileVarlenScheduler
            num_head = mdQ.shape[1]
            num_batch = mCuSeqlensQ.shape[0] - 1

View on GitHub (pinned to 24c60942c5)

Solutions

1. Cast mdQ to half or bfloat16 before the call: mdQ.to(torch.bfloat16)
2. Run under autocast/bf16 so downstream tensors are already the right dtype
3. Keep the fp32 copy separately if needed for loss scaling

Example fix

# before
mdQ = dq_accum_result.to(torch.float32)
post(mdQaccum, mdQ, scale, ...)

# after
mdQ = dq_accum_result.to(torch.bfloat16)   # or torch.float16
post(mdQaccum, mdQ, scale, ...)

assert mdQ.dtype in (torch.float16, torch.bfloat16), mdQ.dtype

def is_supported_q_dtype(t: torch.Tensor) -> bool:
    return t.dtype in (torch.float16, torch.bfloat16)

When it happens

Trigger: Calling the postprocess kernel with mdQ in torch.float32 or float64 — e.g. keeping Q/dQ in fp32 for numerical debugging or because autocast was disabled.

Common situations: Debugging with fp32 tensors then forgetting to cast back; a pipeline stage upcasting dQ to fp32 before postprocess; mismatched autocast contexts.

Related errors

• dQaccum tensor must be Float32
• Only int32 is supported for unique.
• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• Block sparse tensors{context} must have shapes (B, H, M) and
• The layout of mdV is wrong

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/8f73b15d0f2a5051. Report an issue: GitHub.