feat(enginex): CCCL-style algorithm factor replacement engine — 18 operator dispatch system
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.
This commit is contained in:
63
enginex/ops/cache.py
Normal file
63
enginex/ops/cache.py
Normal file
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
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])
|
||||
Reference in New Issue
Block a user