From 39f2e7d97e4ecf2192d94de7b94431df273b66d5 Mon Sep 17 00:00:00 2001 From: zhshgmail Date: Thu, 23 Apr 2026 17:04:15 -0700 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=2024=5FKvCacheUpdateWithRope?= =?UTF-8?q?Backward:=20cache=5Fposition=20=E8=B6=8A=E7=95=8C=20(L2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 与 PR #84 同类问题 —— benchmark Model.get_input_groups() 的默认 int64 tensor 生成用 torch.randint(0, 10000, ...), cache_position 语义上必须 < max_seq (= grad_key_cache.shape[2])。max_seq ∈ {16..2048}，远小于 10000，导致**所有 50 个 case 100% OOB**。 症状：reference Model.forward() 本身在 NPU 上跑 grad_key_cache[:, :, cache_position] fancy indexing 即崩，错误码 507035 / AIV 334 subErrType 0x4 "BIU data to VEC incorrect"，跨 56 核 同一 PC。第一次 OOB 越界访存打坏 NPU stream 后，后续所有 torch.npu. synchronize 都返回 sticky 507035，即便输入合法也无法恢复（需重启进程）。 故一个 50-case 跑下来 ref 只有 1/50 能通过（case 0 碰运气 cp.max 落在 [0, 128) 内）。 修复方法 (与 PR #84 L1 op#24 同构): (1) .json 每个 case 的 cache_position 加 `"range": [0, max_seq - 1]` 其中 max_seq = grad_key_cache.shape[2] (2) model.py 的 int 分支读取 `inp.get('range', [0, default_max - 1])` 并 fallback 到原行为。注意 default fallback 用 `default_max - 1` 以精确匹配原 `torch.randint(0, default_max)` 的 [0, default_max-1] 值域，向后兼容无 range 字段的所有现有 case。 影响范围: - 本 PR 只改 op#24 L2 (其 .py 是 custom 模板，与默认模板不共享)。 - 扫描发现 L2 有 22 个 op 用默认模板，其中少数有 int tensor (如 10_SwigluQuant)。它们的 .py 未改动，行为完全不变。 - 若以后推广到默认模板，fallback `default_max - 1` 保证值域不变。 CPU 侧验证: python3 loads patched model.py + .json 跑 get_input_groups(), 50/50 case 的 cache_position 现在都在 [0, max_seq-1]，OOB 计数从 50/50 降到 0/50。 NPU 侧 full verification pending: 共享服务器 /dev/mapper/ts-root 100% 满，a5ops_dev3 container 退出， NPU 跑 Model.forward 全 50 case 的验证留到磁盘恢复后。但 CPU 侧数据 合法性 + PR #84 同构证据已足够证明此 PR 正确。 Refs: PR #84 (L1 op#24 EmbeddingDenseBackward 同类问题修复) --- .../24_KvCacheUpdateWithRopeBackward.json | 100 +++++++++--------- .../24_KvCacheUpdateWithRopeBackward.py | 11 +- 2 files changed, 59 insertions(+), 52 deletions(-) diff --git a/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.json b/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.json index 63ccf8e1..2a37ce34 100644 --- a/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.json +++ b/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.json @@ -1,50 +1,50 @@ -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 32, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 32, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 32, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 8, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 1024, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 1024, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 128, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 128, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 128, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [128]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 32, 16]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 32, 16]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 4, 16]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 4, 16]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 4, 16]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 128, 256]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 128, 256]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 8, 256]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 2048, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 2048, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 256, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 256, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 256, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [256]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 12, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 12, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 12, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [12]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 512, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 512, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 1, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 128, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 128, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 16, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 16, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 16, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 2, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 1024, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 1024, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 48, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 48, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 48, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [48]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 24, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 24, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 24, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 4, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 128, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 128, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 256]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 256]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 8, 256]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 16, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 16, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 64, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 64, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 128, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 128, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 128, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [128]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 4, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 12, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [12]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 32, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 32, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 2, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 1024, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 1024, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 64, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 64, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 4, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 24, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 24, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 24, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 32, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 32, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 2, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 2, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 2, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 48, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 48, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 48, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [48]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 24, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 24, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 24, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 16, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16]}]} -{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 32, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 32, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 32, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 32, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [4, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 8, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 1024, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 1024, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 8, 128, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 128, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [8, 128, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [128], "range": [0, 1023]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 32, 16]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 32, 16]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 1, 4, 16]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 4, 16]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 4, 16]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4], "range": [0, 31]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 12, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 128, 256]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 128, 256]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 8, 256]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 2048, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 2048, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 256, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 256, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 256, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [256], "range": [0, 2047]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 6, 12, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 12, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [3, 12, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [12], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 512, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 512, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 1, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 24, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 128, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 128, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 16, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 16, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 16, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 2, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 1024, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 1024, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 8, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64], "range": [0, 1023]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 48, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 48, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 48, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [48], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 8, 24, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 24, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [6, 24, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 4, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4], "range": [0, 31]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 128, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 128, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 16, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [4, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 256]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 256]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 8, 256]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 256]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 4, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 16, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 16, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 2, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1], "range": [0, 15]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 8, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [8, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 24, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 64, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 64, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 6, 8, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [2, 8, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 128, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 128, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 128, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [128], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 8, 4, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 12, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 12, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [12], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 32, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 32, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 2, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 2, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2], "range": [0, 31]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 1024, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 1024, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 4, 64, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 64, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 64, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64], "range": [0, 1023]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 32]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 128, 32]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 4, 32]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 32]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [4], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 8, 24, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 24, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [3, 24, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 32, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 6, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 32, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 32, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 8, 2, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 2, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 2, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [2], "range": [0, 31]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 16, 48, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 48, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [2, 48, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [48], "range": [0, 511]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 2, 16, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [1, 16, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 256, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 256, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 4, 32, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 32, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [4, 32, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [32], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 12, 24, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 24, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 24, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [24], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 64, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 8, 1, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 1, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [1], "range": [0, 63]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 128, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 128, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 8, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "bfloat16", "shape": [2, 8, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [8], "range": [0, 127]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 256, 128]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 256, 128]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 4, 16, 128]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 128]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float32", "shape": [1, 16, 128]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [16], "range": [0, 255]}]} +{"inputs": [{"name": "grad_key_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 512, 64]}, {"name": "grad_value_cache", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 512, 64]}, {"name": "key_states", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 16, 64, 64]}, {"name": "cos", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 64, 64]}, {"name": "sin", "type": "tensor", "required": true, "dtype": "float16", "shape": [1, 64, 64]}, {"name": "cache_position", "type": "tensor", "required": true, "dtype": "int64", "shape": [64], "range": [0, 511]}]} diff --git a/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.py b/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.py index 024bdd02..f9c01b24 100644 --- a/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.py +++ b/benchmarks/NPUKernelBench/level2/24_KvCacheUpdateWithRopeBackward.py @@ -99,9 +99,16 @@ def get_input_groups(): elif dtype_str == 'bool': tensors[name] = (torch.rand(shape) > 0.5).to(torch.bool) elif dtype_str in ('int32', 'int64', 'int8'): - max_val = {'int32': 1000, 'int64': 10000, 'int8': 127}.get(dtype_str, 100) + default_max = {'int32': 1000, 'int64': 10000, 'int8': 127}.get(dtype_str, 100) + # Allow per-case `range: [lo, hi]` (INCLUSIVE on both ends) to bound + # int tensor semantically. For indices into another tensor axis (e.g. + # cache_position < max_seq), the un-bounded default would produce OOB + # on NPU fancy-indexing → deterministic crash (AIV 334 BIU-to-VEC). + # Default matches original behavior: `torch.randint(0, default_max)` + # = inclusive `[0, default_max - 1]`. + rng = inp.get('range', [0, default_max - 1]) dtype = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, 'int32': torch.int32, 'int64': torch.int64, 'int8': torch.int8, 'bool': torch.bool}[dtype_str] - tensors[name] = torch.randint(0, max_val, shape, dtype=dtype) + tensors[name] = torch.randint(rng[0], rng[1] + 1, shape, dtype=dtype) else: dtype = {'float32': torch.float32, 'float16': torch.float16, 'bfloat16': torch.bfloat16, 'int32': torch.int32, 'int64': torch.int64, 'int8': torch.int8, 'bool': torch.bool}.get(dtype_str, torch.float32) tensors[name] = torch.randn(shape, dtype=dtype)