INNER CODE UNIT · Python

block_fill_cuda

microsoft/muzic · museformer/museformer/kernels/block_fill/main.py:67

def block_fill_cuda(out, block_ranges, head_masks, fill_value):
    num_blocks, num_heads = _check_input(block_ranges, head_masks, out)

    num_heads_2, seq_len_1, seq_len_2 = out.shape
    assert num_heads_2 == num_heads
    assert block_ranges.is_contiguous()
    assert head_masks.is_contiguous()
    assert out.is_contiguous()

    max_query = (block_ranges[:, 1] - block_ranges[:, 0]).max().item()
    max_key = (block_ranges[:, 3] - block_ranges[:, 2]).max().item()

    if max_query <= 0 or max_key <= 0:
        return out

    module = load_cuda_module()

    module.cuda_forward(block_ranges, head_masks, num_blocks, max_query, max_key, num_heads,

View source record →