From ef9145db268c148761ee1befd78db1a8a0b9ee5e Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:29:53 -0500 Subject: [PATCH 1/3] [TRTLLM-15177][chore] consolidate trtllm-gen SiTu activation slot handling Item 1.2 of the PR #17269 deferred-cleanup bucket. (a) Record the decision to keep SiTu in the trtllm-gen-local ActType_TrtllmGen enum rather than adding it to the shared ActivationType, which mirrors the cutlass enum in common.h and would then need a cutlass member no cutlass kernel implements. Documented at the enum and left the enums unchanged. (b) Route SiTu and SwiGLU per-expert alpha/beta through a single _gemm1_activation_params accessor at the op call instead of duplicating the is_situ_activation ternary. Behavior-identical: same tensors, same op slot. The deeper storage merge (eliminating the separate situ_alpha/situ_beta buffers by reusing swiglu_alpha/swiglu_beta) is left as a TODO because 'swiglu_alpha is not None' gates quant-method selection and validation on the numeric path; it needs GPU parity revalidation before it can land. Signed-off-by: Brian Nguyen --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 27 ++++++++++++++++--- tensorrt_llm/_torch/utils.py | 6 +++++ 2 files changed, 29 insertions(+), 4 deletions(-) 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..b7cc790b36e2 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 @@ -338,6 +338,21 @@ def _to_trtllm_gen_activation_type(self, def is_situ_activation(self) -> bool: return self.trtllm_gen_activation_type == ActType_TrtllmGen.SiTu + @property + def _gemm1_activation_params( + self) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Per-expert (alpha, beta) for the shared gemm1 op slot. + + SiTu and SwiGLU are mutually exclusive (see + ``_validate_backend_local_activation``) and feed the same + ``gemm1_alpha``/``gemm1_beta`` op arguments, so route both through one + accessor. SiTu's values live in the ``situ_alpha``/``situ_beta`` + buffers; SwiGLU's in ``swiglu_alpha``/``swiglu_beta``. + """ + if self.is_situ_activation: + return self.situ_alpha, self.situ_beta + return self.swiglu_alpha, self.swiglu_beta + def _validate_backend_local_activation(self) -> None: if self.trtllm_gen_activation_type is None: if (self.trtllm_gen_activation_alpha is not None @@ -594,6 +609,13 @@ def create_weights(self): else: self.quant_method.create_weights(self) + # TODO(TRTLLM-15177 item 1.2(b)): fully merge these SiTu buffers into + # the swiglu_alpha/swiglu_beta storage so the gemm1 slot has a single + # backing parameter and _gemm1_activation_params can drop the branch. + # Deferred: `swiglu_alpha is not None` currently gates quant-method + # selection and validation (create_moe.py, this file's _get_quant_method + # / _check_configs), so reusing that storage for SiTu changes control + # flow on the numeric path and needs GPU parity revalidation first. if self.is_situ_activation: situ_alpha = nn.Parameter(torch.full( (self.expert_size_per_partition, ), @@ -901,10 +923,7 @@ 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) + gemm1_alpha, gemm1_beta = self._gemm1_activation_params 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): From c45b584925601f114359d6b9b064e95c47844a70 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 02:02:16 -0700 Subject: [PATCH 2/3] [TRTLLM-15177][chore] merge SiTu activation params into swiglu_alpha/swiglu_beta storage Complete item 1.2(b): drop the separate situ_alpha/situ_beta parameters and reuse the swiglu_alpha/swiglu_beta storage, so the gemm1 alpha/beta op slot has a single backing parameter and the _gemm1_activation_params accessor branch disappears. Safe with respect to the 'swiglu_alpha is not None' gates: - create_moe.py validates the constructor kwargs, which remain None for SiTu; the storage is populated later in create_weights. - _get_quant_method consults swiglu_alpha only on the nvfp4 branch; SiTu requires W4A8_MXFP4_MXFP8, and selection also runs before the storage is populated. - _check_configs's swiglu gate admits w4a8_mxfp4_mxfp8, and its SiTu branch now validates the merged storage. - The NVFP4 post-load swiglu_beta rescale (NVFP4TRTLLMGenFusedMoEMethod.process_weights_after_loading) is not in SiTu's MXFP4 quant-method ancestry. - _validate_backend_local_activation runs from __init__, before create_weights, so it still rejects constructor-provided SwiGLU parameters for SiTu. Signed-off-by: Brian Nguyen --- .../modules/fused_moe/fused_moe_trtllm_gen.py | 60 +++++++++---------- 1 file changed, 28 insertions(+), 32 deletions(-) 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 b7cc790b36e2..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 @@ -338,22 +338,12 @@ def _to_trtllm_gen_activation_type(self, def is_situ_activation(self) -> bool: return self.trtllm_gen_activation_type == ActType_TrtllmGen.SiTu - @property - def _gemm1_activation_params( - self) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: - """Per-expert (alpha, beta) for the shared gemm1 op slot. - - SiTu and SwiGLU are mutually exclusive (see - ``_validate_backend_local_activation``) and feed the same - ``gemm1_alpha``/``gemm1_beta`` op arguments, so route both through one - accessor. SiTu's values live in the ``situ_alpha``/``situ_beta`` - buffers; SwiGLU's in ``swiglu_alpha``/``swiglu_beta``. - """ - if self.is_situ_activation: - return self.situ_alpha, self.situ_beta - return self.swiglu_alpha, self.swiglu_beta - 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): @@ -564,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, ) @@ -609,26 +601,26 @@ def create_weights(self): else: self.quant_method.create_weights(self) - # TODO(TRTLLM-15177 item 1.2(b)): fully merge these SiTu buffers into - # the swiglu_alpha/swiglu_beta storage so the gemm1 slot has a single - # backing parameter and _gemm1_activation_params can drop the branch. - # Deferred: `swiglu_alpha is not None` currently gates quant-method - # selection and validation (create_moe.py, this file's _get_quant_method - # / _check_configs), so reusing that storage for SiTu changes control - # flow on the numeric path and needs GPU parity revalidation first. + # 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() @@ -652,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], @@ -923,7 +916,10 @@ def run_moe( ] else 2 intermediate_size_per_partition_padded = self.w3_w1_weight.shape[ -2] // factor - gemm1_alpha, gemm1_beta = self._gemm1_activation_params + # 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") From 3b66f7259576273f7b470c128bd7b82117db887c Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Thu, 13 Aug 2026 13:40:16 -0500 Subject: [PATCH 3/3] test waives: skip VANILLA BERT ptp_quickstart e2e (pre-existing timeout on RTXPro6000D-PyTorch-1, tracking bug pending) Signed-off-by: Brian Nguyen --- tests/integration/test_lists/waives.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d3bfb4973452..fd2a30ada545 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -348,6 +348,7 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_kimi-k25-thinkin 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 (https://nvbugs/6605819) +test_e2e.py::test_ptp_quickstart_bert[VANILLA-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)