fix: correct TRT-LLM MLA labels and GPT-OSS MXFP4 profiles

README.md

@@ -402,6 +402,20 @@ complete test coverage, and release gates. Use [DEVELOPMENT.md](DEVELOPMENT.md)
 for environment and local test details and [CONTRIBUTING.md](CONTRIBUTING.md)
 before sending a change.
 
+Direct Python `ReplaySpec.workload` synthetic workloads accept
+`length_sampler: numpy_random_state` for InferenceX-compatible seeded token
+lengths in Gym replay. The public prediction/recommendation YAML and CLI, and
+the Sweeper `Workload` schema, do not expose this option and reject it.
+The supported direct path must omit `source_type`; all workload-driver inputs
+with `source_type` reject `length_sampler` at the common runner entrypoint,
+including AFD and AFD+PD. Materialized trace replay rejects
+non-default samplers. The default `python_random`
+preserves existing workloads. Both sample the full input vector before output
+lengths; unknown sampler names are rejected. NumPy seeds must fit an unsigned
+32-bit integer; the Python sampler retains unsigned 64-bit seed support.
+Each replay initializes its own seeded sampler, so independent runs reproduce
+the same request lengths.
+
 ### FPM accuracy overview
 
 The [FPM Accuracy Overview](https://ai-dynamo.org/aisimulate/fpm-accuracy/?branch=main)

THIRD_PARTY_NOTICES.md

@@ -36,17 +36,20 @@ at the identified revision is available at:
 https://github.com/ai-dynamo/aiconfigurator/blob/13b5cf2697876692b0a52098266c81162add11fc/LICENSE
 
 The 18 B200 TensorRT-LLM 1.3.0rc20 performance tables under
-`src/aiconfigurator_core/systems/data/b200_sxm/*/trtllm/1.3.0rc20/` include
+`src/aisimulate_core/systems/data/b200_sxm/*/trtllm/1.3.0rc20/` include
 power measurements derived from AIConfigurator commit
-`915f590680d8a79fe9c39f6f3a9ff13bc267fcce` (PR #1584). Sixteen tables are
-unmodified, byte-identical copies. The context-attention and
+`915f590680d8a79fe9c39f6f3a9ff13bc267fcce` (PR #1584). Fourteen tables remain
+unmodified, byte-identical copies after the September 17-18 refresh. The
+context-MLA table was replaced with locally collected measurements, and the
+MoE table has locally refreshed MXFP4 rows with unavailable-power sentinels.
+The context-attention and
 generation-attention tables are modified derivatives: AISimulate preserves
 newer local timing rows and adds the typed `0.0` / `0.0` unavailable sentinel
 to those local-only identities. Source paths, row counts, measured coverage,
 and merge details are recorded in
-`src/aiconfigurator_core/systems/data/b200_sxm/README.md`. The two unmodified
+`src/aisimulate_core/systems/data/b200_sxm/README.md`. The two unmodified
 upstream attention copies under
-`src/aiconfigurator_core/systems/data/b200_sxm/power_upstream/` support focused
+`src/aisimulate_core/systems/data/b200_sxm/power_upstream/` support focused
 import regression tests, which pin source and packaged SHA-256 digests.
 
 The corresponding energy expectations in the repository-root file

docs/trtllm-eager-execution-plan.md

@@ -0,0 +1,89 @@
+# TRT-LLM eager execution: proposed measurement contract
+
+Status: deferred at the user’s request. Retain collection evidence; do not
+extend the data contract or implement a timing correction. The proposal below
+is preserved for a future decision.
+The current GPT-OSS op-based prediction still omits a material host execution
+boundary. See the [trace report](../python/aisimulate/collector/trtllm/gym-gpt-eager-20260918.md)
+for measured evidence and its limits.
+
+## Problem
+
+GPT-OSS-120B, B200 TP2/EP1, 1k/1k, concurrency 4 reproduces client TTFT at
+48.86 ms versus a 16.53 ms prediction. The corrected MXFP4/MXFP8 MoE kernel
+cost is close to the trace. Eager prefill/mixed execution has substantial CPU
+submission gaps and cross-rank waiting; graph decode does not have the same
+boundary. Increasing the shared MoE kernel table would charge that delay to
+graph decode and erase its existing TPOT alignment.
+
+## Proposed coordinated change
+
+Keep the op-based GPU model and add explicit execution-boundary measurements.
+Do not insert an E2E-derived TTFT constant or a model-wide latency multiplier.
+
+1. Extend the collector output contract with an execution-mode identity
+   (`cuda_graph` or `eager`) and CPU dispatch timing. Preserve GPU latency as
+   a separate measurement. Record actual capture success, runtime source,
+   CPU model/affinity, rank membership, warmup policy and raw timing samples.
+   The measured runtime version remains part of the identity: rc14 host timing
+   must not be published as an rc20 measurement.
+   Existing rows without this evidence retain their legacy behavior; they
+   must not be relabeled as measured eager execution.
+2. Collect both modes at identical serving-selected modules, shapes, dtypes
+   and topology. Use phase-correct attention and fused all-reduce/norm
+   boundaries. Preserve the overlap within a module and between CPU and GPU;
+   a profiler's inflated CPU times are diagnostic evidence, not database rows.
+3. Carry the resolved TRT-LLM graph policy from Gym/runner into the native
+   performance model. Mixed/prefill selects eager measurements where the
+   pinned runtime cannot capture those steps. Decode selects graph data only
+   for captured shapes. Graph memory reservation is a separate setting.
+4. Compose host submission and GPU work with explicit dependencies. Do not
+   add CPU time to the sum of all GPU kernels, or separately add collective
+   waiting already caused by a late peer. Validate the composition against
+   complete, independent per-rank traces before enabling the new path.
+5. Keep missing eager coverage visible. Do not silently report a legacy
+   graph lookup as an eager measurement. Preserve the existing path for
+   configurations outside the qualified collection domain.
+
+The exact row schema and composition must be frozen after the paired
+measurement smoke, before formal collection. Approval would cover the
+new execution-mode/CPU-timing contract and its producer/consumer work; it
+is not an accuracy claim for an unmeasured composition algorithm.
+
+## Owners and affected surfaces
+
+- Collector: `python/aisimulate/collector/trtllm/`, the network collector for
+  fused communication, perf schema/metadata, and collector tests.
+- Data loading and native timing: `crates/core/src/perfmodel/perf_database/`,
+  `operators/`, and `engine/runtime.rs`; Python lowering and diagnostics.
+- Deployment identity: Gym normalization, AISimulate runner/provider config,
+  and native engine config. Preserve actual graph capture coverage.
+- Documentation and data: measured rows, immutable provenance, paired results,
+  missing coverage, and unchanged controls.
+
+## Acceptance checks
+
+- The existing TP2/C4 workload retains its exact 40-request token vectors.
+  Repeated unprofiled serving measurements remain independent validation.
+- Prefill, mixed, and graph decode are evaluated separately; complete rank
+  groups and timing boundaries are checked. CPU affinity is retained.
+- Cover neighboring token counts and the affected GPT-OSS topologies, not
+  only the diagnostic 886-token step. Select collection points before looking
+  at validation errors.
+- Replay all 31 GPT-OSS/B200 points. If shared timing or config code changes,
+  rerun all 263 points and retain every TTFT/TPOT regression and failure.
+- Reject double counting with synthetic dependency tests; test graph/eager
+  selection at capture boundaries and explicit missing-data behavior.
+- Report remaining queue/frontend error separately from forward error.
+  Do not claim the entire 32 ms client gap is a kernel or host-only delta.
+
+## Repository approval boundary
+
+`python/aisimulate/.claude/rules/collector/layer_permissions.md` requires:
+
+> Changing the data contract requires explicit human approval.
+
+Adding execution-mode/CPU-timing fields changes the producer/consumer contract.
+The current collection/fix request authorizes diagnostic collection; this
+proposal makes the newly discovered contract extension reviewable before its
+implementation. No collector policy file is changed.

python/aisimulate/THIRD_PARTY_NOTICES.md

@@ -36,17 +36,20 @@ at the identified revision is available at:
 https://github.com/ai-dynamo/aiconfigurator/blob/13b5cf2697876692b0a52098266c81162add11fc/LICENSE
 
 The 18 B200 TensorRT-LLM 1.3.0rc20 performance tables under
-`src/aiconfigurator_core/systems/data/b200_sxm/*/trtllm/1.3.0rc20/` include
+`src/aisimulate_core/systems/data/b200_sxm/*/trtllm/1.3.0rc20/` include
 power measurements derived from AIConfigurator commit
-`915f590680d8a79fe9c39f6f3a9ff13bc267fcce` (PR #1584). Sixteen tables are
-unmodified, byte-identical copies. The context-attention and
+`915f590680d8a79fe9c39f6f3a9ff13bc267fcce` (PR #1584). Fourteen tables remain
+unmodified, byte-identical copies after the September 17-18 refresh. The
+context-MLA table was replaced with locally collected measurements, and the
+MoE table has locally refreshed MXFP4 rows with unavailable-power sentinels.
+The context-attention and
 generation-attention tables are modified derivatives: AISimulate preserves
 newer local timing rows and adds the typed `0.0` / `0.0` unavailable sentinel
 to those local-only identities. Source paths, row counts, measured coverage,
 and merge details are recorded in
-`src/aiconfigurator_core/systems/data/b200_sxm/README.md`. The two unmodified
+`src/aisimulate_core/systems/data/b200_sxm/README.md`. The two unmodified
 upstream attention copies under
-`src/aiconfigurator_core/systems/data/b200_sxm/power_upstream/` support focused
+`src/aisimulate_core/systems/data/b200_sxm/power_upstream/` support focused
 import regression tests, which pin source and packaged SHA-256 digests.
 
 The corresponding energy expectations in the repository-root file

python/aisimulate/collector/README.md

@@ -127,6 +127,9 @@ capabilities.py                        — generation-time capability/denylist f
 The plan is one equation: cases = dedup(base grid ∪ model shapes), then
 intersected with hardware capability floors and minus the hang denylist.
 
+TRT-LLM MLA precision and GPT-OSS TP profile collection evidence is recorded
+in [the September 17 campaign](trtllm/gym-profiles-20260917.md).
+
 ## Whole-forward FPM campaign
 
 Use the dedicated `python3 -m collector.fpm_forward` entry point for vLLM
@@ -664,3 +667,12 @@ Explicit limits must be positive for one-shot module collection.
 Omitting the option keeps SGLang's default. The existing
 wide-EP and DSA paths keep their separate behavior. The ordinary context table
 uses the existing consumer schema; this does not add mixed-precision modeling.
+
+## TRT-LLM MoE autotuning failures
+
+TRT-LLM MoE collection requires either a loaded tactic cache or at least one
+successful autotuning attempt before measuring a tuning-enabled case. Ordinary
+failures may try smaller eligible token counts; if every attempt fails, the
+case raises before timing or publishing rows. CUDA accelerator errors propagate
+immediately. The successful collection path and historical profile values are
+unchanged by this failure guard.

python/aisimulate/collector/op_backend_facts.yaml

@@ -175,7 +175,7 @@ ops:
       - {framework: "trtllm", version: "1.0.0", system: "l40s", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.0.0", system: "l40s", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "b200_sxm", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
-      - {framework: "trtllm", version: "1.3.0rc20", system: "b200_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
+      - {framework: "trtllm", version: "1.3.0rc20", system: "b200_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "fp8", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "b300_sxm", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "b300_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "gb200", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
@@ -185,7 +185,7 @@ ops:
       - {framework: "trtllm", version: "1.3.0rc20", system: "h100_sxm", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "h100_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "h200_sxm", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
-      - {framework: "trtllm", version: "1.3.0rc20", system: "h200_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
+      - {framework: "trtllm", version: "1.3.0rc20", system: "h200_sxm", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "fp8", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "mla_context", kv_cache_dtype: "bfloat16", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
       - {framework: "trtllm", version: "1.3.0rc20", system: "rtx_pro_6000_server", op_name: "mla_context", kv_cache_dtype: "fp8", mla_dtype: "bfloat16", backends: ["trtllm_internal"], kernel_sources: ["default"]}
   - op_file: custom_allreduce_perf

python/aisimulate/collector/trtllm/collect_mla.py

@@ -651,7 +651,13 @@ def kernel_func():
     log_perf(
         item_list=[
             {
-                "mla_dtype": "bfloat16",
+                # TRT-LLM quantizes context Q/K/V internally for FP8 KV on
+                # SM90/100/103/120. The BF16 Python inputs do not describe
+                # the FMHA compute precision (attentionOp.cpp:1228-1231 at
+                # c25c23f71786bad54d192893d696ce8043426eca / v1.3.0rc20).
+                "mla_dtype": (
+                    dtype_str if is_context_phase and get_sm_version() in (90, 100, 103, 120) else "bfloat16"
+                ),
                 "kv_cache_dtype": dtype_str,
                 "num_heads": num_heads,
                 "batch_size": len(context_sequence_lengths),

python/aisimulate/collector/trtllm/collect_moe.py

@@ -357,6 +357,7 @@ def run_moe_torch(
     model_config.mapping = mapping
     model_config.quant_config = quant_config
     model_config.moe_max_num_tokens = num_tokens_lists[-1]  # to avoid multi-chunk auxi stream in cuda-graph mode.
+    is_gpt_oss = model_name in {"openai/gpt-oss-120b", "openai/gpt-oss-20b"}
     swiglu_alpha = None
     swiglu_beta = None
     swiglu_limit = None
@@ -417,15 +418,21 @@ def run_moe_torch(
         "w4afp8": 4,
         "nvfp4": 4,
     }[moe_type]
-    if (inter_size // moe_tp_size) % (256 // _weight_bits) != 0:
+    # Blackwell MXFP4 uses TRTLLMGenFusedMoE, whose weight loader pads
+    # before TP sharding. Its logical intermediate size need not satisfy
+    # the CUTLASS plugin's physical-weight alignment check below.
+    # TensorRT-LLM v1.3.0rc20, c25c23f71786bad54d192893d696ce8043426eca:
+    # _torch/modules/fused_moe/quantization.py:5634-5659,5705-5735.
+    native_mxfp4_padding = moe_type in _MXFP4_MOE_TYPES and 100 <= sm_version < 120
+    if not native_mxfp4_padding and (inter_size // moe_tp_size) % (256 // _weight_bits) != 0:
         raise ValueError(
             f"TRT-LLM fused MoE requires the TP-sharded intermediate size to be a multiple "
             f"of 256/weight_bits = {256 // _weight_bits} for {moe_type} (TLLM_CHECK_WITH_INFO "
             f"weight-layout alignment); got inter_size={inter_size} / moe_tp={moe_tp_size} = "
             f"{inter_size // moe_tp_size}"
         )
 
-    if model_name in ["openai/gpt-oss-120b", "openai/gpt-oss-20b"]:
+    if is_gpt_oss:
         swiglu_alpha = torch.tensor([1.702] * (num_experts // moe_ep_size), dtype=torch.float32).to(
             torch.device(device)
         )
@@ -499,14 +506,24 @@ def run_moe_torch(
         router_logits_dtype = torch.float32
     else:
         # for low latency mode in fp4, experts > 128 is not supported.
-        routing_method = RenormalizeMoeRoutingMethod(topk)
+        # Native GPT-OSS uses BF16 routing weights for TRTLLM, FP32 otherwise:
+        # modeling_gpt_oss.py:159-162@1.3.0rc14 (93cb6518), :165-168@1.3.0rc20.
+        if is_gpt_oss:
+            routing_method = RenormalizeMoeRoutingMethod(
+                topk, output_dtype=torch.bfloat16 if model_config.moe_backend == "trtllm" else torch.float32
+            )
+        else:
+            routing_method = RenormalizeMoeRoutingMethod(topk)
 
     create_moe_kwargs = {
         "routing_method": routing_method,
         "num_experts": num_experts,
         "hidden_size": hidden_size,
         "intermediate_size": inter_size,
         "dtype": dtype,
+        # Native GPT-OSS expert construction requests bias=True:
+        # modeling_gpt_oss.py:183@1.3.0rc14 (93cb6518), :189@1.3.0rc20.
+        "bias": is_gpt_oss,
         # In both low latency and attention dp scenarios, create_moe needs not to do allreduce
         # inside op.
         "reduce_results": False,
@@ -634,7 +651,11 @@ def fp32_to_mxfp4(tensor):
             else:
                 continue
 
-    if moe_type != "w4a16_mxfp4":
+    # TRTLLMGen MXFP4 must warm its tactic cache just like serving:
+    # model_engine.py:827-842@1.3.0rc14 (93cb6518). A B200 GPT-OSS
+    # TP1/256-token graph measurement drops 26% after explicit tuning. Keep the
+    # separate Triton/CUTLASS MXFP4 tuning behavior unchanged.
+    if moe_type != "w4a16_mxfp4" or model_config.moe_backend == "trtllm":
         cleanup_empty_json_files(moe_tune_path)
         # The tuned-tactic cache MUST be SM-scoped: tactic indices are positions
         # in the runner's per-arch config table, and replaying another arch's
@@ -648,6 +669,7 @@ def fp32_to_mxfp4(tensor):
         cache_path = (
             f"{moe_tune_path}/sm{get_sm_version()}_"
             f"{moe_type}_{hidden_size}_{inter_size // moe_tp_size}_{num_experts // moe_ep_size}"
+            f"{'_gptoss_bias' if is_gpt_oss else ''}"
         )
         existing_files = glob.glob(f"{cache_path}*")
         cache_loaded = False
@@ -666,6 +688,8 @@ def fp32_to_mxfp4(tensor):
 
         if not cache_loaded:
             torch.cuda.synchronize()
+            tuning_succeeded = False
+            last_tuning_error = None
             for i in range(len(num_tokens_lists)):
                 max_tokens_for_tuning = num_tokens_lists[-i - 1]
                 if max_tokens_for_tuning > max_tokens:
@@ -688,9 +712,16 @@ def fp32_to_mxfp4(tensor):
                                 do_finalize=not min_latency_mode,
                             )
                         torch.cuda.synchronize()
+                        tuning_succeeded = True
                     except Exception as e:
-                        print(f"tune failed for {max_tokens_for_tuning} tokens: {e}, fallback to samller tokens")
+                        accelerator_error = getattr(torch, "AcceleratorError", None)
+                        if accelerator_error is not None and isinstance(e, accelerator_error):
+                            raise
+                        last_tuning_error = e
+                        print(f"tune failed for {max_tokens_for_tuning} tokens: {e}, trying smaller tokens")
                         continue
+            if not tuning_succeeded:
+                raise RuntimeError("MoE autotuning did not succeed for any eligible token count") from last_tuning_error
 
     del hidden_states_max_tokens, logits_max_tokens
     if moe_type == "fp8_block":