Add fp8 support to reshape_and_cache_flash (#6667)

2026-08-11 11:28:04 +08:00 · 2024-07-24 18:36:52 +00:00
parent ee812580f7
commit 0e63494cf3
8 changed files with 98 additions and 43 deletions
@@ -248,7 +248,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
      "                        Tensor! key_cache,"
      "                        Tensor! value_cache,"
      "                        Tensor slot_mapping,"
-      "                        str kv_cache_dtype) -> ()");
+      "                        str kv_cache_dtype,"
+      "                        float k_scale, float v_scale) -> ()");
  cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
                 &reshape_and_cache_flash);