def make_quant_norm(model):
    """
    Replace all LlamaRMSNorm modules with TritonLlamaRMSNorm modules
    """

    for name, m in model.named_modules():
        if not isinstance(m, LlamaRMSNorm):
            continue

        norm = TritonLlamaRMSNorm(m.weight, m.variance_epsilon)

        if '.' in name:
            parent_name = name.rsplit('.', 1)[0]
            child_name = name[len(parent_name) + 1:]
            parent = model.get_submodule(parent_name)
        else:
            parent_name = ''
            parent = model
            child_name = name

        #print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}")

        setattr(parent, child_name, norm)

    Replace all LlamaRMSNorm modules with TritonLlamaRMSNorm modules
    
quant/triton_norm.py