{"record":{"id":"eb12db26e85ddfdf","repo":"xai-org/x-algorithm","slug":"block-sparse-tensors-context-must-have-shapes-b","errorCode":null,"errorMessage":"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N).","messagePattern":"Block sparse tensors(.+?) must have shapes \\(B, H, M\\) and \\(B, H, M, N\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py","lineNumber":343,"sourceCode":"\n    if sparse_block_size_q % base_m_block != 0:\n        raise ValueError(\n            f\"Block sparse tensors{context} have block size {sparse_block_size_q}, \"\n            f\"which must be a multiple of {base_m_block}.\"\n        )\n\n    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)\n    expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv)\n    q_subtile_factor = sparse_block_size_q // base_m_block\n    expected_count_shape = (batch_size, num_head, expected_m_blocks)\n    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)\n\n    mask_block_cnt = tensors.mask_block_cnt\n    mask_block_idx = tensors.mask_block_idx\n    if mask_block_cnt is None or mask_block_idx is None:\n        raise ValueError(\"mask_block_cnt and mask_block_idx must be provided for block sparsity.\")\n    if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4:\n        raise ValueError(\n            f\"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N).\"\n        )\n    for dim_name, cur, tgt in (\n        (\"batch\", mask_block_cnt.shape[0], expected_count_shape[0]),\n        (\"head\", mask_block_cnt.shape[1], expected_count_shape[1]),\n    ):\n        if cur != tgt and cur != 1:\n            raise ValueError(f\"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.\")\n    for dim_name, cur, tgt in (\n        (\"batch\", mask_block_idx.shape[0], expected_index_shape[0]),\n        (\"head\", mask_block_idx.shape[1], expected_index_shape[1]),\n    ):\n        if cur != tgt and cur != 1:\n            raise ValueError(f\"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.\")\n    if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:\n        raise ValueError(f\"Block sparse tensors{context} must share the same m-block dimension.\")\n    if mask_block_idx.shape[3] > expected_n_blocks:\n        raise ValueError(","sourceCodeStart":325,"sourceCodeEnd":361,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py#L325-L361","documentation":"Raised by infer_block_sparse_expected_shapes when validating block-sparse attention mask tensors. The count tensor mask_block_cnt must be 3-D (B, H, M) and the index tensor mask_block_idx must be 4-D (B, H, M, N); any other rank is rejected before shapes are compared. This is an early structural check so downstream broadcast/expansion logic can assume fixed ranks.","triggerScenarios":"Calling normalize_block_sparse_config (directly or via the ranker FA4 forward path) with BlockSparseTensorsTorch whose mask_block_cnt is not ndim==3 or mask_block_idx is not ndim==4, e.g. passing a 2-D count tensor or a 5-D index tensor built from a custom BlockMask conversion.","commonSituations":"Converting a FlexAttention/BlockMask seqlens or indices tensor without adding the batch/head dims; accidentally stacking or squeezing a dim; mismatch between varlen (2-D) layouts and the dense (B, H, ...) layout this kernel expects.","solutions":["Reshape mask_block_cnt to (B, H, M) and mask_block_idx to (B, H, M, N) before passing them in","If you built these from a BlockMask, use the repo's provided conversion utilities instead of hand-rolling tensor indexing","Print .shape of both tensors right before the call and compare against the message's expected shapes"],"exampleFix":"// before\ncnt = counts.squeeze()        # now (H, M)\nidx = indices                 # (B, H, N)\n\n// after\ncnt = counts.squeeze(0).unsqueeze(0) if counts.ndim == 2 else counts  # (B, H, M)\nidx = indices.unsqueeze(2) if indices.ndim == 3 else indices          # (B, H, M, N)","handlingStrategy":"validation","validationCode":"def check_block_sparse_ranks(tensors):\n    assert tensors.mask_block_cnt is not None and tensors.mask_block_idx is not None\n    assert tensors.mask_block_cnt.ndim == 3, tensors.mask_block_cnt.shape\n    assert tensors.mask_block_idx.ndim == 4, tensors.mask_block_idx.shape","typeGuard":"def has_valid_block_ranks(t: BlockSparseTensorsTorch) -> bool:\n    return (t.mask_block_cnt is not None and t.mask_block_idx is not None\n            and t.mask_block_cnt.ndim == 3 and t.mask_block_idx.ndim == 4)","tryCatchPattern":null,"preventionTips":["Log .shape of both mask tensors before normalization","Build masks only via the repo's BlockMask conversion utilities"],"tags":["block-sparsity","shape-validation","flash-attention","cuda"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}