tests/unittest/_torch/modules/test_fused_shared_expert.py (2)

144-202: ⚡ Quick win

Cover the non-contiguous input copy path too.

fused_sigmoid_gate_mul_add() explicitly normalizes final_hidden_states and shared_expert_output with .contiguous() when their last dimension is strided, but every test here still passes contiguous inputs. A small case with strided inputs plus an explicit output= buffer would lock down that branch.

Suggested test shape

+def test_fused_sigmoid_gate_mul_add_non_contig_inputs():
+    torch.manual_seed(0)
+    num_tokens, hidden = 32, 4096
+    dtype = torch.bfloat16
+    device = "cuda"
+
+    final_base = torch.randn(num_tokens, hidden * 2, dtype=dtype, device=device)
+    shared_base = torch.randn(num_tokens, hidden * 2, dtype=dtype, device=device)
+    final = final_base[:, ::2]
+    shared = shared_base[:, ::2]
+    gate = torch.randn(num_tokens, 1, dtype=dtype, device=device)
+    output = torch.empty(num_tokens, hidden, dtype=dtype, device=device)
+
+    assert final.stride(-1) != 1
+    assert shared.stride(-1) != 1
+
+    expected = _reference(final, gate, shared)
+    out = fused_sigmoid_gate_mul_add(final, gate, shared, output=output)
+    torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/test_fused_shared_expert.py` around lines 144 -
202, Add a new unit test that exercises the non-contiguous input copy path for
fused_sigmoid_gate_mul_add by passing non-contiguous final_hidden_states and
shared_expert_output (e.g., create them via slicing or empty_strided so their
last dimension has stride != 1) along with an explicit non-contiguous output
buffer, then call fused_sigmoid_gate_mul_add(final, gate, shared, output=output)
and assert the returned tensor uses the provided output buffer (out.data_ptr()
== output.data_ptr()) and matches the reference (_reference(final, gate,
shared)) within the existing tolerances; this will trigger the code paths that
call .contiguous() on final_hidden_states and shared_expert_output and lock down
that branch.

---

42-218: Please make sure this lands in perf regression coverage too.

These unit tests are good for correctness, but the PR’s value is throughput/TPOT on Qwen3.5 serving. I don’t see a matching perf-list update in this diff, so please either add or point to the existing entry that exercises this fused path in tests/integration/test_lists/test-db/l0_perf.yml and the relevant tests/integration/test_lists/qa/llm_perf_*.yml.

As per coding guidelines, performance-sensitive paths should have pre-merge and QA perf coverage, and functional-only tests are not enough when a kernel rewrite can regress latency/throughput.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/test_fused_shared_expert.py` around lines 42 -
218, Add perf-regression coverage for the new fused kernel: update
tests/integration/test_lists/test-db/l0_perf.yml (or add a new entry) to include
a workload that exercises fused_sigmoid_gate_mul_add (the same shape used in
test_fused_sigmoid_gate_mul_add_qwen35_shape: hidden=4096, B=64) and add or
reference the corresponding QA perf list in tests/integration/test_lists/qa/
(e.g. llm_perf_*.yml) so the fused path is included in both pre-merge and QA
runs; ensure the perf entry names and payload mirror the unit test’s shape and
dtype (bfloat16) and reference the test module
tests/unittest/_torch/modules/test_fused_shared_expert.py or the
fused_sigmoid_gate_mul_add kernel so CI will run throughput/TPOT measurements.

tensorrt_llm/_torch/models/modeling_qwen3_next.py (1)

241-253: 🏗️ Heavy lift

Add a distributed regression test for this fused TP path.

The new unit suite validates the standalone kernel, but this branch is where the integration risk lives: allocate_output(...) produces a strided buffer, the fused kernel writes into it, and then the result is immediately all-reduced. A regression here would still pass all current tests.

As per coding guidelines, assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/modeling_qwen3_next.py` around lines 241 - 253,
Add a distributed regression test that exercises the fused TP path where
allocate_output returns a strided buffer and the fused kernel writes into it
before an immediate all-reduce: create a test that instantiates the model with
mapping.tp_size > 1 and enable_attention_dp = False, invoke the branch that
calls torch.ops.trtllm.allocate_output, fused_sigmoid_gate_mul_add, and
self.allreduce (use shared_expert_outputs inputs so shared_expert_output and
shared_expert_gate_logits are non-trivial), and assert the final_hidden_states
after the allreduce matches a reference (e.g., CPU or non-fused fallback) across
ranks; include variants for different tp_size values and an edge case with
non-contiguous/strided output buffers to catch stride-handling bugs, and run the
test under torch.distributed (or the repo’s distributed test harness) so it
executes the full allocate->fused-kernel->allreduce integration.