xai-org/x-algorithm · error · TypeError

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

Error message

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

What it means

The SM100 flash-attention forward kernel requires Q, V (and K) to have identical element types since the MMA descriptors and TMA copies are specialized per dtype. __call__ checks self.q_dtype against self.v_dtype and raises TypeError when they differ.

Source

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

            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"]
                    self.num_regs_correction = fp8_tune["num_regs_correction"]
                    self.num_regs_other = 512 - self.num_regs_softmax * 2 - self.num_regs_correction

View on GitHub (pinned to 24c60942c5)

Solutions

1. Cast Q to match V's dtype (or vice versa) before the call: q = q.to(v.dtype)
2. Unify all three tensors to one dtype in the model/config
3. If fp8 is intended, quantize Q, K, and V together with the same fp8 variant
4. Add an upfront assert q.dtype == v.dtype in the caller to fail before kernel launch

Example fix

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

assert q.dtype == v.dtype == k.dtype, f'{q.dtype}, {k.dtype}, {v.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 and V in different dtypes (e.g. Q bf16, V fp8e4m3), triggering the const_expr(q_dtype != v_dtype) branch.

Common situations: Enabling an fp8 V cache while keeping Q in bf16, mixed-precision inference configs that quantize only V, or a projection layer output dtype differing from the Q path.

Related errors

• Type mismatch: {self.q_dtype} != {self.k_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/e657ab8fafb794a7. Report an issue: GitHub.