xai-org/x-algorithm · error · ValueError

{name}_block tensors must live on CUDA

Error message

{name}_block tensors must live on CUDA

What it means

Raised by _check_and_expand_block when either the block count or index tensor of a block-sparsity pair is not on CUDA. The downstream cute-DSL kernels are GPU-only, so CPU metadata is rejected before kernel launch.

Source

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

    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,
    hint: str | Callable[[], str] | None,
    device: torch.device,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Move both block tensors to CUDA: cnt = cnt.cuda(); idx = idx.cuda()
2. Verify a CUDA device is available before calling (torch.cuda.is_available())
3. Keep all block-sparsity metadata on the same GPU as q/k/v

Example fix

# before
mask_block_cnt = torch.zeros(B, H, M, dtype=torch.int32)

# after
mask_block_cnt = torch.zeros(B, H, M, dtype=torch.int32, device='cuda')

assert cnt.is_cuda and idx.is_cuda, "block tensors must be on CUDA"

def on_cuda(*ts) -> bool:
    return all(t is None or t.is_cuda for t in ts)

When it happens

Trigger: Passing mask_block_cnt and/or mask_block_idx that are CPU tensors (e.g. freshly created via torch.zeros on CPU or loaded from a file without .to('cuda')) to normalize_block_sparse_tensors / the ranker API.

Common situations: Prototyping on CPU then switching to the CUDA path without moving metadata tensors; loading pickled sparse schedules from disk; debugging scripts that build indices with numpy and forget the device transfer.

Related errors

• {name} must live on CUDA
• {name}_block_cnt and {name}_block_idx must be on the same de
• {name} must have dtype torch.int32
• {name} must be on the same device as block sparse tensors
• Block sparse tensors{context} must have shapes (B, H, M) and

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