From f453b7ebede13711f2f0f4c85e52593d45fd0571 Mon Sep 17 00:00:00 2001 From: yurekami Date: Fri, 26 Dec 2025 01:45:18 +0900 Subject: [PATCH] fix: correct dtype in quantize_k_cache docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docstring incorrectly stated the return dtype as uint8_t, but the actual implementation returns torch.float8_e4m3fn. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- tests/quant.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/quant.py b/tests/quant.py index 06247599..dc9a06d9 100644 --- a/tests/quant.py +++ b/tests/quant.py @@ -7,7 +7,7 @@ def quantize_k_cache( ) -> torch.Tensor: """ Quantize the k-cache - Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype uint8_t, where t = input_k_cache.element_size() + Return a tensor with shape (num_blocks, block_size, h_k, dv + 4(dv/tile_size) + t(d-dv)) of dtype float8_e4m3fn, where t = input_k_cache.element_size() For more detail about the layout of K/V, please refer to comments in flash_mla_interface.py or README.md """ assert dv % tile_size == 0