From 153e612fbe63e3eb3f63914d78c64e457e7494cf Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Tue, 8 Sep 2026 08:18:01 +0000 Subject: [PATCH 1/7] sla xpu init implement --- .gitignore | 1 + .../minimax_h3_fl2v_turbo_sla_4step.json | 44 ++++ .../common/ops/attn/dynamic_sparse_attn.py | 48 ++++- .../minimax_h3/weights/transformer_weights.py | 5 +- .../runners/minimax_h3/minimax_h3_runner.py | 2 +- lightx2v_kernel_xpu/CMakeLists.txt | 97 +++++++++ lightx2v_kernel_xpu/README.md | 24 +++ lightx2v_kernel_xpu/build.sh | 2 +- lightx2v_kernel_xpu/cute/cute_fmha_torch.cpp | 59 +++++- .../patches/minimax_h3_sparse_kernel.patch | 21 ++ .../cute/patches/minimax_h3_sparse_lut.patch | 92 +++++++++ .../cute/patches/minimax_h3_sparse_tail.patch | 18 ++ .../python/sycl_kernels/__init__.py | 35 ++++ .../python/sycl_kernels/sla.py | 193 ++++++++++++++++++ .../python/sycl_kernels/sla_triton.py | 91 +++++++++ .../test/bench_sla_sparse_attention.py | 76 +++++++ .../test/test_sla_sparse_attention.py | 69 +++++++ .../test/test_xpu_sla_attn_adapter.py | 79 +++++++ .../ops/attn/intel_xpu/__init__.py | 1 + .../ops/attn/intel_xpu/xpu_sla_attn.py | 41 ++++ .../run_minimax_h3_fl2av_turbo_sla_4step.sh | 87 ++++++++ 21 files changed, 1079 insertions(+), 6 deletions(-) create mode 100644 configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json create mode 100644 lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch create mode 100644 lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch create mode 100644 lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_tail.patch create mode 100644 lightx2v_kernel_xpu/python/sycl_kernels/sla.py create mode 100644 lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py create mode 100644 lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py create mode 100644 lightx2v_kernel_xpu/test/test_sla_sparse_attention.py create mode 100644 lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py create mode 100644 lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py create mode 100755 scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh diff --git a/.gitignore b/.gitignore index d6fe813c1..a6139a02c 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,4 @@ lightx2v_ros/build lightx2v_ros/install lightx2v_ros/log .gitnexus +lightx2v_kernel_xpu/_cmake_build diff --git a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json new file mode 100644 index 000000000..9e83f51e9 --- /dev/null +++ b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json @@ -0,0 +1,44 @@ +{ + "infer_steps": 5, + "target_video_length": 362, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "block", + "text_encoder_cpu_offload": true, + "text_encoder_offload_granularity": "block", + "text_encoder_host_pinned": false, + "text_encoder_release_block_offload_buffers": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "dynamic_sparse_attn", + "refiner_attn_type": "intel_xpu_cute_attn", + "dynamic_sparse_attn_setting": { + "sparsity_ratio": 0.85, + "operator": "intel_xpu" + }, + "rms_type": "intel_xpu", + "rope_type": "minimax_h3_xpu_rope", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 6.0, + "audio_flow_shift": 3.0, + "h3_step_update": "training_euler", + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true, + "lora_dynamic_apply": true, + "lora_configs": [ + { + "path": "/llm/models/Minimax-h3-Turbo-SLA/minimax_h3_fl2v_turbo_4step_v0.1_768p_sla_bf16.safetensors", + "strength": 1.0, + "alpha": 128 + } + ] +} diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index eb9281e81..5ace0b9a6 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -77,9 +77,18 @@ def __init__(self, config=None): if not 0.0 <= self.sparsity_ratio < 1.0: raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}") - self.arch = get_cuda_arch(torch.cuda.current_device()) self.topk = 1 - self.sparsity_ratio - if self.operator == "triton": + self.arch = None + if self.operator != "intel_xpu": + self.arch = get_cuda_arch(torch.cuda.current_device()) + + if self.operator == "intel_xpu": + # The optimized MiniMax-H3 kernel consumes BLHD directly. Keep + # this branch ahead of CUDA architecture discovery so merely + # constructing the XPU backend never touches torch.cuda. + self.BLKQ, self.BLKK = 128, 128 + self.apply_func = self.apply_intel_xpu + elif self.operator == "triton": self.BLKQ, self.BLKK = 64, 64 self.apply_func = self.apply_triton elif self.operator == "triton_ar": # triton for AR models @@ -105,6 +114,41 @@ def __init__(self, config=None): # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}") + def apply_intel_xpu( + self, + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + """Run the XPU SLA router and fused sparse attention on one sequence.""" + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("Intel XPU SLA expects q, k and v in [L, H, D] layout") + if q.shape != k.shape or k.shape != v.shape: + raise ValueError("Intel XPU SLA currently requires self-attention with matching q/k/v shapes") + if max_seqlen_q != q.shape[0] or max_seqlen_kv != k.shape[0]: + raise ValueError("Intel XPU SLA currently supports one unpadded sequence per call") + if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: + raise ValueError("Intel XPU SLA currently supports one sequence per call") + if cu_seqlens_kv is not None and cu_seqlens_kv.numel() != 2: + raise ValueError("Intel XPU SLA currently supports one sequence per call") + + from lightx2v_platform.ops.attn.intel_xpu.xpu_sla_attn import sla_sparse_attention + + out = sla_sparse_attention( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + keep_ratio=self.topk, + block_q=self.BLKQ, + block_k=self.BLKK, + ) + return out.reshape(max_seqlen_q, -1) + def apply( self, q, diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..b85c3d6c5 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -74,7 +74,10 @@ def __init__(self, prefix, config, create_cuda_buffer=False): attn_type = config.get("attn_type", "flash_attn3") attention_cls = ATTN_WEIGHT_REGISTER[attn_type] if attn_type == "dynamic_sparse_attn": - calculate = attention_cls(config.get("dynamic_sparse_attn_setting", {})) + sparse_config = config.get("dynamic_sparse_attn_setting", {}) + if sparse_config.get("operator") == "intel_xpu" and config.get("seq_parallel", False): + raise NotImplementedError("Intel XPU SLA does not yet support MiniMax-H3 sequence parallelism") + calculate = attention_cls(sparse_config) else: calculate = attention_cls() if attn_type == "sol_attn": diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index cbc1b33c3..510425f3c 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -61,7 +61,7 @@ def build_minimax_h3_model_with_lora(config, model_kwargs, lora_configs): if not lora_config.get("path"): raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].path") if lora_config.get("alpha") is None: - raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].alpha (use 8 for the MiniMax-H3 Turbo LoRA)") + raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].alpha (use the alpha published with the checkpoint)") model_kwargs.update( lora_path=lora_config["path"], lora_strength=lora_config.get("strength", 1.0), diff --git a/lightx2v_kernel_xpu/CMakeLists.txt b/lightx2v_kernel_xpu/CMakeLists.txt index 7da51157c..133bf8eb6 100644 --- a/lightx2v_kernel_xpu/CMakeLists.txt +++ b/lightx2v_kernel_xpu/CMakeLists.txt @@ -284,6 +284,9 @@ if(ENABLE_CUTE_FMHA) set(H3_OVERLAY_ROOT "${CMAKE_BINARY_DIR}/minimax_h3_cute_overlay") set(H3_OVERLAY_DIR "${H3_OVERLAY_ROOT}/flash_attention_v2/collective") file(MAKE_DIRECTORY "${H3_OVERLAY_DIR}") + # Remove a stale sparse kernel overlay left by older build trees. The + # dense target must resolve the upstream kernel header unchanged. + file(REMOVE "${H3_OVERLAY_ROOT}/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp") set(H3_MAINLOOP_SOURCE "${CUTLASS_SYCL_ROOT}/applications/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp") set(H3_MAINLOOP_PATCH @@ -306,9 +309,69 @@ if(ENABLE_CUTE_FMHA) "${H3_OVERLAY_DIR}/xe_fmha_fwd_mainloop.hpp" COPYONLY) file(REMOVE "${H3_MAINLOOP_TEMP}") + + set(H3_SPARSE_OVERLAY_ROOT "${CMAKE_BINARY_DIR}/minimax_h3_sparse_cute_overlay") + set(H3_SPARSE_OVERLAY_DIR "${H3_SPARSE_OVERLAY_ROOT}/flash_attention_v2/collective") + file(MAKE_DIRECTORY "${H3_SPARSE_OVERLAY_DIR}") + configure_file( + "${H3_OVERLAY_DIR}/xe_fmha_fwd_mainloop.hpp" + "${H3_SPARSE_OVERLAY_DIR}/xe_fmha_fwd_mainloop.hpp" + COPYONLY) + + # Add block-LUT traversal after the long-sequence dataflow patch. The + # sparse path keeps CUTE's DPAS/online-softmax implementation and changes + # only which physical K/V tiles each Q workgroup visits. + set(H3_SPARSE_MAINLOOP_PATCH + "${CMAKE_CURRENT_SOURCE_DIR}/cute/patches/minimax_h3_sparse_lut.patch") + execute_process( + COMMAND patch --batch --forward --fuzz=0 + ${H3_SPARSE_OVERLAY_DIR}/xe_fmha_fwd_mainloop.hpp + ${H3_SPARSE_MAINLOOP_PATCH} + RESULT_VARIABLE H3_SPARSE_MAINLOOP_RESULT + OUTPUT_VARIABLE H3_SPARSE_MAINLOOP_STDOUT + ERROR_VARIABLE H3_SPARSE_MAINLOOP_STDERR) + if(NOT H3_SPARSE_MAINLOOP_RESULT EQUAL 0) + message(FATAL_ERROR + "Failed to apply MiniMax-H3 sparse mainloop patch:\n${H3_SPARSE_MAINLOOP_STDOUT}${H3_SPARSE_MAINLOOP_STDERR}") + endif() + set(H3_SPARSE_TAIL_PATCH + "${CMAKE_CURRENT_SOURCE_DIR}/cute/patches/minimax_h3_sparse_tail.patch") + execute_process( + COMMAND patch --batch --forward --fuzz=0 + ${H3_SPARSE_OVERLAY_DIR}/xe_fmha_fwd_mainloop.hpp ${H3_SPARSE_TAIL_PATCH} + RESULT_VARIABLE H3_SPARSE_TAIL_RESULT + OUTPUT_VARIABLE H3_SPARSE_TAIL_STDOUT + ERROR_VARIABLE H3_SPARSE_TAIL_STDERR) + if(NOT H3_SPARSE_TAIL_RESULT EQUAL 0) + message(FATAL_ERROR + "Failed to apply MiniMax-H3 sparse tail patch:\n${H3_SPARSE_TAIL_STDOUT}${H3_SPARSE_TAIL_STDERR}") + endif() + + set(H3_KERNEL_DIR "${H3_SPARSE_OVERLAY_ROOT}/flash_attention_v2/kernel") + file(MAKE_DIRECTORY "${H3_KERNEL_DIR}") + set(H3_KERNEL_SOURCE + "${CUTLASS_SYCL_ROOT}/applications/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp") + set(H3_KERNEL_PATCH + "${CMAKE_CURRENT_SOURCE_DIR}/cute/patches/minimax_h3_sparse_kernel.patch") + set(H3_KERNEL_TEMP "${H3_KERNEL_DIR}/xe_fmha_fwd_kernel.hpp.new") + execute_process( + COMMAND patch --batch --forward --fuzz=0 + --output=${H3_KERNEL_TEMP} ${H3_KERNEL_SOURCE} ${H3_KERNEL_PATCH} + RESULT_VARIABLE H3_KERNEL_PATCH_RESULT + OUTPUT_VARIABLE H3_KERNEL_PATCH_STDOUT + ERROR_VARIABLE H3_KERNEL_PATCH_STDERR) + if(NOT H3_KERNEL_PATCH_RESULT EQUAL 0) + message(FATAL_ERROR + "Failed to apply MiniMax-H3 sparse kernel patch:\n${H3_KERNEL_PATCH_STDOUT}${H3_KERNEL_PATCH_STDERR}") + endif() + configure_file("${H3_KERNEL_TEMP}" "${H3_KERNEL_DIR}/xe_fmha_fwd_kernel.hpp" COPYONLY) + file(REMOVE "${H3_KERNEL_TEMP}") file(COPY "${CUTLASS_SYCL_ROOT}/applications/flash_attention_v2/collective/fmha_fusion.hpp" DESTINATION "${H3_OVERLAY_DIR}") + file(COPY + "${CUTLASS_SYCL_ROOT}/applications/flash_attention_v2/collective/fmha_fusion.hpp" + DESTINATION "${H3_SPARSE_OVERLAY_DIR}") add_library(cute_fmha_minimax_h3_torch SHARED cute/cute_fmha_torch.cpp) set_target_properties(cute_fmha_minimax_h3_torch PROPERTIES @@ -341,6 +404,39 @@ if(ENABLE_CUTE_FMHA) "-spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate") target_link_libraries(cute_fmha_minimax_h3_torch PRIVATE ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY}) + + add_library(cute_fmha_minimax_h3_sparse_torch SHARED cute/cute_fmha_torch.cpp) + set_target_properties(cute_fmha_minimax_h3_sparse_torch PROPERTIES + PREFIX "" + BUILD_RPATH "${TORCH_INSTALL_PREFIX}/lib" + INSTALL_RPATH "\$ORIGIN;${TORCH_INSTALL_PREFIX}/lib") + target_compile_features(cute_fmha_minimax_h3_sparse_torch PRIVATE cxx_std_17) + target_compile_definitions(cute_fmha_minimax_h3_sparse_torch PRIVATE + CUTLASS_ENABLE_SYCL SYCL_INTEL_TARGET ${XPU_ARCH_DEFINE} + CUTE_FMHA_TORCH_LIBRARY=sycl_kernels_cute_minimax_h3_sparse + CUTE_FMHA_SPARSE=1) + target_include_directories(cute_fmha_minimax_h3_sparse_torch PRIVATE + "${CUTLASS_SYCL_ROOT}/include" + "${CUTLASS_SYCL_ROOT}/tools/util/include" + "${CUTLASS_SYCL_ROOT}/examples/common" + "${H3_SPARSE_OVERLAY_ROOT}" + "${CUTLASS_SYCL_ROOT}/applications" + "${CMAKE_CURRENT_SOURCE_DIR}/cute" + ${TORCH_INCLUDE_DIRS}) + target_compile_options(cute_fmha_minimax_h3_sparse_torch PRIVATE + -O3 -fsycl -fsycl-targets=spir64_gen + -Xsycl-target-backend "-device ${XPU_TARGET}" + -fno-sycl-instrument-device-code -Wno-unknown-pragmas + -Wno-unused-variable -Wno-unused-but-set-variable + -Wno-unused-local-typedef -Wno-uninitialized -Wno-reorder-ctor + -Wno-logical-op-parentheses -Wno-unused-function -Wno-deprecated-copy) + target_link_options(cute_fmha_minimax_h3_sparse_torch PRIVATE + -fsycl -fsycl-targets=spir64_gen + -Xsycl-target-backend "-device ${XPU_TARGET}" + -Xspirv-translator + "-spirv-ext=+SPV_INTEL_split_barrier,+SPV_INTEL_2d_block_io,+SPV_INTEL_subgroup_matrix_multiply_accumulate") + target_link_libraries(cute_fmha_minimax_h3_sparse_torch PRIVATE + ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY}) endif() # ── Install targets for wheel assembly ──────────────────────────────────────── @@ -352,4 +448,5 @@ install(FILES "${ESIMD_RUNTIME}" DESTINATION sycl_kernels) if(ENABLE_CUTE_FMHA) install(TARGETS cute_fmha_torch LIBRARY DESTINATION sycl_kernels) install(TARGETS cute_fmha_minimax_h3_torch LIBRARY DESTINATION sycl_kernels) + install(TARGETS cute_fmha_minimax_h3_sparse_torch LIBRARY DESTINATION sycl_kernels) endif() diff --git a/lightx2v_kernel_xpu/README.md b/lightx2v_kernel_xpu/README.md index 0dcd91e98..fda82870a 100644 --- a/lightx2v_kernel_xpu/README.md +++ b/lightx2v_kernel_xpu/README.md @@ -274,3 +274,27 @@ version `0.0.1` because CUTE FMHA is disabled there. CUTE FMHA is disabled on Windows because the current sycl-tla kernel produces incorrect attention results there. `build.bat` builds only the existing ESIMD/oneDNN extension, and `sycl_kernels.has_cute_fmha()` returns `False`. + +### MiniMax-H3 SLA block-sparse attention + +The Linux BMG build also exposes an SLA forward path for BF16 MiniMax-H3 +self-attention in `[B,L,H,128]` layout: + +```python +lut = sycl_kernels.sla_block_map(q, k, keep_ratio=0.15, block_q=128, block_k=128) +out = sycl_kernels.sparse_block_attention(q, k, v, lut) +# Or route and execute in one call: +out = sycl_kernels.sla_sparse_attention(q, k, v, keep_ratio=0.15) +``` + +The CUTE kernel fuses sparse QK, online softmax, and PV and never materializes +the score matrix. The optimized contract is forward-only, non-causal, B=1, +BF16, D=128, equal Q/K/V head counts, and 128x128 SLA blocks. Other supported +64/128 block or GQA shapes currently use a slower Triton-XPU fallback. + +Benchmark the kernel, router, and dense CUTE baseline separately with: + +```bash +ONEAPI_DEVICE_SELECTOR=level_zero:0 PYTHONPATH=python \ + python test/bench_sla_sparse_attention.py +``` diff --git a/lightx2v_kernel_xpu/build.sh b/lightx2v_kernel_xpu/build.sh index b9817dac3..d3ea4a647 100755 --- a/lightx2v_kernel_xpu/build.sh +++ b/lightx2v_kernel_xpu/build.sh @@ -71,7 +71,7 @@ mkdir -p python/sycl_kernels find _cmake_build -maxdepth 1 -type f -name '_ext*.so' -exec cp -f {} python/sycl_kernels/ \; find _cmake_build -maxdepth 1 -type f -name 'rms_norm_torch*.so' -exec cp -f {} python/sycl_kernels/ \; find _cmake_build -maxdepth 1 -type f -name 'minimax_h3_rope_torch*.so' -exec cp -f {} python/sycl_kernels/ \; -find _cmake_build -maxdepth 1 -type f -name 'cute_fmha_torch*.so' -exec cp -f {} python/sycl_kernels/ \; +find _cmake_build -maxdepth 1 -type f -name 'cute_fmha*.so' -exec cp -f {} python/sycl_kernels/ \; cp -f lgrf_uni/libesimd.unify.lgrf.so python/sycl_kernels/ if [[ "${SKIP_TESTS:-0}" != "1" ]]; then diff --git a/lightx2v_kernel_xpu/cute/cute_fmha_torch.cpp b/lightx2v_kernel_xpu/cute/cute_fmha_torch.cpp index 55e8b8a58..dff13fa94 100644 --- a/lightx2v_kernel_xpu/cute/cute_fmha_torch.cpp +++ b/lightx2v_kernel_xpu/cute/cute_fmha_torch.cpp @@ -210,6 +210,8 @@ template < void run_d128_tile( const void* q_ptr, const void* k_ptr, const void* v_ptr, void* o_ptr, int B, int H, int Lq, int Lkv, int D, float scale, + const int* block_lut = nullptr, int lut_q_blocks = 0, + int lut_topk = 0, int lut_block_tiles = 0, int64_t q_stride_seq = -1, int64_t q_stride_head = -1, int64_t q_stride_batch = -1, int64_t k_stride_seq = -1, int64_t k_stride_head = -1, int64_t k_stride_batch = -1, @@ -294,7 +296,11 @@ void run_d128_tile( nullptr, stride_K, // k_cache nullptr, stride_V, // v_cache }, - {scale, nullptr, 0, nullptr}, + {scale, nullptr, 0, nullptr, +#if defined(CUTE_FMHA_SPARSE) + block_lut, H, lut_q_blocks, lut_topk, lut_block_tiles +#endif + }, {}, hw_info}; @@ -378,12 +384,63 @@ at::Tensor sdp(const at::Tensor& q, const at::Tensor& k, const at::Tensor& v) { return o; } +#if defined(CUTE_FMHA_SPARSE) +at::Tensor sparse_sdp( + const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + const at::Tensor& block_lut) { + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "sycl_kernels sparse CUTE FMHA: expect q/k/v [B,L,H,D]"); + TORCH_CHECK(block_lut.dim() == 4, + "sycl_kernels sparse CUTE FMHA: expect LUT [B,H,Qblocks,topk]"); + TORCH_CHECK(q.device().is_xpu() && k.device().is_xpu() && v.device().is_xpu() && + block_lut.device().is_xpu(), + "sycl_kernels sparse CUTE FMHA: all tensors must be on XPU"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16, + "sycl_kernels sparse CUTE FMHA: q/k/v must be BF16"); + TORCH_CHECK(block_lut.scalar_type() == at::kInt, + "sycl_kernels sparse CUTE FMHA: LUT must be int32"); + + const int B = checked_int(q.size(0), "batch"); + const int L = checked_int(q.size(1), "sequence length"); + const int H = checked_int(q.size(2), "head count"); + const int D = checked_int(q.size(3), "head dimension"); + constexpr int BlockQ = 128; + constexpr int BlockK = 128; + constexpr int CuteKvTile = 32; + const int q_blocks = (L + BlockQ - 1) / BlockQ; + TORCH_CHECK(B == 1 && D == 128, + "sycl_kernels sparse CUTE FMHA: only B=1,D=128 are supported"); + TORCH_CHECK(k.sizes() == q.sizes() && v.sizes() == q.sizes(), + "sycl_kernels sparse CUTE FMHA: q/k/v shapes must match"); + TORCH_CHECK(block_lut.size(0) == B && block_lut.size(1) == H && + block_lut.size(2) == q_blocks && block_lut.size(3) > 0, + "sycl_kernels sparse CUTE FMHA: invalid LUT shape"); + + auto qc = q.contiguous(), kc = k.contiguous(), vc = v.contiguous(); + auto lut = block_lut.contiguous(); + auto output = at::empty_like(qc); + const float scale = 1.0f / std::sqrt(static_cast(D)); + run_d128_tile( + qc.data_ptr(), kc.data_ptr(), vc.data_ptr(), output.data_ptr(), + B, H, L, L, D, scale, lut.const_data_ptr(), q_blocks, + checked_int(lut.size(3), "LUT topk"), BlockK / CuteKvTile); + return output; +} +#endif + } // namespace TORCH_LIBRARY(CUTE_FMHA_TORCH_LIBRARY, m) { m.def("sdp(Tensor q, Tensor k, Tensor v) -> Tensor"); +#if defined(CUTE_FMHA_SPARSE) + m.def("sparse_sdp(Tensor q, Tensor k, Tensor v, Tensor block_lut) -> Tensor"); +#endif } TORCH_LIBRARY_IMPL(CUTE_FMHA_TORCH_LIBRARY, XPU, m) { m.impl("sdp", &sdp); +#if defined(CUTE_FMHA_SPARSE) + m.impl("sparse_sdp", &sparse_sdp); +#endif } diff --git a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch new file mode 100644 index 000000000..b6e75bb83 --- /dev/null +++ b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch @@ -0,0 +1,21 @@ +--- xe_fmha_fwd_kernel.hpp ++++ xe_fmha_fwd_kernel.hpp +@@ -226,7 +226,9 @@ + if (CollectiveMainloop::CausalMask && seq_coord < discard_seq_coord) continue; + const int seq_len_new = CollectiveMainloop::CausalMask ? full_tile_offset + cute::min(seq_len_kv, seq_coord - discard_seq_coord) + q_sg_tile : seq_len_kv; + const int seq_len = seq_len_new + seq_len_kv_cache; +- const int k_blocks = cute::ceil_div(seq_len, get<1>(TileShapeQK{})); ++ const int k_blocks = params.mainloop.ptr_block_lut ++ ? params.mainloop.lut_topk * params.mainloop.lut_block_tiles ++ : cute::ceil_div(seq_len, get<1>(TileShapeQK{})); + + int offset_q = 0, offset_k = 0, offset_v = 0, offset_o = 0; + int offset_k_cache = 0, offset_v_cache = 0; +@@ -287,7 +289,7 @@ + tArA, tA_max, tA_sum, + blk_qv, 0, k_blocks, k_blocks, +- thr_id, seq_len, seq_len_kv_cache, idx_b, ++ thr_id, seq_len, seq_len_kv_cache, idx_b, head_q, + full_tile_offset, discard_seq_coord, + K_cache(_,_,head,l_coord), + V_cache(_,_,head,l_coord)); diff --git a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch new file mode 100644 index 000000000..91ad36793 --- /dev/null +++ b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch @@ -0,0 +1,92 @@ +--- xe_fmha_fwd_mainloop.hpp ++++ xe_fmha_fwd_mainloop.hpp +@@ -150,6 +150,11 @@ + int const* ptr_page_table = nullptr; + int page_size = 0; + int const* num_pages_per_seq = nullptr; ++ int const* ptr_block_lut = nullptr; ++ int lut_num_heads = 0; ++ int lut_q_blocks = 0; ++ int lut_topk = 0; ++ int lut_block_tiles = 0; + }; + + // Kernel-facing parameters +@@ -172,7 +177,9 @@ + Params to_underlying_arguments(Arguments const &args, void * /* workspace */) { + constexpr double kLog2e = 1.4426950408889634074; // log_2(e) + ElementS val = args.scale * static_cast(kLog2e); +- return Params{val, args.ptr_page_table, args.page_size, args.num_pages_per_seq}; ++ return Params{val, args.ptr_page_table, args.page_size, args.num_pages_per_seq, ++ args.ptr_block_lut, args.lut_num_heads, args.lut_q_blocks, ++ args.lut_topk, args.lut_block_tiles}; + } + + CUTLASS_HOST_DEVICE static +@@ -211,6 +218,7 @@ + int seq_len, + int seq_len_kv_cache, + int l_coord, ++ int head_q, + int full_tile_offset, + int discard_seq_coord, + TensorK_cache2D const& K_cache_2D = TensorK_cache2D{}, +@@ -308,6 +316,15 @@ + /* Initialization steps for first block: Q/K prefetch, O init */ + /* TODO: limit D prefetch for large head size, and reorder K prefetches */ + int kblocks_cache = ceil_div(seq_len_kv_cache, get<1>(TileShapeQK{})); ++ auto physical_k_tile = [&](int logical_k) { ++ if (params.ptr_block_lut == nullptr) return logical_k; ++ int slot = logical_k / params.lut_block_tiles; ++ int sub_tile = logical_k - slot * params.lut_block_tiles; ++ int q_block = get<0>(blk_qv); ++ int lut_offset = ((l_coord * params.lut_num_heads + head_q) * ++ params.lut_q_blocks + q_block) * params.lut_topk; ++ return params.ptr_block_lut[lut_offset + slot] * params.lut_block_tiles + sub_tile; ++ }; + for (int D = 0; D < size<3>(pQgQ); D++) { + prefetch(prefetch_q, pQgQ(_,_,_,D)); + } +@@ -323,7 +340,9 @@ + prefetch(prefetch_k_cache, pKgK_cache(_,_,_,K,D)); + } + } else { +- prefetch(prefetch_k, pKgK(_,_,_,K - kblocks_cache,D)); ++ int physical_K = physical_k_tile(K - kblocks_cache); ++ if (physical_K < ceil_div(seq_len, get<1>(TileShapeQK{}))) ++ prefetch(prefetch_k, pKgK(_,_,_,physical_K,D)); + } + } + } +@@ -355 +374 @@ +- k_idx = K - kblocks_cache; ++ k_idx = physical_k_tile(K - kblocks_cache); +@@ -356,0 +376,7 @@ ++ if constexpr (!is_cache) { ++ if (params.ptr_block_lut != nullptr && ++ k_idx >= ceil_div(seq_len, get<1>(TileShapeQK{}))) { ++ barrier_wait(ScopeWorkgroup); ++ return; ++ } ++ } +@@ -399 +425,2 @@ +- if (check_remainder_k && K == total_blk - 1) { ++ int dense_total_blk = ceil_div(seq_len, get<1>(TileShapeQK{})); ++ if (check_remainder_k && k_idx == dense_total_blk - 1) { +@@ -401,2 +428 @@ +- int k_val = get<0>(tKgK_cur(0,0,0,k_idx,0)) + kblocks_cache * get<1>(TileShapeQK{}); +- int k = k_val + get_sub_group().get_local_id()[0]; ++ int k = k_idx * get<1>(TileShapeQK{}) + get_sub_group().get_local_id()[0]; +@@ -449,8 +476,10 @@ + } else { + prefetch(prefetch_k, pKgK(_,_,_,K_next-kblocks_cache,D)); + } +- } else { +- prefetch(prefetch_k, pKgK(_,_,_,K_next-kblocks_cache,D)); ++ } else if (K_next < blk_k1) { ++ int physical_K_next = physical_k_tile(K_next-kblocks_cache); ++ if (physical_K_next < ceil_div(seq_len, get<1>(TileShapeQK{}))) ++ prefetch(prefetch_k, pKgK(_,_,_,physical_K_next,D)); + } + } + barrier_wait(ScopeWorkgroup); diff --git a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_tail.patch b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_tail.patch new file mode 100644 index 000000000..8174bedf1 --- /dev/null +++ b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_tail.patch @@ -0,0 +1,18 @@ +--- xe_fmha_fwd_mainloop.hpp ++++ xe_fmha_fwd_mainloop.hpp +@@ -427,10 +427,8 @@ +- FragSRow k_rem_mask; +- int k = k_idx * get<1>(TileShapeQK{}) + get_sub_group().get_local_id()[0]; ++ Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); ++ Tensor gP = local_tile(cPgP, take<0,2>(TileShapeQK{}), make_coord(get<0>(blk_qv), k_idx)); ++ auto cS_thread = thr_mma_qk.partition_C(gP); + CUTLASS_PRAGMA_UNROLL +- for (int i = 0; i < k_rem_mask.size(); i++, k += intel::sg_size) { +- k_rem_mask(i) = (k < seq_len) ? ElementS(sycl::nan(0u)) : ElementS(-INFINITY); +- } +- CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tSrS.size(); i++) { +- tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); ++ int col_idx = get<1>(cS_thread(i)); ++ if (col_idx >= seq_len) tSrS(i) = ElementS(-INFINITY); + } diff --git a/lightx2v_kernel_xpu/python/sycl_kernels/__init__.py b/lightx2v_kernel_xpu/python/sycl_kernels/__init__.py index 8a721829d..97db6a0cf 100644 --- a/lightx2v_kernel_xpu/python/sycl_kernels/__init__.py +++ b/lightx2v_kernel_xpu/python/sycl_kernels/__init__.py @@ -6,6 +6,7 @@ _pkg_dir = os.path.dirname(os.path.abspath(__file__)) _cute_fmha_loaded = False _cute_fmha_minimax_h3_loaded = False +_cute_fmha_minimax_h3_sparse_loaded = False _rms_norm_loaded = False _minimax_h3_rope_loaded = False @@ -174,6 +175,19 @@ def _load_cute_fmha_minimax_h3(): _cute_fmha_minimax_h3_loaded = True +def _load_cute_fmha_minimax_h3_sparse(): + global _cute_fmha_minimax_h3_sparse_loaded + if _cute_fmha_minimax_h3_sparse_loaded: + return + import torch + + candidates = sorted(glob.glob(os.path.join(_pkg_dir, "cute_fmha_minimax_h3_sparse_torch*.so"))) + if not candidates: + raise ImportError("cute_fmha_minimax_h3_sparse_torch.so not found") + torch.ops.load_library(candidates[0]) + _cute_fmha_minimax_h3_sparse_loaded = True + + def _use_minimax_h3_cute(q, k, v): import torch @@ -222,3 +236,24 @@ def has_cute_fmha(): return True except (ImportError, OSError, RuntimeError): return False + + +def sla_block_map(q, k, keep_ratio=0.2, block_q=128, block_k=128): + """Build an SLA block LUT for BLHD query/key tensors.""" + from .sla import sla_block_map as _sla_block_map + + return _sla_block_map(q, k, keep_ratio, block_q, block_k) + + +def sparse_block_attention(q, k, v, lut, block_q=128, block_k=128, scale=None): + """Run fused BF16 block-sparse attention on Intel XPU.""" + from .sla import sparse_block_attention as _sparse_block_attention + + return _sparse_block_attention(q, k, v, lut, block_q, block_k, scale) + + +def sla_sparse_attention(q, k, v, keep_ratio=0.2, block_q=128, block_k=128, scale=None): + """Build the SLA routing LUT and run fused sparse attention.""" + from .sla import sla_sparse_attention as _sla_sparse_attention + + return _sla_sparse_attention(q, k, v, keep_ratio, block_q, block_k, scale) diff --git a/lightx2v_kernel_xpu/python/sycl_kernels/sla.py b/lightx2v_kernel_xpu/python/sycl_kernels/sla.py new file mode 100644 index 000000000..f99255000 --- /dev/null +++ b/lightx2v_kernel_xpu/python/sycl_kernels/sla.py @@ -0,0 +1,193 @@ +"""SLA routing and fused block-sparse attention for Intel XPU. + +The public layout is BLHD, matching :func:`sycl_kernels.cute_sdp` and the +MiniMax-H3 attention path. Routing is intentionally expressed with PyTorch +XPU operations; the optimized QK/softmax/PV path is one CUTE kernel and does +not materialize the sparse score matrix. Triton is retained as a fallback. +""" + +from __future__ import annotations + +import math + +import torch + + +def _block_mean_blhd(x: torch.Tensor, block_size: int) -> torch.Tensor: + """Return block means as [B, H, ceil(L / block_size), D].""" + if x.ndim != 4: + raise ValueError(f"expected a BLHD tensor, got shape {tuple(x.shape)}") + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + + batch, length, heads, dim = x.shape + full_blocks, tail = divmod(length, block_size) + pieces = [] + if full_blocks: + prefix = x[:, : full_blocks * block_size] + pieces.append(prefix.reshape(batch, full_blocks, block_size, heads, dim).mean(dim=2)) + if tail: + pieces.append(x[:, full_blocks * block_size :].mean(dim=1, keepdim=True)) + if not pieces: + raise ValueError("sequence length must be non-zero") + # [B, blocks, H, D] -> [B, H, blocks, D]. The small pooled tensor is + # made contiguous because it is consumed by batched matmul immediately. + return torch.cat(pieces, dim=1).permute(0, 2, 1, 3).contiguous() + + +def sla_block_map( + q: torch.Tensor, + k: torch.Tensor, + keep_ratio: float = 0.2, + block_q: int = 128, + block_k: int = 128, +) -> torch.Tensor: + """Build the SLA key-block LUT for BLHD Q and K tensors. + + Returns int32 ``[B, Hq, ceil(Lq/block_q), topk]`` indices. Smooth-K is + applied in pooled form: ``mean(block(K)) - mean(sequence(K))``. This is + algebraically identical to centering the full K tensor first, while + avoiding an additional full-sequence allocation. + """ + if q.ndim != 4 or k.ndim != 4: + raise ValueError("q and k must use [B, L, H, D] layout") + if q.device != k.device or q.dtype != k.dtype: + raise ValueError("q and k must share device and dtype") + if q.shape[0] != k.shape[0] or q.shape[3] != k.shape[3]: + raise ValueError("q and k batch/head_dim must match") + if not 0.0 < keep_ratio <= 1.0: + raise ValueError(f"keep_ratio must be in (0, 1], got {keep_ratio}") + hq, hkv = q.shape[2], k.shape[2] + if hq % hkv != 0: + raise ValueError(f"query heads ({hq}) must be divisible by KV heads ({hkv})") + + pooled_q = _block_mean_blhd(q, block_q) + pooled_k = _block_mean_blhd(k, block_k) + pooled_k = pooled_k - k.mean(dim=1).unsqueeze(2) + if hq != hkv: + pooled_k = pooled_k.repeat_interleave(hq // hkv, dim=1) + + scores = torch.matmul(pooled_q, pooled_k.transpose(-1, -2)) + key_blocks = scores.shape[-1] + topk = max(1, min(key_blocks, int(keep_ratio * key_blocks))) + # Sorting is unnecessary for online softmax and costs a measurable amount + # at long sequence lengths. int32 halves LUT traffic in the fused kernel. + return torch.topk(scores, topk, dim=-1, sorted=False).indices.to(torch.int32).contiguous() + + +def sparse_block_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + lut: torch.Tensor, + block_q: int = 128, + block_k: int = 128, + scale: float | None = None, +) -> torch.Tensor: + """Slow, device-independent reference for tests and bring-up.""" + _validate_sparse_inputs(q, k, v, lut, block_q, block_k) + batch, q_len, q_heads, dim = q.shape + kv_len, kv_heads = k.shape[1], k.shape[2] + group_size = q_heads // kv_heads + scale = dim**-0.5 if scale is None else scale + out = torch.empty_like(q) + offsets = torch.arange(block_k, device=lut.device, dtype=torch.long) + + for b in range(batch): + for h in range(q_heads): + kv_h = h // group_size + for qb in range(lut.shape[2]): + q_start = qb * block_q + q_stop = min(q_start + block_q, q_len) + block_ids = lut[b, h, qb].long() + positions = (block_ids[:, None] * block_k + offsets[None, :]).reshape(-1) + positions = positions[positions < kv_len] + q_tile = q[b, q_start:q_stop, h].float() + k_tile = k[b, positions, kv_h].float() + v_tile = v[b, positions, kv_h].float() + probs = torch.softmax(torch.matmul(q_tile, k_tile.transpose(0, 1)) * scale, dim=-1) + out[b, q_start:q_stop, h] = torch.matmul(probs, v_tile).to(out.dtype) + return out + + +def _validate_sparse_inputs(q, k, v, lut, block_q: int, block_k: int) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k and v must use [B, L, H, D] layout") + if not (q.device == k.device == v.device == lut.device): + raise ValueError("q, k, v and lut must be on the same device") + if not (q.dtype == k.dtype == v.dtype): + raise ValueError("q, k and v must share dtype") + if q.shape[0] != k.shape[0] or k.shape != v.shape: + raise ValueError("batch must match and k/v shapes must be identical") + if q.shape[3] != k.shape[3]: + raise ValueError("q and k/v head_dim must match") + if q.shape[2] % k.shape[2] != 0: + raise ValueError("query head count must be divisible by KV head count") + expected_q_blocks = math.ceil(q.shape[1] / block_q) + if lut.ndim != 4 or tuple(lut.shape[:3]) != (q.shape[0], q.shape[2], expected_q_blocks): + raise ValueError( + "lut must have shape [B, Hq, ceil(Lq/block_q), topk], got " + f"{tuple(lut.shape)}" + ) + if lut.shape[3] == 0: + raise ValueError("lut topk dimension must be non-zero") + if block_q <= 0 or block_k <= 0: + raise ValueError("block sizes must be positive") + + +def sparse_block_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + lut: torch.Tensor, + block_q: int = 128, + block_k: int = 128, + scale: float | None = None, +) -> torch.Tensor: + """Run fused block-sparse attention on Intel XPU. + + QK, online softmax, and PV are fused. No score tensor is written to + global memory. The first optimized contract is BF16, D=128, forward-only. + """ + _validate_sparse_inputs(q, k, v, lut, block_q, block_k) + if q.device.type != "xpu": + raise ValueError("sparse_block_attention requires Intel XPU tensors") + if q.dtype != torch.bfloat16 or q.shape[-1] != 128: + raise ValueError("the optimized XPU path currently requires BF16 and head_dim=128") + if block_q not in (64, 128) or block_k not in (64, 128): + raise ValueError("the optimized XPU path supports block sizes 64 or 128") + + if block_q == 128 and block_k == 128 and scale is None and q.shape[2] == k.shape[2]: + # The CUTE path preserves the tuned dense kernel's 2D block loads, + # DPAS pipeline, and online softmax; the LUT only redirects K/V tiles. + from . import _load_cute_fmha_minimax_h3_sparse + + try: + op = torch.ops.sycl_kernels_cute_minimax_h3_sparse.sparse_sdp + except AttributeError: + _load_cute_fmha_minimax_h3_sparse() + op = torch.ops.sycl_kernels_cute_minimax_h3_sparse.sparse_sdp + return op(q, k, v, lut.to(torch.int32).contiguous()) + + # Triton remains as a correctness-oriented fallback for GQA, custom scale, + # and 64-sized blocks while their specialized CUTE variants are developed. + from .sla_triton import launch_sparse_block_attention + + return launch_sparse_block_attention( + q.contiguous(), k.contiguous(), v.contiguous(), lut.to(torch.int32).contiguous(), + block_q, block_k, q.shape[-1] ** -0.5 if scale is None else scale, + ) + + +def sla_sparse_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + keep_ratio: float = 0.2, + block_q: int = 128, + block_k: int = 128, + scale: float | None = None, +) -> torch.Tensor: + """Convenience entry point combining SLA routing and sparse attention.""" + lut = sla_block_map(q, k, keep_ratio, block_q, block_k) + return sparse_block_attention(q, k, v, lut, block_q, block_k, scale) diff --git a/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py b/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py new file mode 100644 index 000000000..c216bc455 --- /dev/null +++ b/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py @@ -0,0 +1,91 @@ +"""Triton-XPU implementation detail for fused SLA block attention.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _sparse_block_attention_fwd( + q_ptr, + k_ptr, + v_ptr, + lut_ptr, + out_ptr, + q_len: tl.constexpr, + kv_len: tl.constexpr, + q_heads: tl.constexpr, + kv_heads: tl.constexpr, + q_blocks: tl.constexpr, + topk: tl.constexpr, + scale, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + qb = tl.program_id(0) + bh = tl.program_id(1) + batch = bh // q_heads + q_head = bh - batch * q_heads + kv_head = q_head // (q_heads // kv_heads) + + offs_q = qb * BLOCK_Q + tl.arange(0, BLOCK_Q) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, HEAD_DIM) + + # Internally Q/K/V are BHLD so every per-head matrix tile is contiguous. + # Direct BLHD access makes successive rows H*D apart and prevents Triton + # XPU from forming efficient block loads / DPAS operands. + q_base = (batch * q_heads + q_head) * q_len * HEAD_DIM + kv_base = (batch * kv_heads + kv_head) * kv_len * HEAD_DIM + q_ptrs = q_ptr + q_base + offs_q[:, None] * HEAD_DIM + offs_d[None, :] + q = tl.load(q_ptrs, mask=offs_q[:, None] < q_len, other=0.0) + + row_max = tl.full([BLOCK_Q], -float("inf"), tl.float32) + row_sum = tl.zeros([BLOCK_Q], tl.float32) + acc = tl.zeros([BLOCK_Q, HEAD_DIM], tl.float32) + lut_base = (bh * q_blocks + qb) * topk + log2e: tl.constexpr = 1.4426950408889634 + + for slot in tl.range(0, topk, num_stages=2): + key_block = tl.load(lut_ptr + lut_base + slot) + key_pos = key_block * BLOCK_K + offs_k + valid_k = key_pos < kv_len + k_ptrs = k_ptr + kv_base + key_pos[:, None] * HEAD_DIM + offs_d[None, :] + v_ptrs = v_ptr + kv_base + key_pos[:, None] * HEAD_DIM + offs_d[None, :] + k_tile = tl.load(k_ptrs, mask=valid_k[:, None], other=0.0) + v_tile = tl.load(v_ptrs, mask=valid_k[:, None], other=0.0) + + scores = tl.dot(q, tl.trans(k_tile)).to(tl.float32) * (scale * log2e) + scores = tl.where(valid_k[None, :], scores, -float("inf")) + tile_max = tl.max(scores, axis=1) + new_max = tl.maximum(row_max, tile_max) + alpha = tl.exp2(row_max - new_max) + probs = tl.exp2(scores - new_max[:, None]) + + acc *= alpha[:, None] + acc += tl.dot(probs.to(v_tile.dtype), v_tile) + row_sum = row_sum * alpha + tl.sum(probs, axis=1) + row_max = new_max + + acc /= row_sum[:, None] + out_ptrs = out_ptr + q_base + offs_q[:, None] * HEAD_DIM + offs_d[None, :] + tl.store(out_ptrs, acc, mask=offs_q[:, None] < q_len) + + +def launch_sparse_block_attention(q, k, v, lut, block_q, block_k, scale): + batch, q_len, q_heads, head_dim = q.shape + kv_len, kv_heads = k.shape[1], k.shape[2] + q_blocks, topk = lut.shape[2], lut.shape[3] + q_bhld = q.permute(0, 2, 1, 3).contiguous() + k_bhld = k.permute(0, 2, 1, 3).contiguous() + v_bhld = v.permute(0, 2, 1, 3).contiguous() + output = torch.empty_like(q_bhld) + grid = (q_blocks, batch * q_heads) + _sparse_block_attention_fwd[grid]( + q_bhld, k_bhld, v_bhld, lut, output, + q_len, kv_len, q_heads, kv_heads, q_blocks, topk, scale, + BLOCK_Q=block_q, BLOCK_K=block_k, HEAD_DIM=head_dim, + num_warps=8, num_stages=3, + ) + return output.permute(0, 2, 1, 3) diff --git a/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py b/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py new file mode 100644 index 000000000..c48f539e4 --- /dev/null +++ b/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +import argparse +import json +import statistics +import time +from pathlib import Path + +import torch + + +def measure(fn, warmup, iterations): + for _ in range(warmup): + fn() + torch.xpu.synchronize() + samples = [] + for _ in range(iterations): + start = time.perf_counter() + fn() + torch.xpu.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return samples + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--build-dir", default="_cmake_build") + parser.add_argument("--sequence-length", type=int, default=19292) + parser.add_argument("--heads", type=int, default=7) + parser.add_argument("--keep-ratio", type=float, default=0.15) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=7) + args = parser.parse_args() + + build_dir = Path(args.build_dir).resolve() + torch.ops.load_library(str(build_dir / "cute_fmha_minimax_h3_sparse_torch.so")) + torch.ops.load_library(str(build_dir / "cute_fmha_minimax_h3_torch.so")) + from sycl_kernels.sla import sla_block_map + + shape = (1, args.sequence_length, args.heads, 128) + q = torch.randn(shape, device="xpu", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + lut = sla_block_map(q, k, args.keep_ratio, 128, 128) + + sparse_ms = measure( + lambda: torch.ops.sycl_kernels_cute_minimax_h3_sparse.sparse_sdp(q, k, v, lut), + args.warmup, + args.iterations, + ) + router_ms = measure( + lambda: sla_block_map(q, k, args.keep_ratio, 128, 128), + args.warmup, + args.iterations, + ) + dense_ms = measure( + lambda: torch.ops.sycl_kernels_cute_minimax_h3.sdp(q, k, v), + args.warmup, + args.iterations, + ) + sparse_median = statistics.median(sparse_ms) + router_median = statistics.median(router_ms) + dense_median = statistics.median(dense_ms) + print(json.dumps({ + "shape": shape, + "keep_ratio": args.keep_ratio, + "lut_shape": list(lut.shape), + "dense_ms": dense_ms, + "sparse_kernel_ms": sparse_ms, + "router_ms": router_ms, + "kernel_speedup": dense_median / sparse_median, + "attention_speedup_including_router": dense_median / (sparse_median + router_median), + }, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py b/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py new file mode 100644 index 000000000..15ec828df --- /dev/null +++ b/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py @@ -0,0 +1,69 @@ +import math + +import pytest +import torch + +from sycl_kernels.sla import ( + sla_block_map, + sparse_block_attention, + sparse_block_attention_reference, +) + + +def _explicit_router(q, k, keep_ratio, block_q, block_k): + # Small test-only implementation matching LightX2V's original BHLD SLA. + q_bhld = q.permute(0, 2, 1, 3) + k_bhld = k.permute(0, 2, 1, 3) + k_bhld = k_bhld - k_bhld.mean(dim=2, keepdim=True) + + def pool(x, block): + chunks = [x[:, :, start : start + block].mean(dim=2) for start in range(0, x.shape[2], block)] + return torch.stack(chunks, dim=2) + + pq, pk = pool(q_bhld, block_q), pool(k_bhld, block_k) + if pq.shape[1] != pk.shape[1]: + pk = pk.repeat_interleave(pq.shape[1] // pk.shape[1], dim=1) + scores = pq @ pk.transpose(-1, -2) + topk = max(1, min(scores.shape[-1], int(keep_ratio * scores.shape[-1]))) + return torch.topk(scores, topk, dim=-1, sorted=False).indices + + +@pytest.mark.parametrize("length", [17, 32, 35]) +def test_sla_router_matches_original_semantics(length): + torch.manual_seed(3) + q = torch.randn(1, length, 4, 8) + k = torch.randn(1, length, 2, 8) + actual = sla_block_map(q, k, keep_ratio=0.5, block_q=8, block_k=8) + expected = _explicit_router(q, k, 0.5, 8, 8) + # topk(sorted=False) order is not contractual; compare selected sets. + torch.testing.assert_close(actual.sort(dim=-1).values.long(), expected.sort(dim=-1).values) + + +def test_full_lut_reference_matches_dense_gqa(): + torch.manual_seed(4) + q = torch.randn(1, 19, 4, 16, dtype=torch.bfloat16) + k = torch.randn(1, 19, 2, 16, dtype=torch.bfloat16) + v = torch.randn_like(k) + block = 8 + blocks = math.ceil(q.shape[1] / block) + lut = torch.arange(blocks, dtype=torch.int32).view(1, 1, 1, blocks).expand(1, 4, blocks, blocks) + actual = sparse_block_attention_reference(q, k, v, lut, block, block) + expected = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.repeat_interleave(2, dim=2).transpose(1, 2), + v.repeat_interleave(2, dim=2).transpose(1, 2), + ).transpose(1, 2) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.xpu.is_available(), reason="XPU is unavailable") +def test_xpu_sparse_kernel_matches_reference(): + torch.manual_seed(5) + q = torch.randn(1, 257, 4, 128, device="xpu", dtype=torch.bfloat16) + k = torch.randn(1, 257, 2, 128, device="xpu", dtype=torch.bfloat16) + v = torch.randn_like(k) + lut = sla_block_map(q, k, keep_ratio=0.5, block_q=128, block_k=128) + actual = sparse_block_attention(q, k, v, lut, 128, 128) + expected = sparse_block_attention_reference(q, k, v, lut, 128, 128) + torch.xpu.synchronize() + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=3e-2) diff --git a/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py b/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py new file mode 100644 index 000000000..14d3d8511 --- /dev/null +++ b/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py @@ -0,0 +1,79 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightx2v_platform.ops.attn.intel_xpu import xpu_sla_attn + + +def test_three_sycl_sla_apis_are_forwarded(monkeypatch): + calls = [] + + def record(name): + def call(*args): + calls.append((name, args)) + return args[0] + + return call + + monkeypatch.setattr( + xpu_sla_attn, + "_sycl_kernels", + SimpleNamespace( + sla_block_map=record("map"), + sparse_block_attention=record("sparse"), + sla_sparse_attention=record("combined"), + ), + ) + q = torch.empty(1, 9, 2, 4) + lut = torch.empty(1, 2, 1, 1, dtype=torch.int32) + + xpu_sla_attn.sla_block_map(q, q, 0.15, 128, 128) + xpu_sla_attn.sparse_block_attention(q, q, q, lut, 128, 128, None) + xpu_sla_attn.sla_sparse_attention(q, q, q, 0.15, 128, 128, None) + + assert [name for name, _ in calls] == ["map", "sparse", "combined"] + assert calls[0][1][2:] == (0.15, 128, 128) + assert calls[2][1][3:] == (0.15, 128, 128, None) + + +def test_missing_sycl_sla_api_has_actionable_error(monkeypatch): + monkeypatch.setattr(xpu_sla_attn, "_sycl_kernels", SimpleNamespace()) + with pytest.raises(RuntimeError, match="rebuild/install lightx2v_kernel_xpu"): + xpu_sla_attn.sla_block_map(torch.empty(1), torch.empty(1)) + + +def test_dynamic_sparse_attention_uses_xpu_combined_api(monkeypatch): + from lightx2v.common.ops.attn.dynamic_sparse_attn import DynamicSparseAttnWeight + + def fail_if_cuda_is_queried(*args, **kwargs): + raise AssertionError("the Intel XPU backend must not query CUDA") + + received = {} + + def fake_sla(q, k, v, keep_ratio, block_q, block_k, scale=None): + received.update( + shape=tuple(q.shape), + keep_ratio=keep_ratio, + block_q=block_q, + block_k=block_k, + scale=scale, + ) + return q + + monkeypatch.setattr(torch.cuda, "current_device", fail_if_cuda_is_queried) + monkeypatch.setattr(xpu_sla_attn, "sla_sparse_attention", fake_sla) + attention = DynamicSparseAttnWeight({"operator": "intel_xpu", "sparsity_ratio": 0.85}) + q = torch.empty(17, 7, 128, dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, 17], dtype=torch.int32) + + output = attention.apply(q, q, q, cu_seqlens, cu_seqlens, 17, 17) + + assert output.shape == (17, 7 * 128) + assert received == { + "shape": (1, 17, 7, 128), + "keep_ratio": pytest.approx(0.15), + "block_q": 128, + "block_k": 128, + "scale": None, + } diff --git a/lightx2v_platform/ops/attn/intel_xpu/__init__.py b/lightx2v_platform/ops/attn/intel_xpu/__init__.py index c32418887..a74fd5fd3 100644 --- a/lightx2v_platform/ops/attn/intel_xpu/__init__.py +++ b/lightx2v_platform/ops/attn/intel_xpu/__init__.py @@ -1,2 +1,3 @@ from .xpu_cute_attn import * from .xpu_flash_attn import * +from .xpu_sla_attn import * diff --git a/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py new file mode 100644 index 000000000..fac8d3342 --- /dev/null +++ b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py @@ -0,0 +1,41 @@ +"""Intel XPU adapters for the SLA routing and sparse-attention kernels.""" + +from __future__ import annotations + +try: + import sycl_kernels as _sycl_kernels +except (ImportError, OSError) as exc: + _sycl_kernels = None + _IMPORT_ERROR = exc +else: + _IMPORT_ERROR = None + + +def _get_sycl_api(name): + if _sycl_kernels is None: + raise RuntimeError( + "Intel XPU SLA requires lightx2v_kernel_xpu's sycl-kernels package " + "with SLA/CUTE sparse attention enabled" + ) from _IMPORT_ERROR + api = getattr(_sycl_kernels, name, None) + if api is None: + raise RuntimeError( + f"sycl_kernels.{name} is unavailable; rebuild/install " + "lightx2v_kernel_xpu with SLA/CUTE sparse attention enabled" + ) + return api + + +def sla_block_map(q, k, keep_ratio=0.2, block_q=128, block_k=128): + """Build an SLA LUT for BLHD tensors using ``sycl_kernels.sla_block_map``.""" + return _get_sycl_api("sla_block_map")(q, k, keep_ratio, block_q, block_k) + + +def sparse_block_attention(q, k, v, lut, block_q=128, block_k=128, scale=None): + """Apply fused sparse QK/softmax/PV using a precomputed SLA LUT.""" + return _get_sycl_api("sparse_block_attention")(q, k, v, lut, block_q, block_k, scale) + + +def sla_sparse_attention(q, k, v, keep_ratio=0.2, block_q=128, block_k=128, scale=None): + """Route and apply sparse attention through ``sycl_kernels.sla_sparse_attention``.""" + return _get_sycl_api("sla_sparse_attention")(q, k, v, keep_ratio, block_q, block_k, scale) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh new file mode 100755 index 000000000..b0cafd376 --- /dev/null +++ b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) + +lightx2v_path=${LIGHTX2V_PATH:-${REPO_ROOT}} +model_path=${MODEL_PATH:-/llm/models/MiniMax-H3} +lora_path=${LORA_PATH:-/llm/models/Minimax-h3-Turbo-SLA/minimax_h3_fl2v_turbo_4step_v0.1_768p_sla_bf16.safetensors} +config_template=${CONFIG_JSON:-${lightx2v_path}/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json} +first_frame=${FIRST_FRAME:-${lightx2v_path}/assets/inputs/imgs/flf2v_input_first_frame-fs8.png} +last_frame=${LAST_FRAME:-${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_frame-fs8.png} +output_path=${OUTPUT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_fl2av_turbo_sla_4step.mp4} +prompt=${PROMPT:-Create a coherent cinematic transition between the two frames with natural synchronized ambient sound.} +seed=${SEED:-42} + +export ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-0} +export PLATFORM=${PLATFORM:-intel_xpu} +export PYTHONFAULTHANDLER=${PYTHONFAULTHANDLER:-1} +export PYTHONUNBUFFERED=${PYTHONUNBUFFERED:-1} +export PYTHONPATH=${PYTHONPATH:-} + +[[ -d "${model_path}" ]] || { echo "Model directory not found: ${model_path}" >&2; exit 1; } +[[ -f "${lora_path}" ]] || { echo "SLA LoRA checkpoint not found: ${lora_path}" >&2; exit 1; } +[[ -f "${config_template}" ]] || { echo "Config file not found: ${config_template}" >&2; exit 1; } +[[ -f "${first_frame}" ]] || { echo "First frame not found: ${first_frame}" >&2; exit 1; } +[[ -f "${last_frame}" ]] || { echo "Last frame not found: ${last_frame}" >&2; exit 1; } + +mkdir -p "$(dirname -- "${output_path}")" +runtime_config=$(mktemp "${TMPDIR:-/tmp}/lightx2v-minimax-h3-sla-XXXXXX.json") +trap 'rm -f -- "${runtime_config}"' EXIT + +# Keep the checked-in config reusable while allowing LORA_PATH to override the +# local checkpoint location without editing JSON. +python - "${config_template}" "${runtime_config}" "${lora_path}" <<'PY' +import json +import sys + +source, destination, lora_path = sys.argv[1:] +with open(source, encoding="utf-8") as handle: + config = json.load(handle) +if config.get("attn_type") != "dynamic_sparse_attn": + raise ValueError("MiniMax-H3 SLA config must use attn_type=dynamic_sparse_attn") +settings = config.get("dynamic_sparse_attn_setting", {}) +if settings.get("operator") != "intel_xpu": + raise ValueError("MiniMax-H3 SLA config must use operator=intel_xpu") +if len(config.get("lora_configs", [])) != 1: + raise ValueError("MiniMax-H3 SLA config must contain exactly one LoRA entry") +config["lora_configs"][0]["path"] = lora_path +with open(destination, "w", encoding="utf-8") as handle: + json.dump(config, handle, indent=2, ensure_ascii=False) +PY + +source "${lightx2v_path}/scripts/base/base.sh" +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +# Fail before loading the 166 GB base model if an older sycl-kernels wheel is +# active or the SLA API was not packaged. +python - <<'PY' +import sycl_kernels + +required = ("sla_block_map", "sparse_block_attention", "sla_sparse_attention") +missing = [name for name in required if not callable(getattr(sycl_kernels, name, None))] +if missing: + raise RuntimeError(f"Installed sycl_kernels is missing SLA APIs: {missing}") +print(f"Using sycl_kernels from {sycl_kernels.__file__}") +PY + +echo "MiniMax-H3 model: ${model_path}" +echo "SLA LoRA: ${lora_path}" +echo "Config: ${config_template}" +echo "XPU: ${ZE_AFFINITY_MASK}" +echo "First frame: ${first_frame}" +echo "Last frame: ${last_frame}" +echo "Output: ${output_path}" + +torchrun --standalone --nproc_per_node=1 -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task fl2av \ + --model_path "${model_path}" \ + --config_json "${runtime_config}" \ + --prompt "${prompt}" \ + --image_path "${first_frame}" \ + --last_frame_path "${last_frame}" \ + --save_result_path "${output_path}" \ + --seed "${seed}" From 5139e8555a97cff060aaf503c285f48fe257d599 Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Tue, 8 Sep 2026 08:59:02 +0000 Subject: [PATCH 2/7] fix: align XPU SLA attention operator naming --- .../minimax_h3_fl2v_turbo_sla_4step.json | 2 +- .../common/ops/attn/dynamic_sparse_attn.py | 19 +++++++------------ .../minimax_h3/weights/transformer_weights.py | 5 +---- .../runners/minimax_h3/minimax_h3_runner.py | 2 +- 4 files changed, 10 insertions(+), 18 deletions(-) diff --git a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json index 9e83f51e9..11d25160e 100644 --- a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json @@ -19,7 +19,7 @@ "refiner_attn_type": "intel_xpu_cute_attn", "dynamic_sparse_attn_setting": { "sparsity_ratio": 0.85, - "operator": "intel_xpu" + "operator": "intel_xpu_cute_attn" }, "rms_type": "intel_xpu", "rope_type": "minimax_h3_xpu_rope", diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index 5ace0b9a6..6bff44f57 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -78,17 +78,9 @@ def __init__(self, config=None): raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}") self.topk = 1 - self.sparsity_ratio - self.arch = None - if self.operator != "intel_xpu": - self.arch = get_cuda_arch(torch.cuda.current_device()) - - if self.operator == "intel_xpu": - # The optimized MiniMax-H3 kernel consumes BLHD directly. Keep - # this branch ahead of CUDA architecture discovery so merely - # constructing the XPU backend never touches torch.cuda. - self.BLKQ, self.BLKK = 128, 128 - self.apply_func = self.apply_intel_xpu - elif self.operator == "triton": + self.arch = get_cuda_arch(torch.cuda.current_device()) if torch.cuda.is_available() else None + + if self.operator == "triton": self.BLKQ, self.BLKK = 64, 64 self.apply_func = self.apply_triton elif self.operator == "triton_ar": # triton for AR models @@ -109,12 +101,15 @@ def __init__(self, config=None): elif self.operator == "magi": self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_magi + elif self.operator == "intel_xpu_cute_attn": + self.BLKQ, self.BLKK = 128, 128 + self.apply_func = self.apply_intel_xpu_cute_attn else: raise NotImplementedError(f"Not supported SLA operator: {self.operator}.") # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}") - def apply_intel_xpu( + def apply_intel_xpu_cute_attn( self, q, k, diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index b85c3d6c5..9634eb595 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -74,10 +74,7 @@ def __init__(self, prefix, config, create_cuda_buffer=False): attn_type = config.get("attn_type", "flash_attn3") attention_cls = ATTN_WEIGHT_REGISTER[attn_type] if attn_type == "dynamic_sparse_attn": - sparse_config = config.get("dynamic_sparse_attn_setting", {}) - if sparse_config.get("operator") == "intel_xpu" and config.get("seq_parallel", False): - raise NotImplementedError("Intel XPU SLA does not yet support MiniMax-H3 sequence parallelism") - calculate = attention_cls(sparse_config) + calculate = attention_cls(config.get("dynamic_sparse_attn_setting", {})) else: calculate = attention_cls() if attn_type == "sol_attn": diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 510425f3c..cbc1b33c3 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -61,7 +61,7 @@ def build_minimax_h3_model_with_lora(config, model_kwargs, lora_configs): if not lora_config.get("path"): raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].path") if lora_config.get("alpha") is None: - raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].alpha (use the alpha published with the checkpoint)") + raise ValueError("MiniMax-H3 dynamic LoRA requires lora_configs[0].alpha (use 8 for the MiniMax-H3 Turbo LoRA)") model_kwargs.update( lora_path=lora_config["path"], lora_strength=lora_config.get("strength", 1.0), From a8ce72f9b90d67e3c6309525ff17de0451cb75c2 Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Wed, 9 Sep 2026 01:46:03 +0000 Subject: [PATCH 3/7] feat: enable multi-card XPU SLA inference --- .../test/test_xpu_sla_attn_adapter.py | 30 ++++++++++- .../run_minimax_h3_fl2av_turbo_sla_4step.sh | 53 ++++++++++++++++--- 2 files changed, 76 insertions(+), 7 deletions(-) diff --git a/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py b/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py index 14d3d8511..08fcce3b9 100644 --- a/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py +++ b/lightx2v_kernel_xpu/test/test_xpu_sla_attn_adapter.py @@ -63,7 +63,7 @@ def fake_sla(q, k, v, keep_ratio, block_q, block_k, scale=None): monkeypatch.setattr(torch.cuda, "current_device", fail_if_cuda_is_queried) monkeypatch.setattr(xpu_sla_attn, "sla_sparse_attention", fake_sla) - attention = DynamicSparseAttnWeight({"operator": "intel_xpu", "sparsity_ratio": 0.85}) + attention = DynamicSparseAttnWeight({"operator": "intel_xpu_cute_attn", "sparsity_ratio": 0.85}) q = torch.empty(17, 7, 128, dtype=torch.bfloat16) cu_seqlens = torch.tensor([0, 17], dtype=torch.int32) @@ -77,3 +77,31 @@ def fake_sla(q, k, v, keep_ratio, block_q, block_k, scale=None): "block_k": 128, "scale": None, } + + +@pytest.mark.parametrize( + ("sequence_length", "local_heads"), + [ + (19292, 28), # TP=2: full sequence, head shard + (19292, 28), # SP=2 after Ulysses: global sequence, head shard + (19292, 14), # TP=2 + SP=2 after Ulysses + ], +) +def test_xpu_sla_accepts_distributed_attention_layout(monkeypatch, sequence_length, local_heads): + """TP/SP both enter SLA as global sequence plus rank-local heads.""" + from lightx2v.common.ops.attn.dynamic_sparse_attn import DynamicSparseAttnWeight + + received = {} + + def fake_sla(q, k, v, keep_ratio, block_q, block_k, scale=None): + received["shape"] = tuple(q.shape) + return q + + monkeypatch.setattr(xpu_sla_attn, "sla_sparse_attention", fake_sla) + attention = DynamicSparseAttnWeight({"operator": "intel_xpu_cute_attn", "sparsity_ratio": 0.85}) + q = torch.empty(sequence_length, local_heads, 128, dtype=torch.bfloat16, device="meta") + + output = attention.apply(q, q, q) + + assert received["shape"] == (1, sequence_length, local_heads, 128) + assert output.shape == (sequence_length, local_heads * 128) diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh index b0cafd376..a13aabe36 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh @@ -13,13 +13,42 @@ last_frame=${LAST_FRAME:-${lightx2v_path}/assets/inputs/imgs/flf2v_input_last_fr output_path=${OUTPUT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_fl2av_turbo_sla_4step.mp4} prompt=${PROMPT:-Create a coherent cinematic transition between the two frames with natural synchronized ambient sound.} seed=${SEED:-42} +parallel_mode=${PARALLEL_MODE:-single} +tp_size=${TP_SIZE:-2} +sp_size=${SP_SIZE:-2} -export ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-0} +case "${parallel_mode}" in + single) tp_size=1; sp_size=1 ;; + tp) sp_size=1 ;; + sp) tp_size=1 ;; + sp_tp) ;; + *) echo "PARALLEL_MODE must be one of: single, tp, sp, sp_tp" >&2; exit 1 ;; +esac +nproc_per_node=$((tp_size * sp_size)) +if ((tp_size < 1 || sp_size < 1)); then + echo "TP_SIZE and SP_SIZE must be positive integers" >&2 + exit 1 +fi + +if [[ -z "${ZE_AFFINITY_MASK:-}" ]]; then + ZE_AFFINITY_MASK=$(seq -s, 0 $((nproc_per_node - 1))) + export ZE_AFFINITY_MASK +fi export PLATFORM=${PLATFORM:-intel_xpu} export PYTHONFAULTHANDLER=${PYTHONFAULTHANDLER:-1} export PYTHONUNBUFFERED=${PYTHONUNBUFFERED:-1} export PYTHONPATH=${PYTHONPATH:-} +if ((sp_size > 1)); then + # oneCCL settings used by the Ulysses all-to-all path on Intel XPU. + export CCL_SYCL_ALLTOALL_ARC_LL=${CCL_SYCL_ALLTOALL_ARC_LL:-1} + export CCL_SYCL_ALLTOALL_TMP_BUF=${CCL_SYCL_ALLTOALL_TMP_BUF:-1} + export CCL_SYCL_CCL_BARRIER=${CCL_SYCL_CCL_BARRIER:-1} + export CCL_SYCL_ALLREDUCE_SIMPLE_THRESHOLD=${CCL_SYCL_ALLREDUCE_SIMPLE_THRESHOLD:-4294967296} + export CCL_SYCL_REDUCE_SCATTER_SIMPLE_THRESHOLD=${CCL_SYCL_REDUCE_SCATTER_SIMPLE_THRESHOLD:-4294967296} + export CCL_SYCL_ALLGATHERV_SIMPLE_THRESHOLD=${CCL_SYCL_ALLGATHERV_SIMPLE_THRESHOLD:-4294967296} +fi + [[ -d "${model_path}" ]] || { echo "Model directory not found: ${model_path}" >&2; exit 1; } [[ -f "${lora_path}" ]] || { echo "SLA LoRA checkpoint not found: ${lora_path}" >&2; exit 1; } [[ -f "${config_template}" ]] || { echo "Config file not found: ${config_template}" >&2; exit 1; } @@ -32,21 +61,32 @@ trap 'rm -f -- "${runtime_config}"' EXIT # Keep the checked-in config reusable while allowing LORA_PATH to override the # local checkpoint location without editing JSON. -python - "${config_template}" "${runtime_config}" "${lora_path}" <<'PY' +python - "${config_template}" "${runtime_config}" "${lora_path}" "${tp_size}" "${sp_size}" <<'PY' import json import sys -source, destination, lora_path = sys.argv[1:] +source, destination, lora_path, tp_size, sp_size = sys.argv[1:] +tp_size, sp_size = int(tp_size), int(sp_size) with open(source, encoding="utf-8") as handle: config = json.load(handle) if config.get("attn_type") != "dynamic_sparse_attn": raise ValueError("MiniMax-H3 SLA config must use attn_type=dynamic_sparse_attn") settings = config.get("dynamic_sparse_attn_setting", {}) -if settings.get("operator") != "intel_xpu": - raise ValueError("MiniMax-H3 SLA config must use operator=intel_xpu") +if settings.get("operator") != "intel_xpu_cute_attn": + raise ValueError("MiniMax-H3 SLA config must use operator=intel_xpu_cute_attn") if len(config.get("lora_configs", [])) != 1: raise ValueError("MiniMax-H3 SLA config must contain exactly one LoRA entry") config["lora_configs"][0]["path"] = lora_path +if tp_size > 1 or sp_size > 1: + parallel = {"tensor_p_size": tp_size, "seq_p_size": sp_size} + if sp_size > 1: + parallel.update({"seq_p_attn_type": "ulysses", "seq_p_a2a_backend": "torch"}) + config["parallel"] = parallel +else: + config.pop("parallel", None) +if tp_size > 1: + config["tp_mm_type"] = "IntelTensorParallel" + config["text_encoder_tensor_parallel"] = True with open(destination, "w", encoding="utf-8") as handle: json.dump(config, handle, indent=2, ensure_ascii=False) PY @@ -71,11 +111,12 @@ echo "MiniMax-H3 model: ${model_path}" echo "SLA LoRA: ${lora_path}" echo "Config: ${config_template}" echo "XPU: ${ZE_AFFINITY_MASK}" +echo "Parallel mode: ${parallel_mode} (TP=${tp_size}, SP=${sp_size}, processes=${nproc_per_node})" echo "First frame: ${first_frame}" echo "Last frame: ${last_frame}" echo "Output: ${output_path}" -torchrun --standalone --nproc_per_node=1 -m lightx2v.infer \ +torchrun --standalone --nproc_per_node="${nproc_per_node}" -m lightx2v.infer \ --model_cls minimax_h3 \ --task fl2av \ --model_path "${model_path}" \ From 199cd0904ad6116bc47bb996b49d43328b5113cf Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Wed, 9 Sep 2026 01:53:29 +0000 Subject: [PATCH 4/7] update --- .../patches/minimax_h3_sparse_kernel.patch | 2 +- .../cute/patches/minimax_h3_sparse_lut.patch | 4 +-- .../python/sycl_kernels/sla.py | 14 ++++++----- .../python/sycl_kernels/sla_triton.py | 21 +++++++++++++--- .../test/bench_sla_sparse_attention.py | 25 +++++++++++-------- .../test/test_sla_sparse_attention.py | 1 - .../ops/attn/intel_xpu/xpu_sla_attn.py | 10 ++------ .../run_minimax_h3_fl2av_turbo_sla_4step.sh | 19 ++++++++++++++ 8 files changed, 64 insertions(+), 32 deletions(-) diff --git a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch index b6e75bb83..5b97f04f0 100644 --- a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch +++ b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_kernel.patch @@ -8,7 +8,7 @@ + const int k_blocks = params.mainloop.ptr_block_lut + ? params.mainloop.lut_topk * params.mainloop.lut_block_tiles + : cute::ceil_div(seq_len, get<1>(TileShapeQK{})); - + int offset_q = 0, offset_k = 0, offset_v = 0, offset_o = 0; int offset_k_cache = 0, offset_v_cache = 0; @@ -287,7 +289,7 @@ diff --git a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch index 91ad36793..3afbd5c00 100644 --- a/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch +++ b/lightx2v_kernel_xpu/cute/patches/minimax_h3_sparse_lut.patch @@ -10,7 +10,7 @@ + int lut_topk = 0; + int lut_block_tiles = 0; }; - + // Kernel-facing parameters @@ -172,7 +177,9 @@ Params to_underlying_arguments(Arguments const &args, void * /* workspace */) { @@ -21,7 +21,7 @@ + args.ptr_block_lut, args.lut_num_heads, args.lut_q_blocks, + args.lut_topk, args.lut_block_tiles}; } - + CUTLASS_HOST_DEVICE static @@ -211,6 +218,7 @@ int seq_len, diff --git a/lightx2v_kernel_xpu/python/sycl_kernels/sla.py b/lightx2v_kernel_xpu/python/sycl_kernels/sla.py index f99255000..4b70c5cb6 100644 --- a/lightx2v_kernel_xpu/python/sycl_kernels/sla.py +++ b/lightx2v_kernel_xpu/python/sycl_kernels/sla.py @@ -125,10 +125,7 @@ def _validate_sparse_inputs(q, k, v, lut, block_q: int, block_k: int) -> None: raise ValueError("query head count must be divisible by KV head count") expected_q_blocks = math.ceil(q.shape[1] / block_q) if lut.ndim != 4 or tuple(lut.shape[:3]) != (q.shape[0], q.shape[2], expected_q_blocks): - raise ValueError( - "lut must have shape [B, Hq, ceil(Lq/block_q), topk], got " - f"{tuple(lut.shape)}" - ) + raise ValueError(f"lut must have shape [B, Hq, ceil(Lq/block_q), topk], got {tuple(lut.shape)}") if lut.shape[3] == 0: raise ValueError("lut topk dimension must be non-zero") if block_q <= 0 or block_k <= 0: @@ -174,8 +171,13 @@ def sparse_block_attention( from .sla_triton import launch_sparse_block_attention return launch_sparse_block_attention( - q.contiguous(), k.contiguous(), v.contiguous(), lut.to(torch.int32).contiguous(), - block_q, block_k, q.shape[-1] ** -0.5 if scale is None else scale, + q.contiguous(), + k.contiguous(), + v.contiguous(), + lut.to(torch.int32).contiguous(), + block_q, + block_k, + q.shape[-1] ** -0.5 if scale is None else scale, ) diff --git a/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py b/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py index c216bc455..1a980d7c7 100644 --- a/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py +++ b/lightx2v_kernel_xpu/python/sycl_kernels/sla_triton.py @@ -83,9 +83,22 @@ def launch_sparse_block_attention(q, k, v, lut, block_q, block_k, scale): output = torch.empty_like(q_bhld) grid = (q_blocks, batch * q_heads) _sparse_block_attention_fwd[grid]( - q_bhld, k_bhld, v_bhld, lut, output, - q_len, kv_len, q_heads, kv_heads, q_blocks, topk, scale, - BLOCK_Q=block_q, BLOCK_K=block_k, HEAD_DIM=head_dim, - num_warps=8, num_stages=3, + q_bhld, + k_bhld, + v_bhld, + lut, + output, + q_len, + kv_len, + q_heads, + kv_heads, + q_blocks, + topk, + scale, + BLOCK_Q=block_q, + BLOCK_K=block_k, + HEAD_DIM=head_dim, + num_warps=8, + num_stages=3, ) return output.permute(0, 2, 1, 3) diff --git a/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py b/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py index c48f539e4..fe20e7e1d 100644 --- a/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py +++ b/lightx2v_kernel_xpu/test/bench_sla_sparse_attention.py @@ -60,16 +60,21 @@ def main(): sparse_median = statistics.median(sparse_ms) router_median = statistics.median(router_ms) dense_median = statistics.median(dense_ms) - print(json.dumps({ - "shape": shape, - "keep_ratio": args.keep_ratio, - "lut_shape": list(lut.shape), - "dense_ms": dense_ms, - "sparse_kernel_ms": sparse_ms, - "router_ms": router_ms, - "kernel_speedup": dense_median / sparse_median, - "attention_speedup_including_router": dense_median / (sparse_median + router_median), - }, indent=2)) + print( + json.dumps( + { + "shape": shape, + "keep_ratio": args.keep_ratio, + "lut_shape": list(lut.shape), + "dense_ms": dense_ms, + "sparse_kernel_ms": sparse_ms, + "router_ms": router_ms, + "kernel_speedup": dense_median / sparse_median, + "attention_speedup_including_router": dense_median / (sparse_median + router_median), + }, + indent=2, + ) + ) if __name__ == "__main__": diff --git a/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py b/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py index 15ec828df..b9814fc2c 100644 --- a/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py +++ b/lightx2v_kernel_xpu/test/test_sla_sparse_attention.py @@ -2,7 +2,6 @@ import pytest import torch - from sycl_kernels.sla import ( sla_block_map, sparse_block_attention, diff --git a/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py index fac8d3342..8c6d8c9ac 100644 --- a/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py +++ b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py @@ -13,16 +13,10 @@ def _get_sycl_api(name): if _sycl_kernels is None: - raise RuntimeError( - "Intel XPU SLA requires lightx2v_kernel_xpu's sycl-kernels package " - "with SLA/CUTE sparse attention enabled" - ) from _IMPORT_ERROR + raise RuntimeError("Intel XPU SLA requires lightx2v_kernel_xpu's sycl-kernels package with SLA/CUTE sparse attention enabled") from _IMPORT_ERROR api = getattr(_sycl_kernels, name, None) if api is None: - raise RuntimeError( - f"sycl_kernels.{name} is unavailable; rebuild/install " - "lightx2v_kernel_xpu with SLA/CUTE sparse attention enabled" - ) + raise RuntimeError(f"sycl_kernels.{name} is unavailable; rebuild/install lightx2v_kernel_xpu with SLA/CUTE sparse attention enabled") return api diff --git a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh index a13aabe36..ecb495b7b 100755 --- a/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +++ b/scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh @@ -1,6 +1,25 @@ #!/usr/bin/env bash set -euo pipefail +# Usage: +# # Single XPU (default) +# bash scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +# +# # Tensor parallel on 2 XPUs +# PARALLEL_MODE=tp TP_SIZE=2 \ +# bash scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +# +# # Sequence parallel on 2 XPUs +# PARALLEL_MODE=sp SP_SIZE=2 \ +# bash scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +# +# # Tensor parallel 2 x sequence parallel 2 (4 XPUs) +# PARALLEL_MODE=sp_tp TP_SIZE=2 SP_SIZE=2 \ +# bash scripts/platforms/intel_xpu/run_minimax_h3_fl2av_turbo_sla_4step.sh +# +# Optional overrides include ZE_AFFINITY_MASK, MODEL_PATH, LORA_PATH, +# CONFIG_JSON, FIRST_FRAME, LAST_FRAME, OUTPUT_PATH, PROMPT, and SEED. + SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) From 7609d1d85d381b18744a8609941e7dbd0e1adb1c Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Wed, 9 Sep 2026 02:02:18 +0000 Subject: [PATCH 5/7] refactor: move XPU SLA adapter to platform layer --- .../common/ops/attn/dynamic_sparse_attn.py | 46 ++++--------------- .../ops/attn/intel_xpu/xpu_sla_attn.py | 41 +++++++++++++++++ 2 files changed, 51 insertions(+), 36 deletions(-) diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index 6bff44f57..1cdddcaf4 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -1,3 +1,5 @@ +from functools import partial + import torch from loguru import logger @@ -103,47 +105,19 @@ def __init__(self, config=None): self.apply_func = self.apply_magi elif self.operator == "intel_xpu_cute_attn": self.BLKQ, self.BLKK = 128, 128 - self.apply_func = self.apply_intel_xpu_cute_attn + from lightx2v_platform.ops.attn.intel_xpu.xpu_sla_attn import apply_intel_xpu_cute_attn + + self.apply_func = partial( + apply_intel_xpu_cute_attn, + keep_ratio=self.topk, + block_q=self.BLKQ, + block_k=self.BLKK, + ) else: raise NotImplementedError(f"Not supported SLA operator: {self.operator}.") # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}") - def apply_intel_xpu_cute_attn( - self, - q, - k, - v, - cu_seqlens_q=None, - cu_seqlens_kv=None, - max_seqlen_q=None, - max_seqlen_kv=None, - **kwargs, - ): - """Run the XPU SLA router and fused sparse attention on one sequence.""" - if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: - raise ValueError("Intel XPU SLA expects q, k and v in [L, H, D] layout") - if q.shape != k.shape or k.shape != v.shape: - raise ValueError("Intel XPU SLA currently requires self-attention with matching q/k/v shapes") - if max_seqlen_q != q.shape[0] or max_seqlen_kv != k.shape[0]: - raise ValueError("Intel XPU SLA currently supports one unpadded sequence per call") - if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: - raise ValueError("Intel XPU SLA currently supports one sequence per call") - if cu_seqlens_kv is not None and cu_seqlens_kv.numel() != 2: - raise ValueError("Intel XPU SLA currently supports one sequence per call") - - from lightx2v_platform.ops.attn.intel_xpu.xpu_sla_attn import sla_sparse_attention - - out = sla_sparse_attention( - q.unsqueeze(0), - k.unsqueeze(0), - v.unsqueeze(0), - keep_ratio=self.topk, - block_q=self.BLKQ, - block_k=self.BLKK, - ) - return out.reshape(max_seqlen_q, -1) - def apply( self, q, diff --git a/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py index 8c6d8c9ac..6dc7dea4a 100644 --- a/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py +++ b/lightx2v_platform/ops/attn/intel_xpu/xpu_sla_attn.py @@ -33,3 +33,44 @@ def sparse_block_attention(q, k, v, lut, block_q=128, block_k=128, scale=None): def sla_sparse_attention(q, k, v, keep_ratio=0.2, block_q=128, block_k=128, scale=None): """Route and apply sparse attention through ``sycl_kernels.sla_sparse_attention``.""" return _get_sycl_api("sla_sparse_attention")(q, k, v, keep_ratio, block_q, block_k, scale) + + +def apply_intel_xpu_cute_attn( + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + *, + keep_ratio=0.2, + block_q=128, + block_k=128, + **kwargs, +): + """Run the XPU SLA router and fused sparse attention on one sequence.""" + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("Intel XPU SLA expects q, k and v in [L, H, D] layout") + if q.shape != k.shape or k.shape != v.shape: + raise ValueError("Intel XPU SLA currently requires self-attention with matching q/k/v shapes") + if max_seqlen_q is None: + max_seqlen_q = q.shape[0] + if max_seqlen_kv is None: + max_seqlen_kv = k.shape[0] + if max_seqlen_q != q.shape[0] or max_seqlen_kv != k.shape[0]: + raise ValueError("Intel XPU SLA currently supports one unpadded sequence per call") + if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: + raise ValueError("Intel XPU SLA currently supports one sequence per call") + if cu_seqlens_kv is not None and cu_seqlens_kv.numel() != 2: + raise ValueError("Intel XPU SLA currently supports one sequence per call") + + out = sla_sparse_attention( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + keep_ratio=keep_ratio, + block_q=block_q, + block_k=block_k, + ) + return out.reshape(max_seqlen_q, -1) From 0cd2362d65e620ff3f8ed32bac5413223c8661ad Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Wed, 9 Sep 2026 02:08:00 +0000 Subject: [PATCH 6/7] update --- .../platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json index 11d25160e..d01f864fe 100644 --- a/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json +++ b/configs/platforms/intel_xpu/minimax_h3_fl2v_turbo_sla_4step.json @@ -33,7 +33,7 @@ "audio_latents_per_second": 40, "audio_channels": 2, "keep_latents_dtype_in_scheduler": true, - "lora_dynamic_apply": true, + "lora_dynamic_apply": false, "lora_configs": [ { "path": "/llm/models/Minimax-h3-Turbo-SLA/minimax_h3_fl2v_turbo_4step_v0.1_768p_sla_bf16.safetensors", From db1b0a0ace1e88d89d221e6244b44c2be4643297 Mon Sep 17 00:00:00 2001 From: qiuxin2012 Date: Wed, 9 Sep 2026 06:01:05 +0000 Subject: [PATCH 7/7] update tp --- .../ops/mm/intel_xpu/mm_weight.py | 48 ++++++++----------- 1 file changed, 20 insertions(+), 28 deletions(-) diff --git a/lightx2v_platform/ops/mm/intel_xpu/mm_weight.py b/lightx2v_platform/ops/mm/intel_xpu/mm_weight.py index 08e58204e..a83622eea 100644 --- a/lightx2v_platform/ops/mm/intel_xpu/mm_weight.py +++ b/lightx2v_platform/ops/mm/intel_xpu/mm_weight.py @@ -49,25 +49,22 @@ def _get_intel_tensor_parallel_class(): tensor_parallel_class = MM_WEIGHT_REGISTER["TensorParallel"] class IntelTensorParallelWeight(tensor_parallel_class): - """Tensor-parallel linear layer with a oneCCL hang workaround. - - oneCCL 2021.15 can hang when a TP all-reduce is used - together with SP collectives. Gather every TP partial instead and - sum them locally to preserve the all-reduce result without entering - the problematic oneCCL all-reduce path. - """ + """Tensor-parallel linear layer with a large-world oneCCL workaround.""" def apply(self, input_tensor): output = self._mm.apply(input_tensor) if self.split_dim == "row" and self.reduce_output and self.tp_size > 1 and self.tp_group is not None: - # Work around the oneCCL 2021.15 TP all-reduce hang by using - # all-gather followed by an equivalent local reduction. - partials = [torch.empty_like(output) for _ in range(self.tp_size)] - dist.all_gather(partials, output.contiguous(), group=self.tp_group) - output.copy_(partials[0]) - for partial in partials[1:]: - output.add_(partial) + if self.tp_size <= 2: + dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.tp_group) + else: + # Avoid the large-world oneCCL all-reduce path that has + # previously hung on eight-device configurations. + partials = [torch.empty_like(output) for _ in range(self.tp_size)] + dist.all_gather(partials, output.contiguous(), group=self.tp_group) + output.copy_(partials[0]) + for partial in partials[1:]: + output.add_(partial) if self._row_split_bias is not None: output = output + self._row_split_bias @@ -205,25 +202,20 @@ def _get_tensor_parallel_rms_class(): tensor_parallel_rms_class = RMS_WEIGHT_REGISTER["TensorParallelFP32"] class TensorParallelRMSWeight(tensor_parallel_rms_class): - """Tensor-parallel RMSNorm with a oneCCL hang workaround. - - oneCCL 2021.15 can hang when a TP all-reduce is used - together with SP collectives. Gather the local squared sums from - every TP rank and add them locally, which is mathematically - equivalent to the original all-reduce. - """ + """Tensor-parallel RMSNorm with a large-world oneCCL workaround.""" def apply(self, input_tensor): input_fp32 = input_tensor.float() local_sum = input_fp32.square().sum(dim=-1, keepdim=True) if self.tp_size > 1 and self.tp_group is not None: - # Avoid the problematic oneCCL all-reduce while retaining - # the same global sum used by RMSNorm. - partials = [torch.empty_like(local_sum) for _ in range(self.tp_size)] - dist.all_gather(partials, local_sum.contiguous(), group=self.tp_group) - local_sum.copy_(partials[0]) - for partial in partials[1:]: - local_sum.add_(partial) + if self.tp_size <= 2: + dist.all_reduce(local_sum, op=dist.ReduceOp.SUM, group=self.tp_group) + else: + partials = [torch.empty_like(local_sum) for _ in range(self.tp_size)] + dist.all_gather(partials, local_sum.contiguous(), group=self.tp_group) + local_sum.copy_(partials[0]) + for partial in partials[1:]: + local_sum.add_(partial) global_hidden_dim = input_tensor.shape[-1] * self.tp_size output = input_fp32 * torch.rsqrt(local_sum / global_hidden_dim + self.eps)