please register them via envs.py

These patterns (including _is_hip) should be avoided.
Decoding is a critical core stage of the system. If we allow platform-specific conditionals to accumulate here, the complexity will grow rapidly and become increasingly hard to manage.

Over time, this will significantly hurt maintainability and extensibility.

Suggestions:

• Introduce a Platform abstraction
• Consider a plugin-based architecture

discussions are welcomed

The codebase already has a well-established pattern for isolating platform-specific logic via Platform polymorphic methods:

• Platform.get_attn_backend_cls_str() — platform decides which attention backend to use
• Platform.enable_dit_layerwise_offload_for_wan_by_default() — platform decides default behavior
• Platform.get_device_communicator_cls() — platform decides communicator

We can directly reuse this Strategy pattern:

1. Add optimize_vae() hook to Platform base class

# interface.py
class Platform:
    ...
    @classmethod
    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
        “”"Apply platform-specific optimizations to VAE after loading.“”"
        return vae

2. RocmPlatform overrides it with AITer GroupNorm replacement

# rocm.py
class RocmPlatform(Platform):
    ...
    @classmethod
    def optimize_vae(cls, vae: torch.nn.Module) -> torch.nn.Module:
        if not envs.SGLANG_USE_ROCM_VAE:
            return vae
        try:
            from aiter.ops.groupnorm import GroupNorm as AITerGroupNorm
            count = cls._replace_groupnorm(vae, AITerGroupNorm)
            logger.info(“replaced %d GroupNorm with AITer GroupNorm”, count)
        except Exception as e:
            logger.warning(“failed to apply AITer GroupNorm: %s”, e)
        return vae

    @staticmethod
    def _replace_groupnorm(module, aiter_gn_cls):
        count = 0
        for name, child in module.named_children():
            if isinstance(child, torch.nn.GroupNorm):
                replacement = aiter_gn_cls(child.num_groups, child.num_channels, ...)
                replacement.weight = child.weight  # zero-copy weight sharing
                replacement.bias = child.bias
                setattr(module, name, replacement)
                count += 1
            else:
                count += RocmPlatform._replace_groupnorm(child, aiter_gn_cls)
        return count

3. Call it from VAELoader — one line

VAELoader already does similar model-level optimizations (e.g. _convert_conv3d_weights_to_channels_last_3d). Just add one line before returning:

# vae_loader.py
vae = current_platform.optimize_vae(vae)
return vae

4. decoding.py — zero changes needed

The core decoding stage stays completely platform-agnostic. No _is_hip, no _use_aiter_vae, no _replace_groupnorm_with_aiter.

Why this is the right approach:

• Follows the existing architecture exactly (same pattern as get_attn_backend_cls_str)
• Separation of concerns: ROCm code lives only in rocm.py; core stages don’t know about platforms
• Extensible: future NPU/MUSA VAE optimizations just override optimize_vae() in their own platform file — no core code changes needed
• Env var SGLANG_USE_ROCM_VAE is registered in envs.py and consumed internally by RocmPlatform
• Minimal diff: ~3 files, a few lines each

@mickqian Thanks for the detailed review and suggestions, I'll follow the Platform abstraction pattern and address this in a follow-up PR.

Hi @mickqian could you help review this PR-20496 #20496?
This PR moves platform-specific AITer GroupNorm replacement out of decoding.py into the Platform polymorphic pattern:

• Register SGLANG_USE_ROCM_VAE in envs.py
• Add optimize_vae() hook to Platform base class
• Override in RocmPlatform with AITer GroupNorm replacement
• Call from VAELoader at load time
• Remove all platform-specific logic from decoding.py

Thanks.