xai-org/x-algorithm · error · TypeError

Type mismatch: {self.q_dtype} != {self.k_dtype}

Error message

Type mismatch: {self.q_dtype} != {self.k_dtype}

What it means

The SM100 flash-attention forward kernel requires Q, K, and V to share the same element type because the tcgen05 MMA and TMA paths are compiled for one dtype. Before launching the inner kernel, __call__ compares self.q_dtype with self.k_dtype and raises TypeError on mismatch.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py:490

                [2, 4, 3, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 3, 2, 0]
            )
            LSE_layout_transpose = [3, 2, 1, 0] if const_expr(mCuSeqlensQ is None) else [2, 1, 0]
            num_splits = mO.shape[0]
        else:
            O_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
            LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
            num_splits = Int32(1)
        mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose))
        mLSE = (
            cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
            if const_expr(mLSE is not None)
            else None
        )
        V_layout_transpose = [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2]
        mV = cute.make_tensor(mV.iterator, cute.select(mV.layout, mode=V_layout_transpose))

        if const_expr(self.q_dtype != self.k_dtype):
            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}")
        if const_expr(self.q_dtype != self.v_dtype):
            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}")
        if const_expr(self.q_dtype.width == 8):
            paged_kv_non_tma = not self.use_tma_KV
            if const_expr(self.head_dim_padded < 96):
                fp8_regs = _FP8_SMALL_HDIM_REGS[paged_kv_non_tma]
                self.num_regs_softmax = fp8_regs["num_regs_softmax"]
                self.num_regs_correction = fp8_regs["num_regs_correction"]
                self.num_regs_other = fp8_regs["num_regs_other"]
            else:
                fp8_tune = _FP8_TUNING_CONFIG.get(
                    (self.use_2cta_instrs, self.is_causal, self.head_dim_padded, self.is_sm103), {}
                )
                if const_expr("ex2_emu_freq" in fp8_tune):
                    self._tune = {**self._tune, **fp8_tune}
                    self.enable_ex2_emu = self._tune["ex2_emu_freq"] > 0
                if const_expr(not paged_kv_non_tma and "num_regs_softmax" in fp8_tune):
                    self.num_regs_softmax = fp8_tune["num_regs_softmax"]

View on GitHub (pinned to 24c60942c5)

Solutions

1. Cast Q to the same dtype as K (usually the KV dtype is the reference): q = q.to(k.dtype)
2. Make all of Q/K/V consistent, e.g. all bfloat16 or all float16, in the model config before calling the kernel
3. If fp8 Q was intended, ensure K and V are also fp8 with the same variant (E4M3FN vs E5M2)
4. Add an assertion in the wrapper that q.dtype == k.dtype == v.dtype to fail early with a clearer message

Example fix

// before
out = flash_attn(q, k, v)  # q: bf16, k: fp16
// after
q = q.to(k.dtype)
out = flash_attn(q, k, v)

assert q.dtype == k.dtype, f'q {q.dtype} != k {k.dtype}'

def same_dtype(*ts):
    d = ts[0].dtype
    return all(t.dtype == d for t in ts)

When it happens

Trigger: Calling the forward kernel with Q in one dtype (e.g. bfloat16) and K in another (e.g. float16 or fp8), so const_expr(q_dtype != k_dtype) is true at compile time.

Common situations: Mixing a bf16 Q embedding with fp16 KV projections, enabling fp8 KV cache while Q stays bf16, or a dataloader/config change that cast one tensor but not the others.

Related errors

• Type mismatch: {self.q_dtype} != {self.v_dtype}
• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• Block sparse tensors{context} must have shapes (B, H, M) and
• Only Float16 or BFloat16 is supported
• The layout of mdV is wrong

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