xai-org/x-algorithm · error · TypeError

dQaccum tensor must be Float32

Error message

dQaccum tensor must be Float32

What it means

When mdQaccum is provided to the flash backward postprocess kernel, it must be Float32 — the accumulator is kept in higher precision and downcast only at the end. Any other dtype fails the const_expr type check.

Source

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

            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
            num_block = cute.ceil_div(mdQ.shape[0], self.tile_m)
        else:
            TileScheduler = SingleTileScheduler

View on GitHub (pinned to 24c60942c5)

Solutions

1. Allocate mdQaccum as torch.float32 (torch.zeros(..., dtype=torch.float32, device=...))
2. Don't let autocast dictate the accumulator dtype — create it explicitly
3. If memory-bound, keep fp32 accumulator but free it promptly after postprocess

Example fix

# before
dqaccum = torch.zeros(q_shape, dtype=q.dtype, device='cuda')  # bf16

# after
dqaccum = torch.zeros(q_shape, dtype=torch.float32, device='cuda')

assert mdQaccum.dtype == torch.float32, mdQaccum.dtype

def is_fp32(t: torch.Tensor) -> bool:
    return t.dtype == torch.float32

When it happens

Trigger: Passing a bf16/fp16 dQ accumulation tensor to the postprocess kernel, e.g. allocating dQaccum with the same dtype as Q to save memory.

Common situations: Memory optimizations that shrink the accumulator; dtype propagation from autocast allocating bf16 buffers; adapting a pipeline that assumed fp16 accumulators.

Related errors

• Only Float16 or BFloat16 is supported
• Only int32 is supported for unique.
• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• {name}_block tensors must have dtype torch.int32
• Block sparse tensors{context} must have shapes (B, H, M) and

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