b4e055e9a9
EngineX replaces the missing corex_gdn/corex_moe/corex_fa2 operator chain
that Sub168 has but our BI-V100 image lacks.
Architecture (mirrors CCCL dispatch/tuning/kernel three-layer system):
Registry (policy_selector) → three-tier dispatch:
Tier 1: Native .so via dlopen (libcorex_gdn.so, libixattn.so)
Tier 2: ixformer Python ops (vendor-provided)
Tier 3: PyTorch fallback (always available)
Critical fixes vs comp 168 docker log:
- moe_topk_softmax: replacement for missing ixformer op
- gdn_prefill: NaN-stable chunked impl (chunk_size=16)
- gdn_decode: state clamp prevents NaN accumulation
18 operators, all tests pass.
64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
"""
|
|
EngineX cache operators.
|
|
|
|
KV cache management for paged attention.
|
|
CCCL parallel: dispatch_batch_memcpy (block copies between cache slots).
|
|
"""
|
|
|
|
from typing import Dict, List
|
|
|
|
import torch
|
|
|
|
|
|
def reshape_and_cache_pytorch(
|
|
key: torch.Tensor, # [num_tokens, num_kv_heads, head_size]
|
|
value: torch.Tensor, # [num_tokens, num_kv_heads, head_size]
|
|
key_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size]
|
|
value_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size]
|
|
slot_mapping: torch.Tensor, # [num_tokens] — maps token → (block, offset)
|
|
kv_cache_dtype: str = "auto",
|
|
k_scale: float = 1.0,
|
|
v_scale: float = 1.0,
|
|
) -> None:
|
|
"""Write new K,V into their assigned cache slots."""
|
|
num_tokens = key.shape[0]
|
|
block_size = key_cache.shape[2]
|
|
|
|
for i in range(num_tokens):
|
|
slot = slot_mapping[i].item()
|
|
if slot < 0:
|
|
continue
|
|
block_idx = slot // block_size
|
|
block_offset = slot % block_size
|
|
key_cache[block_idx, :, block_offset, :] = key[i] * k_scale
|
|
value_cache[block_idx, :, block_offset, :] = value[i] * v_scale
|
|
|
|
|
|
def copy_blocks_pytorch(
|
|
key_caches: List[torch.Tensor],
|
|
value_caches: List[torch.Tensor],
|
|
block_mapping: torch.Tensor, # [num_pairs, 2] src→dst
|
|
) -> None:
|
|
"""Copy cache blocks (used for fork/copy-on-write)."""
|
|
num_pairs = block_mapping.shape[0]
|
|
num_layers = len(key_caches)
|
|
|
|
for i in range(num_pairs):
|
|
src = block_mapping[i, 0].item()
|
|
dst = block_mapping[i, 1].item()
|
|
for layer in range(num_layers):
|
|
key_caches[layer][dst].copy_(key_caches[layer][src])
|
|
value_caches[layer][dst].copy_(value_caches[layer][src])
|
|
|
|
|
|
def swap_blocks_pytorch(
|
|
src: torch.Tensor,
|
|
dst: torch.Tensor,
|
|
block_mapping: torch.Tensor,
|
|
) -> None:
|
|
"""Swap cache blocks between GPU and CPU."""
|
|
for i in range(block_mapping.shape[0]):
|
|
src_idx = block_mapping[i, 0].item()
|
|
dst_idx = block_mapping[i, 1].item()
|
|
dst[dst_idx].copy_(src[src_idx])
|