diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom index 4f55c76bc658..d3f771237b71 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_decode_split_f32_f16_wmma.loom @@ -79,6 +79,7 @@ template.def<@ggml.flash_attention.decode_split.produce_partials.active> device %c0_f32 = scalar.constant 0.0 : f32 %negative_large = scalar.constant -1e+30 : f32 %c0_i32 = scalar.constant 0 : i32 + %masked_value_stage_bytes = index.constant 1024 : offset %c1_i32 = scalar.constant 1 : i32 %attention_scale = config.get @ggml.flash_attention.attention_scale : f32 %c0_f16 = scalar.constant 0.0 : f16 @@ -153,6 +154,13 @@ template.def<@ggml.flash_attention.decode_split.produce_partials.active> device %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer %tail_value_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %visibility_stage_bytes = index.constant 256 : offset + %visibility_stage = buffer.alloca align(16) %visibility_stage_bytes : buffer + %visibility_view = buffer.view %visibility_stage[%c0_offset] : buffer -> view<64xf32> + // Per-subgroup 16x32 V panels for key tiles that mix visible and masked + // keys; the score stage is free once probabilities are published. + %masked_value_stage_offset = index.scale %subgroup, %masked_value_stage_bytes : index, offset -> offset + %masked_value_stage_view = buffer.view %score_stage[%masked_value_stage_offset] : buffer -> view<16x32xf16> %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x[%query_stage_stride]xf16> %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> @@ -274,6 +282,19 @@ template.def<@ggml.flash_attention.decode_split.produce_partials.active> device view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + // Masked keys must reach P*V as V = +0, not just P = 0. The F16 WMMA + // result depended on the sign of V at keys whose probability is zero + // (measured: flipping the sign of one masked V row changed one output + // element by one F16 ulp). Rows past the sequence end hold whatever an + // earlier request left there, so identical requests gave different logits. + // Key tiles with every key visible read V directly, fully masked tiles + // use a zero V fragment, and mixed tiles stage V with masked rows zeroed. + %is_first_subgroup = index.cmp eq, %subgroup, %c0 : index + scf.if %is_first_subgroup { + view.store %lane_mask, %visibility_view[%lane] : f32, view<64xf32> + } + %lane_key_visible = scalar.cmpf ogt, %lane_mask, %negative_large : f32 + %lane_key_tile = index.div %lane, %c16 : index kernel.barrier scope(workgroup) ordering(acq_rel) scf.for %output_tile = [%c0 to %output_tile_count step %c1] unroll { %output_tile_channel = index.mul %output_tile, %c128 : index @@ -319,10 +340,59 @@ template.def<@ggml.flash_attention.decode_split.produce_partials.active> device %value_token, %value_token_end = index.assume %value_token0, %value_token_end0 [le(%value_token_end0, %bounded_key_value_token_count)] : index, index %full_value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%value_token_end]x[%value_width]xf16> %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment0 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel0] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> - %value_fragment1 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel1] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> - %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %key_tile_ordinal = index.div %key_tile, %c16 : index + %lane_in_key_tile = index.cmp eq, %lane_key_tile, %key_tile_ordinal : index + %lane_outside_key_tile = index.cmp ne, %lane_key_tile, %key_tile_ordinal : index + %lane_tile_visible = scalar.ori %lane_outside_key_tile, %lane_key_visible : i1 + %lane_tile_has_visible = scalar.andi %lane_in_key_tile, %lane_key_visible : i1 + %key_tile_all_visible = kernel.subgroup.vote.all %lane_tile_visible : i1 + %key_tile_any_visible = kernel.subgroup.vote.any %lane_tile_has_visible : i1 + %next_product_accumulator0, %next_product_accumulator1 = scf.if %key_tile_all_visible -> (vector<8xf16>, vector<8xf16>) { + %value_fragment0 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel0] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> + %value_fragment1 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel1] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> + %direct_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %direct_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %direct_accumulator0, %direct_accumulator1 : vector<8xf16>, vector<8xf16> + } else { + %masked_accumulator0, %masked_accumulator1 = scf.if %key_tile_any_visible -> (vector<8xf16>, vector<8xf16>) { + scf.for %stage_iteration = [%c0 to %c8 step %c1] unroll { + %stage_linear = index.madd %stage_iteration, %c64, %lane : index + %stage_key_row = index.div %stage_linear, %c32 : index + %stage_column = index.rem %stage_linear, %c32 : index + %stage_second_fragment = index.cmp uge, %stage_column, %c16 : index + %stage_fragment_column = index.rem %stage_column, %c16 : index + %stage_channel_base = scf.select %stage_second_fragment, %clamped_value_channel1, %clamped_value_channel0 : index + %stage_channel0 = index.add %stage_channel_base, %stage_fragment_column : index + %stage_channel = index.assume %stage_channel0 [lt(%stage_channel0, %value_width)] : index + %stage_block_key0 = index.add %key_tile, %stage_key_row : index + %stage_block_key = index.min %stage_block_key0, %c63 : index + %stage_key_mask = view.load %visibility_view[%stage_block_key] : view<64xf32> -> f32 + %stage_key_visible = scalar.cmpf ogt, %stage_key_mask, %negative_large : f32 + %stage_value = scf.if %stage_key_visible -> (f16) { + %stage_token0 = index.add %value_token, %stage_key_row : index + %stage_token = index.assume %stage_token0 [lt(%stage_token0, %value_token_end)] : index + %loaded = view.load %full_value_view[%stage_token, %stage_channel] : view<[%value_token_end]x[%value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %stage_value, %masked_value_stage_view[%stage_key_row, %stage_column] : f16, view<16x32xf16> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + %staged_value_fragment0 = vector.fragment.load %masked_value_stage_view[%c0, %c0] shape [%k, %n] : view<16x32xf16> -> vector<16xf16> + %staged_value_fragment1 = vector.fragment.load %masked_value_stage_view[%c0, %c16] shape [%k, %n] : view<16x32xf16> -> vector<16xf16> + %staged_accumulator0 = vector.mma %probability_fragment, %staged_value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %staged_accumulator1 = vector.mma %probability_fragment, %staged_value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.yield %staged_accumulator0, %staged_accumulator1 : vector<8xf16>, vector<8xf16> + } else { + %zero_value_fragment = vector.fragment %c0_f16x16 shape [%k, %n] : vector<16xf16> + %zero_accumulator0 = vector.mma %probability_fragment, %zero_value_fragment, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %zero_accumulator1 = vector.mma %probability_fragment, %zero_value_fragment, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %zero_accumulator0, %zero_accumulator1 : vector<8xf16>, vector<8xf16> + } + scf.yield %masked_accumulator0, %masked_accumulator1 : vector<8xf16>, vector<8xf16> + } scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> } scf.yield %full_product_fragment0, %full_product_fragment1 : vector<8xf16>, vector<8xf16> @@ -337,7 +407,12 @@ template.def<@ggml.flash_attention.decode_split.produce_partials.active> device %value_token = index.add %value_token0, %key_row : index %value_token_valid = index.cmp ult, %value_token, %bounded_key_value_token_count : index %value_channel_valid = index.cmp ult, %value_channel, %value_head_size : index - %value_valid = scalar.andi %value_token_valid, %value_channel_valid : i1 + %value_valid0 = scalar.andi %value_token_valid, %value_channel_valid : i1 + %block_value_key0 = index.add %key_tile, %key_row : index + %block_value_key = index.min %block_value_key0, %c63 : index + %value_key_mask = view.load %visibility_view[%block_value_key] : view<64xf32> -> f32 + %value_key_visible = scalar.cmpf ogt, %value_key_mask, %negative_large : f32 + %value_valid = scalar.andi %value_valid0, %value_key_visible : i1 %value_element = scf.if %value_valid -> (f16) { %global_value_channel = index.add %value_head_base, %value_channel : index %loaded = view.load %value_view[%value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%value_width]xf16> -> f16 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom index e5cb72adef84..d225ccac32f6 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention_f32_f16_wmma.loom @@ -240,6 +240,12 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %c0_f16x8 = vector.constant 0.0 : vector<8xf16> %c0_f32x4 = vector.constant 0.0 : vector<4xf32> %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> + %positive_f32x4 = vector.constant 1e+30 : vector<4xf32> + %positive_large = scalar.constant 1e+30 : f32 + %c63 = index.constant 63 : index + %staged_value_offset = index.constant 2048 : offset + %key_visibility_bytes = index.constant 1024 : offset + %hidden_summary_bytes = index.constant 16 : offset %m = index.constant 16 : index %n = index.constant 16 : index %k = index.constant 16 : index @@ -327,6 +333,15 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %score_stage = buffer.alloca align(16) %score_stage_bytes : buffer %probability_stage = buffer.alloca align(16) %probability_stage_bytes : buffer %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer + // Per key, the largest probability any query row of this tile gives it + // ([key][subgroup]), and the per-subgroup largest negated mask value. + %key_visibility_stage = buffer.alloca align(16) %key_visibility_bytes : buffer + %key_visibility_view = buffer.view %key_visibility_stage[%c0_offset] : buffer -> view<64x4xf32> + %hidden_summary_stage = buffer.alloca align(16) %hidden_summary_bytes : buffer + %hidden_summary_view = buffer.view %hidden_summary_stage[%c0_offset] : buffer -> view<4xf32> + // V rows staged for blocks with masked keys; shares the score stage with + // the 2 KiB product exchange, after score reads have retired. + %staged_value_view = buffer.view %score_stage[%staged_value_offset] : buffer -> view<16x64xf16> %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x20xf32> %mask_summary_view = buffer.view %score_stage[%c0_offset] : buffer -> view<4xf32> @@ -373,14 +388,22 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %effective_mask = vector.select %query_valid, %mask_summary_f32, %negative_f32x4 : vector<4xf32> %lane_mask_maximum = vector.reduce %effective_mask, %negative_large : vector<4xf32>, f32 %subgroup_mask_maximum = kernel.subgroup.reduce %lane_mask_maximum : f32 + %valid_row_mask = vector.select %query_valid, %mask_summary_f32, %positive_f32x4 : vector<4xf32> + %negated_valid_row_mask = vector.negf %valid_row_mask : vector<4xf32> + %lane_hidden_maximum = vector.reduce %negated_valid_row_mask, %negative_large : vector<4xf32>, f32 + %subgroup_hidden_maximum = kernel.subgroup.reduce %lane_hidden_maximum : f32 %is_subgroup_leader = index.cmp eq, %lane, %c0 : index scf.if %is_subgroup_leader { view.store %subgroup_mask_maximum, %mask_summary_view[%subgroup] : f32, view<4xf32> + view.store %subgroup_hidden_maximum, %hidden_summary_view[%subgroup] : f32, view<4xf32> } kernel.barrier scope(workgroup) ordering(acq_rel) %subgroup_mask_maxima = vector.load %mask_summary_view[%c0] : view<4xf32> -> vector<4xf32> %workgroup_mask_maximum = vector.reduce %subgroup_mask_maxima, %negative_large : vector<4xf32>, f32 %block_has_attention = scalar.cmpf ogt, %workgroup_mask_maximum, %negative_large : f32 + %subgroup_hidden_maxima = vector.load %hidden_summary_view[%c0] : view<4xf32> -> vector<4xf32> + %workgroup_hidden_maximum = vector.reduce %subgroup_hidden_maxima, %negative_large : vector<4xf32>, f32 + %block_fully_visible = scalar.cmpf olt, %workgroup_hidden_maximum, %positive_large : f32 %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 = scf.if %block_has_attention -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { // Four independent wave-level WMMAs produce a 16x64 score tile. %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index @@ -427,6 +450,8 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_layout> view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_layout> view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_layout> + %lane_probability_maximum = vector.reduce %probability, %c0_f32 : vector<4xf32>, f32 + view.store %lane_probability_maximum, %key_visibility_view[%lane, %subgroup] : f32, view<64x4xf32> kernel.barrier scope(workgroup) ordering(acq_rel) // Compute the 64-channel tiles in this output group sequentially. This // keeps one P*V accumulator live per wave and reuses a 2 KiB exchange @@ -453,13 +478,54 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %value_channel0 = index.add %value_head_base, %output_tile_channel : index %value_channel = index.add %value_channel0, %subgroup_product_channel : index %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> + // Masked keys must reach P*V as V = +0, not just P = 0. The F16 WMMA + // result depended on the sign of V at keys whose probability is zero + // (measured: flipping the sign of one masked V row changed one output + // element by one F16 ulp). Rows past the sequence end hold whatever an + // earlier request left there, so identical requests gave different + // logits. Blocks where some valid row masks some key stage V through + // LDS and zero the keys no row of this tile gives a probability. + %product_fragment = scf.if %block_fully_visible -> (vector<8xf16>) { + %direct_product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + scf.yield %direct_product_fragment : vector<8xf16> + } else { + %staged_product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) { + scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %staged_key_row = index.div %linear, %c64 : index + %staged_channel = index.rem %linear, %c64 : index + %staged_block_key0 = index.add %key_tile, %staged_key_row : index + %staged_block_key = index.min %staged_block_key0, %c63 : index + %staged_key_probabilities = vector.load %key_visibility_view[%staged_block_key, %c0] : view<64x4xf32> -> vector<4xf32> + %staged_key_probability = vector.reduce %staged_key_probabilities, %c0_f32 : vector<4xf32>, f32 + %staged_key_visible = scalar.cmpf ogt, %staged_key_probability, %c0_f32 : f32 + %staged_value = scf.if %staged_key_visible -> (f16) { + %staged_token0 = index.add %key_origin, %staged_block_key : index + %staged_token = index.assume %staged_token0 [lt(%staged_token0, %bounded_key_value_token_count)] : index + %staged_value_channel0 = index.add %value_channel0, %staged_channel : index + %staged_value_channel = index.assume %staged_value_channel0 [lt(%staged_value_channel0, %value_width)] : index + %loaded = view.load %value_view[%staged_token, %staged_value_channel] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %staged_value, %staged_value_view[%staged_key_row, %staged_channel] : f16, view<16x64xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %staged_value_view[%c0, %subgroup_product_channel] shape [%k, %n] : view<16x64xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_product_accumulator : vector<8xf16> + } + scf.yield %staged_product_fragment : vector<8xf16> } vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> kernel.barrier scope(workgroup) ordering(acq_rel) @@ -601,6 +667,8 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_layout> view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_layout> view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_layout> + %lane_probability_maximum = vector.reduce %probability, %c0_f32 : vector<4xf32>, f32 + view.store %lane_probability_maximum, %key_visibility_view[%lane, %subgroup] : f32, view<64x4xf32> kernel.barrier scope(workgroup) ordering(acq_rel) // Reuse the tail K scratch for V after probability publication. Each // workgroup stages only its output group, starting at local column zero. @@ -612,7 +680,12 @@ template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_pr %tail_value_channel = index.rem %linear, %tail_value_stage_channel_count : index %tail_value_row_valid = index.cmp ult, %tail_value_row, %tail_key_count : index %tail_value_channel_valid = index.cmp ult, %tail_value_channel, %group_output_channel_count : index - %tail_value_valid = scalar.andi %tail_value_row_valid, %tail_value_channel_valid : i1 + %tail_value_in_range = scalar.andi %tail_value_row_valid, %tail_value_channel_valid : i1 + %tail_value_block_key = index.min %tail_value_row, %c63 : index + %tail_key_probabilities = vector.load %key_visibility_view[%tail_value_block_key, %c0] : view<64x4xf32> -> vector<4xf32> + %tail_key_probability = vector.reduce %tail_key_probabilities, %c0_f32 : vector<4xf32>, f32 + %tail_key_visible = scalar.cmpf ogt, %tail_key_probability, %c0_f32 : f32 + %tail_value_valid = scalar.andi %tail_value_in_range, %tail_key_visible : i1 %output_value_channel = index.add %output_group_channel_base, %tail_value_channel : index %tail_value = scf.if %tail_value_valid -> (f16) { %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 76f6d624486c..77f7e6afc66e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -380,4 +380,9 @@ if (TARGET ggml-hrx) add_executable(test-hrx-moe-split test-hrx-moe-split.cpp) target_link_libraries(test-hrx-moe-split PRIVATE ggml) add_test(NAME test-hrx-moe-split COMMAND test-hrx-moe-split) + + add_executable(test-hrx-fa-masked-v test-hrx-fa-masked-v.cpp) + target_link_libraries(test-hrx-fa-masked-v PRIVATE ggml-hrx ggml hrx::hrx loomc::loomc) + target_include_directories(test-hrx-fa-masked-v PRIVATE ../ggml/include ../ggml/src ../ggml/src/ggml-hrx) + add_test(NAME test-hrx-fa-masked-v COMMAND test-hrx-fa-masked-v) endif() diff --git a/tests/test-hrx-fa-masked-v.cpp b/tests/test-hrx-fa-masked-v.cpp new file mode 100644 index 000000000000..96b7f4e24606 --- /dev/null +++ b/tests/test-hrx-fa-masked-v.cpp @@ -0,0 +1,255 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// FLASH_ATTN_EXT on HRX must not depend on K/V rows that the mask hides. A KV cache keeps whatever an earlier request +// (or a rejected draft) wrote past the current sequence end, so a dependence there makes identical requests give +// different logits. Each case runs the same attention several times on the HRX device, changing only the masked K/V +// rows (random values, +x, -x, +0, -0), and requires bitwise-identical outputs; one run is also checked against the +// CPU backend (normalized MSE). Cases cover the decode-split kernel (1..15 query rows) and the prefill kernel +// (16+ rows), with the sequence end inside a 64-key block and in the final partial block. + +#include "dispatch/dispatch-scheduler.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" +#include "graph/graph.h" +#include "kernel-corpus/kernel-corpus.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: requirement failed: %s\n", __FILE__, __LINE__, #condition); \ + std::abort(); \ + } \ + } while (false) + +static constexpr int64_t kHeadSize = 128; +static constexpr int64_t kQueryHeadCount = 16; +static constexpr int64_t kKeyValueHeadCount = 8; +static constexpr double kMaxNmse = 5e-4; + +static std::string kernel_name_for_id(uint64_t kernel_id) { + const ggml::hrx::KernelResolveResult resolved = + ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), "gfx1151", kernel_id); + REQUIRE(resolved.found()); + return ggml::hrx::kernel_definition_name(*resolved.definition); +} + +// The HRX kernels planned for the graph, or "" when HRX does not take it. +static std::string hrx_plan(ggml_cgraph * graph) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + if (!scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics) || !scheduler.plan().valid()) { + return ""; + } + std::string names; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + const std::string name = kernel_name_for_id(dispatch.kernel.kernel_id); + if (names.find(name) == std::string::npos) { + names += (names.empty() ? "" : "+") + name; + } + } + return names; +} + +struct Case { + int64_t queries; // query rows in the batch + int64_t visible; // cells 0 .. visible-1 hold the sequence; the last query row sees all of them + int64_t cells; // KV cells in the graph (the mask width) +}; + +// Graph like llama.cpp builds it: K/V cache views [head, cell, kv head], Q permuted to [head, token, head]. +struct Graph { + ggml_context * ctx = nullptr; + ggml_cgraph * graph = nullptr; + ggml_tensor * q = nullptr; + ggml_tensor * k = nullptr; + ggml_tensor * v = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * out = nullptr; + ggml_backend_buffer_t buffer = nullptr; +}; + +static Graph build(ggml_backend_t backend, const Case & c) { + Graph g; + ggml_init_params params = { 32 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + g.ctx = ggml_init(params); + REQUIRE(g.ctx != nullptr); + g.q = ggml_new_tensor_3d(g.ctx, GGML_TYPE_F32, kHeadSize, kQueryHeadCount, c.queries); + g.k = ggml_new_tensor_3d(g.ctx, GGML_TYPE_F16, kHeadSize, kKeyValueHeadCount, c.cells); + g.v = ggml_new_tensor_3d(g.ctx, GGML_TYPE_F16, kHeadSize, kKeyValueHeadCount, c.cells); + g.mask = ggml_new_tensor_2d(g.ctx, GGML_TYPE_F16, c.cells, c.queries); + ggml_tensor * k_view = ggml_view_3d(g.ctx, g.k, kHeadSize, c.cells, kKeyValueHeadCount, g.k->nb[2], g.k->nb[1], 0); + ggml_tensor * v_view = ggml_view_3d(g.ctx, g.v, kHeadSize, c.cells, kKeyValueHeadCount, g.v->nb[2], g.v->nb[1], 0); + ggml_tensor * q_permuted = ggml_permute(g.ctx, g.q, 0, 2, 1, 3); + ggml_tensor * attention = ggml_flash_attn_ext(g.ctx, q_permuted, k_view, v_view, g.mask, + 1.0f / std::sqrt(static_cast(kHeadSize)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(attention, GGML_PREC_F32); + g.out = ggml_reshape_2d(g.ctx, attention, kHeadSize * kQueryHeadCount, c.queries); + g.graph = ggml_new_graph(g.ctx); + ggml_build_forward_expand(g.graph, g.out); + g.buffer = ggml_backend_alloc_ctx_tensors(g.ctx, backend); + REQUIRE(g.buffer != nullptr); + return g; +} + +static void release(Graph & g) { + ggml_backend_buffer_free(g.buffer); + ggml_free(g.ctx); +} + +static std::vector compute(ggml_backend_t backend, Graph & g, const std::vector & q, + const std::vector & k, const std::vector & v, + const std::vector & mask) { + ggml_backend_tensor_set(g.q, q.data(), 0, ggml_nbytes(g.q)); + ggml_backend_tensor_set(g.k, k.data(), 0, ggml_nbytes(g.k)); + ggml_backend_tensor_set(g.v, v.data(), 0, ggml_nbytes(g.v)); + ggml_backend_tensor_set(g.mask, mask.data(), 0, ggml_nbytes(g.mask)); + REQUIRE(ggml_backend_graph_compute(backend, g.graph) == GGML_STATUS_SUCCESS); + std::vector result(static_cast(ggml_nelements(g.out))); + ggml_backend_tensor_get(g.out, result.data(), 0, result.size() * sizeof(float)); + return result; +} + +static double nmse(const std::vector & got, const std::vector & expected) { + double err = 0.0; + double ref = 0.0; + for (size_t i = 0; i < got.size(); ++i) { + const double diff = static_cast(got[i]) - expected[i]; + err += diff * diff; + ref += static_cast(expected[i]) * expected[i]; + } + return ref > 0.0 ? err / ref : err; +} + +int main() { + ggml_backend_dev_t device = ggml_backend_dev_by_name("HRX0"); + if (device == nullptr) { + ggml_backend_load_all(); + device = ggml_backend_dev_by_name("HRX0"); + } + if (device == nullptr) { + std::printf("test-hrx-fa-masked-v: no HRX0 device, skipped\n"); + return 0; + } + ggml_backend_t hrx = ggml_backend_dev_init(device, nullptr); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr && cpu != nullptr); + + std::vector cases; + for (int64_t seed_case = 0; seed_case < 8; ++seed_case) { + cases.push_back({ 1 + seed_case % 15, 6 + 7 * seed_case, 256 }); // decode split, end in block 0 + cases.push_back({ 1 + (3 * seed_case) % 15, 70 + 11 * seed_case, 256 }); // decode split, end in block 1 + cases.push_back({ 16 + 9 * seed_case, 30 + 13 * seed_case, 256 }); // prefill, end inside a block + } + cases.push_back({ 1, 150, 200 }); // decode split, partial last block + cases.push_back({ 24, 150, 200 }); // prefill, partial last block (tail path) + cases.push_back({ 40, 40, 512 }); // prefill, fresh prompt in a long cache + + const size_t kv_row = static_cast(kHeadSize * kKeyValueHeadCount); + int failures = 0; + for (size_t ci = 0; ci < cases.size(); ++ci) { + const Case & c = cases[ci]; + REQUIRE(c.visible >= c.queries && c.visible <= c.cells); + std::mt19937 rng(static_cast(1000 + ci)); + std::normal_distribution normal(0.0f, 1.0f); + std::vector q(static_cast(kHeadSize * kQueryHeadCount * c.queries)); + for (float & x : q) { + x = 2.0f * normal(rng); + } + std::vector k(kv_row * c.cells); + std::vector v(kv_row * c.cells); + for (size_t i = 0; i < k.size(); ++i) { + k[i] = ggml_fp32_to_fp16(normal(rng)); + v[i] = ggml_fp32_to_fp16(normal(rng)); + } + // Causal mask: query row r sits at position visible - queries + r. + std::vector mask(static_cast(c.cells * c.queries)); + for (int64_t r = 0; r < c.queries; ++r) { + const int64_t position = c.visible - c.queries + r; + for (int64_t cell = 0; cell < c.cells; ++cell) { + mask[r * c.cells + cell] = ggml_fp32_to_fp16(cell <= position ? 0.0f : -INFINITY); + } + } + + Graph g = build(hrx, c); + const std::string plan = ggml_backend_supports_op(hrx, g.out->src[0]) ? hrx_plan(g.graph) : ""; + if (plan.empty()) { + std::printf("FAIL case %zu (q=%lld visible=%lld cells=%lld): not planned on HRX\n", ci, + static_cast(c.queries), static_cast(c.visible), + static_cast(c.cells)); + ++failures; + release(g); + continue; + } + const std::vector reference = compute(hrx, g, q, k, v, mask); + + // Rewrite only the rows past the sequence end, five ways. + int mismatches = 0; + for (int variant = 0; variant < 5; ++variant) { + std::vector k2 = k; + std::vector v2 = v; + std::mt19937 stale_rng(static_cast(77 + variant)); + for (size_t i = kv_row * c.visible; i < k2.size(); ++i) { + float kv = 0.0f; + float vv = 0.0f; + switch (variant) { + case 0: kv = 3.0f * normal(stale_rng); vv = 3.0f * normal(stale_rng); break; + case 1: kv = 30.0f; vv = 1.0f; break; + case 2: kv = -30.0f; vv = -1.0f; break; + case 3: kv = 0.0f; vv = 0.0f; break; + default: kv = -0.0f; vv = -0.0f; break; + } + k2[i] = ggml_fp32_to_fp16(kv); + v2[i] = ggml_fp32_to_fp16(vv); + } + const std::vector got = compute(hrx, g, q, k2, v2, mask); + if (std::memcmp(got.data(), reference.data(), got.size() * sizeof(float)) != 0) { + ++mismatches; + } + } + release(g); + + Graph gc = build(cpu, c); + const std::vector expected = compute(cpu, gc, q, k, v, mask); + release(gc); + const double error = nmse(reference, expected); + const bool ok = mismatches == 0 && error <= kMaxNmse; + failures += ok ? 0 : 1; + std::printf("%s case %zu q=%lld visible=%lld cells=%lld (%s): masked-row variants differing %d/5, nmse %.3g\n", + ok ? "ok " : "FAIL", ci, static_cast(c.queries), static_cast(c.visible), + static_cast(c.cells), plan.c_str(), mismatches, error); + } + ggml_backend_free(hrx); + ggml_backend_free(cpu); + if (failures != 0) { + std::printf("test-hrx-fa-masked-v: %d of %zu cases failed\n", failures, cases.size()); + return 1; + } + std::printf("test-hrx-fa-masked-v: %zu cases passed\n", cases.size()); + return 0; +}