diff --git a/app/gradio_demo.py b/app/gradio_demo.py index 57e519205..548cdcf65 100644 --- a/app/gradio_demo.py +++ b/app/gradio_demo.py @@ -198,7 +198,7 @@ def run_inference( "boundary_step_index": 2, "boundary": 0.900, "use_image_encoder": False if "wan2.2" in model_cls else True, - "rope_type": "torch", + "rope_type": "torch_complex_rope", "t5_lazy_load": lazy_load, "bucket_shape": { "0.667": [[480, 832], [544, 960], [720, 960]], diff --git a/configs/cosmos3/cosmos3_super_i2v.json b/configs/cosmos3/cosmos3_super_i2v.json index e7e7213bf..801dbf21b 100644 --- a/configs/cosmos3/cosmos3_super_i2v.json +++ b/configs/cosmos3/cosmos3_super_i2v.json @@ -10,7 +10,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json index 63724a4a3..a1d4b19b6 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json +++ b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld.json @@ -15,7 +15,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json index b62d93fd5..79ba9260b 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json +++ b/configs/cosmos3/cosmos3_super_omni_action_fd_agibotworld_multichunk.json @@ -17,7 +17,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_action_id_av.json b/configs/cosmos3/cosmos3_super_omni_action_id_av.json index 05d05ecdf..d426b7f01 100644 --- a/configs/cosmos3/cosmos3_super_omni_action_id_av.json +++ b/configs/cosmos3/cosmos3_super_omni_action_id_av.json @@ -14,7 +14,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_i2av.json b/configs/cosmos3/cosmos3_super_omni_i2av.json index 48602096e..bef82d375 100644 --- a/configs/cosmos3/cosmos3_super_omni_i2av.json +++ b/configs/cosmos3/cosmos3_super_omni_i2av.json @@ -11,7 +11,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_i2v.json b/configs/cosmos3/cosmos3_super_omni_i2v.json index 3072500ad..d1b0884d6 100644 --- a/configs/cosmos3/cosmos3_super_omni_i2v.json +++ b/configs/cosmos3/cosmos3_super_omni_i2v.json @@ -10,7 +10,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_t2av.json b/configs/cosmos3/cosmos3_super_omni_t2av.json index 48602096e..bef82d375 100644 --- a/configs/cosmos3/cosmos3_super_omni_t2av.json +++ b/configs/cosmos3/cosmos3_super_omni_t2av.json @@ -11,7 +11,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_omni_t2v.json b/configs/cosmos3/cosmos3_super_omni_t2v.json index 3072500ad..d1b0884d6 100644 --- a/configs/cosmos3/cosmos3_super_omni_t2v.json +++ b/configs/cosmos3/cosmos3_super_omni_t2v.json @@ -10,7 +10,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/cosmos3/cosmos3_super_t2i.json b/configs/cosmos3/cosmos3_super_t2i.json index e1599caf9..cf4a5eb09 100644 --- a/configs/cosmos3/cosmos3_super_t2i.json +++ b/configs/cosmos3/cosmos3_super_t2i.json @@ -10,7 +10,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "cosmos3_meta_init": true, diff --git a/configs/cosmos3/cosmos3_super_t2v.json b/configs/cosmos3/cosmos3_super_t2v.json index 3072500ad..d1b0884d6 100644 --- a/configs/cosmos3/cosmos3_super_t2v.json +++ b/configs/cosmos3/cosmos3_super_t2v.json @@ -10,7 +10,7 @@ "feature_caching": "NoCaching", "rms_norm_type": "one-pass", "attn_rms_norm_type": "one-pass", - "rope_type": "triton", + "rope_type": "cosmos3_rope", "self_attn_type": "flash_attn3", "causal_self_attn_type": "flash_attn3", "add_resolution_template": false, diff --git a/configs/flux2/flux2_dev.json b/configs/flux2/flux2_dev.json index 5c28ee9b5..416c47c7e 100644 --- a/configs/flux2/flux2_dev.json +++ b/configs/flux2/flux2_dev.json @@ -8,6 +8,6 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "text_encoder_out_layers": [10, 20, 30] } diff --git a/configs/flux2/flux2_dev_offload.json b/configs/flux2/flux2_dev_offload.json index 6f5ba2523..d81e604fc 100644 --- a/configs/flux2/flux2_dev_offload.json +++ b/configs/flux2/flux2_dev_offload.json @@ -8,7 +8,7 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "text_encoder_out_layers": [10, 20, 30], "cpu_offload": true, "offload_granularity": "model" diff --git a/configs/flux2/flux2_klein.json b/configs/flux2/flux2_klein.json index 95143f840..06d9272f2 100644 --- a/configs/flux2/flux2_klein.json +++ b/configs/flux2/flux2_klein.json @@ -8,5 +8,5 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/flux2/flux2_klein_distill.json b/configs/flux2/flux2_klein_distill.json index 2f09e7cba..53ccbffe4 100644 --- a/configs/flux2/flux2_klein_distill.json +++ b/configs/flux2/flux2_klein_distill.json @@ -8,5 +8,5 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/flux2/flux2_klein_distill_fls.json b/configs/flux2/flux2_klein_distill_fls.json index deaac5372..d53c492ef 100644 --- a/configs/flux2/flux2_klein_distill_fls.json +++ b/configs/flux2/flux2_klein_distill_fls.json @@ -8,7 +8,7 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "fls": { "enable": true, "fovea_strength": 10, diff --git a/configs/flux2/flux2_klein_distill_offload.json b/configs/flux2/flux2_klein_distill_offload.json index 31dca84d5..30949fce7 100644 --- a/configs/flux2/flux2_klein_distill_offload.json +++ b/configs/flux2/flux2_klein_distill_offload.json @@ -8,7 +8,7 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "cpu_offload": true, "offload_granularity": "model" } diff --git a/configs/flux2/flux2_klein_i2i_distill.json b/configs/flux2/flux2_klein_i2i_distill.json index 23a23674c..bfc3c57b4 100644 --- a/configs/flux2/flux2_klein_i2i_distill.json +++ b/configs/flux2/flux2_klein_i2i_distill.json @@ -8,5 +8,5 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/flux2/flux2_klein_i2i_distill_edit.json b/configs/flux2/flux2_klein_i2i_distill_edit.json index c33507962..ffb9bd4f7 100644 --- a/configs/flux2/flux2_klein_i2i_distill_edit.json +++ b/configs/flux2/flux2_klein_i2i_distill_edit.json @@ -9,5 +9,5 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/flux2/flux2_klein_i2i_distill_edit_ulysses.json b/configs/flux2/flux2_klein_i2i_distill_edit_ulysses.json index ca4c832c0..b2dfc28ea 100644 --- a/configs/flux2/flux2_klein_i2i_distill_edit_ulysses.json +++ b/configs/flux2/flux2_klein_i2i_distill_edit_ulysses.json @@ -9,7 +9,7 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "parallel": { "seq_p_size": 2, "seq_p_attn_type": "ulysses" diff --git a/configs/flux2/flux2_klein_i2i_distill_ulysses.json b/configs/flux2/flux2_klein_i2i_distill_ulysses.json index b37e604bd..3ad936ab3 100644 --- a/configs/flux2/flux2_klein_i2i_distill_ulysses.json +++ b/configs/flux2/flux2_klein_i2i_distill_ulysses.json @@ -8,7 +8,7 @@ "enable_cfg": false, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "parallel": { "seq_p_size": 2, "seq_p_attn_type": "ulysses" diff --git a/configs/flux2/flux2_klein_i2i_edit.json b/configs/flux2/flux2_klein_i2i_edit.json index f0e7f7928..b59fc3614 100644 --- a/configs/flux2/flux2_klein_i2i_edit.json +++ b/configs/flux2/flux2_klein_i2i_edit.json @@ -9,5 +9,5 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/flux2/flux2_klein_i2i_edit_cfg_parallel.json b/configs/flux2/flux2_klein_i2i_edit_cfg_parallel.json index a5dc4200a..73ef20f37 100644 --- a/configs/flux2/flux2_klein_i2i_edit_cfg_parallel.json +++ b/configs/flux2/flux2_klein_i2i_edit_cfg_parallel.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "parallel": { "cfg_p_size": 2 } diff --git a/configs/flux2/flux2_klein_i2i_edit_cfg_parallel_8step.json b/configs/flux2/flux2_klein_i2i_edit_cfg_parallel_8step.json index 1de924ab6..44f12d611 100644 --- a/configs/flux2/flux2_klein_i2i_edit_cfg_parallel_8step.json +++ b/configs/flux2/flux2_klein_i2i_edit_cfg_parallel_8step.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "parallel": { "cfg_p_size": 2 } diff --git a/configs/flux2/flux2_klein_i2i_inpaint_mask.json b/configs/flux2/flux2_klein_i2i_inpaint_mask.json index 71e545d0e..29560014c 100644 --- a/configs/flux2/flux2_klein_i2i_inpaint_mask.json +++ b/configs/flux2/flux2_klein_i2i_inpaint_mask.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "max_image_area": 1048576, "inpaint_mask_enabled": true } diff --git a/configs/flux2/flux2_klein_i2i_inpaint_mask_cache.json b/configs/flux2/flux2_klein_i2i_inpaint_mask_cache.json index 442173181..2a7a95fca 100644 --- a/configs/flux2/flux2_klein_i2i_inpaint_mask_cache.json +++ b/configs/flux2/flux2_klein_i2i_inpaint_mask_cache.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "max_image_area": 1048576, "inpaint_mask_enabled": true } diff --git a/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel.json b/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel.json index bd064145b..ac85504be 100644 --- a/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel.json +++ b/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "max_image_area": 1048576, "inpaint_mask_enabled": true, "parallel": { diff --git a/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel_cache.json b/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel_cache.json index 633a31d8a..ef5ee3b93 100644 --- a/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel_cache.json +++ b/configs/flux2/flux2_klein_i2i_inpaint_mask_cfg_parallel_cache.json @@ -9,7 +9,7 @@ "enable_cfg": true, "patch_size": 2, "tokenizer_max_length": 512, - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "max_image_area": 1048576, "inpaint_mask_enabled": true, "parallel": { diff --git a/configs/hidream_o1_image/hidream_o1_image_i2i.json b/configs/hidream_o1_image/hidream_o1_image_i2i.json index 0e20bb6dc..421a53ac9 100644 --- a/configs/hidream_o1_image/hidream_o1_image_i2i.json +++ b/configs/hidream_o1_image/hidream_o1_image_i2i.json @@ -5,7 +5,7 @@ "shift": 3.0, "scheduler_name": "default", "attn_type": "flash_attn3", - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "target_height": 2048, "target_width": 2048, "keep_original_aspect": true, diff --git a/configs/hidream_o1_image/hidream_o1_image_i2i_layout.json b/configs/hidream_o1_image/hidream_o1_image_i2i_layout.json index 6677fe04f..a02462d38 100644 --- a/configs/hidream_o1_image/hidream_o1_image_i2i_layout.json +++ b/configs/hidream_o1_image/hidream_o1_image_i2i_layout.json @@ -5,7 +5,7 @@ "shift": 1.0, "scheduler_name": "default", "attn_type": "flash_attn3", - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "target_height": 2048, "target_width": 2048, "keep_original_aspect": false, diff --git a/configs/hidream_o1_image/hidream_o1_image_i2i_multi.json b/configs/hidream_o1_image/hidream_o1_image_i2i_multi.json index 6677fe04f..a02462d38 100644 --- a/configs/hidream_o1_image/hidream_o1_image_i2i_multi.json +++ b/configs/hidream_o1_image/hidream_o1_image_i2i_multi.json @@ -5,7 +5,7 @@ "shift": 1.0, "scheduler_name": "default", "attn_type": "flash_attn3", - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "target_height": 2048, "target_width": 2048, "keep_original_aspect": false, diff --git a/configs/hidream_o1_image/hidream_o1_image_t2i.json b/configs/hidream_o1_image/hidream_o1_image_t2i.json index a273f71e8..4f419e3c4 100644 --- a/configs/hidream_o1_image/hidream_o1_image_t2i.json +++ b/configs/hidream_o1_image/hidream_o1_image_t2i.json @@ -5,7 +5,7 @@ "shift": 3.0, "scheduler_name": "default", "attn_type": "flash_attn3", - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "target_height": 2048, "target_width": 2048, "keep_original_aspect": false, diff --git a/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604.json b/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604.json index 6fdf7676a..650579fbe 100644 --- a/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604.json +++ b/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604.json @@ -11,5 +11,5 @@ "enable_cfg": false, "rms_norm_type": "one-pass", "attn_type": "flash_attn3", - "rope_type": "flashinfer" + "rope_type": "flashinfer_rope" } diff --git a/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604_dist.json b/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604_dist.json index 8db52f902..f77e3487d 100644 --- a/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604_dist.json +++ b/configs/hidream_o1_image/hidream_o1_image_t2i_dev_2604_dist.json @@ -11,7 +11,7 @@ "enable_cfg": false, "rms_norm_type": "one-pass", "attn_type": "flash_attn3", - "rope_type": "flashinfer", + "rope_type": "flashinfer_rope", "parallel": { "seq_p_size": 2, "seq_p_attn_type": "ulysses" diff --git a/configs/hidream_o1_image/mlu/hidream_o1_image_i2i_dist.json b/configs/hidream_o1_image/mlu/hidream_o1_image_i2i_dist.json index 458408a5a..87e065b01 100644 --- a/configs/hidream_o1_image/mlu/hidream_o1_image_i2i_dist.json +++ b/configs/hidream_o1_image/mlu/hidream_o1_image_i2i_dist.json @@ -5,7 +5,7 @@ "shift": 3.0, "scheduler_name": "default", "attn_type": "mlu_flash_attn", - "rope_type": "torch", + "rope_type": "torch_real_rope", "target_height": 2048, "target_width": 2048, "keep_original_aspect": true, diff --git a/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604.json b/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604.json index 534e301ae..c0fe5023d 100644 --- a/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604.json +++ b/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604.json @@ -11,7 +11,7 @@ "enable_cfg": false, "cpu_offload": false, "attn_type": "mlu_flash_attn", - "rope_type": "torch", + "rope_type": "torch_real_rope", "ln_type": "torch", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm" diff --git a/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604_dist.json b/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604_dist.json index 88495b3aa..216e9cc33 100644 --- a/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604_dist.json +++ b/configs/hidream_o1_image/mlu/hidream_o1_image_t2i_dev_2604_dist.json @@ -11,7 +11,7 @@ "enable_cfg": false, "cpu_offload": false, "attn_type": "mlu_flash_attn", - "rope_type": "torch", + "rope_type": "torch_real_rope", "ln_type": "torch", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm", diff --git a/configs/lingbot_fast/lingbot_fast_i2v.json b/configs/lingbot_fast/lingbot_fast_i2v.json index 5a2e64e02..de85c4fcb 100644 --- a/configs/lingbot_fast/lingbot_fast_i2v.json +++ b/configs/lingbot_fast/lingbot_fast_i2v.json @@ -23,5 +23,5 @@ "sink_size": 3, "kv_offload": false }, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/lingbot_fast/lingbot_fast_i2v_kv_kiviquant.json b/configs/lingbot_fast/lingbot_fast_i2v_kv_kiviquant.json index b0c6d07ef..fd431fe7c 100755 --- a/configs/lingbot_fast/lingbot_fast_i2v_kv_kiviquant.json +++ b/configs/lingbot_fast/lingbot_fast_i2v_kv_kiviquant.json @@ -29,5 +29,5 @@ }, "kv_offload": false }, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/ltx2/ltx2.json b/configs/ltx2/ltx2.json index dabc0bc72..e8814ff4f 100755 --- a/configs/ltx2/ltx2.json +++ b/configs/ltx2/ltx2.json @@ -13,6 +13,7 @@ "audio_fps": 24000, "audio_mel_bins":16, "double_precision_rope": true, + "rope_type": "torch_real_rope", "use_tiling_vae": false, - "dit_original_ckpt": "Lightricks/LTX-2/ltx-2-19b-dev.safetensors" + "dit_original_ckpt": "/LTX-2/ltx-2-19b-dev.safetensors" } diff --git a/configs/ltx2/ltx2_3.json b/configs/ltx2/ltx2_3.json index 9fb99082d..b94f3da00 100755 --- a/configs/ltx2/ltx2_3.json +++ b/configs/ltx2/ltx2_3.json @@ -14,7 +14,7 @@ "audio_mel_bins":16, "double_precision_rope": true, "use_tiling_vae": false, - "dit_original_ckpt": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", + "dit_original_ckpt": "/LTX-2.3/ltx-2.3-22b-dev.safetensors", "caption_proj_before_connector": true, "cross_attention_adaln": true, "apply_gated_attention": true, diff --git a/configs/matrix_game2/matrix_game2_gta_drive.json b/configs/matrix_game2/matrix_game2_gta_drive.json index 934708b95..7727784e6 100644 --- a/configs/matrix_game2/matrix_game2_gta_drive.json +++ b/configs/matrix_game2/matrix_game2_gta_drive.json @@ -66,5 +66,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/matrix_game2/matrix_game2_gta_drive_streaming.json b/configs/matrix_game2/matrix_game2_gta_drive_streaming.json index 50a320a7a..283969383 100644 --- a/configs/matrix_game2/matrix_game2_gta_drive_streaming.json +++ b/configs/matrix_game2/matrix_game2_gta_drive_streaming.json @@ -66,5 +66,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/matrix_game2/matrix_game2_templerun.json b/configs/matrix_game2/matrix_game2_templerun.json index adcb3c645..475b00d56 100644 --- a/configs/matrix_game2/matrix_game2_templerun.json +++ b/configs/matrix_game2/matrix_game2_templerun.json @@ -59,5 +59,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/matrix_game2/matrix_game2_templerun_streaming.json b/configs/matrix_game2/matrix_game2_templerun_streaming.json index 39f522a8f..b18aab520 100644 --- a/configs/matrix_game2/matrix_game2_templerun_streaming.json +++ b/configs/matrix_game2/matrix_game2_templerun_streaming.json @@ -66,5 +66,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/matrix_game2/matrix_game2_universal.json b/configs/matrix_game2/matrix_game2_universal.json index b66ebda78..8c9c5fa61 100644 --- a/configs/matrix_game2/matrix_game2_universal.json +++ b/configs/matrix_game2/matrix_game2_universal.json @@ -67,5 +67,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/matrix_game2/matrix_game2_universal_streaming.json b/configs/matrix_game2/matrix_game2_universal_streaming.json index 020c5a748..43de15685 100644 --- a/configs/matrix_game2/matrix_game2_universal_streaming.json +++ b/configs/matrix_game2/matrix_game2_universal_streaming.json @@ -66,5 +66,5 @@ "num_layers": 30, "out_dim": 16, "use_image_encoder": false, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/platforms/ascend_npu/longcat_image_t2i.json b/configs/platforms/ascend_npu/longcat_image_t2i.json index 1dbad9e54..d313af9cc 100644 --- a/configs/platforms/ascend_npu/longcat_image_t2i.json +++ b/configs/platforms/ascend_npu/longcat_image_t2i.json @@ -10,6 +10,6 @@ "dit_quant_scheme": "Default", "rms_norm_type": "torch", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch" } diff --git a/configs/platforms/ascend_npu/qwen_image_i2i_2511.json b/configs/platforms/ascend_npu/qwen_image_i2i_2511.json index 88bb21160..a3dbd9c0d 100755 --- a/configs/platforms/ascend_npu/qwen_image_i2i_2511.json +++ b/configs/platforms/ascend_npu/qwen_image_i2i_2511.json @@ -11,7 +11,7 @@ "cpu_offload": true, "offload_granularity": "model", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/ascend_npu/qwen_image_t2i_2512.json b/configs/platforms/ascend_npu/qwen_image_t2i_2512.json index 88039e37b..8e8eaa8c2 100755 --- a/configs/platforms/ascend_npu/qwen_image_t2i_2512.json +++ b/configs/platforms/ascend_npu/qwen_image_t2i_2512.json @@ -9,7 +9,7 @@ "cpu_offload": true, "offload_granularity": "model", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/ascend_npu/wan_t2v.json b/configs/platforms/ascend_npu/wan_t2v.json index c1f15b77f..d96db47a9 100755 --- a/configs/platforms/ascend_npu/wan_t2v.json +++ b/configs/platforms/ascend_npu/wan_t2v.json @@ -13,7 +13,7 @@ "cpu_offload": true, "offload_granularity": "model", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/ascend_npu/wan_ti2v_t2v.json b/configs/platforms/ascend_npu/wan_ti2v_t2v.json index eb078a7fb..e79464f73 100755 --- a/configs/platforms/ascend_npu/wan_ti2v_t2v.json +++ b/configs/platforms/ascend_npu/wan_ti2v_t2v.json @@ -19,7 +19,7 @@ "cpu_offload": true, "offload_granularity": "model", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "t5_cpu_offload": false, diff --git a/configs/platforms/ascend_npu/z_image_turbo_t2i.json b/configs/platforms/ascend_npu/z_image_turbo_t2i.json index 89c8fadb3..ad483206c 100755 --- a/configs/platforms/ascend_npu/z_image_turbo_t2i.json +++ b/configs/platforms/ascend_npu/z_image_turbo_t2i.json @@ -9,7 +9,7 @@ "cpu_offload": false, "offload_granularity": "model", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/enflame_gcu/wan_t2v.json b/configs/platforms/enflame_gcu/wan_t2v.json index ba6df94ca..74f316b13 100644 --- a/configs/platforms/enflame_gcu/wan_t2v.json +++ b/configs/platforms/enflame_gcu/wan_t2v.json @@ -12,7 +12,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/hygon_dcu/qwen_image_i2i_2511.json b/configs/platforms/hygon_dcu/qwen_image_i2i_2511.json index 33b12a25c..a6b5321d1 100755 --- a/configs/platforms/hygon_dcu/qwen_image_i2i_2511.json +++ b/configs/platforms/hygon_dcu/qwen_image_i2i_2511.json @@ -9,7 +9,7 @@ "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/hygon_dcu/wan_i2v.json b/configs/platforms/hygon_dcu/wan_i2v.json index e0958f0c7..0e3c76cad 100644 --- a/configs/platforms/hygon_dcu/wan_i2v.json +++ b/configs/platforms/hygon_dcu/wan_i2v.json @@ -11,7 +11,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/hygon_dcu/wan_i2v_int8.json b/configs/platforms/hygon_dcu/wan_i2v_int8.json index e15dfc774..a1e0416ac 100644 --- a/configs/platforms/hygon_dcu/wan_i2v_int8.json +++ b/configs/platforms/hygon_dcu/wan_i2v_int8.json @@ -11,7 +11,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "dit_quantized": true, diff --git a/configs/platforms/hygon_dcu/wan_t2v.json b/configs/platforms/hygon_dcu/wan_t2v.json index 3ea6a7594..a2619568a 100755 --- a/configs/platforms/hygon_dcu/wan_t2v.json +++ b/configs/platforms/hygon_dcu/wan_t2v.json @@ -12,7 +12,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3.json b/configs/platforms/intel_xpu/wan_t2v_1_3.json index f4d046aa2..68b2511f5 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3.json @@ -16,7 +16,7 @@ "vae_offload_cache": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_block_offload.json b/configs/platforms/intel_xpu/wan_t2v_1_3_block_offload.json index 8e483b595..65658a5f9 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_block_offload.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_block_offload.json @@ -16,7 +16,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload.json b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload.json index 93da6324b..7a96dc477 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload.json @@ -18,7 +18,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8.json b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8.json index eef8c7e7e..e53ab05c5 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8.json @@ -21,7 +21,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8_block.json b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8_block.json index a965a6d5a..afb369cd7 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8_block.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_lazy_offload_fp8_block.json @@ -21,7 +21,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_nocache.json b/configs/platforms/intel_xpu/wan_t2v_1_3_nocache.json index d87f521a9..78c98f41e 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_nocache.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_nocache.json @@ -15,7 +15,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "feature_caching": "NoCaching" diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_phase_offload.json b/configs/platforms/intel_xpu/wan_t2v_1_3_phase_offload.json index e2cd80070..8ebeb5548 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_phase_offload.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_phase_offload.json @@ -16,7 +16,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_t5_quant.json b/configs/platforms/intel_xpu/wan_t2v_1_3_t5_quant.json index cb7efba32..ecb5b0724 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_t5_quant.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_t5_quant.json @@ -14,7 +14,7 @@ "t5_cpu_offload": false, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "t5_quantized": true, diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_teacache.json b/configs/platforms/intel_xpu/wan_t2v_1_3_teacache.json index 7128d0eb2..0a52e3182 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_teacache.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_teacache.json @@ -15,7 +15,7 @@ "vae_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "feature_caching": "Tea", diff --git a/configs/platforms/intel_xpu/wan_t2v_1_3_xpu_flash_attn.json b/configs/platforms/intel_xpu/wan_t2v_1_3_xpu_flash_attn.json index f6dd4bc82..355487833 100644 --- a/configs/platforms/intel_xpu/wan_t2v_1_3_xpu_flash_attn.json +++ b/configs/platforms/intel_xpu/wan_t2v_1_3_xpu_flash_attn.json @@ -14,7 +14,7 @@ "t5_cpu_offload": true, "use_tiling_vae": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "t5_quantized": true, diff --git a/configs/platforms/intel_xpu/z_image_turbo_t2i.json b/configs/platforms/intel_xpu/z_image_turbo_t2i.json index e8097cad3..e2b80cf4e 100644 --- a/configs/platforms/intel_xpu/z_image_turbo_t2i.json +++ b/configs/platforms/intel_xpu/z_image_turbo_t2i.json @@ -6,7 +6,7 @@ "enable_cfg": false, "sample_guide_scale": 0.0, "patch_size": 2, - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/metax/ltx_2_3.json b/configs/platforms/metax/ltx_2_3.json index cd0fc5de4..82f4d83a7 100644 --- a/configs/platforms/metax/ltx_2_3.json +++ b/configs/platforms/metax/ltx_2_3.json @@ -11,7 +11,8 @@ "offload_granularity": "model", "norm_modulate_backend": "torch", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", + "rope_layout": "split_half", "layer_norm_type": "torch", "rms_norm_type": "torch", "num_channels_latents": 128, diff --git a/configs/platforms/metax/qwen_image_i2i_2511.json b/configs/platforms/metax/qwen_image_i2i_2511.json index e891fd88c..290d6603e 100755 --- a/configs/platforms/metax/qwen_image_i2i_2511.json +++ b/configs/platforms/metax/qwen_image_i2i_2511.json @@ -9,7 +9,7 @@ "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/metax/qwen_image_t2i_2512.json b/configs/platforms/metax/qwen_image_t2i_2512.json index 8d61a7be3..bd78f8278 100755 --- a/configs/platforms/metax/qwen_image_t2i_2512.json +++ b/configs/platforms/metax/qwen_image_t2i_2512.json @@ -9,7 +9,7 @@ "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "cpu_offload": "true", diff --git a/configs/platforms/metax/wan22_t2v.json b/configs/platforms/metax/wan22_t2v.json index e0b76de7d..73575c13a 100755 --- a/configs/platforms/metax/wan22_t2v.json +++ b/configs/platforms/metax/wan22_t2v.json @@ -19,7 +19,7 @@ "vae_cpu_offload": false, "boundary": 0.875, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/metax/wan_t2v.json b/configs/platforms/metax/wan_t2v.json index 6a14c5052..88f5b4642 100755 --- a/configs/platforms/metax/wan_t2v.json +++ b/configs/platforms/metax/wan_t2v.json @@ -12,7 +12,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/mlu/qwen_image_i2i_2511.json b/configs/platforms/mlu/qwen_image_i2i_2511.json index 8b956dc8f..ac3bf66ae 100755 --- a/configs/platforms/mlu/qwen_image_i2i_2511.json +++ b/configs/platforms/mlu/qwen_image_i2i_2511.json @@ -9,7 +9,7 @@ "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm" } diff --git a/configs/platforms/mlu/qwen_image_t2i_2512_distill.json b/configs/platforms/mlu/qwen_image_t2i_2512_distill.json index d019fc40d..e48584126 100755 --- a/configs/platforms/mlu/qwen_image_t2i_2512_distill.json +++ b/configs/platforms/mlu/qwen_image_t2i_2512_distill.json @@ -13,7 +13,7 @@ } ], "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm" } diff --git a/configs/platforms/mlu/qwen_image_t2i_2512_distill_dist.json b/configs/platforms/mlu/qwen_image_t2i_2512_distill_dist.json index 9b02e7990..42df100f5 100755 --- a/configs/platforms/mlu/qwen_image_t2i_2512_distill_dist.json +++ b/configs/platforms/mlu/qwen_image_t2i_2512_distill_dist.json @@ -13,7 +13,7 @@ } ], "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm", "parallel": { diff --git a/configs/platforms/mlu/wan_animate.json b/configs/platforms/mlu/wan_animate.json index 293a60537..3b2f2e91f 100755 --- a/configs/platforms/mlu/wan_animate.json +++ b/configs/platforms/mlu/wan_animate.json @@ -14,7 +14,7 @@ "enable_cfg": false, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm", "refert_num": 1, diff --git a/configs/platforms/mlu/wan_animate_dist.json b/configs/platforms/mlu/wan_animate_dist.json index 6a4c2ec3e..6bcf54de6 100755 --- a/configs/platforms/mlu/wan_animate_dist.json +++ b/configs/platforms/mlu/wan_animate_dist.json @@ -14,7 +14,7 @@ "enable_cfg": false, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm", "refert_num": 1, diff --git a/configs/platforms/mlu/wan_t2v.json b/configs/platforms/mlu/wan_t2v.json index 9f1bbfe65..2b4471063 100755 --- a/configs/platforms/mlu/wan_t2v.json +++ b/configs/platforms/mlu/wan_t2v.json @@ -12,7 +12,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "mlu_rms_norm" } diff --git a/configs/platforms/mlu/z_image_turbo_t2i.json b/configs/platforms/mlu/z_image_turbo_t2i.json index 769d1ea1d..6df820105 100644 --- a/configs/platforms/mlu/z_image_turbo_t2i.json +++ b/configs/platforms/mlu/z_image_turbo_t2i.json @@ -6,6 +6,6 @@ "enable_cfg": false, "sample_guide_scale": 0.0, "patch_size": 2, - "rope_type": "torch", + "rope_type": "torch_real_rope", "rms_norm_type": "mlu_rms_norm" } diff --git a/configs/platforms/mlu/z_image_turbo_t2i_dist.json b/configs/platforms/mlu/z_image_turbo_t2i_dist.json index f95fa499b..2eac08d67 100644 --- a/configs/platforms/mlu/z_image_turbo_t2i_dist.json +++ b/configs/platforms/mlu/z_image_turbo_t2i_dist.json @@ -6,7 +6,7 @@ "enable_cfg": false, "sample_guide_scale": 0.0, "patch_size": 2, - "rope_type": "torch", + "rope_type": "torch_real_rope", "rms_norm_type": "mlu_rms_norm", "parallel": { "seq_p_size": 2, diff --git a/configs/platforms/mthreads_musa/qwen_image_i2i_2511.json b/configs/platforms/mthreads_musa/qwen_image_i2i_2511.json index 43da91d1e..aa15899e7 100755 --- a/configs/platforms/mthreads_musa/qwen_image_i2i_2511.json +++ b/configs/platforms/mthreads_musa/qwen_image_i2i_2511.json @@ -9,7 +9,7 @@ "CONDITION_IMAGE_SIZE": 147456, "USE_IMAGE_ID_IN_PROMPT": true, "modulate_type": "torch", - "rope_type": "torch_naive", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/platforms/mthreads_musa/wan_t2v.json b/configs/platforms/mthreads_musa/wan_t2v.json index 74cd6c8c0..16c5ef13c 100755 --- a/configs/platforms/mthreads_musa/wan_t2v.json +++ b/configs/platforms/mthreads_musa/wan_t2v.json @@ -12,7 +12,7 @@ "enable_cfg": true, "cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch_naive", + "rope_type": "torch_real_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/qwen_image/qwen_image_i2i_layered.json b/configs/qwen_image/qwen_image_i2i_layered.json index a9f8d8520..fff6d7a84 100755 --- a/configs/qwen_image/qwen_image_i2i_layered.json +++ b/configs/qwen_image/qwen_image_i2i_layered.json @@ -13,5 +13,5 @@ "use_en_prompt": true, "use_additional_t_cond": true, "use_layer3d_rope": true, - "rope_type": "torch" + "rope_type": "torch_complex_rope" } diff --git a/configs/seko_talk/mlu/seko_talk_bf16.json b/configs/seko_talk/mlu/seko_talk_bf16.json index 4f05c1136..ed623881f 100644 --- a/configs/seko_talk/mlu/seko_talk_bf16.json +++ b/configs/seko_talk/mlu/seko_talk_bf16.json @@ -13,6 +13,6 @@ "enable_cfg": false, "cpu_offload": false, "use_31_block": true, - "rope_type": "torch", + "rope_type": "torch_complex_rope", "modulate_type": "torch" } diff --git a/configs/seko_talk/mlu/seko_talk_int8.json b/configs/seko_talk/mlu/seko_talk_int8.json index 9107a8381..01341d37c 100644 --- a/configs/seko_talk/mlu/seko_talk_int8.json +++ b/configs/seko_talk/mlu/seko_talk_int8.json @@ -23,7 +23,7 @@ "t5_quantized": true, "t5_quant_scheme": "int8-tmo", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/seko_talk/mlu/seko_talk_int8_dist.json b/configs/seko_talk/mlu/seko_talk_int8_dist.json index 7bba5d947..b32874772 100644 --- a/configs/seko_talk/mlu/seko_talk_int8_dist.json +++ b/configs/seko_talk/mlu/seko_talk_int8_dist.json @@ -23,7 +23,7 @@ "t5_quantized": true, "t5_quant_scheme": "int8-tmo", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "parallel": { diff --git a/configs/seko_talk/npu/seko_talk_01_bf16.json b/configs/seko_talk/npu/seko_talk_01_bf16.json index 4223dd10c..13ebc669b 100644 --- a/configs/seko_talk/npu/seko_talk_01_bf16.json +++ b/configs/seko_talk/npu/seko_talk_01_bf16.json @@ -17,7 +17,7 @@ "audio_adapter_cpu_offload": false, "clip_cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "rms_norm_type": "torch" } diff --git a/configs/seko_talk/npu/seko_talk_02_int8.json b/configs/seko_talk/npu/seko_talk_02_int8.json index 3c32b7d60..765eec4e1 100644 --- a/configs/seko_talk/npu/seko_talk_02_int8.json +++ b/configs/seko_talk/npu/seko_talk_02_int8.json @@ -20,7 +20,7 @@ "t5_quantized": true, "t5_quant_scheme": "int8-npu", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "cpu_offload": true, "clip_cpu_offload": false, diff --git a/configs/seko_talk/npu/seko_talk_03_int8_dist.json b/configs/seko_talk/npu/seko_talk_03_int8_dist.json index afeda92ef..b58b18b67 100644 --- a/configs/seko_talk/npu/seko_talk_03_int8_dist.json +++ b/configs/seko_talk/npu/seko_talk_03_int8_dist.json @@ -23,7 +23,7 @@ "t5_quantized": true, "t5_quant_scheme": "int8-npu", "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "cpu_offload": true, "audio_adapter_cpu_offload": true, diff --git a/configs/seko_talk/npu/seko_talk_04_bf16_dist.json b/configs/seko_talk/npu/seko_talk_04_bf16_dist.json index 9143ee63d..e72b17150 100644 --- a/configs/seko_talk/npu/seko_talk_04_bf16_dist.json +++ b/configs/seko_talk/npu/seko_talk_04_bf16_dist.json @@ -16,7 +16,7 @@ "audio_adapter_cpu_offload": false, "clip_cpu_offload": false, "modulate_type": "torch", - "rope_type": "torch", + "rope_type": "torch_complex_rope", "layer_norm_type": "torch", "rms_norm_type": "torch", "parallel": { diff --git a/configs/seko_talk/shot/stream/f2v.json b/configs/seko_talk/shot/stream/f2v.json index ce6558f7f..ee974f2be 100644 --- a/configs/seko_talk/shot/stream/f2v.json +++ b/configs/seko_talk/shot/stream/f2v.json @@ -11,7 +11,7 @@ "sample_shift": 5, "enable_cfg": false, "use_31_block": true, - "rope_type": "torch", + "rope_type": "torch_complex_rope", "target_video_length": 33, "prev_frame_length": 1, "f2v_process": true, diff --git a/configs/seko_talk/shot/stream/s2v.json b/configs/seko_talk/shot/stream/s2v.json index 627dbc34d..1c8fcfbbb 100644 --- a/configs/seko_talk/shot/stream/s2v.json +++ b/configs/seko_talk/shot/stream/s2v.json @@ -11,7 +11,7 @@ "sample_shift": 5, "enable_cfg": false, "use_31_block": true, - "rope_type": "torch", + "rope_type": "torch_complex_rope", "target_video_length": 33, "prev_frame_length": 5, "cpu_offload": true, diff --git a/configs/self_forcing/wan_t2v_sf.json b/configs/self_forcing/wan_t2v_sf.json index 1647b7412..dda312966 100755 --- a/configs/self_forcing/wan_t2v_sf.json +++ b/configs/self_forcing/wan_t2v_sf.json @@ -11,8 +11,8 @@ "sample_shift": 5.0, "enable_cfg": false, "cpu_offload": false, - "dit_original_ckpt": "/data/nvme4/gushiqiao/Self-Forcing/checkpoints/self_forcing_dmd.pt", - "causal_rope_type": "triton", + "dit_original_ckpt": "/Self-Forcing/checkpoints/self_forcing_dmd.pt", + "causal_rope_type": "wan_causal_rope", "ar_config": { "local_attn_size": -1, "num_frame_per_chunk": 3, diff --git a/configs/self_forcing/wan_t2v_sf_14b.json b/configs/self_forcing/wan_t2v_sf_14b.json index f070df42b..401cc0b11 100755 --- a/configs/self_forcing/wan_t2v_sf_14b.json +++ b/configs/self_forcing/wan_t2v_sf_14b.json @@ -18,5 +18,5 @@ "timesteps_index": [0, 179, 358, 679], "kv_offload": false }, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/self_forcing/wan_t2v_sf_fp8.json b/configs/self_forcing/wan_t2v_sf_fp8.json index 581547846..1bff16ea1 100755 --- a/configs/self_forcing/wan_t2v_sf_fp8.json +++ b/configs/self_forcing/wan_t2v_sf_fp8.json @@ -20,5 +20,5 @@ "dit_quantized": true, "dit_quantized_ckpt": "/path/to/self_forcing_fp8.safetensors", "dit_quant_scheme": "fp8-vllm", - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/self_forcing/wan_t2v_sf_kv_kiviquant.json b/configs/self_forcing/wan_t2v_sf_kv_kiviquant.json index 86ad7d605..b8581bf07 100755 --- a/configs/self_forcing/wan_t2v_sf_kv_kiviquant.json +++ b/configs/self_forcing/wan_t2v_sf_kv_kiviquant.json @@ -24,5 +24,5 @@ }, "kv_offload": false }, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/self_forcing/wan_t2v_sf_kv_kiviquant_14b.json b/configs/self_forcing/wan_t2v_sf_kv_kiviquant_14b.json index 68c2e8985..5dd1423d1 100755 --- a/configs/self_forcing/wan_t2v_sf_kv_kiviquant_14b.json +++ b/configs/self_forcing/wan_t2v_sf_kv_kiviquant_14b.json @@ -24,5 +24,5 @@ }, "kv_offload": false }, - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/self_forcing/wan_t2v_sf_nvfp4.json b/configs/self_forcing/wan_t2v_sf_nvfp4.json index ad9cda1cd..50242416e 100755 --- a/configs/self_forcing/wan_t2v_sf_nvfp4.json +++ b/configs/self_forcing/wan_t2v_sf_nvfp4.json @@ -20,5 +20,5 @@ "dit_quantized": true, "dit_quantized_ckpt": "/path/to/self_forcing_nvfp4.safetensors", "dit_quant_scheme": "nvfp4", - "causal_rope_type": "triton" + "causal_rope_type": "wan_causal_rope" } diff --git a/configs/wan22/wan_moe_i2v_audio.json b/configs/wan22/wan_moe_i2v_audio.json index 72e615137..7c59599eb 100755 --- a/configs/wan22/wan_moe_i2v_audio.json +++ b/configs/wan22/wan_moe_i2v_audio.json @@ -1,4 +1,5 @@ { + "causal_rope_type": "wan_causal_rope", "infer_steps": 6, "target_fps": 16, "video_duration": 16, diff --git a/lightx2v/common/ops/__init__.py b/lightx2v/common/ops/__init__.py index aab2983ef..b4f783ad9 100755 --- a/lightx2v/common/ops/__init__.py +++ b/lightx2v/common/ops/__init__.py @@ -3,4 +3,5 @@ from .embedding import * from .mm import * from .norm import * +from .rope import * from .tensor import * diff --git a/lightx2v/common/ops/rope/__init__.py b/lightx2v/common/ops/rope/__init__.py new file mode 100644 index 000000000..1cdeff55d --- /dev/null +++ b/lightx2v/common/ops/rope/__init__.py @@ -0,0 +1,4 @@ +from .chunked_rope import ChunkedRope +from .flashinfer_rope import FlashInferRope +from .template import RopeLayout, RopeTemplate +from .torch_rope import TorchComplexRope, TorchRealRope diff --git a/lightx2v/common/ops/rope/chunked_rope.py b/lightx2v/common/ops/rope/chunked_rope.py new file mode 100644 index 000000000..1943a49ff --- /dev/null +++ b/lightx2v/common/ops/rope/chunked_rope.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import torch + +from lightx2v.utils.registry_factory import ROPE_REGISTER + +from .template import RopeTemplate + + +def _slice_freqs(freqs, start: int, end: int): + if isinstance(freqs, tuple): + return tuple(item[start:end] for item in freqs) + return freqs[start:end] + + +@ROPE_REGISTER("chunked_rope") +class ChunkedRope(RopeTemplate): + def __init__(self, inner: RopeTemplate, chunk_size: int): + super().__init__(layout=getattr(inner, "layout", "interleaved"), compute_dtype=getattr(inner, "compute_dtype", torch.float32)) + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}.") + self.inner = inner + self.chunk_size = chunk_size + + def set_config(self, config=None): + super().set_config(config) + if hasattr(self.inner, "set_config"): + self.inner.set_config(config) + + def apply_single(self, x: torch.Tensor, freqs, **kwargs): + freq_len = freqs[0].shape[0] if isinstance(freqs, tuple) else freqs.shape[0] + seq_len = min(x.shape[0], freq_len) + output = torch.empty_like(x) + for start in range(0, seq_len, self.chunk_size): + end = min(start + self.chunk_size, seq_len) + output[start:end] = self.inner.apply_single(x[start:end], _slice_freqs(freqs, start, end), **kwargs) + if seq_len < x.shape[0]: + output[seq_len:] = x[seq_len:] + return output + + def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, **kwargs): + freq_len = freqs[0].shape[0] if isinstance(freqs, tuple) else freqs.shape[0] + seq_len = min(q.shape[0], k.shape[0], freq_len) + q_out, k_out = torch.empty_like(q), torch.empty_like(k) + for start in range(0, seq_len, self.chunk_size): + end = min(start + self.chunk_size, seq_len) + q_chunk, k_chunk = self.inner.apply(q[start:end], k[start:end], _slice_freqs(freqs, start, end), **kwargs) + q_out[start:end], k_out[start:end] = q_chunk, k_chunk + if seq_len < q.shape[0]: + q_out[seq_len:], k_out[seq_len:] = q[seq_len:], k[seq_len:] + return q_out, k_out diff --git a/lightx2v/common/ops/rope/flashinfer_rope.py b/lightx2v/common/ops/rope/flashinfer_rope.py new file mode 100644 index 000000000..48bf93d3d --- /dev/null +++ b/lightx2v/common/ops/rope/flashinfer_rope.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import torch + +from lightx2v.common.magi_custom_op_mode import use_magi_custom_ops +from lightx2v.utils.registry_factory import ROPE_REGISTER + +try: + from magi_compiler import magi_register_custom_op +except ImportError: + magi_register_custom_op = None + +from .template import RopeTemplate + +try: + from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace +except ImportError: + apply_rope_with_cos_sin_cache_inplace = None + + +@torch.library.custom_op( + "lightx2v::rope_flashinfer_", + mutates_args=("query", "key"), + device_types="cuda", + schema=("(Tensor positions, Tensor(a!) query, Tensor(b!) key, Tensor cos_sin_cache, int head_size, bool is_neox) -> ()"), +) +def rope_flashinfer_( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + cos_sin_cache: torch.Tensor, + head_size: int, + is_neox: bool, +) -> None: + apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=head_size, + cos_sin_cache=cos_sin_cache, + is_neox=is_neox, + ) + + +@rope_flashinfer_.register_fake +def rope_flashinfer_fake( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + cos_sin_cache: torch.Tensor, + head_size: int, + is_neox: bool, +) -> None: + return None + + +@ROPE_REGISTER("flashinfer_rope") +class FlashInferRope(RopeTemplate): + @staticmethod + def is_available(): + return apply_rope_with_cos_sin_cache_inplace is not None + + def prepare_freqs(self, freqs, rotary_dim: int | None = None): + if torch.is_tensor(freqs) and torch.is_complex(freqs): + while freqs.ndim > 2 and freqs.shape[-2] == 1: + freqs = freqs.squeeze(-2) + return torch.cat([freqs.real, freqs.imag], dim=-1).float().contiguous() + + if isinstance(freqs, tuple): + cos, sin = freqs[:2] + if rotary_dim is None: + raise ValueError("rotary_dim is required for tuple RoPE frequencies.") + if cos.shape[-1] == rotary_dim: + if self.layout == "interleaved": + cos, sin = cos[..., ::2], sin[..., ::2] + else: + cos, sin = cos[..., : rotary_dim // 2], sin[..., : rotary_dim // 2] + 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]}.") + while cos.ndim > 2 and cos.shape[0] == 1: + cos, sin = cos.squeeze(0), sin.squeeze(0) + return torch.cat([cos, sin], dim=-1).float().contiguous() + + if torch.is_tensor(freqs): + return freqs.float().contiguous() + raise TypeError(f"Unsupported RoPE frequency type: {type(freqs)!r}") + + def _apply_eager(self, q: torch.Tensor, k: torch.Tensor, freqs: torch.Tensor, positions: torch.Tensor | None = None): + if apply_rope_with_cos_sin_cache_inplace is None: + raise ImportError("flashinfer is required for FlashInferRope.") + if q.ndim != 3 or k.ndim != 3: + raise ValueError(f"FlashInferRope expects [L, H, D] tensors, got q={q.shape}, k={k.shape}.") + length, q_heads, head_dim = q.shape + k_heads = k.shape[1] + if positions is None: + positions = torch.arange(length, device=q.device, dtype=torch.long) + query = q.reshape(length, q_heads * head_dim).contiguous() + key = k.reshape(length, k_heads * head_dim).contiguous() + is_neox = self.layout == "split_half" + if torch.compiler.is_compiling(): + rope_flashinfer_(positions, query, key, freqs, head_dim, is_neox) + else: + apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=head_dim, + cos_sin_cache=freqs, + is_neox=is_neox, + ) + return query.view_as(q), key.view_as(k) + + def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, positions: torch.Tensor | None = None, **kwargs): + squeeze_batch = q.ndim == 4 and q.shape[0] == 1 and k.ndim == 4 and k.shape[0] == 1 + if squeeze_batch: + q, k = q[0], k[0] + freqs = self.prepare_freqs(freqs, rotary_dim=q.shape[-1]) + if positions is None and self.layout == "interleaved" and use_magi_custom_ops() and magi_register_custom_op is not None and apply_rope_with_cos_sin_cache_inplace is not None: + q_out, k_out = torch.ops.lightx2v.rope_flashinfer(q, k, freqs) + else: + q_out, k_out = self._apply_eager(q, k, freqs, positions=positions) + if squeeze_batch: + return q_out.unsqueeze(0), k_out.unsqueeze(0) + return q_out, k_out + + def apply_single(self, x: torch.Tensor, freqs, **kwargs): + output, _ = self.apply(x, x.clone(), freqs, **kwargs) + return output + + +def _rope_meta(q, k, freqs): + return torch.empty_like(q), torch.empty_like(k) + + +if magi_register_custom_op is not None and apply_rope_with_cos_sin_cache_inplace is not None: + + @magi_register_custom_op( + "lightx2v::rope_flashinfer", + infer_output_meta_fn=_rope_meta, + is_subgraph_boundary=True, + ) + def _rope_flashinfer_custom_op(q: torch.Tensor, k: torch.Tensor, freqs: torch.Tensor): + return FlashInferRope(layout="interleaved")._apply_eager(q.clone(), k.clone(), freqs) diff --git a/lightx2v/common/ops/rope/template.py b/lightx2v/common/ops/rope/template.py new file mode 100644 index 000000000..8a785d1ee --- /dev/null +++ b/lightx2v/common/ops/rope/template.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from typing import Literal + +import torch + +RopeLayout = Literal["interleaved", "split_half"] + + +class RopeTemplate: + def __init__(self, layout: RopeLayout = "interleaved", compute_dtype: torch.dtype = torch.float32): + if layout not in {"interleaved", "split_half"}: + raise ValueError(f"Unsupported RoPE layout: {layout}") + self.layout = layout + self.compute_dtype = compute_dtype + self.config = {} + + def set_config(self, config=None): + if config is not None: + self.config = config + + # RoPE has no checkpoint parameters, but it participates in the same + # WeightModule tree as linear, attention, and norm implementations. + def load(self, weight_dict): + pass + + def to_cpu(self, non_blocking=False): + pass + + def to_cuda(self, non_blocking=False): + pass + + def state_dict(self, destination=None): + return {} if destination is None else destination + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + return {} if destination is None else destination + + def load_state_dict_from_disk(self, block_index, adapter_block_index=None): + pass + + def named_parameters(self, prefix=""): + return iter(()) + + def prepare_freqs(self, freqs): + return freqs + + def apply_single(self, x: torch.Tensor, freqs, **kwargs) -> torch.Tensor: + raise NotImplementedError + + def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, **kwargs): + return self.apply_single(q, freqs, **kwargs), self.apply_single(k, freqs, **kwargs) + + +def broadcast_freqs(freqs: torch.Tensor, target: torch.Tensor, unsqueeze_dim: int = -2): + while freqs.ndim < target.ndim: + freqs = freqs.unsqueeze(unsqueeze_dim) + return freqs diff --git a/lightx2v/common/ops/rope/torch_rope.py b/lightx2v/common/ops/rope/torch_rope.py new file mode 100644 index 000000000..bc016261f --- /dev/null +++ b/lightx2v/common/ops/rope/torch_rope.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import torch + +from lightx2v.common.magi_custom_op_mode import use_magi_custom_ops +from lightx2v.utils.registry_factory import ROPE_REGISTER + +try: + from magi_compiler import magi_register_custom_op +except ImportError: + magi_register_custom_op = None + +from .template import RopeTemplate, broadcast_freqs + + +@ROPE_REGISTER("torch_complex_rope") +class TorchComplexRope(RopeTemplate): + def __init__(self, layout="interleaved", compute_dtype: torch.dtype = torch.float32): + if layout != "interleaved": + raise ValueError("TorchComplexRope only supports interleaved layout.") + super().__init__(layout=layout, compute_dtype=compute_dtype) + + def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, **kwargs): + if q.ndim == 3 and k.ndim == 3 and self.compute_dtype == torch.float32 and torch.is_complex(freqs) and use_magi_custom_ops() and magi_register_custom_op is not None and not kwargs: + return torch.ops.lightx2v.rope_torch_complex(q, k, freqs) + return super().apply(q, k, freqs, **kwargs) + + def apply_single(self, x: torch.Tensor, freqs: torch.Tensor, rotary_dim: int | None = None, unsqueeze_dim: int = -2, **kwargs): + if not torch.is_complex(freqs): + raise TypeError("TorchComplexRope expects a complex frequency tensor.") + rotary_dim = rotary_dim or x.shape[-1] + if rotary_dim % 2: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}.") + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x_complex = torch.view_as_complex(x_rot.to(self.compute_dtype).reshape(*x_rot.shape[:-1], -1, 2).contiguous()) + freqs = broadcast_freqs(freqs, x_complex, unsqueeze_dim) + output = torch.view_as_real(x_complex * freqs).flatten(-2).to(x.dtype) + return torch.cat((output, x_pass), dim=-1) if x_pass.shape[-1] else output + + +@ROPE_REGISTER("torch_real_rope") +class TorchRealRope(RopeTemplate): + def _cos_sin(self, freqs, rotary_dim: int): + if torch.is_tensor(freqs) and torch.is_complex(freqs): + return freqs.real, freqs.imag, True + if isinstance(freqs, tuple): + cos, sin = freqs + return cos, sin, cos.shape[-1] == rotary_dim // 2 + if torch.is_tensor(freqs): + if freqs.shape[-1] != rotary_dim: + raise ValueError(f"Concatenated cos-sin cache must have last dim {rotary_dim}, got {freqs.shape[-1]}.") + return freqs[..., : rotary_dim // 2], freqs[..., rotary_dim // 2 :], True + raise TypeError(f"Unsupported RoPE frequency type: {type(freqs)!r}") + + @staticmethod + def _rotate_interleaved(x: torch.Tensor): + pairs = x.reshape(*x.shape[:-1], -1, 2) + first, second = pairs.unbind(dim=-1) + return torch.stack((-second, first), dim=-1).flatten(-2) + + @staticmethod + def _rotate_split_half(x: torch.Tensor): + first, second = x.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + def apply_single(self, x: torch.Tensor, freqs, rotary_dim: int | None = None, unsqueeze_dim: int = -2, **kwargs): + rotary_dim = rotary_dim or x.shape[-1] + if rotary_dim % 2: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}.") + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + cos, sin, pairwise = self._cos_sin(freqs, rotary_dim) + x_float = x_rot.to(self.compute_dtype) + if pairwise: + if self.layout == "interleaved": + first, second = x_float[..., 0::2], x_float[..., 1::2] + else: + first, second = x_float.chunk(2, dim=-1) + cos = broadcast_freqs(cos.to(self.compute_dtype), first, unsqueeze_dim) + sin = broadcast_freqs(sin.to(self.compute_dtype), first, unsqueeze_dim) + first_out = first * cos - second * sin + second_out = first * sin + second * cos + if self.layout == "interleaved": + output = torch.empty_like(x_float) + output[..., 0::2] = first_out + output[..., 1::2] = second_out + else: + output = torch.cat((first_out, second_out), dim=-1) + else: + cos = broadcast_freqs(cos.to(self.compute_dtype), x_float, unsqueeze_dim) + sin = broadcast_freqs(sin.to(self.compute_dtype), x_float, unsqueeze_dim) + rotate = self._rotate_interleaved if self.layout == "interleaved" else self._rotate_split_half + output = x_float * cos + rotate(x_float) * sin + output = output.to(x.dtype) + return torch.cat((output, x_pass), dim=-1) if x_pass.shape[-1] else output + + +def _rope_meta(q, k, freqs): + return torch.empty_like(q), torch.empty_like(k) + + +if magi_register_custom_op is not None: + + @magi_register_custom_op( + "lightx2v::rope_torch_complex", + infer_output_meta_fn=_rope_meta, + is_subgraph_boundary=True, + ) + def _rope_torch_complex_custom_op(q: torch.Tensor, k: torch.Tensor, freqs: torch.Tensor): + module = TorchComplexRope() + return module.apply_single(q, freqs), module.apply_single(k, freqs) diff --git a/lightx2v/disagg/utils.py b/lightx2v/disagg/utils.py index 0d5a40da3..8aab7a2c3 100644 --- a/lightx2v/disagg/utils.py +++ b/lightx2v/disagg/utils.py @@ -36,7 +36,7 @@ def set_config( model_cls, config_path=None, attn_mode="flash_attn2", - rope_type="torch", + rope_type="torch_complex_rope", infer_steps=50, target_video_length=81, target_height=480, diff --git a/lightx2v/models/networks/bagel/infer/transformer_infer.py b/lightx2v/models/networks/bagel/infer/transformer_infer.py index 3730ce0e7..10dcb7808 100644 --- a/lightx2v/models/networks/bagel/infer/transformer_infer.py +++ b/lightx2v/models/networks/bagel/infer/transformer_infer.py @@ -16,42 +16,6 @@ from lightx2v_platform.base.global_var import AI_DEVICE -# Copied from transformers.models.llama.modeling_llama.rotate_half -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb -def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`, *optional*): - Deprecated and unused. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - class BagelTransformerInfer(BaseTransformerInfer): def __init__(self, config, llm_config): if flash_attn_varlen_func is None: @@ -139,7 +103,7 @@ def self_attn( packed_key_states[packed_vae_token_indexes] = weights.k_norm_moe_gen.apply(packed_key_states[packed_vae_token_indexes]) packed_cos, packed_sin = packed_query_position_embeddings - packed_query_states, packed_key_states = apply_rotary_pos_emb(packed_query_states, packed_key_states, packed_cos, packed_sin, unsqueeze_dim=1) + packed_query_states, packed_key_states = weights.rope.apply(packed_query_states, packed_key_states, (packed_cos, packed_sin), unsqueeze_dim=1) packed_query_states = packed_query_states.to(torch.bfloat16) packed_key_states = packed_key_states.to(torch.bfloat16) diff --git a/lightx2v/models/networks/bagel/weights/transformer_weights.py b/lightx2v/models/networks/bagel/weights/transformer_weights.py index 665388fdb..8d31eb65e 100644 --- a/lightx2v/models/networks/bagel/weights/transformer_weights.py +++ b/lightx2v/models/networks/bagel/weights/transformer_weights.py @@ -1,7 +1,10 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ( MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, ) @@ -135,6 +138,10 @@ def __init__( self.mm_type = mm_type self.task = task self.config = config + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_real_rope")](layout="split_half", compute_dtype=torch.float32), + ) # q self.add_module( "q_proj", diff --git a/lightx2v/models/networks/cosmos3/infer/transformer_infer.py b/lightx2v/models/networks/cosmos3/infer/transformer_infer.py index c4067ccc9..ab69a32d9 100644 --- a/lightx2v/models/networks/cosmos3/infer/transformer_infer.py +++ b/lightx2v/models/networks/cosmos3/infer/transformer_infer.py @@ -4,7 +4,7 @@ from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.cosmos3.infer.module_io import Cosmos3TransformerInferModuleOutput -from lightx2v.models.networks.cosmos3.infer.utils import apply_cosmos3_rotary, build_rotary_embeddings +from lightx2v.models.networks.cosmos3.infer.utils import build_rotary_embeddings class Cosmos3TransformerInfer(BaseTransformerInfer): @@ -18,7 +18,6 @@ def __init__(self, config): self.rope_theta = config.get("rope_theta", 5000000) rope_scaling = config.get("rope_scaling", None) self.rope_axes_dim = tuple(rope_scaling.get("mrope_section", [24, 20, 20]) if rope_scaling else [24, 20, 20]) - self.rope_type = config.get("rope_type", "triton") if self.config.get("seq_parallel", False): self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") else: @@ -73,8 +72,8 @@ def _infer_attention(self, block_weights, und_seq, gen_seq, rotary_emb): k_gen = block_weights.norm_added_k.apply(k_gen) cos_und, sin_und, cos_gen, sin_gen = rotary_emb - q_und, k_und = apply_cosmos3_rotary(q_und, k_und, cos_und, sin_und, rope_type=self.rope_type) - q_gen, k_gen = apply_cosmos3_rotary(q_gen, k_gen, cos_gen, sin_gen, rope_type=self.rope_type) + q_und, k_und = block_weights.rope.apply(q_und, k_und, (cos_und, sin_und)) + q_gen, k_gen = block_weights.rope.apply(q_gen, k_gen, (cos_gen, sin_gen)) k_und_full, v_und_full = self._repeat_kv_for_gqa(k_und, v_und) causal_out = self._infer_attn(block_weights.causal_self_attn, q_und, k_und_full, v_und_full, causal=True) diff --git a/lightx2v/models/networks/cosmos3/infer/utils.py b/lightx2v/models/networks/cosmos3/infer/utils.py index 0e2e8b923..2e29f02af 100644 --- a/lightx2v/models/networks/cosmos3/infer/utils.py +++ b/lightx2v/models/networks/cosmos3/infer/utils.py @@ -3,6 +3,9 @@ import torch import torch.nn.functional as F +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 @@ -10,15 +13,24 @@ triton = None tl = None -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None +@ROPE_REGISTER("cosmos3_rope") +class Cosmos3Rope(RopeTemplate): + def __init__(self, layout="split_half", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "split_half": + raise ValueError("Cosmos3Rope only supports split_half layout.") + self.torch_rope = TorchRealRope(layout=layout, compute_dtype=compute_dtype) -def rotate_half(x: torch.Tensor) -> torch.Tensor: - half = x.shape[-1] // 2 - return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + def apply(self, query, key, freqs, **kwargs): + cos, sin = freqs + if query.is_cuda and triton is not None: + return apply_split_half_rotary_triton(query, cos, sin), apply_split_half_rotary_triton(key, cos, sin) + return self.torch_rope.apply(query, key, freqs, **kwargs) + + def apply_single(self, x, freqs, **kwargs): + output, _ = self.apply(x, x.clone(), freqs, **kwargs) + return output if triton is not None: @@ -89,32 +101,6 @@ def apply_split_half_rotary_triton(x: torch.Tensor, cos: torch.Tensor, sin: torc return output -def apply_cosmos3_rotary(query: torch.Tensor, key: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, rope_type: str = "triton"): - if query.is_cuda and rope_type == "flashinfer" and apply_rope_with_cos_sin_cache_inplace is not None: - seq_len, q_heads, head_dim = query.shape - kv_heads = key.shape[1] - positions = torch.arange(seq_len, device=query.device, dtype=torch.long) - cos_sin_cache = torch.cat([cos[:, : head_dim // 2].float(), sin[:, : head_dim // 2].float()], dim=-1).contiguous() - query_flat = query.reshape(seq_len, q_heads * head_dim).contiguous() - key_flat = key.reshape(seq_len, kv_heads * head_dim).contiguous() - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query_flat, - key=key_flat, - head_size=head_dim, - cos_sin_cache=cos_sin_cache, - is_neox=True, - ) - return query_flat.view(seq_len, q_heads, head_dim), key_flat.view(seq_len, kv_heads, head_dim) - - if query.is_cuda and rope_type in {"flashinfer", "triton"} and triton is not None: - return apply_split_half_rotary_triton(query, cos, sin), apply_split_half_rotary_triton(key, cos, sin) - - cos = cos.unsqueeze(1) - sin = sin.unsqueeze(1) - return query * cos + rotate_half(query) * sin, key * cos + rotate_half(key) * sin - - def get_timestep_embedding( timesteps: torch.Tensor, embedding_dim: int = 256, diff --git a/lightx2v/models/networks/cosmos3/weights/transformer_weights.py b/lightx2v/models/networks/cosmos3/weights/transformer_weights.py index 9fd1e69b3..e4b5694c6 100644 --- a/lightx2v/models/networks/cosmos3/weights/transformer_weights.py +++ b/lightx2v/models/networks/cosmos3/weights/transformer_weights.py @@ -1,5 +1,8 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.models.networks.cosmos3.infer.utils import Cosmos3Rope # noqa: F401 +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER class Cosmos3TransformerWeights(WeightModule): @@ -187,6 +190,10 @@ def __init__( lora_path=None, ): super().__init__() + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "cosmos3_rope")](layout="split_half", compute_dtype=torch.float32), + ) lora_prefix = "layers" attn_type = config.get("self_attn_type", "torch_sdpa") causal_attn_type = config.get("causal_self_attn_type", "torch_sdpa") diff --git a/lightx2v/models/networks/ernie_image/infer/transformer_infer.py b/lightx2v/models/networks/ernie_image/infer/transformer_infer.py index 036e5ff57..3dd8e2223 100644 --- a/lightx2v/models/networks/ernie_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/ernie_image/infer/transformer_infer.py @@ -1,4 +1,3 @@ -import torch import torch.nn.functional as F @@ -11,16 +10,6 @@ def __init__(self, config): def set_scheduler(self, scheduler): self.scheduler = scheduler - @staticmethod - def _apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: - rot_dim = freqs_cis.shape[-1] - x, x_pass = x_in[..., :rot_dim], x_in[..., rot_dim:] - cos_ = torch.cos(freqs_cis).to(x.dtype) - sin_ = torch.sin(freqs_cis).to(x.dtype) - x1, x2 = x.chunk(2, dim=-1) - x_rotated = torch.cat((-x2, x1), dim=-1) - return torch.cat((x * cos_ + x_rotated * sin_, x_pass), dim=-1) - def _attention(self, weights, hidden_states, rotary_pos_emb): query = weights.to_q.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) key = weights.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) @@ -28,8 +17,12 @@ def _attention(self, weights, hidden_states, rotary_pos_emb): query = weights.norm_q.apply(query) key = weights.norm_k.apply(key) - query = self._apply_rotary_emb(query, rotary_pos_emb) - key = self._apply_rotary_emb(key, rotary_pos_emb) + query, key = weights.rope.apply( + query, + key, + (rotary_pos_emb.cos(), rotary_pos_emb.sin()), + rotary_dim=rotary_pos_emb.shape[-1], + ) hidden_states = weights.attn.apply( query, diff --git a/lightx2v/models/networks/ernie_image/weights/transformer_weights.py b/lightx2v/models/networks/ernie_image/weights/transformer_weights.py index d86757332..3564b468c 100644 --- a/lightx2v/models/networks/ernie_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/ernie_image/weights/transformer_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER class ErnieImageTransformerWeights(WeightModule): @@ -40,6 +42,10 @@ def __init__( lora_path=None, ): super().__init__() + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_real_rope")](layout="split_half", compute_dtype=torch.float32), + ) prefix = f"layers.{block_index}" eps = config.get("eps", 1e-6) self.add_module( diff --git a/lightx2v/models/networks/flux2/infer/pre_infer.py b/lightx2v/models/networks/flux2/infer/pre_infer.py index e4fe3b9ef..5bcbd41e3 100644 --- a/lightx2v/models/networks/flux2/infer/pre_infer.py +++ b/lightx2v/models/networks/flux2/infer/pre_infer.py @@ -60,11 +60,6 @@ def infer(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img freqs_cos = torch.cat([text_rope[0], image_rope[0]], dim=0) freqs_sin = torch.cat([text_rope[1], image_rope[1]], dim=0) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half = freqs_cos[:, ::2].contiguous() # [L, D/2] - sin_half = freqs_sin[:, ::2].contiguous() # [L, D/2] - image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1) # [L, D] - else: image_rotary_emb = (freqs_cos, freqs_sin) return Flux2PreInferModuleOutput( @@ -118,11 +113,6 @@ def infer(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img freqs_cos = torch.cat([text_rope[0], image_rope[0]], dim=0) freqs_sin = torch.cat([text_rope[1], image_rope[1]], dim=0) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half = freqs_cos[:, ::2].contiguous() - sin_half = freqs_sin[:, ::2].contiguous() - image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1) - else: image_rotary_emb = (freqs_cos, freqs_sin) return Flux2PreInferModuleOutput( diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py index b4d120990..e65a4f398 100644 --- a/lightx2v/models/networks/flux2/infer/transformer_infer.py +++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py @@ -4,8 +4,6 @@ from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer -from .utils import apply_rope_with_flashinfer, apply_rope_with_torch - class Flux2TransformerInfer(BaseTransformerInfer): def __init__(self, config): @@ -26,13 +24,6 @@ def __init__(self, config): self.seq_p_fp4_comm = False self.enable_head_parallel = False - rope_funcs = { - "flashinfer": apply_rope_with_flashinfer, - "torch": apply_rope_with_torch, - } - rope_type = config.get("rope_type", "flashinfer") - self.apply_rope_func = rope_funcs.get(rope_type, apply_rope_with_torch) - def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -93,7 +84,7 @@ def infer_double_stream_block( key = torch.cat([txt_key, img_key], dim=0) value = torch.cat([txt_value, img_value], dim=0) - query, key = self.apply_rope_func(query, key, image_rotary_emb) + query, key = block_weights.rope.apply(query, key, image_rotary_emb) total_len = query.shape[0] cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32) @@ -194,7 +185,7 @@ def infer_single_stream_block( query = block_weights.norm_q.apply(query) key = block_weights.norm_k.apply(key) - query, key = self.apply_rope_func(query, key, image_rotary_emb) + query, key = block_weights.rope.apply(query, key, image_rotary_emb) total_len = query.shape[0] cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32) diff --git a/lightx2v/models/networks/flux2/infer/utils.py b/lightx2v/models/networks/flux2/infer/utils.py index f36bfff26..4af3d510f 100644 --- a/lightx2v/models/networks/flux2/infer/utils.py +++ b/lightx2v/models/networks/flux2/infer/utils.py @@ -1,74 +1 @@ -from typing import Tuple - -import torch - -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - - -def apply_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - """Apply rotary position embedding using flashinfer. - - Args: - xq: Query tensor [L, H, D] - xk: Key tensor [L, H, D] - cos_sin_cache: Cosine and sine cache [L, D] where first half is cos, second half is sin - - Returns: - Tuple of (xq, xk) with rotary embedding applied - """ - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous() - key = xk.reshape(L, H * D).contiguous() - - positions = torch.arange(L, device=xq.device, dtype=torch.long) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out - - -def apply_rope_with_torch( - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cis: Tuple[torch.Tensor, torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - """Apply rotary position embedding using PyTorch. - Args: - xq: Query tensor [L, H, D] - xk: Key tensor [L, H, D] - freqs_cis: Tuple of (cos, sin) each [L, D] - - Returns: - Tuple of (xq, xk) with rotary embedding applied - """ - cos, sin = freqs_cis # [L, D] - - cos = cos[:, None, :] - sin = sin[:, None, :] - - def _apply_rope(x, cos, sin): - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [L, H, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - return out - - xq_out = _apply_rope(xq, cos, sin) - xk_out = _apply_rope(xk, cos, sin) - return xq_out, xk_out +"""RoPE is selected by Flux2 attention weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py index da3ed9141..25a473113 100644 --- a/lightx2v/models/networks/flux2/weights/transformer_weights.py +++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER class Flux2DoubleBlockWeights(WeightModule): @@ -13,6 +15,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.mm_type = config.get("dit_quant_scheme", "Default") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) p = f"transformer_blocks.{self.block_idx}" @@ -191,6 +197,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.mm_type = config.get("dit_quant_scheme", "Default") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) p = f"single_transformer_blocks.{self.block_idx}" diff --git a/lightx2v/models/networks/hidream_o1_image/infer/rope.py b/lightx2v/models/networks/hidream_o1_image/infer/rope.py index 367fd57fc..73d0b9748 100644 --- a/lightx2v/models/networks/hidream_o1_image/infer/rope.py +++ b/lightx2v/models/networks/hidream_o1_image/infer/rope.py @@ -1,45 +1 @@ -import torch - -from lightx2v.models.networks.hidream_o1_image.qwen3_vl import apply_rotary_pos_emb - -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - - -def apply_hidream_rope_with_torch(q, k, rope_cos_sin): - cos, sin = rope_cos_sin[:2] - q_rope = q.transpose(1, 2) - k_rope = k.transpose(1, 2) - q_rope, k_rope = apply_rotary_pos_emb(q_rope, k_rope, cos, sin) - return q_rope.transpose(1, 2).contiguous(), k_rope.transpose(1, 2).contiguous() - - -def apply_hidream_rope_with_flashinfer(q, k, rope_cos_sin): - if apply_rope_with_cos_sin_cache_inplace is None: - raise ImportError("flashinfer is required when hidream_o1_image rope_type='flashinfer'.") - if q.shape[0] != 1: - raise NotImplementedError("HiDream flashinfer RoPE currently expects batch=1 CFG forwards.") - - cos, sin = rope_cos_sin[:2] - seq_len, q_heads, head_dim = q.shape[1], q.shape[2], q.shape[3] - kv_heads = k.shape[2] - rotary_dim = head_dim // 2 - cos_sin_cache = torch.cat([cos[0, :, :rotary_dim].float(), sin[0, :, :rotary_dim].float()], dim=-1).contiguous() - if len(rope_cos_sin) > 2: - positions = rope_cos_sin[2] - else: - positions = torch.arange(seq_len, device=q.device, dtype=torch.long) - - query = q[0].reshape(seq_len, q_heads * head_dim).contiguous() - key = k[0].reshape(seq_len, kv_heads * head_dim).contiguous() - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=head_dim, - cos_sin_cache=cos_sin_cache, - is_neox=True, - ) - return query.view(1, seq_len, q_heads, head_dim), key.view(1, seq_len, kv_heads, head_dim) +"""HiDream RoPE is selected by decoder-block weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/hidream_o1_image/infer/transformer_infer.py b/lightx2v/models/networks/hidream_o1_image/infer/transformer_infer.py index d8a517300..b334ba088 100644 --- a/lightx2v/models/networks/hidream_o1_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/hidream_o1_image/infer/transformer_infer.py @@ -2,20 +2,11 @@ import torch.distributed as dist from lightx2v.models.networks.hidream_o1_image.infer.module_io import HidreamTransformerInferOutput -from lightx2v.models.networks.hidream_o1_image.infer.rope import apply_hidream_rope_with_flashinfer, apply_hidream_rope_with_torch class HidreamO1ImageTransformerInfer: def __init__(self, config): self.config = config - rope_type = config.get("rope_type", "flashinfer") - rope_funcs = { - "flashinfer": apply_hidream_rope_with_flashinfer, - "torch": apply_hidream_rope_with_torch, - } - if rope_type not in rope_funcs: - raise ValueError(f"Unsupported HiDream rope_type={rope_type!r}. Supported values: {sorted(rope_funcs)}") - self.apply_rope_func = rope_funcs[rope_type] if self.config["seq_parallel"]: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") self.seq_p_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False) @@ -27,6 +18,10 @@ def __init__(self, config): self.seq_p_fp4_comm = False self.enable_head_parallel = False + def _apply_rope(self, rope, q, k, rope_cos_sin): + positions = rope_cos_sin[2] if len(rope_cos_sin) > 2 else None + return rope.apply(q, k, tuple(rope_cos_sin[:2]), positions=positions, unsqueeze_dim=-2) + def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -173,7 +168,7 @@ def _project_qkv(self, weights, hidden_states, rope_cos_sin): v = weights.v_proj.apply(flat_hidden).reshape(batch, seq_len, weights.kv_heads, weights.head_dim) q = weights.q_norm.apply(q) k = weights.k_norm.apply(k) - q, k = self.apply_rope_func(q, k, rope_cos_sin) + q, k = self._apply_rope(weights.rope, q, k, rope_cos_sin) return q, k, v def _two_pass_attn(self, weights, q, k, v, idx_ar): diff --git a/lightx2v/models/networks/hidream_o1_image/infer/vision_infer.py b/lightx2v/models/networks/hidream_o1_image/infer/vision_infer.py index 47d87207d..b62149819 100644 --- a/lightx2v/models/networks/hidream_o1_image/infer/vision_infer.py +++ b/lightx2v/models/networks/hidream_o1_image/infer/vision_infer.py @@ -1,23 +1,9 @@ import torch from transformers.activations import ACT2FN +from lightx2v.common.ops.rope import TorchRealRope -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb_vision(q, k, cos, sin): - q_dtype = q.dtype - k_dtype = k.dtype - q = q.float() - k = k.float() - cos = cos.unsqueeze(-2).float() - sin = sin.unsqueeze(-2).float() - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed.to(q_dtype), k_embed.to(k_dtype) +_HIDREAM_VISION_ROPE = TorchRealRope(layout="split_half") class HidreamO1ImageVisionInfer: @@ -145,7 +131,7 @@ def _infer_attn(self, weights, hidden_states, cu_seqlens, position_embeddings): qkv = weights.qkv.apply(hidden_states).reshape(seq_len, 3, weights.num_heads, weights.head_dim).permute(1, 0, 2, 3) q, k, v = qkv.unbind(0) cos, sin = position_embeddings - q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + q, k = _HIDREAM_VISION_ROPE.apply(q, k, (cos, sin)) lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() outs = [] start = 0 diff --git a/lightx2v/models/networks/hidream_o1_image/qwen3_vl.py b/lightx2v/models/networks/hidream_o1_image/qwen3_vl.py index 1e24425fe..fc6cc482f 100644 --- a/lightx2v/models/networks/hidream_o1_image/qwen3_vl.py +++ b/lightx2v/models/networks/hidream_o1_image/qwen3_vl.py @@ -5,13 +5,10 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig -USE_BF16_ROPE = os.environ.get("USE_BF16_ROPE", "0") - +from lightx2v.common.ops.rope import TorchRealRope -def rotate_half(x): - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) +USE_BF16_ROPE = os.environ.get("USE_BF16_ROPE", "0") +_QWEN3_VL_ROPE = TorchRealRope(layout="split_half") class Qwen3VLTextRotaryEmbedding(nn.Module): @@ -74,8 +71,5 @@ def forward(self, x, position_ids): def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed + del position_ids + return _QWEN3_VL_ROPE.apply(q, k, (cos, sin), unsqueeze_dim=unsqueeze_dim) diff --git a/lightx2v/models/networks/hidream_o1_image/weights/transformer_weights.py b/lightx2v/models/networks/hidream_o1_image/weights/transformer_weights.py index 60d82d4c9..30b156427 100644 --- a/lightx2v/models/networks/hidream_o1_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/hidream_o1_image/weights/transformer_weights.py @@ -1,7 +1,8 @@ +import torch from transformers.activations import ACT2FN from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -46,6 +47,10 @@ def __init__(self, block_index, config, mm_type, rms_norm_type, attn_type): super().__init__() self.block_index = block_index self.config = config + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="split_half", compute_dtype=torch.float32), + ) self.heads = None self.kv_heads = None self.head_dim = None diff --git a/lightx2v/models/networks/hunyuan_video/infer/posemb_layers.py b/lightx2v/models/networks/hunyuan_video/infer/posemb_layers.py index 948c0a6e0..69ea03822 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/posemb_layers.py +++ b/lightx2v/models/networks/hunyuan_video/infer/posemb_layers.py @@ -65,109 +65,6 @@ def get_meshgrid_nd(start, *args, dim=2): # https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80 -def reshape_for_broadcast( - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], - x: torch.Tensor, -): - """ - Reshape frequency tensor for broadcasting it with another tensor. - - This function reshapes the frequency tensor to have the same shape as the target tensor 'x' - for the purpose of broadcasting the frequency tensor during element-wise operations. - - Notes: - When using FlashMHAModified, head_first should be False. - When using Attention, head_first should be True. - - Args: - freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped. - x (torch.Tensor): Target tensor for broadcasting compatibility. - head_first (bool): head dimension first (except batch dim) or not. - - Returns: - torch.Tensor: Reshaped frequency tensor. - - Raises: - AssertionError: If the frequency tensor doesn't match the expected shape. - AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions. - """ - ndim = x.ndim - shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] - return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape) - - -def rotate_half(x): - x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - return torch.stack([-x_imag, x_real], dim=-1).flatten(3) - - -def apply_rotary_emb( - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Apply rotary embeddings to input tensors using the given frequency tensor. - - This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided - frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor - is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are - returned as real tensors. - - Args: - xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] - xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] - freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. - head_first (bool): head dimension first (except batch dim) or not. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. - - """ - cos, sin = reshape_for_broadcast(freqs_cis, xq) # [S, D] - # real * cos - imag * sin - # imag * cos + real * sin - xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq) - xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk) - return xq_out, xk_out - - -def rotate_half_force_bf16(x): - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - return torch.stack([-x_imag, x_real], dim=-1).flatten(3) - - -def apply_rotary_emb_force_bf16( - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Apply rotary embeddings to input tensors using the given frequency tensor. - - This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided - frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor - is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are - returned as real tensors. - - Args: - xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] - xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] - freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. - head_first (bool): head dimension first (except batch dim) or not. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. - - """ - cos, sin = reshape_for_broadcast(freqs_cis, xq) # [S, D] - # real * cos - imag * sin - # imag * cos + real * sin - xq_out = xq * cos + rotate_half_force_bf16(xq) * sin - xk_out = xk * cos + rotate_half_force_bf16(xk) * sin - return xq_out, xk_out - - def get_nd_rotary_pos_embed( rope_dim_list, start, diff --git a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py index 87341b4f6..cba596177 100755 --- a/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/hunyuan_video/infer/transformer_infer.py @@ -1,14 +1,7 @@ -from typing import Tuple - import torch import torch.nn.functional as F from einops import rearrange -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except Exception as e: - apply_rope_with_cos_sin_cache_inplace = None - from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from .module_io import HunyuanVideo15ImgBranchOutput, HunyuanVideo15TxtBranchOutput @@ -38,61 +31,12 @@ def apply_gate(x, gate=None, tanh=False): return x * gate.unsqueeze(1) -def apply_hunyuan_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - B, L, H, D = xq.shape - - query = xq.reshape(B * L, H * D).contiguous() - key = xk.reshape(B * L, H * D).contiguous() - - positions = torch.arange(B * L, device=xq.device, dtype=torch.long) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(B, L, H, D) - xk_out = key.view(B, L, H, D) - return xq_out, xk_out - - -def apply_hunyuan_rope_with_torch( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - B, L, H, D = xq.shape - - cos = cos_sin_cache[:, : D // 2] - sin = cos_sin_cache[:, D // 2 :] - - def _apply_rope(x: torch.Tensor) -> torch.Tensor: - x_flat = x.view(B * L, H, D) - x1 = x_flat[..., ::2] - x2 = x_flat[..., 1::2] - - cos_ = cos.unsqueeze(1) - sin_ = sin.unsqueeze(1) - - o1 = x1.float() * cos_ - x2.float() * sin_ - o2 = x2.float() * cos_ + x1.float() * sin_ - - out = torch.empty_like(x_flat) - out[..., ::2] = o1 - out[..., 1::2] = o2 - return out.view(B, L, H, D) - - xq_out = _apply_rope(xq) - xk_out = _apply_rope(xk) - return xq_out, xk_out +def _apply_hunyuan_rope(module, xq, xk, cos_sin_cache): + q_shape, k_shape = xq.shape, xk.shape + q = xq.reshape(-1, xq.shape[-2], xq.shape[-1]) + k = xk.reshape(-1, xk.shape[-2], xk.shape[-1]) + q, k = module.apply(q, k, cos_sin_cache) + return q.view(q_shape), k.view(k_shape) class HunyuanVideo15TransformerInfer(BaseTransformerInfer): @@ -115,10 +59,6 @@ def __init__(self, config): self.modulate_func = fuse_scale_shift_kernel else: self.modulate_func = modulate - if self.config.get("rope_type", "flashinfer") == "flashinfer": - self.apply_rope_func = apply_hunyuan_rope_with_flashinfer - else: - self.apply_rope_func = apply_hunyuan_rope_with_torch def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -173,7 +113,7 @@ def _infer_img_branch_before_attn(self, weights, infer_module_out): img_v = rearrange(img_v, "L (H D) -> L H D", H=self.heads_num) img_q = weights.img_branch.img_attn_q_norm.apply(img_q) img_k = weights.img_branch.img_attn_k_norm.apply(img_k) - img_q, img_k = self.apply_rope_func(img_q.unsqueeze(0), img_k.unsqueeze(0), cos_sin_cache=infer_module_out.cos_sin) + img_q, img_k = _apply_hunyuan_rope(weights.rope, img_q.unsqueeze(0), img_k.unsqueeze(0), infer_module_out.cos_sin) return ( img_q, img_k, diff --git a/lightx2v/models/networks/hunyuan_video/weights/transformer_weights.py b/lightx2v/models/networks/hunyuan_video/weights/transformer_weights.py index 67bcb33ff..bcf9a95c6 100755 --- a/lightx2v/models/networks/hunyuan_video/weights/transformer_weights.py +++ b/lightx2v/models/networks/hunyuan_video/weights/transformer_weights.py @@ -1,9 +1,12 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ( ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, ) @@ -52,6 +55,10 @@ def __init__(self, block_index, task, config, block_prefix="double_blocks", crea self.block_index = block_index self.task = task self.config = config + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.create_cuda_buffer = create_cuda_buffer self.create_cpu_buffer = create_cpu_buffer diff --git a/lightx2v/models/networks/lingbot_video/infer/transformer_infer.py b/lightx2v/models/networks/lingbot_video/infer/transformer_infer.py index 3de889f52..01b68d2b2 100644 --- a/lightx2v/models/networks/lingbot_video/infer/transformer_infer.py +++ b/lightx2v/models/networks/lingbot_video/infer/transformer_infer.py @@ -2,7 +2,6 @@ import torch.nn.functional as F from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer -from lightx2v.models.networks.lingbot_video.infer.utils import apply_rotary_emb from lightx2v.utils.envs import GET_DTYPE @@ -29,8 +28,11 @@ def _attention(self, weights, hidden_states, rotary_emb): k = weights.attn.to_k.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) v = weights.attn.to_v.apply(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) - q = apply_rotary_emb(weights.attn.norm_q.apply(q), rotary_emb) - k = apply_rotary_emb(weights.attn.norm_k.apply(k), rotary_emb) + q, k = weights.attn.rope.apply( + weights.attn.norm_q.apply(q), + weights.attn.norm_k.apply(k), + rotary_emb, + ) seq_len = q.shape[0] cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=q.device) hidden_states = weights.attn.calculate.apply( diff --git a/lightx2v/models/networks/lingbot_video/infer/utils.py b/lightx2v/models/networks/lingbot_video/infer/utils.py index e109ac8f3..3b83c5c7e 100644 --- a/lightx2v/models/networks/lingbot_video/infer/utils.py +++ b/lightx2v/models/networks/lingbot_video/infer/utils.py @@ -31,13 +31,6 @@ def get_timestep_embedding( return emb -def apply_rotary_emb(x, freqs_cis): - with torch.amp.autocast(x.device.type, enabled=False): - x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - out = torch.view_as_real(x_c * freqs_cis.unsqueeze(1)).flatten(2) - return out.type_as(x) - - def precompute_freqs_cis(dim, end, theta): freqs_cis = [] for d, e in zip(dim, end): diff --git a/lightx2v/models/networks/lingbot_video/weights/transformer_weights.py b/lightx2v/models/networks/lingbot_video/weights/transformer_weights.py index 7a957b979..a979efcbd 100644 --- a/lightx2v/models/networks/lingbot_video/weights/transformer_weights.py +++ b/lightx2v/models/networks/lingbot_video/weights/transformer_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER, TENSOR_REGISTER class LingBotVideoTransformerWeights(WeightModule): @@ -29,6 +31,10 @@ def __init__(self, block_index, config): class LingBotVideoAttentionWeights(WeightModule): def __init__(self, prefix, config): super().__init__() + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float32), + ) mm_type = config.get("dit_quant_scheme", "Default") if mm_type != "Default": raise NotImplementedError("LingBot-Video t2i currently supports original BF16 transformer weights only.") diff --git a/lightx2v/models/networks/longcat_image/infer/transformer_infer.py b/lightx2v/models/networks/longcat_image/infer/transformer_infer.py index b8e70d36d..fed2d9e53 100755 --- a/lightx2v/models/networks/longcat_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/longcat_image/infer/transformer_infer.py @@ -3,8 +3,6 @@ from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer -from .utils import apply_longcat_rope_with_flashinfer, apply_longcat_rope_with_torch - class LongCatImageTransformerInfer(BaseTransformerInfer): """Transformer inference for LongCat Image model. @@ -29,14 +27,6 @@ def __init__(self, config): self.seq_p_fp4_comm = False self.enable_head_parallel = False - # RoPE function selection - rope_funcs = { - "flashinfer": apply_longcat_rope_with_flashinfer, - "torch": apply_longcat_rope_with_torch, - } - rope_type = config.get("rope_type", "flashinfer") - self.apply_rope_func = rope_funcs.get(rope_type, apply_longcat_rope_with_torch) - def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -136,7 +126,7 @@ def infer_double_stream_block( value = torch.cat([txt_value, img_value], dim=0) # Apply rotary embedding: [L, H, D] - query, key = self.apply_rope_func(query, key, image_rotary_emb) + query, key = block_weights.rope.apply(query, key, image_rotary_emb) # Calculate cu_seqlens for flash attention (batch_size=1) total_len = query.shape[0] @@ -263,7 +253,7 @@ def infer_single_stream_block( key = block_weights.norm_k.apply(key) # Apply rotary embedding: [L, H, D] - query, key = self.apply_rope_func(query, key, image_rotary_emb) + query, key = block_weights.rope.apply(query, key, image_rotary_emb) # Calculate cu_seqlens for flash attention (batch_size=1) total_len = query.shape[0] diff --git a/lightx2v/models/networks/longcat_image/infer/utils.py b/lightx2v/models/networks/longcat_image/infer/utils.py index b4cb5d737..aa894f845 100644 --- a/lightx2v/models/networks/longcat_image/infer/utils.py +++ b/lightx2v/models/networks/longcat_image/infer/utils.py @@ -1,80 +1 @@ -from typing import Tuple - -import torch - -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - - -def apply_longcat_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - """Apply rotary position embedding using flashinfer. - - Args: - xq: Query tensor [L, H, D] - xk: Key tensor [L, H, D] - cos_sin_cache: Cosine and sine cache [L, D] where first half is cos, second half is sin - - Returns: - Tuple of (xq, xk) with rotary embedding applied - """ - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous() - key = xk.reshape(L, H * D).contiguous() - - # Create positions directly on GPU to avoid CPU-GPU sync - positions = torch.arange(L, device=xq.device, dtype=torch.long) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out - - -def apply_longcat_rope_with_torch( - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cis: Tuple[torch.Tensor, torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - """Apply rotary position embedding using PyTorch. - - Follows the diffusers implementation for LongCat/Flux. - - Args: - xq: Query tensor [L, H, D] - xk: Key tensor [L, H, D] - freqs_cis: Tuple of (cos, sin) each [L, D] - - Returns: - Tuple of (xq, xk) with rotary embedding applied - """ - cos, sin = freqs_cis # [L, D] - - # Expand for heads: [L, D] -> [L, 1, D] - cos = cos[:, None, :] - sin = sin[:, None, :] - - def _apply_rope(x, cos, sin): - # Split into real and imaginary parts (interleaved format) - x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [L, H, D//2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - return out - - xq_out = _apply_rope(xq, cos, sin) - xk_out = _apply_rope(xk, cos, sin) - return xq_out, xk_out +"""RoPE is selected by LongCat attention weights and applied directly during inference.""" diff --git a/lightx2v/models/networks/longcat_image/weights/transformer_weights.py b/lightx2v/models/networks/longcat_image/weights/transformer_weights.py index 7c081d964..19de63a36 100755 --- a/lightx2v/models/networks/longcat_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/longcat_image/weights/transformer_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER class LongCatImageDoubleBlockWeights(WeightModule): @@ -13,6 +15,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.mm_type = config.get("dit_quant_scheme", "Default") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) p = f"transformer_blocks.{self.block_idx}" @@ -218,6 +224,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe self.mm_type = config.get("dit_quant_scheme", "Default") self.rms_norm_type = config.get("rms_norm_type", "torch") self.attn_type = config.get("attn_type", "flash_attn3") + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) p = f"single_transformer_blocks.{self.block_idx}" diff --git a/lightx2v/models/networks/ltx2/infer/transformer_infer.py b/lightx2v/models/networks/ltx2/infer/transformer_infer.py index ad0404bbd..f1fd3d5bd 100644 --- a/lightx2v/models/networks/ltx2/infer/transformer_infer.py +++ b/lightx2v/models/networks/ltx2/infer/transformer_infer.py @@ -16,7 +16,7 @@ from lightx2v.models.networks.ltx2.infer.module_io import LTX2PreInferModuleOutput from lightx2v.models.networks.ltx2.infer.triton_ops import fuse_scale_shift_kernel, fused_rmsnorm_modulate -from lightx2v.models.networks.ltx2.infer.utils import apply_rotary_emb, modulate_torch_naive, modulate_with_rmsnorm_torch_naive, rmsnorm_torch_naive +from lightx2v.models.networks.ltx2.infer.utils import modulate_torch_naive, modulate_with_rmsnorm_torch_naive, rmsnorm_torch_naive from lightx2v.models.networks.wan.infer.triton_ops import norm_infer @@ -35,7 +35,6 @@ def __init__(self, config): config: Model configuration dictionary """ self.config = config - self.rope_type = config["rope_type"] self.blocks_num = config.get("num_layers", 48) self.v_num_heads = config.get("num_attention_heads", 32) self.v_head_dim = config.get("attention_head_dim", 128) @@ -75,6 +74,15 @@ def __init__(self, config): self.reset_infer_states() self.reset_guidance_perturbation() + def _apply_rope(self, rope, tensor, freqs): + cos, _ = freqs + if tensor.ndim != 4 and cos.ndim == 4: + batch, heads, tokens, _ = cos.shape + tensor = tensor.reshape(batch, tokens, heads, -1).swapaxes(1, 2) + output = rope.apply_single(tensor, freqs) + return output.swapaxes(1, 2).reshape(batch, tokens, -1) + return rope.apply_single(tensor, freqs) + def set_scheduler(self, scheduler): """Set the scheduler for inference.""" self.scheduler = scheduler @@ -266,8 +274,8 @@ def _infer_attn( q = attn_phase.q_norm.apply(q) k = attn_phase.k_norm.apply(k) if pe is not None: - q = apply_rotary_emb(q, pe, self.rope_type) - k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type) + q = self._apply_rope(attn_phase.rope, q, pe) + k = self._apply_rope(attn_phase.rope, k, pe if k_pe is None else k_pe) q = q.view(-1, num_heads_effective, head_dim) k = k.view(-1, num_heads_effective, head_dim) diff --git a/lightx2v/models/networks/ltx2/infer/utils.py b/lightx2v/models/networks/ltx2/infer/utils.py index cfc35bf7a..cf9524860 100755 --- a/lightx2v/models/networks/ltx2/infer/utils.py +++ b/lightx2v/models/networks/ltx2/infer/utils.py @@ -1,10 +1,9 @@ import functools import math -from typing import Callable, Tuple +from typing import Callable import numpy as np import torch -from einops import rearrange def rmsnorm_torch_naive(x, weight=None, bias=None, eps=1e-6): @@ -70,55 +69,6 @@ def get_timestep_embedding( return emb -def apply_rotary_emb( - input_tensor: torch.Tensor, - freqs_cis: Tuple[torch.Tensor, torch.Tensor], - rope_type: str = "split", -) -> torch.Tensor: - if rope_type == "interleaved": - return apply_interleaved_rotary_emb(input_tensor, *freqs_cis) - elif rope_type == "split": - return apply_split_rotary_emb(input_tensor, *freqs_cis) - else: - raise ValueError(f"Invalid rope type: {rope_type}") - - -def apply_interleaved_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor) -> torch.Tensor: - t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) - t1, t2 = t_dup.unbind(dim=-1) - t_dup = torch.stack((-t2, t1), dim=-1) - input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") - - out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs - - return out - - -def apply_split_rotary_emb(input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor) -> torch.Tensor: - needs_reshape = False - if input_tensor.ndim != 4 and cos_freqs.ndim == 4: - b, h, t, _ = cos_freqs.shape - input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2) - needs_reshape = True - - split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) - first_half_input = split_input[..., :1, :] - second_half_input = split_input[..., 1:, :] - - output = split_input * cos_freqs.unsqueeze(-2) - first_half_output = output[..., :1, :] - second_half_output = output[..., 1:, :] - - first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input) - second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input) - - output = rearrange(output, "... d r -> ... (d r)") - if needs_reshape: - output = output.swapaxes(1, 2).reshape(b, t, -1) - - return output - - @functools.lru_cache(maxsize=5) def generate_freq_grid_np(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.Tensor: theta = positional_embedding_theta diff --git a/lightx2v/models/networks/ltx2/weights/transformer_weights.py b/lightx2v/models/networks/ltx2/weights/transformer_weights.py index 9688a8fff..3c2f1a92d 100755 --- a/lightx2v/models/networks/ltx2/weights/transformer_weights.py +++ b/lightx2v/models/networks/ltx2/weights/transformer_weights.py @@ -1,3 +1,4 @@ +import torch import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList @@ -5,6 +6,7 @@ ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, TENSOR_REGISTER, ) @@ -329,6 +331,10 @@ def __init__( self.lazy_load_file = lazy_load_file self.attn_rms_norm_type = self.config.get("rms_norm_type", "sgl-kernel") self.apply_gated_attention = self.config.get("apply_gated_attention", False) + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_real_rope")](layout="split_half", compute_dtype=torch.float32), + ) block_lora_prefix = "model.diffusion_model.blocks" model_prefix = "model.diffusion_model" @@ -527,6 +533,10 @@ def __init__( self.lazy_load_file = lazy_load_file self.attn_rms_norm_type = self.config.get("rms_norm_type", "sgl-kernel") self.apply_gated_attention = self.config.get("apply_gated_attention", False) + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_real_rope")](layout="split_half", compute_dtype=torch.float32), + ) block_lora_prefix = "model.diffusion_model.blocks" model_prefix = "model.diffusion_model" diff --git a/lightx2v/models/networks/motus/infer/transformer_infer.py b/lightx2v/models/networks/motus/infer/transformer_infer.py index e00cfa393..334e25493 100644 --- a/lightx2v/models/networks/motus/infer/transformer_infer.py +++ b/lightx2v/models/networks/motus/infer/transformer_infer.py @@ -1,18 +1,10 @@ -from functools import partial - import torch from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.wan.infer.triton_ops import fuse_scale_shift_kernel -from lightx2v.models.networks.wan.infer.utils import ( - apply_wan_rope_with_chunk, - apply_wan_rope_with_flashinfer, - apply_wan_rope_with_torch, - apply_wan_rope_with_torch_naive, -) from lightx2v.models.networks.wan.weights.motus import apply_mm from lightx2v.utils.envs import GET_DTYPE, GET_SENSITIVE_DTYPE -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, ROPE_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE torch_device_module = getattr(torch, AI_DEVICE) @@ -38,30 +30,6 @@ def __init__(self, model, config): self.head_dim = config["dim"] // config["num_heads"] self.modulate_func = fuse_scale_shift_kernel if config.get("modulate_type", "triton") == "triton" else modulate - rope_funcs = { - "flashinfer": apply_wan_rope_with_flashinfer, - "torch": apply_wan_rope_with_torch, - "torch_naive": apply_wan_rope_with_torch_naive, - } - rope_type = config.get("rope_type", "flashinfer") - if rope_type in ROPE_REGISTER: - rope_class = ROPE_REGISTER[rope_type] - self.rope_instance = rope_class() - - def rope_wrapper(xq, xk, cos_sin_cache): - return self.rope_instance.apply(xq, xk, cos_sin_cache) - - rope_func = rope_wrapper - else: - rope_func = rope_funcs.get(rope_type, apply_wan_rope_with_torch) - if config.get("rope_chunk", False): - rope_func = partial( - apply_wan_rope_with_chunk, - chunk_size=config.get("rope_chunk_size", 100), - rope_func=rope_func, - ) - self.apply_rope_func = rope_func - if self.config["seq_parallel"]: self.seq_p_group = self.config.get("device_mesh").get_group(mesh_dim="seq_p") self.seq_p_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False) @@ -170,13 +138,13 @@ def _apply_video_rope(self, q, k, cos_sin_cache): raise ValueError("Motus video rope expects q/k with shape [B, L, H, D].") if q.shape[0] == 1: - q_out, k_out = self.apply_rope_func(q.squeeze(0), k.squeeze(0), cos_sin_cache) + q_out, k_out = self.weights.rope.apply(q.squeeze(0), k.squeeze(0), cos_sin_cache) return q_out.unsqueeze(0), k_out.unsqueeze(0) q_list = [] k_list = [] for batch_idx in range(q.shape[0]): - q_i, k_i = self.apply_rope_func(q[batch_idx], k[batch_idx], cos_sin_cache) + q_i, k_i = self.weights.rope.apply(q[batch_idx], k[batch_idx], cos_sin_cache) q_list.append(q_i) k_list.append(k_i) return torch.stack(q_list, dim=0), torch.stack(k_list, dim=0) diff --git a/lightx2v/models/networks/neopp/infer/pre_infer.py b/lightx2v/models/networks/neopp/infer/pre_infer.py index 18a11ba20..b9a0cc5af 100644 --- a/lightx2v/models/networks/neopp/infer/pre_infer.py +++ b/lightx2v/models/networks/neopp/infer/pre_infer.py @@ -46,34 +46,12 @@ def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None): return abs_x, abs_y -def apply_rotary_emb_1d( - x: torch.Tensor, - cos_cached: torch.Tensor, - sin_cached: torch.Tensor, - positions: torch.Tensor, -): - """对输入张量的一部分应用1D RoPE。""" - # x: (..., seq_len, dim_part) - # positions: (..., seq_len) - # cos_cached: (max_pos, dim_part / 2) - - cos = cos_cached[positions] # Shape: (positions.shape, dim_part / 2) - sin = sin_cached[positions] # Shape: (positions.shape, dim_part / 2) - - x1 = x[..., 0::2] - x2 = x[..., 1::2] - - rotated_x1 = x1 * cos - x2 * sin - rotated_x2 = x1 * sin + x2 * cos - - x_rotated = torch.empty_like(x) - x_rotated[..., 0::2] = rotated_x1 - x_rotated[..., 1::2] = rotated_x2 - return x_rotated +def apply_rotary_emb_1d(rope, x, cos_cached, sin_cached, positions): + return rope.apply_single(x, (cos_cached[positions], sin_cached[positions])) def apply_2d_rotary_pos_emb( - x: torch.Tensor, cos_cached_x: torch.Tensor, sin_cached_x: torch.Tensor, cos_cached_y: torch.Tensor, sin_cached_y: torch.Tensor, abs_positions_x: torch.Tensor, abs_positions_y: torch.Tensor + rope, x: torch.Tensor, cos_cached_x: torch.Tensor, sin_cached_x: torch.Tensor, cos_cached_y: torch.Tensor, sin_cached_y: torch.Tensor, abs_positions_x: torch.Tensor, abs_positions_y: torch.Tensor ): """应用2D RoPE到输入张量x。""" dim = x.shape[-1] @@ -85,9 +63,9 @@ def apply_2d_rotary_pos_emb( x_part_2 = x[..., dim_half:] # 将与 abs_positions_x 相关的旋转应用于 x_part_1 - rotated_part_1 = apply_rotary_emb_1d(x_part_1, cos_cached_x, sin_cached_x, abs_positions_x) + rotated_part_1 = apply_rotary_emb_1d(rope, x_part_1, cos_cached_x, sin_cached_x, abs_positions_x) # 将与 abs_positions_y 相关的旋转应用于 x_part_2 - rotated_part_2 = apply_rotary_emb_1d(x_part_2, cos_cached_y, sin_cached_y, abs_positions_y) + rotated_part_2 = apply_rotary_emb_1d(rope, x_part_2, cos_cached_y, sin_cached_y, abs_positions_y) # 将它们重新拼接起来。确保顺序与你分割时一致。 return torch.cat((rotated_part_1, rotated_part_2), dim=-1) @@ -162,12 +140,13 @@ def patchify(self, images, patch_size, channel_first=False): x = x.reshape(shape=(images.shape[0], h * w, patch_size**2 * 3)) return x - def _apply_2d_rotary_pos_emb(self, patch_embeds, grid_hw): + def _apply_2d_rotary_pos_emb(self, rope, patch_embeds, grid_hw): """ Apply 2D Rotary Position Embedding to the patch embeddings. """ abs_pos_x, abs_pos_y = build_abs_positions_from_grid_hw(grid_hw, device=patch_embeds.device) embeddings = apply_2d_rotary_pos_emb( + rope, patch_embeds.to(torch.float32), # RoPE calculations are often more stable in float32 self.cos_cached_x, self.sin_cached_x, @@ -190,7 +169,7 @@ def extract_feature(self, weights, pixel_values: torch.FloatTensor, grid_hw=None self.sin_cached_x = self.sin_cached_x.to(patch_embeds.device) self.cos_cached_y = self.cos_cached_y.to(patch_embeds.device) self.sin_cached_y = self.sin_cached_y.to(patch_embeds.device) - patch_embeds = self._apply_2d_rotary_pos_emb(patch_embeds, grid_hw) # [28072, 1024] + patch_embeds = self._apply_2d_rotary_pos_emb(weights.rope, patch_embeds, grid_hw) # [28072, 1024] assert (grid_hw[:, 0] * grid_hw[:, 1]).sum() == patch_embeds.shape[0] patches_list = [] diff --git a/lightx2v/models/networks/neopp/weights/pre_weights.py b/lightx2v/models/networks/neopp/weights/pre_weights.py index f91a80209..e6efc49c6 100755 --- a/lightx2v/models/networks/neopp/weights/pre_weights.py +++ b/lightx2v/models/networks/neopp/weights/pre_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule -from lightx2v.utils.registry_factory import CONV2D_WEIGHT_REGISTER, MM_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import CONV2D_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, ROPE_REGISTER class NeoppPreWeights(WeightModule): @@ -9,6 +11,10 @@ def __init__(self, config): self.patch_size = config.get("patch_size", 16) self.merge_size = config.get("merge_size", 2) self.mm_type = "Default" + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "torch_real_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.add_module( "vision_model_mot_gen_patch_embedding", diff --git a/lightx2v/models/networks/qwen_image/infer/transformer_infer.py b/lightx2v/models/networks/qwen_image/infer/transformer_infer.py index 3c896b6c8..dd6628c84 100755 --- a/lightx2v/models/networks/qwen_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/qwen_image/infer/transformer_infer.py @@ -11,7 +11,6 @@ fuse_scale_shift_gate_select01_kernel, fuse_scale_shift_kernel, ) -from .utils import apply_qwen_rope_with_flashinfer, apply_qwen_rope_with_torch, apply_qwen_rope_with_torch_naive try: from magi_compiler import magi_compile @@ -51,14 +50,6 @@ def __init__(self, config): self.modulate_func = fuse_scale_shift_kernel else: self.modulate_func = lambda x, scale, shift: x * (1 + scale) + shift - rope_funcs = { - "flashinfer": apply_qwen_rope_with_flashinfer, - "torch": apply_qwen_rope_with_torch, - "torch_naive": apply_qwen_rope_with_torch_naive, - } - rope_type = config.get("rope_type", "flashinfer") - self.apply_rope_func = rope_funcs.get(rope_type, apply_qwen_rope_with_torch) - self.img_qkv_len1 = None self.cu_seqlens_qkv1 = None self.img_qkv_len2 = None @@ -134,7 +125,7 @@ def infer_img_qkv( if img_attn_phase.norm_k is not None: img_key = img_attn_phase.norm_k.apply(img_key) - img_query, img_key = self.apply_rope_func(img_query, img_key, img_freqs) + img_query, img_key = img_attn_phase.rope.apply(img_query, img_key, img_freqs) return img_query, img_key, img_value, img_gate1, img_mod2 @@ -161,7 +152,7 @@ def infer_txt_qkv(self, txt_attn_phase, encoder_hidden_states, temb_txt_silu, tx if txt_attn_phase.norm_added_k is not None: txt_key = txt_attn_phase.norm_added_k.apply(txt_key) - txt_query, txt_key = self.apply_rope_func(txt_query, txt_key, txt_freqs) + txt_query, txt_key = txt_attn_phase.rope.apply(txt_query, txt_key, txt_freqs) return txt_query, txt_key, txt_value, seq_txt, txt_gate1, txt_mod2 diff --git a/lightx2v/models/networks/qwen_image/infer/utils.py b/lightx2v/models/networks/qwen_image/infer/utils.py index 1b1851b84..8792883aa 100755 --- a/lightx2v/models/networks/qwen_image/infer/utils.py +++ b/lightx2v/models/networks/qwen_image/infer/utils.py @@ -1,143 +1 @@ -from typing import Tuple - -import torch - -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - -try: - from magi_compiler import magi_register_custom_op -except ImportError: - magi_register_custom_op = None - -# Re-export for transformer_infer init. -from lightx2v.common.magi_custom_op_mode import ( - set_magi_custom_op_mode, # noqa: F401 - use_magi_custom_ops, -) - - -def _qwen_rope_meta(xq, xk, cos_sin_cache): - return torch.empty_like(xq), torch.empty_like(xk) - - -def _apply_qwen_rope_with_flashinfer_eager(xq, xk, cos_sin_cache): - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous() - key = xk.reshape(L, H * D).contiguous() - - positions = torch.arange(L, device="cpu", dtype=torch.long).to(xq.device, non_blocking=True) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out - - -def _apply_qwen_rope_with_flashinfer_magi(xq, xk, cos_sin_cache): - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous().clone() - key = xk.reshape(L, H * D).contiguous().clone() - - positions = torch.arange(L, device="cpu", dtype=torch.long).to(xq.device, non_blocking=True) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out - - -def _apply_qwen_rope_with_torch_impl(xq, xk, cos_sin_cache): - xq_rotated = torch.view_as_complex(xq.float().unflatten(-1, (-1, 2))) - xk_rotated = torch.view_as_complex(xk.float().unflatten(-1, (-1, 2))) - freqs_cis = cos_sin_cache.unsqueeze(1) - xq_out = torch.view_as_real(xq_rotated * freqs_cis).flatten(-2) - xk_out = torch.view_as_real(xk_rotated * freqs_cis).flatten(-2) - return xq_out.type_as(xq), xk_out.type_as(xk) - - -if magi_register_custom_op is not None and apply_rope_with_cos_sin_cache_inplace is not None: - - @magi_register_custom_op( - "lightx2v::qwen_rope_flashinfer", - infer_output_meta_fn=_qwen_rope_meta, - is_subgraph_boundary=True, - ) - def _qwen_rope_flashinfer_custom_op(xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - return _apply_qwen_rope_with_flashinfer_magi(xq, xk, cos_sin_cache) - - -if magi_register_custom_op is not None: - - @magi_register_custom_op( - "lightx2v::qwen_rope_torch", - infer_output_meta_fn=_qwen_rope_meta, - is_subgraph_boundary=True, - ) - def _qwen_rope_torch_custom_op(xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - return _apply_qwen_rope_with_torch_impl(xq, xk, cos_sin_cache) - - -def apply_qwen_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - if use_magi_custom_ops() and magi_register_custom_op is not None and apply_rope_with_cos_sin_cache_inplace is not None: - return torch.ops.lightx2v.qwen_rope_flashinfer(xq, xk, cos_sin_cache) - return _apply_qwen_rope_with_flashinfer_eager(xq, xk, cos_sin_cache) - - -def apply_qwen_rope_with_torch( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - if use_magi_custom_ops() and magi_register_custom_op is not None: - return torch.ops.lightx2v.qwen_rope_torch(xq, xk, cos_sin_cache) - return _apply_qwen_rope_with_torch_impl(xq, xk, cos_sin_cache) - - -def apply_qwen_rope_with_torch_naive( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - cos = cos_sin_cache.real.unsqueeze(1) - sin = cos_sin_cache.imag.unsqueeze(1) - - def _rotate(x: torch.Tensor) -> torch.Tensor: - x_even = x[..., 0::2] - x_odd = x[..., 1::2] - - x_rot_even = x_even * cos - x_odd * sin - x_rot_odd = x_even * sin + x_odd * cos - - x_out = torch.empty_like(x) - x_out[..., 0::2] = x_rot_even - x_out[..., 1::2] = x_rot_odd - return x_out - - xq_out = _rotate(xq) - xk_out = _rotate(xk) - return xq_out, xk_out +from lightx2v.common.magi_custom_op_mode import set_magi_custom_op_mode, use_magi_custom_ops # noqa: F401 diff --git a/lightx2v/models/networks/qwen_image/weights/transformer_weights.py b/lightx2v/models/networks/qwen_image/weights/transformer_weights.py index 6dfdb7fe1..3e5552b32 100755 --- a/lightx2v/models/networks/qwen_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/qwen_image/weights/transformer_weights.py @@ -1,9 +1,12 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ( ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, ) @@ -210,6 +213,10 @@ def __init__( self.layer_norm_type = config.get("layer_norm_type", "Triton") self.lazy_load = lazy_load self.lazy_load_file = lazy_load_file + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.add_module( "img_mod", @@ -331,6 +338,10 @@ def __init__( self.layer_norm_type = config.get("layer_norm_type", "Triton") self.lazy_load = lazy_load self.lazy_load_file = lazy_load_file + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.add_module( "txt_mod", diff --git a/lightx2v/models/networks/seedvr/infer/transformer_infer.py b/lightx2v/models/networks/seedvr/infer/transformer_infer.py index 26607a9be..890fd8c66 100644 --- a/lightx2v/models/networks/seedvr/infer/transformer_infer.py +++ b/lightx2v/models/networks/seedvr/infer/transformer_infer.py @@ -7,7 +7,6 @@ from lightx2v.models.networks.seedvr.utils import na from lightx2v.models.networks.seedvr.utils.attention import FlashAttentionVarlen from lightx2v.models.networks.seedvr.utils.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv, safe_pad_operation -from lightx2v.models.networks.seedvr.utils.rope import get_na_rope from lightx2v.models.networks.seedvr.utils.window import get_window_op from .utils import apply_adaln_single, norm_no_weight @@ -22,11 +21,8 @@ def __init__(self, config): self.norm_type = config.get("norm", "fusedrms") self.qk_norm_type = config.get("qk_norm", "fusedrms") self.norm_eps = config.get("norm_eps", 1.0e-5) - self.rope_type = config.get("rope_type", None) - self.rope_dim = config.get("rope_dim", None) self.mlp_type = config.get("mlp_type", "swiglu") - self.rope = get_na_rope(rope_type=self.rope_type, dim=self.rope_dim) if self.rope_type else None self.attn = FlashAttentionVarlen() def set_scheduler(self, scheduler): @@ -60,11 +56,10 @@ def _attn_forward(self, block_weight, vid, txt, vid_shape, txt_shape, cache): vid_k = norm_k_vid.apply(vid_k) txt_k = norm_k_txt.apply(txt_k) - if self.rope is not None: - if self.rope.mm: - vid_q, vid_k, txt_q, txt_k = self.rope(vid_q, vid_k, vid_shape, txt_q, txt_k, txt_shape, cache) - else: - vid_q, vid_k = self.rope(vid_q, vid_k, vid_shape, cache) + if block_weight.rope.multimodal: + vid_q, vid_k, txt_q, txt_k = block_weight.rope.apply(vid_q, vid_k, vid_shape, txt_q=txt_q, txt_k=txt_k, txt_shape=txt_shape, cache=cache) + else: + vid_q, vid_k = block_weight.rope.apply(vid_q, vid_k, vid_shape, cache=cache) vid_len = cache("vid_len", lambda: vid_shape.prod(-1)) txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) @@ -139,34 +134,33 @@ def make_window(x: torch.Tensor): all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win) concat_win, unconcat_win = cache_win("mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)) - if self.rope is not None: - if self.rope.mm: - _, num_h, _ = txt_q.shape - txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)") - txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape) - txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)] - txt_q_repeat = [t for sub in txt_q_repeat for t in sub] - txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat) - txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h) - - txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)") - txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape) - txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)] - txt_k_repeat = [t for sub in txt_k_repeat for t in sub] - txt_k_repeat, _ = na.flatten(txt_k_repeat) - txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h) - - vid_q, vid_k, txt_q, txt_k = self.rope( - vid_q, - vid_k, - window_shape, - txt_q_repeat, - txt_k_repeat, - txt_shape_repeat, - cache_win, - ) - else: - vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win) + if block_weight.rope.multimodal: + _, num_h, _ = txt_q.shape + txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)") + txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape) + txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)] + txt_q_repeat = [t for sub in txt_q_repeat for t in sub] + txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat) + txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h) + + txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)") + txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape) + txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)] + txt_k_repeat = [t for sub in txt_k_repeat for t in sub] + txt_k_repeat, _ = na.flatten(txt_k_repeat) + txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h) + + vid_q, vid_k, txt_q, txt_k = block_weight.rope.apply( + vid_q, + vid_k, + window_shape, + txt_q=txt_q_repeat, + txt_k=txt_k_repeat, + txt_shape=txt_shape_repeat, + cache=cache_win, + ) + else: + vid_q, vid_k = block_weight.rope.apply(vid_q, vid_k, window_shape, cache=cache_win) out = self.attn( q=concat_win(vid_q, txt_q).bfloat16(), diff --git a/lightx2v/models/networks/seedvr/utils/rope.py b/lightx2v/models/networks/seedvr/utils/rope.py index d6fb90a35..0f38ec44a 100755 --- a/lightx2v/models/networks/seedvr/utils/rope.py +++ b/lightx2v/models/networks/seedvr/utils/rope.py @@ -1,10 +1,13 @@ from functools import lru_cache -from typing import Optional, Tuple +from typing import Tuple import torch from einops import rearrange from torch import nn +from lightx2v.common.ops.rope import RopeTemplate +from lightx2v.utils.registry_factory import ROPE_REGISTER + from .cache import Cache from .rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb @@ -161,13 +164,35 @@ def get_freqs( return torch.cat(vid_freq_list, dim=0), torch.cat(txt_freq_list, dim=0) -def get_na_rope(rope_type: Optional[str], dim: int): - if rope_type is None: - return None - if rope_type == "torch": - return None - if rope_type == "rope3d": - return NaRotaryEmbedding3d(dim=dim) - if rope_type == "mmrope3d": - return NaMMRotaryEmbedding3d(dim=dim) - raise NotImplementedError(f"{rope_type} is not supported.") +class SeedVRRopeBase(RopeTemplate): + multimodal = False + implementation_class = None + + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("SeedVR RoPE only supports interleaved layout.") + self.implementation = None + + def set_config(self, config=None): + super().set_config(config) + if config is not None: + self.implementation = self.implementation_class(dim=config["rope_dim"]) + + def apply(self, q, k, shape, *, cache, txt_q=None, txt_k=None, txt_shape=None, **kwargs): + if self.implementation is None: + raise RuntimeError("SeedVR RoPE must be configured before apply().") + if self.multimodal: + return self.implementation(q, k, shape, txt_q, txt_k, txt_shape, cache) + return self.implementation(q, k, shape, cache) + + +@ROPE_REGISTER("rope3d") +class SeedVRRope3D(SeedVRRopeBase): + implementation_class = NaRotaryEmbedding3d + + +@ROPE_REGISTER("mmrope3d") +class SeedVRMMRope3D(SeedVRRopeBase): + multimodal = True + implementation_class = NaMMRotaryEmbedding3d diff --git a/lightx2v/models/networks/seedvr/utils/rotary_embedding_torch.py b/lightx2v/models/networks/seedvr/utils/rotary_embedding_torch.py index 5957eadea..2291e378b 100644 --- a/lightx2v/models/networks/seedvr/utils/rotary_embedding_torch.py +++ b/lightx2v/models/networks/seedvr/utils/rotary_embedding_torch.py @@ -9,6 +9,10 @@ from torch.amp import autocast from torch.nn import Module +from lightx2v.common.ops.rope import TorchRealRope + +_SEEDVR_ROPE = TorchRealRope(layout="interleaved") + # helper functions @@ -38,13 +42,6 @@ def slice_at_dim(t, dim_slice: slice, *, dim): # rotary embedding helper functions -def rotate_half(x): - x = rearrange(x, "... (d r) -> ... d r", r=2) - x1, x2 = x.unbind(dim=-1) - x = torch.stack((-x2, x1), dim=-1) - return rearrange(x, "... d r -> ... (d r)") - - @autocast("cuda", enabled=False) def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2, freqs_seq_dim=None): dtype = t.dtype @@ -68,7 +65,11 @@ def apply_rotary_emb(freqs, t, start_index=0, scale=1.0, seq_dim=-2, freqs_seq_d t_right = t[..., end_index:] # Apply rotary embeddings without modifying t in place - t_transformed = (t_middle * freqs.cos() * scale) + (rotate_half(t_middle) * freqs.sin() * scale) + t_transformed = _SEEDVR_ROPE.apply_single( + t_middle, + (freqs.cos() * scale, freqs.sin() * scale), + unsqueeze_dim=0, + ) out = torch.cat((t_left, t_transformed, t_right), dim=-1) diff --git a/lightx2v/models/networks/seedvr/weights/transformer_weights.py b/lightx2v/models/networks/seedvr/weights/transformer_weights.py index 9fe6ead6c..b4c2d4604 100755 --- a/lightx2v/models/networks/seedvr/weights/transformer_weights.py +++ b/lightx2v/models/networks/seedvr/weights/transformer_weights.py @@ -1,5 +1,8 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER +from lightx2v.models.networks.seedvr.utils import rope as _seedvr_rope # noqa: F401 +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER, TENSOR_REGISTER class SeedVRTransformerWeights(WeightModule): @@ -24,6 +27,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): last_layer_vid_only = bool(config.get("last_layer_vid_only")) blocks = WeightModuleList( SeedVRTransformerBlockWeights( + config=config, block_index=i, shared_weights=not ((i < mm_layers) if isinstance(mm_layers, int) else mm_layers[i]), vid_only=(last_layer_vid_only and i == self.blocks_num - 1), @@ -44,6 +48,7 @@ class SeedVRTransformerBlockWeights(WeightModule): def __init__( self, *, + config, block_index: int, shared_weights: bool, vid_only: bool, @@ -64,6 +69,11 @@ def __init__( self.mlp_type = mlp_type self.norm_eps = norm_eps self.rms_norm_type = rms_norm_type + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "rope3d")](layout="interleaved", compute_dtype=torch.float32), + ) + self.rope.set_config(config) branches = ["all"] if shared_weights else ["vid", "txt"] self.branches = branches diff --git a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py index 5a416d5f4..6b6ad6fc9 100755 --- a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py @@ -12,7 +12,6 @@ from lightx2v.models.input_encoders.hf.seko_audio.audio_adapter import align_hidden_states_and_mask, calculate_n_query_tokens, get_qk_lens_audio_range from lightx2v.models.networks.wan.infer.offload.transformer_infer import WanOffloadTransformerInfer from lightx2v.models.networks.wan.infer.self_forcing.transformer_infer import WanSFTransformerInfer -from lightx2v.models.networks.wan.infer.triton_ops import apply_audio_cache_rope from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER from lightx2v_platform.base.global_var import AI_DEVICE @@ -213,6 +212,7 @@ def _rope_freqs_for_cache_range( def _apply_rope_with_cache_range( self, + phase, x, freqs, h, @@ -226,38 +226,20 @@ def _apply_rope_with_cache_range( global_end=None, sink_tokens=0, ): - orig_dtype = x.dtype - if self.config.get("causal_rope_type", "triton") == "triton": - return apply_audio_cache_rope( - x, - freqs, - h=h, - w=w, - token_start=token_start, - ref_tokens=ref_tokens, - local_per_frame=local_per_frame, - world_size=world_size, - rank=rank, - global_end=global_end, - sink_tokens=sink_tokens, - ).to(orig_dtype) - pos_freqs = self._rope_freqs_for_cache_range( + return phase.causal_rope.apply_audio_cache( + x, freqs, - h, - w, - world_size, - rank, - token_start, - token_end, - ref_tokens, - local_per_frame, + h=h, + w=w, + world_size=world_size, + rank=rank, + token_start=token_start, + token_end=token_end, + ref_tokens=ref_tokens, + local_per_frame=local_per_frame, global_end=global_end, sink_tokens=sink_tokens, ) - n = x.size(1) - x_c = torch.view_as_complex(x.float().reshape(x.size(0), n, -1, 2)) - out = torch.view_as_real(x_c * pos_freqs.to(torch.complex64)).flatten(2) - return out.to(orig_dtype) def infer_block_with_kvcache(self, block, x, pre_infer_out): if hasattr(block.compute_phases[0], "before_proj"): @@ -343,15 +325,15 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh if seq_parallel: if replicated_ref_prefill: - q_rope = self._apply_rope_with_cache_range(q, freqs, h, w, 1, 0, local_current_start, local_current_end, local_ref_tokens, local_per_frame) - k_rope = self._apply_rope_with_cache_range(k, freqs, h, w, 1, 0, local_current_start, local_current_end, local_ref_tokens, local_per_frame) + q_rope = self._apply_rope_with_cache_range(phase, q, freqs, h, w, 1, 0, local_current_start, local_current_end, local_ref_tokens, local_per_frame) + k_rope = self._apply_rope_with_cache_range(phase, k, freqs, h, w, 1, 0, local_current_start, local_current_end, local_ref_tokens, local_per_frame) shard_heads = self.num_heads // sp_world_size h0 = sp_rank * shard_heads h1 = h0 + shard_heads kv_cache.store_kv(k_rope[:, h0:h1], v[:, h0:h1], local_start_idx, local_end_idx, self.block_idx) else: start_frame = self.kv_cache_manager.ref_num_frames + segment_idx * frames - q_rope, k_rope = self._apply_rope_sp(q, k, grid_sizes, freqs, start_frame) + q_rope, k_rope = self._apply_rope_sp(phase, q, k, grid_sizes, freqs, start_frame) use_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False) use_fp4_comm = self.config["parallel"].get("seq_p_fp4_comm", False) k_to_store = all2all_seq2head( diff --git a/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py b/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py index d6893265a..fcaf02d55 100755 --- a/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/causvid/transformer_infer.py @@ -5,7 +5,7 @@ from lightx2v.models.networks.wan.infer.offload.transformer_infer import WanOffloadTransformerInfer from lightx2v.utils.envs import * -from ..utils import apply_rotary_emb, compute_freqs_causvid +from ..utils import compute_freqs_causvid class WanTransformerInferCausVid(WanOffloadTransformerInfer): @@ -109,8 +109,7 @@ def infer_self_attn(self, weights, grid_sizes, embed, x, embed0, seq_lens, freqs # TODO: Implement parallel attention for causvid inference raise NotImplementedError("Parallel attention is not implemented for causvid inference") - q = apply_rotary_emb(q, freqs_i) - k = apply_rotary_emb(k, freqs_i) + q, k = weights.rope.apply(q, k, freqs_i) self.kv_cache[block_idx]["k"][kv_start:kv_end] = k self.kv_cache[block_idx]["v"][kv_start:kv_end] = v diff --git a/lightx2v/models/networks/wan/infer/dreamzero/transformer_infer.py b/lightx2v/models/networks/wan/infer/dreamzero/transformer_infer.py index 1bb94baa2..ea6914005 100644 --- a/lightx2v/models/networks/wan/infer/dreamzero/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/dreamzero/transformer_infer.py @@ -3,11 +3,45 @@ import torch import torch.nn.functional as F +from lightx2v.common.ops.rope import FlashInferRope, RopeTemplate, TorchComplexRope from lightx2v.models.networks.wan.infer.dreamzero.pre_infer import _category_linear from lightx2v.models.networks.wan.infer.transformer_infer import WanTransformerInfer from lightx2v.models.networks.wan.infer.triton_ops import apply_rotary_embedding -from lightx2v.models.networks.wan.infer.utils import apply_rope_with_cos_sin_cache_inplace from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.registry_factory import ROPE_REGISTER + + +@ROPE_REGISTER("dreamzero_rope") +class DreamZeroRope(RopeTemplate): + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("DreamZeroRope only supports interleaved layout.") + self.torch_rope = TorchComplexRope(compute_dtype=compute_dtype) + self.flashinfer_rope = FlashInferRope(layout=layout, compute_dtype=compute_dtype) + + def apply(self, q, k, freqs, **kwargs): + if freqs.shape[0] != q.shape[0]: + raise ValueError(f"DreamZero RoPE length mismatch: freqs={freqs.shape[0]}, q={q.shape[0]}.") + if q.is_cuda and self.flashinfer_rope.is_available(): + cos_sin = torch.cat( + [freqs.real.reshape(freqs.shape[0], -1), freqs.imag.reshape(freqs.shape[0], -1)], + dim=-1, + ).contiguous() + positions = torch.arange(q.shape[0], device=q.device, dtype=torch.long) + return self.flashinfer_rope.apply(q, k, cos_sin, positions=positions) + if q.is_cuda: + cos = freqs.real.reshape(freqs.shape[0], -1).contiguous() + sin = freqs.imag.reshape(freqs.shape[0], -1).contiguous() + return ( + apply_rotary_embedding(q.contiguous(), cos, sin, interleaved=False), + apply_rotary_embedding(k.contiguous(), cos, sin, interleaved=False), + ) + return self.torch_rope.apply(q, k, freqs.to(q.device)) + + def apply_single(self, x, freqs, **kwargs): + output, _ = self.apply(x, x.clone(), freqs, **kwargs) + return output class DreamZeroTransformerInfer(WanTransformerInfer): @@ -24,8 +58,6 @@ def __init__(self, config): self.kv_caches = {} self.cross_attn_kv_caches = {} self._cu_seqlens_cache = {} - self._rope_cache = {} - self.dreamzero_rope_type = config.get("dreamzero_rope_type", config.get("rope_type", "flashinfer")) def create_empty_kv_cache(self, dtype, device): capacity = max(int(self.max_attention_size), 1) @@ -145,7 +177,6 @@ def clear_cache(self, cache_name=None): self.kv_caches.clear() self.cross_attn_kv_caches.clear() self._cu_seqlens_cache.clear() - self._rope_cache.clear() else: cache_names = {cache_name, f"{cache_name}_cond", f"{cache_name}_uncond"} for name in cache_names: @@ -172,34 +203,6 @@ def _get_cu_seqlens(self, length, device): self._cu_seqlens_cache[key] = cached return cached - def _get_rope_cos_sin(self, freqs): - key = ("cos_sin", freqs.data_ptr(), tuple(freqs.shape), str(freqs.device), str(freqs.dtype)) - cached = self._rope_cache.get(key) - if cached is not None: - return cached - cos = freqs.real.reshape(freqs.shape[0], -1).contiguous() - sin = freqs.imag.reshape(freqs.shape[0], -1).contiguous() - self._rope_cache[key] = (cos, sin) - return cos, sin - - def _get_flashinfer_cos_sin(self, freqs): - key = ("flashinfer", freqs.data_ptr(), tuple(freqs.shape), str(freqs.device), str(freqs.dtype)) - cached = self._rope_cache.get(key) - if cached is not None: - return cached - cos, sin = self._get_rope_cos_sin(freqs) - cos_sin = torch.cat([cos, sin], dim=-1).contiguous() - self._rope_cache[key] = cos_sin - return cos_sin - - def _get_rope_positions(self, length, device): - key = ("positions", int(length), str(device)) - cached = self._rope_cache.get(key) - if cached is None: - cached = torch.arange(int(length), device=device, dtype=torch.long) - self._rope_cache[key] = cached - return cached - def _modulate(self, x, scale, shift): if x.is_cuda and x.is_contiguous(): scale_arg = scale @@ -210,41 +213,6 @@ def _modulate(self, x, scale, shift): return self.modulate_func(x, scale=scale_arg, shift=shift_arg).reshape_as(x) return (x.float() * (1.0 + scale.float()) + shift.float()).to(x.dtype) - @staticmethod - def _apply_rope_polar(x, freqs): - x_dtype = x.dtype - seq_len, num_heads, _ = x.shape - x_complex = torch.view_as_complex(x.to(torch.float32).reshape(seq_len, num_heads, -1, 2)) - out = torch.view_as_real(x_complex * freqs.to(x.device)).flatten(2) - return out.to(x_dtype) - - def _apply_rope(self, q, k, freqs): - if freqs.shape[0] != q.shape[0]: - raise ValueError(f"DreamZero RoPE length mismatch: freqs={freqs.shape[0]}, q={q.shape[0]}.") - - if q.is_cuda and self.dreamzero_rope_type == "flashinfer" and apply_rope_with_cos_sin_cache_inplace is not None: - seq_len, num_heads, head_dim = q.shape - query = q.reshape(seq_len, num_heads * head_dim).contiguous() - key = k.reshape(seq_len, num_heads * head_dim).contiguous() - apply_rope_with_cos_sin_cache_inplace( - positions=self._get_rope_positions(seq_len, q.device), - query=query, - key=key, - head_size=head_dim, - cos_sin_cache=self._get_flashinfer_cos_sin(freqs), - is_neox=False, - ) - return query.view(seq_len, num_heads, head_dim), key.view(seq_len, num_heads, head_dim) - - if q.is_cuda and self.dreamzero_rope_type in {"flashinfer", "triton"}: - cos, sin = self._get_rope_cos_sin(freqs) - return ( - apply_rotary_embedding(q.contiguous(), cos, sin, interleaved=False), - apply_rotary_embedding(k.contiguous(), cos, sin, interleaved=False), - ) - - return self._apply_rope_polar(q, freqs), self._apply_rope_polar(k, freqs) - def infer_self_attn_with_cache(self, phase, x, shift_msa, scale_msa, pre_infer_out, kv_cache): seq_total = x.shape[0] norm1_out = phase.norm1.apply(x) @@ -257,7 +225,7 @@ def infer_self_attn_with_cache(self, phase, x, shift_msa, scale_msa, pre_infer_o q = phase.self_attn_norm_q.apply(phase.self_attn_q.apply(norm1_out)).view(seq_total, self.num_heads, self.head_dim) k = phase.self_attn_norm_k.apply(phase.self_attn_k.apply(norm1_out)).view(seq_total, self.num_heads, self.head_dim) v = phase.self_attn_v.apply(norm1_out).view(seq_total, self.num_heads, self.head_dim) - q, k = self._apply_rope(q, k, pre_infer_out.freqs) + q, k = phase.dreamzero_rope.apply(q, k, pre_infer_out.freqs) action_k = action_v = None if pre_infer_out.action_register_length is not None: diff --git a/lightx2v/models/networks/wan/infer/fastwam/transformer_infer.py b/lightx2v/models/networks/wan/infer/fastwam/transformer_infer.py index 33e504cd7..d6b7517d1 100644 --- a/lightx2v/models/networks/wan/infer/fastwam/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/fastwam/transformer_infer.py @@ -6,13 +6,6 @@ def modulate(x, shift, scale): return x * (1 + scale) + shift -def rope_apply(x, freqs): - seq_len, num_heads, head_dim = x.shape - x_complex = torch.view_as_complex(x.to(torch.float64).reshape(seq_len, num_heads, head_dim // 2, 2)) - x_out = torch.view_as_real(x_complex * freqs).flatten(2) - return x_out.to(x.dtype) - - class FastWAMTransformerInfer: def __init__(self, config): self.config = config @@ -46,8 +39,7 @@ def _build_self_attention_io(self, block, x, freqs, t_mod): q = self._reshape_heads(self_attn.norm_q.apply(self_attn.q.apply(attn_input))) k = self._reshape_heads(self_attn.norm_k.apply(self_attn.k.apply(attn_input))) v = self._reshape_heads(self_attn.v.apply(attn_input)) - q = rope_apply(q, freqs) - k = rope_apply(k, freqs) + q, k = self_attn.rope.apply(q, k, freqs) return q, k, v, x, gate_msa, shift_mlp, scale_mlp, gate_mlp def _cross_attn(self, block, x, context, context_mask): diff --git a/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py b/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py index 0b39cf6be..e44609992 100644 --- a/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py @@ -13,12 +13,6 @@ def linear_interpolation(features, seq_len): return output_features.transpose(1, 2) -def rotate_half(x): - x = rearrange(x, "... (d r) -> ... d r", r=2) - x1, x2 = x.unbind(dim=-1) - return rearrange(torch.stack((-x2, x1), dim=-1), "... d r -> ... (d r)") - - def normalize_and_scale(column, source_range, target_range, epsilon=1e-8): source_min, source_max = source_range new_min, new_max = target_range @@ -31,14 +25,13 @@ def __init__(self, head_dim): self.head_dim = head_dim self.base = 10000 - def __call__(self, x, pos_indices): + def __call__(self, rope, x, pos_indices): freqs = 1.0 / (self.base ** (torch.arange(0, self.head_dim, 2, device=pos_indices.device).float() / self.head_dim)) freqs = torch.einsum("..., f -> ... f", pos_indices.float(), freqs) freqs = repeat(freqs, "... n -> ... (n r)", r=2) cos = rearrange(freqs.cos().float(), "n d -> 1 1 n d").to(x.device) sin = rearrange(freqs.sin().float(), "n d -> 1 1 n d").to(x.device) - x_float = x.float() - return (x_float * cos + rotate_half(x_float) * sin).type_as(x) + return rope.apply_single(x, (cos, sin)) class WanInfiniteTalkTransformerInfer(WanOffloadTransformerInfer): @@ -120,7 +113,7 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, pre_infer_out): q = phase.self_attn_norm_q.apply(phase.self_attn_q.apply(norm1_out)).view(s, n, d) k = phase.self_attn_norm_k.apply(phase.self_attn_k.apply(norm1_out)).view(s, n, d) v = phase.self_attn_v.apply(norm1_out).view(s, n, d) - q, k = self.apply_rope_func(q, k, cos_sin) + q, k = phase.rope.apply(q, k, cos_sin) x_ref_attn_map = None ref_target_masks = pre_infer_out.adapter_args.get("ref_target_masks") @@ -251,7 +244,7 @@ def _apply_multi_human_audio_rope(self, q, encoder_k, x_ref_attn_map, grid_t): normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] q_rope = rearrange(q, "t s h d -> 1 h (t s) d") - q_rope = self.rope_1d(q_rope, normalized_pos) + q_rope = self.rope_1d(phase.rope_1d, q_rope, normalized_pos) q = rearrange(q_rope, "1 h (t s) d -> t s h d", t=grid_t) audio_tokens = encoder_k.shape[1] @@ -260,6 +253,6 @@ def _apply_multi_human_audio_rope(self, q, encoder_k, x_ref_attn_map, grid_t): per_frame[audio_tokens // 2 :] = (self.rope_h2[0] + self.rope_h2[1]) / 2 encoder_pos = torch.concat([per_frame] * grid_t, dim=0) k_rope = rearrange(encoder_k, "t s h d -> 1 h (t s) d") - k_rope = self.rope_1d(k_rope, encoder_pos) + k_rope = self.rope_1d(phase.rope_1d, k_rope, encoder_pos) encoder_k = rearrange(k_rope, "1 h (t s) d -> t s h d", t=grid_t) return q, encoder_k diff --git a/lightx2v/models/networks/wan/infer/lingbot_va/transformer_infer.py b/lightx2v/models/networks/wan/infer/lingbot_va/transformer_infer.py index 2c79307b0..dad1329e3 100644 --- a/lightx2v/models/networks/wan/infer/lingbot_va/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/lingbot_va/transformer_infer.py @@ -23,15 +23,6 @@ def _split_modulation(self, modulation, timestep_proj): values = table + timestep_proj return [_token_modulation(item) for item in values.chunk(6, dim=1)] - @staticmethod - def _apply_rotary(x: torch.Tensor, rotary_emb: torch.Tensor) -> torch.Tensor: - x_dtype = x.dtype - seq_len, num_heads, head_dim = x.shape - freqs = rotary_emb.to(x.device) - x_complex = torch.view_as_complex(x.to(torch.float64).reshape(seq_len, num_heads, -1, 2)) - out = torch.view_as_real(x_complex * freqs).flatten(2) - return out.to(x_dtype) - def infer_self_attn_with_kvcache(self, phase, x, shift_msa, scale_msa, rotary_emb, update_cache, cache_name): query_len = x.shape[0] norm1_out = phase.norm1.apply(x).to(x.dtype) @@ -40,8 +31,7 @@ def infer_self_attn_with_kvcache(self, phase, x, shift_msa, scale_msa, rotary_em q = phase.self_attn_norm_q.apply(phase.self_attn_q.apply(norm1_out)).view(query_len, self.num_heads, self.head_dim) k = phase.self_attn_norm_k.apply(phase.self_attn_k.apply(norm1_out)).view(query_len, self.num_heads, self.head_dim) v = phase.self_attn_v.apply(norm1_out).view(query_len, self.num_heads, self.head_dim) - q = self._apply_rotary(q, rotary_emb) - k = self._apply_rotary(k, rotary_emb) + q, k = phase.lingbot_rope.apply(q, k, rotary_emb.to(q.device)) cache = self.kv_cache_manager.get_self_attn_kv_cache(cache_name) if self.kv_cache_manager is not None else None slots = None diff --git a/lightx2v/models/networks/wan/infer/matrix_game2/posemb_layers.py b/lightx2v/models/networks/wan/infer/matrix_game2/posemb_layers.py index 549ea44da..f23c86ff4 100644 --- a/lightx2v/models/networks/wan/infer/matrix_game2/posemb_layers.py +++ b/lightx2v/models/networks/wan/infer/matrix_game2/posemb_layers.py @@ -126,57 +126,32 @@ def reshape_for_broadcast( return freqs_cis.view(*shape) -def rotate_half(x): - x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2] - return torch.stack([-x_imag, x_real], dim=-1).flatten(3) - - def apply_rotary_emb( + rope, xq: torch.Tensor, xk: torch.Tensor, freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], head_first: bool = False, start_offset: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Apply rotary embeddings to input tensors using the given frequency tensor. - - This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided - frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor - is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are - returned as real tensors. - - Args: - xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D] - xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D] - freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential. - head_first (bool): head dimension first (except batch dim) or not. - - Returns: - Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings. - - """ - # print(freqs_cis[0].shape, xq.shape, xk.shape) - xk_out = None - assert isinstance(freqs_cis, tuple) if isinstance(freqs_cis, tuple): - cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D] - cos, sin = cos.to(xq.device), sin.to(xq.device) - # real * cos - imag * sin - # imag * cos + real * sin - xq_out = (xq.float() * cos[:, start_offset : start_offset + xq.shape[1], :, :] + rotate_half(xq.float()) * sin[:, start_offset : start_offset + xq.shape[1], :, :]).type_as(xq) - xk_out = (xk.float() * cos[:, start_offset : start_offset + xk.shape[1], :, :] + rotate_half(xk.float()) * sin[:, start_offset : start_offset + xk.shape[1], :, :]).type_as(xk) + cos, sin = freqs_cis + if head_first: + shape = [1] * xq.ndim + shape[-2], shape[-1] = cos.shape + cos, sin = cos.view(*shape), sin.view(*shape) + else: + cos = cos[start_offset : start_offset + xq.shape[1]].view(1, xq.shape[1], 1, cos.shape[-1]) + sin = sin[start_offset : start_offset + xq.shape[1]].view(1, xq.shape[1], 1, sin.shape[-1]) + return rope.apply(xq, xk, (cos.to(xq.device), sin.to(xq.device))) + + if head_first: + shape = [1] * xq.ndim + shape[-2], shape[-1] = freqs_cis.shape + freqs_cis = freqs_cis.view(*shape) else: - # view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex) - xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2] - freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2] - # (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin) - # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq) - xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # [B, S, H, D//2] - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk) - - return xq_out, xk_out + freqs_cis = freqs_cis.view(1, freqs_cis.shape[0], 1, freqs_cis.shape[1]) + return rope.apply(xq, xk, freqs_cis.to(xq.device)) def get_nd_rotary_pos_embed( diff --git a/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py b/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py index 48e4ca27a..82a8a9764 100755 --- a/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py @@ -176,7 +176,7 @@ def _infer_action_model_with_kvcache( q = q0 * torch.rsqrt(q0.pow(2).mean(dim=-1, keepdim=True) + 1e-6) k = k0 * torch.rsqrt(k0.pow(2).mean(dim=-1, keepdim=True) + 1e-6) - q, k = apply_rotary_emb(q, k, freqs_cis, start_offset=start_frame, head_first=False) + q, k = apply_rotary_emb(phase.rope, q, k, freqs_cis, start_offset=start_frame, head_first=False) if is_causal: action_mouse_kv_cache = self.kv_cache_manager.action_mouse_kv_cache @@ -250,7 +250,7 @@ def _infer_action_model_with_kvcache( B, TS, H, D = q.shape T_ = TS // S q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D) - q, k = apply_rotary_emb(q, k, freqs_cis, start_offset=start_frame, head_first=False) + q, k = apply_rotary_emb(phase.rope, q, k, freqs_cis, start_offset=start_frame, head_first=False) k = k.expand(S, k.shape[1], k.shape[2], k.shape[3]) v = v.expand(S, v.shape[1], v.shape[2], v.shape[3]) diff --git a/lightx2v/models/networks/wan/infer/matrix_game3/transformer_infer.py b/lightx2v/models/networks/wan/infer/matrix_game3/transformer_infer.py index 2cebd34f4..8ef324305 100644 --- a/lightx2v/models/networks/wan/infer/matrix_game3/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/matrix_game3/transformer_infer.py @@ -32,7 +32,7 @@ torch_device_module = getattr(torch, AI_DEVICE) -def rope_apply_with_indices(x, grid_sizes, freqs, indices): +def rope_apply_with_indices(rope, x, grid_sizes, freqs, indices): """Apply RoPE using explicit frame indices (for memory-aware attention). Rather than assuming sequential frame positions 0..F-1, this uses the @@ -85,11 +85,7 @@ def rope_apply_with_indices(x, grid_sizes, freqs, indices): else: raise ValueError(f"Unexpected freqs shape: {freqs.shape}") - cos_sin = cos_sin.to(x.device) - # Apply RoPE - x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - out = torch.view_as_real(x_complex * cos_sin).flatten(-2) - return out.type_as(x) + return rope.apply_single(x, cos_sin.to(x.device), unsqueeze_dim=0) class WanMtxg3TransformerInfer(WanTransformerInfer): @@ -301,8 +297,8 @@ def _infer_self_attn_mg3(self, phase, x, shift_msa, scale_msa, pre_infer_out): # RoPE with explicit indices mem_indices = memory_latent_idx if memory_latent_idx is not None else list(range(memory_length)) - q_memory = rope_apply_with_indices(q_memory, grid_sizes_mem, self.freqs, mem_indices) - k_memory = rope_apply_with_indices(k_memory, grid_sizes_mem, self.freqs, mem_indices) + q_memory = rope_apply_with_indices(phase.indexed_rope, q_memory, grid_sizes_mem, self.freqs, mem_indices) + k_memory = rope_apply_with_indices(phase.indexed_rope, k_memory, grid_sizes_mem, self.freqs, mem_indices) if predict_latent_idx is not None: if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2: @@ -312,8 +308,8 @@ def _infer_self_attn_mg3(self, phase, x, shift_msa, scale_msa, pre_infer_out): else: pred_indices = list(range(grid_sizes_pred[0, 0].item())) - q_pred = rope_apply_with_indices(q_pred, grid_sizes_pred, self.freqs, pred_indices) - k_pred = rope_apply_with_indices(k_pred, grid_sizes_pred, self.freqs, pred_indices) + q_pred = rope_apply_with_indices(phase.indexed_rope, q_pred, grid_sizes_pred, self.freqs, pred_indices) + k_pred = rope_apply_with_indices(phase.indexed_rope, k_pred, grid_sizes_pred, self.freqs, pred_indices) q = torch.cat([q_memory.squeeze(0), q_pred.squeeze(0)], dim=0) k = torch.cat([k_memory.squeeze(0), k_pred.squeeze(0)], dim=0) @@ -333,8 +329,8 @@ def _infer_self_attn_mg3(self, phase, x, shift_msa, scale_msa, pre_infer_out): pred_indices = predict_latent_idx else: pred_indices = list(range(grid_sizes.tuple[0])) - q = rope_apply_with_indices(q_unsq, grid_sizes_t, self.freqs, pred_indices).squeeze(0) - k = rope_apply_with_indices(k_unsq, grid_sizes_t, self.freqs, pred_indices).squeeze(0) + q = rope_apply_with_indices(phase.indexed_rope, q_unsq, grid_sizes_t, self.freqs, pred_indices).squeeze(0) + k = rope_apply_with_indices(phase.indexed_rope, k_unsq, grid_sizes_t, self.freqs, pred_indices).squeeze(0) img_qkv_len = q.shape[0] @@ -436,19 +432,19 @@ def _infer_action_module(self, phase, x, pre_infer_out): if memory_length > 0: freqs_memory = self._get_action_rotary_pos_embed(memory_length, mouse_head_dim, self.mouse_qk_dim_list) - q_mem, k_mem = apply_rotary_emb(q_m[:, :memory_length], k_m[:, :memory_length], freqs_memory, head_first=False) + q_mem, k_mem = apply_rotary_emb(phase.rope, q_m[:, :memory_length], k_m[:, :memory_length], freqs_memory, head_first=False) q_m[:, :memory_length] = q_mem k_m[:, :memory_length] = k_mem pred_length = tt - memory_length if pred_length > 0: freqs_pred = self._get_action_rotary_pos_embed(pred_length, mouse_head_dim, self.mouse_qk_dim_list) - q_pred, k_pred = apply_rotary_emb(q_m[:, memory_length:], k_m[:, memory_length:], freqs_pred, head_first=False) + q_pred, k_pred = apply_rotary_emb(phase.rope, q_m[:, memory_length:], k_m[:, memory_length:], freqs_pred, head_first=False) q_m[:, memory_length:] = q_pred k_m[:, memory_length:] = k_pred else: freqs = self._get_action_rotary_pos_embed(tt, mouse_head_dim, self.mouse_qk_dim_list) - q_m, k_m = apply_rotary_emb(q_m, k_m, freqs, head_first=False) + q_m, k_m = apply_rotary_emb(phase.rope, q_m, k_m, freqs, head_first=False) mouse_attn = self._run_flash_attention(q_m, k_m, v_m, causal=False) mouse_attn = rearrange(mouse_attn, "(b s) t h d -> b (t s) (h d)", b=batch_size, s=spatial_tokens) @@ -512,19 +508,19 @@ def _infer_action_module(self, phase, x, pre_infer_out): q_k = rearrange(q_k, "b (t s) h d -> (b s) t h d", s=spatial_tokens) if memory_length > 0: freqs_memory = self._get_action_rotary_pos_embed(memory_length, keyboard_head_dim, self.mouse_qk_dim_list) - q_mem, k_mem = apply_rotary_emb(q_k[:, :memory_length], k_k[:, :memory_length], freqs_memory, head_first=False) + q_mem, k_mem = apply_rotary_emb(phase.rope, q_k[:, :memory_length], k_k[:, :memory_length], freqs_memory, head_first=False) q_k[:, :memory_length] = q_mem k_k[:, :memory_length] = k_mem pred_length = tt - memory_length if pred_length > 0: freqs_pred = self._get_action_rotary_pos_embed(pred_length, keyboard_head_dim, self.mouse_qk_dim_list) - q_pred, k_pred = apply_rotary_emb(q_k[:, memory_length:], k_k[:, memory_length:], freqs_pred, head_first=False) + q_pred, k_pred = apply_rotary_emb(phase.rope, q_k[:, memory_length:], k_k[:, memory_length:], freqs_pred, head_first=False) q_k[:, memory_length:] = q_pred k_k[:, memory_length:] = k_pred else: freqs = self._get_action_rotary_pos_embed(tt, keyboard_head_dim, self.rope_dim_list) - q_k, k_k = apply_rotary_emb(q_k, k_k, freqs, head_first=False) + q_k, k_k = apply_rotary_emb(phase.rope, q_k, k_k, freqs, head_first=False) k_k = k_k.repeat(spatial_tokens, 1, 1, 1) v_k = v_k.repeat(spatial_tokens, 1, 1, 1) diff --git a/lightx2v/models/networks/wan/infer/pre_infer.py b/lightx2v/models/networks/wan/infer/pre_infer.py index dabdbf2a6..07d9f85c6 100755 --- a/lightx2v/models/networks/wan/infer/pre_infer.py +++ b/lightx2v/models/networks/wan/infer/pre_infer.py @@ -65,30 +65,15 @@ def prepare_cos_sin(self, grid_sizes, freqs): ], dim=-1, ) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_sin = cos_sin.reshape(seq_len, -1) - # Extract cos and sin parts separately and concatenate - cos_half = cos_sin.real.contiguous() - sin_half = cos_sin.imag.contiguous() - cos_sin = torch.cat([cos_half, sin_half], dim=-1) - if self.seq_p_group is not None: - world_size = dist.get_world_size(self.seq_p_group) - cur_rank = dist.get_rank(self.seq_p_group) - seqlen = cos_sin.shape[0] - padding_size = (world_size - (seqlen % world_size)) % world_size - if padding_size > 0: - cos_sin = F.pad(cos_sin, (0, 0, 0, padding_size)) - cos_sin = torch.chunk(cos_sin, world_size, dim=0)[cur_rank] - else: - cos_sin = cos_sin.reshape(seq_len, 1, -1) - if self.seq_p_group is not None: - world_size = dist.get_world_size(self.seq_p_group) - cur_rank = dist.get_rank(self.seq_p_group) - seqlen = cos_sin.shape[0] - padding_size = (world_size - (seqlen % world_size)) % world_size - if padding_size > 0: - cos_sin = F.pad(cos_sin, (0, 0, 0, 0, 0, padding_size)) - cos_sin = torch.chunk(cos_sin, world_size, dim=0)[cur_rank] + cos_sin = cos_sin.reshape(seq_len, 1, -1) + if self.seq_p_group is not None: + world_size = dist.get_world_size(self.seq_p_group) + cur_rank = dist.get_rank(self.seq_p_group) + seqlen = cos_sin.shape[0] + padding_size = (world_size - (seqlen % world_size)) % world_size + if padding_size > 0: + cos_sin = F.pad(cos_sin, (0, 0, 0, 0, 0, padding_size)) + cos_sin = torch.chunk(cos_sin, world_size, dim=0)[cur_rank] return cos_sin @torch.no_grad() diff --git a/lightx2v/models/networks/wan/infer/s2v/rope.py b/lightx2v/models/networks/wan/infer/s2v/rope.py index 162f31844..deea3e522 100644 --- a/lightx2v/models/networks/wan/infer/s2v/rope.py +++ b/lightx2v/models/networks/wan/infer/s2v/rope.py @@ -56,23 +56,3 @@ def rope_precompute(x, grid_sizes, freqs, start=None): output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i seq_bucket.append(seq_bucket[-1] + seq_len) return output - - -@torch.amp.autocast("cuda", enabled=False) -def rope_apply(x, precomputed_freqs): - """Match Wan2.2/wan/modules/s2v/model_s2v.py rope_apply for precomputed freqs.""" - output = [] - for i in range(x.size(0)): - s = x.size(1) - x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, x.size(2), -1, 2)) - freqs_i = precomputed_freqs[i, :s] - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - if s < x.size(1): - x_i = torch.cat([x_i, x[i, s:]]) - output.append(x_i) - return torch.stack(output).float() - - -@torch.amp.autocast("cuda", enabled=False) -def apply_precomputed_rope(x, precomputed_freqs): - return rope_apply(x, precomputed_freqs) diff --git a/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py b/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py index affe80ec5..5878558a2 100644 --- a/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/s2v/transformer_infer.py @@ -2,7 +2,6 @@ import torch.distributed as dist from lightx2v.models.networks.wan.infer.s2v.audio_inject import apply_audio_inject -from lightx2v.models.networks.wan.infer.s2v.rope import apply_precomputed_rope from lightx2v.models.networks.wan.infer.s2v.wan_ops import ( cross_attn_forward, mm_weight_autocast_nd, @@ -95,8 +94,8 @@ def _s2v_self_attn(self, phase0, norm_x, seq_lens, freqs): q = wan_rms_norm(phase0.self_attn_norm_q, mm_weight_autocast_nd(phase0.self_attn_q, norm_x)).view(b, s, n, d) k = wan_rms_norm(phase0.self_attn_norm_k, mm_weight_autocast_nd(phase0.self_attn_k, norm_x)).view(b, s, n, d) v = mm_weight_autocast_nd(phase0.self_attn_v, norm_x).view(b, s, n, d) - q = apply_precomputed_rope(q, freqs) - k = apply_precomputed_rope(k, freqs) + q, k = phase0.s2v_rope.apply(q, k, freqs) + q, k = q.float(), k.float() attn_out = ( phase0.self_attn_1_parallel.apply( q=q.squeeze(0).to(self.infer_dtype), @@ -118,7 +117,7 @@ def _s2v_self_attn(self, phase0, norm_x, seq_lens, freqs): ) return mm_weight_autocast_nd(phase0.self_attn_o, attn_out) - return s2v_self_attn_forward(phase0, norm_x, seq_lens, freqs, self.num_heads, self.head_dim, apply_precomputed_rope) + return s2v_self_attn_forward(phase0, norm_x, seq_lens, freqs, self.num_heads, self.head_dim) @torch.no_grad() def infer_s2v_block(self, block, x, pre_infer_out): diff --git a/lightx2v/models/networks/wan/infer/s2v/wan_ops.py b/lightx2v/models/networks/wan/infer/s2v/wan_ops.py index 0602abdce..9746f5619 100644 --- a/lightx2v/models/networks/wan/infer/s2v/wan_ops.py +++ b/lightx2v/models/networks/wan/infer/s2v/wan_ops.py @@ -141,7 +141,7 @@ def half(x): return x.unflatten(0, (b, lq)).type(out_dtype) -def s2v_self_attn_forward(phase0, norm_x, seq_lens, freqs, num_heads, head_dim, rope_apply_fn): +def s2v_self_attn_forward(phase0, norm_x, seq_lens, freqs, num_heads, head_dim): """Mirror of WanS2VSelfAttention.forward (model_s2v.py + model.py).""" b, s, n, d = norm_x.size(0), norm_x.size(1), num_heads, head_dim @@ -149,9 +149,10 @@ def s2v_self_attn_forward(phase0, norm_x, seq_lens, freqs, num_heads, head_dim, k = wan_rms_norm(phase0.self_attn_norm_k, mm_weight_autocast_nd(phase0.self_attn_k, norm_x)).view(b, s, n, d) v = mm_weight_autocast_nd(phase0.self_attn_v, norm_x).view(b, s, n, d) + q, k = phase0.s2v_rope.apply(q, k, freqs) attn = flash_attention( - q=rope_apply_fn(q, freqs), - k=rope_apply_fn(k, freqs), + q=q.float(), + k=k.float(), v=v, k_lens=seq_lens, ) diff --git a/lightx2v/models/networks/wan/infer/self_forcing/transformer_infer.py b/lightx2v/models/networks/wan/infer/self_forcing/transformer_infer.py index b8b75ec14..943f42354 100755 --- a/lightx2v/models/networks/wan/infer/self_forcing/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/self_forcing/transformer_infer.py @@ -6,8 +6,6 @@ from lightx2v.common.offload.manager import WeightAsyncStreamManager from lightx2v.common.ops.attn.utils.all2all import all2all_seq2head from lightx2v.models.networks.wan.infer.transformer_infer import WanTransformerInfer -from lightx2v.models.networks.wan.infer.triton_ops import causal_rope_apply_triton -from lightx2v.models.networks.wan.infer.utils import causal_rope_apply from lightx2v_platform.base.global_var import AI_DEVICE torch_device_module = getattr(torch, AI_DEVICE) @@ -49,18 +47,13 @@ def __init__(self, config): else: self.infer_func = self.infer_with_kvcache - if self.config.get("causal_rope_type", "torch") == "triton": - self.causal_rope_apply_func = causal_rope_apply_triton - else: - self.causal_rope_apply_func = causal_rope_apply - def _calculate_q_k_len(self, q, k_lens): q_lens = torch.tensor([q.size(0)], dtype=torch.int32) cu_seqlens_q = torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32) cu_seqlens_k = torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32) return cu_seqlens_q, cu_seqlens_k - def _apply_rope_sp(self, q, k, grid_sizes, freqs, start_frame): + def _apply_rope_sp(self, phase, q, k, grid_sizes, freqs, start_frame): f, h, w = grid_sizes[0].tolist() full_seq_len = f * h * w c = q.size(-1) // 2 @@ -82,13 +75,7 @@ def _apply_rope_sp(self, q, k, grid_sizes, freqs, start_frame): pos_freqs = F.pad(pos_freqs, (0, 0, 0, 0, 0, padding_size)) pos_freqs = torch.chunk(pos_freqs, world_size, dim=0)[cur_rank][: q.size(0)] - n = q.size(1) - q_c = torch.view_as_complex(q.float().reshape(q.size(0), n, -1, 2)) - k_c = torch.view_as_complex(k.float().reshape(k.size(0), n, -1, 2)) - pos_freqs = pos_freqs.to(torch.complex64) - q = torch.view_as_real(q_c * pos_freqs).flatten(2).type_as(q) - k = torch.view_as_real(k_c * pos_freqs).flatten(2).type_as(k) - return q, k + return phase.causal_rope.apply(q, k, pos_freqs.to(torch.complex64)) def infer_with_kvcache(self, blocks, x, pre_infer_out): """Run all transformer blocks with the rolling self-attention KV cache.""" @@ -223,10 +210,16 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh current_start_frame = seg_index * self.num_frame_per_chunk if self.config.get("seq_parallel", False): - q, k = self._apply_rope_sp(q, k, grid_sizes, freqs, current_start_frame) + q, k = self._apply_rope_sp(phase, q, k, grid_sizes, freqs, current_start_frame) else: - q = self.causal_rope_apply_func(q.unsqueeze(0), grid_sizes, freqs, start_frame=current_start_frame).type_as(v)[0] - k = self.causal_rope_apply_func(k.unsqueeze(0), grid_sizes, freqs, start_frame=current_start_frame).type_as(v)[0] + q, k = phase.causal_rope.apply( + q.unsqueeze(0), + k.unsqueeze(0), + freqs, + grid_sizes=grid_sizes, + start_frame=current_start_frame, + ) + q, k = q.type_as(v)[0], k.type_as(v)[0] kv_cache = self.kv_cache_manager.self_attn_kv_cache seq_parallel = self.config.get("seq_parallel", False) diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py index 6fa4614d5..51225f381 100755 --- a/lightx2v/models/networks/wan/infer/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/transformer_infer.py @@ -1,5 +1,3 @@ -from functools import partial - import torch from loguru import logger @@ -10,7 +8,6 @@ from .mxfp8_fuse import WanMxfp8FuseMixin, scaled_mxfp8_modulate_quant from .triton_ops import fuse_scale_shift_kernel -from .utils import apply_wan_rope_with_chunk, apply_wan_rope_with_flashinfer, apply_wan_rope_with_torch, apply_wan_rope_with_torch_naive torch_device_module = getattr(torch, AI_DEVICE) @@ -34,29 +31,6 @@ def __init__(self, config): self.modulate_func = fuse_scale_shift_kernel else: self.modulate_func = modulate - rope_funcs = { - "flashinfer": apply_wan_rope_with_flashinfer, - "torch": apply_wan_rope_with_torch, - "torch_naive": apply_wan_rope_with_torch_naive, - } - rope_type = self.config.get("rope_type", "flashinfer") - # Try to get rope function from registry first (for platform-specific implementations) - if rope_type in ROPE_REGISTER: - rope_class = ROPE_REGISTER[rope_type] - self.rope_instance = rope_class() - - # Create a wrapper function that matches the expected signature - def rope_wrapper(xq, xk, cos_sin_cache): - return self.rope_instance.apply(xq, xk, cos_sin_cache) - - rope_func = rope_wrapper - else: - # Fallback to hardcoded functions - rope_func = rope_funcs.get(rope_type, apply_wan_rope_with_torch) - if self.config.get("rope_chunk", False): - self.apply_rope_func = partial(apply_wan_rope_with_chunk, chunk_size=self.config.get("rope_chunk_size", 100), rope_func=rope_func) - else: - self.apply_rope_func = rope_func self.clean_cuda_cache = self.config.get("clean_cuda_cache", False) self.mxfp8_fuse_enable = self.config.get("mxfp8_fuse_enable", True) self.infer_dtype = GET_DTYPE() @@ -254,7 +228,7 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None): q = phase.self_attn_norm_q.apply(phase.self_attn_q.apply(norm1_out)).view(s, n, d) k = phase.self_attn_norm_k.apply(phase.self_attn_k.apply(norm1_out)).view(s, n, d) v = phase.self_attn_v.apply(norm1_out).view(s, n, d) - q, k = self.apply_rope_func(q, k, cos_sin) + q, k = phase.rope.apply(q, k, cos_sin) img_qkv_len = q.shape[0] if self.clean_cuda_cache: diff --git a/lightx2v/models/networks/wan/infer/utils.py b/lightx2v/models/networks/wan/infer/utils.py index 1168f1a6b..4b8b9abd3 100755 --- a/lightx2v/models/networks/wan/infer/utils.py +++ b/lightx2v/models/networks/wan/infer/utils.py @@ -1,166 +1,112 @@ import torch import torch.distributed as dist -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - +from lightx2v.common.ops.rope import RopeTemplate, TorchComplexRope from lightx2v.utils.envs import * - - -# FlashInfer's Python wrapper performs JIT/cache bookkeeping that Dynamo cannot -# trace. Keep it behind an opaque, in-place torch operator while compiling. -@torch.library.custom_op( - "lightx2v::wan_rope_flashinfer_", - mutates_args=("query", "key"), - device_types="cuda", - schema=("(Tensor positions, Tensor(a!) query, Tensor(b!) key, Tensor cos_sin_cache, int head_size) -> ()"), -) -def wan_rope_flashinfer_( - positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - cos_sin_cache: torch.Tensor, - head_size: int, -) -> None: - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=head_size, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - -@wan_rope_flashinfer_.register_fake -def wan_rope_flashinfer_fake( - positions: torch.Tensor, - query: torch.Tensor, - key: torch.Tensor, - cos_sin_cache: torch.Tensor, - head_size: int, -) -> None: - return None - - -def apply_wan_rope_with_torch( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - n = xq.size(1) - seq_len = cos_sin_cache.size(0) - - xq = torch.view_as_complex(xq[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2)) - xk = torch.view_as_complex(xk[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2)) - # Apply rotary embedding - xq = torch.view_as_real(xq * cos_sin_cache).flatten(2) - xk = torch.view_as_real(xk * cos_sin_cache).flatten(2) - xq = torch.cat([xq, xq[seq_len:]]) - xk = torch.cat([xk, xk[seq_len:]]) - - return xq.to(GET_DTYPE()), xk.to(GET_DTYPE()) - - -def apply_wan_rope_with_torch_naive( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - n = xq.size(1) - seq_len = cos_sin_cache.size(0) - total_len = xq.size(0) - - cos = cos_sin_cache.real.expand(-1, n, -1) - sin = cos_sin_cache.imag.expand(-1, n, -1) - - def _rotate_part(x_part: torch.Tensor): - x_even = x_part[..., 0::2] - x_odd = x_part[..., 1::2] - - x_rot_even = x_even * cos - x_odd * sin - x_rot_odd = x_even * sin + x_odd * cos - - x_part[..., 0::2] = x_rot_even - x_part[..., 1::2] = x_rot_odd - - # Q - xq_part = xq[:seq_len] - _rotate_part(xq_part) - - # K - xk_part = xk[:seq_len] - _rotate_part(xk_part) - - if seq_len < total_len: - xq_out = torch.cat([xq_part, xq[seq_len:]], dim=0) - xk_out = torch.cat([xk_part, xk[seq_len:]], dim=0) - else: - xq_out = xq_part - xk_out = xk_part - - return xq_out.to(GET_DTYPE()), xk_out.to(GET_DTYPE()) - - -def apply_wan_rope_with_chunk( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, - chunk_size: int, - rope_func, -): - seq_len = cos_sin_cache.size(0) - x_q = torch.empty_like(xq) - x_k = torch.empty_like(xk) - - for start in range(0, seq_len, chunk_size): - end = min(start + chunk_size, seq_len) - xq_chunk = xq[start:end] - xk_chunk = xk[start:end] - cos_sin_chunk = cos_sin_cache[start:end] - xq_chunk_out, xk_chunk_out = rope_func(xq_chunk, xk_chunk, cos_sin_chunk) - x_q[start:end].copy_(xq_chunk_out, non_blocking=True) - x_k[start:end].copy_(xk_chunk_out, non_blocking=True) - del xq_chunk_out, xk_chunk_out - - target_dtype = GET_DTYPE() - if x_q.dtype != target_dtype: - x_q = x_q.to(target_dtype) - if x_k.dtype != target_dtype: - x_k = x_k.to(target_dtype) - - return x_q, x_k - - -def apply_wan_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous() - key = xk.reshape(L, H * D).contiguous() - - positions = torch.arange(L, device=xq.device, dtype=torch.long) - - if torch.compiler.is_compiling(): - wan_rope_flashinfer_(positions, query, key, cos_sin_cache, D) - else: - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, +from lightx2v.utils.registry_factory import ROPE_REGISTER + + +@ROPE_REGISTER("wan_causal_rope") +class WanCausalRope(RopeTemplate): + def __init__(self, layout="interleaved", compute_dtype=torch.float64): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("WanCausalRope only supports interleaved layout.") + self.torch_rope = TorchComplexRope(compute_dtype=compute_dtype) + + def apply(self, q, k, freqs, grid_sizes=None, start_frame=0, **kwargs): + if grid_sizes is None: + return self.torch_rope.apply(q, k, freqs, **kwargs) + return ( + self.apply_single(q, freqs, grid_sizes=grid_sizes, start_frame=start_frame), + self.apply_single(k, freqs, grid_sizes=grid_sizes, start_frame=start_frame), ) - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out + def apply_single(self, x, freqs, grid_sizes=None, start_frame=0, **kwargs): + if grid_sizes is None: + return self.torch_rope.apply_single(x, freqs, **kwargs) + if x.is_cuda: + from lightx2v.models.networks.wan.infer.triton_ops import causal_rope_apply_triton + + return causal_rope_apply_triton(x, grid_sizes, freqs, start_frame=start_frame) + c = x.size(3) // 2 + freq_parts = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + output = [] + for i, (frames, height, width) in enumerate(grid_sizes.tolist()): + seq_len = frames * height * width + pos_freqs = torch.cat( + [ + freq_parts[0][start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1), + freq_parts[1][:height].view(1, height, 1, -1).expand(frames, height, width, -1), + freq_parts[2][:width].view(1, 1, width, -1).expand(frames, height, width, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + output.append(self.torch_rope.apply_single(x[i], pos_freqs)) + return torch.stack(output).type_as(x) + + def apply_audio_cache( + self, + x, + freqs, + *, + h, + w, + token_start, + token_end, + ref_tokens, + local_per_frame, + world_size=1, + rank=0, + global_end=None, + sink_tokens=0, + ): + if x.is_cuda: + from lightx2v.models.networks.wan.infer.triton_ops import apply_audio_cache_rope + + return apply_audio_cache_rope( + x, + freqs, + h=h, + w=w, + token_start=token_start, + ref_tokens=ref_tokens, + local_per_frame=local_per_frame, + world_size=world_size, + rank=rank, + global_end=global_end, + sink_tokens=sink_tokens, + ).to(x.dtype) + + c = x.size(-1) // 2 + temporal_dim = c - 2 * (c // 3) + freq_parts = freqs.split([temporal_dim, c // 3, c // 3], dim=1) + spatial = torch.cat( + [ + freq_parts[1][:h].view(h, 1, -1).expand(h, w, -1), + freq_parts[2][:w].view(1, w, -1).expand(h, w, -1), + ], + dim=-1, + ).reshape(h * w, -1) + if world_size > 1: + padding = (world_size - spatial.size(0) % world_size) % world_size + if padding: + spatial = torch.cat([spatial, torch.ones(padding, spatial.size(1), dtype=spatial.dtype, device=spatial.device)], dim=0) + spatial = torch.chunk(spatial, world_size, dim=0)[rank][:local_per_frame] + + positions = torch.arange(token_start, token_end, device=freqs.device, dtype=torch.long) + if global_end is not None: + recent_local = max(0, token_end - int(sink_tokens)) + recent_start = int(global_end) - recent_local + recent = recent_start + positions - int(sink_tokens) + positions = torch.where(positions < int(sink_tokens), positions, recent) + is_ref = positions < ref_tokens + gen_idx = torch.clamp(positions - ref_tokens, min=0) + ref_frames = ref_tokens // local_per_frame + frame_idx = torch.where(is_ref, positions // local_per_frame, ref_frames + gen_idx // local_per_frame) + spatial_idx = torch.where(is_ref, positions % local_per_frame, gen_idx % local_per_frame) + pos_freqs = torch.cat([freq_parts[0][frame_idx], spatial[spatial_idx]], dim=-1).unsqueeze(1) + return self.torch_rope.apply_single(x, pos_freqs.to(torch.complex64)).to(x.dtype) def compute_freqs(c, grid_sizes, freqs): @@ -224,50 +170,6 @@ def pad_freqs(original_tensor, target_len): return padded_tensor -def apply_rotary_emb(x, freqs_i): - n = x.size(1) - seq_len = freqs_i.size(0) - - x_i = torch.view_as_complex(x[:seq_len].to(torch.float32).reshape(seq_len, n, -1, 2)) - # Apply rotary embedding - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[seq_len:]]) - return x_i.to(GET_DTYPE()) - - -def apply_rotary_emb_chunk(x, freqs_i, chunk_size, remaining_chunk_size=100): - n = x.size(1) - seq_len = freqs_i.size(0) - - output_chunks = [] - for start in range(0, seq_len, chunk_size): - end = min(start + chunk_size, seq_len) - x_chunk = x[start:end] - freqs_chunk = freqs_i[start:end] - - x_chunk_complex = torch.view_as_complex(x_chunk.to(torch.float32).reshape(end - start, n, -1, 2)) - x_chunk_embedded = torch.view_as_real(x_chunk_complex * freqs_chunk).flatten(2).to(GET_DTYPE()) - output_chunks.append(x_chunk_embedded) - del x_chunk_complex, x_chunk_embedded - torch.cuda.empty_cache() - - result = [] - for chunk in output_chunks: - result.append(chunk) - del output_chunks - torch.cuda.empty_cache() - - for start in range(seq_len, x.size(0), remaining_chunk_size): - end = min(start + remaining_chunk_size, x.size(0)) - result.append(x[start:end]) - - x_i = torch.cat(result, dim=0) - del result - torch.cuda.empty_cache() - - return x_i.to(GET_DTYPE()) - - def rope_params(max_seq_len, dim, theta=10000): assert dim % 2 == 0 freqs = torch.outer( @@ -316,21 +218,3 @@ def guidance_scale_embedding(w, embedding_dim=256, cfg_range=(1.0, 6.0), target_ emb = torch.nn.functional.pad(emb, (0, 1).to(w.device)) assert emb.shape == (w.shape[0], embedding_dim) return emb - - -def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): - n, c = x.size(2), x.size(3) // 2 - freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) - output = [] - - for i, (f, h, w) in enumerate(grid_sizes.tolist()): - seq_len = f * h * w - x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)) - freqs_i = torch.cat( - [freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)], - dim=-1, - ).reshape(seq_len, 1, -1) - x_i = torch.view_as_real(x_i * freqs_i).flatten(2) - x_i = torch.cat([x_i, x[i, seq_len:]]) - output.append(x_i) - return torch.stack(output).type_as(x) diff --git a/lightx2v/models/networks/wan/weights/audio/transformer_weights.py b/lightx2v/models/networks/wan/weights/audio/transformer_weights.py index 528067880..d91dd62dc 100755 --- a/lightx2v/models/networks/wan/weights/audio/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/audio/transformer_weights.py @@ -1,8 +1,11 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights from lightx2v.utils.registry_factory import ( LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, + ROPE_REGISTER, TENSOR_REGISTER, ) @@ -10,6 +13,12 @@ class WanAudioTransformerWeights(WanTransformerWeights): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__(config, lazy_load_path, lora_path) + for phase in self.iter_self_attention_phases(): + if not hasattr(phase, "causal_rope"): + phase.add_module( + "causal_rope", + ROPE_REGISTER[config.get("causal_rope_type", "wan_causal_rope")](layout="interleaved", compute_dtype=torch.float64), + ) for i in range(self.blocks_num): self.blocks[i].compute_phases.append( WanAudioAdapterCA( diff --git a/lightx2v/models/networks/wan/weights/dreamzero/transformer_weights.py b/lightx2v/models/networks/wan/weights/dreamzero/transformer_weights.py index 3d89e23c6..8673a66e9 100644 --- a/lightx2v/models/networks/wan/weights/dreamzero/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/dreamzero/transformer_weights.py @@ -2,9 +2,10 @@ from lightx2v.common.modules.weight_module import WeightModule from lightx2v.common.ops.attn.flash_attn import flash_attn_varlen_func_v2, flash_attn_varlen_func_v3 +from lightx2v.models.networks.wan.infer.dreamzero.transformer_infer import DreamZeroRope # noqa: F401 from lightx2v.models.networks.wan.weights.dreamzero.pre_weights import DreamZeroCategoryLinearWeights from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER @ATTN_WEIGHT_REGISTER("dreamzero_fa2") @@ -163,6 +164,11 @@ def __init__(self): class DreamZeroTransformerWeights(WanTransformerWeights): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__(config, lazy_load_path=lazy_load_path, lora_path=lora_path) + for phase in self.iter_self_attention_phases(): + phase.add_module( + "dreamzero_rope", + ROPE_REGISTER[config.get("dreamzero_rope_type", "dreamzero_rope")](layout="interleaved", compute_dtype=torch.float32), + ) for block in self.blocks: self._ensure_image_cross_attn(block.compute_phases[1]) if getattr(self, "offload_block_cuda_buffers", None) is not None: diff --git a/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py b/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py index d7385b3a9..a812b0e84 100644 --- a/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/fastwam/transformer_weights.py @@ -1,10 +1,16 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER, TENSOR_REGISTER class FastWAMSelfAttentionWeights(WeightModule): def __init__(self, prefix, block_index, config): super().__init__() + self.add_module( + "rope", + ROPE_REGISTER[config.get("fastwam_rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float64), + ) block = f"{prefix}.blocks.{block_index}" rms_type = config.get("rms_norm_type", "torch") eps = float(config.get("eps", 1e-6)) diff --git a/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py b/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py index 303b68ca5..3481ce93f 100755 --- a/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/infinitetalk/transformer_weights.py @@ -1,6 +1,8 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.models.networks.wan.weights.transformer_weights import WanCrossAttention, WanFFN, WanSelfAttention -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, TENSOR_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, ROPE_REGISTER, TENSOR_REGISTER class WanInfiniteTalkTransformerWeights(WeightModule): @@ -184,6 +186,10 @@ def __init__( super().__init__() self.block_index = block_index self.config = config + self.add_module( + "rope_1d", + ROPE_REGISTER[config.get("audio_rope_type", "torch_real_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.add_module( "norm_x", LN_WEIGHT_REGISTER["torch"]( diff --git a/lightx2v/models/networks/wan/weights/lingbot_va/transformer_weights.py b/lightx2v/models/networks/wan/weights/lingbot_va/transformer_weights.py index c59f1cb9e..48ff85aa6 100644 --- a/lightx2v/models/networks/wan/weights/lingbot_va/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/lingbot_va/transformer_weights.py @@ -1,5 +1,7 @@ +import torch + from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights -from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER, ROPE_REGISTER class LingbotVATransformerWeights(WanTransformerWeights): @@ -7,6 +9,11 @@ class LingbotVATransformerWeights(WanTransformerWeights): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__(config, lazy_load_path=lazy_load_path, lora_path=lora_path) + for phase in self.iter_self_attention_phases(): + phase.add_module( + "lingbot_rope", + ROPE_REGISTER[config.get("lingbot_rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float64), + ) mm_type = config.get("dit_quant_scheme", "Default") if mm_type != "Default": assert config.get("dit_quantized") is True diff --git a/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py b/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py index fc006fca7..14f1016b0 100755 --- a/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/matrix_game2/transformer_weights.py @@ -1,3 +1,5 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.models.networks.wan.weights.transformer_weights import ( WanFFN, @@ -9,6 +11,7 @@ LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, TENSOR_REGISTER, ) @@ -98,6 +101,10 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): self.quant_method = config.get("quant_method", None) self.attn_rms_norm_type = self.config.get("rms_norm_type", "self_forcing") + self.add_module( + "rope", + ROPE_REGISTER[config.get("action_rope_type", "torch_real_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.add_module( "keyboard_embed_0", diff --git a/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py b/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py index 40e91b8f3..53fc01133 100644 --- a/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/matrix_game3/transformer_weights.py @@ -1,3 +1,5 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.models.networks.wan.weights.transformer_weights import ( WanFFN, @@ -8,6 +10,7 @@ LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, TENSOR_REGISTER, ) @@ -77,8 +80,13 @@ def __init__(self, block_index, task, mm_type, config, has_action=False, block_p self.config = config self.has_action = has_action + self_attn = WanSelfAttention(block_index, block_prefix, task, mm_type, config) + self_attn.add_module( + "indexed_rope", + ROPE_REGISTER[config.get("indexed_rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float32), + ) phases = [ - WanSelfAttention(block_index, block_prefix, task, mm_type, config), + self_attn, WanMtxg3CamInjection(block_index, block_prefix, mm_type, config), WanMtxg3CrossAttention(block_index, block_prefix, task, mm_type, config), ] @@ -219,6 +227,10 @@ def __init__(self, block_index, block_prefix, task, mm_type, config): self.block_index = block_index self.mm_type = mm_type self.config = config + self.add_module( + "rope", + ROPE_REGISTER[config.get("action_rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float32), + ) attn_rms_norm_type = config.get("rms_norm_type", "sgl-kernel") diff --git a/lightx2v/models/networks/wan/weights/motus/transformer_weights.py b/lightx2v/models/networks/wan/weights/motus/transformer_weights.py index 63f771ee0..006b0d9c1 100644 --- a/lightx2v/models/networks/wan/weights/motus/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/motus/transformer_weights.py @@ -1,5 +1,8 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights +from lightx2v.utils.registry_factory import ROPE_REGISTER from ._shared import MotusJointExpertBlockWeights, MotusJointExpertTransformerWeights, load_prefixed_submodules from .pre_weights import build_motus_expert_configs @@ -78,6 +81,10 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): expert_configs = build_motus_expert_configs(config) wan_num_heads = config["num_heads"] wan_head_dim = config["dim"] // wan_num_heads + rope = ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32) + if config.get("rope_chunk", False): + rope = ROPE_REGISTER["chunked_rope"](inner=rope, chunk_size=config.get("rope_chunk_size", 100)) + self.add_module("rope", rope) self.add_module("video", MotusWanTransformerWeights(config)) self.add_module( "action", diff --git a/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py b/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py index b8058c70d..a28dd8f81 100644 --- a/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/s2v/transformer_weights.py @@ -1,6 +1,8 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule from lightx2v.models.networks.wan.weights.transformer_weights import WanTransformerWeights -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER class WanS2VAudioInjectPhase(WeightModule): @@ -48,6 +50,11 @@ def __init__(self, injector_index, mm_type, config): class WanS2VTransformerWeights(WanTransformerWeights): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__(config, lazy_load_path, lora_path) + for phase in self.iter_self_attention_phases(): + phase.add_module( + "s2v_rope", + ROPE_REGISTER[config.get("s2v_rope_type", "torch_complex_rope")](layout="interleaved", compute_dtype=torch.float64), + ) inject_layers = config.get("audio_inject_layers", []) self.injected_block_id = {layer_idx: idx for idx, layer_idx in enumerate(inject_layers)} for layer_idx, injector_idx in self.injected_block_id.items(): diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py index 50654ef6e..13658fbe3 100755 --- a/lightx2v/models/networks/wan/weights/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/transformer_weights.py @@ -1,9 +1,13 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.models.networks.wan.infer.utils import WanCausalRope # noqa: F401 from lightx2v.utils.registry_factory import ( ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, TENSOR_REGISTER, ) @@ -141,6 +145,22 @@ def non_block_weights_to_cpu(self): self.head.to_cpu() self.head_modulation.to_cpu() + def iter_self_attention_phases(self): + for block in self.blocks: + yield block.compute_phases[0] + for name in ("offload_block_cuda_buffers", "offload_block_cpu_buffers"): + buffers = getattr(self, name, None) + if buffers is not None: + for block in buffers: + yield block.compute_phases[0] + phases = getattr(self, "offload_phase_cuda_buffers", None) + if phases is not None: + yield phases[0] + phase_buffers = getattr(self, "offload_phase_cpu_buffers", None) + if phase_buffers is not None: + for phases in phase_buffers: + yield phases[0] + class WanTransformerAttentionBlock(WeightModule): def __init__( @@ -239,6 +259,15 @@ def __init__( self.lazy_load = lazy_load self.lazy_load_file = lazy_load_file self.attn_rms_norm_type = self.config.get("rms_norm_type", "sgl-kernel") + rope = ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32) + if config.get("rope_chunk", False): + rope = ROPE_REGISTER["chunked_rope"](inner=rope, chunk_size=config.get("rope_chunk_size", 100)) + self.add_module("rope", rope) + if config.get("causal_rope_type") is not None: + self.add_module( + "causal_rope", + ROPE_REGISTER[config.get("causal_rope_type", "wan_causal_rope")](layout="interleaved", compute_dtype=torch.float64), + ) self.add_module( "modulation", diff --git a/lightx2v/models/networks/worldmirror/models/layers/norm_rope.py b/lightx2v/models/networks/worldmirror/models/layers/norm_rope.py index df5bd2968..bcd32364e 100644 --- a/lightx2v/models/networks/worldmirror/models/layers/norm_rope.py +++ b/lightx2v/models/networks/worldmirror/models/layers/norm_rope.py @@ -6,6 +6,10 @@ import torch.nn as nn from torch import Tensor +from lightx2v.common.ops.rope import TorchRealRope + +_WORLDMIRROR_NORMALIZED_ROPE = TorchRealRope(layout="split_half") + class PositionGetter: """Generates and caches 2D spatial positions for patches in a grid.""" @@ -23,11 +27,6 @@ def __call__(self, batch_size: int, height: int, width: int, device: torch.devic return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone() -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - - class NormalizedRotaryPositionEmbedding2D(nn.Module): """DINOv3-aligned 2D Rotary Position Embedding.""" @@ -136,4 +135,4 @@ def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor flat_indices = indices.view(-1) gathered_sin = sin[flat_indices].view(B, 1, N, C_head) gathered_cos = cos[flat_indices].view(B, 1, N, C_head) - return (tokens * gathered_cos) + (_rotate_half(tokens) * gathered_sin) + return _WORLDMIRROR_NORMALIZED_ROPE.apply_single(tokens, (gathered_cos, gathered_sin)) diff --git a/lightx2v/models/networks/worldmirror/models/layers/rope.py b/lightx2v/models/networks/worldmirror/models/layers/rope.py index 07eb7e83e..d72566777 100755 --- a/lightx2v/models/networks/worldmirror/models/layers/rope.py +++ b/lightx2v/models/networks/worldmirror/models/layers/rope.py @@ -14,6 +14,10 @@ import torch.nn as nn import torch.nn.functional as F +from lightx2v.common.ops.rope import TorchRealRope + +_WORLDMIRROR_ROPE = TorchRealRope(layout="split_half") + class PositionGetter: """Generates and caches 2D spatial positions for patches in a grid. @@ -112,20 +116,6 @@ def _compute_frequency_components(self, dim: int, seq_len: int, device: torch.de return self.frequency_cache[cache_key] - @staticmethod - def _rotate_features(x: torch.Tensor) -> torch.Tensor: - """Performs feature rotation by splitting and recombining feature dimensions. - - Args: - x: Input tensor to rotate. - - Returns: - Rotated feature tensor. - """ - feature_dim = x.shape[-1] - x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :] - return torch.cat((-x2, x1), dim=-1) - def _apply_1d_rope(self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp: torch.Tensor, sin_comp: torch.Tensor) -> torch.Tensor: """Applies 1D rotary position embeddings along one dimension. @@ -142,8 +132,7 @@ def _apply_1d_rope(self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp cos = F.embedding(positions, cos_comp)[:, None, :, :] sin = F.embedding(positions, sin_comp)[:, None, :, :] - # Apply rotation - return (tokens * cos) + (self._rotate_features(tokens) * sin) + return _WORLDMIRROR_ROPE.apply_single(tokens, (cos, sin)) def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: """Applies 2D rotary position embeddings to input tokens. diff --git a/lightx2v/models/networks/worldplay/infer/ar_transformer_infer.py b/lightx2v/models/networks/worldplay/infer/ar_transformer_infer.py index 92da97745..98ba16ca7 100644 --- a/lightx2v/models/networks/worldplay/infer/ar_transformer_infer.py +++ b/lightx2v/models/networks/worldplay/infer/ar_transformer_infer.py @@ -12,6 +12,7 @@ ) from lightx2v.models.networks.hunyuan_video.infer.transformer_infer import ( HunyuanVideo15TransformerInfer, + _apply_hunyuan_rope, apply_gate, ) from lightx2v.models.networks.worldplay.prope.camera_rope import prope_qkv @@ -235,7 +236,7 @@ def _infer_img_branch_before_attn(self, weights, infer_module_out): img_k_pre_rope = img_k.unsqueeze(0) # Apply RoPE for standard attention branch - img_q, img_k = self.apply_rope_func(img_q.unsqueeze(0), img_k.unsqueeze(0), cos_sin_cache=self.scheduler.cos_sin) + img_q, img_k = _apply_hunyuan_rope(weights.rope, img_q.unsqueeze(0), img_k.unsqueeze(0), self.scheduler.cos_sin) return ( img_q, diff --git a/lightx2v/models/networks/worldplay/infer/bi_transformer_infer.py b/lightx2v/models/networks/worldplay/infer/bi_transformer_infer.py index 3f9226c8d..f9d3ff51d 100644 --- a/lightx2v/models/networks/worldplay/infer/bi_transformer_infer.py +++ b/lightx2v/models/networks/worldplay/infer/bi_transformer_infer.py @@ -9,6 +9,7 @@ ) from lightx2v.models.networks.hunyuan_video.infer.transformer_infer import ( HunyuanVideo15TransformerInfer, + _apply_hunyuan_rope, apply_gate, ) from lightx2v.models.networks.worldplay.prope.camera_rope import prope_qkv @@ -147,7 +148,7 @@ def _infer_img_branch_before_attn(self, weights, infer_module_out): img_k_pre_rope = img_k.unsqueeze(0) # Apply RoPE for standard attention branch - img_q, img_k = self.apply_rope_func(img_q.unsqueeze(0), img_k.unsqueeze(0), cos_sin_cache=self.scheduler.cos_sin) + img_q, img_k = _apply_hunyuan_rope(weights.rope, img_q.unsqueeze(0), img_k.unsqueeze(0), self.scheduler.cos_sin) return ( img_q, diff --git a/lightx2v/models/networks/worldplay/infer/transformer_infer.py b/lightx2v/models/networks/worldplay/infer/transformer_infer.py index 04d3c10db..6366cacf4 100644 --- a/lightx2v/models/networks/worldplay/infer/transformer_infer.py +++ b/lightx2v/models/networks/worldplay/infer/transformer_infer.py @@ -9,6 +9,7 @@ ) from lightx2v.models.networks.hunyuan_video.infer.transformer_infer import ( HunyuanVideo15TransformerInfer, + _apply_hunyuan_rope, apply_gate, ) from lightx2v.models.networks.worldplay.prope.camera_rope import prope_qkv @@ -137,7 +138,7 @@ def _infer_img_branch_before_attn(self, weights, infer_module_out): img_k_pre_rope = img_k.unsqueeze(0) # Apply RoPE for standard attention branch - img_q, img_k = self.apply_rope_func(img_q.unsqueeze(0), img_k.unsqueeze(0), cos_sin_cache=self.scheduler.cos_sin) + img_q, img_k = _apply_hunyuan_rope(weights.rope, img_q.unsqueeze(0), img_k.unsqueeze(0), self.scheduler.cos_sin) return ( img_q, diff --git a/lightx2v/models/networks/z_image/infer/transformer_infer.py b/lightx2v/models/networks/z_image/infer/transformer_infer.py index 26657f8a0..515d6c76b 100755 --- a/lightx2v/models/networks/z_image/infer/transformer_infer.py +++ b/lightx2v/models/networks/z_image/infer/transformer_infer.py @@ -2,9 +2,6 @@ import torch.nn.functional as F from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer -from lightx2v.utils.registry_factory import ROPE_REGISTER - -from .utils import apply_rotary_emb_qwen, apply_wan_rope_with_flashinfer class ZImageTransformerInfer(BaseTransformerInfer): @@ -28,27 +25,6 @@ def __init__(self, config): self.enable_head_parallel = False self.seq_p_tensor_fusion = False - rope_funcs = { - "flashinfer": apply_wan_rope_with_flashinfer, - "torch": apply_rotary_emb_qwen, - } - - rope_type = self.config.get("rope_type", "flashinfer") - if rope_type in ROPE_REGISTER: - rope_class = ROPE_REGISTER[rope_type] - self.rope_instance = rope_class() - - # Create a wrapper function that matches the expected signature - def rope_wrapper(xq, xk, cos_sin_cache): - return self.rope_instance.apply(xq, xk, cos_sin_cache) - - rope_func = rope_wrapper - else: - if rope_type not in rope_funcs: - raise ValueError(f"Unsupported z-image rope_type: {rope_type}") - rope_func = rope_funcs[rope_type] - self.apply_rope_func = rope_func - def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -94,7 +70,7 @@ def infer_attn( if attn_phase.norm_k is not None: key = attn_phase.norm_k.apply(key) - query, key = self.apply_rope_func(query, key, freqs_cis) + query, key = attn_phase.rope.apply(query, key, freqs_cis) total_seq_len = query.shape[0] cu_seqlens = torch.tensor([0, total_seq_len], dtype=torch.int32, device="cpu") diff --git a/lightx2v/models/networks/z_image/infer/utils.py b/lightx2v/models/networks/z_image/infer/utils.py index 6f596466f..e524e4afa 100755 --- a/lightx2v/models/networks/z_image/infer/utils.py +++ b/lightx2v/models/networks/z_image/infer/utils.py @@ -1,59 +1,5 @@ -from typing import Tuple - import torch -try: - from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace -except ImportError: - apply_rope_with_cos_sin_cache_inplace = None - - -def apply_wan_rope_with_flashinfer( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -): - L, H, D = xq.shape - - query = xq.reshape(L, H * D).contiguous() - key = xk.reshape(L, H * D).contiguous() - - positions = torch.arange(L, device="cpu", dtype=torch.long).to(xq.device, non_blocking=True) - - apply_rope_with_cos_sin_cache_inplace( - positions=positions, - query=query, - key=key, - head_size=D, - cos_sin_cache=cos_sin_cache, - is_neox=False, - ) - - xq_out = query.view(L, H, D) - xk_out = key.view(L, H, D) - return xq_out, xk_out - - -def apply_rotary_emb_qwen( - xq: torch.Tensor, - xk: torch.Tensor, - cos_sin_cache: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - head_dim = xq.shape[-1] - cos, sin = cos_sin_cache[..., : head_dim // 2], cos_sin_cache[..., head_dim // 2 :] - cos = cos.unsqueeze(1).to(dtype=xq.dtype) - sin = sin.unsqueeze(1).to(dtype=xq.dtype) - - def rotate(x): - x_even = x[..., 0::2] - x_odd = x[..., 1::2] - x_out = torch.empty_like(x) - x_out[..., 0::2] = x_even * cos - x_odd * sin - x_out[..., 1::2] = x_odd * cos + x_even * sin - return x_out - - return rotate(xq), rotate(xk) - def patchify(hidden_states: torch.Tensor, patch_size: int = 2, f_patch_size: int = 1) -> torch.Tensor: B, C, H, W = hidden_states.shape diff --git a/lightx2v/models/networks/z_image/weights/transformer_weights.py b/lightx2v/models/networks/z_image/weights/transformer_weights.py index b252027aa..5e64da9da 100755 --- a/lightx2v/models/networks/z_image/weights/transformer_weights.py +++ b/lightx2v/models/networks/z_image/weights/transformer_weights.py @@ -1,8 +1,11 @@ +import torch + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ( ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, + ROPE_REGISTER, ) @@ -260,6 +263,10 @@ def __init__( self.config = config self.attn_type = config.get("attn_type", "flash_attn3") self.rms_norm_type = config.get("rms_norm_type", "one-pass") + self.add_module( + "rope", + ROPE_REGISTER[config.get("rope_type", "flashinfer_rope")](layout="interleaved", compute_dtype=torch.float32), + ) self.lazy_load = lazy_load self.lazy_load_file = lazy_load_file diff --git a/lightx2v/models/schedulers/ernie_image/scheduler.py b/lightx2v/models/schedulers/ernie_image/scheduler.py index 7c026a133..2608d19b1 100644 --- a/lightx2v/models/schedulers/ernie_image/scheduler.py +++ b/lightx2v/models/schedulers/ernie_image/scheduler.py @@ -34,9 +34,8 @@ def set_timesteps(self): 1.0, 0.0, self.config["infer_steps"] + 1, - device=AI_DEVICE, dtype=torch.float32, - ) + ).tolist() self.scheduler.set_timesteps(sigmas=sigmas[:-1], device=AI_DEVICE) self.timesteps = self.scheduler.timesteps self.infer_steps = len(self.timesteps) diff --git a/lightx2v/models/schedulers/longcat_image/scheduler.py b/lightx2v/models/schedulers/longcat_image/scheduler.py index e64615596..fad48ee12 100755 --- a/lightx2v/models/schedulers/longcat_image/scheduler.py +++ b/lightx2v/models/schedulers/longcat_image/scheduler.py @@ -284,11 +284,6 @@ def chunk_image_emb(image_emb): image_emb = F.pad(image_emb, (0, 0, 0, padding_size)) return torch.chunk(image_emb, world_size, dim=0)[cur_rank] - if self.config.get("rope_type", "flashinfer") == "flashinfer": - txt_emb = rotary_emb[:txt_seq_len] - img_emb = chunk_image_emb(rotary_emb[txt_seq_len:]) - return torch.cat([txt_emb, img_emb], dim=0) - freqs_cos, freqs_sin = rotary_emb txt_cos, img_cos = freqs_cos[:txt_seq_len], freqs_cos[txt_seq_len:] txt_sin, img_sin = freqs_sin[:txt_seq_len], freqs_sin[txt_seq_len:] @@ -388,15 +383,7 @@ def prepare(self, input_info): ids = torch.cat([txt_ids, img_ids], dim=0).to(AI_DEVICE, dtype=torch.float32) freqs_cos, freqs_sin = self.pos_embed(ids) - # Convert to flashinfer format if needed - if self.config.get("rope_type", "flashinfer") == "flashinfer": - # pos_embed returns interleaved format: [c0, c0, c1, c1, ...] - # flashinfer needs: [c0, c1, ..., s0, s1, ...] - cos_half = freqs_cos[:, ::2].contiguous() # [L, D/2] - sin_half = freqs_sin[:, ::2].contiguous() # [L, D/2] - self.image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1) # [L, D] - else: - self.image_rotary_emb = (freqs_cos, freqs_sin) + self.image_rotary_emb = (freqs_cos, freqs_sin) self.image_rotary_emb = self._seq_parallel_rope(self.image_rotary_emb, txt_seq_len) # Handle CFG: prepare negative embeddings rotary @@ -406,12 +393,7 @@ def prepare(self, input_info): neg_ids = torch.cat([neg_txt_ids, img_ids], dim=0).to(AI_DEVICE, dtype=torch.float32) neg_freqs_cos, neg_freqs_sin = self.pos_embed(neg_ids) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - neg_cos_half = neg_freqs_cos[:, ::2].contiguous() - neg_sin_half = neg_freqs_sin[:, ::2].contiguous() - self.negative_image_rotary_emb = torch.cat([neg_cos_half, neg_sin_half], dim=-1) - else: - self.negative_image_rotary_emb = (neg_freqs_cos, neg_freqs_sin) + self.negative_image_rotary_emb = (neg_freqs_cos, neg_freqs_sin) self.negative_image_rotary_emb = self._seq_parallel_rope(self.negative_image_rotary_emb, neg_txt_seq_len) def step_pre(self, step_index): @@ -469,13 +451,7 @@ def prepare_i2i(self, input_info, input_image, vae): ids = torch.cat([txt_ids, combined_img_ids], dim=0).to(AI_DEVICE, dtype=torch.float32) freqs_cos, freqs_sin = self.pos_embed(ids) - # Convert to flashinfer format if needed - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half = freqs_cos[:, ::2].contiguous() - sin_half = freqs_sin[:, ::2].contiguous() - self.image_rotary_emb = torch.cat([cos_half, sin_half], dim=-1) - else: - self.image_rotary_emb = (freqs_cos, freqs_sin) + self.image_rotary_emb = (freqs_cos, freqs_sin) self.image_rotary_emb = self._seq_parallel_rope(self.image_rotary_emb, txt_seq_len) # Store output sequence length for later truncation @@ -488,12 +464,7 @@ def prepare_i2i(self, input_info, input_image, vae): neg_ids = torch.cat([neg_txt_ids, combined_img_ids], dim=0).to(AI_DEVICE, dtype=torch.float32) neg_freqs_cos, neg_freqs_sin = self.pos_embed(neg_ids) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - neg_cos_half = neg_freqs_cos[:, ::2].contiguous() - neg_sin_half = neg_freqs_sin[:, ::2].contiguous() - self.negative_image_rotary_emb = torch.cat([neg_cos_half, neg_sin_half], dim=-1) - else: - self.negative_image_rotary_emb = (neg_freqs_cos, neg_freqs_sin) + self.negative_image_rotary_emb = (neg_freqs_cos, neg_freqs_sin) self.negative_image_rotary_emb = self._seq_parallel_rope(self.negative_image_rotary_emb, neg_txt_seq_len) def _encode_image(self, image): diff --git a/lightx2v/models/schedulers/qwen_image/scheduler.py b/lightx2v/models/schedulers/qwen_image/scheduler.py index 21278dd45..c7b707ed9 100755 --- a/lightx2v/models/schedulers/qwen_image/scheduler.py +++ b/lightx2v/models/schedulers/qwen_image/scheduler.py @@ -667,13 +667,6 @@ def prepare(self, input_info): self.prepare_i2i_denoise_strength_latents(input_info) self.image_rotary_emb = self.pos_embed(self.input_info.image_shapes, input_info.txt_seq_lens[0], device=AI_DEVICE) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half_img = self.image_rotary_emb[0].real.contiguous() - sin_half_img = self.image_rotary_emb[0].imag.contiguous() - cos_half_txt = self.image_rotary_emb[1].real.contiguous() - sin_half_txt = self.image_rotary_emb[1].imag.contiguous() - self.image_rotary_emb[0] = torch.cat([cos_half_img, sin_half_img], dim=-1) - self.image_rotary_emb[1] = torch.cat([cos_half_txt, sin_half_txt], dim=-1) if self.seq_p_group is not None: world_size = dist.get_world_size(self.seq_p_group) cur_rank = dist.get_rank(self.seq_p_group) @@ -685,13 +678,6 @@ def prepare(self, input_info): if self.config["enable_cfg"]: self.negative_image_rotary_emb = self.pos_embed(self.input_info.image_shapes, input_info.txt_seq_lens[1], device=AI_DEVICE) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half_img = self.negative_image_rotary_emb[0].real.contiguous() - sin_half_img = self.negative_image_rotary_emb[0].imag.contiguous() - cos_half_txt = self.negative_image_rotary_emb[1].real.contiguous() - sin_half_txt = self.negative_image_rotary_emb[1].imag.contiguous() - self.negative_image_rotary_emb[0] = torch.cat([cos_half_img, sin_half_img], dim=-1) - self.negative_image_rotary_emb[1] = torch.cat([cos_half_txt, sin_half_txt], dim=-1) if self.seq_p_group is not None: world_size = dist.get_world_size(self.seq_p_group) cur_rank = dist.get_rank(self.seq_p_group) diff --git a/lightx2v/models/schedulers/z_image/scheduler.py b/lightx2v/models/schedulers/z_image/scheduler.py index 7edfcaa6b..be1d4ae20 100755 --- a/lightx2v/models/schedulers/z_image/scheduler.py +++ b/lightx2v/models/schedulers/z_image/scheduler.py @@ -507,9 +507,7 @@ def generate_freqs_cis_from_position_ids(self, position_ids: torch.Tensor, devic device = position_ids.device position_ids = position_ids.to(device) - rope_type = self.config.get("rope_type", "flashinfer") cache_key = ( - rope_type, str(device), str(position_ids.dtype), tuple(position_ids.shape), @@ -520,11 +518,7 @@ def generate_freqs_cis_from_position_ids(self, position_ids: torch.Tensor, devic if cached_freqs_cis is not None: return cached_freqs_cis - if rope_type in {"flashinfer", "torch"}: - freqs_cos, freqs_sin = self.rope_embedder(position_ids, return_real=True) - freqs_cis = torch.cat([freqs_cos, freqs_sin], dim=-1).float() - else: - freqs_cis = self.rope_embedder(position_ids) + freqs_cis = self.rope_embedder(position_ids) self.freqs_cis_cache[cache_key] = freqs_cis return freqs_cis @@ -577,13 +571,6 @@ def prepare(self, input_info): self.image_rotary_emb = self.pos_embed(self.input_info.image_shapes, input_info.txt_seq_lens[0], device=AI_DEVICE) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half_img = self.image_rotary_emb[0].real.contiguous() - sin_half_img = self.image_rotary_emb[0].imag.contiguous() - cos_half_txt = self.image_rotary_emb[1].real.contiguous() - sin_half_txt = self.image_rotary_emb[1].imag.contiguous() - self.image_rotary_emb[0] = torch.cat([cos_half_img, sin_half_img], dim=-1) - self.image_rotary_emb[1] = torch.cat([cos_half_txt, sin_half_txt], dim=-1) if self.seq_p_group is not None: world_size = dist.get_world_size(self.seq_p_group) cur_rank = dist.get_rank(self.seq_p_group) @@ -595,13 +582,6 @@ def prepare(self, input_info): if self.config["enable_cfg"]: self.negative_image_rotary_emb = self.pos_embed(self.input_info.image_shapes, input_info.txt_seq_lens[1], device=AI_DEVICE) - if self.config.get("rope_type", "flashinfer") == "flashinfer": - cos_half_img = self.negative_image_rotary_emb[0].real.contiguous() - sin_half_img = self.negative_image_rotary_emb[0].imag.contiguous() - cos_half_txt = self.negative_image_rotary_emb[1].real.contiguous() - sin_half_txt = self.negative_image_rotary_emb[1].imag.contiguous() - self.negative_image_rotary_emb[0] = torch.cat([cos_half_img, sin_half_img], dim=-1) - self.negative_image_rotary_emb[1] = torch.cat([cos_half_txt, sin_half_txt], dim=-1) if self.seq_p_group is not None: world_size = dist.get_world_size(self.seq_p_group) cur_rank = dist.get_rank(self.seq_p_group) diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py index 351bb31dd..41ab93634 100755 --- a/lightx2v/pipeline.py +++ b/lightx2v/pipeline.py @@ -171,7 +171,7 @@ def create_generator( boundary_step_index=2, denoising_step_list=[1000, 750, 500, 250], config_json=None, - rope_type="torch", + rope_type="torch_complex_rope", resize_mode=None, audio_fps=24000, double_precision_rope=True, diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index 97155fa10..afa0e0301 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -199,7 +199,11 @@ def auto_calc_config(config): elif os.path.exists(os.path.join(config["model_path"], "transformer", "config.json")): with open(os.path.join(config["model_path"], "transformer", "config.json"), "r") as f: model_config = json.load(f) - if config["model_cls"] == "z_image": + if config["model_cls"] == "ltx2": + # Upstream LTX2 uses rope_type for the layout name ("split"), + # while LightX2V uses it as the registered RoPE implementation. + model_config.pop("rope_type", None) + elif config["model_cls"] == "z_image": # https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/blob/main/transformer/config.json z_image_patch_size = model_config.pop("all_patch_size", [2]) z_image_f_patch_size = model_config.pop("all_f_patch_size", [1]) @@ -228,6 +232,19 @@ def auto_calc_config(config): if "infer_steps" not in config and "num_inference_steps" in config: config["infer_steps"] = config["num_inference_steps"] + if config["model_cls"] == "lingbot_va" and "target_video_length" not in config: + ar_config = config.get("ar_config", {}) + required_keys = ("num_frame_per_chunk", "num_chunks") + missing_keys = [key for key in required_keys if key not in ar_config] + if missing_keys: + raise ValueError(f"LingBot-VA requires ar_config.{', ar_config.'.join(missing_keys)} to derive target_video_length.") + latent_frames = int(ar_config["num_frame_per_chunk"]) * int(ar_config["num_chunks"]) + temporal_stride = int(config["vae_stride"][0]) + if latent_frames <= 0 or temporal_stride <= 0: + raise ValueError(f"LingBot-VA requires positive latent frame count and VAE temporal stride, got latent_frames={latent_frames}, temporal_stride={temporal_stride}.") + config["target_video_length"] = (latent_frames - 1) * temporal_stride + 1 + logger.info(f"Auto-set LingBot-VA target_video_length={config['target_video_length']} from {latent_frames} latent frames and temporal stride {temporal_stride}.") + if config["task"] in ["i2v", "t2av", "i2av", "i2va", "s2v", "rs2v", "ltx2_s2v", "v2av"]: if config["target_video_length"] % config["vae_stride"][0] != 1: logger.warning(f"`num_frames - 1` has to be divisible by {config['vae_stride'][0]}. Rounding to the nearest number.") diff --git a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py index 346209283..5bfc1d895 100644 --- a/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py +++ b/lightx2v_platform/ops/rope/ascend_npu/npu_rope.py @@ -36,8 +36,10 @@ def GET_SENSITIVE_DTYPE(): @PLATFORM_ROPE_REGISTER("npu_rope") class NpuRope(RopeTemplate): - def __init__(self): - super().__init__() + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("NpuRope only supports interleaved layout.") self.sensitive_layer_dtype = GET_SENSITIVE_DTYPE() def _apply_rope_fp32(self, xq, xk, cos_sin_cache): diff --git a/lightx2v_platform/ops/rope/enflame_gcu/wan_rope.py b/lightx2v_platform/ops/rope/enflame_gcu/wan_rope.py index 4497f01e9..5de40f43f 100644 --- a/lightx2v_platform/ops/rope/enflame_gcu/wan_rope.py +++ b/lightx2v_platform/ops/rope/enflame_gcu/wan_rope.py @@ -6,8 +6,10 @@ @PLATFORM_ROPE_REGISTER("gcu_wan_rope") class GcuWanRope(RopeTemplate): - def __init__(self): - super().__init__() + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("GcuWanRope only supports interleaved layout.") def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor): """ diff --git a/lightx2v_platform/ops/rope/iluvatar_cuda/wan_rope.py b/lightx2v_platform/ops/rope/iluvatar_cuda/wan_rope.py index d5d93ad9a..45a21e19b 100644 --- a/lightx2v_platform/ops/rope/iluvatar_cuda/wan_rope.py +++ b/lightx2v_platform/ops/rope/iluvatar_cuda/wan_rope.py @@ -10,8 +10,10 @@ @PLATFORM_ROPE_REGISTER("iluvatar_wan_rope") class IluvatarWanRope(RopeTemplate): - def __init__(self): - super().__init__() + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + super().__init__(layout=layout, compute_dtype=compute_dtype) + if layout != "interleaved": + raise ValueError("IluvatarWanRope only supports interleaved layout.") def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor): """ diff --git a/lightx2v_platform/ops/rope/rope_template.py b/lightx2v_platform/ops/rope/rope_template.py index b90e1ee76..7ea59858f 100644 --- a/lightx2v_platform/ops/rope/rope_template.py +++ b/lightx2v_platform/ops/rope/rope_template.py @@ -25,7 +25,11 @@ def GET_DTYPE(): class RopeTemplate(metaclass=ABCMeta): - def __init__(self): + def __init__(self, layout="interleaved", compute_dtype=torch.float32): + if layout not in {"interleaved", "split_half"}: + raise ValueError(f"Unsupported RoPE layout: {layout}") + self.layout = layout + self.compute_dtype = compute_dtype self.infer_dtype = GET_DTYPE() self.config = {} @@ -47,3 +51,31 @@ def apply(self, xq: torch.Tensor, xk: torch.Tensor, cos_sin_cache: torch.Tensor) def set_config(self, config=None): if config is not None: self.config = config + + def load(self, weight_dict): + pass + + def to_cpu(self, non_blocking=False): + pass + + def to_cuda(self, non_blocking=False): + pass + + def state_dict(self, destination=None): + return {} if destination is None else destination + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + return {} if destination is None else destination + + def load_state_dict_from_disk(self, block_index, adapter_block_index=None): + pass + + def named_parameters(self, prefix=""): + return iter(()) + + def prepare_freqs(self, freqs): + return freqs + + def apply_single(self, x: torch.Tensor, freqs, **kwargs): + output, _ = self.apply(x, x.clone(), freqs, **kwargs) + return output