Hugging Face's logo

Kernels:
flashrt
/
fa2-seqused-runtime

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
55a186d
·
verified ·
1 Parent(s): db6aa8a

Uploaded using `kernel-builder`.

Browse files
Files changed (4) hide show
  1. build/torch211-cxx11-cu128-x86_64-linux/__init__.py +7 -1
  2. build/torch211-cxx11-cu128-x86_64-linux/{_fa2_seqused_runtime_cuda_99d26a1.abi3.so → _fa2_seqused_runtime_cuda_00cce1c.abi3.so} +2 -2
  3. build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
  4. build/torch211-cxx11-cu128-x86_64-linux/metadata.json +4 -4
build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED
@@ -13,7 +13,9 @@ from ._ops import add_op_namespace_prefix, ops
13
 
14
  SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
15
  COMPILED_HEAD_DIM_BUCKETS = (64, 96, 128, 256)
16
- SPLIT_HEAD_DIMS = SUPPORTED_HEAD_DIMS
 
 
17
 
18
 
19
  @dataclass(frozen=True)
@@ -45,6 +47,8 @@ def recommended_num_splits(
45
  raise ValueError("all shape values and num_sms must be positive")
46
  if int(head_dim) not in SUPPORTED_HEAD_DIMS:
47
  raise ValueError("head_dim must be a positive multiple of 8 at most 256")
 
 
48
  block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
49
  n_blocks = _ceildiv(seqlen_k, block_n)
50
  m_blocks = _ceildiv(seqlen_q, 64)
@@ -85,6 +89,8 @@ def allocate_workspace(
85
  raise ValueError("q and k must have shape (B, S, H, D)")
86
  if q.shape[-1] not in SUPPORTED_HEAD_DIMS:
87
  raise ValueError("head_dim must be a positive multiple of 8 at most 256")
 
 
88
  if num_sms is None:
89
  num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
90
  splits = recommended_num_splits(
 
13
 
14
  SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
15
  COMPILED_HEAD_DIM_BUCKETS = (64, 96, 128, 256)
16
+ # Vendored FA2 split-KV partial-head tiles are valid only in these logical
17
+ # ranges. Other supported dimensions use the correct no-split kernel.
18
+ SPLIT_HEAD_DIMS = tuple(range(40, 129, 8)) + tuple(range(232, 257, 8))
19
 
20
 
21
  @dataclass(frozen=True)
 
47
  raise ValueError("all shape values and num_sms must be positive")
48
  if int(head_dim) not in SUPPORTED_HEAD_DIMS:
49
  raise ValueError("head_dim must be a positive multiple of 8 at most 256")
50
+ if int(head_dim) not in SPLIT_HEAD_DIMS:
51
+ return 1
52
  block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
53
  n_blocks = _ceildiv(seqlen_k, block_n)
54
  m_blocks = _ceildiv(seqlen_q, 64)
 
89
  raise ValueError("q and k must have shape (B, S, H, D)")
90
  if q.shape[-1] not in SUPPORTED_HEAD_DIMS:
91
  raise ValueError("head_dim must be a positive multiple of 8 at most 256")
92
+ if q.shape[-1] not in SPLIT_HEAD_DIMS:
93
+ return None
94
  if num_sms is None:
95
  num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
96
  splits = recommended_num_splits(
build/torch211-cxx11-cu128-x86_64-linux/{_fa2_seqused_runtime_cuda_99d26a1.abi3.so → _fa2_seqused_runtime_cuda_00cce1c.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:824df1266b7f0c7859f9132a97e3f034a9d8b1e434a340386d21a5b533986977
3
- size 422846752
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f70d0474e811f655ed17c01408b97590485edf886f6f5075ce61af4376a0f6f1
3
+ size 422846784
build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fa2_seqused_runtime_cuda_99d26a1
3
- ops = torch.ops._fa2_seqused_runtime_cuda_99d26a1
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fa2_seqused_runtime_cuda_99d26a1::{op_name}"
 
1
  import torch
2
+ from . import _fa2_seqused_runtime_cuda_00cce1c
3
+ ops = torch.ops._fa2_seqused_runtime_cuda_00cce1c
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fa2_seqused_runtime_cuda_00cce1c::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fa2-seqused-runtime",
3
- "id": "_fa2_seqused_runtime_cuda_99d26a1",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
@@ -16,9 +16,9 @@
16
  "digest": {
17
  "algorithm": "sha256",
18
  "files": {
19
- "__init__.py": "zaIQDw3dhB5xbCQPrZ9tGWr/Z559+ooyZOaGIa2/+LQ=",
20
- "_fa2_seqused_runtime_cuda_99d26a1.abi3.so": "gk3xJmt/DHhZ+RMql+PwNKnYseQ0o0A4bSGltTOYaXc=",
21
- "_ops.py": "3zmyZsIXRYN9PWXNasLSKultcD9S0IaBbtFEFMbUth0=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }
 
1
  {
2
  "name": "fa2-seqused-runtime",
3
+ "id": "_fa2_seqused_runtime_cuda_00cce1c",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
 
16
  "digest": {
17
  "algorithm": "sha256",
18
  "files": {
19
+ "__init__.py": "CN+Ar/JeWCSd1O49S4rR3XPtrtlRXwoY54NuRMsiFDM=",
20
+ "_fa2_seqused_runtime_cuda_00cce1c.abi3.so": "9w0EdOgR9lXtF8AUCLl1kEhe34hvb1B1zmGvQ3ag9vE=",
21
+ "_ops.py": "ERvskvnU6x9hhnDBgzGtOv1CM0j5hy1NqqMrW3nZ1D8=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }