INNER CODE UNIT · Python

nb

FareedKhan-dev/kimi-k3-in-c · tools/budget.py:137

            nb = b - a
            ntensor += 1
            total += nb
            c = classify(name)
            by_class[c] += nb
            cnt_class[c] += 1
            by_dtype[e["dtype"]] += nb

            m = re.search(r"layers\.(\d+)\.self_attn\.(\w+)", name)
            if m:
                # A KDA layer has A_log and conv1d weights; an MLA layer has kv_a_proj.
                if m.group(2) in ("A_log",):
                    kda_layers.add(int(m.group(1)))
                if "kv_a_proj" in m.group(2) or "kv_b_proj" in m.group(2):
                    mla_layers.add(int(m.group(1)))

            if ".block_sparse_moe.experts." in name:
                routed_bytes += nb

View source record →