Skip to content

Commit d54ee0b

Browse files
authored
feat(xpu): support MiniMax-H3 FP8 inference (#1431)
## Summary - Add MiniMax-H3 FP8 inference support for Intel XPU. - Add FP8 configuration and launch script. - Port large-matrix chunking from the day0 implementation. - Add oneDNN primitive caching. - Fix missing bias in FP8 linear operations. - Move MMWeightFp8IntelXpu into platform-specific ops. - Add related tests and documentation. ## Validation - XPU extension builds successfully. - Python syntax checks pass. - git diff --check passes.
1 parent 25cf2a2 commit d54ee0b

10 files changed

Lines changed: 546 additions & 183 deletions

File tree

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
{
2+
"infer_steps": 30,
3+
"target_video_length": 124,
4+
"target_height": 544,
5+
"target_width": 960,
6+
"fps": 24,
7+
"target_fps": 24,
8+
"enable_cfg": false,
9+
"cpu_offload": true,
10+
"offload_granularity": "block",
11+
"text_encoder_cpu_offload": true,
12+
"text_encoder_offload_granularity": "block",
13+
"text_encoder_release_block_offload_buffers": true,
14+
"vae_cpu_offload": true,
15+
"lazy_load": false,
16+
"unload_modules": false,
17+
"attn_type": "intel_xpu_cute_attn",
18+
"rms_type": "intel_xpu",
19+
"rope_type": "minimax_h3_xpu_rope",
20+
"feature_caching": "NoCaching",
21+
"use_compile": false,
22+
"video_flow_shift": 12.0,
23+
"audio_flow_shift": 3.0,
24+
"vae_spatial_scale_factor": 16,
25+
"audio_sampling_rate": 32000,
26+
"audio_latents_per_second": 40,
27+
"audio_channels": 2,
28+
"keep_latents_dtype_in_scheduler": true,
29+
"dit_quantized": true,
30+
"dit_quant_scheme": "fp8-intel-xpu",
31+
"dit_quantized_ckpt": "/llm/models/MiniMax-H3/quantized/fp8/minimax_h3_fp8_xpu.safetensors"
32+
}

lightx2v/common/ops/mm/mm_weight.py

Lines changed: 0 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -124,11 +124,6 @@ def _fp8_scaled_mm(
124124
except ImportError:
125125
marlin_cuda_quant = None
126126

127-
try:
128-
import sycl_kernels
129-
except ImportError:
130-
sycl_kernels = None
131-
132127

133128
class MMWeightTemplate(metaclass=ABCMeta):
134129
def __init__(
@@ -2634,64 +2629,3 @@ def apply(self, input_tensor):
26342629
def unwrap_tp_weight(module):
26352630
"""Return the concrete tensor-owning MM implementation from a TP wrapper."""
26362631
return module._mm if isinstance(module, MMWeightTP) else module
2637-
2638-
2639-
@MM_WEIGHT_REGISTER("fp8-intel-xpu")
2640-
class MMWeightFp8IntelXpu(MMWeightQuantTemplate):
2641-
"""
2642-
Name: W-fp8-channel-sym-A-fp16-Intel-XPU
2643-
2644-
Intel XPU optimized FP8 kernel:
2645-
Weight Storage: fp8 (torch.float8_e4m3fn) - saves 50% memory
2646-
Computation: fp16 using PyTorch native ops
2647-
- Dynamically dequantize FP8 → FP16 during forward
2648-
- Use torch.nn.functional.linear (compatible with Intel XPU)
2649-
2650-
Benefits:
2651-
- Memory efficient: FP8 storage (8-bit)
2652-
- Compatible: FP16 compute using PyTorch native ops
2653-
- Intel XPU friendly: No CUDA-specific kernels
2654-
2655-
Usage in config:
2656-
{
2657-
"dit_quant_scheme": "fp8-intel-xpu",
2658-
"weight_auto_quant": true,
2659-
"dit_quantized": true
2660-
}
2661-
"""
2662-
2663-
def __init__(self, weight_name, bias_name, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, lazy_load_file=None, is_post_adapter=False, lora_prefix=None, lora_path=""):
2664-
super().__init__(weight_name, bias_name, create_cuda_buffer, create_cpu_buffer, lazy_load, lazy_load_file, is_post_adapter, lora_prefix, lora_path)
2665-
2666-
self.load_func = self.load_fp8_perchannel_sym
2667-
self.weight_need_transpose = False # We'll handle transpose in apply
2668-
2669-
def apply(self, input_tensor):
2670-
# # """
2671-
# Forward pass with FP8 → FP16 dequantization
2672-
2673-
# Steps:
2674-
# 1. Dequantize weight: fp8 → fp16 (weight * scale)
2675-
# 2. Compute: torch.nn.functional.linear(input_fp16, weight_fp16, bias)
2676-
# """
2677-
# # Ensure input is FP16
2678-
# # print(input_tensor.dtype)
2679-
2680-
if sycl_kernels is not None:
2681-
try:
2682-
return sycl_kernels.onednn_w8a16_fp8(input_tensor, self.weight, self.weight_scale.to(torch.float))
2683-
except RuntimeError:
2684-
pass # Fall through to torch dequantization path
2685-
2686-
infer_dtype = self.infer_dtype
2687-
squeeze_output = False
2688-
if input_tensor.dim() == 3 and input_tensor.shape[0] == 1:
2689-
input_tensor = input_tensor.squeeze(0)
2690-
squeeze_output = True
2691-
input_tensor = input_tensor.to(infer_dtype)
2692-
weight_fp16 = self.weight.to(infer_dtype) * self.weight_scale.to(infer_dtype)
2693-
bias_fp16 = self.bias.to(infer_dtype) if hasattr(self, "bias") and self.bias is not None else None
2694-
output = torch.nn.functional.linear(input_tensor, weight_fp16, bias_fp16)
2695-
if squeeze_output:
2696-
output = output.unsqueeze(0)
2697-
return output

lightx2v/models/networks/minimax_h3/model.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
"fp8-torchao",
2929
"fp8-triton",
3030
"fp8-vllm",
31+
"fp8-intel-xpu",
3132
"int8-q8f",
3233
"int8-sgl",
3334
"int8-torchao",

lightx2v_kernel_xpu/README.md

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ Exposed as the Python package `sycl_kernels`:
1919
|----------|-------------|
2020
| `sdp(Q, K, V)` | ESIMD Flash Attention — `[B, L, H, 128]` fp16/bf16, PTL-H doubleGRF |
2121
| `onednn_w8a8_int8(x, qweight, scales[, bias])` | Dynamic W8A8 GEMM — rowwise INT8 activation quantization × INT8 weights |
22-
| `onednn_w8a16_fp8(x, qweight, scales[, bias])` | W8A16 GEMM — fp16/bf16 activations × FP8_E4M3 weights, per-column scale |
22+
| `onednn_w8a16_fp8(x, qweight, scales[, bias])` | Cached W8A16 GEMM — fp16/bf16/fp32 activations × FP8 E4M3/E5M2 weights, per-column scale |
2323
| `onednn_w4a16(x, weight, scales, zeros[, bias])` | W4A16 GEMM — fp16/bf16 activations × INT4 packed weights |
2424

2525
Tested on **Intel Arc B390 GPU** (PTL-H / Xe2), PyTorch 2.9.1+xpu, oneAPI 2025.2.
@@ -136,13 +136,16 @@ out = sycl_kernels.onednn_w8a8_int8(x, qweight, scales)
136136
out = sycl_kernels.onednn_w8a8_int8(x, qweight, scales, bias)
137137

138138
# ── W8A16 FP8 GEMM ────────────────────────────────────────────────────────────
139-
# x : [M, K] fp16 or bf16 on XPU
140-
# qweight : [N, K] float8_e4m3fn on XPU
139+
# x : [M, K] fp16/bf16/fp32 on XPU
140+
# qweight : [N, K] float8_e4m3fn or float8_e5m2 on XPU
141141
# scales : [N, 1] fp32 on XPU (per-output-channel absmax scale)
142-
# bias : [N] fp16/bf16 on XPU (optional)
142+
# bias : [N] same dtype as x on XPU (optional)
143143
# Returns : [M, N] same dtype as x
144144
out = sycl_kernels.onednn_w8a16_fp8(x, qweight, scales)
145145
out = sycl_kernels.onednn_w8a16_fp8(x, qweight, scales, bias)
146+
# Repeated shapes reuse cached oneDNN primitives. Known MiniMax-H3 large
147+
# projections are split along N to avoid unsupported full-size primitives.
148+
hits, misses, size = sycl_kernels.fp8_cache_stats()
146149

147150
# ── W4A16 GEMM ────────────────────────────────────────────────────────────────
148151
# x : [M, K] fp16 or bf16 on XPU

lightx2v_kernel_xpu/csrc/entry.cpp

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
//
1616

1717
#include <optional>
18+
#include <tuple>
1819
#include <torch/extension.h>
1920

2021
torch::Tensor onednn_w4a16(
@@ -39,6 +40,10 @@ torch::Tensor onednn_w8a16_fp8(
3940
std::optional<torch::Tensor> bias
4041
);
4142

43+
void fp8_cache_clear();
44+
std::tuple<int64_t, int64_t, int64_t> fp8_cache_stats();
45+
std::tuple<int64_t, int64_t, int64_t> fp8_failure_cache_stats();
46+
4247
torch::Tensor sdp_torch(
4348
torch::Tensor Q,
4449
torch::Tensor K,
@@ -52,9 +57,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
5257
py::arg("x"), py::arg("weight"), py::arg("scales"),
5358
py::arg("bias") = py::none());
5459
m.def("onednn_w8a16_fp8", &onednn_w8a16_fp8,
55-
"onednn FP16 x FP8_E4M3 per-N-scale gemm",
60+
"oneDNN W8A16 FP8 per-N-scale GEMM",
5661
py::arg("x"), py::arg("weight"), py::arg("scales"),
5762
py::arg("bias") = py::none());
63+
m.def("fp8_cache_clear", &fp8_cache_clear,
64+
"Clear cached oneDNN FP8 primitives and counters");
65+
m.def("fp8_cache_stats", &fp8_cache_stats,
66+
"Return FP8 primitive cache (hits, misses, size)");
67+
m.def("fp8_failure_cache_stats", &fp8_failure_cache_stats,
68+
"Return FP8 failure cache (failures, negative hits, size)");
5869
m.def("sdp", &sdp_torch,
5970
"ESIMD Flash Attention SDP [B,L,H,128] PTL-H (fp16/bf16)",
6071
py::arg("Q"), py::arg("K"), py::arg("V"));

0 commit comments

Comments
 (0)