sgl-project/sglang · error · ValueError

{name}_block tensors must live on CUDA

Error message

{name}_block tensors must live on CUDA

What it means

Raised when mask/full block count or index tensors are not on CUDA. The FA4 cute block-sparse kernels read these tensors directly from GPU memory, so CPU-resident tensors are rejected during normalization.

Source

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

    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


def _check_and_expand_metadata_tensor(
    name: str,
    tensor: torch.Tensor | None,

View on GitHub (pinned to 0132848349)

Solutions

  1. Move all block sparse tensors to CUDA: cnt = cnt.to('cuda')
  2. Create them directly on GPU: torch.zeros(shape, dtype=torch.int32, device='cuda')
  3. Check tensor.is_cuda in a helper before calling the attention op

Example fix

// before
cnt = torch.zeros((B,H,M), dtype=torch.int32)  # CPU
// after
cnt = torch.zeros((B,H,M), dtype=torch.int32, device='cuda')
Defensive patterns

Strategy: validation

Validate before calling

assert mask_block_cnt.is_cuda and mask_block_idx.is_cuda, 'block sparse tensors must be CUDA'

Prevention

When it happens

Trigger: Passing mask_block_cnt or mask_block_idx (or full_block_* / dq_write_order) that was created with torch.randint(..., device='cpu') or from a numpy conversion, without calling .cuda()/.to('cuda').

Common situations: Prototyping a BlockMask on CPU, loading sparse metadata from disk/numpy, or forgetting the device= argument when generating test tensors.

Related errors


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