diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 12f61cddea18c..99d627b2f7851 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -831,6 +831,9 @@ class GQAAttentionBase { args.buffer = reinterpret_cast(flash_buffer_alloc); args.buffer_size_per_thread = buffer_size_per_thread; args.query = Q; + args.q_batch_stride = packed_qkv + ? static_cast(packed_batch_stride) + : static_cast(SafeInt(num_heads_) * sequence_length * head_size); args.k_cache = present_key_data; args.v_cache = present_value_data; args.k_scale = k_scale; @@ -874,7 +877,11 @@ class GQAAttentionBase { args.buffer_size_per_thread = buffer_size_per_thread; // Offset Q and output for this batch - args.query = Q + static_cast(b) * num_heads_ * sequence_length * head_size; + const ptrdiff_t q_batch_stride_elems = packed_batch_stride > 0 + ? packed_batch_stride + : static_cast(SafeInt(num_heads_) * sequence_length * head_size); + args.query = Q + static_cast(SafeInt(b) * static_cast(q_batch_stride_elems)); + args.q_batch_stride = static_cast(q_batch_stride_elems); args.k_cache = present_key_data + static_cast(b) * kv_num_heads_ * seqlen_present_kv_cache * packed_row_bytes; args.v_cache = present_value_data + @@ -884,10 +891,13 @@ class GQAAttentionBase { args.output = output->MutableData() + static_cast(b) * sequence_length * hidden_size; - // Slice attention bias for this batch (the kernel sees batch_size=1, so batch_idx=0 inside) + // Slice attention bias for this batch (the kernel sees batch_size=1, so batch_idx=0 inside). + // Bias shape is [batch|1, num_heads|1, S, T]; the batch stride uses the actual head + // extent (1 when the head dim is broadcast). const float* batch_bias = attention_bias_data; if (attention_bias_data != nullptr && !attention_bias_broadcast_batch) { - batch_bias += static_cast(b) * num_heads_ * sequence_length * attention_bias_seqlen_stride; + const size_t bias_head_extent = attention_bias_broadcast_head ? 1 : static_cast(num_heads_); + batch_bias += static_cast(SafeInt(b) * bias_head_extent * sequence_length * attention_bias_seqlen_stride); } args.attention_bias = batch_bias; args.attention_bias_seqlen_stride = attention_bias_seqlen_stride; diff --git a/onnxruntime/core/mlas/inc/mlas_qkv_quant.h b/onnxruntime/core/mlas/inc/mlas_qkv_quant.h index f6a5a48e6ccc7..29f794e90affd 100644 --- a/onnxruntime/core/mlas/inc/mlas_qkv_quant.h +++ b/onnxruntime/core/mlas/inc/mlas_qkv_quant.h @@ -249,46 +249,48 @@ MlasSVGemm( * It avoids materializing the full [S, T] attention probability matrix. */ struct MlasFlashAttentionQuantizedKVArgs { - int batch_size; - int num_heads; // Q heads - int kv_num_heads; // KV heads (for GQA sharing) - int sequence_length; // Q sequence length (new tokens) - int total_seqlen; // Total KV sequence length (past + new) - int head_size; - int past_seqlen; // For computing causal positions - int local_window_size; // -1 = disabled - int seqlen_present_kv; // Buffer dimension for present KV (may be > total_seqlen) - int q_block_size; // Br (query block size) - int kv_block_size; // Bc (KV block size) - float scale; // 1/sqrt(head_size) or user-specified + int batch_size = 0; + int num_heads = 0; // Q heads + int kv_num_heads = 0; // KV heads (for GQA sharing) + int sequence_length = 0; // Q sequence length (new tokens) + int total_seqlen = 0; // Total KV sequence length (past + new) + int head_size = 0; + int past_seqlen = 0; // For computing causal positions + int local_window_size = -1; // -1 = disabled + int seqlen_present_kv = 0; // Buffer dimension for present KV (may be > total_seqlen) + int q_block_size = 0; // Br (query block size) + int kv_block_size = 0; // Bc (KV block size) + float scale = 0.0f; // 1/sqrt(head_size) or user-specified - MLAS_KV_QUANT_TYPE quant_type; - bool per_channel_k; // Whether K uses per-channel scales - bool per_channel_v; // Whether V uses per-channel scales + MLAS_KV_QUANT_TYPE quant_type = MLAS_KV_QUANT_TYPE::S8_PerTensor; + bool per_channel_k = false; // Whether K uses per-channel scales + bool per_channel_v = false; // Whether V uses per-channel scales - int thread_count; - float* buffer; - size_t buffer_size_per_thread; + int thread_count = 1; + float* buffer = nullptr; + size_t buffer_size_per_thread = 0; - const float* query; // [B, N, S, H] FP32 - const uint8_t* k_cache; // [B, kv_N, seqlen_present, packed_row_bytes] quantized - const uint8_t* v_cache; // [B, kv_N, seqlen_present, packed_row_bytes] quantized - const float* k_scale; // Scalar or per-channel scales for K - const float* v_scale; // Scalar or per-channel scales for V - float* output; // [B, S, N, H] FP32 + const float* query = nullptr; // [B, N, S, H] FP32 + size_t q_batch_stride = 0; // element stride between consecutive batches in `query` + // (num_heads*S*H for unpacked, (num_heads+2*kv_num_heads)*S*H for packed QKV) + const uint8_t* k_cache = nullptr; // [B, kv_N, seqlen_present, packed_row_bytes] quantized + const uint8_t* v_cache = nullptr; // [B, kv_N, seqlen_present, packed_row_bytes] quantized + const float* k_scale = nullptr; // Scalar or per-channel scales for K + const float* v_scale = nullptr; // Scalar or per-channel scales for V + float* output = nullptr; // [B, S, N, H] FP32 // Attention bias (additive, applied after QK GEMM before masking/softmax). // Shape: [B|1, N|1, S, T] where dimensions of size 1 are broadcast. - const float* attention_bias; // nullptr if no bias - int attention_bias_seqlen_stride; // stride along the T (total_seqlen) dimension = shape[3] - bool attention_bias_broadcast_batch; // true if shape[0] == 1 - bool attention_bias_broadcast_head; // true if shape[1] == 1 + const float* attention_bias = nullptr; // nullptr if no bias + int attention_bias_seqlen_stride = 0; // stride along the T (total_seqlen) dimension = shape[3] + bool attention_bias_broadcast_batch = true; // true if shape[0] == 1 + bool attention_bias_broadcast_head = true; // true if shape[1] == 1 // Flash decoding fields (used when sequence_length == 1 and KV is split across threads). // Partials buffer stores per-(batch, head, kv_chunk) intermediate results: // [m_partial, l_partial, output_partial[head_size]] for each chunk. - float* flash_decoding_partials; // nullptr to disable flash decoding - int kv_chunk_count; // number of KV chunks = ceil(total_seqlen / kv_block_size) + float* flash_decoding_partials = nullptr; // nullptr to disable flash decoding + int kv_chunk_count = 0; // number of KV chunks = ceil(total_seqlen / kv_block_size) }; /** diff --git a/onnxruntime/core/mlas/lib/flashattn_qkv.cpp b/onnxruntime/core/mlas/lib/flashattn_qkv.cpp index 364011fe26e26..24fbfc4aab497 100644 --- a/onnxruntime/core/mlas/lib/flashattn_qkv.cpp +++ b/onnxruntime/core/mlas/lib/flashattn_qkv.cpp @@ -127,10 +127,12 @@ MlasFlashAttentionQuantizedKVThreaded( ? args->v_scale + kv_head_idx * static_cast(head_size) : args->v_scale; - // Q pointer: layout [batch, num_heads, seq, head_size] or packed + // Q pointer: layout [batch, num_heads, seq, head_size]. The batch stride is + // supplied separately (args->q_batch_stride) so the kernel works with both the + // standard BNSH layout and packed-QKV input where Q/K/V are interleaved per batch. const float* q_ptr = args->query + - (static_cast(batch_idx) * static_cast(num_heads) + - static_cast(head_idx)) * static_cast(sequence_length) * static_cast(head_size) + + static_cast(batch_idx) * args->q_batch_stride + + static_cast(head_idx) * static_cast(sequence_length) * static_cast(head_size) + static_cast(q_idx) * static_cast(head_size); // Iterate over KV blocks @@ -162,10 +164,14 @@ MlasFlashAttentionQuantizedKVThreaded( static_cast(args->attention_bias_seqlen_stride); const ptrdiff_t bias_matrix_size = static_cast(sequence_length) * bias_seqlen_stride; + // The bias tensor has shape [batch|1, num_heads|1, S, T]; the batch + // stride uses the actual head extent (1 when the head dim is broadcast). + const ptrdiff_t bias_head_extent = + args->attention_bias_broadcast_head ? 1 : static_cast(num_heads); ptrdiff_t bias_offset = 0; if (!args->attention_bias_broadcast_batch) { bias_offset += static_cast(batch_idx) * - static_cast(num_heads) * bias_matrix_size; + bias_head_extent * bias_matrix_size; } if (!args->attention_bias_broadcast_head) { bias_offset += static_cast(head_idx) * bias_matrix_size; @@ -378,10 +384,11 @@ MlasFlashDecodingQuantizedKVThreaded( ? args->v_scale + kv_head_idx * static_cast(head_size) : args->v_scale; - // Q pointer: layout [batch, num_heads, 1, head_size] (sequence_length=1) + // Q pointer: layout [batch, num_heads, 1, head_size] (sequence_length=1). + // The batch stride is supplied separately to support packed-QKV input. const float* q_ptr = args->query + - (static_cast(batch_idx) * static_cast(num_heads) + - static_cast(head_idx)) * static_cast(head_size); + static_cast(batch_idx) * args->q_batch_stride + + static_cast(head_idx) * static_cast(head_size); // Step 1: QK^T GEMM for this KV chunk const uint8_t* k_block = k_cache_head + static_cast(ir) * packed_row_bytes; @@ -405,10 +412,14 @@ MlasFlashDecodingQuantizedKVThreaded( const ptrdiff_t bias_seqlen_stride = static_cast(args->attention_bias_seqlen_stride); const ptrdiff_t bias_matrix_size = bias_seqlen_stride; // S=1 + // The bias tensor has shape [batch|1, num_heads|1, S, T]; the batch stride + // uses the actual head extent (1 when the head dim is broadcast). + const ptrdiff_t bias_head_extent = + args->attention_bias_broadcast_head ? 1 : static_cast(num_heads); ptrdiff_t bias_offset = 0; if (!args->attention_bias_broadcast_batch) { bias_offset += static_cast(batch_idx) * - static_cast(num_heads) * bias_matrix_size; + bias_head_extent * bias_matrix_size; } if (!args->attention_bias_broadcast_head) { bias_offset += static_cast(head_idx) * bias_matrix_size; diff --git a/onnxruntime/test/mlas/bench/bench_qkv_quant.cpp b/onnxruntime/test/mlas/bench/bench_qkv_quant.cpp index 23ca591ba6ed2..08945dbb5bf94 100644 --- a/onnxruntime/test/mlas/bench/bench_qkv_quant.cpp +++ b/onnxruntime/test/mlas/bench/bench_qkv_quant.cpp @@ -554,6 +554,7 @@ static void BM_GQA_Flash(benchmark::State& state) { args.buffer = buffer.data(); args.buffer_size_per_thread = buffer_size_per_thread; args.query = query.data(); + args.q_batch_stride = static_cast(num_heads) * seq_len * head_size; args.k_cache = k_cache.data(); args.v_cache = v_cache.data(); args.k_scale = k_scale.data(); diff --git a/onnxruntime/test/mlas/unittest/test_qkv_quant.cpp b/onnxruntime/test/mlas/unittest/test_qkv_quant.cpp index 5f0b18fa2cac8..d2474673d0892 100644 --- a/onnxruntime/test/mlas/unittest/test_qkv_quant.cpp +++ b/onnxruntime/test/mlas/unittest/test_qkv_quant.cpp @@ -493,6 +493,7 @@ class MlasFlashAttentionQuantizedKVTest : public MlasTestBase { args.buffer = flash_buffer; args.buffer_size_per_thread = buffer_size_per_thread; args.query = Q; + args.q_batch_stride = seq_len * head_size; args.k_cache = k_quant; args.v_cache = v_quant; args.k_scale = k_scale; @@ -592,6 +593,7 @@ class MlasFlashAttentionQuantizedKVTest : public MlasTestBase { args.buffer = flash_buffer; args.buffer_size_per_thread = buffer_size_per_thread; args.query = Q; + args.q_batch_stride = head_size; args.k_cache = k_quant; args.v_cache = v_quant; args.k_scale = k_scale_buf; diff --git a/onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py b/onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py index 4a4d3e6ff43e8..4481b071c80e0 100644 --- a/onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py +++ b/onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py @@ -164,6 +164,7 @@ def create_quantized_gqa_graph( bit_width, buffer_seq_len=None, is_past=False, + packed_qkv=False, ): """Create an ONNX graph for GroupQueryAttention with quantized KV cache.""" if buffer_seq_len is None: @@ -171,6 +172,7 @@ def create_quantized_gqa_graph( hidden_size = num_heads * head_size kv_hidden_size = kv_num_heads * head_size + query_hidden_size = (num_heads + 2 * kv_num_heads) * head_size if packed_qkv else hidden_size packed_head_size = head_size // 2 if bit_width == 4 else head_size cache_ort_type = TensorProto.UINT8 if bit_width == 4 else TensorProto.INT8 @@ -186,8 +188,8 @@ def create_quantized_gqa_graph( # Inputs inputs = [ "query", - "key", - "value", + "" if packed_qkv else "key", + "" if packed_qkv else "value", "past_key", "past_value", "seqlens_k", @@ -220,20 +222,29 @@ def create_quantized_gqa_graph( # Graph inputs graph_input = [ - helper.make_tensor_value_info("query", TensorProto.FLOAT, [batch_size, seq_len, hidden_size]), - helper.make_tensor_value_info("key", TensorProto.FLOAT, [batch_size, seq_len, kv_hidden_size]), - helper.make_tensor_value_info("value", TensorProto.FLOAT, [batch_size, seq_len, kv_hidden_size]), - helper.make_tensor_value_info( - "past_key", cache_ort_type, [batch_size, kv_num_heads, past_kv_seqlen, packed_head_size] - ), - helper.make_tensor_value_info( - "past_value", cache_ort_type, [batch_size, kv_num_heads, past_kv_seqlen, packed_head_size] - ), - helper.make_tensor_value_info("seqlens_k", TensorProto.INT32, [batch_size]), - helper.make_tensor_value_info("total_sequence_length", TensorProto.INT32, [1]), - helper.make_tensor_value_info("k_scale", TensorProto.FLOAT, None), - helper.make_tensor_value_info("v_scale", TensorProto.FLOAT, None), + helper.make_tensor_value_info("query", TensorProto.FLOAT, [batch_size, seq_len, query_hidden_size]), ] + if not packed_qkv: + graph_input.extend( + [ + helper.make_tensor_value_info("key", TensorProto.FLOAT, [batch_size, seq_len, kv_hidden_size]), + helper.make_tensor_value_info("value", TensorProto.FLOAT, [batch_size, seq_len, kv_hidden_size]), + ] + ) + graph_input.extend( + [ + helper.make_tensor_value_info( + "past_key", cache_ort_type, [batch_size, kv_num_heads, past_kv_seqlen, packed_head_size] + ), + helper.make_tensor_value_info( + "past_value", cache_ort_type, [batch_size, kv_num_heads, past_kv_seqlen, packed_head_size] + ), + helper.make_tensor_value_info("seqlens_k", TensorProto.INT32, [batch_size]), + helper.make_tensor_value_info("total_sequence_length", TensorProto.INT32, [1]), + helper.make_tensor_value_info("k_scale", TensorProto.FLOAT, None), + helper.make_tensor_value_info("v_scale", TensorProto.FLOAT, None), + ] + ) # Graph outputs graph_output = [ @@ -468,6 +479,110 @@ def run_quantized_gqa_prompt_test( ) +def run_quantized_gqa_packed_qkv_test( + batch_size, seq_len, num_heads, kv_num_heads, head_size, quant_type, bit_width, atol=None +): + """Run a packed-QKV quantized GQA prompt test and compare against FP32 reference with quantization noise.""" + np.random.seed(43) + + hidden_size = num_heads * head_size + kv_hidden_size = kv_num_heads * head_size + + query = np.random.uniform(-0.5, 0.5, (batch_size, seq_len, hidden_size)).astype(np.float32) + key_input = np.random.uniform(-0.5, 0.5, (batch_size, seq_len, kv_hidden_size)).astype(np.float32) + value_input = np.random.uniform(-0.5, 0.5, (batch_size, seq_len, kv_hidden_size)).astype(np.float32) + packed_qkv = np.concatenate([query, key_input, value_input], axis=2) + + k_bnsh = key_input.reshape(batch_size, seq_len, kv_num_heads, head_size).transpose(0, 2, 1, 3) + v_bnsh = value_input.reshape(batch_size, seq_len, kv_num_heads, head_size).transpose(0, 2, 1, 3) + + if bit_width == 8: + if quant_type == "PER_TENSOR": + _, k_scale = quantize_int8_per_tensor(k_bnsh) + _, v_scale = quantize_int8_per_tensor(v_bnsh) + else: + _, k_scale = quantize_int8_per_channel(k_bnsh) + _, v_scale = quantize_int8_per_channel(v_bnsh) + else: + if quant_type == "PER_TENSOR": + _, k_scale = quantize_int4_per_tensor(k_bnsh) + _, v_scale = quantize_int4_per_tensor(v_bnsh) + else: + _, k_scale = quantize_int4_per_channel(k_bnsh) + _, v_scale = quantize_int4_per_channel(v_bnsh) + + packed_head_size = head_size // 2 if bit_width == 4 else head_size + if bit_width == 4: + past_k = np.zeros((batch_size, kv_num_heads, seq_len, packed_head_size), dtype=np.uint8) + past_v = np.zeros((batch_size, kv_num_heads, seq_len, packed_head_size), dtype=np.uint8) + else: + past_k = np.zeros((batch_size, kv_num_heads, seq_len, packed_head_size), dtype=np.int8) + past_v = np.zeros((batch_size, kv_num_heads, seq_len, packed_head_size), dtype=np.int8) + + seqlens_k = np.array([seq_len - 1] * batch_size, dtype=np.int32) + total_seq = np.array([seq_len], dtype=np.int32) + + onnx_model_str = create_quantized_gqa_graph( + batch_size, seq_len, num_heads, kv_num_heads, head_size, quant_type, bit_width, packed_qkv=True + ) + sess_options = SessionOptions() + sess = InferenceSession(onnx_model_str, sess_options, providers=["CPUExecutionProvider"]) + + feeds = { + "query": packed_qkv, + "past_key": past_k, + "past_value": past_v, + "seqlens_k": seqlens_k, + "total_sequence_length": total_seq, + "k_scale": k_scale, + "v_scale": v_scale, + } + + outputs = sess.run(None, feeds) + out_ort = outputs[0] + + if bit_width == 8 and quant_type == "PER_TENSOR": + k_q = np.clip(np.round(k_bnsh / k_scale[0]), -128, 127).astype(np.int8) + v_q = np.clip(np.round(v_bnsh / v_scale[0]), -128, 127).astype(np.int8) + k_deq = dequantize_int8_per_tensor(k_q, k_scale[0]) + v_deq = dequantize_int8_per_tensor(v_q, v_scale[0]) + elif bit_width == 8 and quant_type == "PER_CHANNEL": + k_q = np.clip(np.round(k_bnsh / k_scale.reshape(1, kv_num_heads, 1, head_size)), -128, 127).astype(np.int8) + v_q = np.clip(np.round(v_bnsh / v_scale.reshape(1, kv_num_heads, 1, head_size)), -128, 127).astype(np.int8) + k_deq = dequantize_int8_per_channel(k_q, k_scale, kv_num_heads, head_size) + v_deq = dequantize_int8_per_channel(v_q, v_scale, kv_num_heads, head_size) + elif bit_width == 4 and quant_type == "PER_TENSOR": + k_q = np.clip(np.round(k_bnsh / k_scale[0]), -8, 7).astype(np.int8) + v_q = np.clip(np.round(v_bnsh / v_scale[0]), -8, 7).astype(np.int8) + k_deq = k_q.astype(np.float32) * k_scale[0] + v_deq = v_q.astype(np.float32) * v_scale[0] + elif bit_width == 4 and quant_type == "PER_CHANNEL": + k_q = np.clip(np.round(k_bnsh / k_scale.reshape(1, kv_num_heads, 1, head_size)), -8, 7).astype(np.int8) + v_q = np.clip(np.round(v_bnsh / v_scale.reshape(1, kv_num_heads, 1, head_size)), -8, 7).astype(np.int8) + k_deq = k_q.astype(np.float32) * k_scale.reshape(1, kv_num_heads, 1, head_size) + v_deq = v_q.astype(np.float32) * v_scale.reshape(1, kv_num_heads, 1, head_size) + else: + raise ValueError(f"Unsupported config: bit_width={bit_width}, quant_type={quant_type}") + + out_ref = reference_gqa(query, k_deq, v_deq, num_heads, kv_num_heads, head_size, causal=True) + + if atol is None: + atol = 0.15 if bit_width == 4 else 0.05 + + if np.any(np.isnan(out_ort)): + raise AssertionError(f"NaN in output (quant={quant_type}, bit={bit_width}, packed QKV)") + if np.allclose(out_ort, 0.0): + raise AssertionError(f"Output is all zeros (quant={quant_type}, bit={bit_width}, packed QKV)") + + np.testing.assert_allclose( + out_ort, + out_ref, + atol=atol, + rtol=0.1, + err_msg=f"Packed-QKV quantized GQA output mismatch (quant={quant_type}, bit={bit_width})", + ) + + # ---- Test class ---- @@ -529,6 +644,17 @@ def test_int8_multi_batch(self): bit_width=8, ) + def test_int8_packed_qkv_multi_batch(self): + run_quantized_gqa_packed_qkv_test( + batch_size=3, + seq_len=8, + num_heads=4, + kv_num_heads=2, + head_size=16, + quant_type="PER_TENSOR", + bit_width=8, + ) + def test_int4_multi_batch(self): run_quantized_gqa_prompt_test( batch_size=2, @@ -830,6 +956,24 @@ def test_int8_bias_broadcast_head(self): bias_broadcast_head=True, ) + def test_int8_bias_broadcast_head_multi_batch(self): + """Bias shape [B, 1, S, T] with batch_size > 1 and num_heads > 1. + + Regression test: the bias batch stride must use the head extent (1 when the + head dimension is broadcast), not num_heads. With batch_size == 1 the bug is + masked because batch_idx is always 0. + """ + run_quantized_gqa_bias_test( + batch_size=3, + seq_len=8, + num_heads=4, + kv_num_heads=2, + head_size=16, + quant_type="PER_TENSOR", + bit_width=8, + bias_broadcast_head=True, + ) + def test_int8_bias_broadcast_both(self): """Bias shape [1, 1, S, T] with batch_size > 1 and num_heads > 1.""" run_quantized_gqa_bias_test(