import pytest import torch from lmdeploy.messages import QuantPolicy # Import common TurboQuant utilities from turboquant_utils from .turboquant_utils import ( _div_up, ) class TestFlattenKVCache: @pytest.fixture def out_dtype(self): yield torch.float16 @pytest.fixture def num_heads(self): yield 4 @pytest.fixture def head_dim(self): yield 32 @pytest.fixture def block_size(self): yield 16 @pytest.fixture def kv_lens(self): yield [2, 24, 47, 48] @pytest.fixture def batch_size(self, kv_lens): yield len(kv_lens) @pytest.fixture def num_blocks_per_input(self, kv_lens, block_size): yield [_div_up(kv_len, block_size) for kv_len in kv_lens] @pytest.fixture def max_num_blocks(self, num_blocks_per_input): yield max(num_blocks_per_input) @pytest.fixture def out_size(self, kv_lens): yield sum(kv_lens) @pytest.fixture def kv_seqlens(self, kv_lens): yield torch.tensor(kv_lens).cuda() @pytest.fixture def k_caches(self, batch_size, max_num_blocks, block_size, num_heads, head_dim, out_dtype): shape = (batch_size * max_num_blocks, block_size, num_heads, head_dim) yield torch.rand(shape, dtype=out_dtype, device='cuda') @pytest.fixture def v_caches(self, k_caches): yield torch.rand_like(k_caches) @pytest.fixture def block_offsets(self, num_blocks_per_input): batch_size = len(num_blocks_per_input) max_num_blocks = max(num_blocks_per_input) batch_ids = torch.arange(batch_size) ret = torch.arange(max_num_blocks) ret = batch_ids[:, None] + ret[None, :] * batch_size yield ret.cuda() @pytest.fixture def gt(self, k_caches, v_caches, kv_lens, block_offsets, block_size, num_heads, out_size, head_dim): k_states = k_caches.new_empty(num_heads, out_size, head_dim) v_states = v_caches.new_empty(num_heads, out_size, head_dim) start_loc = 0 for kv_len, block_offs in zip(kv_lens, block_offsets): remain_len = kv_len for idx, _ in enumerate(range(0, kv_len, block_size)): b_off = block_offs[idx] block_len = min(block_size, remain_len) end_loc = start_loc + block_len k_block = k_caches[b_off, :block_len] v_block = v_caches[b_off, :block_len] k_states[:, start_loc:end_loc] = k_block.transpose(0, 1) v_states[:, start_loc:end_loc] = v_block.transpose(0, 1) start_loc = end_loc remain_len -= block_len yield k_states, v_states def test_flatten_kv_cache(self, k_caches, v_caches, kv_seqlens, block_offsets, out_size, gt): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache k_states, v_states = flatten_kv_cache(k_caches, v_caches, kv_seqlens, block_offsets, out_size=out_size) torch.testing.assert_close(k_states, gt[0]) torch.testing.assert_close(v_states, gt[1]) def precise_round(x: torch.Tensor): return x.sign() * (x.abs() + 0.5).floor() def quant(kv: torch.Tensor, nbits: int = 8): """Quant kv on the head_dim.""" amax = kv.amax(dim=-1, keepdim=True) amin = kv.amin(dim=-1, keepdim=True) scales = (amax - amin) / (2**nbits - 1) zeros = -amin / scales q_kv = (kv / scales + zeros + 0.5).to(torch.uint8) if nbits == 4: q_kv1, q_kv2 = q_kv.split(q_kv.shape[-1] // 2, -1) q_kv = q_kv1 + q_kv2 * 16 return q_kv, torch.cat([scales, zeros], dim=-1) class TestFlattenKVCacheQuant8(TestFlattenKVCache): @pytest.fixture def nbits(self): yield 8 @pytest.fixture def atol(self): yield 4e-3 @pytest.fixture def rtol(self): yield 1e-5 @pytest.fixture def k_quant(self, k_caches, nbits): yield quant(k_caches, nbits) @pytest.fixture def v_quant(self, v_caches, nbits): yield quant(v_caches, nbits) def test_flatten_kv_cache(self, k_quant, v_quant, kv_seqlens, block_offsets, out_size, out_dtype, nbits, gt, atol, rtol): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache k_caches, k_sz = k_quant v_caches, v_sz = v_quant k_sz = k_sz.to(out_dtype) v_sz = v_sz.to(out_dtype) k_states, v_states = flatten_kv_cache(k_caches, v_caches, kv_seqlens, block_offsets, out_size=out_size, out_dtype=out_dtype, k_scales_zeros=k_sz, v_scales_zeros=v_sz, quant_policy=nbits) torch.testing.assert_close(k_states, gt[0], atol=atol, rtol=rtol) torch.testing.assert_close(v_states, gt[1], atol=atol, rtol=rtol) class TestFlattenKVCacheQuant4(TestFlattenKVCacheQuant8): @pytest.fixture def nbits(self): yield 4 @pytest.fixture def atol(self): yield 0.05 @pytest.fixture def rtol(self): yield 1e-3 @pytest.mark.skipif(torch.cuda.get_device_capability()[0] < 9, reason='require device with cc>=9.0') class TestFlattenKVCacheMLAFP8(TestFlattenKVCache): @pytest.fixture def out_dtype(self): yield torch.bfloat16 @pytest.fixture def num_heads(self): yield 1 @pytest.fixture def head_dim(self): yield 576 @pytest.fixture def block_size(self): yield 64 @pytest.fixture def k_cache_mla(self, k_caches): from lmdeploy.pytorch.kernels.cuda.blocked_gemm_fp8 import quant_fp8 num_blocks, block_size, num_heads, _ = k_caches.shape k_cache_pe = k_caches[:, :, :, 512:] k_cache_nope = k_caches[:, :, :, :512].flatten(0, -2) k_cache_nope, k_cache_scale = quant_fp8(k_cache_nope, group_size=128) k_cache_nope = k_cache_nope.view(num_blocks, block_size, num_heads, -1) k_cache_scale = k_cache_scale.reshape(num_blocks, block_size, num_heads, -1).to(torch.float32) dtype = k_cache_nope.dtype out = torch.cat([k_cache_nope, k_cache_scale.view(dtype), k_cache_pe.view(dtype)], dim=-1) yield out def _dequant(self, k_cache_mla): k_cache_nope = k_cache_mla[..., :512].to(torch.float32) k_cache_scale = k_cache_mla[..., 512:512 + 16].view(torch.float32) k_cache_pe = k_cache_mla[..., 512 + 16:].view(torch.bfloat16) k_cache_nope = k_cache_nope.unflatten(-1, (-1, 128)) k_cache_scale = k_cache_scale[..., None] k_cache_nope *= k_cache_scale k_cache_nope = k_cache_nope.flatten(-2, -1).to(k_cache_pe.dtype) k_cache = torch.cat([k_cache_nope, k_cache_pe], dim=-1) return k_cache @pytest.fixture def gt(self, k_cache_mla, kv_lens, block_offsets, block_size, num_heads, out_size, head_dim): k_caches = self._dequant(k_cache_mla) k_states = k_caches.new_empty(num_heads, out_size, head_dim) start_loc = 0 for kv_len, block_offs in zip(kv_lens, block_offsets): remain_len = kv_len for idx, _ in enumerate(range(0, kv_len, block_size)): b_off = block_offs[idx] block_len = min(block_size, remain_len) end_loc = start_loc + block_len k_block = k_caches[b_off, :block_len] k_states[:, start_loc:end_loc] = k_block.transpose(0, 1) start_loc = end_loc remain_len -= block_len yield k_states def test_flatten_kv_cache(self, k_cache_mla, kv_seqlens, block_offsets, out_size, out_dtype, gt): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache_mla_fp8 k_states = flatten_kv_cache_mla_fp8(k_cache_mla, kv_seqlens, block_offsets, out_size=out_size, out_dtype=out_dtype) torch.testing.assert_close(k_states, gt) # ============================================================================= # Tests for quant_policy=QuantPolicy.TURBO_QUANT (TurboQuant) flatten_kv_cache # ============================================================================= class TestFlattenKVCacheQuant42: """Test flatten_kv_cache with quant_policy=QuantPolicy.TURBO_QUANT (TurboQuant). quant_policy=QuantPolicy.TURBO_QUANT uses: - K: QJL4 (3bit MSE + 1bit QJL), stored in rotate domain - V: TurboQuant MSE int2, stored in rotate domain The flatten function should output rotate-domain KV that can be used directly for attention computation in the rotate domain. """ @pytest.fixture def num_heads(self): yield 4 @pytest.fixture def head_dim(self): yield 64 @pytest.fixture def head_dim_v(self): yield 64 @pytest.fixture def block_size(self): yield 16 @pytest.fixture def kv_lens(self): yield [8, 24, 48, 32] @pytest.fixture def batch_size(self, kv_lens): yield len(kv_lens) @pytest.fixture def num_blocks_per_input(self, kv_lens, block_size): yield [(kv_len + block_size - 1) // block_size for kv_len in kv_lens] @pytest.fixture def max_num_blocks(self, num_blocks_per_input): yield max(num_blocks_per_input) @pytest.fixture def out_size(self, kv_lens): yield sum(kv_lens) @pytest.fixture def kv_seqlens(self, kv_lens): yield torch.tensor(kv_lens).cuda() @pytest.fixture def packed_k_dim(self, head_dim): yield head_dim // 2 @pytest.fixture def packed_v_dim(self, head_dim_v): yield head_dim_v // 4 @pytest.fixture def k_caches(self, batch_size, max_num_blocks, block_size, num_heads, packed_k_dim): """Create quantized K cache (uint8). Note: The cache size is based on max_num_blocks, but the actual data is only kv_lens long. The flatten function should only output the actual data length. """ shape = (batch_size * max_num_blocks, block_size, num_heads, packed_k_dim) yield torch.randint(0, 256, shape, dtype=torch.uint8, device='cuda') @pytest.fixture def v_caches(self, batch_size, max_num_blocks, block_size, num_heads, packed_v_dim): """Create quantized V cache (uint8).""" shape = (batch_size * max_num_blocks, block_size, num_heads, packed_v_dim) yield torch.randint(0, 256, shape, dtype=torch.uint8, device='cuda') @pytest.fixture def k_scales_zeros(self, batch_size, max_num_blocks, block_size, num_heads): """K meta: [mse_norm, qjl_norm] for each position.""" shape = (batch_size * max_num_blocks, block_size, num_heads, 2) yield torch.rand(shape, dtype=torch.float16, device='cuda') @pytest.fixture def v_scales_zeros(self, batch_size, max_num_blocks, block_size, num_heads): """V meta: [norm] for each position.""" shape = (batch_size * max_num_blocks, block_size, num_heads, 1) yield torch.rand(shape, dtype=torch.float16, device='cuda') @pytest.fixture def block_offsets(self, num_blocks_per_input): batch_size = len(num_blocks_per_input) max_num_blocks = max(num_blocks_per_input) batch_ids = torch.arange(batch_size) ret = torch.arange(max_num_blocks) ret = batch_ids[:, None] + ret[None, :] * batch_size yield ret.cuda() @pytest.fixture def out_dtype(self): yield torch.float32 def test_flatten_kv_cache_quant42(self, k_caches, v_caches, kv_seqlens, block_offsets, k_scales_zeros, v_scales_zeros, out_dtype, head_dim, head_dim_v, num_heads): """Test flatten_kv_cache with quant_policy=QuantPolicy.TURBO_QUANT. This test verifies that: 1. The flatten function runs without error 2. Output shape is correct 3. Output is in the rotate domain (verified by dequantizing) """ from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache from lmdeploy.pytorch.kernels.cuda.turbo_quant import ( hadamard_rotate_inv, ) # Run flatten with quant_policy=QuantPolicy.TURBO_QUANT k_states, v_states = flatten_kv_cache( k_caches, v_caches, kv_seqlens, block_offsets, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, quant_policy=QuantPolicy.TURBO_QUANT, kv_layout='bshd', flatten_kv_layout='shd', out_dtype=out_dtype, ) # Get actual output size (may differ from expected due to cache padding) actual_out_size = k_states.shape[0] # Verify output shapes - use actual size from flatten output assert k_states.shape == (actual_out_size, num_heads, head_dim), f'K shape mismatch: {k_states.shape}' assert v_states.shape == (actual_out_size, num_heads, head_dim_v), f'V shape mismatch: {v_states.shape}' # Verify output is in rotate domain by checking that inverse rotation # produces reasonable values (not all zeros or NaNs) k_orig = hadamard_rotate_inv(k_states.float()) v_orig = hadamard_rotate_inv(v_states.float()) # Check that inverse rotation produces non-zero values assert k_orig.abs().max() > 1e-6, 'K inverse rotation produced all zeros' assert v_orig.abs().max() > 1e-6, 'V inverse rotation produced all zeros' print(f'flatten_kv_cache quant42: K shape={k_states.shape}, V shape={v_states.shape}') print(f' K rotate domain: mean={k_states.abs().mean():.4f}, max={k_states.abs().max():.4f}') print(f' V rotate domain: mean={v_states.abs().mean():.4f}, max={v_states.abs().max():.4f}') print(f' K orig domain: mean={k_orig.abs().mean():.4f}, max={k_orig.abs().max():.4f}') print(f' V orig domain: mean={v_orig.abs().mean():.4f}, max={v_orig.abs().max():.4f}')