platform

[Watebear]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces sequence parallel support for the LongCat image model, optimizes precision for domestic hardware platforms (using float32 instead of float64 in Wan's pre-inference), and implements platform-specific padding workarounds for bfloat16 on NPUs/MLUs. It also refines attention modules and adds configuration options to disable distributed VAE decoding. The review feedback highlights several critical issues: a performance bottleneck in _infer_token_refiner_attn due to batch-looping and CPU-GPU synchronization, potential runtime crashes in wan_sf_runner.py on non-GPU platforms, a missing fallback check for flash_attn_func_v2 in the worldmirror attention layer, and an invalid value argument being passed to F.pad when using non-constant padding modes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The current implementation of _infer_token_refiner_attn loops over the batch dimension and performs CPU-GPU synchronization (valid_mask.sum().item()) inside the loop. This can introduce significant performance overhead on GPU/NPU platforms.

Since we are using unpadded/varlen attention, we can vectorize this operation across the entire batch by concatenating the valid tokens and constructing a single cu_seqlens tensor for the whole batch. This avoids the Python loop and reduces CPU-GPU synchronization to a single .item() call.

    def _infer_token_refiner_attn(self, attention_module, q, k, v, mask):
        if mask is None:
            mask = torch.ones(q.shape[:2], dtype=torch.bool, device=q.device)
        else:
            mask = mask.to(device=q.device, dtype=torch.bool)
            if mask.ndim == 1:
                mask = mask.unsqueeze(0)
            if mask.shape[0] == 1 and q.shape[0] > 1:
                mask = mask.expand(q.shape[0], -1)

        batch, seqlen, heads, head_dim = q.shape
        out = torch.zeros(batch, seqlen, heads * head_dim, dtype=q.dtype, device=q.device)
        
        valid_lens = mask.sum(dim=-1)
        max_len = int(valid_lens.max().item())
        if max_len == 0:
            return out

        cu_seqlens = torch.zeros(batch + 1, dtype=torch.int32, device=q.device)
        cu_seqlens[1:] = valid_lens.cumsum(dim=0, dtype=torch.int32)

        q_valid = q[mask].contiguous()
        k_valid = k[mask].contiguous()
        v_valid = v[mask].contiguous()

        attn = attention_module.apply(
            q=q_valid,
            k=k_valid,
            v=v_valid,
            cu_seqlens_q=cu_seqlens,
            cu_seqlens_kv=cu_seqlens,
            max_seqlen_q=max_len,
            max_seqlen_kv=max_len,
        )
        out[mask] = attn.reshape(-1, heads * head_dim)
        return out

[gemini-code-assist]

Unconditionally calling getattr(torch, AI_DEVICE) can raise an AttributeError if AI_DEVICE is set to "cpu" or any other platform where the corresponding module does not exist in torch. Additionally, calling torch_device_module.empty_cache() will crash if the module does not have an empty_cache attribute. We should wrap this in a safe helper class to prevent runtime crashes on non-GPU/NPU platforms.

class _DeviceModule:
    @staticmethod
    def empty_cache():
        device_module = getattr(torch, AI_DEVICE, None)
        if device_module is not None and hasattr(device_module, "empty_cache"):
            device_module.empty_cache()

torch_device_module = _DeviceModule()

[gemini-code-assist]

Setting flash_attn_func_v2 = None when flash_attn is not installed will cause a runtime crash in _apply_attention if q.dtype is float16 or bfloat16, because _apply_attention unconditionally calls flash_attn_func_v2 in that branch. Please ensure that _apply_attention is updated to check if flash_attn_func_v2 is not None before calling it, falling back to F.scaled_dot_product_attention if it is unavailable.

[gemini-code-assist]

In PyTorch, F.pad only supports the value argument when mode is "constant". Passing value when mode is "replicate" can raise a ValueError or warning depending on the PyTorch version. We should only pass value when mode == "constant".

Suggested change

	def _pad(x: torch.Tensor, pad, mode: str = "constant", value: float = 0.0) -> torch.Tensor:
	if mode == "replicate" and PLATFORM in _REPLICATE_PAD_FLOAT32_PLATFORMS and x.dtype == torch.bfloat16:
	return F.pad(x.float(), pad, mode=mode, value=value).to(dtype=x.dtype)
	return F.pad(x, pad, mode=mode, value=value)
	def _pad(x: torch.Tensor, pad, mode: str = "constant", value: float = 0.0) -> torch.Tensor:
	if mode == "replicate" and PLATFORM in _REPLICATE_PAD_FLOAT32_PLATFORMS and x.dtype == torch.bfloat16:
	return F.pad(x.float(), pad, mode=mode).to(dtype=x.dtype)
	if mode == "constant":
	return F.pad(x, pad, mode=mode, value=value)
	return F.pad(x, pad, mode=mode)