biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

INDEXER PROBE: instrumentation prints for compressed key width investigation

Browse Source
This commit is contained in:
biondizzle
2026-06-02 04:44:47 +00:00
parent 510eaf4a26
commit 06c92f208f
1 changed files with 38 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
single_shot_inference.py
View File

@@ -359,6 +359,11 @@ class Compressor:
bi = torch.arange(n_complete, device=dev)
pos_idx = ((bi + 1) * r - 1).clamp(max=positions.numel() - 1)
comp_pos = positions[pos_idx]
# INDEXER PROBE: Compressor output shape
ident = f"hd={self.hd} kv_dim={self.kv_dim} ratio={self.ratio} is_csa={self.is_csa}"
print(f" COMPRESSOR OUT [{ident}]: compressed.shape={tuple(compressed.shape)} "
f"dtype={compressed.dtype} stride={compressed.stride()} "
f"contig={compressed.is_contiguous()}", flush=True)
return compressed, comp_pos, torch.zeros(1, T, n_complete, dtype=torch.float32, device=dev)
# =====================================================================
@@ -394,12 +399,41 @@ class Indexer:
self.compressor = Compressor(4, self.ihd, 7168, dev)
self.compressor.load(w, f"{pfx}.compressor", dev)
def forward(self, q_lora, hidden_states, comp_indexer_kv, positions):
if self.q_b_lin is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0: return None
def forward(self, q_lora, hidden_states, comp_indexer_kv, positions, layer_idx=None):
if self.q_b_lin is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0:
print(f" INDEXER SKIP L{layer_idx}: q_b_lin={self.q_b_lin is not None} "
f"comp_idx_kv={tuple(comp_indexer_kv.shape) if comp_indexer_kv is not None else None}", flush=True)
return None
dev = q_lora.device; T = q_lora.shape[0]; n_comp = comp_indexer_kv.shape[0]
# INDEXER PROBE: print shapes at layer_idx==0 only
li = layer_idx
if li == 0:
print(f"\n=== INDEXER PROBE L0 ===", flush=True)
print(f" q_lora: shape={tuple(q_lora.shape)} dtype={q_lora.dtype}", flush=True)
print(f" comp_idx_kv: shape={tuple(comp_indexer_kv.shape)} "
f"dtype={comp_indexer_kv.dtype} stride={comp_indexer_kv.stride()} "
f"contig={comp_indexer_kv.is_contiguous()}", flush=True)
print(f" self.n_ih={self.n_ih} self.ihd={self.ihd} n_ih*ihd={self.n_ih * self.ihd}", flush=True)
print(f" self.q_b_lin.in_features={self.q_b_lin.in_features} out_features={self.q_b_lin.out_features}", flush=True)
print(f" self.wp_lin.in_features={self.wp_lin.in_features} out_features={self.wp_lin.out_features}", flush=True)
if self.compressor is not None:
print(f" self.compressor.kv_dim={self.compressor.kv_dim} ratio={self.compressor.ratio} hd={self.compressor.hd}", flush=True)
q_idx = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd)
w_h = self.wp_lin(hidden_states) # (T, n_ih)
k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd)
# INDEXER PROBE: try reshape, catch failure
try:
k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd)
except RuntimeError as e:
print(f" !!! RESHAPE FAILURE L{li} !!!", flush=True)
print(f" comp_indexer_kv.shape = {tuple(comp_indexer_kv.shape)}", flush=True)
print(f" tried to reshape to ({n_comp}, {self.n_ih}, {self.ihd})", flush=True)
print(f" total elements: have {comp_indexer_kv.numel()}, need {n_comp * self.n_ih * self.ihd}", flush=True)
raise
if li == 0:
print(f"--- INDEXER L0 SCORING TENSORS ---", flush=True)
print(f" q_idx: shape={tuple(q_idx.shape)} dtype={q_idx.dtype}", flush=True)
print(f" k_idx: shape={tuple(k_idx.shape)} dtype={k_idx.dtype}", flush=True)
print(f" w_h: shape={tuple(w_h.shape)} dtype={w_h.dtype}", flush=True)
scores = torch.einsum('tnd,cnd->tnc', q_idx.float(), k_idx.float())
scores = F.relu(scores); total = (scores * w_h.unsqueeze(-1).float()).sum(1)
tk = min(self.top_k, n_comp); _, idx = total.topk(tk, -1); return idx
@@ -558,7 +592,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
# 4. Indexer top-k (CSA)
topk_idx = None
if indexer is not None and ratio == 4:
topk_idx = indexer.forward(q_a, x_normed, kv_cache.comp_idx_kv, positions)
topk_idx = indexer.forward(q_a, x_normed, kv_cache.comp_idx_kv, positions, layer_idx=li)
# 5. Gather KV
_pt('gather_start')