[TRTLLM-16460][perf] Qwen3.8-Flash-Next NVFP4 modeling_v2 target: GB300 Pareto optimizations for ISL 8K / OSL 1K

tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py

@@ -1591,6 +1591,7 @@ def _prepare_next_new_tokens(
         batch_indices_cuda,
         batch_size,
         num_accepted_tokens,
+        fused: bool = False,
     ):
         """The rows the last draft-step kernel wrote when these draft tokens are that loop's."""
         drafted, self._drafted_rows = self._drafted_rows, None
@@ -1602,6 +1603,7 @@ def _prepare_next_new_tokens(
             batch_indices_cuda,
             batch_size,
             num_accepted_tokens,
+            fused=fused,
         )
 
     def _sample_and_accept_draft_tokens_base(

tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/activation/fused_sigmoid_mul_inplace.py

@@ -4,7 +4,11 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
+import triton
+import triton.language as tl
 
 # Registered by the module that defines it: after `import
 # tensorrt_llm._torch.custom_ops` alone, torch.ops.trtllm has no
@@ -16,6 +20,70 @@
 _CKPT = dict(heads=24, head_dim=256, dtype=torch.bfloat16, gate="interleaved")
 
 
+@triton.jit
+def _sigmoid_mul_amax_kernel(
+    out_ptr,
+    gate_ptr,
+    parts_ptr,
+    out_stride_token,
+    gate_stride_token,
+    gate_stride_head,
+    num_parts,
+    hidden_size: tl.constexpr,
+    head_dim: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    """The registered `_fused_sigmoid_mul_kernel`'s arithmetic in place, plus the maximum of
+    the stored `|out|` folded into `parts[program % num_parts]`."""
+    token = tl.program_id(0).to(tl.int64)
+    block = tl.program_id(1)
+    offsets = block * block_size + tl.arange(0, block_size)
+    mask = offsets < hidden_size
+    head = offsets // head_dim
+    dim = offsets - head * head_dim
+    attention = tl.load(out_ptr + token * out_stride_token + offsets, mask=mask, other=0.0).to(
+        tl.float32
+    )
+    gate = tl.load(
+        gate_ptr + token * gate_stride_token + head * gate_stride_head + dim, mask=mask, other=0.0
+    ).to(tl.float32)
+    y = (attention * tl.sigmoid(gate)).to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + token * out_stride_token + offsets, y, mask=mask)
+    peak = tl.max(tl.where(mask, tl.abs(y.to(tl.float32)), 0.0), axis=0)
+    program = token * tl.num_programs(1) + block
+    tl.atomic_max(parts_ptr + program % num_parts, peak)
+
+
+def _sigmoid_mul_amax(attention_output, gate, amax_parts) -> None:
+    # The registered op's layout checks, as loud.
+    assert attention_output.dim() == 2 and attention_output.stride(-1) == 1
+    num_tokens, hidden_size = attention_output.shape
+    assert gate.stride(-1) == 1
+    if gate.dim() == 3:
+        assert gate.shape[0] == num_tokens and gate.shape[1] * gate.shape[2] == hidden_size
+        head_dim, stride_token, stride_head = gate.shape[2], gate.stride(0), gate.stride(1)
+    else:
+        assert gate.dim() == 2 and gate.shape == attention_output.shape
+        head_dim, stride_token, stride_head = hidden_size, gate.stride(0), hidden_size
+    assert amax_parts.dtype == torch.float32 and amax_parts.dim() == 1
+    if num_tokens == 0:
+        return
+    block_size = min(triton.next_power_of_2(hidden_size), 1024 if num_tokens < 1024 else 2048)
+    _sigmoid_mul_amax_kernel[(num_tokens, triton.cdiv(hidden_size, block_size))](
+        attention_output,
+        gate,
+        amax_parts,
+        attention_output.stride(0),
+        stride_token,
+        stride_head,
+        amax_parts.shape[0],
+        hidden_size=hidden_size,
+        head_dim=head_dim,
+        block_size=block_size,
+        num_warps=4,
+    )
+
+
 class FusedSigmoidMulInplace(OpWrapper):
     """`attention_output[t, i] *= sigmoid(gate[t, i // D, i % D])` in place, `D = gate.shape[-1]`.
 
@@ -26,8 +94,14 @@ class FusedSigmoidMulInplace(OpWrapper):
     no repack (bitwise equal to its compacted copy, measured). A 2-D `gate` of
     `attention_output`'s shape is one head spanning the row.
 
-    Fusion boundary: sigmoid and multiply only. The caller owns producing both
-    operands and everything downstream (the o projection). `gate` is read-only.
+    With `amax_parts` (a zeroed fp32 buffer, `fp8_per_tensor_mm.amax_partials`) the same
+    arithmetic runs in this target's copy of the kernel, which also folds the maximum of the
+    stored `|attention_output|` into the buffer (atomic maxima, one slot per program modulo
+    its length) for the per-tensor FP8 o projection.
+
+    Fusion boundary: sigmoid and multiply (and the output maximum) only. The caller owns
+    producing both operands and everything downstream (the o projection). `gate` is
+    read-only.
     """
 
     ARCHS = frozenset({Arch.SM_103})
@@ -45,6 +119,17 @@ class FusedSigmoidMulInplace(OpWrapper):
             why="MTP decode: 1 token, 24x256 gate as the strided half of q_proj",
             spec=dict(tokens=1, **_CKPT),
         ),
+        Cell(
+            why="prefill 16384 tokens folding the output maximum for the FP8 o projection",
+            spec=dict(
+                tokens=16384,
+                heads=24,
+                head_dim=256,
+                dtype=torch.bfloat16,
+                gate="fused_row",
+                amax=True,
+            ),
+        ),
     )
 
     note = """
@@ -82,11 +167,24 @@ class FusedSigmoidMulInplace(OpWrapper):
     fp16/fp32 operands, other hidden sizes, CUDA-graph capture, torch.compile.
     """
 
-    def __call__(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
-        torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate)
+    def __call__(
+        self,
+        attention_output: torch.Tensor,
+        gate: torch.Tensor,
+        amax_parts: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if amax_parts is None:
+            torch.ops.trtllm.fused_sigmoid_mul_inplace(attention_output, gate)
+        else:
+            _sigmoid_mul_amax(attention_output, gate, amax_parts)
         return attention_output
 
-    def reference(self, attention_output: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    def reference(
+        self,
+        attention_output: torch.Tensor,
+        gate: torch.Tensor,
+        amax_parts: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """What `attention_output` holds after the call; neither input is mutated.
 
         `gate.reshape(T, hidden)` is the per-head `(t, i // D, i % D)` read, for