diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 79c2f064e1ac..ba01e7ad536a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -339,6 +339,11 @@ def is_situ_activation(self) -> bool: return self.trtllm_gen_activation_type == ActType_TrtllmGen.SiTu def _validate_backend_local_activation(self) -> None: + # Runs from __init__, before create_weights, so the swiglu_* attributes + # checked below are still the constructor-provided values. For SiTu, + # create_weights later reuses the swiglu_alpha/swiglu_beta storage for + # the backend-local activation parameters (SiTu and SwiGLU are mutually + # exclusive and feed the same gemm1_alpha/gemm1_beta op slots). if self.trtllm_gen_activation_type is None: if (self.trtllm_gen_activation_alpha is not None or self.trtllm_gen_activation_beta is not None): @@ -549,7 +554,9 @@ def _check_configs(self): raise ValueError( "TRTLLM-Gen SiTu requires MXFP4 scaling vector size 32, " f"got {self.scaling_vector_size}.") - for name in ("situ_alpha", "situ_beta"): + # For SiTu these hold the backend-local activation parameters + # (populated by create_weights, which runs before this check). + for name in ("swiglu_alpha", "swiglu_beta"): value = getattr(self, name) if (value.dtype != torch.float32 or value.shape != (self.expert_size_per_partition, ) @@ -594,19 +601,26 @@ def create_weights(self): else: self.quant_method.create_weights(self) + # SiTu reuses the swiglu_alpha/swiglu_beta storage: SiTu and SwiGLU are + # mutually exclusive (constructor-provided SwiGLU parameters are + # rejected by _validate_backend_local_activation) and feed the same + # gemm1_alpha/gemm1_beta op slots. Safe with respect to the + # `swiglu_alpha is not None` gates: create_moe.py checks the + # constructor kwargs (None for SiTu); _get_quant_method consults + # swiglu_alpha only on the nvfp4 branch (SiTu requires + # W4A8_MXFP4_MXFP8) and has already run above; _check_configs runs + # after this point and its swiglu gate admits w4a8_mxfp4_mxfp8. if self.is_situ_activation: - situ_alpha = nn.Parameter(torch.full( + self.swiglu_alpha = nn.Parameter(torch.full( (self.expert_size_per_partition, ), float(self.trtllm_gen_activation_alpha), dtype=torch.float32), - requires_grad=False) - situ_beta = nn.Parameter(torch.full( + requires_grad=False) + self.swiglu_beta = nn.Parameter(torch.full( (self.expert_size_per_partition, ), float(self.trtllm_gen_activation_beta), dtype=torch.float32), - requires_grad=False) - self.register_parameter("situ_alpha", situ_alpha) - self.register_parameter("situ_beta", situ_beta) + requires_grad=False) self._weights_created = True self._check_configs() @@ -630,8 +644,9 @@ def cache_derived_state(self) -> None: if self.is_situ_activation: # Reinitialize constants after meta-device materialization. These # are backend configuration, not checkpoint weights. - self.situ_alpha.data.fill_(float(self.trtllm_gen_activation_alpha)) - self.situ_beta.data.fill_(float(self.trtllm_gen_activation_beta)) + self.swiglu_alpha.data.fill_(float( + self.trtllm_gen_activation_alpha)) + self.swiglu_beta.data.fill_(float(self.trtllm_gen_activation_beta)) def load_weights(self, weights: List[Dict], @@ -901,10 +916,10 @@ def run_moe( ] else 2 intermediate_size_per_partition_padded = self.w3_w1_weight.shape[ -2] // factor - gemm1_alpha = (self.situ_alpha - if self.is_situ_activation else self.swiglu_alpha) - gemm1_beta = (self.situ_beta - if self.is_situ_activation else self.swiglu_beta) + # Holds SwiGLU's per-expert alpha/beta, or SiTu's backend-local + # activation parameters (which reuse this storage; see + # create_weights). + gemm1_alpha, gemm1_beta = self.swiglu_alpha, self.swiglu_beta output1_scale_scalar = self._get_data_or_none("fc31_scale_c") output1_scale_gate_scalar = self._get_data_or_none("fc31_alpha") diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index cb62ec99f76b..3899c560eb69 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -63,6 +63,12 @@ class ActivationType(IntEnum): Relu2 = 8 +# TRTLLM-Gen-local activation encoding, kept separate from the shared +# ActivationType above ON PURPOSE: ActivationType mirrors the cutlass enum in +# common.h and drives cutlass MoE kernels, whereas SiTu exists only in the +# trtllm-gen batched-GEMM kernels. Adding SiTu to the shared ActivationType +# would force a matching cutlass enum member that no cutlass kernel implements. +# So SiTu stays here (TRTLLM-15177 item 1.2(a): decided keep-backend-local). # Keep this in sync with the ActType enum in # cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h class ActType_TrtllmGen(IntEnum): diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 82bbc5e7f0db..bd7595f02b6f 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -345,6 +345,7 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_c perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinking-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb416_mtp3_ccb-NIXL] SKIP (https://nvbugs/6490049) test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] SKIP (https://nvbugs/6373561) test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_w4afp8_8gpus[DeepSeek-R1-W4AFP8-DeepSeek-R1/DeepSeek-R1-W4AFP8] SKIP (https://nvbugs/5836830) +test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) test_e2e.py::test_trtllm_bench_llmapi_launch[pytorch_backend-llama-v3-llama3-8b] SKIP (https://nvbugs/6568058) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py::test_model[TRTLLM-llama-3.1-model/Llama-3.1-8B-Instruct-pytorch] SKIP (https://nvbugs/6602094) @@ -391,6 +392,7 @@ unittest/_torch/visual_gen/test_cosmos3_distilled.py::TestDistilledDenoiseLoop:: unittest/auto_deploy/multigpu/custom_ops SKIP (https://nvbugs/6403920) unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py::test_glm4_moe_decoder_layer_equivalence[1-dtype0-1-8] SKIP (https://nvbugs/6566730) unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling.py::test_glm4_moe_moe_equivalence[dtype0-1-8] SKIP (https://nvbugs/6566730) +unittest/auto_deploy/singlegpu/shim SKIP (bug pending, tracked in PR 17414) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_wait_in_progress_on_zero_timeout SKIP (https://nvbugs/6517836) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_write_transfer_gpu_tensor SKIP (https://nvbugs/6517836) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_write_transfer_multiple_chunks SKIP (https://nvbugs/6517836)