[Bug]: torch.opcheck fails for `_C.rms_norm_per_block_quant`

[ProExpertProg]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

In the unit test for torch.ops._C.rms_norm_per_block_quant custom kernel, for some reason opcheck fails because it thinks the weight tensor got mutated. A closer look reveals a weird issue: the cloned weight arg is the one that gets modified, and the original weight arg stays intact. I could not find a memory issue, I manually confirmed the original weight stays intact when not using opcheck, and E2E evals look good.

torch.testing._internal.optests.generate_tests.OpCheckError: opcheck(op, ...): test_schema failed with Argument weight is not defined as mutable but was mutated (scroll up for stack trace)

vllm/tests/kernels/core/test_fused_quant_layernorm.py

Lines 291 to 304 in 0ebf4e9

	opcheck(
	torch.ops._C.rms_norm_per_block_quant,
	(
	output,
	x,
	layer.weight,
	scales,
	1e-5,
	scale_ub,
	residual,
	group_size[1],
	True, # is_scale_transposed
	),
	)

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.