fix(cake_sage_attn): cover Sage combinations and pass tensor maps by value

flashinfer/cute_dsl/sparse/bsa_attn_sm120.py

@@ -334,11 +334,11 @@ def _bsa_attn_sm120_blk64_sage_fwd_cake(
 ) -> torch.Tensor:
     """Run prequantized SM120 Sage block-sparse attention.
 
-    All CUDA storage is caller-owned. ``out`` is contiguous BF16 BHSD and
-    ``tma_descriptor_workspace`` is contiguous, 128-byte-aligned CUDA uint8
-    storage. Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS,
-    and the operation supports MHA, head dimension 128, non-causal forward
-    without LSE.
+    All CUDA storage is caller-owned. ``out`` is contiguous BF16 BHSD.
+    Tensor-map descriptors are passed by value as grid-constant kernel
+    parameters; ``tma_descriptor_workspace`` is accepted but unused.
+    Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS, and the
+    operation supports MHA, head dimension 128, non-causal forward without LSE.
     """
     import tvm_ffi  # noqa: PLC0415
 
@@ -367,10 +367,9 @@ def _bsa_attn_sm120_blk64_sage_fwd_cake(
         "v_scale": v_scale,
         "q2k_block_index": q2k_block_index,
         "out": out,
-        "tma_descriptor_workspace": tma_descriptor_workspace,
     }
     if not all(isinstance(value, torch.Tensor) for value in tensors.values()):
-        raise TypeError("all input, output, and workspace values must be torch.Tensor")
+        raise TypeError("all input and output values must be torch.Tensor")
     device = q_int8.device
     if device.type != "cuda" or device.index != torch.cuda.current_device():
         raise ValueError("q_int8 must use the current CUDA device")
@@ -396,11 +395,6 @@ def _bsa_attn_sm120_blk64_sage_fwd_cake(
     ):
         if tensor.dtype != torch.float32:
             raise TypeError(f"{name} must use torch.float32")
-    if tma_descriptor_workspace.dtype != torch.uint8:
-        raise TypeError("tma_descriptor_workspace must use torch.uint8")
-    if tma_descriptor_workspace.data_ptr() % 128:
-        raise ValueError("tma_descriptor_workspace must be 128-byte aligned")
-
     if q_int8.ndim != 4:
         raise ValueError("q_int8 must be rank-4 BHSD")
     batch, heads, seqlen_q, head_dim = (int(value) for value in q_int8.shape)
@@ -484,11 +478,6 @@ def _bsa_attn_sm120_blk64_sage_fwd_cake(
         uniform_nonempty,
         int(contiguous_block_indices),
     )
-    required_workspace = int(record["tma_workspace_bytes"])
-    if tma_descriptor_workspace.numel() < required_workspace:
-        raise ValueError(
-            f"tma_descriptor_workspace requires at least {required_workspace} bytes"
-        )
     bindings = {
         "Q_map": q_int8,
         "K_map": k_int8,
@@ -507,7 +496,6 @@ def _bsa_attn_sm120_blk64_sage_fwd_cake(
         "num_k_blocks": k_blocks,
         "block_sparse_num": block_sparse_num,
         "softmax_scale": float(softmax_scale),
-        "tma_descriptor_workspace": tma_descriptor_workspace,
         "grid_x": q_blocks,
         "grid_y": heads,
         "grid_z": batch,
@@ -545,11 +533,11 @@ def bsa_attn_sm120_blk64_sage_fwd(
 ) -> torch.Tensor:
     """Run prequantized SM120 Sage block-sparse attention.
 
-    All CUDA storage is caller-owned. ``out`` is contiguous BF16 BHSD and
-    ``tma_descriptor_workspace`` is contiguous, 128-byte-aligned CUDA uint8
-    storage. Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS,
-    and the operation supports MHA, head dimension 128, non-causal forward
-    without LSE.
+    All CUDA storage is caller-owned. ``out`` is contiguous BF16 BHSD.
+    Tensor-map descriptors are passed by value as grid-constant kernel
+    parameters; ``tma_descriptor_workspace`` is accepted but unused.
+    Q/K use contiguous INT8 BHSD, V uses contiguous FP8 E4M3 HDS, and the
+    operation supports MHA, head dimension 128, non-causal forward without LSE.
 
     Parameters
     ----------
@@ -585,8 +573,10 @@ def bsa_attn_sm120_blk64_sage_fwd(
     out : torch.Tensor
         Caller-owned contiguous BF16 output with shape ``[B, H, Sq, 128]``.
     tma_descriptor_workspace : torch.Tensor
-        Caller-owned contiguous CUDA uint8 workspace, aligned to 128 bytes
-        and large enough for the selected generated kernel.
+        Accepted for compatibility and unused. An empty tensor is sufficient;
+        this operation neither reads nor writes its storage. Tensor-map
+        descriptors are passed by value in the kernel launch, so workspace
+        contents, lifetime, and reuse across streams do not affect attention.
     uniform_block_count : bool
         Whether every query block uses ``block_sparse_num`` selected blocks.
         This is a caller-provided guarantee: when true, ``q2k_block_nums``
@@ -604,6 +594,21 @@ def bsa_attn_sm120_blk64_sage_fwd(
     -------
     torch.Tensor
         The caller-owned ``out`` tensor after attention output is written.
+
+    Notes
+    -----
+    Non-contiguous indices support both per-row ``q2k_block_nums`` and the
+    scalar ``block_sparse_num`` count, with every supported ``block_sizes``
+    layout and with aligned or partial final KV blocks. Omitting
+    ``q2k_block_nums`` uses the scalar count even when
+    ``uniform_block_count=False``. Setting ``uniform_block_count=True``
+    uses the scalar count and ignores any supplied per-row counts. A row
+    selecting zero blocks produces an exactly zero output.
+
+    Initialize the JIT module before CUDA Graph capture. Keep captured input
+    and output tensors alive while their graph is in use. Each captured launch
+    retains its own tensor-map descriptors, so launches with different tensor
+    bindings can be interleaved without shared descriptor storage.
     """
     if backend != "cake":
         raise ValueError(f"unsupported SM120 Sage backend: {backend!r}")

tests/attention/test_cake_sage_block_sparse_attention.py

@@ -298,6 +298,77 @@
     }
 ]""")
 
+# These choices are independent in the public API.  In particular, sorted
+# top-k indices still require the non-contiguous route with either count API.
+_COUNT_MODES = (
+    ("per_row", {"varying_block_counts": True, "empty_first_row": True}),
+    ("omitted", {"omit_block_nums": True, "uniform_block_count": False}),
+    ("uniform", {"uniform_block_count": True}),
+    ("omitted_uniform", {"omit_block_nums": True, "uniform_block_count": True}),
+    (
+        "empty_omitted",
+        {"omit_block_nums": True, "uniform_block_count": False, "block_sparse_num": 0},
+    ),
+    ("empty_uniform", {"uniform_block_count": True, "block_sparse_num": 0}),
+    (
+        "empty_contiguous",
+        {
+            "omit_block_nums": True,
+            "uniform_block_count": True,
+            "contiguous_block_indices": True,
+            "noncontiguous_topk": False,
+            "block_sparse_num": 0,
+        },
+    ),
+)
+_CASES.extend(
+    {
+        "name": f"topk_{count_name}_sizes{mode}_k{seqlen_k}",
+        "params": {
+            "batch_size": 2,
+            "num_heads": 2,
+            "seqlen_q": 100,
+            "seqlen_k": seqlen_k,
+            "selected_blocks": 3,
+            "softmax_scale": "default",
+            "block_sizes_mode": mode,
+            "valid_last_block": 40,
+            "empty_first_row": False,
+            "position_sensitive": False,
+            "noncontiguous_topk": True,
+            "contiguous_block_indices": False,
+            "seed": 508300 + mode,
+            **count_options,
+        },
+    }
+    for mode in range(4)
+    for seqlen_k in (320, 257)
+    for count_name, count_options in _COUNT_MODES
+)
+_CASES.extend(
+    {
+        "name": f"topk_scalar_count_below_capacity_uniform{uniform}",
+        "params": {
+            "batch_size": 1,
+            "num_heads": 2,
+            "seqlen_q": 100,
+            "seqlen_k": 257,
+            "selected_blocks": 3,
+            "softmax_scale": "default",
+            "block_sizes_mode": 0,
+            "empty_first_row": False,
+            "position_sensitive": False,
+            "noncontiguous_topk": True,
+            "contiguous_block_indices": False,
+            "uniform_block_count": uniform,
+            "omit_block_nums": not uniform,
+            "block_sparse_num": 2,
+            "seed": 508399,
+        },
+    }
+    for uniform in (False, True)
+)
+
 
 def _ceil_div(value: int, divisor: int) -> int:
     return (value + divisor - 1) // divisor
@@ -311,15 +382,21 @@ def _make_block_sizes(params, *, batch: int, heads: int, k_blocks: int, device):
     if mode == 1:
         sizes = torch.full((k_blocks,), _BLOCK, dtype=torch.int32, device=device)
         sizes[-1] = valid_last
-        return sizes
-    if mode == 2:
+    elif mode == 2:
         sizes = torch.full((batch, k_blocks), _BLOCK, dtype=torch.int32, device=device)
         sizes[:, -1] = valid_last
-        return sizes
-    sizes = torch.full(
-        (batch, heads, k_blocks), _BLOCK, dtype=torch.int32, device=device
-    )
-    sizes[:, :, -1] = valid_last
+    else:
+        sizes = torch.full(
+            (batch, heads, k_blocks), _BLOCK, dtype=torch.int32, device=device
+        )
+        sizes[:, :, -1] = valid_last
+    if params.get("noncontiguous_topk", False):
+        sizes[..., 0] = 0
+        if mode >= 2:
+            offsets = torch.arange(batch, device=device) * 7
+            if mode == 3:
+                offsets = offsets[:, None] + torch.arange(heads, device=device) * 11
+            sizes[..., -1] = (valid_last - offsets).clamp(min=0)
     return sizes
 
 
@@ -433,9 +510,20 @@ def _make_inputs(params, *, device):
             + 0.25
         )
 
+    block_sizes = _make_block_sizes(
+        params, batch=batch, heads=heads, k_blocks=k_blocks, device=device
+    )
     if bool(params.get("noncontiguous_topk", False)):
         # Make attention to padded tokens fail the numerical comparison.
         logical_v[..., seqlen_k:] = 64.0
+        if block_sizes is not None:
+            sizes = block_sizes.reshape(
+                batch if block_sizes.ndim >= 2 else 1,
+                heads if block_sizes.ndim == 3 else 1,
+                k_blocks,
+            )
+            invalid = torch.arange(_BLOCK, device=device) >= sizes.unsqueeze(-1)
+            logical_v.masked_fill_(invalid.reshape(*sizes.shape[:2], 1, padded_k), 64.0)
 
     inverse_perm = torch.tensor(_V_INVERSE_PERM, dtype=torch.long, device=device)
     v_fp8 = logical_v.view(batch, heads, _HEAD_DIM, k_blocks * 4, 16)
@@ -476,17 +564,27 @@ def _make_inputs(params, *, device):
         )
         q2k = torch.cat((selected, tail), dim=-1).sort(dim=-1).values.to(torch.int32)
     q2k = q2k.contiguous()
+    block_sparse_num = int(params.get("block_sparse_num", selected_blocks))
     q2k_nums = torch.full(
         (batch, heads, q_blocks),
-        selected_blocks,
+        block_sparse_num,
         dtype=torch.int32,
         device=device,
     )
+    if params.get("varying_block_counts", False):
+        q2k_nums = (
+            (
+                torch.arange(batch * heads * q_blocks, device=device)
+                % (selected_blocks + 1)
+            )
+            .reshape(batch, heads, q_blocks)
+            .to(torch.int32)
+        )
     if bool(params["empty_first_row"]):
         q2k_nums[0, 0, 0] = 0
 
-    block_sizes = _make_block_sizes(
-        params, batch=batch, heads=heads, k_blocks=k_blocks, device=device
+    uniform_contiguous = not params["empty_first_row"] and not params.get(
+        "noncontiguous_topk", False
     )
     softmax_scale = (
         _HEAD_DIM**-0.5
@@ -502,9 +600,15 @@ def _make_inputs(params, *, device):
         "K_scale": k_scale.contiguous(),
         "V_scale": v_scale.contiguous(),
         "q2k_block_index": q2k,
-        "q2k_block_nums": q2k_nums.contiguous(),
+        "q2k_block_nums": None
+        if params.get("omit_block_nums", False)
+        else q2k_nums.contiguous(),
+        "uniform_block_count": params.get("uniform_block_count", uniform_contiguous),
+        "contiguous_block_indices": params.get(
+            "contiguous_block_indices", uniform_contiguous
+        ),
         "block_sizes": block_sizes,
-        "block_sparse_num": selected_blocks,
+        "block_sparse_num": block_sparse_num,
         "softmax_scale": float(softmax_scale),
         "O": torch.empty(
             (batch, heads, seqlen_q, _HEAD_DIM),
@@ -561,8 +665,11 @@ def _reference(inputs):
             for q_block in range(_ceil_div(seqlen_q, _BLOCK)):
                 q_start = q_block * _BLOCK
                 q_end = min(seqlen_q, q_start + _BLOCK)
-                count = int(
-                    inputs["q2k_block_nums"][batch_idx, head_idx, q_block].item()
+                counts = inputs["q2k_block_nums"]
+                count = (
+                    int(inputs["block_sparse_num"])
+                    if counts is None or inputs["uniform_block_count"]
+                    else int(counts[batch_idx, head_idx, q_block].item())
                 )
                 if count == 0:
                     continue
@@ -603,7 +710,9 @@ def _reference(inputs):
 
 
 def _aligned_workspace(*, device):
-    storage = torch.empty(_WORKSPACE_BYTES + 127, dtype=torch.uint8, device=device)
+    storage = torch.full(
+        (_WORKSPACE_BYTES + 127,), 0xA5, dtype=torch.uint8, device=device
+    )
     offset = (-int(storage.data_ptr())) % 128
     workspace = storage[offset : offset + _WORKSPACE_BYTES]
     assert workspace.numel() == _WORKSPACE_BYTES
@@ -626,7 +735,16 @@ def test_cake_sage_block_sparse_attention(case):
         tail_block = _ceil_div(int(params["seqlen_k"]), _BLOCK) - 1
         assert torch.all(q2k[..., -1] == tail_block).item()
         assert torch.all(torch.any(q2k[..., 1:] - q2k[..., :-1] > 1, dim=-1)).item()
-    uniform_contiguous = not bool(params["empty_first_row"]) and not noncontiguous_topk
+    if inputs["uniform_block_count"] and inputs["q2k_block_nums"] is not None:
+        # A true uniform guarantee makes the scalar count authoritative, even
+        # if the supplied per-row counts describe a different selection.
+        inputs["q2k_block_nums"].fill_(1 if inputs["block_sparse_num"] == 0 else 0)
+    elif inputs["q2k_block_nums"] is not None:
+        slots = torch.arange(
+            inputs["q2k_block_index"].shape[-1], device=inputs["Q"].device
+        )
+        inactive = slots >= inputs["q2k_block_nums"].unsqueeze(-1)
+        inputs["q2k_block_index"].masked_fill_(inactive, -1)
 
     returned = bsa_attn_sm120_blk64_sage_fwd(
         inputs["Q"],
@@ -642,16 +760,19 @@ def test_cake_sage_block_sparse_attention(case):
         softmax_scale=softmax_scale,
         out=inputs["O"],
         tma_descriptor_workspace=workspace,
-        uniform_block_count=uniform_contiguous,
-        contiguous_block_indices=uniform_contiguous,
+        uniform_block_count=inputs["uniform_block_count"],
+        contiguous_block_indices=inputs["contiguous_block_indices"],
         backend="cake",
     )
 
-    assert workspace_storage.data_ptr() != 0
+    # The compatibility workspace is unused by every exported specialization.
+    assert torch.all(workspace_storage == 0xA5).item()
     assert returned.data_ptr() == inputs["O"].data_ptr()
     if bool(params["empty_first_row"]):
         first_q_block = returned[0, 0, : min(_BLOCK, int(params["seqlen_q"]))]
         assert torch.count_nonzero(first_q_block).item() == 0
+    if inputs["block_sparse_num"] == 0:
+        assert torch.count_nonzero(returned).item() == 0
     atol, rtol = (
         (3.0e-3, 0.0) if bool(params["position_sensitive"]) else (1.0e-2, 1.0e-2)
     )