⚡️ Speed up method MoeWNA16Config.override_quantization_method by 189%

[codeflash-ai]

📄 189% (1.89x) speedup for MoeWNA16Config.override_quantization_method in python/sglang/srt/layers/quantization/moe_wna16.py

⏱️ Runtime : 60.4 milliseconds → 20.9 milliseconds (best of 160 runs)

📝 Explanation and details

The optimized code achieves a 188% speedup by implementing strategic early exits and eliminating redundant expensive calls:

Key Optimizations:

1. Early User Quant Check: The most critical optimization checks user_quant != "moe_wna16" first and returns None immediately. This eliminates the expensive is_moe_wna16_compatible call for ~66% of cases (2023/3032 calls), reducing execution from ~45μs to ~225ns per call.

2. Eliminated Redundant Device Capability Calls: The original code called get_device_capability() for every compatibility check (~112μs per call in the profiler). The optimized version only calls it for AWQ methods that actually need device capability validation, avoiding this expensive operation for GPTQ and unsupported methods.

3. Fast Path Static Method: Extracted compatibility logic into _fast_is_moe_wna16_compatible with optimized control flow:

   • Checks quant_method not in ("gptq", "awq") first to exit early for unsupported methods
   • For GPTQ, validates without any device capability calls
   • For AWQ, performs device capability check only when needed

4. Reduced String Operations: Optimized string handling by checking for empty/None quant_method before calling .lower(), avoiding unnecessary string operations.

Performance Impact by Test Case:

• GPTQ compatible cases: 4000-5000% faster (from ~70μs to ~1-2μs)
• Unsupported methods: 5000-6000% faster (from ~65μs to ~1μs)
• AWQ cases: Still require device capability checks, so smaller improvements (~30% slower to ~25% faster)
• Early exit cases (wrong user_quant): Minimal overhead (~350ns vs ~355ns)

The optimization is particularly effective for validation-heavy workloads where most configurations are either unsupported or use GPTQ (which doesn't need device capability checks), making it ideal for configuration validation pipelines and batch processing scenarios.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3033 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config

# Minimal stubs for external dependencies to keep tests deterministic and isolated.

class DummyAWQConfig:
    @classmethod
    def get_min_capability(cls):
        # AWQ requires device capability >= 75 (Turing or newer)
        return 75

# ---------------------------
# Basic Test Cases
# ---------------------------

def test_gptq_compatible_4bit():
    # Should succeed for GPTQ, 4 bits, desc_act False
    quant_cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 91.7μs -> 2.05μs (4374% faster)

def test_gptq_compatible_8bit():
    # Should succeed for GPTQ, 8 bits, desc_act False
    quant_cfg = {"quant_method": "gptq", "bits": 8, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 72.0μs -> 1.50μs (4704% faster)

def test_awq_compatible_4bit_high_capability():
    # Should succeed for AWQ, 4 bits, high enough device capability
    quant_cfg = {"quant_method": "awq", "bits": 4, "_device_capability": 80}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 67.0μs -> 97.5μs (31.3% slower)

def test_user_quant_not_moe_wna16():
    # Should return None if user_quant is not 'moe_wna16'
    quant_cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "awq") # 355ns -> 367ns (3.27% slower)

def test_quant_method_not_supported():
    # Should return None for unsupported quant_method
    quant_cfg = {"quant_method": "foo", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 77.1μs -> 1.28μs (5908% faster)

# ---------------------------
# Edge Test Cases
# ---------------------------

def test_gptq_desc_act_true():
    # Should return None if desc_act is True (not compatible)
    quant_cfg = {"quant_method": "gptq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 67.6μs -> 1.21μs (5483% faster)

def test_gptq_invalid_bits():
    # Should return None if bits not in [4,8]
    quant_cfg = {"quant_method": "gptq", "bits": 6, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 61.8μs -> 1.32μs (4576% faster)

def test_awq_low_device_capability():
    # Should return None if AWQ device capability is too low
    quant_cfg = {"quant_method": "awq", "bits": 4, "_device_capability": 70}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 63.8μs -> 89.1μs (28.4% slower)

def test_awq_bits_not_4():
    # Should return None if AWQ bits != 4
    quant_cfg = {"quant_method": "awq", "bits": 8, "_device_capability": 80}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 62.9μs -> 1.24μs (4989% faster)

def test_awq_exact_min_capability():
    # Should succeed if device capability is exactly at minimum
    quant_cfg = {"quant_method": "awq", "bits": 4, "_device_capability": 75}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 61.5μs -> 80.1μs (23.1% slower)

def test_missing_quant_method():
    # Should return None if quant_method missing
    quant_cfg = {"bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 62.0μs -> 964ns (6327% faster)

def test_missing_bits():
    # Should return None if bits missing
    quant_cfg = {"quant_method": "gptq", "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 62.8μs -> 1.46μs (4214% faster)

def test_missing_desc_act_gptq():
    # Should return None if desc_act missing for GPTQ
    quant_cfg = {"quant_method": "gptq", "bits": 4}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 60.9μs -> 1.65μs (3591% faster)

def test_missing_device_capability_awq():
    # Should default to device_capability 80 and succeed for AWQ
    quant_cfg = {"quant_method": "awq", "bits": 4}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 61.9μs -> 81.6μs (24.1% slower)

def test_case_insensitive_quant_method():
    # Should handle quant_method case-insensitively
    quant_cfg = {"quant_method": "GptQ", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 60.8μs -> 1.82μs (3237% faster)

# ---------------------------
# Large Scale Test Cases
# ---------------------------

def test_large_batch_gptq():
    # Test with a large batch of GPTQ configs
    for i in range(500):
        quant_cfg = {"quant_method": "gptq", "bits": 4 if i % 2 == 0 else 8, "desc_act": False}
        codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 9.75ms -> 198μs (4818% faster)
        # Now test with an incompatible config
        quant_cfg_bad = {"quant_method": "gptq", "bits": 6, "desc_act": False}
        codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg_bad, "moe_wna16")

def test_large_batch_awq_varied_capability():
    # Test with a large batch of AWQ configs, some above, some below min capability
    for i in range(500):
        cap = 70 + (i % 10)  # 70..79
        quant_cfg = {"quant_method": "awq", "bits": 4, "_device_capability": cap}
        expected = "moe_wna16" if cap >= DummyAWQConfig.get_min_capability() else None
        codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16") # 9.71ms -> 9.72ms (0.037% slower)


def test_mutation_wrong_return_value():
    # If override_quantization_method returns wrong string, test should fail
    quant_cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16"); result = codeflash_output # 98.4μs -> 2.32μs (4137% faster)

def test_mutation_wrong_none_return():
    # If override_quantization_method returns "moe_wna16" for bad config, test should fail
    quant_cfg = {"quant_method": "gptq", "bits": 6, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(quant_cfg, "moe_wna16"); result = codeflash_output # 72.9μs -> 1.39μs (5163% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config

# --- Minimal stubs for dependencies ---

# Simulate CUDA device capability for tests
class DeviceCapabilitySimulator:
    def __init__(self, major=None, minor=None):
        self.major = major
        self.minor = minor
    def __call__(self, device_id=0):
        return (self.major, self.minor)

# Simulate AWQConfig.get_min_capability
class AWQConfig:
    @classmethod
    def get_min_capability(cls):
        return 75

# Patchable global for device capability
get_device_capability = DeviceCapabilitySimulator(8, 0)  # Default: CUDA 8.0 (capability=80)

# --- Unit tests ---

# ----------- BASIC TEST CASES -----------

def test_basic_gptq_compatible_returns_name():
    # Basic: Compatible GPTQ config, user_quant is "moe_wna16"
    cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 73.6μs -> 1.61μs (4460% faster)

def test_basic_awq_compatible_returns_name():
    # Basic: Compatible AWQ config, user_quant is "moe_wna16", device capability sufficient
    get_device_capability.major, get_device_capability.minor = 8, 0  # capability=80 >= 75
    cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 67.4μs -> 97.2μs (30.6% slower)

def test_basic_incompatible_quant_method_returns_none():
    # Basic: Incompatible quant_method, should return None
    cfg = {"quant_method": "other", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 65.8μs -> 1.26μs (5102% faster)

def test_basic_user_quant_not_moe_wna16_returns_none():
    # Basic: user_quant is not "moe_wna16", should return None
    cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "awq") # 342ns -> 354ns (3.39% slower)

def test_basic_gptq_8bit_returns_name():
    # Basic: Compatible GPTQ config with 8 bits
    cfg = {"quant_method": "gptq", "bits": 8, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 73.6μs -> 1.85μs (3867% faster)

# ----------- EDGE TEST CASES -----------

def test_edge_gptq_desc_act_true_returns_none():
    # Edge: GPTQ config with desc_act=True (should NOT be compatible)
    cfg = {"quant_method": "gptq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 66.3μs -> 1.30μs (4992% faster)

def test_edge_gptq_bits_not_4_or_8_returns_none():
    # Edge: GPTQ config with bits=6 (unsupported)
    cfg = {"quant_method": "gptq", "bits": 6, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 67.0μs -> 1.29μs (5101% faster)

def test_edge_awq_bits_not_4_returns_none():
    # Edge: AWQ config with bits=8 (unsupported)
    get_device_capability.major, get_device_capability.minor = 8, 0
    cfg = {"quant_method": "awq", "bits": 8, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 65.4μs -> 1.20μs (5330% faster)

def test_edge_awq_low_device_capability_returns_none():
    # Edge: AWQ config, but device capability too low
    get_device_capability.major, get_device_capability.minor = 7, 0  # capability=70 < 75
    cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 64.4μs -> 92.5μs (30.3% slower)

def test_edge_awq_no_device_capability_returns_none():
    # Edge: AWQ config, but device capability is None
    get_device_capability.major, get_device_capability.minor = None, None
    cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 62.4μs -> 70.4μs (11.4% slower)

def test_edge_missing_keys_returns_none():
    # Edge: Missing keys in config
    cfg = {}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 62.7μs -> 967ns (6383% faster)

def test_edge_quant_method_case_insensitive():
    # Edge: quant_method is uppercase
    cfg = {"quant_method": "GPTQ", "bits": 4, "desc_act": False}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 63.3μs -> 1.77μs (3472% faster)

def test_edge_extra_keys_ignored():
    # Edge: Extra keys in config should not affect result
    cfg = {"quant_method": "gptq", "bits": 4, "desc_act": False, "foo": "bar"}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 63.6μs -> 1.41μs (4401% faster)

def test_edge_device_capability_exact_minimum():
    # Edge: Device capability exactly at minimum for AWQ
    get_device_capability.major, get_device_capability.minor = 7, 5  # capability=75
    cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 62.9μs -> 82.7μs (23.9% slower)

def test_edge_device_capability_negative():
    # Edge: Device capability negative
    get_device_capability.major, get_device_capability.minor = -1, -1
    cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
    codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 60.7μs -> 68.7μs (11.7% slower)

# ----------- LARGE SCALE TEST CASES -----------

def test_large_scale_gptq_configs():
    # Large scale: Many GPTQ configs, all compatible
    configs = [
        {"quant_method": "gptq", "bits": 4, "desc_act": False}
        for _ in range(500)
    ]
    for cfg in configs:
        codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 9.77ms -> 192μs (4976% faster)

def test_large_scale_awq_configs_varied_device_capability():
    # Large scale: Many AWQ configs, half with sufficient device capability, half not
    configs = []
    for i in range(500):
        if i % 2 == 0:
            get_device_capability.major, get_device_capability.minor = 8, 0  # capability=80
        else:
            get_device_capability.major, get_device_capability.minor = 7, 0  # capability=70
        cfg = {"quant_method": "awq", "bits": 4, "desc_act": True}
        configs.append((dict(cfg), get_device_capability.major, get_device_capability.minor))
    for idx, (cfg, major, minor) in enumerate(configs):
        get_device_capability.major, get_device_capability.minor = major, minor
        if major * 10 + minor >= AWQConfig.get_min_capability():
            codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16")
        else:
            codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16")

def test_large_scale_incompatible_configs():
    # Large scale: Many incompatible configs
    configs = [
        {"quant_method": "gptq", "bits": 6, "desc_act": False}
        for _ in range(250)
    ] + [
        {"quant_method": "awq", "bits": 8, "desc_act": True}
        for _ in range(250)
    ]
    get_device_capability.major, get_device_capability.minor = 8, 0
    for cfg in configs:
        codeflash_output = MoeWNA16Config.override_quantization_method(cfg, "moe_wna16") # 9.71ms -> 171μs (5557% faster)

To edit these changes git checkout codeflash/optimize-MoeWNA16Config.override_quantization_method-mhoyi1m3 and push.