During multi-sequence inference with unified KV cache, performance degraded significantly due to the overhead of masking out tokens from other sequences, causing slowdowns on large batch workloads.

The unified KV cache required filtering tokens per sequence via masking, which introduced expensive per-token conditional checks and prevented efficient memory access patterns in the attention kernel, especially in the CUDA backend.

Introduced a new flash attention extension called flash_attn_ext_rows that allows queries to share KV rows and masks compactly via an indirect lookup structure (kv_rows). This enables sending a more efficient representation of relevant tokens per sequence, reducing filtering overhead while maintaining unified KV benefits. The implementation involved changes across CUDA kernels and metal backends to support this optimized data structure.

When dealing with multi-sequence attention caching, avoid per-token masking overhead by restructuring data to enable indirect indexing and compact representation of relevant tokens, which can unlock significant performance gains on GPU backends.