video_sparse_attn_h3 ¶

def prepare_for_regional_compile(self, device: torch.device) -> str | None:
    """Resolve the inference-only sm_100a route before fullgraph capture.

    The ordinary eager route probes the environment, extension, device,
    and tensor contract at every call so it can warn and fall back.  Those
    Python/device-capability checks are not safe inside a regional
    ``fullgraph=True`` block.  Probe one representative tile-64 input on
    the loaded model's device now, then let ``forward`` specialize on the
    resulting plain bool while Dynamo is compiling.
    """
    requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
    enabled = False
    reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention"
    if requested:
        if _sm100a is None:
            reason = "fastvideo_kernel.block_sparse_attn_sm100a is not installed"
        elif not _sm100a_has_compile_safe_mask_route(_sm100a):
            reason = ("neither a native block_sparse_attn_sm100a_from_mask entry nor the raw sm100a "
                      "kernel plus map_to_index compatibility route is installed")
        else:
            # Two 64-token blocks exercise the exact sm_100a inference
            # specialization while keeping the one-time probe tiny.  The
            # kernel predicate checks extension presence, CUDA capability,
            # dtype/layout, head size, block size, and even block count
            # without reading metadata tensor contents.
            probe_query = torch.empty((1, 1, 128, self.head_size), device=device, dtype=torch.bfloat16)
            probe_block_sizes = torch.full((2, ), 64, device=device, dtype=torch.int32)
            reason = _sm100a_unavailable_reason(
                _sm100a,
                probe_query,
                probe_block_sizes,
                grad_mode=False,
            )
            enabled = reason is None

    self._regional_compile_sm100a_enabled = enabled
    # Keep this marker unset when preparation fails. Generic/training
    # torch.compile must retain the established Triton attention route.
    self._regional_compile_layer_idx = (torch.tensor(self.layer_idx, device=device, dtype=torch.int64)
                                        if enabled else None)
    if enabled:
        route = ("native fastvideo-kernel mask entry" if callable(
            getattr(_sm100a, "block_sparse_attn_sm100a_from_mask", None)) else
                 "FastVideo compatibility mask adapter")
        logger.info_once(f"VSA-H3 regional compile mask route: {route}")
    if requested and reason is not None:
        logger.warning_once(f"VSA-H3 regional compile is unavailable and will stay eager: {reason}")
    return reason

def tile(self, x: torch.Tensor, attn_metadata: MiniMaxH3VSAMetadata) -> torch.Tensor:
    """Scatter rows into the padded tile buffer (pad positions stay zero).

    The returned tensor aliases the builder-owned buffer; callers must
    consume it before the next ``tile()`` (both call sites in
    ``forward()`` read it immediately). Odd tile-64 no-grad sm100a
    requests carry one additional all-zero tile internally; metadata and
    all observable outputs retain the logical geometry.
    """
    if x.shape[1] != attn_metadata.total_seq_length:
        raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, "
                         f"got {x.shape[1]}. A non-packed sequence (e.g. the token refiner) is "
                         "routed to the VSA-H3 backend; exclude it from the supported backends.")
    n_tiles = attn_metadata.variable_block_sizes.numel()
    grad_mode = torch.is_grad_enabled() and x.requires_grad
    compiling = torch.compiler.is_compiling()
    regional_compiling = compiling and self._regional_compile_layer_idx is not None
    if regional_compiling:
        sm100a_requested = bool(self._regional_compile_sm100a_enabled)
    elif compiling:
        # Training/generic compile keeps the long-standing Triton route.
        sm100a_requested = False
    else:
        sm100a_requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
    needs_sm100a_pair = (attn_metadata.tile_elems == 64 and n_tiles % 2 != 0 and not grad_mode and sm100a_requested)
    kernel_tiles = n_tiles + int(needs_sm100a_pair)
    target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1])

    # ``untile_combined_index`` maps each packed row to a logical tile
    # slot. Different geometries can share one transport shape; clear a
    # reused allocation once when the mapping identity changes so no old
    # valid row can survive as padding.
    holder = attn_metadata.tile_buf_holder
    if holder is None:
        raise RuntimeError("VSA-H3 metadata has no builder-owned tile buffer holder")
    buffer_matches = (holder.buffer is not None and holder.buffer.shape == target_shape
                      and holder.buffer.dtype == x.dtype and holder.buffer.device == x.device)
    if buffer_matches and holder.untile_geometry is not attn_metadata.untile_combined_index:
        holder.buffer.zero_()
    holder.buffer = scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, holder.buffer)
    holder.untile_geometry = attn_metadata.untile_combined_index
    if needs_sm100a_pair:
        # A prior even geometry can reuse this allocation and may have
        # written the last tile as logical data.
        holder.buffer[:, n_tiles * attn_metadata.tile_elems:].zero_()
    return holder.buffer