xai-org/x-algorithm · error · ValueError

All block sparse tensors must be on the same device

Error message

All block sparse tensors must be on the same device

What it means

All block-sparse tensors in one config must live on the same CUDA device. Here the optional full-block count tensor's device differs from mask_block_cnt's device, which would cause cross-device copies or kernel faults later.

Source

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

        expected_count_shape,
        expected_index_shape,
        context,
        hint,
    )
    if mask_cnt is None or mask_idx is None:
        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")

    full_cnt, full_idx = _check_and_expand_block(
        "full",
        tensors.full_block_cnt,
        tensors.full_block_idx,
        expected_count_shape,
        expected_index_shape,
        context,
        hint,
    )
    if full_cnt is not None and mask_cnt.device != full_cnt.device:
        raise ValueError("All block sparse tensors must be on the same device")

    diag_cnt, diag_idx = _check_and_expand_block(
        "diag",
        tensors.diag_block_cnt,
        tensors.diag_block_idx,
        expected_count_shape,
        expected_index_shape,
        context,
        hint,
    )
    if diag_cnt is not None and mask_cnt.device != diag_cnt.device:
        raise ValueError("All block sparse tensors must be on the same device")

    dq_write_order = _check_and_expand_metadata_tensor(
        "dq_write_order",
        tensors.dq_write_order,
        tuple(mask_idx.shape),
        context,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Move all block-sparse tensors to the same device: tensors.to(device) or per-tensor .to(mask_block_cnt.device)
2. Standardize on creating/loading all masks inside the module's device context
3. Add a device consistency assert in your data pipeline

Example fix

# before
full_cnt = torch.load('full_cnt.pt')          # cpu / cuda:0
mask_cnt = make_mask(device='cuda:1')

# after
device = mask_cnt.device
full_cnt = torch.load('full_cnt.pt', map_location=device)
# or: tensors = tensors.to(device) before normalize_*

dev = tensors.mask_block_cnt.device
for t in (tensors.full_block_cnt, tensors.diag_block_cnt):
    if t is not None:
        assert t.device == dev, (t.device, dev)

def all_same_device(t: BlockSparseTensorsTorch) -> bool:
    dev = t.mask_block_cnt.device
    return all(x is None or x.device == dev for x in
               (t.full_block_cnt, t.full_block_idx, t.diag_block_cnt, t.diag_block_idx))

When it happens

Trigger: Loading full_block_cnt from disk (defaults to CPU or cuda:0) while the mask tensors live on cuda:1; mixing tensors created in different device contexts in multi-GPU training.

Common situations: Multi-GPU models with device_map sharding; masks precomputed and cached on CPU then partially moved to GPU; tensor created inside a torch.cuda.device(idx) context differing from the rest.

Related errors

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

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