RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

minimax: early detect qkv vs q,k,v (#15581) 

avoid a commit charge surge on non-dynamic windows due to double linear
creation.
This commit is contained in:
rattus 2026-08-14 03:08:08 +10:00 committed by GitHub
parent 86aedfd943
commit ddbaa87528
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 42 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
comfy/ldm/minimax_music/ar.py
View File

@ -43,15 +43,17 @@ def sample_topk(logits, top_k, generator):
class RVQAttention(nn.Module):
def __init__(self, hidden_size, num_heads, dtype, device, operations):
def __init__(self, hidden_size, num_heads, merged_qkv, dtype, device, operations):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.merged_qkv = None
self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.merged_qkv = merged_qkv
if merged_qkv:
self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
else:
self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
def forward(self, x):
@ -81,12 +83,14 @@ class RVQRMSNorm(nn.Module):
class RVQMLP(nn.Module):
def __init__(self, hidden_size, intermediate_size, dtype, device, operations):
def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations):
super().__init__()
self.merged_mlp = None
self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.merged_mlp = merged_mlp
if merged_mlp:
self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
else:
self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device)
def forward(self, x):
@ -96,12 +100,12 @@ class RVQMLP(nn.Module):
class RVQDecoderBlock(nn.Module):
def __init__(self, hidden_size, num_heads, intermediate_size, dtype, device, operations):
def __init__(self, hidden_size, num_heads, intermediate_size, merged_qkv, merged_mlp, dtype, device, operations):
super().__init__()
self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device)
self.self_attn = RVQAttention(hidden_size, num_heads, dtype, device, operations)
self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations)
self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device)
self.mlp = RVQMLP(hidden_size, intermediate_size, dtype, device, operations)
self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, dtype, device, operations)
def forward(self, x):
x = x + self.self_attn(self.input_layernorm(x))
@ -113,6 +117,8 @@ class RVQDepthDecoder(nn.Module):
super().__init__()
hidden_size = int(config["hidden_size"])
audio_vocab_size = int(config["audio_vocab_size"])
merged_qkv = config.get("decoder_merged_qkv", False)
merged_mlp = config.get("decoder_merged_mlp", False)
num_codebooks = int(config["audio_num_codebooks"])
self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device)
@ -125,6 +131,8 @@ class RVQDepthDecoder(nn.Module):
hidden_size,
int(config["decoder_num_heads"]),
int(config["decoder_intermediate_size"]),
merged_qkv,
merged_mlp,
dtype,
device,
operations,
@ -148,8 +156,6 @@ class MiniMaxMusic3AR(nn.Module):
qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields})
qwen_config.lm_head = False
qwen_config.fixed_kv = True
qwen_config.merged_qkv = None
qwen_config.merged_mlp = None
self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations)
self.model.prefetch_dynamic_vbars = True
self.model.graph_dynamic_vbar_blocks = True

1
comfy/sd.py
View File

@ -1716,6 +1716,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
if quant is not None:
model_options = model_options.copy()
model_options["quantization_metadata"] = quant
clip_target.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0])
clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel
clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer
elif te_model == TEModel.CLIP_G:

5
comfy/supported_models.py
View File

@ -2218,7 +2218,10 @@ class MiniMaxMusic3(supported_models_base.BASE):
return model_base.ModelType.FLOW
def clip_target(self, state_dict={}):
return supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0])
target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
target.params["projection_config"] = detect
return target
class LongCatImage(supported_models_base.BASE):

15
comfy/text_encoders/llama.py
View File

@ -17,13 +17,6 @@ import comfy.clip_model
from . import qwen_vl
def detect_merged_config(state_dict, prefix="", layer_prefix="model.layers.0."):
return {
"merged_qkv": "{}{}self_attn.qkv_proj.weight".format(prefix, layer_prefix) in state_dict,
"merged_mlp": "{}{}mlp.gate_up_proj.weight".format(prefix, layer_prefix) in state_dict,
}
@dataclass
class FixedKV:
key: torch.Tensor
@ -528,9 +521,9 @@ class Attention(nn.Module):
ops = ops or nn
self.kv_size = self.num_kv_heads * self.head_dim
self.merged_qkv = getattr(config, "merged_qkv", False)
if self.merged_qkv is not False:
if self.merged_qkv:
self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype)
if self.merged_qkv is not True:
else:
self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.k_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
self.v_proj = ops.Linear(config.hidden_size, self.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
@ -627,9 +620,9 @@ class MLP(nn.Module):
super().__init__()
intermediate_size = intermediate_size or config.intermediate_size
self.merged_mlp = getattr(config, "merged_mlp", False)
if self.merged_mlp is not False:
if self.merged_mlp:
self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype)
if self.merged_mlp is not True:
else:
self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype)

34
comfy/text_encoders/minimax_music.py
View File

@ -2,7 +2,6 @@ import torch
from tokenizers import Tokenizer
import comfy.ops
import comfy.text_encoders.llama
from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR
from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt
@ -26,6 +25,15 @@ MODEL_CONFIG = {
}
def detect_merged_config(state_dict, prefix=""):
return {
"merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
"merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
"decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
"decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
}
class MiniMaxMusic3Tokenizer:
def __init__(self, embedding_directory=None, tokenizer_data={}):
tokenizer_json = tokenizer_data.get("tokenizer_json")
@ -58,13 +66,13 @@ class MiniMaxMusic3Tokenizer:
class MiniMaxMusic3TEModel(MiniMaxMusic3AR):
def __init__(self, device="cpu", dtype=None, model_options={}):
def __init__(self, device="cpu", dtype=None, model_options={}, projection_config=None):
dtype = torch.bfloat16
quant_config = model_options.get("quantization_metadata", None)
operations = model_options.get("custom_operations", None)
if operations is None:
operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast
super().__init__(MODEL_CONFIG, dtype, device, operations)
super().__init__({**MODEL_CONFIG, **(projection_config or {})}, dtype, device, operations)
self.dtypes = {dtype}
self.execution_device = device
@ -91,26 +99,6 @@ class MiniMaxMusic3TEModel(MiniMaxMusic3AR):
return hidden.unsqueeze(0), None, {}
def load_state_dict(self, state_dict, strict=True, assign=False):
def select_projections(layers, config):
for layer in layers:
if layer.self_attn.merged_qkv is None:
if config["merged_qkv"]:
del layer.self_attn.q_proj, layer.self_attn.k_proj, layer.self_attn.v_proj
else:
del layer.self_attn.qkv_proj
layer.self_attn.merged_qkv = config["merged_qkv"]
if layer.mlp.merged_mlp is None:
if config["merged_mlp"]:
del layer.mlp.gate_proj, layer.mlp.up_proj
else:
del layer.mlp.gate_up_proj
layer.mlp.merged_mlp = config["merged_mlp"]
select_projections(self.model.layers, comfy.text_encoders.llama.detect_merged_config(state_dict))
select_projections(
self.model.audio_decoder.layers,
comfy.text_encoders.llama.detect_merged_config(state_dict, layer_prefix="model.audio_decoder.layers.0."),
)
if self.model.pruned_embedding is None:
self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict
if self.model.pruned_embedding: