Add HPU FP8 Linear GEMM support for Wan transformer blocks

[Copilot]

Adds --fp8 flag to Wan inference pipelines (i2v, t2v, ti2v, s2v) that quantises all nn.Linear layers inside model.blocks to per-output-channel FP8 at init time, freeing BF16 weights immediately. Forward pass uses dynamic per-row activation quantisation and HPU fp8_gemm_v2. All other Linear layers (embeddings, heads, etc.) remain BF16.

New: wan/utils/fp8_linear.py

• WanFP8Linear — nn.Module wrapper for nn.Linear; precomputes weight_fp8 ([out, in], float8_e4m3fn) and weight_scale ([out, 1], float32) as buffers, deletes original BF16 weight at init
• apply_fp8_linear_hpu — HPU FP8 GEMM via torch.ops.hpu.fp8_gemm_v2 with dynamic per-row activation quant (dynamic_quant)
• wrap_blocks_linear_fp8(model) — replaces every nn.Linear in model.blocks in-place; no-ops on anything outside blocks

Pipeline changes (WanI2V, WanT2V, WanTI2V, WanS2V)

• __init__ and _configure_model accept fp8=False; when enabled, wrap_blocks_linear_fp8 runs before device placement so buffers are moved to HPU alongside the rest of the model

CLI (generate.py)

• --fp8 flag wired through to all four pipeline constructors

# Example: enable FP8 for i2v inference
python generate.py --task i2v-A14B --fp8 --ckpt_dir /path/to/ckpt ...

Constraint: weight compression is strictly a one-time init-time operation — forward only quantises activations dynamically.

Original prompt

In repository wenbinc-Bin/Wan2.2, create a pull request targeting branch hpu to add HPU FP8 Linear GEMM support for Wan inference models.

Requirements:

• Add a new runtime/model initialization flag --fp8.
• When --fp8 is enabled during model initialization, convert eligible Linear layers into an FP8-wrapped implementation at initialization time.
• Do not quantize/compress weights during runtime forward; weight compression must happen only once during model initialization.
• Only wrap Linear layers inside self.blocks for:
  • WanModel
  • WanModel_S2V
• Do not convert other Linear layers in the model; they should remain BF16.
• The wrapped Linear should support FP8 GEMM forward on HPU using logic analogous to the provided apply_fp8_linear_hpu(...) helper and dynamic quantization path for activations.
• During initialization, compress Linear weights into per-channel FP8 GEMM weights/scales and delete/drop the original BF16 Linear weights to save memory.
• Consider implementing the wrapping/configuration in _configure_model if appropriate.
• The implementation should follow the spirit of the provided example:
  • Wrap torch.nn.Linear
  • Precompute/cache per-output-channel FP8 weights and scales at init time
  • Use HPU FP8 GEMM in forward
• Preserve existing behavior when --fp8 is not provided.
• Update any CLI/config plumbing needed so the option is usable for image2video / speech2video / text2video / textimage2video.

Relevant context from user:

• Reference file discussed: wan/textimage2video.py on branch hpu
• Suggested helper behavior:
  • apply_fp8_linear_hpu(input, weight, weight_scale, input_scale=None, bias=None, trans_B=True)
  • dynamic_quant(data, single_scale=False)
• The original prototype also used a wrapped Linear that precomputes per-channel FP8 weights/scales and frees BF16 weights.

Deliverables:

• Code changes in the repo on a new branch from hpu
• Open a pull request back to hpu
• PR title should clearly mention HPU FP8 Linear GEMM support for Wan blocks
• Include a concise PR description summarizing scope and constraints

The following is the prior conversation context from the user's chat exploration (may be truncated):

User: ```
class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule):
def init(
self,
module: torch.nn.Linear,
offload_dtype: torch.dtype = None,
offload_device: Union[str, torch.device] = None,
onload_dtype: torch.dtype = None,
onload_device: Union[str, torch.device] = None,
preparing_dtype: torch.dtype = None,
preparing_device: Union[str, torch.device] = None,
computation_dtype: torch.dtype = None,
computation_device: Union[str, torch.device] = None,
vram_limit: float = None,
name: str = "",
disk_map: DiskMap = None,
**kwargs
):
with skip_model_initialization():
super().init(
in_features=module.in_features,
out_features=module.out_features,
bias=module.bias is not None,
)
self.set_dtype_and_device(
offload_dtype,
offload_device,
onload_dtype,
onload_device,
preparing_dtype,
preparing_device,
computation_dtype,
computation_device,
vram_limit,
)
self.weight = module.weight
self.bias = module.bias
self.state = 0
self.name = name
self.lora_A_weights = []
self.lora_B_weights = []
self.lora_merger = None

    # Always enable FP8 GEMM path
    self.enable_fp8 = True

    # FP8 GEMM dtype is decoupled from computation_dtype
    # so computation_dtype can stay BF16 while GEMM uses FP8
    self.fp8_gemm_dtype = torch.float8_e4m3fn

    self.computation_device_type = parse_device_type(self.computation_device)
    self.fp8_weight_scale_b = None
    self.fp8_weight = None
    self.fp8_weight_cache_valid = False

    if offload_dtype == "disk":
        self.disk_map = disk_map
        self.disk_offload = True
    else:
        self.disk_offload = False

def get_fp8_max(self):
    fp8_max = 448.0
    if self.fp8_gemm_dtype == torch.float8_e4m3fnuz:
        fp8_max = fp8_max / 2.0
    return fp8_max

def invalidate_fp8_weight_cache(self):
    self.fp8_weight_scale_b = None
    self.fp8_weight = None
    self.fp8_weight_cache_valid = False

def prepare_fp8_weight_cache(self, weight: torch.Tensor):
    if weight is None:
        self.invalidate_fp8_weight_cache()
        return

    device = weight.device
    fp8_max = self.get_fp8_max()

    # PyTorch Linear weight shape: [out_features, in_features] == [n, k]
    # per output channel scale => one scale per row => [n, 1]
    w_max = torch.max(torch.abs(weight), dim=1, keepdim=True).values

...

This pull request was created from Copilot chat.

[mengker33]

There is no fp8 for q in cross-attn. And i asked copilot, does the answer below make sense?

[Wei-Lin-Intel]

There is fp8 for q in cross-attn because cross-attn is in transformer blocks. But in cross-attn, q Linear uses x as the input, and k/v Linears use context as the input, so the dequant of context can be shared for k/v, unlike q/k/v shares the same input tensor in self-attn.
All the Linears in transformer blocks has been wrapped by WanFP8Linear. Its forward function would perform dequant for input tensor. Since some Linears share the input in self/cross-attn, I seperate the dequant OP for them.