xai-org/x-algorithm · error · NotImplementedError

Block sparsity + paged KV not supported on SM100

Error message

Block sparsity + paged KV not supported on SM100

What it means

The SM100 flash-attention forward kernel does not implement the combination of block sparsity (blocksparse_tensors provided) and paged KV (mPageTable provided). The check use_block_sparsity and mPageTable is not None aborts with NotImplementedError at launch time.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/flash_fwd_sm100.py:783

        self.shared_storage = SharedStorage

        softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(
            softmax_scale, self.score_mod
        )
        window_size_left = Int32(window_size_left) if window_size_left is not None else None
        window_size_right = Int32(window_size_right) if window_size_right is not None else None
        fastdiv_mods = utils.compute_fastdiv_mods(
            mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_tensors, mPageTable
        )

        head_divmod = None
        if cutlass.const_expr(self.pack_gqa):
            head_divmod = FastDivmodDivisor(self.qhead_per_kvhead)

        self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None)
        if cutlass.const_expr(self.use_block_sparsity and mPageTable is not None):
            raise NotImplementedError("Block sparsity + paged KV not supported on SM100")
        if cutlass.const_expr(self.use_block_sparsity and self.is_varlen_q):
            assert const_expr(blocksparse_tensors.cu_total_m_blocks is not None), (
                "blocksparse_tensors.cu_total_m_blocks must be provided for varlen blocksparsity"
            )

        self.kernel(
            mQ,
            mK,
            mV,
            mO,
            mLSE,
            mCuSeqlensQ,
            mCuSeqlensK,
            mSeqUsedQ,
            mSeqUsedK,
            mPageTable,
            tma_atom_Q,
            tma_atom_K,

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Disable block sparsity for this call (pass blocksparse_tensors=None)
  2. Or disable paged KV (pass mPageTable=None and use dense KV tensors)
  3. Materialize the paged KV into a contiguous/dense tensor first, then use block sparsity
  4. Request/port an implementation combining both features for SM100 upstream

Example fix

// before
out = flash_attn(q, k, v, mPageTable=page_table, blocksparse_tensors=bs)
// after
k_dense, v_dense = unpaged_from_table(k, v, page_table)
out = flash_attn(q, k_dense, v_dense, blocksparse_tensors=bs)
Defensive patterns

Strategy: validation

Validate before calling

if blocksparse_tensors is not None and page_table is not None:
    raise ValueError('block sparsity + paged KV unsupported; densify KV or drop sparsity')

Try / catch

try:
    out = flash_attn(q, k, v, mPageTable=pt, blocksparse_tensors=bs)
except NotImplementedError:
    k, v = gather_paged(k, v, pt)
    out = flash_attn(q, k, v, blocksparse_tensors=bs)

Prevention

When it happens

Trigger: Passing both a non-None blocksparse_tensors argument and a non-None mPageTable (page table for paged KV cache) to __call__.

Common situations: Serving stacks that use vLLM-style paged KV caches combined with block-sparse attention masks; upgrading a workflow that previously used only one of the two features on a different architecture.

Related errors


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