| |
| |
|
|
| import torch |
| import triton |
| from ._triton_kernels.kv_cache import _cat_and_cache_mla_kernel |
| from ._aiter_compat.torch_guard import torch_compile_guard |
| from .utils.logger import AiterTritonLogger |
| from .utils.types import e4m3_dtype |
|
|
| _LOGGER = AiterTritonLogger() |
|
|
|
|
| def cat_and_cache_mla_fake_tensor( |
| k_nope: torch.Tensor, |
| k_pe: torch.Tensor, |
| kv_cache: torch.Tensor, |
| slot_mapping: torch.Tensor, |
| k_scale: torch.Tensor, |
| apply_scale: bool = True, |
| shuffled_kv_cache: bool = False, |
| ) -> None: |
| return None |
|
|
|
|
| @torch_compile_guard(gen_fake=cat_and_cache_mla_fake_tensor) |
| def cat_and_cache_mla( |
| k_nope: torch.Tensor, |
| k_pe: torch.Tensor, |
| kv_cache: torch.Tensor, |
| slot_mapping: torch.Tensor, |
| k_scale: torch.Tensor, |
| apply_scale: bool = True, |
| shuffled_kv_cache: bool = False, |
| ) -> None: |
| """ |
| Perform concat k_nope and k_pe to kv_cache inplace |
| |
| Key parameters: |
| - k_nope: Matrix X with shape (B_slot, KH, D1). |
| - k_pe: Matrix W with shape (B_slot, KH, D2). |
| - kv_cache: Matrix W with shape (B_cache, KH, D1 + D2). |
| - slot_mapping: Matrix W with shape (B_slot, ). |
| |
| B is the number of decode tokens, B_slot is the number of prefill + decode tokens, B_cahce is the max number of tokens of kv_cache |
| QH must be multiple of KH |
| |
| Returns: |
| - kv_cache: The output matrix with shape (B_max, KH, D1 + D2) (inplace). |
| """ |
| _LOGGER.info( |
| f"CAT_AND_CACHE_MLA: k_nope={tuple(k_nope.shape)} k_pe={tuple(k_pe.shape)} " |
| + f"kv_cache={tuple(kv_cache.shape)} slot_mapping={tuple(slot_mapping.shape)}" |
| ) |
|
|
| b, kh, d_nope = k_nope.shape |
| bk, kh2, d_rope = k_pe.shape |
| kv_cache_dtype = kv_cache.dtype |
| assert kv_cache_dtype in [ |
| torch.bfloat16, |
| e4m3_dtype, |
| torch.uint8, |
| ], "KV cache dtype must be BF16, FP8 or packed FP4" |
|
|
| block_size = 1 |
| SCALE_K_WIDTH_NOPE = 4 |
| SCALE_K_WIDTH_ROPE = 4 |
| if kv_cache_dtype == torch.uint8: |
| assert shuffled_kv_cache, "shuffle_kv_cache must be True for FP4 KV cache" |
| b_cache, h_cache, block_size, d_cache = kv_cache.shape |
| SCALE_K_LORA = d_nope // 16 |
| SCALE_K_ROPE = d_rope // 16 |
| SCALE_K_WIDTH_NOPE = ( |
| min(16, triton.next_power_of_2(SCALE_K_LORA)) |
| if SCALE_K_LORA >= 4 |
| else SCALE_K_LORA |
| ) |
| SCALE_K_WIDTH_ROPE = ( |
| min(16, triton.next_power_of_2(SCALE_K_ROPE)) |
| if SCALE_K_ROPE >= 4 |
| else SCALE_K_ROPE |
| ) |
| else: |
| if shuffled_kv_cache: |
| b_cache, h_cache, block_size, d_cache = kv_cache.shape |
| else: |
| b_cache, h_cache, d_cache = kv_cache.shape |
| (b_slot,) = slot_mapping.shape |
|
|
| assert ( |
| b == bk and b_slot == b_slot |
| ), "K batch dimensions and slot_mapping should be identical (bk == bk == b_slot)" |
| assert kh == kh2 == h_cache, "K head should be identical" |
| if kv_cache.dtype == torch.uint8: |
| assert ( |
| (d_nope + d_rope) // 2 + (d_nope + d_rope) // 16 |
| ) == d_cache, "The D dimension of kv_cache should be (d_nope + d_rope) // 2 + (d_nope + d_rope) // 16 for FP4 KV cache" |
| else: |
| assert ( |
| d_nope + d_rope == d_cache |
| ), "D dimension of k_nope and k_pe should be summed up to be the D dimension of kv_cache" |
| if isinstance(k_scale, torch.Tensor): |
| assert k_scale.numel() == 1, "k_scale should be a single-element torch.Tensor" |
|
|
| if shuffled_kv_cache: |
| kv_cache_stride_b = kv_cache.stride(0) |
| kv_cache_stride_h = kv_cache.stride(1) |
| kv_cache_stride_d = kv_cache.stride(3) |
| else: |
| kv_cache_stride_b = kv_cache.stride(0) |
| kv_cache_stride_h = kv_cache.stride(1) |
| kv_cache_stride_d = kv_cache.stride(2) |
|
|
| assert ( |
| kv_cache_stride_d == 1 |
| ), "The stride of the last dimension of KV cache must be 1" |
|
|
| _cat_and_cache_mla_kernel[(b * kh,)]( |
| k_nope, |
| k_pe, |
| kv_cache, |
| slot_mapping, |
| *k_nope.stride(), |
| *k_pe.stride(), |
| kv_cache_stride_b, |
| kv_cache_stride_h, |
| kv_cache_stride_d, |
| k_scale_ptr=k_scale, |
| KH=kh, |
| BLOCK_D_nope=d_nope, |
| BLOCK_D_pe=d_rope, |
| BLOCK_SIZE=block_size, |
| SHUFFLED_KV_CACHE=shuffled_kv_cache, |
| SCALE_K_WIDTH_NOPE=SCALE_K_WIDTH_NOPE, |
| SCALE_K_WIDTH_ROPE=SCALE_K_WIDTH_ROPE, |
| HAVE_K_SCALE=(k_scale is not None and apply_scale), |
| num_warps=1, |
| ) |
|
|
