Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 3035-3042: The cache key tuple used when caching DenseGEMM kernels
(the variable named cache_key built from self.weight_per_expert, mma_tiler_mn,
cluster_shape_mn, self.scaling_vector_size, self.expert_count, and the
alpha_post flag) must include the output element type; add output_dtype (or the
name used in the method that represents the requested output element type) into
that tuple so BF16-compiled kernels cannot be reused for FP16/FP32/FP4 calls.
Apply the same change to the other occurrence that builds a DenseGEMM cache_key
(the second cache_key construction in the class) so both cache keys include
output_dtype.
- Around line 2895-2900: The mapping that sets c_cutlass_dtype currently uses
.get(self.output_dtype, cutlass.BFloat16) and silently falls back to BF16;
change this to explicitly validate self.output_dtype and raise a clear exception
(e.g., ValueError) when the dtype is unsupported instead of defaulting to
cutlass.BFloat16. Locate the assignments to c_cutlass_dtype in the DenseGEMM
runners (the tactic-probing and forward paths for the FC1 and FC2 runners) where
the mapping dict is used and replace the .get default with an explicit
lookup+error path that names self.output_dtype and the allowed dtypes. Apply the
same fix to the other occurrences referenced in the review (the FC1/FC2
tactic-probing and forward blocks) so both probing and runtime use fail-fast
behavior.

In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py`:
- Around line 2288-2302: The swizzled SF layouts use floor division (m // 128, n
// 128) which drops partial 128-wide tiles and can produce zero-sized tensors;
replace those with ceil-div tile counts (e.g. compute m_tiles = (m + 127)//128
and n_tiles = (n + 127)//128) and use m_tiles/n_tiles in the calls to
cute.make_ordered_layout for a_sf and b_sf (or add explicit 128-tile guards in
the validation path), so the layouts always account for partial tiles and never
produce zero-length dimensions when m<128 or n not divisible by 128.
- Around line 1-27: The file
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/moe_as_dense_gemm/fc2.py
currently has a BSD-3-Clause header; replace that entire BSD header block with
the repository-standard NVIDIA Apache-2.0 license header using the latest
modification year (2026) and include any necessary upstream attribution
separately (do not mix BSD text into the file header). Ensure the new header
matches the project template used across *.py sources (Apache-2.0 text and
NVIDIA copyright line), and leave the rest of fc2.py unchanged.

In `@tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py`:
- Around line 1095-1100: The dispatch logic in configurable_moe.py fails to pass
backend-specific kwargs to DenseGEMMFusedMoE because the earlier tuple check
(CutlassFusedMoE, DeepGemmFusedMoE, CuteDslFusedMoE, DenseGEMMFusedMoE) is a
no-op and later branches use exact class equality (self.backend.__class__ ==) so
DenseGEMMFusedMoE is skipped; fix by either adding an explicit branch for
DenseGEMMFusedMoE that passes enable_alltoall and moe_output (e.g., elif
self.backend.__class__ == DenseGEMMFusedMoE: ...) or, preferably, change the
equality checks to isinstance(self.backend, CutlassFusedMoE) /
isinstance(self.backend, DenseGEMMFusedMoE) (or a common base) so inherited
classes receive the correct kwargs in the dispatch that sets enable_alltoall and
moe_output for the backend.

In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py`:
- Around line 67-75: The fallback to compute expert_size as
token_final_scales.shape[1] when alpha is None is wrong because that is top_k
not total experts and causes scatter_ indexing OOB; change the logic in the
block that builds fc2_alpha (around token_selected_experts, alpha,
token_final_scales, fc2_alpha, scatter_) so expert_size is the actual expert
count (either use an explicit expert_count argument or compute expert_count =
int(token_selected_experts.max().item()) + 1) when alpha is None, allocate
fc2_alpha with that expert_size, and then perform the scatter_. Ensure
token_selected_experts.long() still indexes into the resized fc2_alpha.

In
`@tests/scripts/cute_dsl_kernels/moe_as_dense_gemm/run_moe_as_dense_gemm_fc1.py`:
- Around line 180-182: The code computes weight_per_expert = n // expert_count
but does not validate divisibility, causing silent truncation (e.g.,
mnkl=(512,256,256,1) with expert_count=257 yields 0); update the logic in this
module to first assert or raise a clear error if n % expert_count != 0 and
choose consistent default values so defaults are divisible by expert_count
(adjust mnkl and/or expert_count), and apply the same validation to the other
occurrence of the calculation (the block around the second occurrence computing
weight_per_expert at lines referenced in the review); reference the variables
weight_per_expert, n, expert_count and the default mnkl tuple when making these
changes.
- Around line 1-27: The file run_moe_as_dense_gemm_fc1.py contains an incorrect
BSD-3-Clause header and a 2025 copyright year; replace the existing top-of-file
license block with the repository's standard NVIDIA Apache-2.0 Python header and
update the copyright year to 2026 (the latest meaningful modification), ensuring
the new header sits at the very top of run_moe_as_dense_gemm_fc1.py before any
imports or code so it matches the repo policy for .py sources.
- Around line 59-65: The fallback import in the try/except block around
importing tensorrt_llm._torch.cute_dsl_kernels.blackwell.moe_as_dense_gemm (the
code that inserts into sys.path before importing blackwell.moe_as_dense_gemm as
kernel_module) uses Path(__file__).parents[3], which points to tests/, so change
the inserted path computation to use Path(__file__).parents[4] so the
sys.path.insert(0, str(...)) points to the repository
root/tensorrt_llm/_torch/cute_dsl_kernels and the fallback import succeeds.
- Around line 947-952: The parser arguments for the flags "--vectorized_f32" and
"--use_cupti" use action="store_true" together with default=True so they always
parse to True; update their parser.add_argument calls so the flags can be
disabled at the CLI: either set default=False when using action="store_true" (so
passing the flag enables the feature) or add complementary flags using
action="store_false" (e.g., "--no-vectorized_f32" / "--no-use_cupti") with
default=True to allow turning them off; modify the parser.add_argument calls for
the entries that reference "--vectorized_f32" and "--use_cupti" accordingly.

In
`@tests/scripts/cute_dsl_kernels/moe_as_dense_gemm/run_moe_as_dense_gemm_fc2.py`:
- Around line 1-27: Replace the current BSD-3-Clause header in
run_moe_as_dense_gemm_fc2.py with the repository-standard NVIDIA Apache-2.0
header (using the latest modification year), or if the BSD notice must be
preserved, prepend the official Apache-2.0 header above the existing BSD block;
ensure the file begins with the full Apache-2.0 license text and the NVIDIA
copyright line per project guidelines.
- Around line 62-68: The fallback sys.path insertion uses
Path(__file__).parents[3] which points to tests/ (one directory too shallow)
causing the fallback import of blackwell.moe_as_dense_gemm.fc2 to fail; update
the fallback to insert the correct parent directory by changing
Path(__file__).parents[3] to Path(__file__).parents[4] (so the inserted path is
.../tensorrt_llm/_torch/cute_dsl_kernels) before importing kernel_module,
ensuring the direct python run_moe_as_dense_gemm_fc2.py flow can locate the
module.

In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 471-473: The test mutates process-wide env var
TRTLLM_MOE_FUSED_FC2_ALPHA when backend_type == MoeBackendType.DENSEGEMM and
never restores it; change the code to isolate this override by using pytest's
monkeypatch (e.g., monkeypatch.setenv("TRTLLM_MOE_FUSED_FC2_ALPHA","0") within
the test or fixture that sets backend_type) or save the original os.environ
value and restore it in a finally/teardown block so the env is returned to its
prior state after the parametrized case; update the code around the backend_type
check (the branch referencing MoeBackendType.DENSEGEMM and
TRTLLM_MOE_FUSED_FC2_ALPHA) to use monkeypatch or explicit restore.

In `@tests/unittest/_torch/modules/moe/test_moe_module.py`:
- Around line 1059-1061: The test currently sets
os.environ["TRTLLM_MOE_FUSED_FC2_ALPHA"]="0" unconditionally which leaks into
the pytest process; replace that with the pytest monkeypatch fixture to localize
the change: when checking moe_backend == MoeBackendType.DENSEGEMM.value, call
monkeypatch.setenv("TRTLLM_MOE_FUSED_FC2_ALPHA", "0") instead of writing to
os.environ so the environment is restored after the test (reference
TRTLLM_MOE_FUSED_FC2_ALPHA, moe_backend, and MoeBackendType.DENSEGEMM.value).

In `@tests/unittest/_torch/thop/parallel/test_moe_densegemm.py`:
- Around line 1-2: The file test_moe_densegemm.py currently contains only an
SPDX short-form header; replace it with the repository’s full NVIDIA Apache-2.0
header block (include the full multi-line copyright/Apache 2.0 license text and
the latest modification year, e.g., 2026) so the top of test_moe_densegemm.py
matches the standard header used across Python sources in the repo.

---

Outside diff comments:
In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 148-153: The assertion that guards MoE load-balancer support
excludes the new DenseGEMMFusedMoE class, causing calls that forward
init_load_balancer (e.g., via create_moe_backend when moe_load_balancer is set)
to raise; update the allowlist in the assertion that references WideEPMoE,
CutlassFusedMoE, TRTLLMGenFusedMoE, CuteDslFusedMoE, DeepGemmFusedMoE to also
include DenseGEMMFusedMoE, and make the same change for the second identical
assertion later in the file so DenseGEMMFusedMoE is permitted wherever load
balancing is checked.