Refine scheme #82
Refine scheme #82
Conversation
Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: yiliu30 <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the quantization scheme implementation, focusing on MXFP4. It introduces a new fp4_utils.py file to centralize FP4-related operations and cleans up commented-out code and improves code style in several places. My review identifies a critical bug in an assertion that fails to check for tensor contiguity correctly, and a high-severity performance issue due to a redundant and inefficient function call which can be removed.
| orig_shape = result_shape | ||
| else: | ||
| raise AssertionError("unsupported") | ||
| assert data_lp.is_contiguous, f"Data must be contiguous, got {data_lp.stride()}" |
There was a problem hiding this comment.
This assertion has a syntax error and introduces a functional regression.
- Buggy Assertion:
is_contiguousis a method and should be called asdata_lp.is_contiguous(). The current codeassert data_lp.is_contiguous, ...is alwaysTruebecause it asserts on the method object itself, not its result. This means the contiguity check is effectively disabled, which could lead to silent errors or incorrect behavior downstream if a non-contiguous tensor is passed. - Functional Regression: The previous implementation handled non-contiguous tensors by transposing them. This new implementation removes that logic and intends to enforce contiguity via an assertion. This is a significant change in behavior that could break callers if they pass non-contiguous tensors.
Please fix the assertion and consider if dropping support for non-contiguous tensors is intended and safe.
| assert data_lp.is_contiguous, f"Data must be contiguous, got {data_lp.stride()}" | |
| assert data_lp.is_contiguous(), f"Data must be contiguous, got {data_lp.stride()}" |
|
|
||
| data_lp = data_lp.reshape(orig_shape) | ||
| orig_shape = [*orig_shape[:-1], orig_shape[-1] // 2] | ||
| data_lp = cast_to_fp4(data_lp) |
There was a problem hiding this comment.
The call to cast_to_fp4 on this line is redundant. The pack_fp4_to_uint8 function, called on the next line, already handles rounding the float tensor to the nearest FP4 value before packing. Removing this line will improve performance and simplify the code. The cast_to_fp4 function itself is also implemented inefficiently.
| def cast_to_fp4(x): | ||
| sign = torch.sign(x) | ||
| x = torch.abs(x) | ||
| x[(x >= 0.0) & (x <= 0.25)] = 0.0 | ||
| x[(x > 0.25) & (x < 0.75)] = 0.5 | ||
| x[(x >= 0.75) & (x <= 1.25)] = 1.0 | ||
| x[(x > 1.25) & (x < 1.75)] = 1.5 | ||
| x[(x >= 1.75) & (x <= 2.5)] = 2.0 | ||
| x[(x > 2.5) & (x < 3.5)] = 3.0 | ||
| x[(x >= 3.5) & (x <= 5.0)] = 4.0 | ||
| x[x > 5.0] = 6.0 | ||
| return x * sign No newline at end of file |
There was a problem hiding this comment.
Following the removal of its only call site in mxfp4_qdq_utils.py, this function cast_to_fp4 is now unused and can be removed. Additionally, its implementation using multiple masked assignments is inefficient on GPUs.
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
Purpose
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)