Diff Coverage

        query = query.transpose(1, 2).reshape(-1, query.shape[1], query.shape[-1])
        key = key.transpose(1, 2).reshape(-1, key.shape[1], key.shape[-1])
        value = value.transpose(1, 2).reshape(-1, value.shape[1], value.shape[-1])
        if attention_mask is None and is_causal:
            npu_mask = _get_compressed_causal_mask(query.device)
            sparse_mode = 3
        else:
            npu_mask = None if attention_mask is None else _npu_attention_mask(attention_mask)
        return query, key, value, "TND", npu_mask, sparse_mode

) -> bool:
    """Gather valid QKV tokens directly into TND, retaining output indices."""
    attention_mask = context.attention_mask
    if attention_mask.shape[-1] != context.query_length:
        raise ValueError("the padding mask length must equal the query/key sequence length")
    attention_mask = attention_mask.to(device=context.query.device, dtype=torch.bool).expand(context.batch_size, -1)
    context.attention_mask = attention_mask
    if bool(attention_mask.all()):
        context.attention_mask = None