Repository: vllm-project/vllm
File: tests/kernels/core/test_fused_quant_layernorm.py (Python)

Code before:
            x, weight, EPS, quant_dtype, scale_ub, residual
        )
    return out, scales, residual


def ops_impl(
    weight: torch.Tensor,
    x: torch.Tensor,
    quant_dtype: torch.dtype,
    residual: torch.Tensor | None,
    scale_ub: torch.Tensor | None,
    group_size: list[int] | None,
    tma_alignment: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    return ops_dynamic_per_token_or_block_quant(
        weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
    )


@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
@pytest.mark.parametrize("group_size", GROUP_SIZES)
@pytest.mark.parametrize("tma_alignment", TMA_ALIGNMENTS)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_rms_norm(
    default_vllm_config,
    num_tokens: int,
    hidden_size: int,
    add_residual: bool,
    has_scale_ub: bool,
    dtype: torch.dtype,
    quant_dtype: torch.dtype,
    group_size: list[int] | None,
    tma_alignment: int,
    seed: int,
    device: str,
) -> None:
    torch.random.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
    torch.set_default_device(device)
    torch.cuda.set_device(device)

    if group_size is not None and hidden_size % group_size[1] != 0:
        # skip
        return

Review comment (suggestion):
```suggestion
@pytest.mark.parametrize("group_size, tma_alignment", [
    (None, 0),
    *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)
)
```

Code after:
    else:
        out, scales = ops.rms_norm_dynamic_per_token_quant(
            x, weight, EPS, quant_dtype, scale_ub, residual
        )
    return out, scales, residual


def ops_impl(
    weight: torch.Tensor,
    x: torch.Tensor,
    quant_dtype: torch.dtype,
    residual: torch.Tensor | None,
    scale_ub: torch.Tensor | None,
    group_size: list[int] | None,
    tma_alignment: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    return ops_dynamic_per_token_or_block_quant(
        weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
    )


@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
@pytest.mark.parametrize(
    "group_size, tma_alignment",
    [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_rms_norm(
    default_vllm_config,
    num_tokens: int,
    hidden_size: int,
    add_residual: bool,
    has_scale_ub: bool,
    dtype: torch.dtype,
    quant_dtype: torch.dtype,
    group_size: list[int] | None,
    tma_alignment: int,
    seed: int,
    device: str,
) -> None:
    torch.random.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
    torch.set_default_device(device)
    torch.cuda.set_device(device)