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)