Fix: w8a8_block_fp8_matmul_triton issues in specific scenarios.

[solrex]

Motivation

Currently, w8a8_block_fp8_matmul_triton is the only block-wise FP8 kernel available for sm89 (DeepGEMM is not used). This PR addresses 2 persistent issues with this kernel:

1. The kernel triggers a series of torch.compile warnings, , as shown in the logs:

[rank0]:W0821 14:14:21.867000 202313 torch/_dynamo/variables/builtin.py:898] [13/1] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler
[rank0]:W0821 14:14:26.136000 202313 torch/_dynamo/variables/builtin.py:898] [13/2] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler
[rank0]:W0821 14:14:29.005000 202313 torch/_dynamo/variables/builtin.py:898] [13/3] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler
[rank0]:W0821 14:14:30.633000 202313 torch/_dynamo/variables/builtin.py:898] [13/4] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler
[rank0]:W0821 14:15:56.041000 202313 torch/_dynamo/variables/builtin.py:898] [13/6] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler
[rank0]:W0821 14:16:00.433000 202313 torch/_dynamo/variables/builtin.py:898] [13/7] incorrect arg count <bound method BuiltinVariable._call_min_max of BuiltinVariable(min)> got an unexpected keyword argument 'key' and no constant handler

2. The test/srt/quant/test_fp8_kernel.py unit test fails, as demonstrated by the following output when running with SGL_ENABLE_JIT_DEEPGEMM=0:

# SGL_ENABLE_JIT_DEEPGEMM=0 python test/srt/quant/test_fp8_kernel.py 
.E
======================================================================
ERROR: test_w8a8_block_fp8_matmul (__main__.TestW8A8BlockFP8Matmul.test_w8a8_block_fp8_matmul)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/utils.py", line 2199, in retry
    return fn()
           ^^^^
  File "/sgl-workspace/sglang/python/sglang/test/test_utils.py", line 1365, in <lambda>
    lambda: super(CustomTestCase, self)._callTestMethod(method),
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: tensor model parallel group is not initialized

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/test/test_utils.py", line 1364, in _callTestMethod
    retry(
  File "/sgl-workspace/sglang/python/sglang/srt/utils.py", line 2202, in retry
    raise Exception(f"retry() exceed maximum number of retries.")
Exception: retry() exceed maximum number of retries.

----------------------------------------------------------------------
Ran 2 tests in 1.128s

Modifications

• Refactor code to resolve _call_min_max of BuiltinVariable(min) torch.compile warnings.
• Fix log_info_on_rank0 to ensure w8a8_block_fp8_matmul_triton unit tests run successfully.

You can verify the fix for the torch.compile issue with the following code snippet:

import torch
import json

with open("/sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/N=7168,K=256,device_name=NVIDIA_H20,dtype=int8_w8a8,block_shape=[128, 128].json") as f:
    configs = {int(key): val for key, val in json.load(f).items()}

def F():
    config = configs[min(configs.keys(), key=lambda x: abs(x - 64))]
    return config

def F_new():
    config = configs[min([(abs(k - 64), k) for k in configs])[1]]
    return config

CF = torch.compile(F)
print("Compiled F:", CF())
CF_new = torch.compile(F_new)
print("Compiled F_new:", CF_new())

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Summary of Changes

Hello @solrex, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses critical issues within the w8a8_block_fp8_matmul_triton kernel, which is currently the sole block-wise FP8 kernel available for sm89. The changes resolve persistent torch.compile warnings related to argument handling and fix a unit test failure caused by an uninitialized tensor model parallel group, ensuring the kernel's stability and testability.

Highlights

• Resolved torch.compile warnings: The method for finding the optimal configuration in fp8_kernel.py was refactored to prevent torch.compile from issuing incorrect arg count warnings when using min with a key argument. This involved changing the lookup logic to explicitly create tuples for comparison.
• Fixed unit test failure: The log_info_on_rank0 utility function in utils.py was updated to correctly handle scenarios where the tensor model parallel group is not initialized. This modification prevents an AssertionError that previously caused a unit test to fail, ensuring the test suite runs successfully.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively addresses two separate issues: a torch.compile warning and a unit test failure due to an uninitialized parallel group. The fixes are correct and well-targeted. I have one suggestion to improve the memory efficiency of the torch.compile fix by using a generator expression instead of a list comprehension.