xai-org/x-algorithm · error · ValueError

{name} must be on the same device as block sparse tensors

Error message

{name} must be on the same device as block sparse tensors

What it means

Raised by _check_and_expand_metadata_tensor when a metadata tensor's device differs from the device of the block-sparse tensors. Everything must be co-located on the same CUDA device for the kernel to dereference them safely.

Source

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

        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,
    hint: str | Callable[[], str] | None,
    device: torch.device,
) -> torch.Tensor | None:
    if tensor is None:
        return None
    if tensor.dtype != torch.int32:
        raise ValueError(f"{name} must have dtype torch.int32")
    if tensor.device != device:
        raise ValueError(f"{name} must be on the same device as block sparse tensors")
    if not tensor.is_cuda:
        raise ValueError(f"{name} must live on CUDA")
    return _expand_sparsity_tensor(tensor, expected_shape, name, context, hint)


def get_block_sparse_expected_shapes(
    batch_size: int,
    num_head: int,
    seqlen_q: int,
    seqlen_k: int,
    m_block_size: int,
    n_block_size: int,
    q_stage: int,
) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
    m_block_size_effective = q_stage * m_block_size
    expected_m_blocks = ceildiv(seqlen_q, m_block_size_effective)
    expected_n_blocks = ceildiv(seqlen_k, n_block_size)
    expected_count_shape = (batch_size, num_head, expected_m_blocks)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Move the metadata tensor to the block tensors' device: t = t.to(cnt.device)
2. Centralize a single device variable and use it for all inputs
3. Assert t.device == cnt.device before the call

Example fix

# before
meta = compute_meta().cuda(0)
out = ranker(q.to('cuda:1'), ..., meta=meta)

# after
dev = q.device
meta = compute_meta().to(dev)
out = ranker(q.to(dev), ..., meta=meta)

device = mask_block_cnt.device
assert meta.device == device, f"meta on {meta.device}, blocks on {device}"

def matches_device(t: torch.Tensor, dev: torch.device) -> bool:
    return t.device == dev

When it happens

Trigger: Block tensors on 'cuda:0' but the metadata tensor on CPU or on 'cuda:1'; typically after moving q/k/v with .to(model.device) but leaving metadata behind.

Common situations: Multi-GPU training with device_per_rank; heterogeneous data pipelines where metadata is computed on CPU; moving only some inputs in .to() calls.

Related errors

• {name}_block_cnt and {name}_block_idx must be on the same de
• {name}_block tensors must live on CUDA
• {name} must have dtype torch.int32
• {name} must live on CUDA
• All block sparse tensors must be on the same device

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