def qkv_indexer_norm_rope(
query: torch.Tensor,
key: torch.Tensor,
weights: torch.Tensor,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
k_bias: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
eps: float,
zero_centered: bool,
native_rope: bool = False,
*,
main_inputs,
decode_cache=None,
decode_ring=None,
decode_z: torch.Tensor | None = None,
decode_slots: torch.Tensor | None = None,
) -> tuple[
torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]:
"""Prepare indexer/attention tensors and write their separate cache states."""
(
main_qkv,
main_qw,
main_kw,
main_rope,
main_slots,
main_kc,
main_vc,
main_qh,
main_kh,
main_rotary,
main_eps,
main_bias,
main_k_scale,
main_v_scale,
) = main_inputs
fp8_cache = main_kc.dtype in (torch.uint8, torch.float8_e4m3fn)
if fp8_cache:
main_kc = main_kc.view(torch.float8_e4m3fn)
main_vc = main_vc.view(torch.float8_e4m3fn)
tokens, width = query.shape
main_q = main_qkv.new_empty((tokens, main_qh * 192))
main_k = main_qkv.new_empty((tokens, main_kh * 192))
main_v = torch.empty_like(main_k)
main_group = 4 if main_rotary > 128 else 8
if (decode_cache is None) != (decode_ring is None):
raise ValueError("CSA fusion needs both the summary cache and the ring")
if decode_cache is not None:
if decode_z is None or decode_slots is None:
raise ValueError("CSA fusion needs decode z and slots")
elif decode_z is not None or decode_slots is not None:
raise ValueError("CSA decode inputs need a summary cache")
if decode_slots is not None and decode_slots.numel() > tokens:
raise ValueError("CSA decode rows must fit the indexer token count")
# One indexer group per local KV head; outputs are group-major.
groups = main_kh
if width % (groups * 256):
raise ValueError("Indexer query heads must divide evenly into KV groups")
heads = width // (groups * 256)
padded_heads = max(heads, 8)
q = torch.empty(
(groups, tokens, padded_heads, 256),
device=query.device,
dtype=torch.float8_e4m3fn,
)
k = torch.empty((tokens, 1, 256), device=key.device, dtype=key.dtype)
w = torch.empty(
(groups, tokens, padded_heads), device=weights.device, dtype=torch.float32
)
if tokens:
indexer_programs = groups * padded_heads + 1
main_programs = triton.cdiv(main_qh + 2 * main_kh, main_group)
_joint_prepare[(tokens, indexer_programs + main_programs)](
query,
key,
weights,
q_weight,
k_weight,
k_bias,
cos_sin_cache,
positions,
q,
k,
w,
decode_z,
decode_slots,
decode_cache.mean if decode_cache is not None else None,
decode_ring.rows if decode_ring is not None else None,
decode_ring.metadata().base if decode_ring is not None else None,
decode_cache is not None,
decode_slots.numel() if decode_slots is not None else 0,
decode_z.stride(0) if decode_z is not None else 0,
decode_slots.stride(0) if decode_slots is not None else 0,
decode_ring.capacity if decode_ring is not None else 8,
*(
summary_layout(decode_cache.mean)
if decode_cache is not None
else (0, 8)
),
query.stride(0),
key.stride(0),
weights.stride(0),
heads,
padded_heads,
groups,
tokens,
cos_sin_cache.shape[1],
cos_sin_cache.stride(0),
eps,
float(zero_centered),
native_rope,
main_qkv,
main_qw,
main_kw,
main_rope,
main_slots,
main_q,
main_k,
main_v,
main_kc,
main_vc,
main_qh,
main_kh,
main_qkv.stride(0),
main_rotary,
main_rope.stride(0),
main_slots.numel(),
main_kc.shape[1],
main_kc.stride(0),
main_kc.stride(1),
main_kc.stride(2),
main_eps,
main_bias,
main_k_scale,
main_v_scale,
fp8_cache,
num_warps=4,
enable_fp_fusion=False,
)
return q.transpose(0, 1), k, w.transpose(0, 1), main_q, main_k, main_v