Add seqlens_k bounds validation in GroupQueryAttention to prevent GEMM OOB

Merged
VR
6mo ago
+342-10

Description

Validate seqlens_k tensor values in the CPU GroupQueryAttention operator before they are used as GEMM dimensions. Without this check, a crafted model can supply negative or oversized seqlens_k values that cause out-of-bounds reads in the K/V present cache buffers.

Changes

  • group_query_attention.cc: Add validation loop in Compute() before any seqlens_k access:
    • seqlens_k[b] >= 0 (prevents unsigned wraparound in static_cast<size_t>)
    • seqlens_k[b] + 1 <= present_kv_seqlen (prevents GEMM reading past K/V buffer)
    • For non-first-prompt: seqlens_k[b] + 1 >= sequence_length (prevents underflow in past_seqlen = total_seqlen - sequence_length)
  • group_query_attention_helper.h: Fix seqlens_k shape validation (&& to ||) so wrong-length tensors are correctly rejected
  • Tests: 4 regression tests covering negative, oversized, multi-batch, and boundary-valid seqlens_k values

Motivation and Context

MSRC case 108962: A crafted model can set seqlens_k values that, when cast to size_t and used as GEMM N dimension, cause heap OOB reads from the present K/V cache buffers.