[Question]: Can I add more supported head sizes for attention_v1?

Open
#759 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Stale
Tech stack
cpp
Domain
ai, hpc

Research direction

Start in aiter/csrc/kernels/attention_v1.cu and inspect the CALL_CUSTOM_LAUNCHER_BLK_HEAD macro and its existing head_size cases. Determine whether the attention_v1 kernel supports head size 256 and identify the relevant validation path; done means support is verified or the limitation is clearly documented.

Written by the indexing model from the issue text.

Description

Problem Description

In this file aiter/csrc/kernels/attention_v1.cu, can I add a case of 256 for the head_size?

#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE)                       \
    switch(head_size)                                                         \
    {                                                                         \
    case 64: CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64); break;           \
    case 128: CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128); break;         \
    default: TORCH_CHECK(false, "Unsupported head size: ", head_size); break; \
    }
Operating System

22.04.5 LTS (Jammy Jellyfish)

CPU

Intel(R) Xeon(R) Platinum 8468V

GPU

AMD Instinct MI300X

ROCm Version

ROCM 6.3.1

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

Dominant language
Python
Stars
565
Forks
585
Avg merge
3d 4h
Merged PRs (30d)
366

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from ROCm/aiter

All issues in ROCm/aiter

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.