Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -989,6 +989,274 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s)
s->use_transposed_mask_once = true;
break;

case 59:
s->head_size = 64;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 4;
s->has_waves_per_eu = true;
s->waves_per_eu = 2;
s->use_v_double_buffer = true;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_mfma32_skip_legacy_qreg = true;
s->use_transposed_mask_limit = true;
s->use_fast_paged_kv_desc = true;
break;

case 60:
s->head_size = 64;
s->block_size = 32;
s->num_query_heads = 32;
s->num_kv_heads = 32;
s->dtype = "fp16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 4;
s->has_waves_per_eu = true;
s->waves_per_eu = 2;
s->use_v_double_buffer = true;
s->use_staggered_iter_wait = true;
s->has_tile_size = true;
s->tile_size = 128;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_mfma32_skip_legacy_qreg = true;
s->use_transposed_mask_limit = true;
break;

case 61:
s->head_size = 64;
s->block_size = 32;
s->num_query_heads = 8;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 2;
s->use_v_double_buffer = true;
break;

case 62:
s->head_size = 64;
s->block_size = 32;
s->num_query_heads = 8;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 2;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_v_double_buffer = true;
s->use_staggered_iter_wait = true;
break;

case 63:
s->head_size = 128;
s->block_size = 32;
s->num_query_heads = 8;
s->num_kv_heads = 8;
s->dtype = "fp16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 1;
s->use_v_double_buffer = true;
s->use_early_v_schedule = true;
s->use_sched_barrier = true;
break;

case 64:
s->head_size = 256;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 2;
s->num_warps = 2;
s->has_waves_per_eu = true;
s->waves_per_eu = 2;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_kq_lds_pad = true;
s->kq_lds_pad_halves = 16;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_transposed_mask_limit = true;
s->use_k_single_buffer = true;
s->use_q_direct_reg = true;
s->use_softmax_mfma_interleave = true;
s->softmax_interleave_mode = 2;
s->use_mask_phase_split = true;
break;

case 65:
s->head_size = 256;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 2;
s->num_warps = 2;
s->has_waves_per_eu = true;
s->waves_per_eu = 2;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_kq_lds_pad = true;
s->kq_lds_pad_halves = 16;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_transposed_mask_limit = true;
s->use_k_single_buffer = true;
s->use_q_direct_reg = true;
s->use_softmax_mfma_interleave = true;
s->softmax_interleave_mode = 2;
s->use_mask_phase_split = true;
s->use_transposed_invariant_hoist = true;
break;

case 66:
s->head_size = 128;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 2;
s->num_warps = 4;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_transposed_mask_limit = true;
s->use_q_direct_reg = true;
break;

case 67:
s->head_size = 128;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "fp16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 2;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_transposed_mask_limit = true;
s->use_k_single_buffer = true;
s->use_kq_lds_pad = true;
s->use_mask_phase_split = true;
s->use_softmax_mfma_interleave = true;
s->softmax_interleave_mode = 2;
s->softmax_interleave_groups = 2;
break;

case 68:
s->head_size = 256;
s->block_size = 16;
s->num_query_heads = 32;
s->num_kv_heads = 32;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 1;
s->num_warps = 2;
s->has_waves_per_eu = true;
s->waves_per_eu = 2;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_kq_lds_pad = true;
s->kq_lds_pad_halves = 16;
s->use_transposed_scalar_state = true;
s->use_transposed_mask_once = true;
s->use_transposed_half_local_pv = true;
s->use_transposed_mask_limit = true;
s->use_k_single_buffer = true;
s->use_q_direct_reg = true;
s->use_softmax_mfma_interleave = true;
s->softmax_interleave_mode = 2;
s->use_mask_phase_split = true;
break;

case 69:
s->head_size = 64;
s->block_size = 32;
s->num_query_heads = 64;
s->num_kv_heads = 8;
s->dtype = "bf16";
s->use_sinks = false;
s->sliding_window = 0;
s->has_softcap = false;
s->num_seqs = 2;
s->num_warps = 4;
s->has_tile_size = true;
s->tile_size = 64;
s->block_m_per_warp = 32;
s->use_mfma_32x32 = true;
s->use_transposed_qk_32x32 = true;
s->use_transposed_scalar_state = true;
s->use_transposed_half_local_pv = true;
s->kv_storage_dtype = "fp8e4m3";
s->use_fp8_mfma_qk = true;
break;

default:
return -1;
}
Expand Down
Loading
Loading