File size: 1,430 Bytes
8e9f35a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """Frozen Hopper kernel recipe."""
import cutlass
from .fwd import SolAttnSplitForwardSm90
from .mainloop import SolAttnMainloopSm90
def make_kernel(tokens: int, kv_splits: int):
blocks = (tokens + 63) // 64
full_groups, tail = divmod(blocks, 64)
has_full_groups = tail == 0
has_full_blocks = tokens % 64 == 0
kernel = SolAttnMainloopSm90 if kv_splits == 1 else SolAttnSplitForwardSm90
return kernel(
cutlass.BFloat16,
head_dim=128,
head_dim_v=128,
qhead_per_kvhead=1,
is_causal=False,
is_local=False,
pack_gqa=False,
tile_m=64,
tile_n=64,
num_stages=1,
num_threads=128,
sol_attn_assume_lane_group_route_reduce=(
has_full_blocks and has_full_groups
),
sol_attn_assume_full_k_exact_blocks=has_full_blocks,
sol_attn_tail_exact_words1=0 < tail <= 8,
sol_attn_assume_full_route_groups=has_full_groups,
sol_attn_static_num_full_route_groups=(
-1 if has_full_groups else full_groups
),
sol_attn_static_tail_valid_count=(-1 if has_full_groups else tail),
sol_attn_tail_physical_tile16=0 < tail <= 16,
sol_attn_exact_mask_seqlen_last_only=(
not has_full_blocks
),
sol_attn_tail16_lane_group_route_reduce=tail == 16,
sol_attn_num_splits=kv_splits,
)
__all__ = ["make_kernel"]
|