[main][Refactor] Quantization Module Refactor

[SlightwindSec]

Summary

This PR refactors the vllm_ascend/quantization module to improve code organization, maintainability, and extensibility. The refactoring introduces a clear separation of concerns with a registry-based scheme discovery pattern, abstract base classes for quantization schemes, and dedicated wrapper classes.

Key Changes

1. Modular Directory Structure

Before	After
Flat file structure with mixed responsibilities	Organized into methods/ subpackage for schemes
Single quant_config.py (600+ lines)	Separate config files: modelslim_config.py, compressed_tensors_config.py
utils.py with scheme lookup logic	methods/registry.py with decorator-based registration

2. Registry-Based Scheme Discovery

Replaced hardcoded ASCEND_QUANTIZATION_METHOD_MAP dictionary with a decorator-based registry pattern:

# Before: Manual dictionary mapping
ASCEND_QUANTIZATION_METHOD_MAP = {
    "W8A8_DYNAMIC": {"linear": AscendW8A8DynamicLinearMethod, ...},
    ...
}

# After: Decorator-based registration
@register_scheme("W8A8_DYNAMIC", "linear")
class AscendW8A8DynamicLinearMethod(AscendLinearScheme):
    ...

3. Abstract Base Classes

Introduced three abstract base classes in methods/base.py:

• AscendLinearScheme - Base for linear layer quantization
• AscendMoEScheme - Base for MoE layer quantization
• AscendAttentionScheme - Base for attention layer quantization

4. Separated Config and Wrapper Classes

• Config classes (AscendModelSlimConfig, AscendCompressedTensorsConfig): Handle config parsing and scheme selection
• Wrapper classes (AscendLinearMethod, AscendFusedMoEMethod, etc.): Implement vLLM interfaces and delegate to schemes

5. Cleaner Public API

# New clean module interface
from vllm_ascend.quantization import (
    AscendModelSlimConfig,
    AscendCompressedTensorsConfig,
)
from vllm_ascend.quantization.methods import get_scheme_class

Architecture Diagram

classDiagram
    direction TB
    
    class QuantizationConfig {
        <<vLLM Interface>>
        +get_quant_method()
    }
    
    class AscendModelSlimConfig {
        +quant_description
        +get_quant_method()
        -create_scheme_for_layer()
    }
    
    class AscendCompressedTensorsConfig {
        +target_scheme_map
        +get_quant_method()
        -_get_scheme_from_parts()
    }
    
    class AscendLinearMethod {
        <<Wrapper>>
        +quant_method: AscendLinearScheme
        +create_weights()
        +apply()
    }
    
    class AscendFusedMoEMethod {
        <<Wrapper>>
        +quant_method: AscendMoEScheme
        +create_weights()
        +apply()
    }
    
    class AscendLinearScheme {
        <<Abstract>>
        +get_weight()*
        +apply()*
        +get_pertensor_param()
        +get_perchannel_param()
    }
    
    class AscendMoEScheme {
        <<Abstract>>
        +get_weight()*
        +get_dynamic_quant_param()*
        +apply()*
    }
    
    class W8A8DynamicLinear {
        +get_weight()
        +apply()
    }
    
    class W8A8DynamicMoE {
        +get_weight()
        +apply()
    }
    
    QuantizationConfig <|-- AscendModelSlimConfig
    QuantizationConfig <|-- AscendCompressedTensorsConfig
    
    AscendModelSlimConfig ..> AscendLinearMethod : creates
    AscendModelSlimConfig ..> AscendFusedMoEMethod : creates
    AscendCompressedTensorsConfig ..> AscendLinearMethod : creates
    AscendCompressedTensorsConfig ..> AscendFusedMoEMethod : creates
    
    AscendLinearMethod o-- AscendLinearScheme : delegates to
    AscendFusedMoEMethod o-- AscendMoEScheme : delegates to
    
    AscendLinearScheme <|-- W8A8DynamicLinear
    AscendMoEScheme <|-- W8A8DynamicMoE

Loading

Scheme Registration Flow

sequenceDiagram
    participant Module as Scheme Module
    participant Registry as _SCHEME_REGISTRY
    participant Config as QuantConfig
    participant Wrapper as Wrapper Class
    
    Note over Module: At import time
    Module->>Registry: @register_scheme("W8A8_DYNAMIC", "linear")
    Registry->>Registry: Store (quant_type, layer_type) -> Class
    
    Note over Config: At runtime
    Config->>Config: Determine quant_type from description
    Config->>Registry: get_scheme_class(quant_type, layer_type)
    Registry-->>Config: Return scheme class
    Config->>Config: scheme = scheme_cls()
    Config->>Wrapper: Create wrapper with scheme
    Wrapper-->>Config: Return wrapper instance

Loading

File Changes Summary

Original Files	Refactored Files
__init__.py (empty)	__init__.py (exports public API)
quant_config.py	modelslim_config.py + wrappers.py
compressed_tensors/	compressed_tensors_config.py
utils.py	methods/registry.py
w8a8_dynamic.py	methods/w8a8_dynamic.py
w8a8.py	methods/w8a8_static.py
w4a4_flatquant_dynamic.py	methods/w4a4_flatquant.py
...	methods/base.py (new)

Benefits

1. Extensibility: Adding new quantization schemes only requires implementing the base class and adding @register_scheme decorator
2. Maintainability: Clear separation between config parsing, wrapper logic, and scheme implementation
3. Testability: Abstract base classes enable easier unit testing and mocking
4. Discoverability: Registry pattern makes it easy to list all supported schemes
5. Reduced Coupling: Config classes no longer need to know about all scheme implementations

---

• vLLM version: v0.13.0
• vLLM main: vllm-project/vllm@2f4e654

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-structured refactoring of the quantization module. The introduction of base classes for quantization schemes, a central registry, and the separation of scheme implementations from wrapper logic are excellent architectural improvements. This will greatly enhance modularity and make it easier to add new quantization algorithms in the future. I've identified a couple of areas where the design could be made more robust to prevent potential issues. Overall, this is a high-quality contribution.

[gemini-code-assist]

The get_pertensor_param method in AscendLinearScheme lacks the input_size parameter, which is necessary for schemes like W4A4_FLATQUANT_DYNAMIC that depend on it. Currently, W4A4_FLATQUANT_DYNAMIC uses a class attribute to pass input_size from get_weight to get_pertensor_param, which is not thread-safe and is a fragile design. To make the interface more robust, input_size should be passed directly.

    def get_pertensor_param(self, input_size: int,
                            params_dtype: torch.dtype) -> Dict[str, Any]:
        """Return per-tensor parameter specifications (e.g., input_scale).
        
        Args:
            input_size: Input dimension of the linear layer.
            params_dtype: Data type for parameters.
            
        Returns:
            Dictionary mapping parameter names to empty tensors.
        """
        return {}

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[weijinqian0]

why not return QuantType.W4A8 directly?

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.