Fix LRU cache pollution causing BLOCK_SIZE_S3 KeyError in gemm_afp4wfp4

[Duyi-Wang]

Motivation

Fix KeyError: 'Keyword argument BLOCK_SIZE_S3 was specified but unrecognised' crash in gemm_afp4wfp4 triggered when using quark_w4a4_mxfp4 quantization in SGLang.

Technical Details

get_gemm_config() in aiter/ops/triton/utils/gemm_config_utils.py uses @functools.lru_cache, which caches the returned (dict, bool) tuple by reference. The fused kernel fused_gemm_afp4wfp4_split_cat mutates the cached dict in-place by injecting BLOCK_SIZE_S3:

config["BLOCK_SIZE_S3"] = triton.next_power_of_2(...)

Subsequent calls to gemm_afp4wfp4 with the same (M, N, K) receive the same polluted dict from the cache. When unpacked via **config into a Triton kernel that doesn't accept BLOCK_SIZE_S3, it raises a KeyError.

Fix: Added config.pop("BLOCK_SIZE_S3", None) in the three affected functions in gemm_afp4wfp4.py (gemm_afp4wfp4_, gemm_afp4wfp4_preshuffled_scales, gemm_afp4wfp4_preshuffle) immediately after obtaining the config dict. This strips the polluted key before it reaches the kernel call, regardless of whether the cache was polluted or not.

This is a minimal workaround. A more thorough fix would be to have get_gemm_config() return defensive copies to prevent cache mutation entirely.

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

Addresses a Triton kernel launch crash (KeyError: '... BLOCK_SIZE_S3 ... unrecognised') caused by BLOCK_SIZE_S3 being injected into (and then reused from) a cached GEMM config when using MXFP4 quantization paths.

Changes:

• Strip BLOCK_SIZE_S3 from the GEMM config dict after config retrieval in gemm_afp4wfp4_.
• Strip BLOCK_SIZE_S3 from the GEMM config dict after config retrieval in gemm_afp4wfp4_preshuffled_scales.
• Strip BLOCK_SIZE_S3 from the GEMM config dict after config retrieval in gemm_afp4wfp4_preshuffle.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[k50112113]

Instead of popping BLOCK_SIZE_S3 with a default to None in the gemm_afp4wfp4 API, my suggestion would be to add .copy() at _get_config() at /root/aiter/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_afp4wfp4.py

like the following:

def _get_config(
    M: int,
    N: int,
    K: int,
    shuffle: bool = False,
):
    # Note: Config files use K=2*K in their naming
    K = 2 * K
    if shuffle:
        return get_gemm_config("GEMM-AFP4WFP4_PRESHUFFLED", M, N, K).copy()
    else:
        return get_gemm_config("GEMM-AFP4WFP4", M, N, K).copy()

to copy the config dict

so that the configs of both the gemm_afp4wfp4 and the fused_gemm_afp4wfp4_split_cat

Please let me know if this works

Thanks,
Shao-Chun Lee

[Duyi-Wang]

Instead of popping BLOCK_SIZE_S3 with a default to None in the gemm_afp4wfp4 API, my suggestion would be to add .copy() at _get_config() at /root/aiter/aiter/ops/triton/_triton_kernels/gemm/basic/gemm_afp4wfp4.py

like the following:

def _get_config(
    M: int,
    N: int,
    K: int,
    shuffle: bool = False,
):
    # Note: Config files use K=2*K in their naming
    K = 2 * K
    if shuffle:
        return get_gemm_config("GEMM-AFP4WFP4_PRESHUFFLED", M, N, K).copy()
    else:
        return get_gemm_config("GEMM-AFP4WFP4", M, N, K).copy()

to copy the config dict

so that the configs of both the gemm_afp4wfp4 and the fused_gemm_afp4wfp4_split_cat

Please let me know if this works

Thanks, Shao-Chun Lee

Returning a deep copy in _get_config() would be a better solution to the cache pollution issue and it works, but I’m unsure about its impact on performance or other components. Therefore, I’m using this as a workaround for my current workload.
The fix based on deep-copy has been pushed.

[Duyi-Wang]

Similar one merged #2173. However, this is the same workaround and is only applicable to gemm_afp4wfp4. @k50112113

[k50112113]

Yes, this looks good, thanks for the work