Shape of slash mismatch when input batchsize > 1

[polarispw]

Describe the bug

When evaluating Llama3-8B with Minfer on lm-eval's zero-shot tasks, it went the following error:

Then I tried to explore the shape inside vertical_and_slash_kernel(q, k, v, vertical_size, slash_size)

Seems that the batch dim is always 1 in slash?

Steps to reproduce

import lm_eval
from lm_eval.models.huggingface import HFLM
from lm_eval.evaluator import request_caching_arg_to_dict
from lm_eval.tasks import TaskManager
from lm_eval.utils import make_table

from minference import MInference
minference_patch = MInference("minference", args.model)
model = minference_patch(model)

hflm = HFLM(pretrained=model, tokenizer=tokenizer, batch_size=args.zeroshot_bsz)

        manager = TaskManager()
        task_names = manager.match_tasks(args.tasks)
        logging.info(f"Selected Tasks: {task_names}")

        cache_args = request_caching_arg_to_dict(cache_requests=args.cache_processed_data)
        results = lm_eval.simple_evaluate(hflm, tasks=task_names, num_fewshot=0, batch_size=args.zeroshot_bsz, **cache_args)
        logging.info("\n"+make_table(results))
        if "groups" in results:
            logging.info(make_table(results, "groups"))

Above is my script, but I think circumstances with batch > 1 will work

Something about batch was mentioned in ##46.

Expected Behavior

Logs

No response

Additional Information

No response