xai-org/x-algorithm · error · ValueError

{name}_block_cnt and {name}_block_idx must be on the same de

Error message

{name}_block_cnt and {name}_block_idx must be on the same device

What it means

Raised by _check_and_expand_block (called from normalize_block_sparse_tensors) when the block-sparsity count tensor and index tensor for the same named pair (e.g. mask_block_cnt / mask_block_idx) are on different torch devices. The library requires both tensors of a pair to be co-located because the CUDA kernels index them together.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py:242

def _check_and_expand_block(
    name: str,
    cnt: torch.Tensor | None,
    idx: torch.Tensor | None,
    expected_count_shape: Tuple[int, ...],
    expected_index_shape: Tuple[int, ...],
    context: str | None,
    hint: str | Callable[[], str] | None,
) -> Tuple[torch.Tensor | None, torch.Tensor | None]:
    if (cnt is None) != (idx is None):
        raise ValueError(
            f"{name}_block_cnt and {name}_block_idx must both be provided or both be None"
        )
    if cnt is None or idx is None:
        return None, None
    if cnt.dtype != torch.int32 or idx.dtype != torch.int32:
        raise ValueError(f"{name}_block tensors must have dtype torch.int32")
    if cnt.device != idx.device:
        raise ValueError(f"{name}_block_cnt and {name}_block_idx must be on the same device")
    if not cnt.is_cuda or not idx.is_cuda:
        raise ValueError(f"{name}_block tensors must live on CUDA")
    expanded_cnt = _expand_sparsity_tensor(
        cnt, expected_count_shape, f"{name}_block_cnt", context, hint
    )
    if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]:
        expected_index_shape = (*expected_index_shape[:3], idx.shape[3])
    expanded_idx = _expand_sparsity_tensor(
        idx, expected_index_shape, f"{name}_block_idx", context, hint
    )
    return expanded_cnt, expanded_idx


def _check_and_expand_metadata_tensor(
    name: str,
    tensor: torch.Tensor | None,
    expected_shape: Tuple[int, ...],
    context: str | None,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Move both tensors to the same device, e.g. cnt = cnt.to(idx.device) before calling the API
2. Ensure cnt.device == idx.device in a pre-flight assert
3. Check for accidental .cuda() on only one tensor in data loading code

Example fix

# before
mask_block_cnt = torch.from_numpy(cnt_np).cuda()
mask_block_idx = torch.from_numpy(idx_np)  # still on CPU

# after
mask_block_idx = torch.from_numpy(idx_np).to(mask_block_cnt.device)

assert cnt is None or idx is None or cnt.device == idx.device, f"device mismatch: {cnt.device} vs {idx.device}"

def same_device(a: torch.Tensor, b: torch.Tensor) -> bool:
    return a.device == b.device

try:
    out = api(...)
except ValueError as e:
    if "same device" in str(e):
        cnt = cnt.to(idx.device)
        out = api(...)
    else:
        raise

When it happens

Trigger: Calling the ranker/attention API with mask_block_cnt on 'cuda:0' but mask_block_idx on 'cpu' (or on a different GPU, e.g. 'cuda:1'). Also happens when one tensor is created with device=... and the other comes from a cached/preloaded tensor that was never moved.

Common situations: Loading sparse metadata from disk (numpy/CPU tensors) and only moving one of the two tensors to GPU; multi-GPU pipelines where tensors are pinned to different devices; mixing tensors produced by different pipeline stages.

Related errors

• {name} must have dtype torch.int32
• {name} must be on the same device as block sparse tensors
• {name}_block tensors must live on CUDA
• {name} must live on CUDA
• mask_block_cnt and mask_block_idx must be provided for block

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