diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 2c6fe00303e5..713da2bdeab6 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1978,6 +1978,11 @@ def _fp8_block_scaling_gemm_sm100_constraint(inputs: List[List[int]]) -> int: return inputs[0][0] +def _fp8_block_scaling_gemm_sm120_constraint(inputs: List[List[int]]) -> int: + # SM120 activation scales are [padded_m, k // 512]. + return fp4_utils.pad_up(inputs[0][0], 4) + + def _fp8_quantize_1x128_sm90_constraint(inputs: List[List[int]]) -> int: # The implementation aligns with the fp8_quantize_1x128 custom op. pad_m = fp4_utils.pad_up(inputs[0][0], 4) @@ -1988,6 +1993,9 @@ def _fp8_quantize_1x128_sm90_constraint(inputs: List[List[int]]) -> int: @lru_cache(maxsize=None) def _get_fp8_block_scaling_gemm_constraint_spec( sm_version: int) -> Tuple[ConstraintSpec, ...]: + if sm_version == 120: + return (ConstraintSpec(2, 0, + _fp8_block_scaling_gemm_sm120_constraint), ) if sm_version >= 100: return (ConstraintSpec(2, 1, _fp8_block_scaling_gemm_sm100_constraint), ) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 264fa36fb1b3..13f402129b7a 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -59,7 +59,6 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutl accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] SKIP (https://nvbugs/5596343) -accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 SKIP (https://nvbugs/6647310) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/5616182) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/6437412) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[pp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/6278337) @@ -271,7 +270,6 @@ full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp1pp2] SKIP (https://nvbugs/6649818) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) -full:RTXPro6000D/accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6652876) full:RTXPro6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/5929339) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6273850) full:RTX_6000D/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp SKIP (https://nvbugs/6275856)