diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json index 5c52a73ec..c2ed5838d 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "flash_attn3", "rms_type": "torch_native", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "lora_dynamic_apply": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json index f9db2b856..30ab8bd98 100644 --- a/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json +++ b/configs/minimax_h3/dmd/minimax_h3_bf16_4step_sol.json @@ -27,12 +27,12 @@ "strict": true }, "rms_type": "torch_native", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "lora_dynamic_apply": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json index ad427bfff..099349c98 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json index 19819f2e2..5ccbb7915 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090.json @@ -17,13 +17,13 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json index 035390763..a8f21d9de 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8.json @@ -20,7 +20,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, @@ -30,7 +30,7 @@ "vae_attn_type": "sage_attn2", "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json index 2a65cffc2..036f541f8 100755 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sla.json @@ -23,7 +23,7 @@ "operator": "sage2" }, "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, @@ -31,7 +31,7 @@ "vae_attn_type": "sage_attn2", "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json index 995bd43eb..c58aa6b11 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_4step_5090_vae_fp8_sol.json @@ -32,7 +32,7 @@ "strict": true }, "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, @@ -47,7 +47,7 @@ "vae_attn_type": "sage_attn2", "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json index e1feb03aa..0581ee0ca 100644 --- a/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_fp8_8step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json index a76bbe433..20c14d4e6 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_4step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json index c2e603c37..3bc5d8fb1 100644 --- a/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json +++ b/configs/minimax_h3/dmd/minimax_h3_int8_convrot_8step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 6.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "dit_quantized": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json index 42e082753..2fe26c43d 100755 --- a/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json +++ b/configs/minimax_h3/dmd/minimax_h3_ref2av_4step.json @@ -16,12 +16,12 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", + "h3_sampling_profile": "dmd", "vae_spatial_scale_factor": 16, "audio_sampling_rate": 32000, "text_encoder_tensor_parallel": true, diff --git a/configs/minimax_h3/dmd/minimax_h3_ref2av_turbo_4step_tp8.json b/configs/minimax_h3/dmd/minimax_h3_ref2av_turbo_4step_tp8.json new file mode 100644 index 000000000..ceef1da09 --- /dev/null +++ b/configs/minimax_h3/dmd/minimax_h3_ref2av_turbo_4step_tp8.json @@ -0,0 +1,40 @@ +{ + "infer_steps": 4, + "target_video_length": 362, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "reference_image_resize_mode": "match", + "offload_granularity": "model", + "use_adaln_cache": true, + "adaln_cache_dir": "/data/wushuo1/.cache/lightx2v/adaln", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "sgl_exact_neox_rope", + "feature_caching": "NoCaching", + "use_compile": false, + "warmup": false, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "h3_sampling_profile": "dmd", + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "text_encoder_tensor_parallel": true, + "lora_dynamic_apply": false, + "lora_configs": [ + { + "path": "/data/wushuo1/models/Minimax-h3-Turbo/Minimax-h3-Turbo/minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors", + "strength": 1.0, + "alpha": 8 + } + ], + "parallel": { + "tensor_p_size": 8 + } +} diff --git a/configs/minimax_h3/fp8/minimax_h3.json b/configs/minimax_h3/fp8/minimax_h3.json index 61a955ed9..3c45df0d8 100644 --- a/configs/minimax_h3/fp8/minimax_h3.json +++ b/configs/minimax_h3/fp8/minimax_h3.json @@ -15,7 +15,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json index b23bcd897..55124b9b3 100644 --- a/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_encoder_fp8.json @@ -18,7 +18,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/fp8/minimax_h3_sp_5090.json b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json index 82e9268fd..3d3fc5ca9 100644 --- a/configs/minimax_h3/fp8/minimax_h3_sp_5090.json +++ b/configs/minimax_h3/fp8/minimax_h3_sp_5090.json @@ -17,7 +17,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, diff --git a/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json index c83671b84..ad8576564 100644 --- a/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json +++ b/configs/minimax_h3/fp8/minimax_h3_vae_fp8.json @@ -17,7 +17,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, diff --git a/configs/minimax_h3/minimax_h3.json b/configs/minimax_h3/minimax_h3.json index 6d956e57c..faa322c37 100644 --- a/configs/minimax_h3/minimax_h3.json +++ b/configs/minimax_h3/minimax_h3.json @@ -6,16 +6,18 @@ "fps": 24, "enable_cfg": false, "cpu_offload": true, + "reference_image_resize_mode": "diffusers", "offload_granularity": "model", "use_adaln_cache": true, "adaln_cache_dir": "~/.cache/lightx2v/adaln", "text_encoder_cpu_offload": true, "vae_cpu_offload": true, + "vae_encode_fp32": true, "lazy_load": false, "unload_modules": false, - "attn_type": "sage_attn2", - "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "attn_type": "torch_sdpa", + "rms_type": "torch_native", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/minimax_h3_block_offload.json b/configs/minimax_h3/minimax_h3_block_offload.json index 73b84dfbb..59af22400 100644 --- a/configs/minimax_h3/minimax_h3_block_offload.json +++ b/configs/minimax_h3/minimax_h3_block_offload.json @@ -16,7 +16,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/minimax_h3_compile.json b/configs/minimax_h3/minimax_h3_compile.json index 942467485..cd5f34993 100644 --- a/configs/minimax_h3/minimax_h3_compile.json +++ b/configs/minimax_h3/minimax_h3_compile.json @@ -15,7 +15,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "torch_real_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": true, "warmup": true, diff --git a/configs/minimax_h3/minimax_h3_sol_block_offload.json b/configs/minimax_h3/minimax_h3_sol_block_offload.json index 8f60bc649..d221f21ab 100644 --- a/configs/minimax_h3/minimax_h3_sol_block_offload.json +++ b/configs/minimax_h3/minimax_h3_sol_block_offload.json @@ -27,7 +27,7 @@ "strict": true }, "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/minimax_h3_sp.json b/configs/minimax_h3/minimax_h3_sp.json index 3bd107ecc..6921ff29f 100644 --- a/configs/minimax_h3/minimax_h3_sp.json +++ b/configs/minimax_h3/minimax_h3_sp.json @@ -16,7 +16,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/minimax_h3_tp.json b/configs/minimax_h3/minimax_h3_tp.json index d952b11a9..7b88e19f1 100644 --- a/configs/minimax_h3/minimax_h3_tp.json +++ b/configs/minimax_h3/minimax_h3_tp.json @@ -17,7 +17,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/configs/minimax_h3/minimax_h3_tp_sp.json b/configs/minimax_h3/minimax_h3_tp_sp.json index 9eb0e7cd4..962abfb98 100644 --- a/configs/minimax_h3/minimax_h3_tp_sp.json +++ b/configs/minimax_h3/minimax_h3_tp_sp.json @@ -17,7 +17,7 @@ "unload_modules": false, "attn_type": "sage_attn2", "rms_type": "sgl-kernel", - "rope_type": "minimax_h3_triton_rope", + "rope_type": "sgl_exact_neox_rope", "feature_caching": "NoCaching", "use_compile": false, "video_flow_shift": 12.0, diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index 1cdddcaf4..cfe548a2a 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -129,6 +129,35 @@ def apply( max_seqlen_kv=None, **kwargs, ): + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must either both be set or both be None") + if q.ndim == 3 and cu_seqlens_q is not None and cu_seqlens_q.numel() > 2: + q_bounds = cu_seqlens_q.tolist() + kv_bounds = cu_seqlens_kv.tolist() + if ( + len(q_bounds) != len(kv_bounds) + or q_bounds[0] != 0 + or kv_bounds[0] != 0 + or q_bounds[-1] != q.shape[0] + or kv_bounds[-1] != k.shape[0] + ): + raise ValueError("Invalid packed dynamic sparse attention boundaries") + output = q.new_empty((q.shape[0], q.shape[1] * v.shape[-1])) + for q_start, q_stop, kv_start, kv_stop in zip(q_bounds, q_bounds[1:], kv_bounds, kv_bounds[1:]): + if q_start == q_stop: + continue + segment = self.apply_func( + q[q_start:q_stop], + k[kv_start:kv_stop], + v[kv_start:kv_stop], + None, + None, + q_stop - q_start, + kv_stop - kv_start, + **kwargs, + ) + output[q_start:q_stop].copy_(segment) + return output if max_seqlen_q is None: max_seqlen_q = q.shape[0] if max_seqlen_kv is None: diff --git a/lightx2v/common/ops/attn/flash_attn.py b/lightx2v/common/ops/attn/flash_attn.py index 2d19874d6..6090658cf 100755 --- a/lightx2v/common/ops/attn/flash_attn.py +++ b/lightx2v/common/ops/attn/flash_attn.py @@ -34,6 +34,20 @@ from .template import AttnWeightTemplate +def _uses_varlen(q, cu_seqlens_q, cu_seqlens_kv): + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must either both be set or both be None") + return (q.ndim == 4 and q.shape[0] > 1) or (cu_seqlens_q is not None and cu_seqlens_q.numel() > 2) + + +def _flatten_varlen_qkv(q, k, v): + if q.ndim == 4: + q = q.reshape(-1, q.shape[-2], q.shape[-1]) + k = k.reshape(-1, k.shape[-2], k.shape[-1]) + v = v.reshape(-1, v.shape[-2], v.shape[-1]) + return q, k, v + + @ATTN_WEIGHT_REGISTER("flash_attn2") class FlashAttn2Weight(AttnWeightTemplate): def __init__(self): @@ -52,27 +66,25 @@ def apply( ): causal = kwargs.get("causal", False) softmax_scale = kwargs.get("softmax_scale", None) - if len(q.shape) == 3: - bs = 1 - elif len(q.shape) == 4: - bs = q.shape[0] - total_seqlen = bs * max_seqlen_q - - if bs == 1: + total_seqlen = q.shape[0] if q.ndim == 3 else q.shape[0] * q.shape[1] + if not _uses_varlen(q, cu_seqlens_q, cu_seqlens_kv): if len(q.shape) == 3: q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) x = flash_attn_func_v2(q, k, v, softmax_scale=softmax_scale, causal=causal).reshape(total_seqlen, -1) else: + if cu_seqlens_q is None: + batch_size, sequence_length = q.shape[:2] + cu_seqlens_q = torch.arange(batch_size + 1, device=q.device, dtype=torch.int32) * sequence_length + cu_seqlens_kv = torch.arange(batch_size + 1, device=k.device, dtype=torch.int32) * k.shape[1] + max_seqlen_q = sequence_length + max_seqlen_kv = k.shape[1] if cu_seqlens_q.is_cpu: cu_seqlens_q = cu_seqlens_q.to(q.device, non_blocking=True) if cu_seqlens_kv.is_cpu: cu_seqlens_kv = cu_seqlens_kv.to(k.device, non_blocking=True) - if len(q.shape) == 4: - q = q.reshape(-1, q.shape[-2], q.shape[-1]) - k = k.reshape(-1, k.shape[-2], k.shape[-1]) - v = v.reshape(-1, v.shape[-2], v.shape[-1]) + q, k, v = _flatten_varlen_qkv(q, k, v) x = flash_attn_varlen_func_v2( q, k, @@ -127,27 +139,25 @@ def apply( ): causal = kwargs.get("causal", False) softmax_scale = kwargs.get("softmax_scale", None) - if len(q.shape) == 3: - bs = 1 - elif len(q.shape) == 4: - bs = q.shape[0] - total_seqlen = bs * max_seqlen_q - - if bs == 1: + total_seqlen = q.shape[0] if q.ndim == 3 else q.shape[0] * q.shape[1] + if not _uses_varlen(q, cu_seqlens_q, cu_seqlens_kv): if len(q.shape) == 3: q = q.unsqueeze(0) k = k.unsqueeze(0) v = v.unsqueeze(0) x = flash_attn_func_v3(q, k, v, softmax_scale=softmax_scale, causal=causal).reshape(total_seqlen, -1) else: + if cu_seqlens_q is None: + batch_size, sequence_length = q.shape[:2] + cu_seqlens_q = torch.arange(batch_size + 1, device=q.device, dtype=torch.int32) * sequence_length + cu_seqlens_kv = torch.arange(batch_size + 1, device=k.device, dtype=torch.int32) * k.shape[1] + max_seqlen_q = sequence_length + max_seqlen_kv = k.shape[1] if cu_seqlens_q.is_cpu: cu_seqlens_q = cu_seqlens_q.to(q.device, non_blocking=True) if cu_seqlens_kv.is_cpu: cu_seqlens_kv = cu_seqlens_kv.to(k.device, non_blocking=True) - if len(q.shape) == 4: - q = q.reshape(-1, q.shape[-2], q.shape[-1]) - k = k.reshape(-1, k.shape[-2], k.shape[-1]) - v = v.reshape(-1, v.shape[-2], v.shape[-1]) + q, k, v = _flatten_varlen_qkv(q, k, v) x = flash_attn_varlen_func_v3( q, k, diff --git a/lightx2v/common/ops/attn/sage_attn.py b/lightx2v/common/ops/attn/sage_attn.py index 07945b9f5..7af97ec12 100755 --- a/lightx2v/common/ops/attn/sage_attn.py +++ b/lightx2v/common/ops/attn/sage_attn.py @@ -38,6 +38,11 @@ logger.info("sageattn not found, please install sageattention first") sageattn = None +try: + from sageattention import sageattn_varlen +except ImportError: + sageattn_varlen = None + if magi_register_custom_op is not None and sageattn is not None: @@ -88,6 +93,31 @@ def apply( **kwargs, ): q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must either both be set or both be None") + if q.ndim == 3 and cu_seqlens_q is not None and cu_seqlens_q.numel() > 2: + if sageattn_varlen is None: + raise ImportError("Packed SageAttention2 requires sageattn_varlen") + cu_seqlens_q = cu_seqlens_q.to(q.device, non_blocking=True).contiguous() + cu_seqlens_kv = cu_seqlens_kv.to(k.device, non_blocking=True).contiguous() + if max_seqlen_q is None: + max_seqlen_q = int((cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item()) + if max_seqlen_kv is None: + max_seqlen_kv = int((cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item()) + x = sageattn_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + is_causal=kwargs.get("causal", False), + sm_scale=kwargs.get("softmax_scale"), + smooth_k=False, + ) + return x.flatten(1) + if len(q.shape) == 3: bs = 1 q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) diff --git a/lightx2v/common/ops/attn/sol_attn.py b/lightx2v/common/ops/attn/sol_attn.py index 7d07d89ac..caec38046 100644 --- a/lightx2v/common/ops/attn/sol_attn.py +++ b/lightx2v/common/ops/attn/sol_attn.py @@ -467,6 +467,37 @@ def apply( max_seqlen_kv=None, **kwargs, ): + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must either both be set or both be None") + if q.ndim == 3 and cu_seqlens_q is not None and cu_seqlens_q.numel() > 2: + q_bounds = cu_seqlens_q.tolist() + kv_bounds = cu_seqlens_kv.tolist() + if ( + len(q_bounds) != len(kv_bounds) + or q_bounds[0] != 0 + or kv_bounds[0] != 0 + or q_bounds[-1] != q.shape[0] + or kv_bounds[-1] != k.shape[0] + ): + raise ValueError("Invalid packed Sol-Attn boundaries") + output = q.new_empty((q.shape[0], q.shape[1] * v.shape[-1])) + for q_start, q_stop, kv_start, kv_stop in zip(q_bounds, q_bounds[1:], kv_bounds, kv_bounds[1:]): + if q_start == q_stop: + continue + segment_mask = None if attn_mask is None else attn_mask[..., q_start:q_stop, kv_start:kv_stop] + segment = self.apply( + q=q[q_start:q_stop], + k=k[kv_start:kv_stop], + v=v[kv_start:kv_stop], + drop_rate=drop_rate, + attn_mask=segment_mask, + causal=causal, + max_seqlen_q=q_stop - q_start, + max_seqlen_kv=kv_stop - kv_start, + **kwargs, + ) + output[q_start:q_stop].copy_(segment) + return output scale = kwargs.get("softmax_scale", kwargs.get("scale")) dense_kwargs = { "drop_rate": drop_rate, diff --git a/lightx2v/common/ops/attn/torch_sdpa.py b/lightx2v/common/ops/attn/torch_sdpa.py index fd12abd80..3e28a6a29 100644 --- a/lightx2v/common/ops/attn/torch_sdpa.py +++ b/lightx2v/common/ops/attn/torch_sdpa.py @@ -13,6 +13,13 @@ class TorchSDPAWeight(AttnWeightTemplate): def __init__(self): self.config = {} + @staticmethod + def _cu_bounds(cu_seqlens, seq_len, name): + bounds = cu_seqlens.tolist() + if len(bounds) < 2 or bounds[0] != 0 or bounds[-1] != seq_len or any(start > stop for start, stop in zip(bounds, bounds[1:])): + raise ValueError(f"Invalid cu_seqlens_{name}") + return bounds + def apply( self, q, @@ -27,35 +34,73 @@ def apply( max_seqlen_kv=None, **kwargs, ): - if q.ndim == 3: - q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) + softmax_scale = kwargs.get("softmax_scale") if attn_mask is not None and attn_mask.dtype != torch.bool: attn_mask = attn_mask.to(q.dtype) - # Hunyuan3D upstream Attention uses SDPA flash kernel (see hy3dshape hunyuandit.py). - # Matching this context is required for bit-identical attention vs the reference. - sdpa_ctx = nullcontext() - if kwargs.get("model_cls") == "hunyuan3d": - sdpa_ctx = torch.backends.cuda.sdp_kernel( - enable_flash=True, - enable_math=False, - enable_mem_efficient=True, - ) - with sdpa_ctx: - # q/k/v are (B, H, S, D) here, so head count is dim 1. GQA models such as - # neopp (32 q heads, 8 kv heads) need SDPA to broadcast the kv groups. - x = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=attn_mask, - dropout_p=drop_rate, - is_causal=causal, - enable_gqa=q.shape[1] != k.shape[1], + + def run_sdpa(query, key, value, mask): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + # Match Hunyuan3D's flash-only upstream dispatch. + sdpa_ctx = nullcontext() + if kwargs.get("model_cls") == "hunyuan3d": + sdpa_ctx = torch.backends.cuda.sdp_kernel( + enable_flash=True, + enable_math=False, + enable_mem_efficient=True, + ) + with sdpa_ctx: + return F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + dropout_p=drop_rate, + is_causal=causal, + scale=softmax_scale, + enable_gqa=query.shape[1] != key.shape[1], + ).transpose(1, 2) + + if (cu_seqlens_q is None) != (cu_seqlens_kv is None): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must either both be set or both be None") + + use_packed = q.ndim == 3 and cu_seqlens_q is not None + if use_packed: + if cu_seqlens_q.numel() != cu_seqlens_kv.numel(): + raise ValueError("cu_seqlens_q and cu_seqlens_kv must describe the same number of sequences") + if cu_seqlens_q.numel() < 2: + raise ValueError("cu_seqlens must contain at least two boundaries") + use_packed = cu_seqlens_q.numel() > 2 + + # A 4D tensor already has an explicit batch dimension. A two-entry + # cu_seqlens tensor describes one sequence, so it uses the same dense path. + if not use_packed: + if q.ndim == 3: + q, k, v = q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0) + return run_sdpa(q, k, v, attn_mask).flatten(2).squeeze(0) + + if any(x.ndim != 3 for x in (q, k, v)): + raise ValueError("Packed Torch SDPA expects 3D q/k/v") + + q_bounds = self._cu_bounds(cu_seqlens_q, q.shape[0], "q") + kv_bounds = self._cu_bounds(cu_seqlens_kv, k.shape[0], "kv") + if len(q_bounds) != len(kv_bounds): + raise ValueError("Packed q and kv must contain the same number of sequences") + if v.shape[0] != k.shape[0]: + raise ValueError("Packed k and v sequence lengths must match") + + # PyTorch SDPA has no varlen API, so process each cu_seqlens segment independently. + output = q.new_empty((q.shape[0], q.shape[1], v.shape[-1])) + for q_start, q_stop, kv_start, kv_stop in zip(q_bounds, q_bounds[1:], kv_bounds, kv_bounds[1:]): + if q_start == q_stop: + continue + segment_mask = None if attn_mask is None else attn_mask[..., q_start:q_stop, kv_start:kv_stop] + segment = run_sdpa( + q[q_start:q_stop].unsqueeze(0), + k[kv_start:kv_stop].unsqueeze(0), + v[kv_start:kv_stop].unsqueeze(0), + segment_mask, ) - x = x.transpose(1, 2) - b, s, a, d = x.shape - out = x.reshape(b, s, -1) - return out.squeeze(0) + output[q_start:q_stop].copy_(segment[0]) + return output.flatten(1) diff --git a/lightx2v/common/ops/attn/ulysses_attn.py b/lightx2v/common/ops/attn/ulysses_attn.py index 21b530ee6..a9077dda6 100755 --- a/lightx2v/common/ops/attn/ulysses_attn.py +++ b/lightx2v/common/ops/attn/ulysses_attn.py @@ -113,8 +113,9 @@ def apply_new( ``None`` selects QKV-only attention; passing only ``aux_q=None`` selects the q-only cross-attention form. Set ``aux_first=True`` when auxiliary tokens precede the A2A tokens in the attention sequence. - This interface supports one logical Q sequence and one logical KV - sequence only. Packed-varlen batches are not supported. + Packed boundaries in ``attention_kwargs`` must describe the full + sequence reconstructed after A2A, including the auxiliary prefix when + ``aux_first=True``. ``tensor_fusion`` packs Q/K/V into one Ulysses communication payload. The return value is ``(output, aux_output)``. """ diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 61613dd5a..6ffecb46f 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -628,9 +628,10 @@ def act_quant_fp8_perchannel_sym_vllm(self, x): return input_tensor_quant, input_tensor_scale def act_quant_fp8_perchannel_sym_sgl(self, x): + x = x.contiguous() m, k = x.shape - input_tensor_quant = torch.empty((m, k), dtype=torch.float8_e4m3fn, device="cuda", requires_grad=False) - input_tensor_scale = torch.empty((m, 1), dtype=torch.float32, device="cuda", requires_grad=False) + input_tensor_quant = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device, requires_grad=False) + input_tensor_scale = torch.empty((m, 1), dtype=torch.float32, device=x.device, requires_grad=False) sgl_kernel.sgl_per_token_quant_fp8(x, input_tensor_quant, input_tensor_scale) return input_tensor_quant, input_tensor_scale @@ -659,9 +660,10 @@ def act_quant_fp8_perchannelgroup128_sym_deepgemm(self, x): return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, n), (x_amax / 448.0).view(m, -1) def act_quant_fp8_perchannelgroup128_sym_sgl(self, x): + x = x.contiguous() m, k = x.shape - input_tensor_quant = torch.empty((m, k), dtype=torch.float8_e4m3fn, device="cuda", requires_grad=False) - input_tensor_scale = torch.empty((m, k // 128), dtype=torch.float32, device="cuda", requires_grad=False) + input_tensor_quant = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device, requires_grad=False) + input_tensor_scale = torch.empty((m, k // 128), dtype=torch.float32, device=x.device, requires_grad=False) sgl_kernel.sgl_per_token_group_quant_fp8( x, input_tensor_quant, @@ -2025,6 +2027,7 @@ def __init__( self.scale_force_fp32 = True def apply(self, input_tensor): + input_tensor = input_tensor.contiguous() shape = (input_tensor.shape[0], self.weight.shape[1]) dtype = input_tensor.dtype device = input_tensor.device @@ -2592,6 +2595,7 @@ def __init__( lora_path="", reduce_output=True, lora_column_chunks=1, + fp32_reduce=False, ): super().__init__( weight_name, @@ -2610,6 +2614,7 @@ def __init__( self.split_dim = split_dim # "col" for column split, "row" for row split self.reduce_output = reduce_output self.lora_column_chunks = lora_column_chunks + self.fp32_reduce = bool(fp32_reduce) assert split_dim in ["col", "row"], f"split_dim must be 'col' or 'row', got {split_dim}" assert lora_column_chunks >= 1, f"lora_column_chunks must be positive, got {lora_column_chunks}" @@ -2727,7 +2732,12 @@ def apply(self, input_tensor): # For row split, need all-reduce to combine results from all ranks if self.split_dim == "row" and self.reduce_output and self.tp_size > 1 and self.tp_group is not None: + output_dtype = output.dtype + if self.fp32_reduce: + output = output.float() dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group) + if self.fp32_reduce: + output = output.to(output_dtype) # Add bias after all-reduce (bias is not split for row split) if self._row_split_bias is not None: output = output + self._row_split_bias diff --git a/lightx2v/common/ops/norm/__init__.py b/lightx2v/common/ops/norm/__init__.py index 1e7247973..507abdd3d 100755 --- a/lightx2v/common/ops/norm/__init__.py +++ b/lightx2v/common/ops/norm/__init__.py @@ -1,2 +1,3 @@ +from .sgl_exact_rms_norm_128 import * from .layer_norm_weight import * from .rms_norm_weight import * diff --git a/lightx2v/common/ops/norm/sgl_exact_rms_norm_128.py b/lightx2v/common/ops/norm/sgl_exact_rms_norm_128.py new file mode 100644 index 000000000..9886f0d39 --- /dev/null +++ b/lightx2v/common/ops/norm/sgl_exact_rms_norm_128.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from lightx2v.utils.registry_factory import RMS_WEIGHT_REGISTER + +from .rms_norm_weight import RMSWeightTemplate + +# Adapted from SGLang commit 8ef646a5c65bd2f8922483057dddc02e2b0de18c. + + +def _require_nvidia_triton(tensor: torch.Tensor) -> None: + if tensor.device.type == "cuda" and getattr(torch.version, "hip", None) is not None: + raise RuntimeError("SGL exact RMSNorm 128 kernel supports NVIDIA CUDA only") + + +@triton.jit +def _mul_rn_f32(x, y): + return tl.inline_asm_elementwise( + asm="mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _fma_rn_f32(x, y, z): + return tl.inline_asm_elementwise( + asm="fma.rn.f32 $0, $1, $2, $3;", + constraints="=f,f,f,f", + args=[x, y, z], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _rsqrt_approx_f32(x): + return tl.inline_asm_elementwise( + asm="rsqrt.approx.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _sgl_exact_rms_norm_128_kernel( + x_ptr, + weight_ptr, + num_heads, + token_stride, + head_stride, + EPS: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + token = row // num_heads + head = row % num_heads + lane = tl.arange(0, 32) + base = x_ptr + token * token_stride + head * head_stride + lane * 4 + + x0 = tl.load(base).to(tl.float32) + x1 = tl.load(base + 1).to(tl.float32) + x2 = tl.load(base + 2).to(tl.float32) + x3 = tl.load(base + 3).to(tl.float32) + accumulator = _fma_rn_f32(x0, x0, 0.0) + accumulator = _fma_rn_f32(x1, x1, accumulator) + accumulator = _fma_rn_f32(x2, x2, accumulator) + accumulator = _fma_rn_f32(x3, x3, accumulator) + + # Match the CUDA warp's SHFL.BFLY reduction order: 16, 8, 4, 2, 1. + accumulator = tl.sum(tl.reshape(accumulator, (2, 16), can_reorder=False), axis=0) + accumulator = tl.sum(tl.reshape(accumulator, (2, 8), can_reorder=False), axis=0) + accumulator = tl.sum(tl.reshape(accumulator, (2, 4), can_reorder=False), axis=0) + accumulator = tl.sum(tl.reshape(accumulator, (2, 2), can_reorder=False), axis=0) + accumulator = tl.sum(tl.reshape(accumulator, (2, 1), can_reorder=False), axis=0) + sum_of_squares = tl.sum(accumulator) + rstd = _rsqrt_approx_f32(_fma_rn_f32(sum_of_squares, 0.0078125, EPS)) + + weight_base = weight_ptr + lane * 4 + w0 = tl.load(weight_base).to(tl.float32) + w1 = tl.load(weight_base + 1).to(tl.float32) + w2 = tl.load(weight_base + 2).to(tl.float32) + w3 = tl.load(weight_base + 3).to(tl.float32) + y0 = _mul_rn_f32(_mul_rn_f32(x0, rstd), w0) + y1 = _mul_rn_f32(_mul_rn_f32(x1, rstd), w1) + y2 = _mul_rn_f32(_mul_rn_f32(x2, rstd), w2) + y3 = _mul_rn_f32(_mul_rn_f32(x3, rstd), w3) + tl.store(base, y0) + tl.store(base + 1, y1) + tl.store(base + 2, y2) + tl.store(base + 3, y3) + + +def _apply_sgl_exact_rms_norm_128( + hidden_states: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + head_dim = 128 + if hidden_states.ndim != 3 or hidden_states.shape[-1] != head_dim: + raise ValueError(f"SGL exact RMSNorm 128 expects [tokens, heads, 128], got {hidden_states.shape}") + if hidden_states.device != weight.device: + raise ValueError("SGL exact RMSNorm 128 tensors must be on one device") + if hidden_states.device.type != "cuda": + return F.rms_norm(hidden_states.float(), (head_dim,), weight.float(), eps).to(hidden_states.dtype) + _require_nvidia_triton(hidden_states) + if hidden_states.dtype is not torch.bfloat16 or weight.dtype is not torch.bfloat16: + raise TypeError("SGL exact RMSNorm 128 requires BF16 activations and weights") + if hidden_states.shape[1] <= 0: + raise ValueError(f"SGL exact RMSNorm 128 requires at least one head, got {hidden_states.shape}") + if hidden_states.stride(-1) != 1 or hidden_states.stride(-2) != head_dim: + raise ValueError(f"Unsupported SGL exact RMSNorm 128 strides: {hidden_states.stride()}") + if hidden_states.shape[0] > 1 and hidden_states.stride(0) < hidden_states.shape[1] * head_dim: + raise ValueError(f"Overlapping SGL exact RMSNorm 128 token strides: {hidden_states.stride()}") + if weight.shape != (head_dim,) or not weight.is_contiguous(): + raise ValueError("SGL exact RMSNorm 128 weights must be a contiguous [128] tensor") + with torch.cuda.device(hidden_states.device): + if hidden_states.numel(): + _sgl_exact_rms_norm_128_kernel[(hidden_states.shape[0] * hidden_states.shape[1],)]( + hidden_states, + weight, + hidden_states.shape[1], + hidden_states.stride(0), + hidden_states.stride(1), + EPS=float(eps), + num_warps=1, + ) + return hidden_states + + +@RMS_WEIGHT_REGISTER("sgl_exact_rms_norm_128") +class SGLExactRMSNorm128(RMSWeightTemplate): + def apply(self, input_tensor: torch.Tensor) -> torch.Tensor: + return _apply_sgl_exact_rms_norm_128(input_tensor, self._get_actual_weight(), self.eps) + + +__all__ = ["SGLExactRMSNorm128"] diff --git a/lightx2v/common/ops/rope/__init__.py b/lightx2v/common/ops/rope/__init__.py index 1cdeff55d..a819c2000 100644 --- a/lightx2v/common/ops/rope/__init__.py +++ b/lightx2v/common/ops/rope/__init__.py @@ -1,4 +1,5 @@ from .chunked_rope import ChunkedRope from .flashinfer_rope import FlashInferRope +from .sgl_exact_neox_rope import SGLExactNeoXRope from .template import RopeLayout, RopeTemplate from .torch_rope import TorchComplexRope, TorchRealRope diff --git a/lightx2v/common/ops/rope/sgl_exact_neox_rope.py b/lightx2v/common/ops/rope/sgl_exact_neox_rope.py new file mode 100644 index 000000000..3ccbc91c5 --- /dev/null +++ b/lightx2v/common/ops/rope/sgl_exact_neox_rope.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import triton +import triton.language as tl + +from lightx2v.utils.registry_factory import ROPE_REGISTER + +from .template import RopeTemplate + +# Adapted from SGLang commit 8ef646a5c65bd2f8922483057dddc02e2b0de18c. + + +def _require_nvidia_triton(tensor: torch.Tensor) -> None: + if tensor.device.type == "cuda" and getattr(torch.version, "hip", None) is not None: + raise RuntimeError("SGL exact NeoX RoPE kernel supports NVIDIA CUDA only") + + +@triton.jit +def _round_bf16_to_fp32(value): + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _round_fp16_to_fp32(value): + rounded = tl.inline_asm_elementwise( + asm="cvt.rn.f16.f32 $0, $1;", + constraints="=h,f", + args=[value], + dtype=tl.float16, + is_pure=True, + pack=1, + ) + return rounded.to(tl.float32) + + +@triton.jit +def _qk_neox_rope_kernel( + q_ptr, + k_ptr, + cache_ptr, + positions_ptr, + q_rows, + q_heads, + k_heads, + head_dim, + q_token_stride, + q_head_stride, + k_token_stride, + k_head_stride, + position_count, + ROPE_DIM: tl.constexpr, + BLOCK_HALF: tl.constexpr, + IS_BF16: tl.constexpr, +): + pid = tl.program_id(0) + is_k = pid >= q_rows + row = pid - q_rows if is_k else pid + heads = k_heads if is_k else q_heads + token = row // heads + head = row % heads + tensor_ptr = k_ptr if is_k else q_ptr + token_stride = k_token_stride if is_k else q_token_stride + head_stride = k_head_stride if is_k else q_head_stride + + half = ROPE_DIM // 2 + offsets = tl.arange(0, BLOCK_HALF) + mask = offsets < half + base = token * token_stride + head * head_stride + position = tl.load(positions_ptr + token % position_count) + cache_base = position * ROPE_DIM + + first = tl.load(tensor_ptr + base + offsets, mask=mask, other=0.0).to(tl.float32) + second = tl.load(tensor_ptr + base + half + offsets, mask=mask, other=0.0).to(tl.float32) + cos = tl.load(cache_ptr + cache_base + offsets, mask=mask, other=0.0).to(tl.float32) + sin = tl.load(cache_ptr + cache_base + half + offsets, mask=mask, other=0.0).to(tl.float32) + + # Round each product to the activation dtype before the final add/sub. + # The helpers are optimization barriers, so Triton cannot contract an FMA. + if IS_BF16: + first_cos = _round_bf16_to_fp32(first * cos) + second_sin = _round_bf16_to_fp32(second * sin) + second_cos = _round_bf16_to_fp32(second * cos) + first_sin = _round_bf16_to_fp32(first * sin) + else: + first_cos = _round_fp16_to_fp32(first * cos) + second_sin = _round_fp16_to_fp32(second * sin) + second_cos = _round_fp16_to_fp32(second * cos) + first_sin = _round_fp16_to_fp32(first * sin) + out_first = first_cos - second_sin + out_second = second_cos + first_sin + + tl.store(tensor_ptr + base + offsets, out_first, mask=mask) + tl.store(tensor_ptr + base + half + offsets, out_second, mask=mask) + + +def _apply_neox_rope_fallback( + hidden_states: torch.Tensor, + cache: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + head_dim = hidden_states.shape[-1] + rotary_dim = cache.shape[-1] + half = rotary_dim // 2 + flat = hidden_states.reshape(-1, hidden_states.shape[-2], head_dim) + if flat.shape[0] % positions.numel(): + raise ValueError(f"RoPE position count {positions.numel()} does not divide token count {flat.shape[0]}") + repeated_positions = positions.repeat(flat.shape[0] // positions.numel()) + selected = cache.index_select(0, repeated_positions) + cos = selected[:, None, :half] + sin = selected[:, None, half:] + first = flat[..., :half] + second = flat[..., half:rotary_dim] + first_cos = (first * cos).to(flat.dtype) + second_sin = (second * sin).to(flat.dtype) + second_cos = (second * cos).to(flat.dtype) + first_sin = (first * sin).to(flat.dtype) + rotated = torch.cat(((first_cos - second_sin).to(flat.dtype), (second_cos + first_sin).to(flat.dtype), flat[..., rotary_dim:]), dim=-1) + return rotated.reshape(hidden_states.shape) + + +def _prepare_qk_neox_rope_inputs( + q: torch.Tensor, + k: torch.Tensor, + cache: torch.Tensor, + positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if q.ndim < 2 or k.ndim < 2 or cache.ndim != 2 or positions.ndim != 1: + raise ValueError(f"SGL exact NeoX RoPE expects Q/K [..., heads, dim], cache [positions, rotary_dim], and positions [tokens]; got {q.shape}, {k.shape}, {cache.shape}, and {positions.shape}") + if q.dtype not in (torch.float16, torch.bfloat16) or k.dtype != q.dtype or cache.dtype != q.dtype: + raise TypeError(f"SGL exact NeoX RoPE requires matching FP16/BF16 Q/K/cache tensors, got {q.dtype}, {k.dtype}, and {cache.dtype}") + if positions.dtype is not torch.long: + raise TypeError(f"SGL exact NeoX RoPE positions must use torch.long, got {positions.dtype}") + if q.device != k.device or q.device != cache.device or q.device != positions.device: + raise ValueError("SGL exact NeoX RoPE tensors must be on one device") + if q.shape[-1] != k.shape[-1] or q.shape[-2] <= 0 or k.shape[-2] <= 0: + raise ValueError(f"Invalid Q/K shapes for SGL exact NeoX RoPE: {q.shape}, {k.shape}") + rotary_dim = cache.shape[-1] + if cache.shape[0] == 0: + raise ValueError("SGL exact NeoX RoPE cache must contain at least one position") + if rotary_dim <= 0 or rotary_dim % 2 or rotary_dim > q.shape[-1]: + raise ValueError(f"Invalid rotary dimension {rotary_dim} for head dimension {q.shape[-1]}") + if positions.numel() == 0: + raise ValueError("SGL exact NeoX RoPE positions must not be empty") + q_tokens = q.numel() // (q.shape[-2] * q.shape[-1]) + k_tokens = k.numel() // (k.shape[-2] * k.shape[-1]) + if q_tokens % positions.numel() or k_tokens % positions.numel(): + raise ValueError(f"RoPE position count {positions.numel()} must divide Q/K token counts {q_tokens}/{k_tokens}") + return cache.contiguous(), positions.contiguous() + + +def _apply_qk_neox_rope( + q: torch.Tensor, + k: torch.Tensor, + cache: torch.Tensor, + positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + cache, positions = _prepare_qk_neox_rope_inputs(q, k, cache, positions) + rotary_dim = cache.shape[-1] + if q.device.type != "cuda": + return _apply_neox_rope_fallback(q, cache, positions), _apply_neox_rope_fallback(k, cache, positions) + _require_nvidia_triton(q) + q_shape = q.shape + k_shape = k.shape + q = q.reshape(-1, q.shape[-2], q.shape[-1]) + k = k.reshape(-1, k.shape[-2], k.shape[-1]) + + def safe_row_layout(tensor: torch.Tensor) -> bool: + heads = tensor.shape[1] + head_dim = tensor.shape[2] + head_stride = tensor.stride(1) + token_span = (heads - 1) * head_stride + head_dim + return tensor.stride(2) == 1 and (heads <= 1 or head_stride >= head_dim) and (tensor.shape[0] <= 1 or tensor.stride(0) >= token_span) + + if not safe_row_layout(q): + q = q.contiguous() + if not safe_row_layout(k): + k = k.contiguous() + q_rows = q.shape[0] * q.shape[-2] + k_rows = k.shape[0] * k.shape[-2] + if q_rows + k_rows == 0: + return q.reshape(q_shape), k.reshape(k_shape) + with torch.cuda.device(q.device): + _qk_neox_rope_kernel[(q_rows + k_rows,)]( + q, + k, + cache, + positions, + q_rows, + q.shape[-2], + k.shape[-2], + q.shape[-1], + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + positions.numel(), + ROPE_DIM=rotary_dim, + BLOCK_HALF=triton.next_power_of_2(rotary_dim // 2), + IS_BF16=q.dtype is torch.bfloat16, + num_warps=1, + ) + return q.reshape(q_shape), k.reshape(k_shape) + + +@ROPE_REGISTER("sgl_exact_neox_rope") +class SGLExactNeoXRope(RopeTemplate): + def __init__(self, layout="split_half", compute_dtype=torch.bfloat16): + if layout != "split_half": + raise ValueError("SGL exact NeoX RoPE requires split_half layout") + super().__init__(layout=layout, compute_dtype=compute_dtype) + + def prepare_freqs(self, freqs, rotary_dim: int | None = None): + if not isinstance(freqs, tuple) or len(freqs) != 2: + raise TypeError("SGL exact NeoX RoPE expects a (cos, sin) tuple") + if rotary_dim is None: + raise ValueError("rotary_dim is required for tuple RoPE frequencies") + if rotary_dim <= 0 or rotary_dim % 2: + raise ValueError(f"rotary_dim must be a positive even integer, got {rotary_dim}") + cos, sin = freqs + if not torch.is_tensor(cos) or not torch.is_tensor(sin): + raise TypeError("SGL exact NeoX RoPE cos/sin entries must be tensors") + if cos.ndim != 2 or sin.ndim != 2: + raise ValueError(f"SGL exact NeoX RoPE expects 2D cos/sin tensors, got {cos.shape} and {sin.shape}") + if cos.shape != sin.shape or cos.device != sin.device: + raise ValueError(f"SGL exact NeoX RoPE cos/sin tensors must match, got {cos.shape} and {sin.shape}") + if cos.shape[-1] == rotary_dim: + half = rotary_dim // 2 + cos, sin = cos[:, :half], sin[:, :half] + elif cos.shape[-1] != rotary_dim // 2: + raise ValueError(f"RoPE frequency width must be {rotary_dim // 2} or {rotary_dim}, got {cos.shape[-1]}") + cache = torch.cat((cos, sin), dim=-1) + cache = cache.to(dtype=self.compute_dtype).contiguous() + positions = torch.arange(cache.shape[0], device=cache.device, dtype=torch.long) + return cache, positions + + @staticmethod + def _is_prepared(freqs) -> bool: + return isinstance(freqs, tuple) and len(freqs) == 2 and torch.is_tensor(freqs[0]) and torch.is_tensor(freqs[1]) and freqs[0].ndim == 2 and freqs[1].dtype == torch.long + + def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, **kwargs): + if kwargs.get("materialize", False): + q, k = q.contiguous(), k.contiguous() + if not self._is_prepared(freqs): + freqs = self.prepare_freqs(freqs, rotary_dim=kwargs.get("rotary_dim")) + cache, positions = freqs + return _apply_qk_neox_rope(q, k, cache, positions) + + def validate_inputs(self, q: torch.Tensor, k: torch.Tensor, freqs): + if not self._is_prepared(freqs): + freqs = self.prepare_freqs(freqs) + cache, positions = freqs + return _prepare_qk_neox_rope_inputs(q, k, cache, positions) + + def apply_single(self, x: torch.Tensor, freqs, **kwargs) -> torch.Tensor: + return self.apply(x, torch.empty_like(x), freqs, **kwargs)[0] + + +__all__ = ["SGLExactNeoXRope"] diff --git a/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py index 7b82dafc2..6d94aa352 100644 --- a/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py +++ b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py @@ -23,10 +23,9 @@ from __future__ import annotations -import gc import json import math -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from pathlib import Path import torch @@ -47,6 +46,33 @@ def _empty_device_cache(device: torch.device) -> None: backend.empty_cache() +@contextmanager +def _deterministic_audio_decode_context(device: torch.device): + if device.type != "cuda": + yield + return + backends = torch.backends + previous = ( + backends.cudnn.allow_tf32, + backends.cuda.matmul.allow_tf32, + backends.cudnn.deterministic, + backends.cudnn.benchmark, + ) + backends.cudnn.allow_tf32 = False + backends.cuda.matmul.allow_tf32 = False + backends.cudnn.deterministic = True + backends.cudnn.benchmark = False + try: + yield + finally: + ( + backends.cudnn.allow_tf32, + backends.cuda.matmul.allow_tf32, + backends.cudnn.deterministic, + backends.cudnn.benchmark, + ) = previous + + def _component_dir(model_path: str | Path, component: str) -> Path: model_path = Path(model_path) nested = model_path / component @@ -62,6 +88,14 @@ def _wn_conv1d(*args, **kwargs) -> nn.Module: return weight_norm(nn.Conv1d(*args, **kwargs)) +@torch.jit.script +def _snakebeta(hidden_states: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor: + shape = hidden_states.shape + hidden_states = hidden_states.reshape(shape[0], shape[1], -1) + hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2) + return hidden_states.reshape(shape) + + def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: """Kaiser-windowed sinc filter, arithmetically identical to the release.""" @@ -93,7 +127,7 @@ def __init__(self, channels: int) -> None: def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: alpha = torch.exp(self.alpha.unsqueeze(0).unsqueeze(-1)) beta = torch.exp(self.beta.unsqueeze(0).unsqueeze(-1)) - return hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2) + return _snakebeta(hidden_states, alpha, beta) class MiniMaxH3AudioSnake1d(nn.Module): @@ -489,16 +523,14 @@ def encode(self, waveform: torch.Tensor, *, return_cpu: bool = True) -> torch.Te def _activate(self) -> torch.device: if self.cpu_offload: self.to(self.execution_device) - device = next(self.parameters()).device - dtype = next(self.parameters()).dtype - if dtype != torch.float32: - raise RuntimeError(f"MiniMax-H3 audio VAE weights must remain float32 for parity; found {dtype}. Move the module by device only, without a dtype cast.") - return device + parameter = next(self.parameters()) + if parameter.dtype != torch.float32: + raise RuntimeError(f"MiniMax-H3 audio VAE weights must remain float32 for parity; found {parameter.dtype}. Move the module by device only, without a dtype cast.") + return parameter.device def offload(self) -> None: self.to("cpu") _empty_device_cache(self.execution_device) - gc.collect() def _prepare_stereo_latents( self, @@ -537,7 +569,7 @@ def _run_decode( # Disable an ambient CUDA autocast: the released DAC/BigVGAN weights # and arithmetic stay FP32 (BF16 decodes are roughly 20 dB quieter). autocast_context = torch.autocast(device_type="cuda", enabled=False) if device.type == "cuda" else nullcontext() - with torch.no_grad(), autocast_context: + with torch.no_grad(), _deterministic_audio_decode_context(device), autocast_context: decoded = self.decoder(self.dec_in_proj(latents)).float() if stereo_groups is not None: diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py index 0d2c03c2f..57dc06289 100644 --- a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py @@ -139,6 +139,7 @@ def _qwen_linear(config, weight_name, *, tp_group, tp_rank, tp_size, split_dim, tp_size=tp_size, split_dim=split_dim, create_cuda_buffer=create_cuda_buffer, + fp32_reduce=bool(config.get("qwen3vl_fp32_reduce", False)), ) @@ -1027,7 +1028,12 @@ def load_vision_encoder(self): model_config = self._read_model_config(text_encoder_path) vision_config = dict(model_config["vision_config"]) logger.info(f"Building native MiniMax-H3 Qwen3-VL vision tower from {text_encoder_path}") - self.vision_encoder = MiniMaxH3Qwen3VLVisionTower.from_pretrained(text_encoder_path, vision_config) + self.vision_encoder = MiniMaxH3Qwen3VLVisionTower.from_pretrained( + text_encoder_path, + vision_config, + tp_group=self.tp_group, + fp32_reduce=bool(self.config.get("qwen3vl_fp32_reduce", False)), + ) return self.vision_encoder def unload_text_encoder(self): diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py index 9362cb098..970d716e9 100644 --- a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl_vision.py @@ -6,6 +6,7 @@ from pathlib import Path import torch +import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from loguru import logger @@ -65,6 +66,39 @@ def _rotate_half(value): return torch.cat((-second, first), dim=-1) +def _replace_parameter(module, name, value): + setattr(module, name, nn.Parameter(value.contiguous(), requires_grad=False)) + + +def _column_shard(value, tp_rank, tp_size): + if value.shape[0] % tp_size: + raise ValueError(f"Cannot column-shard shape {tuple(value.shape)} over vision TP size {tp_size}") + shard_size = value.shape[0] // tp_size + return value.narrow(0, tp_rank * shard_size, shard_size).clone() + + +def _row_shard(value, tp_rank, tp_size): + if value.shape[1] % tp_size: + raise ValueError(f"Cannot row-shard shape {tuple(value.shape)} over vision TP size {tp_size}") + shard_size = value.shape[1] // tp_size + return value.narrow(1, tp_rank * shard_size, shard_size) + + +def _row_parallel_linear(module, hidden_states, tp_group, tp_rank, tp_size, *, fp32_reduce=False): + if tp_size == 1: + return module(hidden_states) + # Add the row-parallel bias once before all-reduce, matching SGLang. + bias = module.bias if tp_rank == 0 else None + output = F.linear(hidden_states, module.weight, bias) + output_dtype = output.dtype + if fp32_reduce: + output = output.float() + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=tp_group) + if fp32_reduce: + output = output.to(output_dtype) + return output + + class _PatchEmbed(nn.Module): def __init__(self, config): super().__init__() @@ -80,14 +114,40 @@ def forward(self, pixels): class _VisionAttention(nn.Module): - def __init__(self, config): + def __init__(self, config, tp_group=None, tp_rank=0, tp_size=1, fp32_reduce=False): super().__init__() - self.num_heads = config["num_heads"] - self.head_dim = config["hidden_size"] // self.num_heads + self.tp_group = tp_group + self.tp_rank = tp_rank + self.tp_size = tp_size + self.fp32_reduce = fp32_reduce + total_num_heads = config["num_heads"] + if total_num_heads % tp_size: + raise ValueError(f"Qwen3-VL vision heads ({total_num_heads}) must be divisible by TP size ({tp_size})") + self.num_heads = total_num_heads // tp_size + self.head_dim = config["hidden_size"] // total_num_heads self.scaling = self.head_dim**-0.5 self.qkv = nn.Linear(config["hidden_size"], config["hidden_size"] * 3, bias=True) self.proj = nn.Linear(config["hidden_size"], config["hidden_size"], bias=True) + def shard_for_tensor_parallel(self): + if self.tp_size == 1: + return + hidden_size = self.qkv.weight.shape[1] + if hidden_size % self.tp_size: + raise ValueError(f"Qwen3-VL vision hidden size ({hidden_size}) must be divisible by TP size ({self.tp_size})") + local_size = hidden_size // self.tp_size + qkv_weight = torch.cat( + [self.qkv.weight.narrow(0, component * hidden_size + self.tp_rank * local_size, local_size) for component in range(3)], + dim=0, + ) + qkv_bias = torch.cat( + [self.qkv.bias.narrow(0, component * hidden_size + self.tp_rank * local_size, local_size) for component in range(3)], + dim=0, + ) + _replace_parameter(self.qkv, "weight", qkv_weight) + _replace_parameter(self.qkv, "bias", qkv_bias) + _replace_parameter(self.proj, "weight", _row_shard(self.proj.weight, self.tp_rank, self.tp_size)) + def forward(self, hidden_states, cu_seqlens, cos, sin): length = hidden_states.shape[0] query, key, value = self.qkv(hidden_states).reshape(length, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0) @@ -109,26 +169,42 @@ def forward(self, hidden_states, cu_seqlens, cos, sin): scale=self.scaling, ) outputs.append(out.transpose(1, 2).reshape(end - start, -1)) - return self.proj(torch.cat(outputs, dim=0)) + return _row_parallel_linear(self.proj, torch.cat(outputs, dim=0), self.tp_group, self.tp_rank, self.tp_size, fp32_reduce=self.fp32_reduce) class _VisionMLP(nn.Module): - def __init__(self, config): + def __init__(self, config, tp_group=None, tp_rank=0, tp_size=1, fp32_reduce=False): super().__init__() + self.tp_group = tp_group + self.tp_rank = tp_rank + self.tp_size = tp_size + self.fp32_reduce = fp32_reduce self.linear_fc1 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=True) self.linear_fc2 = nn.Linear(config["intermediate_size"], config["hidden_size"], bias=True) + def shard_for_tensor_parallel(self): + if self.tp_size == 1: + return + _replace_parameter(self.linear_fc1, "weight", _column_shard(self.linear_fc1.weight, self.tp_rank, self.tp_size)) + _replace_parameter(self.linear_fc1, "bias", _column_shard(self.linear_fc1.bias, self.tp_rank, self.tp_size)) + _replace_parameter(self.linear_fc2, "weight", _row_shard(self.linear_fc2.weight, self.tp_rank, self.tp_size)) + def forward(self, hidden_states): - return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states), approximate="tanh")) + hidden_states = F.gelu(self.linear_fc1(hidden_states), approximate="tanh") + return _row_parallel_linear(self.linear_fc2, hidden_states, self.tp_group, self.tp_rank, self.tp_size, fp32_reduce=self.fp32_reduce) class _VisionBlock(nn.Module): - def __init__(self, config): + def __init__(self, config, tp_group=None, tp_rank=0, tp_size=1, fp32_reduce=False): super().__init__() self.norm1 = nn.LayerNorm(config["hidden_size"], eps=1e-6) self.norm2 = nn.LayerNorm(config["hidden_size"], eps=1e-6) - self.attn = _VisionAttention(config) - self.mlp = _VisionMLP(config) + self.attn = _VisionAttention(config, tp_group, tp_rank, tp_size, fp32_reduce) + self.mlp = _VisionMLP(config, tp_group, tp_rank, tp_size, fp32_reduce) + + def shard_for_tensor_parallel(self): + self.attn.shard_for_tensor_parallel() + self.mlp.shard_for_tensor_parallel() def forward(self, hidden_states, cu_seqlens, cos, sin): hidden_states = hidden_states + self.attn(self.norm1(hidden_states), cu_seqlens, cos, sin) @@ -136,8 +212,12 @@ def forward(self, hidden_states, cu_seqlens, cos, sin): class _PatchMerger(nn.Module): - def __init__(self, config, postshuffle=False): + def __init__(self, config, postshuffle=False, tp_group=None, tp_rank=0, tp_size=1, fp32_reduce=False): super().__init__() + self.tp_group = tp_group + self.tp_rank = tp_rank + self.tp_size = tp_size + self.fp32_reduce = fp32_reduce merged_size = config["hidden_size"] * config["spatial_merge_size"] ** 2 self.merged_size = merged_size self.postshuffle = postshuffle @@ -145,28 +225,65 @@ def __init__(self, config, postshuffle=False): self.linear_fc1 = nn.Linear(merged_size, merged_size) self.linear_fc2 = nn.Linear(merged_size, config["out_hidden_size"]) + def shard_for_tensor_parallel(self): + if self.tp_size == 1: + return + _replace_parameter(self.linear_fc1, "weight", _column_shard(self.linear_fc1.weight, self.tp_rank, self.tp_size)) + _replace_parameter(self.linear_fc1, "bias", _column_shard(self.linear_fc1.bias, self.tp_rank, self.tp_size)) + _replace_parameter(self.linear_fc2, "weight", _row_shard(self.linear_fc2.weight, self.tp_rank, self.tp_size)) + def forward(self, hidden_states): if self.postshuffle: hidden_states = self.norm(hidden_states.view(-1, self.merged_size)) else: hidden_states = self.norm(hidden_states).view(-1, self.merged_size) - return self.linear_fc2(F.gelu(self.linear_fc1(hidden_states))) + hidden_states = F.gelu(self.linear_fc1(hidden_states)) + return _row_parallel_linear(self.linear_fc2, hidden_states, self.tp_group, self.tp_rank, self.tp_size, fp32_reduce=self.fp32_reduce) class MiniMaxH3Qwen3VLVisionTower(nn.Module): - def __init__(self, config): + def __init__(self, config, tp_group=None, fp32_reduce=False): super().__init__() self.config = dict(config) + self.tp_size = dist.get_world_size(tp_group) if tp_group is not None else 1 + self.tp_rank = dist.get_rank(tp_group) if tp_group is not None else 0 self.spatial_merge_size = int(config["spatial_merge_size"]) self.patch_embed = _PatchEmbed(config) self.pos_embed = nn.Embedding(config["num_position_embeddings"], config["hidden_size"]) - self.blocks = nn.ModuleList([_VisionBlock(config) for _ in range(config["depth"])]) - self.merger = _PatchMerger(config) + self.blocks = nn.ModuleList([_VisionBlock(config, tp_group, self.tp_rank, self.tp_size, fp32_reduce) for _ in range(config["depth"])]) + self.merger = _PatchMerger( + config, + tp_group=tp_group, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + fp32_reduce=fp32_reduce, + ) self.deepstack_visual_indexes = list(config["deepstack_visual_indexes"]) - self.deepstack_merger_list = nn.ModuleList([_PatchMerger(config, postshuffle=True) for _ in self.deepstack_visual_indexes]) + self.deepstack_merger_list = nn.ModuleList( + [ + _PatchMerger( + config, + postshuffle=True, + tp_group=tp_group, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + fp32_reduce=fp32_reduce, + ) + for _ in self.deepstack_visual_indexes + ] + ) head_dim = config["hidden_size"] // config["num_heads"] self.register_buffer("rotary_inv_freq", 1.0 / (10000.0 ** (torch.arange(0, head_dim // 2, 2).float() / (head_dim // 2))), persistent=False) + def shard_for_tensor_parallel(self): + if self.tp_size == 1: + return + for block in self.blocks: + block.shard_for_tensor_parallel() + self.merger.shard_for_tensor_parallel() + for merger in self.deepstack_merger_list: + merger.shard_for_tensor_parallel() + def forward(self, pixels, grid_thw): grid_thw = grid_thw.to(device=pixels.device) indices, weights = _bilinear_indices_weights(grid_thw, int(math.sqrt(self.config["num_position_embeddings"])), self.spatial_merge_size) @@ -187,10 +304,10 @@ def forward(self, pixels, grid_thw): return self.merger(hidden_states), deepstack @classmethod - def from_pretrained(cls, text_encoder_path, vision_config): + def from_pretrained(cls, text_encoder_path, vision_config, tp_group=None, fp32_reduce=False): root = Path(text_encoder_path) with torch.device("meta"): - model = cls(vision_config) + model = cls(vision_config, tp_group=tp_group, fp32_reduce=fp32_reduce) with (root / "model.safetensors.index.json").open("r", encoding="utf-8") as handle: weight_map = json.load(handle)["weight_map"] prefix = "model.visual." @@ -204,12 +321,13 @@ def from_pretrained(cls, text_encoder_path, vision_config): with safe_open(root / shard, framework="pt", device="cpu") as checkpoint: for name in shard_names: state[name[len(prefix) :]] = checkpoint.get_tensor(name) - missing, unexpected = model.load_state_dict(state, strict=False, assign=True) - # rotary_inv_freq is non-persistent, so every persistent tensor must match. - if missing or unexpected: - raise RuntimeError(f"Qwen3-VL vision checkpoint mismatch: missing={missing}, unexpected={unexpected}") + model.load_state_dict(state, assign=True) + model.shard_for_tensor_parallel() + del state head_dim = vision_config["hidden_size"] // vision_config["num_heads"] model.rotary_inv_freq = 1.0 / (10000.0 ** (torch.arange(0, head_dim // 2, 2, dtype=torch.float32) / (head_dim // 2))) + if model.tp_size > 1: + logger.info("Sharded native Qwen3-VL vision tower over rank {}/{}", model.tp_rank, model.tp_size) return model.eval().requires_grad_(False) diff --git a/lightx2v/models/networks/minimax_h3/config.py b/lightx2v/models/networks/minimax_h3/config.py new file mode 100644 index 000000000..f44e66770 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/config.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +_REMOVED_EXECUTION_OPTIONS = ( + "sgl_aligned", + "h3_packed_sequence_alignment", + "h3_rng_mode", + "h3_step_update", + "sglang_compatible_export", + "keep_latents_dtype_in_scheduler", +) + + +@dataclass(frozen=True) +class MiniMaxH3ExecutionProfile: + sampling_profile: str + packed_sequence_alignment: int = 64 + + +def resolve_minimax_h3_execution_profile(config: Mapping[str, Any]) -> MiniMaxH3ExecutionProfile: + """Resolve the fixed H3 execution contract and its checkpoint sampling profile.""" + if "h3_sglang_parity_ops" in config: + raise ValueError("MiniMax-H3 h3_sglang_parity_ops was removed; the aligned operators are always enabled") + if "h3_ops" in config: + raise ValueError("MiniMax-H3 h3_ops was removed; leaf backends use their standard registries") + + removed = [name for name in _REMOVED_EXECUTION_OPTIONS if name in config] + if removed: + names = ", ".join(removed) + raise ValueError( + f"MiniMax-H3 execution settings {names} were removed; aligned execution is the only implementation. " + "Use h3_sampling_profile='base' or 'dmd' to select the checkpoint sampling contract." + ) + + sampling_profile = config.get("h3_sampling_profile", "base") + if sampling_profile not in {"base", "dmd"}: + raise ValueError(f"MiniMax-H3 h3_sampling_profile must be 'base' or 'dmd', got {sampling_profile!r}") + return MiniMaxH3ExecutionProfile(sampling_profile=sampling_profile) + + +__all__ = [ + "MiniMaxH3ExecutionProfile", + "resolve_minimax_h3_execution_profile", +] diff --git a/lightx2v/models/networks/minimax_h3/infer/module_io.py b/lightx2v/models/networks/minimax_h3/infer/module_io.py index f19ca11d9..e8aa3141a 100644 --- a/lightx2v/models/networks/minimax_h3/infer/module_io.py +++ b/lightx2v/models/networks/minimax_h3/infer/module_io.py @@ -24,7 +24,9 @@ class MiniMaxH3PreInferOutput: video_indices: torch.Tensor audio_indices: torch.Tensor text_indices: torch.Tensor + cu_seqlens: torch.Tensor norm_out_modulation: torch.Tensor | None = None + prepared_rotary_emb: object | None = None sequence_parallel_state: MiniMaxH3SequenceParallelState | None = None diff --git a/lightx2v/models/networks/minimax_h3/infer/post_infer.py b/lightx2v/models/networks/minimax_h3/infer/post_infer.py index 9143df184..8e17c0c30 100644 --- a/lightx2v/models/networks/minimax_h3/infer/post_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/post_infer.py @@ -1,16 +1,31 @@ +import torch.distributed as dist import torch.nn.functional as F from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3VelocityOutput +from lightx2v.models.networks.minimax_h3.infer.sgl_exact_ops import sgl_exact_indexed_scale_shift +from lightx2v.models.networks.minimax_h3.infer.tensor_parallel import all_gather_last_dim from lightx2v.utils.envs import GET_DTYPE class MiniMaxH3PostInfer: def __init__(self, config): self.config = config + self.tp_group = None + self.tp_size = 1 + if config.get("tensor_parallel", False): + self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + self.tp_size = dist.get_world_size(self.tp_group) def set_scheduler(self, scheduler): self.scheduler = scheduler + def _gather_tp_last_dim(self, tensor): + return all_gather_last_dim(tensor, self.tp_group, self.tp_size) + + @staticmethod + def _apply_modulation(hidden_states, shift, scale, indices): + return sgl_exact_indexed_scale_shift(hidden_states, shift, scale, indices) + def infer(self, weights, hidden_states, pre_infer_out): modulation = pre_infer_out.norm_out_modulation if modulation is None: @@ -19,15 +34,17 @@ def infer(self, weights, hidden_states, pre_infer_out): if pre_infer_out.temb is None: raise RuntimeError("MiniMax-H3 final-norm modulation is missing") modulation = weights.norm_out_linear.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())) + modulation = self._gather_tp_last_dim(modulation) shift, scale = modulation.chunk(2, dim=-1) indices = pre_infer_out.timestep_indices hidden_states = weights.norm_out.apply(hidden_states) - hidden_states = hidden_states * (1.0 + scale.index_select(0, indices)) - hidden_states = hidden_states + shift.index_select(0, indices) + hidden_states = self._apply_modulation(hidden_states, shift, scale, indices) - # Both released output heads are fp32 and run over all packed rows - # before modality selection. hidden_states = hidden_states.float() - video = weights.proj_out.apply(hidden_states).index_select(0, pre_infer_out.video_indices) - audio = weights.audio_proj_out.apply(hidden_states).index_select(0, pre_infer_out.audio_indices) + video = weights.proj_out.apply(hidden_states) + audio = weights.audio_proj_out.apply(hidden_states) + video = video.index_select(0, pre_infer_out.video_indices) + audio = audio.index_select(0, pre_infer_out.audio_indices) + video = self._gather_tp_last_dim(video) + audio = self._gather_tp_last_dim(audio) return MiniMaxH3VelocityOutput(video=video, audio=audio) diff --git a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py index f3ae92323..71533a650 100644 --- a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py @@ -5,6 +5,8 @@ import torch.nn.functional as F from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3PreInferOutput +from lightx2v.models.networks.minimax_h3.infer.sgl_exact_ops import _silu_mul_with_activation_rounding_inplace +from lightx2v.models.networks.minimax_h3.infer.tensor_parallel import all_gather_last_dim, row_parallel_linear from lightx2v.utils.envs import GET_DTYPE @@ -26,12 +28,15 @@ class MiniMaxH3PreInfer: def __init__(self, config): self.config = config global_num_heads = int(config.get("num_attention_heads", 56)) + self.tp_group = None + self.tp_rank = 0 if config.get("tensor_parallel", False): - tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") - tp_size = dist.get_world_size(tp_group) + self.tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + self.tp_rank = dist.get_rank(self.tp_group) + self.tp_size = dist.get_world_size(self.tp_group) else: - tp_size = 1 - self.num_heads = global_num_heads // tp_size + self.tp_size = 1 + self.num_heads = global_num_heads // self.tp_size self.head_dim = int(config.get("attention_head_dim", 128)) self.hidden_size = int(config.get("hidden_size", 5376)) self.rope_freq_dim = int(config.get("rope_freq_dim", 16)) @@ -42,10 +47,16 @@ def __init__(self, config): def set_scheduler(self, scheduler): self.scheduler = scheduler + @staticmethod + def _project_qkv(weights, hidden_states): + projected = weights.qkv.apply(hidden_states) + return weights.qkv.split_qkv(projected) + def _attention(self, weights, hidden_states): - q = weights.to_q.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) - k = weights.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) - v = weights.to_v.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + q, k, v = self._project_qkv(weights, hidden_states) + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + k = k.unflatten(-1, (self.num_heads, self.head_dim)) + v = v.unflatten(-1, (self.num_heads, self.head_dim)) q = weights.norm_q.apply(q) k = weights.norm_k.apply(k) seq_len = q.shape[0] @@ -59,13 +70,21 @@ def _attention(self, weights, hidden_states): max_seqlen_q=seq_len, max_seqlen_kv=seq_len, causal=False, + softmax_scale=self.head_dim**-0.5, ) return weights.to_out.apply(out.to(GET_DTYPE())) @staticmethod def _ff(weights, hidden_states): - value, gate = weights.in_proj.apply(hidden_states).chunk(2, dim=-1) - return weights.out_proj.apply(value * F.silu(gate)) + hidden_states = weights.in_proj.apply(hidden_states) + hidden_states = _silu_mul_with_activation_rounding_inplace(hidden_states) + return weights.out_proj.apply(hidden_states) + + def _gather_tp_last_dim(self, tensor): + return all_gather_last_dim(tensor, self.tp_group, self.tp_size) + + def _apply_time_linear_2(self, module, hidden_states): + return row_parallel_linear(module, hidden_states, self.tp_group, self.tp_rank, self.tp_size) def _refine_text(self, weights, text_embeds): for block in weights.refiner_blocks: @@ -98,9 +117,14 @@ def infer(self, weights, prompt_embeds): layout = self.scheduler.layout bulk_dtype = GET_DTYPE() - video_embeds = weights.proj_in.apply(self.scheduler.video_latents.float()).to(bulk_dtype) - audio_embeds = weights.audio_proj_in.apply(self.scheduler.audio_latents.float()).to(bulk_dtype) + video_embeds = weights.proj_in.apply(self.scheduler.video_latents.float()) + audio_embeds = weights.audio_proj_in.apply(self.scheduler.audio_latents.float()) text_embeds = weights.context_embedder.apply(prompt_embeds.to(bulk_dtype)) + video_embeds = self._gather_tp_last_dim(video_embeds) + audio_embeds = self._gather_tp_last_dim(audio_embeds) + text_embeds = self._gather_tp_last_dim(text_embeds) + video_embeds = video_embeds.to(bulk_dtype) + audio_embeds = audio_embeds.to(bulk_dtype) text_embeds = self._refine_text(weights, text_embeds) hidden_states = text_embeds.new_zeros((layout.sequence_length, self.hidden_size)) @@ -114,7 +138,8 @@ def infer(self, weights, prompt_embeds): # or dtype must also be made in the offline AdaLN cache builder and # followed by regenerating the cache when cached values can change. temb = timestep_embedding(self.scheduler.unique_timesteps, self.freq_dim) - temb = weights.time_linear_2.apply(F.silu(weights.time_linear_1.apply(temb.float()))) + time_hidden = F.silu(weights.time_linear_1.apply(temb.float())) + temb = self._apply_time_linear_2(weights.time_linear_2, time_hidden) timestep_indices = self.scheduler.timestep_indices adaln_indices = timestep_indices * 3 + layout.token_tags.clamp(min=0) @@ -127,4 +152,5 @@ def infer(self, weights, prompt_embeds): video_indices=layout.video_indices, audio_indices=layout.audio_indices, text_indices=layout.text_indices, + cu_seqlens=layout.cu_seqlens, ) diff --git a/lightx2v/models/networks/minimax_h3/infer/sgl_exact_ops.py b/lightx2v/models/networks/minimax_h3/infer/sgl_exact_ops.py new file mode 100644 index 000000000..d8122bb19 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/sgl_exact_ops.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical kernels used by MiniMax-H3 DiT inference.""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +# The numerical kernels below are adapted from SGLang commit +# 8ef646a5c65bd2f8922483057dddc02e2b0de18c (Apache-2.0). Keeping the +# H3-specific subset here avoids importing an SGLang checkout at runtime. + + +def _require_nvidia_triton(tensor: torch.Tensor, operation: str) -> None: + if tensor.device.type != "cuda": + return + if getattr(torch.version, "hip", None) is not None: + raise RuntimeError(f"{operation} exact parity kernel supports NVIDIA CUDA only") + + +@triton.jit +def _round_bf16_to_fp32(value): + """RNE-round FP32 to BF16 precision while retaining an FP32 register.""" + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _indexed_scale_shift_bf16_kernel( + x_ptr, + shift_ptr, + scale_ptr, + indices_ptr, + hidden_size, + stride_x_row, + stride_shift_row, + stride_scale_row, + stride_indices, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, BLOCK_N) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32) + one_plus_scale = _round_bf16_to_fp32(1.0 + scale) + scaled = _round_bf16_to_fp32(x * one_plus_scale) + tl.store(x_ptr + row * stride_x_row + columns, scaled + shift, mask=mask) + + +@triton.jit +def _indexed_gate_bf16_kernel( + x_ptr, + gate_ptr, + other_ptr, + indices_ptr, + hidden_size, + stride_x_row, + stride_gate_row, + stride_other_row, + stride_indices, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, BLOCK_N) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32) + other = tl.load(other_ptr + row * stride_other_row + columns, mask=mask, other=0.0).to(tl.float32) + gated = _round_bf16_to_fp32(gate * other) + tl.store(x_ptr + row * stride_x_row + columns, x + gated, mask=mask) + + +@triton.jit +def _packed_silu_mul_kernel( + output_ptr, + x_ptr, + num_rows, + row_stride, + output_row_stride, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + block = tl.program_id(1).to(tl.int64) + columns = block * BLOCK + tl.arange(0, BLOCK) + mask = (row < num_rows) & (columns < D) + row_base = row * row_stride + gate = tl.load(x_ptr + row_base + columns, mask=mask, other=0.0).to(tl.float32) + value = tl.load(x_ptr + row_base + D + columns, mask=mask, other=0.0).to(tl.float32) + activated = _round_bf16_to_fp32(gate * tl.sigmoid(gate)) + tl.store(output_ptr + row * output_row_stride + columns, activated * value, mask=mask) + + +def _silu_mul_with_activation_rounding_inplace(hidden_states: torch.Tensor) -> torch.Tensor: + if hidden_states.shape[-1] % 2: + raise ValueError(f"SwiGLU input width must be even, got {hidden_states.shape[-1]}") + hidden_size = hidden_states.shape[-1] // 2 + if hidden_states.is_cuda and hidden_states.dtype is torch.bfloat16 and hidden_states.is_contiguous() and hidden_states.numel(): + _require_nvidia_triton(hidden_states, "H3 SwiGLU") + rows = hidden_states.numel() // hidden_states.shape[-1] + with torch.cuda.device(hidden_states.device): + _packed_silu_mul_kernel[(rows, triton.cdiv(hidden_size, 1024))]( + hidden_states, + hidden_states, + rows, + hidden_states.shape[-1], + hidden_states.shape[-1], + D=hidden_size, + BLOCK=1024, + ) + return hidden_states[..., :hidden_size] + + gate, value = hidden_states.chunk(2, dim=-1) + F.silu(gate, inplace=True) + return gate.mul_(value) + + +def _validate_indexed_modulation_inputs( + operation: str, + x: torch.Tensor, + indices: torch.Tensor, + lookup_tensors: tuple[tuple[str, torch.Tensor], ...], + row_tensors: tuple[tuple[str, torch.Tensor], ...] = (), +) -> None: + if x.ndim != 2: + raise ValueError(f"{operation} expects a two-dimensional activation, got {x.shape}") + if x.dtype is not torch.bfloat16: + raise TypeError(f"{operation} requires BF16 activations, got {x.dtype}") + if indices.ndim != 1 or indices.shape[0] != x.shape[0]: + raise ValueError(f"{operation} indices must have shape ({x.shape[0]},), got {indices.shape}") + if indices.dtype is not torch.long: + raise TypeError(f"{operation} indices must use torch.long, got {indices.dtype}") + if indices.device != x.device: + raise ValueError(f"{operation} tensors must be on one device") + + hidden_size = x.shape[1] + tensors = (("activation", x), *lookup_tensors, *row_tensors) + for name, tensor in tensors: + if tensor.device != x.device: + raise ValueError(f"{operation} tensors must be on one device; {name} is on {tensor.device}") + if tensor.dtype is not torch.bfloat16: + raise TypeError(f"{operation} requires BF16 {name}, got {tensor.dtype}") + if tensor.ndim != 2 or tensor.shape[1] != hidden_size: + raise ValueError(f"{operation} {name} must be two-dimensional with width {hidden_size}, got {tensor.shape}") + if tensor.stride(-1) != 1 or (tensor.shape[0] > 1 and tensor.stride(0) < hidden_size): + raise ValueError(f"{operation} {name} must have a non-overlapping contiguous last dimension, got stride {tensor.stride()}") + + lookup_rows = {tensor.shape[0] for _, tensor in lookup_tensors} + if len(lookup_rows) > 1 or (x.shape[0] and lookup_rows == {0}): + raise ValueError(f"{operation} lookup tensors must have the same nonzero row count") + # H3 builds these indices from torch.unique(return_inverse=True); AdaLN + # indices additionally include a bounded token tag. Their range is + # guaranteed at the producer, while min/max here would synchronize the GPU + # on every transformer block. + for name, tensor in row_tensors: + if tensor.shape != x.shape: + raise ValueError(f"{operation} {name} must match activation shape {x.shape}, got {tensor.shape}") + + +def sgl_exact_indexed_scale_shift(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + _validate_indexed_modulation_inputs( + "H3 indexed scale/shift", + x, + indices, + (("shift", shift), ("scale", scale)), + ) + if x.numel() == 0: + return x + if x.device.type != "cuda": + selected_scale = scale.index_select(0, indices).float() + selected_shift = shift.index_select(0, indices).float() + one_plus_scale = (1.0 + selected_scale).to(torch.bfloat16).float() + scaled = (x.float() * one_plus_scale).to(torch.bfloat16).float() + return x.copy_((scaled + selected_shift).to(x.dtype)) + _require_nvidia_triton(x, "H3 indexed scale/shift") + rows, hidden_size = x.shape + with torch.cuda.device(x.device): + _indexed_scale_shift_bf16_kernel[(rows,)]( + x, + shift, + scale, + indices, + hidden_size, + x.stride(0), + shift.stride(0), + scale.stride(0), + indices.stride(0), + BLOCK_N=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return x + + +def sgl_exact_indexed_gate(x: torch.Tensor, gate: torch.Tensor, other: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + _validate_indexed_modulation_inputs( + "H3 indexed gate", + x, + indices, + (("gate", gate),), + (("other", other),), + ) + if x.numel() == 0: + return x + if x.device.type != "cuda": + selected_gate = gate.index_select(0, indices).float() + gated = (selected_gate * other.float()).to(torch.bfloat16).float() + return x.copy_((x.float() + gated).to(x.dtype)) + _require_nvidia_triton(x, "H3 indexed gate") + rows, hidden_size = x.shape + with torch.cuda.device(x.device): + _indexed_gate_bf16_kernel[(rows,)]( + x, + gate, + other, + indices, + hidden_size, + x.stride(0), + gate.stride(0), + other.stride(0), + indices.stride(0), + BLOCK_N=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return x diff --git a/lightx2v/models/networks/minimax_h3/infer/tensor_parallel.py b/lightx2v/models/networks/minimax_h3/infer/tensor_parallel.py new file mode 100644 index 000000000..35ba8599d --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/infer/tensor_parallel.py @@ -0,0 +1,41 @@ +import torch +import torch.distributed as dist + +from lightx2v.common.ops.mm.mm_weight import unwrap_tp_weight + + +def all_gather_last_dim(tensor, group, world_size): + if world_size == 1: + return tensor + tensor = tensor.contiguous() + input_shape = list(tensor.shape) + gathered_shape = input_shape.copy() + gathered_shape[0] *= world_size + gathered = torch.empty(gathered_shape, dtype=tensor.dtype, device=tensor.device) + dist.all_gather_into_tensor(gathered, tensor, group=group) + gathered = gathered.reshape([world_size] + input_shape) + gathered = gathered.movedim(0, tensor.dim() - 1) + output_shape = input_shape.copy() + output_shape[-1] *= world_size + return gathered.reshape(output_shape) + + +def row_parallel_linear(module, tensor, group, rank, world_size): + if world_size == 1: + return module.apply(tensor) + + # MMWeightTP keeps row-parallel bias outside its concrete backend. Run the + # backend locally so quantization, diff, and LoRA are all included, then + # place the replicated bias on rank 0 before the reduction to match H3's + # reference operation order. + concrete = unwrap_tp_weight(module) + output = concrete.apply(tensor) + if rank == 0: + bias = concrete._get_actual_bias(module._row_split_bias) + if bias is not None: + output = output + bias + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=group) + return output + + +__all__ = ["all_gather_last_dim", "row_parallel_linear"] diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index 299d4c3ad..ba56cf19d 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -4,6 +4,12 @@ from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.minimax_h3.adaln_cache import load_persistent_adaln_cache +from lightx2v.models.networks.minimax_h3.infer.sgl_exact_ops import ( + _silu_mul_with_activation_rounding_inplace, + sgl_exact_indexed_gate, + sgl_exact_indexed_scale_shift, +) +from lightx2v.models.networks.minimax_h3.infer.tensor_parallel import all_gather_last_dim from lightx2v.utils.envs import GET_DTYPE from lightx2v_platform.base.global_var import AI_DEVICE @@ -50,41 +56,52 @@ def _cache_device(): return torch.device(AI_DEVICE, device_module.current_device()) def _gather_tp_last_dim(self, tensor): - if self.tp_size == 1: - return tensor - gathered = [torch.empty_like(tensor) for _ in range(self.tp_size)] - dist.all_gather(gathered, tensor.contiguous(), group=self.tp_group) - return torch.cat(gathered, dim=-1) + return all_gather_last_dim(tensor, self.tp_group, self.tp_size) - def _attention(self, weights, hidden_states, pre_infer_out): - q = weights.to_q.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) - k = weights.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) - v = weights.to_v.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + @staticmethod + def _project_qkv(weights, hidden_states): + projected = weights.qkv.apply(hidden_states) + return weights.qkv.split_qkv(projected) + + def _apply_qk_norm_rope(self, weights, q, k, pre_infer_out): + if pre_infer_out.prepared_rotary_emb is None: + pre_infer_out.prepared_rotary_emb = weights.rope.prepare_freqs( + pre_infer_out.rotary_emb, + rotary_dim=pre_infer_out.rotary_emb[0].shape[-1], + ) + pre_infer_out.prepared_rotary_emb = weights.rope.validate_inputs(q, k, pre_infer_out.prepared_rotary_emb) q = weights.norm_q.apply(q) k = weights.norm_k.apply(k) - q, k = weights.rope.apply( + return weights.rope.apply( q, k, - pre_infer_out.rotary_emb, + pre_infer_out.prepared_rotary_emb, rotary_dim=pre_infer_out.rotary_emb[0].shape[-1], ) + + def _attention(self, weights, hidden_states, pre_infer_out): + q, k, v = self._project_qkv(weights, hidden_states) + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + k = k.unflatten(-1, (self.num_heads, self.head_dim)) + v = v.unflatten(-1, (self.num_heads, self.head_dim)) + q, k = self._apply_qk_norm_rope(weights, q, k, pre_infer_out) sp_state = pre_infer_out.sequence_parallel_state + used_seq_len = self.scheduler.layout.used_sequence_length attention_kwargs = { "causal": False, "scheduler": self.scheduler, "block_idx": self.block_idx, + "softmax_scale": self.head_dim**-0.5, + "cu_seqlens_q": pre_infer_out.cu_seqlens, + "cu_seqlens_kv": pre_infer_out.cu_seqlens, + "max_seqlen_q": used_seq_len, + "max_seqlen_kv": used_seq_len, } if sp_state is None: - seq_len = q.shape[0] - cu_seqlens = torch.tensor((0, seq_len), dtype=torch.int32, device=q.device) out = weights.calculate.apply( q=q, k=k, v=v, - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=seq_len, - max_seqlen_kv=seq_len, **attention_kwargs, ) else: @@ -112,11 +129,19 @@ def _attention(self, weights, hidden_states, pre_infer_out): @staticmethod def _ff(weights, hidden_states): - value, gate = weights.in_proj.apply(hidden_states).chunk(2, dim=-1) - return weights.out_proj.apply(value * F.silu(gate)) + hidden_states = weights.in_proj.apply(hidden_states) + hidden_states = _silu_mul_with_activation_rounding_inplace(hidden_states) + return weights.out_proj.apply(hidden_states) + + @staticmethod + def _apply_modulation(hidden_states, shift, scale, indices): + return sgl_exact_indexed_scale_shift(hidden_states, shift, scale, indices) + + @staticmethod + def _apply_residual(residual, gate, branch, indices): + return sgl_exact_indexed_gate(residual, gate, branch, indices) def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None): - # Keep the Python cache lookup outside the compiled block. if modulation is None: modulation = self._compute_adaln_table(weights, pre_infer_out) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1) @@ -124,16 +149,15 @@ def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None): residual = hidden_states normed = weights.norm1.apply(hidden_states) - normed = normed * (1.0 + scale_msa.index_select(0, indices)) - normed = normed + shift_msa.index_select(0, indices) - hidden_states = residual + gate_msa.index_select(0, indices) * self._attention(weights.attn, normed, pre_infer_out) + normed = self._apply_modulation(normed, shift_msa, scale_msa, indices) + attention_output = self._attention(weights.attn, normed, pre_infer_out) + hidden_states = self._apply_residual(residual, gate_msa, attention_output, indices) residual = hidden_states normed = weights.norm2.apply(hidden_states) - normed = normed * (1.0 + scale_mlp.index_select(0, indices)) - normed = normed + shift_mlp.index_select(0, indices) - hidden_states = residual + gate_mlp.index_select(0, indices) * self._ff(weights.ff, normed) - return hidden_states + normed = self._apply_modulation(normed, shift_mlp, scale_mlp, indices) + ff_output = self._ff(weights.ff, normed) + return self._apply_residual(residual, gate_mlp, ff_output, indices) def _compute_adaln_table(self, weights, pre_infer_out): # ADALN CACHE SYNC: This projection is reproduced by the offline builder. diff --git a/lightx2v/models/networks/minimax_h3/infer/triton_ops.py b/lightx2v/models/networks/minimax_h3/infer/triton_ops.py deleted file mode 100644 index 1d4c618d0..000000000 --- a/lightx2v/models/networks/minimax_h3/infer/triton_ops.py +++ /dev/null @@ -1,137 +0,0 @@ -import torch - -from lightx2v.common.ops.rope import RopeTemplate, TorchRealRope -from lightx2v.utils.registry_factory import ROPE_REGISTER - -try: - import triton # type: ignore - import triton.language as tl # type: ignore -except ImportError: - triton = None - tl = None - - -@ROPE_REGISTER("minimax_h3_triton_rope") -class MiniMaxH3TritonRope(RopeTemplate): - """Partial split-half RoPE used by MiniMax-H3. - - H3 rotates only the leading RoPE dimensions of each attention head and - leaves the remaining channels unchanged. CUDA tensors use the local - Triton kernel; other devices (or environments without Triton) fall back - to the shared real-valued RoPE implementation with identical layout. - """ - - def __init__(self, layout="split_half", compute_dtype=torch.float32): - super().__init__(layout=layout, compute_dtype=compute_dtype) - if layout != "split_half": - raise ValueError("MiniMaxH3TritonRope only supports split_half layout") - self.torch_rope = TorchRealRope(layout=layout, compute_dtype=compute_dtype) - - def apply(self, query, key, freqs, rotary_dim=None, **kwargs): - cos, sin = freqs - rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim - if query.is_cuda and key.is_cuda and triton is not None: - return ( - apply_partial_split_half_rotary_triton(query, cos, sin, rotary_dim), - apply_partial_split_half_rotary_triton(key, cos, sin, rotary_dim), - ) - return self.torch_rope.apply(query, key, freqs, rotary_dim=rotary_dim, **kwargs) - - def apply_single(self, x, freqs, rotary_dim=None, **kwargs): - cos, sin = freqs - rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim - if x.is_cuda and triton is not None: - return apply_partial_split_half_rotary_triton(x, cos, sin, rotary_dim) - return self.torch_rope.apply_single(x, freqs, rotary_dim=rotary_dim, **kwargs) - - -if triton is not None: - - @triton.jit - def _partial_split_half_rotary_kernel( - output_ptr, - x_ptr, - cos_ptr, - sin_ptr, - num_heads, - num_tokens, - stride_x_row, - stride_cos_row, - stride_sin_row, - HEAD_SIZE: tl.constexpr, - ROTARY_DIM: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - ): - row_idx = tl.program_id(0) - token_idx = (row_idx // num_heads) % num_tokens - - offsets = tl.arange(0, BLOCK_SIZE) - head_mask = offsets < HEAD_SIZE - rotary_mask = offsets < ROTARY_DIM - rotary_half = ROTARY_DIM // 2 - - x_row_ptr = x_ptr + row_idx * stride_x_row - output_row_ptr = output_ptr + row_idx * stride_x_row - cos_row_ptr = cos_ptr + token_idx * stride_cos_row - sin_row_ptr = sin_ptr + token_idx * stride_sin_row - - x = tl.load(x_row_ptr + offsets, mask=head_mask, other=0.0) - partner_offsets = tl.where(offsets < rotary_half, offsets + rotary_half, offsets - rotary_half) - partner = tl.load(x_row_ptr + partner_offsets, mask=rotary_mask, other=0.0) - rotated = tl.where(offsets < rotary_half, -partner, partner) - cos = tl.load(cos_row_ptr + offsets, mask=rotary_mask, other=1.0) - sin = tl.load(sin_row_ptr + offsets, mask=rotary_mask, other=0.0) - - x_fp32 = x.to(tl.float32) - rotated_fp32 = rotated.to(tl.float32) - rotated_output = x_fp32 * cos.to(tl.float32) + rotated_fp32 * sin.to(tl.float32) - output = tl.where(rotary_mask, rotated_output, x_fp32) - tl.store(output_row_ptr + offsets, output.to(x.dtype), mask=head_mask) - - -def apply_partial_split_half_rotary_triton( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - rotary_dim: int | None = None, -) -> torch.Tensor: - if triton is None: - raise RuntimeError("Triton is required for MiniMax-H3 Triton RoPE") - if not x.is_cuda: - raise ValueError("MiniMax-H3 Triton RoPE requires a CUDA tensor") - if x.ndim != 3: - raise ValueError(f"MiniMax-H3 Triton RoPE expects [L, H, D], got {tuple(x.shape)}") - if cos.shape != sin.shape or cos.ndim != 2: - raise ValueError(f"cos and sin must have matching [L, R] shapes, got {tuple(cos.shape)} and {tuple(sin.shape)}") - - num_tokens, num_heads, head_size = x.shape - rotary_dim = cos.shape[-1] if rotary_dim is None else int(rotary_dim) - if cos.shape[0] != num_tokens: - raise ValueError(f"RoPE token count ({cos.shape[0]}) does not match input ({num_tokens})") - if rotary_dim != cos.shape[-1]: - raise ValueError(f"rotary_dim ({rotary_dim}) must match the H3 cos/sin width ({cos.shape[-1]})") - if rotary_dim <= 0 or rotary_dim > head_size or rotary_dim % 2: - raise ValueError(f"rotary_dim must be positive, even, and <= head_size; got rotary_dim={rotary_dim}, head_size={head_size}") - - x = x.contiguous() - cos = cos.to(device=x.device).contiguous() - sin = sin.to(device=x.device).contiguous() - output = torch.empty_like(x) - block_size = triton.next_power_of_2(head_size) - grid = (num_tokens * num_heads,) - with torch.cuda.device(x.device): - _partial_split_half_rotary_kernel[grid]( - output, - x, - cos, - sin, - num_heads, - num_tokens, - x.stride(1), - cos.stride(0), - sin.stride(0), - HEAD_SIZE=head_size, - ROTARY_DIM=rotary_dim, - BLOCK_SIZE=block_size, - ) - return output diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 94d02ed4c..d0e0df6e1 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -9,6 +9,7 @@ from lightx2v.models.networks.base_model import BaseTransformerModel from lightx2v.models.networks.minimax_h3.adaln_cache import validate_adaln_cache_config +from lightx2v.models.networks.minimax_h3.config import resolve_minimax_h3_execution_profile from lightx2v.models.networks.minimax_h3.infer.module_io import MiniMaxH3SequenceParallelState from lightx2v.models.networks.minimax_h3.infer.offload import MiniMaxH3OffloadTransformerInfer from lightx2v.models.networks.minimax_h3.infer.post_infer import MiniMaxH3PostInfer @@ -39,9 +40,20 @@ "int8-convrot", } +_H3_TP_SPLITS = { + "proj_in": "col", + "audio_proj_in": "col", + "context_embedder": "col", + "time_embedder.linear_1": "col", + "time_embedder.linear_2": "row", + "norm_out.linear": "col", + "proj_out": "col", + "audio_proj_out": "col", +} + class MiniMaxH3Model(BaseTransformerModel): - """LightX2V-native MiniMax-H3 joint audio/video transformer.""" + """MiniMax-H3 joint audio/video transformer.""" pre_weight_class = MiniMaxH3PreWeights transformer_weight_class = MiniMaxH3TransformerWeights @@ -49,6 +61,7 @@ class MiniMaxH3Model(BaseTransformerModel): def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0, lora_alpha=None): self.lora_alpha = lora_alpha + self.execution_profile = resolve_minimax_h3_execution_profile(config) self.use_adaln_cache = bool(config.get("use_adaln_cache", False)) if config.get("cpu_offload", False) and not self.use_adaln_cache: separator = "=" * 88 @@ -80,7 +93,7 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 self.prepost_resident = self.block_offload and config.get("dit_prepost_resident", False) if GET_DTYPE() != torch.bfloat16: raise ValueError( - "MiniMax-H3 requires DTYPE=BF16. The native loader preserves the released checkpoint's 626 BF16 tensors and 12 FP32 projection/time/head tensors without dtype conversion." + "MiniMax-H3 requires DTYPE=BF16. The loader preserves the released checkpoint's 626 BF16 tensors and 12 FP32 projection/time/head tensors without dtype conversion." ) if config.get("cfg_parallel", False) or config.get("enable_cfg", False): raise ValueError("MiniMax-H3 is guidance-distilled and does not have a CFG/unconditional branch") @@ -283,7 +296,15 @@ def _register_dynamic_lora_weights(self, lora_weights, strength): tensor = tensor.to("cpu") setattr(weight, attr_name, tensor.pin_memory()) - registered = {weight.weight_name for weight in weights if getattr(weight, "has_lora_branch", False)} + registered = set() + for weight in weights: + if not getattr(weight, "has_lora_branch", False): + continue + source_names = getattr(weight, "registered_source_weight_names", None) + if source_names is None: + registered.add(weight.weight_name) + else: + registered.update(source_names) missing = sorted(self._pending_dynamic_lora_model_keys - registered) if missing: self._remove_lora() @@ -331,8 +352,10 @@ def _validate_tensor_parallel_config(self): details = ", ".join(f"{name}={value}" for name, value in invalid.items()) raise ValueError(f"MiniMax-H3 TP size {self.tp_size} must divide {details}") - @staticmethod - def _tp_split_type(key): + def _tp_split_type(self, key): + for prefix, split_type in _H3_TP_SPLITS.items(): + if key == prefix or key.startswith(f"{prefix}."): + return split_type if ".attn.to_q." in key or ".attn.to_k." in key or ".attn.to_v." in key: return "col" if ".attn.to_out.0." in key: @@ -592,6 +615,8 @@ def _seq_parallel_post_process(self, output, pre_infer_out): def to_cpu(self): super().to_cpu() + if self.cpu_offload: + self.transformer_infer._clear_adaln_cache() if hasattr(self.transformer_infer, "offload_manager"): # Full teardown moves the active aliases away from the persistent # device buffers. Force buffer 0 to be populated again next run. diff --git a/lightx2v/models/networks/minimax_h3/packing.py b/lightx2v/models/networks/minimax_h3/packing.py index 9024a3382..eebff558b 100644 --- a/lightx2v/models/networks/minimax_h3/packing.py +++ b/lightx2v/models/networks/minimax_h3/packing.py @@ -15,6 +15,7 @@ VIDEO_TAG = 0 TEXT_TAG = 1 AUDIO_TAG = 2 +PADDING_TAG = -1 FPS = 24 AUDIO_LATENTS_PER_SECOND = 40 @@ -47,9 +48,23 @@ class MiniMaxH3PackedSequence: video_indices: torch.Tensor audio_indices: torch.Tensor text_indices: torch.Tensor + used_sequence_length: int num_condition_video_rows: int = 0 num_condition_audio_rows: int = 0 + @property + def cu_seqlens(self) -> torch.Tensor: + bounds = (0, self.used_sequence_length) + if self.used_sequence_length < self.sequence_length: + bounds += (self.sequence_length,) + return torch.tensor(bounds, dtype=torch.int32, device=self.position_ids.device) + + +def _align_sequence_length(sequence_length: int, alignment: int) -> int: + if alignment < 1: + raise ValueError(f"MiniMax-H3 packed-sequence alignment must be positive, got {alignment}") + return ((sequence_length + alignment - 1) // alignment) * alignment + def resolve_canvas_size(aspect_width: float, aspect_height: float) -> tuple[int, int]: """Resolve an aspect ratio to the released 768p/area-capped H3 canvas.""" @@ -197,8 +212,9 @@ def build_t2av_packed_sequence( latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int] = (1, 2, 2), + sequence_alignment: int = 1, ) -> MiniMaxH3PackedSequence: - """Build the exact padless ``[text | audio | video]`` T2AV layout.""" + """Build the exact ``[text | audio | video | optional padding]`` T2AV layout.""" return build_packed_sequence( torch.full((num_text_tokens,), TEXT_TAG, dtype=torch.long), num_latent_frames, @@ -206,6 +222,7 @@ def build_t2av_packed_sequence( latent_width, num_audio_latents, patch_size, + sequence_alignment=sequence_alignment, ) @@ -217,8 +234,9 @@ def build_packed_sequence( num_audio_latents: int, patch_size: tuple[int, int, int] = (1, 2, 2), keyframe_anchors: tuple[str, ...] = (), + sequence_alignment: int = 1, ) -> MiniMaxH3PackedSequence: - """Build ``[Qwen rows | keyframe rows | target audio | target video]``.""" + """Build ``[Qwen rows | keyframes | target audio | target video | padding]``.""" _, patch_h, patch_w = patch_size rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) num_text_tokens = int(text_token_tags.shape[0]) @@ -228,7 +246,8 @@ def build_packed_sequence( condition_start = num_text_tokens audio_start = condition_start + num_condition_rows video_start = audio_start + num_audio_rows - sequence_length = video_start + num_video_rows + used_sequence_length = video_start + num_video_rows + sequence_length = _align_sequence_length(used_sequence_length, sequence_alignment) position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) @@ -261,12 +280,17 @@ def build_packed_sequence( video_positions = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) video_positions[:, :, 0] = _temporal_position_grid(num_latent_frames, float(num_text_tokens))[:, None] video_positions[:, :, 1:] = frame_grid[None] - position_ids[video_start:] = video_positions.reshape(-1, 3) + position_ids[video_start:used_sequence_length] = video_positions.reshape(-1, 3) text_indices = torch.arange(num_text_tokens, dtype=torch.long) audio_indices = torch.arange(audio_start, video_start, dtype=torch.long) - video_indices = torch.cat((torch.arange(condition_start, audio_start, dtype=torch.long), torch.arange(video_start, sequence_length, dtype=torch.long))) - token_tags = torch.empty(sequence_length, dtype=torch.long) + video_indices = torch.cat( + ( + torch.arange(condition_start, audio_start, dtype=torch.long), + torch.arange(video_start, used_sequence_length, dtype=torch.long), + ) + ) + token_tags = torch.full((sequence_length,), PADDING_TAG, dtype=torch.long) token_tags[text_indices] = text_token_tags.to(torch.long) token_tags[audio_indices] = AUDIO_TAG token_tags[video_indices] = VIDEO_TAG @@ -280,6 +304,7 @@ def build_packed_sequence( text_indices=text_indices, num_condition_video_rows=num_condition_rows, num_condition_audio_rows=0, + used_sequence_length=used_sequence_length, ) diff --git a/lightx2v/models/networks/minimax_h3/packing_ref2av.py b/lightx2v/models/networks/minimax_h3/packing_ref2av.py index 26b762d91..a1a3da308 100644 --- a/lightx2v/models/networks/minimax_h3/packing_ref2av.py +++ b/lightx2v/models/networks/minimax_h3/packing_ref2av.py @@ -16,9 +16,11 @@ FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK, + PADDING_TAG, TEXT_TAG, VIDEO_TAG, MiniMaxH3PackedSequence, + _align_sequence_length, _spatial_position_grid, _temporal_position_grid, resolve_canvas_size, @@ -161,15 +163,17 @@ def build_ref2av_packed_sequence( latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int] = (1, 2, 2), + sequence_alignment: int = 1, ) -> MiniMaxH3PackedSequence: - """Build ``[presentation | ordered references | target audio | target video]``.""" + """Build ``[presentation | references | target audio | target video | padding]``.""" _, patch_h, patch_w = patch_size num_text_tokens = int(text_token_tags.shape[0]) num_target_video_rows = num_latent_frames * (latent_height // patch_h) * (latent_width // patch_w) num_target_audio_rows = num_audio_latents * AUDIO_CHANNELS num_reference_video_rows = sum(ref.num_video_rows for ref in references if ref.kind != "audio") num_reference_audio_rows = sum(ref.num_audio_rows for ref in references) - sequence_length = num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows + num_target_video_rows + used_sequence_length = num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows + num_target_video_rows + sequence_length = _align_sequence_length(used_sequence_length, sequence_alignment) position_ids = torch.zeros(sequence_length, 3, dtype=torch.float64) position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) @@ -211,24 +215,25 @@ def build_ref2av_packed_sequence( video_start = audio_start + num_target_audio_rows _fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time, target_width_grid) frame_time = _temporal_position_grid(num_latent_frames, rotary_time) - position_ids[video_start:, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0]) - position_ids[video_start:, 1:] = target_frame_grid.repeat(num_latent_frames, 1) - video_indices = torch.cat(video_indices + [torch.arange(video_start, sequence_length)]) + position_ids[video_start:used_sequence_length, 0] = frame_time.repeat_interleave(target_frame_grid.shape[0]) + position_ids[video_start:used_sequence_length, 1:] = target_frame_grid.repeat(num_latent_frames, 1) + video_indices = torch.cat(video_indices + [torch.arange(video_start, used_sequence_length)]) audio_indices = torch.cat(audio_indices + [torch.arange(audio_start, video_start)]) text_indices = torch.arange(num_text_tokens) - token_tags = torch.empty(sequence_length, dtype=torch.long) + token_tags = torch.full((sequence_length,), PADDING_TAG, dtype=torch.long) token_tags[text_indices] = text_token_tags.long() token_tags[audio_indices] = AUDIO_TAG token_tags[video_indices] = VIDEO_TAG return MiniMaxH3PackedSequence( - sequence_length, - position_ids, - token_tags, - video_indices, - audio_indices, - text_indices, - num_reference_video_rows, - num_reference_audio_rows, + sequence_length=sequence_length, + position_ids=position_ids, + token_tags=token_tags, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices, + num_condition_video_rows=num_reference_video_rows, + num_condition_audio_rows=num_reference_audio_rows, + used_sequence_length=used_sequence_length, ) diff --git a/lightx2v/models/networks/minimax_h3/weights/merged_qkv.py b/lightx2v/models/networks/minimax_h3/weights/merged_qkv.py new file mode 100644 index 000000000..afcf37c87 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/weights/merged_qkv.py @@ -0,0 +1,349 @@ +import torch + +from lightx2v.common.ops.mm.mm_weight import MMWeight, MMWeightTP +from lightx2v.common.ops.utils import create_pin_tensor, resolve_block_name +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER + + +class _LoRAOnlyMMWeight(MMWeight): + """Use the common LoRA implementation without owning a base matrix.""" + + def __init__(self, weight_name, lora_prefix): + super().__init__(weight_name=weight_name, bias_name=None, lora_prefix=lora_prefix) + self.base_attrs = [] + self._lora_target_device = torch.device("cpu") + + def set_lora_target_device(self, device): + self._lora_target_device = torch.device(device) + + def _get_lora_target_device(self): + return self._lora_target_device + + +class _SourceLoRAWeight(MMWeightTP): + """A source-named LoRA branch with the common TP sharding rules.""" + + def __init__( + self, + weight_name, + lora_prefix, + tp_group=None, + tp_rank=0, + tp_size=1, + lora_column_chunks=1, + ): + super().__init__( + weight_name=weight_name, + bias_name=None, + mm_type="Default", + tp_group=tp_group, + tp_rank=tp_rank, + tp_size=tp_size, + split_dim="col", + reduce_output=False, + lora_column_chunks=lora_column_chunks, + lora_prefix=lora_prefix, + ) + self._mm = _LoRAOnlyMMWeight(weight_name, lora_prefix) + self._pending_lora_strength = 1.0 + + @property + def active(self): + return self._mm.has_lora_branch + + def set_target_device(self, device): + self._mm.set_lora_target_device(device) + + def register_lora(self, weight_dict, lora_strength=1): + self._pending_lora_strength = float(lora_strength) + super().register_lora(weight_dict, lora_strength) + + def update_lora(self, weight_dict, lora_strength=1): + self._pending_lora_strength = float(lora_strength) + super().update_lora(weight_dict, lora_strength) + + def pin_lora(self): + for name in ("lora_down", "lora_up", "lora_alpha", "lora_scale"): + tensor = getattr(self._mm, name, None) + if isinstance(tensor, torch.Tensor): + setattr(self._mm, name, create_pin_tensor(tensor.to("cpu"))) + + def apply_lora(self, input_tensor): + return self._mm.apply_lora(input_tensor) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + down_name = resolve_block_name(self._mm.lora_down_name, block_index) + if down_name not in destination: + self._mm.has_lora_branch = False + return destination + + if not hasattr(self._mm, "lora_down"): + local_weights = {} + for attr, name_attr in self._mm.lora_attrs.items(): + if attr in ("weight_diff", "bias_diff"): + continue + source_name = resolve_block_name(getattr(self._mm, name_attr), block_index) + if source_name in destination: + local_weights[getattr(self._mm, name_attr)] = destination[source_name] + self._mm.register_lora(local_weights, self._pending_lora_strength) + else: + self._mm.has_lora_branch = True + self._mm.load_state_dict(destination, block_index, adapter_block_index) + + if hasattr(self._mm, "lora_alpha"): + self._mm.lora_scale = self._mm.lora_alpha / self._mm.lora_down.shape[0] + else: + self._mm.lora_scale = torch.tensor(1.0, device=self._mm.lora_down.device) + return destination + + +def _operator_device(operator): + for name in ("weight", "pin_weight", "weight_cuda_buffer"): + tensor = getattr(operator, name, None) + if isinstance(tensor, torch.Tensor): + return tensor.device + return torch.device("cpu") + + +def _source_attr_name(weight_name, attr_name, bias_name=None): + if attr_name == "weight": + return weight_name + if attr_name == "bias": + if bias_name is None: + raise ValueError(f"No source bias was provided for {weight_name}") + return bias_name + return weight_name.removesuffix(".weight") + f".{attr_name}" + + +def _validate_tensors(tensors, names, *, output_rows=None): + if len({tensor.dtype for tensor in tensors}) != 1: + raise TypeError(f"Packed tensors must use one dtype: {dict(zip(names, (tensor.dtype for tensor in tensors)))}") + if len({tensor.device for tensor in tensors}) != 1: + raise ValueError(f"Packed tensors must be on one device: {dict(zip(names, (tensor.device for tensor in tensors)))}") + if len({tuple(tensor.shape[1:]) for tensor in tensors}) != 1: + raise ValueError(f"Packed tensors must have matching trailing shapes: {dict(zip(names, (tuple(tensor.shape) for tensor in tensors)))}") + if output_rows is not None and any(tensor.ndim == 0 or tensor.shape[0] != rows for tensor, rows in zip(tensors, output_rows)): + raise ValueError(f"Packed metadata must follow the weight output rows: {dict(zip(names, (tuple(tensor.shape) for tensor in tensors)))}") + + +class MiniMaxH3MergedQKVWeight: + """Pack Q/K/V once, then use the selected common matrix-multiply operator.""" + + supports_block_offload = True + + def __init__( + self, + weight_names, + bias_name=None, + bias_names=None, + mm_type="Default", + tp_group=None, + tp_rank=0, + tp_size=1, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + is_post_adapter=False, + lora_prefix="transformer_blocks", + lora_path="", + config=None, + ): + self.source_weight_names = tuple(weight_names) + if len(self.source_weight_names) != 3: + raise ValueError(f"MiniMax-H3 merged QKV expects three source weights, got {self.source_weight_names}") + if bias_name is not None: + raise ValueError("Pass Q/K/V source biases through bias_names") + self.source_bias_names = None if bias_names is None else tuple(bias_names) + if self.source_bias_names is not None and len(self.source_bias_names) != 3: + raise ValueError(f"MiniMax-H3 merged QKV expects three source biases, got {self.source_bias_names}") + if create_cpu_buffer or lazy_load: + raise NotImplementedError("Packed MiniMax-H3 weights do not support CPU buffers or disk lazy loading") + + q_weight_name = self.source_weight_names[0] + if not q_weight_name.endswith(".to_q.weight"): + raise ValueError(f"Unexpected MiniMax-H3 Q weight name {q_weight_name!r}") + self.weight_name = q_weight_name.removesuffix(".to_q.weight") + ".qkv_packed.weight" + self.bias_name = self.weight_name.removesuffix(".weight") + ".bias" if self.source_bias_names is not None else None + self.mm_type = mm_type + self.create_cuda_buffer = create_cuda_buffer + self._mm = MM_WEIGHT_REGISTER[mm_type]( + weight_name=self.weight_name, + bias_name=self.bias_name, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=lazy_load_file, + is_post_adapter=is_post_adapter, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + if config is not None: + self._mm.set_config(config) + + self._source_loras = tuple( + _SourceLoRAWeight( + source_name, + lora_prefix, + tp_group=tp_group, + tp_rank=tp_rank, + tp_size=tp_size, + ) + for source_name in self.source_weight_names + ) + self.local_qkv_dim = None + self._source_output_rows = None + + def __getattr__(self, name): + operator = self.__dict__.get("_mm") + if operator is None: + raise AttributeError(name) + return getattr(operator, name) + + @property + def has_lora_branch(self): + return any(branch.active for branch in self._source_loras) + + @property + def has_diff(self): + return self._mm.has_diff + + @property + def registered_source_weight_names(self): + return tuple(branch.weight_name for branch in self._source_loras if branch.active) + + def set_config(self, config=None): + self._mm.set_config({} if config is None else config) + + def _pack_source_weights(self, weight_dict): + missing = [name for name in self.source_weight_names if name not in weight_dict] + if missing: + raise KeyError(f"MiniMax-H3 merged QKV is missing source weights: {missing}") + weights = tuple(weight_dict[name] for name in self.source_weight_names) + shapes = [tuple(weight.shape) for weight in weights] + if any(weight.ndim != 2 for weight in weights): + raise ValueError(f"MiniMax-H3 Q/K/V weights must be two-dimensional, got {shapes}") + if len({weight.shape for weight in weights}) != 1: + raise ValueError(f"MiniMax-H3 Q/K/V weights must have identical shapes, got {shapes}") + _validate_tensors(weights, self.source_weight_names) + self.local_qkv_dim = weights[0].shape[0] + self._source_output_rows = tuple(weight.shape[0] for weight in weights) + return torch.cat(weights, dim=0).contiguous() + + def _pack_attr(self, weight_dict, attr_name): + bias_names = self.source_bias_names or (None,) * 3 + names = tuple( + _source_attr_name(weight_name, attr_name, source_bias_name) + for weight_name, source_bias_name in zip(self.source_weight_names, bias_names) + ) + missing = [name for name in names if name not in weight_dict] + if missing: + raise KeyError(f"MiniMax-H3 merged QKV is missing {attr_name}: {missing}") + tensors = tuple(weight_dict[name] for name in names) + if all(tensor.ndim == 0 for tensor in tensors): + if any(tensor.item() != tensors[0].item() for tensor in tensors[1:]): + raise ValueError(f"MiniMax-H3 Q/K/V {attr_name} values must match") + return tensors[0].clone(), names + _validate_tensors(tensors, names, output_rows=self._source_output_rows) + return torch.cat(tensors, dim=0).contiguous(), names + + def load(self, weight_dict): + packed = {self.weight_name: self._pack_source_weights(weight_dict)} + consumed = set(self.source_weight_names) + for packed_name, attr_name, _ in self._mm.base_attrs: + if attr_name == "weight": + continue + value, source_names = self._pack_attr(weight_dict, attr_name) + packed[packed_name] = value + consumed.update(source_names) + self._mm.load(packed) + if not self.create_cuda_buffer: + for name in consumed: + weight_dict.pop(name, None) + + def apply(self, input_tensor): + projected = self._mm.apply(input_tensor) + if not self.has_lora_branch: + return projected + for index, branch in enumerate(self._source_loras): + if branch.active: + start = index * self.local_qkv_dim + projected[..., start : start + self.local_qkv_dim].add_(branch.apply_lora(input_tensor)) + return projected + + def split_qkv(self, projected): + if self.local_qkv_dim is None: + raise RuntimeError("MiniMax-H3 merged QKV has not been loaded") + expected_width = 3 * self.local_qkv_dim + if projected.shape[-1] != expected_width: + raise ValueError(f"MiniMax-H3 merged QKV output width must be {expected_width}, got {projected.shape[-1]}") + return projected.split(self.local_qkv_dim, dim=-1) + + def _adapter_target_device(self): + return _operator_device(self._mm) + + def _pin_source_loras_if_needed(self): + if getattr(self._mm, "pin_weight", None) is not None and getattr(self._mm, "weight", None) is None: + for branch in self._source_loras: + if branch.active: + branch.pin_lora() + + def _set_source_lora_devices(self): + target = self._adapter_target_device() + for branch in self._source_loras: + branch.set_target_device(target) + + def register_lora(self, weight_dict, strength): + self._set_source_lora_devices() + for branch in self._source_loras: + branch.register_lora(weight_dict, strength) + self._pin_source_loras_if_needed() + + def update_lora(self, weight_dict, strength): + self._set_source_lora_devices() + for branch in self._source_loras: + branch.update_lora(weight_dict, strength) + self._pin_source_loras_if_needed() + + def remove_lora(self): + for branch in self._source_loras: + branch.remove_lora() + + def register_diff(self, weight_dict): + diff_names = { + name + for branch in self._source_loras + for name in (branch._mm.weight_diff_name, branch._mm.bias_diff_name) + } + present = sorted(diff_names.intersection(weight_dict)) + if present: + raise NotImplementedError(f"Packed MiniMax-H3 QKV does not support diff adapters: {present[:3]}") + + def state_dict(self, destination=None): + destination = self._mm.state_dict(destination) + for branch in self._source_loras: + branch.state_dict(destination) + return destination + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + self._mm.load_state_dict(destination, block_index, adapter_block_index) + self._set_source_lora_devices() + for branch in self._source_loras: + branch.load_state_dict(destination, block_index, adapter_block_index) + return destination + + def to_cuda(self, non_blocking=False): + self._mm.to_cuda(non_blocking) + for branch in self._source_loras: + branch.to_cuda(non_blocking) + + def to_cpu(self, non_blocking=False): + self._mm.to_cpu(non_blocking) + for branch in self._source_loras: + branch.to_cpu(non_blocking) + + def load_state_dict_from_disk(self, block_index, adapter_block_index=None): + raise NotImplementedError("Packed MiniMax-H3 QKV does not support disk lazy loading") + + +__all__ = ["MiniMaxH3MergedQKVWeight"] diff --git a/lightx2v/models/networks/minimax_h3/weights/post_weights.py b/lightx2v/models/networks/minimax_h3/weights/post_weights.py index 322e7b8a9..08a9340a6 100644 --- a/lightx2v/models/networks/minimax_h3/weights/post_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/post_weights.py @@ -1,7 +1,26 @@ +import torch.distributed as dist + from lightx2v.common.modules.weight_module import WeightModule from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +def _linear(name, bias, force_fp32, config, tp_split=None): + kind = "Default-ForceFp32" if force_fp32 else "Default" + if config.get("tensor_parallel", False) and tp_split is not None: + tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + tp_mm_type = config.get("tp_mm_type", "TensorParallel") + return MM_WEIGHT_REGISTER[tp_mm_type]( + weight_name=f"{name}.weight", + bias_name=f"{name}.bias" if bias else None, + mm_type=kind, + tp_group=tp_group, + tp_rank=dist.get_rank(tp_group), + tp_size=dist.get_world_size(tp_group), + split_dim=tp_split, + ) + return MM_WEIGHT_REGISTER[kind](f"{name}.weight", f"{name}.bias" if bias else None) + + def _rms(config, name, eps): return RMS_WEIGHT_REGISTER[config.get("rms_type", "torch_native")](name, eps=eps) @@ -9,6 +28,7 @@ def _rms(config, name, eps): class MiniMaxH3PostWeights(WeightModule): def __init__(self, config): super().__init__() + col = "col" self.add_module( "norm_out", _rms(config, "norm_out.norm.weight", eps=float(config.get("final_norm_eps", 1e-5))), @@ -18,13 +38,13 @@ def __init__(self, config): # its output; update the offline builder if its definition changes. self.add_module( "norm_out_linear", - MM_WEIGHT_REGISTER["Default"]("norm_out.linear.weight", "norm_out.linear.bias"), + _linear("norm_out.linear", bias=True, force_fp32=False, config=config, tp_split=col), ) self.add_module( "proj_out", - MM_WEIGHT_REGISTER["Default-ForceFp32"]("proj_out.weight", "proj_out.bias"), + _linear("proj_out", bias=True, force_fp32=True, config=config, tp_split=col), ) self.add_module( "audio_proj_out", - MM_WEIGHT_REGISTER["Default-ForceFp32"]("audio_proj_out.weight", "audio_proj_out.bias"), + _linear("audio_proj_out", bias=True, force_fp32=True, config=config, tp_split=col), ) diff --git a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py b/lightx2v/models/networks/minimax_h3/weights/pre_weights.py index 76b46f91c..16294e116 100644 --- a/lightx2v/models/networks/minimax_h3/weights/pre_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/pre_weights.py @@ -1,6 +1,9 @@ import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.common.ops.norm import SGLExactRMSNorm128 # noqa: F401 +from lightx2v.models.networks.minimax_h3.weights.merged_qkv import MiniMaxH3MergedQKVWeight +from lightx2v.models.networks.minimax_h3.weights.reordered_mlp import MiniMaxH3ReorderedMLPWeight from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER @@ -24,28 +27,58 @@ def _linear(name, bias=False, force_fp32=False, config=None, tp_split=None): return MM_WEIGHT_REGISTER[kind](f"{name}.weight", f"{name}.bias" if bias else None, **lora_kwargs) -def _rms(config, name, eps): - return RMS_WEIGHT_REGISTER[config.get("rms_type", "torch_native")](name, eps=eps) +def _packed_linear_kwargs(config): + kwargs = { + "mm_type": "Default", + "tp_group": None, + "tp_rank": 0, + "tp_size": 1, + "config": config, + } + if config.get("tensor_parallel", False): + group = config["device_mesh"].get_group(mesh_dim="tensor_p") + kwargs.update( + tp_group=group, + tp_rank=dist.get_rank(group), + tp_size=dist.get_world_size(group), + ) + return kwargs + + +def _rms(config, name, eps, kind=None): + return RMS_WEIGHT_REGISTER[kind or config.get("rms_type", "torch_native")](name, eps=eps) class MiniMaxH3RefinerAttentionWeights(WeightModule): def __init__(self, prefix, config): super().__init__() - self.add_module("to_q", _linear(f"{prefix}.to_q", config=config, tp_split="col")) - self.add_module("to_k", _linear(f"{prefix}.to_k", config=config, tp_split="col")) - self.add_module("to_v", _linear(f"{prefix}.to_v", config=config, tp_split="col")) + self.add_module( + "qkv", + MiniMaxH3MergedQKVWeight( + weight_names=tuple(f"{prefix}.to_{name}.weight" for name in ("q", "k", "v")), + lora_prefix="token_refiner", + **_packed_linear_kwargs(config), + ), + ) + qk_norm_kind = "sgl_exact_rms_norm_128" self.add_module( "norm_q", - _rms(config, f"{prefix}.norm_q.weight", eps=float(config.get("qk_norm_eps", 1e-5))), + _rms( + config, + f"{prefix}.norm_q.weight", + eps=float(config.get("qk_norm_eps", 1e-5)), + kind=qk_norm_kind, + ), ) self.add_module( "norm_k", - _rms(config, f"{prefix}.norm_k.weight", eps=float(config.get("qk_norm_eps", 1e-5))), + _rms( + config, + f"{prefix}.norm_k.weight", + eps=float(config.get("qk_norm_eps", 1e-5)), + kind=qk_norm_kind, + ), ) - # H3's text refiner attends over a short text-only sequence, while the - # main transformer attends over the much longer packed AV sequence. - # Allow sparse main attention without paying its setup/quality cost in - # the refiner. Existing configs retain their previous shared backend. attn_type = config.get("refiner_attn_type", config.get("attn_type", "flash_attn3")) attention_cls = ATTN_WEIGHT_REGISTER[attn_type] if attn_type == "dynamic_sparse_attn": @@ -61,7 +94,12 @@ def __init__(self, prefix, config): class MiniMaxH3FeedForwardWeights(WeightModule): def __init__(self, prefix, config): super().__init__() - self.add_module("in_proj", _linear(f"{prefix}.net.0.proj", config=config, tp_split="col")) + in_proj = MiniMaxH3ReorderedMLPWeight( + weight_name=f"{prefix}.net.0.proj.weight", + lora_prefix="token_refiner", + **_packed_linear_kwargs(config), + ) + self.add_module("in_proj", in_proj) self.add_module("out_proj", _linear(f"{prefix}.net.2", config=config, tp_split="row")) @@ -79,16 +117,27 @@ def __init__(self, index, config): class MiniMaxH3PreWeights(WeightModule): def __init__(self, config): super().__init__() - # The released checkpoint deliberately keeps the two media projections - # and timestep MLP in fp32. The text projection/refiner stay bf16. - self.add_module("proj_in", _linear("proj_in", bias=True, force_fp32=True)) - self.add_module("audio_proj_in", _linear("audio_proj_in", bias=True, force_fp32=True)) - self.add_module("context_embedder", _linear("context_embedder", bias=True)) + col, row = "col", "row" + self.add_module("proj_in", _linear("proj_in", bias=True, force_fp32=True, config=config, tp_split=col)) + self.add_module( + "audio_proj_in", + _linear("audio_proj_in", bias=True, force_fp32=True, config=config, tp_split=col), + ) + self.add_module( + "context_embedder", + _linear("context_embedder", bias=True, config=config, tp_split=col), + ) if not config.get("use_adaln_cache", False): # ADALN CACHE SYNC: The offline builder reads these keys and mirrors # their FP32 semantics; update the offline builder if either changes. - self.add_module("time_linear_1", _linear("time_embedder.linear_1", bias=True, force_fp32=True)) - self.add_module("time_linear_2", _linear("time_embedder.linear_2", bias=True, force_fp32=True)) + self.add_module( + "time_linear_1", + _linear("time_embedder.linear_1", bias=True, force_fp32=True, config=config, tp_split=col), + ) + self.add_module( + "time_linear_2", + _linear("time_embedder.linear_2", bias=True, force_fp32=True, config=config, tp_split=row), + ) self.add_module( "refiner_blocks", WeightModuleList([MiniMaxH3TokenRefinerBlockWeights(i, config) for i in range(int(config.get("num_refiner_layers", 2)))]), diff --git a/lightx2v/models/networks/minimax_h3/weights/reordered_mlp.py b/lightx2v/models/networks/minimax_h3/weights/reordered_mlp.py new file mode 100644 index 000000000..28876e1e0 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/weights/reordered_mlp.py @@ -0,0 +1,189 @@ +import torch + +from lightx2v.models.networks.minimax_h3.weights.merged_qkv import ( + _SourceLoRAWeight, + _operator_device, + _source_attr_name, +) +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER + + +class MiniMaxH3ReorderedMLPWeight: + """Store the H3 SwiGLU projection in the runtime's gate/value row order.""" + + supports_block_offload = True + + def __init__( + self, + weight_name, + bias_name=None, + mm_type="Default", + tp_group=None, + tp_rank=0, + tp_size=1, + create_cuda_buffer=False, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=None, + is_post_adapter=False, + lora_prefix="transformer_blocks", + lora_path="", + config=None, + ): + if create_cpu_buffer or lazy_load: + raise NotImplementedError("Packed MiniMax-H3 weights do not support CPU buffers or disk lazy loading") + if not weight_name.endswith(".weight"): + raise ValueError(f"Unexpected MiniMax-H3 SwiGLU weight name {weight_name!r}") + + self.source_weight_name = weight_name + self.source_bias_name = bias_name + self.weight_name = weight_name.removesuffix(".weight") + ".reordered.weight" + self.bias_name = self.weight_name.removesuffix(".weight") + ".bias" if bias_name is not None else None + self.mm_type = mm_type + self.create_cuda_buffer = create_cuda_buffer + self._mm = MM_WEIGHT_REGISTER[mm_type]( + weight_name=self.weight_name, + bias_name=self.bias_name, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=False, + lazy_load=False, + lazy_load_file=lazy_load_file, + is_post_adapter=is_post_adapter, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + if config is not None: + self._mm.set_config(config) + + self._source_lora = _SourceLoRAWeight( + weight_name, + lora_prefix, + tp_group=tp_group, + tp_rank=tp_rank, + tp_size=tp_size, + lora_column_chunks=2, + ) + self.local_inner_dim = None + self._source_output_rows = None + + def __getattr__(self, name): + operator = self.__dict__.get("_mm") + if operator is None: + raise AttributeError(name) + return getattr(operator, name) + + @property + def has_lora_branch(self): + return self._source_lora.active + + @property + def has_diff(self): + return self._mm.has_diff + + @property + def registered_source_weight_names(self): + return (self.source_weight_name,) if self._source_lora.active else () + + def set_config(self, config=None): + self._mm.set_config({} if config is None else config) + + def _reorder_output_rows(self, tensor, name): + if tensor.ndim == 0: + return tensor.clone() + if self._source_output_rows is not None and tensor.shape[0] != self._source_output_rows: + raise ValueError( + f"MiniMax-H3 SwiGLU {name} must follow the weight output rows " + f"({self._source_output_rows}), got {tuple(tensor.shape)}" + ) + if tensor.shape[0] % 2: + raise ValueError(f"Invalid MiniMax-H3 fused SwiGLU {name} shape {tuple(tensor.shape)}") + value, gate = tensor.chunk(2, dim=0) + return torch.cat((gate, value), dim=0).contiguous() + + def _reorder_source_weight(self, weight): + if weight.ndim != 2 or weight.shape[0] % 2: + raise ValueError(f"Invalid MiniMax-H3 fused SwiGLU weight shape {tuple(weight.shape)}") + self._source_output_rows = weight.shape[0] + self.local_inner_dim = weight.shape[0] // 2 + return self._reorder_output_rows(weight, "weight") + + def load(self, weight_dict): + if self.source_weight_name not in weight_dict: + raise KeyError(f"MiniMax-H3 reordered SwiGLU is missing {self.source_weight_name}") + packed = {self.weight_name: self._reorder_source_weight(weight_dict[self.source_weight_name])} + consumed = {self.source_weight_name} + for packed_name, attr_name, _ in self._mm.base_attrs: + if attr_name == "weight": + continue + source_name = _source_attr_name(self.source_weight_name, attr_name, self.source_bias_name) + if source_name not in weight_dict: + raise KeyError(f"MiniMax-H3 reordered SwiGLU is missing {source_name}") + packed[packed_name] = self._reorder_output_rows(weight_dict[source_name], attr_name) + consumed.add(source_name) + self._mm.load(packed) + if not self.create_cuda_buffer: + for name in consumed: + weight_dict.pop(name, None) + + def apply(self, input_tensor): + projected = self._mm.apply(input_tensor) + if not self._source_lora.active: + return projected + value, gate = self._source_lora.apply_lora(input_tensor).chunk(2, dim=-1) + projected.add_(torch.cat((gate, value), dim=-1)) + return projected + + def _set_source_lora_device(self): + self._source_lora.set_target_device(_operator_device(self._mm)) + + def _pin_source_lora_if_needed(self): + if ( + self._source_lora.active + and getattr(self._mm, "pin_weight", None) is not None + and getattr(self._mm, "weight", None) is None + ): + self._source_lora.pin_lora() + + def register_lora(self, weight_dict, strength): + self._set_source_lora_device() + self._source_lora.register_lora(weight_dict, strength) + self._pin_source_lora_if_needed() + + def update_lora(self, weight_dict, strength): + self._set_source_lora_device() + self._source_lora.update_lora(weight_dict, strength) + self._pin_source_lora_if_needed() + + def remove_lora(self): + self._source_lora.remove_lora() + + def register_diff(self, weight_dict): + diff_names = (self._source_lora._mm.weight_diff_name, self._source_lora._mm.bias_diff_name) + present = sorted(set(diff_names).intersection(weight_dict)) + if present: + raise NotImplementedError(f"Packed MiniMax-H3 SwiGLU does not support diff adapters: {present[:2]}") + + def state_dict(self, destination=None): + destination = self._mm.state_dict(destination) + self._source_lora.state_dict(destination) + return destination + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + self._mm.load_state_dict(destination, block_index, adapter_block_index) + self._set_source_lora_device() + self._source_lora.load_state_dict(destination, block_index, adapter_block_index) + return destination + + def to_cuda(self, non_blocking=False): + self._mm.to_cuda(non_blocking) + self._source_lora.to_cuda(non_blocking) + + def to_cpu(self, non_blocking=False): + self._mm.to_cpu(non_blocking) + self._source_lora.to_cpu(non_blocking) + + def load_state_dict_from_disk(self, block_index, adapter_block_index=None): + raise NotImplementedError("Packed MiniMax-H3 SwiGLU does not support disk lazy loading") + + +__all__ = ["MiniMaxH3ReorderedMLPWeight"] diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index d7fa24758..3375e9a98 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -2,7 +2,10 @@ import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.models.networks.minimax_h3.infer.triton_ops import MiniMaxH3TritonRope # noqa: F401 +from lightx2v.common.ops.norm import SGLExactRMSNorm128 # noqa: F401 +from lightx2v.common.ops.rope import SGLExactNeoXRope # noqa: F401 +from lightx2v.models.networks.minimax_h3.weights.merged_qkv import MiniMaxH3MergedQKVWeight +from lightx2v.models.networks.minimax_h3.weights.reordered_mlp import MiniMaxH3ReorderedMLPWeight from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER @@ -31,8 +34,26 @@ def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): ) -def _rms(config, name, eps, create_cuda_buffer=False): - return RMS_WEIGHT_REGISTER[config.get("rms_type", "torch_native")]( +def _packed_linear_kwargs(config): + kwargs = { + "mm_type": config.get("dit_quant_scheme", "Default"), + "tp_group": None, + "tp_rank": 0, + "tp_size": 1, + "config": config, + } + if config.get("tensor_parallel", False): + group = config["device_mesh"].get_group(mesh_dim="tensor_p") + kwargs.update( + tp_group=group, + tp_rank=dist.get_rank(group), + tp_size=dist.get_world_size(group), + ) + return kwargs + + +def _rms(config, name, eps, create_cuda_buffer=False, kind=None): + return RMS_WEIGHT_REGISTER[kind or config.get("rms_type", "torch_native")]( name, create_cuda_buffer=create_cuda_buffer, eps=eps, @@ -42,10 +63,17 @@ def _rms(config, name, eps, create_cuda_buffer=False): class MiniMaxH3AttentionWeights(WeightModule): def __init__(self, prefix, config, create_cuda_buffer=False): super().__init__() - self.add_module("to_q", _linear(config, f"{prefix}.to_q", create_cuda_buffer=create_cuda_buffer, tp_split="col")) - self.add_module("to_k", _linear(config, f"{prefix}.to_k", create_cuda_buffer=create_cuda_buffer, tp_split="col")) - self.add_module("to_v", _linear(config, f"{prefix}.to_v", create_cuda_buffer=create_cuda_buffer, tp_split="col")) + self.add_module( + "qkv", + MiniMaxH3MergedQKVWeight( + weight_names=tuple(f"{prefix}.to_{name}.weight" for name in ("q", "k", "v")), + create_cuda_buffer=create_cuda_buffer, + **_packed_linear_kwargs(config), + ), + ) + qk_eps = float(config.get("qk_norm_eps", 1e-5)) + qk_norm_kind = "sgl_exact_rms_norm_128" self.add_module( "norm_q", _rms( @@ -53,6 +81,7 @@ def __init__(self, prefix, config, create_cuda_buffer=False): f"{prefix}.norm_q.weight", create_cuda_buffer=create_cuda_buffer, eps=qk_eps, + kind=qk_norm_kind, ), ) self.add_module( @@ -62,13 +91,15 @@ def __init__(self, prefix, config, create_cuda_buffer=False): f"{prefix}.norm_k.weight", create_cuda_buffer=create_cuda_buffer, eps=qk_eps, + kind=qk_norm_kind, ), ) + rope_kind = config.get("rope_type", "sgl_exact_neox_rope") self.add_module( "rope", - ROPE_REGISTER[config.get("rope_type", "torch_real_rope")]( + ROPE_REGISTER[rope_kind]( layout="split_half", - compute_dtype=torch.float32, + compute_dtype=torch.bfloat16, ), ) attn_type = config.get("attn_type", "flash_attn3") @@ -92,7 +123,13 @@ def __init__(self, prefix, config, create_cuda_buffer=False): class MiniMaxH3FeedForwardWeights(WeightModule): def __init__(self, prefix, config, create_cuda_buffer=False): super().__init__() - self.add_module("in_proj", _linear(config, f"{prefix}.net.0.proj", create_cuda_buffer=create_cuda_buffer, tp_split="col")) + in_proj = MiniMaxH3ReorderedMLPWeight( + weight_name=f"{prefix}.net.0.proj.weight", + create_cuda_buffer=create_cuda_buffer, + lora_prefix="transformer_blocks", + **_packed_linear_kwargs(config), + ) + self.add_module("in_proj", in_proj) self.add_module("out_proj", _linear(config, f"{prefix}.net.2", create_cuda_buffer=create_cuda_buffer, tp_split="row")) @@ -139,8 +176,6 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.blocks = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config) for i in range(int(config.get("num_layers", 50)))]) if config.get("cpu_offload", False) and config.get("offload_granularity", "model") == "block": self.offload_block_cuda_buffers = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config, create_cuda_buffer=True) for i in range(2)]) - # Register device buffers before source blocks: buffer allocation - # needs checkpoint metadata that normal CPU loading consumes. self.add_module("offload_block_cuda_buffers", self.offload_block_cuda_buffers) self.offload_phase_cuda_buffers = None self.add_module("blocks", self.blocks) diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 81a16a63b..1fa3da9ee 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -9,6 +9,7 @@ from lightx2v.models.audio_encoders.hf.minimax_h3 import MiniMaxH3AudioVAE from lightx2v.models.input_encoders.hf.minimax_h3 import MiniMaxH3Qwen3VLTextEncoder +from lightx2v.models.networks.minimax_h3.config import resolve_minimax_h3_execution_profile from lightx2v.models.networks.minimax_h3.lora import MiniMaxH3LoraAdapter from lightx2v.models.networks.minimax_h3.model import MiniMaxH3Model from lightx2v.models.networks.minimax_h3.packing import ( @@ -45,7 +46,7 @@ from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import DTYPE_MAP, GET_RECORDER_MODE from lightx2v.utils.input_info import INPUT_INFO_TYPES -from lightx2v.utils.ltx2_media_io import encode_video +from lightx2v.utils.ltx2_media_io import encode_video_sglang_compatible from lightx2v.utils.profiler import ProfilingContext4DebugL1, ProfilingContext4DebugL2 from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -81,12 +82,12 @@ def build_minimax_h3_model_with_lora(config, model_kwargs, lora_configs): @RUNNER_REGISTER("minimax_h3") class MiniMaxH3Runner(DefaultRunner): - """Native MiniMax-H3 audio-video runner. + """MiniMax-H3 audio-video runner. Transformer, text-encoder, and VAE residency are configured independently. ``cpu_offload`` controls only the transformer, while - ``text_encoder_cpu_offload`` and ``vae_cpu_offload`` control the native - Qwen3-VL conditioner and both native VAEs. This mirrors Wan's component + ``text_encoder_cpu_offload`` and ``vae_cpu_offload`` control the + Qwen3-VL conditioner and both VAEs. This mirrors Wan's component offload behavior while keeping Diffusers out of the runtime dependency graph. """ @@ -105,6 +106,8 @@ class MiniMaxH3Runner(DefaultRunner): } def __init__(self, config): + self.execution_profile = resolve_minimax_h3_execution_profile(config) + self.video_vae_class = MiniMaxH3VideoVAE if config.get("lazy_load", False) or config.get("unload_modules", False): raise NotImplementedError("MiniMax-H3 does not support lazy_load or unload_modules yet; use the released sharded checkpoint with model or block CPU offload.") super().__init__(config) @@ -267,7 +270,7 @@ def load_vae(self): video_vae_quant_scheme = self.config["video_vae_quant_scheme"] if video_vae_quantized else None video_vae_quantized_ckpt = self.config["video_vae_quantized_ckpt"] if video_vae_quantized else None vae_sensitive_layer_dtype = DTYPE_MAP[self.config.get("vae_sensitive_layer_dtype", "fp32")] - video_vae = MiniMaxH3VideoVAE.from_pretrained( + video_vae = self.video_vae_class.from_pretrained( self.config["model_path"], device=AI_DEVICE, cpu_offload=cpu_offload, @@ -276,6 +279,7 @@ def load_vae(self): sensitive_layer_dtype=vae_sensitive_layer_dtype, use_compile=self.config.get("vae_use_compile", False), attn_type=self.config.get("vae_attn_type", "torch_sdpa"), + rope_type=self.config.get("rope_type", "sgl_exact_neox_rope"), ) self._vae_decode_tile_shapes = self.config.get("vae_decode_tile_shape", {}) self._validate_vae_decode_tile_shapes(self._vae_decode_tile_shapes, video_vae) @@ -469,10 +473,18 @@ def _prepare_references(self): raise ValueError(f"MiniMax-H3 ref2av accepts at most {MAX_REFERENCE_AUDIOS} audio-bearing references") return references + def _reference_pixels(self, value, *, video: bool) -> torch.Tensor: + pixels = torch.from_numpy(np.asarray(value).copy()) + if video: + pixels = pixels.permute(3, 0, 1, 2)[None] + else: + pixels = pixels.permute(2, 0, 1)[None, :, None] + return self.video_vae.prepare_reference_pixels(pixels) + def _encode_keyframes(self, keyframes): latents = [] for image in keyframes: - pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1)[None, :, None].float().div_(255.0) + pixels = self._reference_pixels(image, video=False) latents.append(self.video_vae.encode_condition(pixels, video=False)) return latents @@ -481,11 +493,11 @@ def _encode_references(self, references): for reference in references: if reference.kind != "audio": if reference.kind == "image": - pixels = torch.from_numpy(np.asarray(reference.image).copy()).permute(2, 0, 1)[None, :, None].float().div_(255.0) + pixels = self._reference_pixels(reference.image, video=False) latent = self.video_vae.encode_condition(pixels, video=False) else: frames = reference.frames[: trim_reference_num_frames(reference.frames.shape[0])] - pixels = torch.from_numpy(frames.copy()).permute(3, 0, 1, 2)[None].float().div_(255.0) + pixels = self._reference_pixels(frames, video=True) latent = self.video_vae.encode_condition(pixels, video=True) reference.num_latent_frames = latent.shape[2] reference.latent_height, reference.latent_width = latent.shape[3:] @@ -564,7 +576,7 @@ def init_run(self): if not self.config.get("cpu_offload", False): logger.info("MiniMax-H3 transformer is resident on the accelerator") elif self.config.get("offload_granularity", "model") == "model": - logger.info("Moving the native MiniMax-H3 transformer to the accelerator") + logger.info("Moving the MiniMax-H3 transformer to the accelerator") self.model.to_cuda() else: logger.info("MiniMax-H3 block offload enabled; keeping source blocks on CPU and using two accelerator buffers") @@ -642,11 +654,8 @@ def run_vae_decoder(self, video_rows, audio_rows): audio = self.audio_vae.decode(audio_latents) return video, audio - @staticmethod - def _video_to_uint8_frames(video): - if video.ndim != 5 or video.shape[0] != 1 or video.shape[1] != 3: - raise ValueError(f"decoded H3 video must be [1,3,F,H,W], got {tuple(video.shape)}") - return (video[0].permute(1, 2, 3, 0).float() * 255.0).round().to(torch.uint8).contiguous().cpu() + def _video_to_uint8_frames(self, video): + return self.video_vae.to_uint8_frames(video) def process_images_after_vae_decoder(self): if self.video_vae.decode_parallel and dist.get_rank() != 0: @@ -674,13 +683,13 @@ def process_images_after_vae_decoder(self): ) logger.info(f"Saving MiniMax-H3 audio-video output to {output_path}") with ProfilingContext4DebugL2("Save Audio-Video Output"): - encode_video( + encode_video_sglang_compatible( video=frames, fps=int(self.config.get("fps", 24)), audio=audio, output_path=output_path, - video_chunks_number=1, - video_codec_options=self.config.get("video_codec_options"), + crf=self.config.get("sglang_export_crf", 25), + threads=self.config.get("sglang_export_threads", 24), ) logger.info(f"MiniMax-H3 output saved to {output_path}") return {"video": None, "audio": None} diff --git a/lightx2v/models/schedulers/minimax_h3/scheduler.py b/lightx2v/models/schedulers/minimax_h3/scheduler.py index 4b1a0b0df..677c7605d 100644 --- a/lightx2v/models/schedulers/minimax_h3/scheduler.py +++ b/lightx2v/models/schedulers/minimax_h3/scheduler.py @@ -2,6 +2,7 @@ import torch +from lightx2v.models.networks.minimax_h3.config import resolve_minimax_h3_execution_profile from lightx2v.models.networks.minimax_h3.packing import ( AUDIO_CHANNELS, KEYFRAME_NOISE_AUG, @@ -43,12 +44,12 @@ class MiniMaxH3Scheduler(BaseScheduler): def __init__(self, config): super().__init__(config) + execution_profile = resolve_minimax_h3_execution_profile(config) infer_steps = int(config["infer_steps"]) self.video_shift = float(config.get("video_flow_shift", 12.0)) self.audio_shift = float(config.get("audio_flow_shift", 3.0)) - self.step_update = config.get("h3_step_update", "reference_blend") - if self.step_update not in {"reference_blend", "training_euler"}: - raise ValueError(f"MiniMax-H3 h3_step_update must be 'reference_blend' or 'training_euler', got {self.step_update!r}") + self.packed_sequence_alignment = execution_profile.packed_sequence_alignment + self.sampling_profile = execution_profile.sampling_profile if self.video_shift <= 0 or self.audio_shift <= 0: raise ValueError("MiniMax-H3 flow shifts must be positive") self.video_sigmas, self.video_timesteps = _make_schedule(infer_steps, self.video_shift, AI_DEVICE) @@ -83,24 +84,27 @@ def prepare( num_audio_latents = audio_latent_num_frames(num_frames) patch_size = tuple(self.config.get("patch_size", (1, 2, 2))) - # The released pipeline uses one CPU random stream even when inference - # runs on CUDA: float32 video noise first, then channel-major audio. - self.generator = torch.Generator(device="cpu").manual_seed(int(seed)) condition_video_latents = condition_video_latents or [] condition_audio_latents = condition_audio_latents or [] + # Video and audio use separate seeded CPU FP32 streams. condition_video_rows = [] + condition_count = len(condition_video_latents) for clean in condition_video_latents: - noise = torch.randn(clean.shape, generator=self.generator, device="cpu", dtype=torch.float32) - clean_rows = patchify_video_latents(clean.float(), patch_size).to(AI_DEVICE) - noise_rows = patchify_video_latents(noise.to(AI_DEVICE), patch_size) - # Match Diffusers' ``scheduler.scale_noise`` exactly: the - # conditioning VAE rows are moved first and mixed on the execution - # device, with the scalar represented in the sample dtype. Doing - # this FP32 operation on CPU differs by an ulp on CUDA and that - # perturbation is amplified by the 48-layer denoiser. - timestep = torch.tensor(KEYFRAME_NOISE_AUG, dtype=clean_rows.dtype, device=clean_rows.device) + clean_cpu = clean.detach().to(device="cpu", dtype=torch.float32) + condition_t, condition_h, condition_w = clean_cpu.shape[-3:] + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn( + (1, int(self.config.get("in_channels", 24)), latent_frames + condition_count, condition_h, condition_w), + generator=generator, + device="cpu", + dtype=torch.float32, + )[:, :, :condition_t] + clean_rows = patchify_video_latents(clean_cpu, patch_size) + noise_rows = patchify_video_latents(noise, patch_size) + timestep = torch.tensor(KEYFRAME_NOISE_AUG, dtype=torch.float32, device="cpu") condition_video_rows.append(timestep * clean_rows + (1.0 - timestep) * noise_rows) + self.generator = torch.Generator(device="cpu").manual_seed(int(seed)) video_noise = torch.randn( (1, int(self.config.get("in_channels", 24)), latent_frames, latent_height, latent_width), generator=self.generator, @@ -108,23 +112,51 @@ def prepare( dtype=torch.float32, ) target_video_rows = patchify_video_latents(video_noise, patch_size) - self.video_latents = torch.cat(condition_video_rows + [target_video_rows.to(AI_DEVICE)]) + self.video_latents = torch.cat(condition_video_rows + [target_video_rows]).to(AI_DEVICE) + + audio_generator = torch.Generator(device="cpu").manual_seed(int(seed)) target_audio_rows = torch.randn( - ( - num_audio_latents * AUDIO_CHANNELS, - int(self.config.get("audio_in_channels", 32)), - ), - generator=self.generator, + (num_audio_latents * AUDIO_CHANNELS, int(self.config.get("audio_in_channels", 32))), + generator=audio_generator, device="cpu", dtype=torch.float32, ) - condition_audio_rows = [latent.transpose(1, 2).reshape(-1, latent.shape[1]).float() for latent in condition_audio_latents] + audio_noise_aug = float(self.config.get("audio_condition_noise_aug", 1.0)) + if not 0.0 <= audio_noise_aug <= 1.0: + raise ValueError(f"MiniMax-H3 audio_condition_noise_aug must be in [0, 1], got {audio_noise_aug}") + condition_audio_rows = [] + for latent in condition_audio_latents: + clean_rows = latent.detach().transpose(1, 2).reshape(-1, latent.shape[1]).to(device="cpu", dtype=torch.float32) + if audio_noise_aug < 1.0: + generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1) + noise_rows = torch.randn(clean_rows.shape, generator=generator, device="cpu", dtype=torch.float32) + timestep = torch.tensor(audio_noise_aug, dtype=torch.float32, device="cpu") + clean_rows = timestep * clean_rows + (1.0 - timestep) * noise_rows + condition_audio_rows.append(clean_rows) self.audio_latents = torch.cat(condition_audio_rows + [target_audio_rows]).to(AI_DEVICE) if references is None: - self.layout_cpu = build_packed_sequence(text_token_tags.cpu(), latent_frames, latent_height, latent_width, num_audio_latents, patch_size, keyframe_anchors) + self.layout_cpu = build_packed_sequence( + text_token_tags.cpu(), + latent_frames, + latent_height, + latent_width, + num_audio_latents, + patch_size, + keyframe_anchors, + sequence_alignment=self.packed_sequence_alignment, + ) else: - self.layout_cpu = build_ref2av_packed_sequence(text_token_tags.cpu(), references, latent_frames, latent_height, latent_width, num_audio_latents, patch_size) + self.layout_cpu = build_ref2av_packed_sequence( + text_token_tags.cpu(), + references, + latent_frames, + latent_height, + latent_width, + num_audio_latents, + patch_size, + sequence_alignment=self.packed_sequence_alignment, + ) self.layout = _layout_to_device(self.layout_cpu, AI_DEVICE) self.num_frames = num_frames self.height = height @@ -149,18 +181,25 @@ def step_pre(self, step_index): self.timestep_indices = inverse.to(AI_DEVICE) @staticmethod - def _step(sample, model_output, timestep, sigmas, step_index, step_update): - # H3 predicts a data-ward velocity. Keep the round trip through - # timestep separate from the stored sigma grid to match the reference. - sigma_from_timestep = 1.0 - timestep.to(device=sample.device, dtype=sample.dtype) - denoised = sample + sigma_from_timestep * model_output + def _step(sample, model_output, timestep, sigmas, step_index, sampling_profile): sigma = sigmas[step_index].to(device=sample.device, dtype=torch.float32) sigma_next = sigmas[step_index + 1].to(device=sample.device, dtype=torch.float32) - if step_update == "training_euler": - # Match MiniMaxH3T2AVDmdTrainer.run_back_simulation exactly. + if sampling_profile == "dmd": return sample.float() + (sigma - sigma_next) * model_output.float() + + # Rebuild sigma from the timestep to preserve the base model's + # reference rounding and operation order. + sigma_from_timestep = 1.0 - timestep.to(device=sample.device, dtype=sample.dtype) ratio = sigma_next / sigma - return ratio * sample.float() + (1.0 - ratio) * denoised.float() + state = sample.float() + velocity = model_output.float() + denoised_scratch = torch.empty_like(state) + torch.mul(sigma_from_timestep, velocity, out=denoised_scratch) + torch.add(state, denoised_scratch, out=denoised_scratch) + torch.mul(1.0 - ratio, denoised_scratch, out=velocity) + torch.mul(ratio, state, out=state) + torch.add(state, velocity, out=state) + return state def step_post(self): if self.video_noise_pred is None or self.audio_noise_pred is None: @@ -173,7 +212,7 @@ def step_post(self): self.video_timesteps[self.step_index], self.video_sigmas, self.step_index, - self.step_update, + self.sampling_profile, ) self.audio_latents[condition_audio_rows:] = self._step( self.audio_latents[condition_audio_rows:], @@ -181,7 +220,7 @@ def step_post(self): self.audio_timesteps[self.step_index], self.audio_sigmas, self.step_index, - self.step_update, + self.sampling_profile, ) def clear(self): diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/__init__.py b/lightx2v/models/video_encoders/hf/minimax_h3/__init__.py index 873260421..60bb1b9b1 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/__init__.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/__init__.py @@ -1 +1,3 @@ -from .video_vae import AutoencoderKLMiniMaxH3Native, MiniMaxH3VideoVAE +from .video_vae import MiniMaxH3VideoVAE + +__all__ = ["MiniMaxH3VideoVAE"] diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/sgl_exact_ops.py b/lightx2v/models/video_encoders/hf/minimax_h3/sgl_exact_ops.py new file mode 100644 index 000000000..b8146e9ad --- /dev/null +++ b/lightx2v/models/video_encoders/hf/minimax_h3/sgl_exact_ops.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical kernels used by the MiniMax-H3 video VAE.""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +# The numerical kernels below are adapted from SGLang commit +# 8ef646a5c65bd2f8922483057dddc02e2b0de18c (Apache-2.0). Keeping the +# H3-specific subset here avoids importing an SGLang checkout at runtime. + + +def _supports_nvidia_triton(tensor: torch.Tensor) -> bool: + return tensor.device.type == "cuda" and getattr(torch.version, "hip", None) is None + + +def _require_nvidia_triton(tensor: torch.Tensor, operation: str) -> None: + if tensor.device.type != "cuda": + return + if getattr(torch.version, "hip", None) is not None: + raise RuntimeError(f"{operation} exact parity kernel supports NVIDIA CUDA only") + + +@triton.jit +def _round_bf16_to_fp32(value): + """RNE-round FP32 to BF16 precision while retaining an FP32 register.""" + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _mul_rn_f32(x, y): + """Correctly-rounded FP32 multiply which cannot contract into an FMA.""" + return tl.inline_asm_elementwise( + asm="mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _packed_silu_mul_kernel( + output_ptr, + x_ptr, + num_rows, + row_stride, + output_row_stride, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + block = tl.program_id(1).to(tl.int64) + columns = block * BLOCK + tl.arange(0, BLOCK) + mask = (row < num_rows) & (columns < D) + row_base = row * row_stride + gate = tl.load(x_ptr + row_base + columns, mask=mask, other=0.0).to(tl.float32) + value = tl.load(x_ptr + row_base + D + columns, mask=mask, other=0.0).to(tl.float32) + activated = _round_bf16_to_fp32(gate * tl.sigmoid(gate)) + tl.store(output_ptr + row * output_row_stride + columns, activated * value, mask=mask) + + +@triton.jit +def _scaled_residual_add_exact_kernel( + output_ptr, + residual_ptr, + x_ptr, + scale_ptr, + numel: tl.constexpr, + width: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32) + scale = tl.load(scale_ptr + offsets % width, mask=mask) + residual = tl.load(residual_ptr + offsets, mask=mask) + tl.store(output_ptr + offsets, residual + _mul_rn_f32(x, scale), mask=mask) + + +def _try_scaled_residual_add_exact( + residual: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor | None: + if ( + torch.is_grad_enabled() + or torch.compiler.is_compiling() + or residual.dtype != torch.float32 + or x.dtype not in (torch.float16, torch.bfloat16) + or scale.dtype != torch.float32 + or not residual.is_cuda + or not _supports_nvidia_triton(residual) + or residual.device != x.device + or residual.device != scale.device + or residual.shape != x.shape + or scale.shape != (x.shape[-1],) + or not residual.is_contiguous() + or not x.is_contiguous() + or not scale.is_contiguous() + or x.numel() == 0 + ): + return None + output = torch.empty_like(residual) + block_size = 1024 + with torch.cuda.device(x.device): + _scaled_residual_add_exact_kernel[(triton.cdiv(x.numel(), block_size),)]( + output, + residual, + x, + scale, + numel=x.numel(), + width=x.shape[-1], + BLOCK_SIZE=block_size, + ) + return output + + +def _silu_mul_with_activation_rounding(hidden_states: torch.Tensor) -> torch.Tensor: + if hidden_states.shape[-1] % 2: + raise ValueError(f"SwiGLU input width must be even, got {hidden_states.shape[-1]}") + hidden_size = hidden_states.shape[-1] // 2 + if hidden_states.is_cuda and hidden_states.dtype is torch.bfloat16 and hidden_states.is_contiguous() and hidden_states.numel(): + _require_nvidia_triton(hidden_states, "H3 VAE SwiGLU") + rows = hidden_states.numel() // hidden_states.shape[-1] + output = hidden_states.new_empty(*hidden_states.shape[:-1], hidden_size) + with torch.cuda.device(hidden_states.device): + _packed_silu_mul_kernel[(rows, triton.cdiv(hidden_size, 1024))]( + output, + hidden_states, + rows, + hidden_states.shape[-1], + hidden_size, + D=hidden_size, + BLOCK=1024, + ) + return output + + gate, value = hidden_states.chunk(2, dim=-1) + return F.silu(gate).mul_(value) + + +def sgl_exact_vae_silu_mul(hidden_states: torch.Tensor) -> torch.Tensor: + return _silu_mul_with_activation_rounding(hidden_states) + + +def sgl_exact_vae_scaled_residual_add( + residual: torch.Tensor, + hidden_states: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + fused = _try_scaled_residual_add_exact(residual, hidden_states, scale) + return residual + hidden_states * scale if fused is None else fused diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 318348b2d..62e619c1c 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -13,12 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Native MiniMax-H3 video VAE encoder/decoder. - -The architecture and decode recipe are ported from the MiniMax-H3 implementation -pinned with the released checkpoint. This module intentionally depends only on -PyTorch and safetensors: no third-party model implementation is imported at -runtime. +"""MiniMax-H3 video VAE encoder/decoder with reference-aligned inference semantics. There are two explicit decode boundaries: @@ -31,9 +26,11 @@ from __future__ import annotations +import functools import gc import json import math +from contextlib import nullcontext from pathlib import Path from typing import NamedTuple @@ -43,11 +40,16 @@ import torch.nn.functional as F from loguru import logger +from lightx2v.common.ops.rope import SGLExactNeoXRope as _registered_rope # noqa: F401 +from lightx2v.models.video_encoders.hf.minimax_h3.sgl_exact_ops import ( + sgl_exact_vae_scaled_residual_add, + sgl_exact_vae_silu_mul, +) from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, load_safetensors_subset, ) -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, ROPE_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE MINIMAX_H3_PIXEL_MEAN = (0.485, 0.456, 0.406) @@ -81,6 +83,58 @@ def _component_dir(model_path: str | Path, component: str) -> Path: raise FileNotFoundError(f"Cannot find MiniMax-H3 {component!r} below {model_path}") +def _cuda_autocast_disabled(tensor: torch.Tensor): + return torch.autocast("cuda", enabled=False) if tensor.is_cuda else nullcontext() + + +def _linear_with_module_dtype( + linear: nn.Module, + tensor: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + if isinstance(linear, nn.Linear): + tensor = tensor.to(linear.weight.dtype) + elif linear.bias is not None: + tensor = tensor.to(linear.bias.dtype) + return linear(tensor).to(out_dtype) + + +def _apply_qk_norm(module: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + if ( + isinstance(module, (nn.LayerNorm, nn.RMSNorm)) + and module.weight is None + and (not isinstance(module, nn.LayerNorm) or module.bias is None) + and hidden_states.is_cuda + and hidden_states.dtype in (torch.float16, torch.bfloat16) + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + ): + with torch.autocast("cuda", enabled=False): + return module(hidden_states) + return module(hidden_states.float()).to(hidden_states.dtype) + + +@functools.lru_cache(maxsize=1) +def _is_sm120() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12 + + +def _linear_without_fused_bias(linear: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + if linear.bias is None or not hidden_states.is_cuda or hidden_states.dtype != linear.weight.dtype or not _is_sm120(): + return linear(hidden_states) + output = torch.matmul(hidden_states, linear.weight.t()) + output += linear.bias + return output + + +def _replace_module_tensor(module: nn.Module, name: str, value: torch.Tensor) -> None: + current = getattr(module, name) + value = value.contiguous() + if isinstance(current, nn.Parameter): + value = nn.Parameter(value, requires_grad=current.requires_grad) + setattr(module, name, value) + + class _SwiGLU(nn.Module): """Checkpoint-compatible SwiGLU used by the ViT decoder.""" @@ -88,28 +142,41 @@ def __init__(self, dim_in: int, dim_out: int, bias: bool = True) -> None: super().__init__() self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias) + def _pack_after_load(self) -> None: + if getattr(self, "_weights_packed", False): + raise RuntimeError("MiniMax-H3 video VAE SwiGLU weights were already packed") + value_weight, gate_weight = self.proj.weight.chunk(2, dim=0) + _replace_module_tensor(self.proj, "weight", torch.cat((gate_weight, value_weight), dim=0)) + if hasattr(self.proj, "weight_scale"): + value_scale, gate_scale = self.proj.weight_scale.chunk(2, dim=0) + _replace_module_tensor(self.proj, "weight_scale", torch.cat((gate_scale, value_scale), dim=0)) + if self.proj.bias is not None: + value_bias, gate_bias = self.proj.bias.chunk(2, dim=0) + _replace_module_tensor(self.proj, "bias", torch.cat((gate_bias, value_bias), dim=0)) + self._weights_packed = True + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) - return hidden_states * F.silu(gate) + return sgl_exact_vae_silu_mul(self.proj(hidden_states)) class _FeedForward(nn.Module): + swiglu_cls = _SwiGLU + def __init__(self, dim: int, mult: int = 4, bias: bool = True) -> None: super().__init__() inner_dim = int(dim * mult) # Keep the original ``net.0.proj`` and ``net.2`` parameter names. self.net = nn.ModuleList( [ - _SwiGLU(dim, inner_dim, bias=bias), + self.swiglu_cls(dim, inner_dim, bias=bias), nn.Dropout(0.0), nn.Linear(inner_dim, dim, bias=bias), ] ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - for module in self.net: - hidden_states = module(hidden_states) - return hidden_states + hidden_states = self.net[0](hidden_states) + return _linear_without_fused_bias(self.net[2], hidden_states) class MiniMaxH3VideoCausalConv3d(nn.Conv3d): @@ -249,25 +316,46 @@ def forward(self, hidden_states): class MiniMaxH3VideoRotaryPosEmbed(nn.Module): """Three-axis rotary embedding used by the non-causal ViT decoder.""" - def __init__(self, dim: int, theta: float = 100.0, num_axes: int = 3) -> None: + def __init__(self, dim: int, theta: float = 100.0, num_axes: int = 3, rope_type: str = "sgl_exact_neox_rope") -> None: super().__init__() if dim % (2 * num_axes) != 0: raise ValueError(f"dim={dim} must be divisible by 2 * num_axes={2 * num_axes}") + self.rope_type = rope_type self.dim = dim self.theta = theta self.num_axes = num_axes - - def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - inv_freq = 1.0 / self.theta ** torch.arange( + inv_freq = 1 / self.theta ** torch.arange( 0, 1, 2 * self.num_axes / self.dim, dtype=torch.float32, - device=position_ids.device, + device="cpu", ) - angles = 2.0 * math.pi * position_ids[:, :, :, None] * inv_freq[None, None, None, :] - angles = angles.flatten(2, 3).tile(2).unsqueeze(2) - return angles.cos(), angles.sin() + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.shape[-1] != self.num_axes: + raise ValueError(f"Expected {self.num_axes} dimensions, got {position_ids.shape[-1]}") + with _cuda_autocast_disabled(position_ids): + angles = 2.0 * math.pi * position_ids[:, :, :, None] + angles = angles * self.inv_freq.to(position_ids.device)[None, None, None, :] + angles = angles.flatten(2, 3).tile(2).unsqueeze(2) + cos = torch.cos(angles) + sin = torch.sin(angles) + return cos.to(dtype=position_ids.dtype), sin.to(dtype=position_ids.dtype) + + def prepare( + self, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + *, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + cos, sin = rotary_emb + if cos.ndim != 4 or cos.shape[0] != 1 or cos.shape[2] != 1: + raise ValueError(f"Expected MiniMax-H3 VAE RoPE frequencies [1, tokens, 1, rotary_dim], got {cos.shape}") + rotary_emb = cos[0, :, 0], sin[0, :, 0] + rope = ROPE_REGISTER[self.rope_type](compute_dtype=dtype) + return rope.prepare_freqs(rotary_emb, rotary_dim=rotary_emb[0].shape[-1]) class MiniMaxH3VideoAttention(nn.Module): @@ -280,12 +368,14 @@ def __init__( bias: bool = True, sensitive_layer_dtype: torch.dtype = torch.float32, attn_type: str = "torch_sdpa", + rope_type: str = "sgl_exact_neox_rope", ) -> None: super().__init__() self.heads = heads self.dim_head = dim_head self.inner_dim = heads * dim_head self.sensitive_layer_dtype = sensitive_layer_dtype + self.rope = ROPE_REGISTER[rope_type]() self.calculate = ATTN_WEIGHT_REGISTER[attn_type]() self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=False) @@ -296,67 +386,62 @@ def __init__( self.to_qkv = None self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=bias), nn.Dropout(0.0)]) - def _pack_fp8_qkv(self) -> None: + def _pack_after_load(self) -> None: + if getattr(self, "_weights_packed", False): + raise RuntimeError("MiniMax-H3 video VAE QKV weights were already packed") linears = (self.to_q, self.to_k, self.to_v) linear_cls = type(linears[0]) if any(type(linear) is not linear_cls for linear in linears[1:]): - raise TypeError("MiniMax-H3 video VAE Q/K/V projections must use one FP8 linear class") + raise TypeError("MiniMax-H3 video VAE Q/K/V projections must use one linear class") + in_features = linears[0].in_features with torch.device("meta"): self.to_qkv = linear_cls( - linears[0].in_features, - sum(linear.out_features for linear in linears), + in_features, + self.inner_dim * 3, bias=linears[0].bias is not None, dtype=linears[0].bias.dtype if linears[0].bias is not None else torch.float16, ) - self.to_qkv.weight = torch.cat([linear.weight for linear in linears], dim=0) - self.to_qkv.weight_scale = torch.cat([linear.weight_scale for linear in linears], dim=0) + + def pack_output_rows(tensors: tuple[torch.Tensor, ...]) -> torch.Tensor: + trailing_shape = tensors[0].shape[1:] + packed = torch.stack( + tuple(tensor.reshape(self.heads, self.dim_head, *trailing_shape) for tensor in tensors), + dim=1, + ) + return packed.reshape(self.inner_dim * 3, *trailing_shape) + + _replace_module_tensor(self.to_qkv, "weight", pack_output_rows(tuple(linear.weight for linear in linears))) + if hasattr(linears[0], "weight_scale"): + if any(not hasattr(linear, "weight_scale") for linear in linears[1:]): + raise TypeError("MiniMax-H3 video VAE Q/K/V projections must use one quantization layout") + _replace_module_tensor( + self.to_qkv, + "weight_scale", + pack_output_rows(tuple(linear.weight_scale for linear in linears)), + ) if linears[0].bias is not None: - self.to_qkv.bias = torch.cat([linear.bias for linear in linears], dim=0) + _replace_module_tensor(self.to_qkv, "bias", pack_output_rows(tuple(linear.bias for linear in linears))) self.to_q = None self.to_k = None self.to_v = None - - @staticmethod - def _apply_rotary( - hidden_states: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - ) -> torch.Tensor: - rotary_dim = cos.shape[-1] - rotary, passthrough = hidden_states[..., :rotary_dim], hidden_states[..., rotary_dim:] - first, second = rotary.chunk(2, dim=-1) - rotated = torch.cat([-second, first], dim=-1) - return torch.cat([rotary * cos + rotated * sin, passthrough], dim=-1) + self._weights_packed = True def forward( self, hidden_states: torch.Tensor, - rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + rotary_emb: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape if self.to_qkv is None: - query = self.to_q(hidden_states) - key = self.to_k(hidden_states) - value = self.to_v(hidden_states) - else: - query, key, value = self.to_qkv(hidden_states).split(self.inner_dim, dim=-1) - query = query.unflatten(2, (self.heads, self.dim_head)) - key = key.unflatten(2, (self.heads, self.dim_head)) - value = value.unflatten(2, (self.heads, self.dim_head)) - - infer_dtype = query.dtype - if self.sensitive_layer_dtype != infer_dtype: - query = query.to(self.sensitive_layer_dtype) - key = key.to(self.sensitive_layer_dtype) - query = self.norm_q(query) - key = self.norm_k(key) - if self.sensitive_layer_dtype != infer_dtype: - query = query.to(infer_dtype) - key = key.to(infer_dtype) + raise RuntimeError("MiniMax-H3 video VAE QKV weights must be packed after checkpoint loading") + qkv = self.to_qkv(hidden_states) + qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head) + query, key, value = torch.chunk(qkv, 3, dim=-1) + query = _apply_qk_norm(self.norm_q, query) + key = _apply_qk_norm(self.norm_k, key) if rotary_emb is not None: - cos, sin = (value.to(query.dtype) for value in rotary_emb) - query = self._apply_rotary(query, cos, sin) - key = self._apply_rotary(key, cos, sin) + query, key = self.rope.apply(query, key, rotary_emb, materialize=True) hidden_states = self.calculate.apply( query, @@ -364,11 +449,15 @@ def forward( value, max_seqlen_q=query.shape[1], max_seqlen_kv=key.shape[1], - ).view(query.shape[0], query.shape[1], self.inner_dim) + softmax_scale=self.dim_head**-0.5, + ).view(batch_size, seq_len, self.inner_dim) return self.to_out[0](hidden_states) class MiniMaxH3VideoTransformerBlock(nn.Module): + attention_cls = MiniMaxH3VideoAttention + feed_forward_cls = _FeedForward + def __init__( self, dim: int, @@ -380,12 +469,13 @@ def __init__( infer_dtype: torch.dtype = torch.float16, sensitive_layer_dtype: torch.dtype = torch.float32, attn_type: str = "torch_sdpa", + rope_type: str = "sgl_exact_neox_rope", ) -> None: super().__init__() self.infer_dtype = infer_dtype self.sensitive_layer_dtype = sensitive_layer_dtype self.norm1 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) - self.attn = MiniMaxH3VideoAttention( + self.attn = self.attention_cls( dim=dim, heads=heads, dim_head=dim_head, @@ -393,37 +483,33 @@ def __init__( bias=bias, sensitive_layer_dtype=sensitive_layer_dtype, attn_type=attn_type, + rope_type=rope_type, ) self.scale1 = nn.Parameter(torch.zeros(dim)) self.norm2 = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) - self.ff = _FeedForward(dim, mult=ffn_mult, bias=bias) + self.ff = self.feed_forward_cls(dim, mult=ffn_mult, bias=bias) self.scale2 = nn.Parameter(torch.zeros(dim)) def forward( self, hidden_states: torch.Tensor, - rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + rotary_emb: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: - norm_hidden_states = self.norm1(hidden_states) - if self.sensitive_layer_dtype != self.infer_dtype: - norm_hidden_states = norm_hidden_states.to(self.infer_dtype) + norm_hidden_states = self.norm1(hidden_states.float()).to(self.infer_dtype) attention_output = self.attn(norm_hidden_states, rotary_emb) - if self.sensitive_layer_dtype != self.infer_dtype: - attention_output = attention_output.to(self.sensitive_layer_dtype) - hidden_states = hidden_states + attention_output * self.scale1 + hidden_states = sgl_exact_vae_scaled_residual_add(hidden_states, attention_output, self.scale1) - norm_hidden_states = self.norm2(hidden_states) - if self.sensitive_layer_dtype != self.infer_dtype: - norm_hidden_states = norm_hidden_states.to(self.infer_dtype) + norm_hidden_states = self.norm2(hidden_states.float()).to(self.infer_dtype) feed_forward_output = self.ff(norm_hidden_states) - if self.sensitive_layer_dtype != self.infer_dtype: - feed_forward_output = feed_forward_output.to(self.sensitive_layer_dtype) - return hidden_states + feed_forward_output * self.scale2 + return sgl_exact_vae_scaled_residual_add(hidden_states, feed_forward_output, self.scale2) class MiniMaxH3VideoViTDecoder3d(nn.Module): """Non-causal ViT decoder with register and zero class tokens.""" + rope_cls = MiniMaxH3VideoRotaryPosEmbed + block_cls = MiniMaxH3VideoTransformerBlock + def __init__( self, in_channels: int = 24, @@ -442,6 +528,7 @@ def __init__( sensitive_layer_dtype: torch.dtype = torch.float32, use_compile: bool = False, attn_type: str = "torch_sdpa", + rope_type: str = "sgl_exact_neox_rope", ) -> None: super().__init__() dim = num_attention_heads * attention_head_dim @@ -454,12 +541,12 @@ def __init__( self.use_compile = use_compile self.compiled_blocks = {} - self.rope = MiniMaxH3VideoRotaryPosEmbed(int(attention_head_dim * rope_dim_ratio), theta=rope_theta) + self.rope = self.rope_cls(int(attention_head_dim * rope_dim_ratio), theta=rope_theta, rope_type=rope_type) self.proj_in = nn.Linear(in_channels, dim) self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, dim)) self.transformer_blocks = nn.ModuleList( [ - MiniMaxH3VideoTransformerBlock( + self.block_cls( dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, @@ -468,6 +555,7 @@ def __init__( infer_dtype=infer_dtype, sensitive_layer_dtype=sensitive_layer_dtype, attn_type=attn_type, + rope_type=rope_type, ) for _ in range(num_layers) ] @@ -480,7 +568,7 @@ def _run_block( block_index: int, block: MiniMaxH3VideoTransformerBlock, hidden_states: torch.Tensor, - rotary_emb: tuple[torch.Tensor, torch.Tensor], + rotary_emb: tuple[torch.Tensor, ...], ) -> torch.Tensor: if not self.use_compile: return block(hidden_states, rotary_emb) @@ -493,52 +581,79 @@ def _run_block( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, num_channels, num_frames, height, width = hidden_states.shape - hidden_states = hidden_states.permute(0, 2, 3, 4, 1).reshape(batch_size, num_frames * height * width, num_channels) - hidden_states = self.proj_in(hidden_states) - num_patches = hidden_states.shape[1] - if self.sensitive_layer_dtype != self.infer_dtype: - hidden_states = hidden_states.to(self.sensitive_layer_dtype) + input_dtype = hidden_states.dtype + hidden_states = hidden_states.view( + batch_size, + num_channels, + num_frames, + 1, + height, + 1, + width, + 1, + ) + hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7) + hidden_states = hidden_states.reshape(batch_size, num_frames * height * width, num_channels) - register_tokens = self.register_tokens.expand(batch_size, -1, -1) - cls_token = torch.zeros_like(hidden_states[:, :1, :]) - hidden_states = torch.cat([hidden_states, register_tokens, cls_token], dim=1) + with _cuda_autocast_disabled(hidden_states): + hidden_states = _linear_with_module_dtype(self.proj_in, hidden_states, input_dtype) + num_patches = hidden_states.shape[1] + hidden_states = torch.cat( + ( + hidden_states, + self.register_tokens.expand(batch_size, -1, -1), + torch.zeros_like(hidden_states[:, 0:1, :]), + ), + dim=1, + ) - grids = [2.0 * (torch.arange(0.5, size, dtype=torch.float32, device=hidden_states.device) / size) - 1.0 for size in (num_frames, height, width)] - position_ids = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=-1).flatten(0, 2) - position_ids = position_ids.unsqueeze(0).expand(batch_size, -1, -1) - suffix_ids = position_ids.new_zeros((batch_size, self.num_register_tokens + 1, 3)) - rotary_emb = self.rope(torch.cat([position_ids, suffix_ids], dim=1)) + coords = [] + for size in (num_frames, height, width): + axis = torch.arange(0.5, size, dtype=input_dtype, device=hidden_states.device) + axis = axis / size + axis = 2.0 * axis - 1.0 + coords.append(axis) + position_ids = torch.stack(torch.meshgrid(*coords, indexing="ij"), dim=-1) + position_ids = position_ids.flatten(0, 2).unsqueeze(0).expand(batch_size, -1, -1) + suffix_ids = torch.zeros( + (batch_size, self.num_register_tokens + 1, 3), + device=hidden_states.device, + dtype=position_ids.dtype, + ) + position_ids = torch.cat((position_ids, suffix_ids), dim=1) + rotary_dtype = torch.get_autocast_dtype("cuda") if hidden_states.is_cuda and torch.is_autocast_enabled("cuda") else self.infer_dtype + rotary_emb = self.rope.prepare(self.rope(position_ids), dtype=rotary_dtype) for block_index, block in enumerate(self.transformer_blocks): hidden_states = self._run_block(block_index, block, hidden_states, rotary_emb) hidden_states = self.norm_out(hidden_states) - if self.sensitive_layer_dtype != self.infer_dtype: - hidden_states = hidden_states.to(self.infer_dtype) - hidden_states = self.proj_out(hidden_states)[:, :num_patches, :] - patch_size, patch_size_t = self.patch_size, self.patch_size_t - hidden_states = hidden_states.view( + with _cuda_autocast_disabled(hidden_states): + output = _linear_with_module_dtype(self.proj_out, hidden_states, hidden_states.dtype) + output = output[:, :num_patches, :] + + video_frames = num_frames * self.patch_size_t + video_height = height * self.patch_size + video_width = width * self.patch_size + output = output.view( batch_size, num_frames, height, width, self.out_channels, - patch_size_t, - patch_size, - patch_size, - ) - hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous() - return hidden_states.reshape( - batch_size, - self.out_channels, - num_frames * patch_size_t, - height * patch_size, - width * patch_size, + self.patch_size_t, + self.patch_size, + self.patch_size, ) + output = output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous() + return output.reshape(batch_size, self.out_channels, video_frames, video_height, video_width) class MiniMaxH3VideoVAE(nn.Module): - """H3 video VAE with original or quantized checkpoint loading.""" + """H3 video VAE with BF16 or FP8 checkpoint loading.""" + + decoder_cls = MiniMaxH3VideoViTDecoder3d + encoder_infer_dtype = torch.float32 def __init__( self, @@ -550,12 +665,15 @@ def __init__( sensitive_layer_dtype: torch.dtype = torch.float32, use_compile: bool = False, attn_type: str = "torch_sdpa", + rope_type: str = "sgl_exact_neox_rope", ) -> None: super().__init__() if quant_scheme not in {None, "fp8-musa", "fp8-sgl"}: raise NotImplementedError(f"Unsupported MiniMax-H3 video VAE quantization scheme: {quant_scheme!r}") if attn_type not in {"torch_sdpa", "sage_attn2"}: raise ValueError(f"Unsupported MiniMax-H3 video VAE attention type: {attn_type!r}; expected torch_sdpa or sage_attn2") + if sensitive_layer_dtype != torch.float32: + raise ValueError("MiniMax-H3 video VAE requires vae_sensitive_layer_dtype='fp32'") self.config = dict(config) self.execution_device = torch.device(device or AI_DEVICE) self.cpu_offload = cpu_offload @@ -584,13 +702,13 @@ def __init__( norm_num_groups=int(config.get("norm_num_groups", 32)), norm_eps=float(config.get("norm_eps", 1e-6)), spatial_padding_mode=config.get("spatial_padding_mode", "reflect"), - infer_dtype=self.infer_dtype, + infer_dtype=self.encoder_infer_dtype, sensitive_layer_dtype=self.sensitive_layer_dtype, ) self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) - self.decoder = MiniMaxH3VideoViTDecoder3d( + self.decoder = self.decoder_cls( in_channels=latent_channels, out_channels=out_channels, patch_size=self.spatial_compression_ratio, @@ -607,6 +725,7 @@ def __init__( sensitive_layer_dtype=self.sensitive_layer_dtype, use_compile=use_compile, attn_type=attn_type, + rope_type=rope_type, ) if quant_scheme is not None: self._replace_decoder_linears_with_fp8(self.decoder.transformer_blocks) @@ -645,10 +764,6 @@ def _replace_decoder_linears_with_fp8(self, module: nn.Module) -> None: else: self._replace_decoder_linears_with_fp8(child) - def _pack_decoder_fp8_qkv(self) -> None: - for block in self.decoder.transformer_blocks: - block.attn._pack_fp8_qkv() - def _make_fp8_linear(self, linear: nn.Linear) -> nn.Module: if self.quant_scheme == "fp8-musa": from lightx2v.models.input_encoders.hf.q_linear import MusaQuantLinearFp8 as linear_cls @@ -675,16 +790,29 @@ def _reset_runtime_buffers(self) -> None: self._buffers["pixel_mean"] = torch.tensor(MINIMAX_H3_PIXEL_MEAN, dtype=self.sensitive_layer_dtype) self._buffers["pixel_std"] = torch.tensor(MINIMAX_H3_PIXEL_STD, dtype=self.sensitive_layer_dtype) + def _post_load(self) -> None: + for block in self.decoder.transformer_blocks: + block.attn._pack_after_load() + block.ff.net[0]._pack_after_load() + def _prepare_inference_dtypes(self) -> None: - # Keep normalization, residuals, and encoder boundaries in the - # sensitive dtype; bulk convolution and matrix multiplication use FP16. - for module in self.encoder.down_blocks.modules(): - if isinstance(module, nn.Conv3d): - module.to(dtype=self.infer_dtype) - self.post_quant_conv.to(dtype=self.infer_dtype) - for module in self.decoder.modules(): - if isinstance(module, nn.Linear): - module.to(dtype=self.infer_dtype) + for block in self.decoder.transformer_blocks: + for linear in ( + block.attn.to_qkv, + block.attn.to_out[0], + block.ff.net[0].proj, + block.ff.net[2], + ): + linear.to(dtype=self.infer_dtype) + + def _cast_decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + return latents + + def _decode_context(self, latents: torch.Tensor): + return torch.autocast("cuda", dtype=self.infer_dtype) if latents.is_cuda else nullcontext() + + def _return_cpu_by_default(self) -> bool: + return False @classmethod def from_pretrained( @@ -698,6 +826,7 @@ def from_pretrained( sensitive_layer_dtype: torch.dtype = torch.float32, use_compile: bool = False, attn_type: str = "torch_sdpa", + rope_type: str = "sgl_exact_neox_rope", ) -> "MiniMaxH3VideoVAE": vae_dir = _component_dir(model_path, "vae") if (checkpoint_path is None) != (quant_scheme is None): @@ -705,8 +834,7 @@ def from_pretrained( weight_path = checkpoint_path if checkpoint_path is not None else vae_dir with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: config = json.load(handle) - - # The released decoder is several GiB. Constructing it on meta avoids + # The released decoder is several GiB. Constructing it on meta avoids # allocating and then immediately overwriting random initialized weights. with torch.device("meta"): model = cls( @@ -717,12 +845,11 @@ def from_pretrained( sensitive_layer_dtype=sensitive_layer_dtype, use_compile=use_compile, attn_type=attn_type, + rope_type=rope_type, ) model._reset_runtime_buffers() model.load_report = load_safetensors_subset(model, weight_path) - if quant_scheme is not None: - # Pack only after loading the checkpoint's original Q/K/V keys. - model._pack_decoder_fp8_qkv() + model._post_load() model._prepare_inference_dtypes() model.eval().requires_grad_(False) if not cpu_offload: @@ -791,7 +918,18 @@ def _split_tiles(self, length: int, tile_size: int, min_overlap: int) -> tuple[l return starts, [tile_size] * num_tiles, overlaps @staticmethod - def _blend(a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int) -> torch.Tensor: + def _blend_values(a: torch.Tensor, b: torch.Tensor, weight_a: torch.Tensor, weight_b: torch.Tensor) -> torch.Tensor: + blended = a * weight_a + blended.add_(b * weight_b) + return blended + + def _blend( + self, + a: torch.Tensor, + b: torch.Tensor, + blend_extent: int, + dim: int, + ) -> torch.Tensor: blend_extent = min(a.shape[dim], b.shape[dim], blend_extent) if blend_extent <= 0: return b @@ -805,7 +943,7 @@ def _blend(a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int) -> tor slice_a[dim] = slice(-blend_extent, None) slice_b = [slice(None)] * b.ndim slice_b[dim] = slice(0, blend_extent) - blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b + blended = self._blend_values(a[tuple(slice_a)], b[tuple(slice_b)], weight_a, weight_b) if blend_extent == b.shape[dim]: return blended @@ -977,12 +1115,13 @@ def _encode_parallel(self, pixels: torch.Tensor, video: bool) -> torch.Tensor: @staticmethod def _sample_posterior(moments: torch.Tensor, generator: torch.Generator) -> torch.Tensor: - mean, logvar = torch.chunk(moments, 2, dim=1) + parameters = moments.to(dtype=torch.float32) + mean, logvar = torch.chunk(parameters, 2, dim=1) logvar = torch.clamp(logvar, -30.0, 20.0) - # Diffusers' randn_tensor preserves a CPU generator by drawing on CPU - # and moving the result, even when the posterior itself is on CUDA. - noise = torch.randn(mean.shape, generator=generator, device="cpu", dtype=mean.dtype).to(mean.device) - return mean + torch.exp(0.5 * logvar) * noise + std = logvar.mul(0.5).exp_() + noise = torch.randn(mean.shape, generator=generator) + noise = noise.to(device=parameters.device) + return noise.mul_(std).add_(mean) def _sample_condition_latents(self, moments: torch.Tensor) -> torch.Tensor: generator = torch.Generator(device="cpu").manual_seed(42) @@ -990,14 +1129,33 @@ def _sample_condition_latents(self, moments: torch.Tensor) -> torch.Tensor: return self.normalize_latents(latents) def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor: - mean = self.latents_mean.to(latents.device).view(1, -1, 1, 1, 1) - std = self.latents_std.to(latents.device).view(1, -1, 1, 1, 1) - return (latents.to(self.sensitive_layer_dtype) - mean) / std + result_device = latents.device + latents_cpu = latents.detach().to(device="cpu", dtype=torch.float32) + mean = self.latents_mean.detach().to(device="cpu", dtype=torch.float32).view(1, -1, 1, 1, 1) + std = self.latents_std.detach().to(device="cpu", dtype=torch.float32).view(1, -1, 1, 1, 1) + return latents_cpu.sub_(mean).div_(std).to(result_device) + + @staticmethod + def prepare_reference_pixels(pixels: torch.Tensor) -> torch.Tensor: + return pixels + + def preprocess(self, pixels: torch.Tensor, *, video: bool = False) -> torch.Tensor: + if pixels.dtype == torch.uint8: + if video: + frames = pixels[0].transpose(0, 1).to(torch.float32).div_(255.0) + mean = self.pixel_mean.to(frames.device).view(1, -1, 1, 1) + std = self.pixel_std.to(frames.device).view(1, -1, 1, 1) + frames.sub_(mean).div_(std) + return frames.contiguous().transpose(0, 1).unsqueeze(0) + images = pixels.squeeze(2).to(torch.float32).div_(255.0) + mean = self.pixel_mean.to(images.device).view(1, -1, 1, 1) + std = self.pixel_std.to(images.device).view(1, -1, 1, 1) + images.sub_(mean).div_(std) + return images.contiguous().unsqueeze(2) - def preprocess(self, pixels: torch.Tensor) -> torch.Tensor: mean = self.pixel_mean.to(pixels.device).view(1, -1, 1, 1, 1) std = self.pixel_std.to(pixels.device).view(1, -1, 1, 1, 1) - return (pixels.to(self.sensitive_layer_dtype) - mean) / std + return pixels.to(self.sensitive_layer_dtype).sub_(mean).div_(std) def encode_condition(self, pixels: torch.Tensor, *, video: bool = False, return_cpu: bool = True) -> torch.Tensor: """Encode an RGB ``[1,3,F,H,W]`` reference with the released seed-42 posterior.""" @@ -1005,7 +1163,8 @@ def encode_condition(self, pixels: torch.Tensor, *, video: bool = False, return_ if pixels.ndim != 5 or pixels.shape[0] != 1 or pixels.shape[1] != 3: raise ValueError(f"reference pixels must be [1,3,F,H,W], got {tuple(pixels.shape)}") device = self._activate() - pixels = self.preprocess(pixels.to(device=device, dtype=self.sensitive_layer_dtype)) + pixels = pixels.to(device=device) + pixels = self.preprocess(pixels, video=video) with torch.no_grad(): if self.encode_parallel: latents = self._encode_parallel(pixels, video) @@ -1050,7 +1209,7 @@ def _gather_tiles(local_tiles: list[torch.Tensor], task_counts: list[int]) -> li # Each worker stacks and sends its tiles exactly once. Rank 0 keeps its # local list and submits one receive per worker; there is no per-tile P2P. if rank != 0: - local_tile_batch = torch.stack(local_tiles) + local_tile_batch = torch.stack(local_tiles).contiguous() send_op = dist.P2POp(dist.isend, local_tile_batch, 0) for request in dist.batch_isend_irecv([send_op]): request.wait() @@ -1091,6 +1250,43 @@ def _decode_parallel( return self._gather_tiles(local_tiles, task_counts) + def _decode_temporal_frame_plan( + self, + latents: torch.Tensor, + num_clips: int, + pad_tokens: int, + ) -> tuple[int, int, int]: + """Compute logical, padding, and output frame counts.""" + chunk_num_frames = self.tokens_chunk_size * self.temporal_compression_ratio + split_count = int(self.token_drop > 0) + 1 + logical_frames = 0 + final_overlap_frames = 0 + + for clip_index in range(num_clips): + token_start = clip_index * self.tokens_chunk_size + token_end = token_start + self.tokens_chunk_size + self.token_overlap + clip_token_count = max( + 0, + min(token_end, latents.shape[2]) - min(token_start, latents.shape[2]), + ) + clip_frame_count = clip_token_count * self.temporal_compression_ratio + for overlap_index in range(split_count): + frame_start = overlap_index * chunk_num_frames + frame_end = min(frame_start + chunk_num_frames, clip_frame_count) + part_frames = max(0, frame_end - frame_start - self.frame_pre_padding) + if overlap_index == 0: + logical_frames += part_frames + else: + final_overlap_frames = part_frames + logical_frames += final_overlap_frames + + pad_frames = 0 + if pad_tokens > 0: + intra_tail = self.clip_length % self.temporal_compression_ratio + num_tokens_before_pad = latents.shape[2] - pad_tokens + pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) % self.tokens_chunk_size == 0 else self.temporal_compression_ratio for offset in range(pad_tokens)) + return logical_frames, pad_frames, logical_frames - pad_frames + def _decode(self, latents: torch.Tensor) -> torch.Tensor | None: """Decode tiled latents and assemble the full video on rank 0. @@ -1124,7 +1320,36 @@ def _decode(self, latents: torch.Tensor) -> torch.Tensor | None: if dist.get_rank() != 0: return None - decoded_chunks: list[torch.Tensor] = [] + logical_frames, pad_frames, output_frames = self._decode_temporal_frame_plan( + latents, + num_clips, + pad_tokens, + ) + if output_frames <= 0: + raise ValueError(f"Video VAE decode planned non-positive output frame count {output_frames} (logical={logical_frames}, pad={pad_frames})") + + decoded = None + write_pos = 0 + observed_frames = 0 + dropped_frames = 0 + + def write_part(part: torch.Tensor) -> None: + nonlocal decoded, write_pos, observed_frames, dropped_frames + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + observed_frames += part_frames + if decoded is None: + output_shape = list(part.shape) + output_shape[2] = output_frames + decoded = torch.empty(output_shape, dtype=part.dtype, device=part.device) + remaining = output_frames - write_pos + copy_frames = min(part_frames, max(0, remaining)) + if copy_frames > 0: + decoded[:, :, write_pos : write_pos + copy_frames].copy_(part[:, :, :copy_frames]) + write_pos += copy_frames + dropped_frames += part_frames - copy_frames + overlap = None for clip_index in range(num_clips): clip_start = clip_index * tiles_per_clip @@ -1138,6 +1363,10 @@ def _decode(self, latents: torch.Tensor) -> torch.Tensor | None: spatial_layout.height_overlaps, spatial_layout.width_overlaps, ) + if self.decode_parallel: + # Release completed views into P2P receive buffers. + all_tiles[clip_start:clip_end] = [None] * tiles_per_clip + del clip_tiles for overlap_index in range(int(self.token_drop > 0) + 1): frame_start = overlap_index * chunk_num_frames @@ -1146,18 +1375,21 @@ def _decode(self, latents: torch.Tensor) -> torch.Tensor | None: if overlap_index == 0: if overlap is not None: chunk = self._blend(overlap, chunk, self.frame_overlap, dim=-3) - decoded_chunks.append(chunk) + overlap = None + write_part(chunk) else: - overlap = chunk + # Break the view's reference to the full decoded clip. + overlap = chunk.contiguous() + del chunk, clip if overlap is not None: - decoded_chunks.append(overlap) - - decoded = torch.cat(decoded_chunks, dim=2) - if pad_tokens > 0: - intra_tail = self.clip_length % temporal_ratio - num_tokens_before_pad = latents.shape[2] - pad_tokens - pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) % tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens)) - decoded = decoded[:, :, :-pad_frames] + write_part(overlap) + overlap = None + if decoded is None: + raise RuntimeError("Video VAE temporal assembly produced no output tensor") + if observed_frames != logical_frames or dropped_frames != pad_frames or write_pos != output_frames: + raise RuntimeError( + f"Video VAE temporal frame plan mismatch: observed={observed_frames} logical={logical_frames}, dropped={dropped_frames} pad={pad_frames}, written={write_pos} output={output_frames}" + ) return decoded def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: @@ -1166,9 +1398,20 @@ def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: return latents.to(self.sensitive_layer_dtype) * std + mean def postprocess(self, video: torch.Tensor) -> torch.Tensor: - mean = self.pixel_mean.to(device=video.device).view(1, -1, 1, 1, 1) - std = self.pixel_std.to(device=video.device).view(1, -1, 1, 1, 1) - return (video.to(self.sensitive_layer_dtype) * std + mean).clamp_(0, 1) + batch_size, channels, frames, height, width = video.shape + inverse_mean = video.new_tensor(tuple(-mean / std for mean, std in zip(MINIMAX_H3_PIXEL_MEAN, MINIMAX_H3_PIXEL_STD))) + inverse_std = video.new_tensor(tuple(1.0 / std for std in MINIMAX_H3_PIXEL_STD)) + video = video.permute(0, 2, 1, 3, 4).reshape(batch_size * frames, channels, height, width) + video = video.clone().sub_(inverse_mean[:, None, None]).div_(inverse_std[:, None, None]) + video.clamp_(0, 1) + return video.reshape(batch_size, frames, channels, height, width).permute(0, 2, 1, 3, 4).contiguous() + + @staticmethod + def to_uint8_frames(video: torch.Tensor) -> torch.Tensor: + if video.ndim != 5 or video.shape[0] != 1 or video.shape[1] != 3: + raise ValueError(f"decoded H3 video must be [1,3,F,H,W], got {tuple(video.shape)}") + pixels = video[0].permute(1, 2, 3, 0).float() * 255.0 + return pixels.clamp_(0, 255).to(torch.uint8).contiguous().cpu() def _activate(self) -> torch.device: if self.cpu_offload: @@ -1195,9 +1438,9 @@ def _run_decode( latents = latents.to(device=device, dtype=self.sensitive_layer_dtype) if denormalize: latents = self.denormalize_latents(latents) - if self.sensitive_layer_dtype != self.infer_dtype: - latents = latents.to(self.infer_dtype) - with torch.no_grad(): + latents = self._cast_decode_latents(latents) + decode_context = self._decode_context(latents) + with torch.no_grad(), decode_context: video = self._decode(latents) if video is None: return None @@ -1208,7 +1451,7 @@ def _run_decode( video = video.float() if return_cpu is None: - return_cpu = self.cpu_offload + return_cpu = self._return_cpu_by_default() if return_cpu: video = video.cpu() return video @@ -1231,7 +1474,7 @@ def decode(self, latents: torch.Tensor, *, return_cpu: bool | None = None) -> to Args: latents: ``[B, 24, T, H, W]`` diffusion-space video latents. - return_cpu: Move the result to CPU. Defaults to ``cpu_offload``. + return_cpu: Move the result to CPU. Defaults to ``False``. """ return self._run_decode( @@ -1240,7 +1483,3 @@ def decode(self, latents: torch.Tensor, *, return_cpu: bool | None = None) -> to postprocess=True, return_cpu=return_cpu, ) - - -# Explicit alias for callers that prefer the upstream autoencoder naming. -AutoencoderKLMiniMaxH3Native = MiniMaxH3VideoVAE diff --git a/lightx2v/utils/ltx2_media_io.py b/lightx2v/utils/ltx2_media_io.py index eb03353e5..3db54965b 100755 --- a/lightx2v/utils/ltx2_media_io.py +++ b/lightx2v/utils/ltx2_media_io.py @@ -1,11 +1,15 @@ import logging import math +import shutil +import subprocess +import tempfile from collections.abc import Generator, Iterator, Mapping from fractions import Fraction from io import BytesIO from pathlib import Path import av +import imageio_ffmpeg import numpy as np import torch from PIL import ExifTags, Image, ImageCms @@ -273,6 +277,89 @@ def all_tiles(first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch. logger.info(f"Video saved to {output_path}") +def encode_video_sglang_compatible( + video: torch.Tensor, + fps: int, + audio: Audio, + output_path: str, + *, + ffmpeg_exe: str | None = None, + crf: int = 25, + threads: int = 24, +) -> None: + if video.ndim != 4 or video.shape[-1] != 3 or video.dtype != torch.uint8: + raise ValueError(f"Expected uint8 video [frames,height,width,3], got {tuple(video.shape)} {video.dtype}") + _, height, width, _ = video.shape + if ffmpeg_exe is None: + try: + ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe() + except RuntimeError: + ffmpeg_exe = shutil.which("ffmpeg") + resolved_ffmpeg = str(Path(ffmpeg_exe).resolve()) if ffmpeg_exe and Path(ffmpeg_exe).is_file() else shutil.which(ffmpeg_exe or "") + if resolved_ffmpeg is None: + raise FileNotFoundError(f"SGLang-compatible ffmpeg was not found: {ffmpeg_exe!r}") + ffmpeg_exe = resolved_ffmpeg + + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + from scipy.io import wavfile + + waveform = audio.waveform.detach().float().clamp(-1.0, 1.0).transpose(0, 1).cpu().numpy() + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as handle: + temporary_wav = handle.name + + try: + wavfile.write(temporary_wav, audio.sampling_rate, waveform) + command = [ + ffmpeg_exe, + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{width}x{height}", + "-pix_fmt", + "rgb24", + "-r", + f"{fps:.02f}", + "-i", + "pipe:0", + "-i", + temporary_wav, + ] + + command += ["-vcodec", "libx264", "-pix_fmt", "yuv420p", "-crf", str(crf)] + if width % 16 or height % 16: + output_width = width if width % 16 == 0 else width + 16 - width % 16 + output_height = height if height % 16 == 0 else height + 16 - height % 16 + command += ["-vf", f"scale={output_width}:{output_height}"] + command += ["-threads", str(threads), "-acodec", "aac", "-map", "0:v:0", "-map", "1:a:0"] + command += ["-v", "warning", output_path] + + with tempfile.TemporaryFile() as stderr_file: + process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=stderr_file) + assert process.stdin is not None + try: + video_cpu = video.contiguous().cpu().numpy() + for frame in video_cpu: + process.stdin.write(frame.tobytes()) + process.stdin.close() + process.stdin = None + returncode = process.wait() + finally: + if process.stdin is not None: + process.stdin.close() + if process.poll() is None: + process.kill() + process.wait() + if returncode: + stderr_file.seek(0) + raise subprocess.CalledProcessError(returncode, command, stderr=stderr_file.read()) + finally: + Path(temporary_wav).unlink(missing_ok=True) + logger.info(f"Video saved through SGLang-compatible ffmpeg to {output_path}") + + def _ltx25_bt709_yuv420p(frames: torch.Tensor) -> torch.Tensor: """Match LTX-2.5's GPU float-RGB to limited-range BT.709 I420 conversion.""" if frames.ndim != 4 or frames.shape[-1] != 3: diff --git a/scripts/minimax_h3/run_minimax_h3_ref2av_turbo_4step_tp8.sh b/scripts/minimax_h3/run_minimax_h3_ref2av_turbo_4step_tp8.sh new file mode 100755 index 000000000..54bc3ab51 --- /dev/null +++ b/scripts/minimax_h3/run_minimax_h3_ref2av_turbo_4step_tp8.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +set -euo pipefail + +lightx2v_path="${LIGHTX2V_PATH:-/data/wushuo1/LightX2V}" +model_path="${MODEL_PATH:-/data/wushuo1/models/wyr_models/minimax_h3/h3_hf_bf16}" +test_case_path="${TEST_CASE_PATH:-/data/wushuo1/商汤H3测试数据/Case1_mecha1}" +config_path="${CONFIG_PATH:-${lightx2v_path}/configs/minimax_h3/dmd/minimax_h3_ref2av_turbo_4step_tp8.json}" +output_path="${OUTPUT_PATH:-${test_case_path}/output_minimax_h3_ref2av_turbo_4step_bf16.mp4}" +log_path="${LOG_PATH:-${test_case_path}/run_minimax_h3_ref2av_turbo_4step_bf16_tp8.log}" + +export PLATFORM="${PLATFORM:-cuda}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 +export PYTHONPATH="${PYTHONPATH:-}" + +source "${lightx2v_path}/scripts/base/base.sh" + +cache_manifest="/data/wushuo1/.cache/lightx2v/adaln/minimax_h3/ref2av_04steps/manifest.json" +if [[ ! -f "${cache_manifest}" ]]; then + CUDA_VISIBLE_DEVICES=0 python "${lightx2v_path}/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py" \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --task ref2av +fi + +prompt="$(python3 -c 'from pathlib import Path; import sys; print(Path(sys.argv[1]).read_text(encoding="utf-8").strip(), end="")' "${test_case_path}/prompt_h3.txt")" +image_path="${test_case_path}/Picture_1.jpg,${test_case_path}/Picture_2.png,${test_case_path}/Picture_3.jpg,${test_case_path}/Picture_4.png" + +mkdir -p "$(dirname "${output_path}")" "$(dirname "${log_path}")" + +torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task ref2av \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "${prompt}" \ + --image_path "${image_path}" \ + --save_result_path "${output_path}" \ + --seed 42 \ + 2>&1 | tee "${log_path}" diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_tp_sparse.sh b/scripts/minimax_h3/run_minimax_h3_t2av_tp_sparse.sh new file mode 100755 index 000000000..7f382ed46 --- /dev/null +++ b/scripts/minimax_h3/run_minimax_h3_t2av_tp_sparse.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +lightx2v_path=/data/nvme6/gushiqiao/codes/LightX2V +model_path=/data/wushuo1/models/wyr_models/minimax_h3 + +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +source ${lightx2v_path}/scripts/base/base.sh +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +prompt='In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' + +torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \ +--model_cls minimax_h3 \ +--task t2av \ +--model_path ${model_path} \ +--config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_tp.json \ +--prompt "$prompt" \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp.mp4 \ +--seed 0 \ +--warmup > ${lightx2v_path}/save_results/minimax_h3_t2av_544p_124_8gpu_tp8.log 2>&1 &