sgl-project/sglang · 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 when validating block-sparse attention metadata: the mask/full block count tensor and index tensor for the same group must reside on the same torch device. The library enforces this before expanding/normalizing the tensors for the FA4 cute block-sparse kernels, since the kernel consumes them as paired device pointers.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/block_sparsity.py:250

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
    )
    # [Note] Allow Compact block sparse indices
    # Allow the last dimension (n_blocks) of idx to be <= expected, since
    # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact
    # index tensors that avoid O(N^2) memory at long sequence lengths.
    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

View on GitHub (pinned to 0132848349)

Solutions

  1. Move both cnt and idx to the same device: idx = idx.to(cnt.device)
  2. Verify no tensor was left on CPU (e.g. built from a numpy array) before passing BlockSparseTensorsTorch
  3. In multi-GPU/TP setups, move all block sparse tensors to the current rank's device

Example fix

// before
tensors = BlockSparseTensorsTorch(mask_block_cnt=cnt_cpu, mask_block_idx=idx_cuda)
// after
tensors = BlockSparseTensorsTorch(mask_block_cnt=cnt_cpu.to(idx_cuda.device), mask_block_idx=idx_cuda)
Defensive patterns

Strategy: validation

Validate before calling

dev = mask_block_cnt.device
assert mask_block_idx.device == dev, f'idx on {mask_block_idx.device}, cnt on {dev}'

Prevention

When it happens

Trigger: Calling normalize_block_sparse_tensors (or normalize_block_sparse_config / _bwd) with e.g. mask_block_cnt on cuda:0 and mask_block_idx on cuda:1 or on CPU.

Common situations: Tensors created at different times (one from a cached CPU BlockMask conversion, one moved to GPU), or multi-GPU runs where tensors were pinned to different cuda devices.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/b48b587646a36a72. Report an issue: GitHub.