[bug] Dropless DeepEP makes GPT-OSS 20B logits depend on other EP-rank samples

[yaoyu-33]

Problem

With Megatron Bridge's GPT-OSS 20B model and Megatron-Core's DeepEP MoE token dispatcher in dropless mode, rank 0's raw logits for a fixed token sequence X depend on the unrelated sequences processed by the other expert-parallel ranks.

The comparison is:

• homogeneous EP batch: every EP rank processes X
• heterogeneous EP batch: rank 0 processes the byte-identical X; every other EP rank processes a distinct Y

No labels, loss, backward pass, optimizer, or dataset pipeline are involved. The test uses a real GPT-OSS 20B Megatron checkpoint and directly compares the model's raw logits.

Observed result:

effective_config dispatcher=flex backend=deepep capacity_factor=None pad_to_capacity=False ep=8
logits_shape=(1, 128, 201088) logits_dtype=torch.float32
repeat=0 target_max_abs_delta=0.25
repeat=1 target_max_abs_delta=0.25
repeat=2 target_max_abs_delta=0.25
status=INVARIANT_VIOLATED mode=dropless max_abs_delta=0.25 homogeneous_repeat_delta=0 atol=0 module_global_buffer_reused=True

The difference is deterministic and bit-reproducible. Repeating the homogeneous forward gives an exact delta of zero, while changing only the other EP ranks' sequences gives an exact max logit delta of 0.25.

This is a correctness issue for gradient accumulation and any two-pass method that assumes a sample's forward is a function only of its own tokens and the model weights.

A smaller module-level harness that exercised one _DeepepManager -> fused_dispatch -> expert-local permute -> token-wise expert op -> reverse permute -> fused_combine round did not reproduce the difference. The full 24-layer GPT-OSS model does, suggesting the trigger requires repeated full-model MoE execution, actual grouped experts, or another model-level interaction. Therefore the exact low-level root cause is not yet confirmed.

Minimal repro

Complete standalone reproducer:

https://gist.github.com/yaoyu-33/c8847fa364d3bb2406b608c863cd465b

Prepare a GPT-OSS 20B Megatron checkpoint using the existing example under examples/models/gpt_oss, then run:

export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0
export CUBLAS_WORKSPACE_CONFIG=:4096:8
export HF_HUB_OFFLINE=1
export TRANSFORMERS_OFFLINE=1

uv run python -m torch.distributed.run \
  --nproc_per_node=8 \
  scripts/training/gpt_oss_deepep_batch_invariance_reproducer.py \
  --hf-model-path /path/to/openai-gpt-oss-20b-snapshot \
  --checkpoint-path /path/to/megatron-gpt-oss-20b \
  --seq-length 128 \
  --repeats 3

Relevant effective configuration:

tensor_model_parallel_size = 1
pipeline_model_parallel_size = 1
expert_model_parallel_size = 8
expert_tensor_parallel_size = 1
micro_batch_size = 1

moe_token_dispatcher_type = "flex"
moe_flex_dispatcher_backend = "deepep"
moe_expert_capacity_factor = None
moe_pad_expert_input_to_capacity = False
moe_router_dtype = "fp32"
moe_grouped_gemm = True

The script loads real weights, runs primer -> homogeneous -> primer -> heterogeneous, and compares rank 0's complete raw-logit tensor.

Expected behavior

Rank 0 receives byte-identical input_ids and position IDs in both measured forwards. Its raw logits should therefore be byte-identical regardless of the unrelated sequences processed by other EP ranks:

target_max_abs_delta=0

Affected area

area:training

The failure is observed through Megatron Bridge's GPT-OSS provider and checkpoint loader, while the likely implementation boundary is the Megatron-Core/DeepEP MoE dispatcher.

Regression?

Not sure.

Environment

Megatron Bridge: b4fc0887eb6ef76a808f6c5f830d6fada9b9c96a
Megatron-Core:   8b3d8a5d9719c201efa7add2a8a112b0d095e846
Container:       mbridge-260522.sqsh
Hardware:        1 node, 8 GPUs
Precision:       BF16 model parameters; FP32 router probabilities
Model:           GPT-OSS 20B, real Megatron checkpoint
Parallelism:     TP=1, PP=1, CP=1, EP=8
Sequence length: 128

The run was offline against a pre-populated Hugging Face cache and a real 39 GB Megatron checkpoint.

Control and current limitation

The original suspected workaround is finite expert capacity plus padding. On current ToT, the same-backend control cannot be run because Megatron-Core explicitly rejects it during config finalization:

ValueError: Flex token dispatcher with deepep backend does not support
moe_pad_expert_input_to_capacity

Therefore this report confirms the dropless failure across a second model family, but does not claim that the capacity+pad workaround has been validated on current ToT.

Additional context

• A one-dispatch module-level DeepEP harness was invariant on 2 and 8 GPUs.
• The real 24-layer GPT-OSS model is not invariant.
• The module-global deep_ep.Buffer object was reused throughout the failing run.
• No issue matching “DeepEP dropless logits/batch invariance” was found in either Megatron-Bridge or Megatron-LM at filing time.

[yaoyu-33]

Padded-control follow-up on the same GPT-OSS 20B checkpoint and ToT revisions:

effective_config dispatcher=alltoall backend=None capacity_factor=1.0 pad_to_capacity=True ep=8
repeat=0 target_max_abs_delta=0
repeat=1 target_max_abs_delta=0
repeat=2 target_max_abs_delta=0
status=INVARIANT_HOLDS mode=capacity_pad max_abs_delta=0 homogeneous_repeat_delta=0 atol=0

This is an exact-zero result, not a tolerance-based pass. It used the same real weights, TP=1/PP=1/CP=1/EP=8, sequence length 128, primer/measurement ordering, and deterministic settings as the failing DeepEP run.

Important limitation: current Megatron-Core ToT rejects flex + deepep + moe_pad_expert_input_to_capacity=True during config finalization. Therefore the supported padded control necessarily uses standard alltoall; it changes both dispatcher and capacity mode and does not prove that padding alone fixes DeepEP. It establishes:

• flex + deepep + dropless: deterministic max logit delta 0.25
• alltoall + capacity_factor=1.0 + pad: exact delta 0

The reproducer has also been extended locally with an explicit --dispatcher alltoall --capacity-factor 1.0 --pad-to-capacity control. Next I am instrumenting module outputs in the failing full-model DeepEP run to identify the first module/layer where the two forwards diverge.

[yaoyu-33]

Module-boundary trace follow-up (--trace-modules is now included in the linked gist):

effective_config dispatcher=flex backend=deepep capacity_factor=None pad_to_capacity=False ep=8
repeat=0 target_max_abs_delta=0.25
activation_trace key=module.decoder.layers.2.mlp.combine max_abs_delta=0.00048828125
activation_trace key=module.decoder.layers.2.mlp.postprocess max_abs_delta=0.00048828125
activation_trace key=module.decoder.layers.5.mlp.combine max_abs_delta=0.001953125
activation_trace key=module.decoder.layers.10.mlp.combine max_abs_delta=0.00390625
activation_trace key=module.decoder.layers.14.mlp.combine max_abs_delta=0.03125
activation_trace key=module.decoder.layers.14.layer_output[0] max_abs_delta=0.03125
activation_trace status=FIRST_DIVERGENCE key=module.decoder.layers.2.mlp.combine max_abs_delta=0.00048828125 compared_tensors=240 atol=0
status=INVARIANT_VIOLATED mode=dropless max_abs_delta=0.25 module_global_buffer_reused=True

The instrumentation preserves the original final-logit delta (0.25). On rank 0, layer 2 attention, pre-MLP layernorm, router output/routing map, and MoE preprocess are byte-identical. The first rank-local difference appears when layer 2 MoELayer.combine() returns. Its postprocess() output has the same delta. The small layer 2/5/10 MoE differences round away in the BF16 residual path; layer 14 is the first difference that survives into the transformer-layer output and then propagates.

Current MCore call chain for the first observable local-token divergence:

MoELayer.combine()
  -> MoEFlexTokenDispatcher.token_combine()
  -> _DeepepManager.combine()
  -> fused_combine() / FusedCombine.forward()
  -> deep_ep.Buffer.combine()

This narrows the first observable invariant break to the DeepEP combine boundary. It does not yet distinguish whether the bad value is introduced inside Buffer.combine() or whether the matching X expert value was already corrupted on the dispatch/expert side and only becomes rank-local/observable after combine. The dispatch-side tensors contain the intentionally different global X/Y token mixture, so a raw whole-tensor comparison there would not be a valid invariant test.

[yaoyu-33]

True single-module isolation and correction to the earlier root-cause hypothesis:

I exported rank-local weights plus the real layer-2 pre-MLP inputs, then reproduced this without constructing or running GPTModel. The new script directly builds one Megatron-Core MoELayer, strict-loads the exported router/expert weights, and forwards only that module.

full GPT-OSS model, flex + DeepEP + dropless: final-logit delta = 0.25
single layer 2, flex + DeepEP + dropless:       output delta = 0.00048828125
single layer 2, homogeneous control:            output delta = 0

The single-layer result is deterministic across repeats and exactly matches the first divergence seen by the full-model module trace.

Knob matrix

All entries below use the same real layer-2 weights and exported activations unless marked synthetic.

Configuration	Max delta	Result
DeepEP dropless, grouped experts (baseline)	0.00048828125	reproduces
DeepEP dropless, no permute fusion	0.00048828125	reproduces
DeepEP dropless, reset fused_a2a._buffer before every forward	0.00048828125	reproduces
Standard alltoall dropless	0.00048828125	reproduces
DeepEP capacity factor 1.0, no padding	0.001953125	reproduces
Standard alltoall, capacity factor 1.0 + padding	0	invariant
DeepEP dropless, tokenwise-affine expert stand-in	0	invariant
DeepEP dropless, SequentialMLP with the same expert weights	0.00390625	reproduces
DeepEP dropless, top-k 2 / top-k 3	0 / 0	invariant for this fixture
DeepEP dropless, top-k 8	0.00390625	reproduces
DeepEP dropless, 2 SMs instead of 20	0.00048828125	reproduces
Synthetic inputs, equal lengths	0	invariant
Synthetic inputs, X length 128 / neighbors length 384	0	invariant
Real homogeneous neighbors	0	invariant

Unsupported controls: top-k 1 requires moe_router_pre_softmax in this configuration; DeepEP rejects an odd num_sms=1.

Corrected interpretation

This evidence falsifies the strongest version of the stale module-global DeepEP buffer hypothesis:

• Buffer reuse is not necessary: forcibly replacing the buffer before each forward does not change the delta.
• DeepEP is not necessary: standard alltoall dropless gives the same single-layer delta.
• Permute fusion is not necessary.
• Dispatch/combine alone is not sufficient: replacing the real experts with a per-token affine function makes the result exactly invariant.
• GroupedLinear specifically is not necessary: SequentialMLP also reproduces.
• Fixed padded expert shapes are invariant in the tested supported configuration; finite capacity without padding is not.

The current evidence points to deterministic finite-precision differences in real expert GEMMs when dropless routing changes per-expert token counts and therefore GEMM shapes/kernel tiling. The combine boundary is still where the changed expert result first becomes observable for rank 0, but it is not evidence that deep_ep.Buffer.combine() corrupts data. The full transformer amplifies the layer-2 BF16 difference from 4.88e-4 to a final-logit delta of 0.25.

Top-k and synthetic-input outcomes show that variable counts alone are not sufficient for every activation/routing pattern, so I would not infer a universal top-k threshold from this one fixture.

Updated gist (now contains both the full-model/export script and the true single-MoELayer reproducer):
https://gist.github.com/yaoyu-33/c8847fa364d3bb2406b608c863cd465b

[yaoyu-33]

Follow-up: a concrete path for making dropless mode batch-invariant without dropping tokens.

I added two controls that keep all of the following unchanged:

• real GPT-OSS layer-2 router and expert weights;
• real exported pre-MLP activations;
• DeepEP flex dispatch/combine;
• dropless routing and the original heterogeneous token counts;
• the real expert MLP, including both GEMMs and its activation.

The only change is how each local expert's tokens are presented to the expert GEMMs.

Expert execution	Shape policy	Max delta
SequentialMLP baseline	one dynamic M=tokens_per_expert[e] GEMM per expert	0.00390625
Real SequentialMLP, tokenwise proof	fixed M=1	0
Real SequentialMLP, chunked prototype	fixed M=32; zero-pad every tail chunk and remove padding afterward	0

The fixed-32 result is exact zero, not a tolerance pass. Routing counts still differ by as much as 173 tokens between the homogeneous and heterogeneous cases, and module_global_buffer_reused=True. Thus DeepEP communication remains genuinely dropless; no routed token is capped or discarded.

Example controls in the updated single-layer reproducer:

# Slow proof: every real expert GEMM has M=1.
--expert-mode sequential_tokenwise

# More practical prototype: real expert GEMMs run in fixed M=32 chunks.
--expert-mode sequential_fixed_shape --expert-chunk-size 32

Updated reproducer:
https://gist.github.com/yaoyu-33/c8847fa364d3bb2406b608c863cd465b

Proposed fix

Add an expert-compute padding mode independent of router capacity/drop semantics:

1. Keep the existing dropless dispatcher and exact tokens_per_expert.
2. After dispatch/permutation, split each expert's token segment into fixed-size chunks (for example 32/64/128 rows).
3. Zero-pad the final chunk to exactly that fixed row count and set its dummy routing probabilities to zero.
4. Run FC1, activation, and FC2 with the same fixed GEMM shape for every chunk.
5. Strip dummy rows and pass the original token order/counts to combine.

This differs importantly from moe_expert_capacity_factor + moe_pad_expert_input_to_capacity: capacity mode can drop overflow tokens, while fixed-shape expert chunks only add temporary compute padding and therefore preserve true dropless semantics.

For an efficient TE grouped implementation, each expert chunk should be treated as a fixed-M GEMM group/tile. Merely rounding each expert from N to ceil(N/32)*32 and executing one variable-M GEMM may still select a different kernel/reduction path, so it would need the same invariance test. The proof above executes every chunk at exactly M=32.

Practical options

1. Correctness-first fix: add fixed-shape expert chunks in MCore/TE. This is the only tested option that is both truly dropless and exact-zero for this reproducer.
2. Current workaround: standard alltoall + finite capacity + padding, with capacity chosen high enough that observed dropped-token count is zero. This was exact-zero here, but it is not formally dropless unless overflow is impossible, and it adds padded compute.
3. Kernel-level alternative to investigate: pin a GEMM algorithm whose reduction order is invariant across M, or use a higher-precision accumulation/output boundary. Ordinary deterministic-kernel settings only guarantee repeatability for the same shape; our result shows they do not guarantee equality across different M values.

The chunked prototype is forward-only validation. A production patch still needs backward parity, zero-token experts, TP/EP combinations, grouped-GEMM integration, FP8, performance, and CUDA-graph coverage. But the exact-zero M=1 and M=32 results provide direct evidence that fixed-shape expert compute can make DeepEP dropless work without token dropping.

[yaoyu-33]

Follow-up: post-dispatch padding to a multiple of 64 is sufficient for this real GPT-OSS 20B reproducer, while keeping the dispatcher genuinely dropless.

I implemented a small MCore prototype with a new moe_expert_padding_multiple config. It reuses the existing Fp8Padding / Fp8Unpadding machinery in TEGroupedMLP, but enables it for BF16 independently of quantization and independently of expert capacity:

1. DeepEP dispatch and the original tokens_per_expert remain unchanged.
2. Each local expert segment is zero-padded to ceil(tokens / 64) * 64 immediately before expert compute.
3. Dummy rows have zero routing probability.
4. Dummy rows are removed before DeepEP combine.

There is no capacity, no overflow path, and no token dropping (moe_expert_capacity_factor=None).

Results

Test	Result
Baseline real layer-2 TEGroupedMLP, dynamic M	0.00048828125
Grouped expert, rounded to multiple of 32	0.00048828125
Grouped expert, rounded to multiple of 64	0
Native MCore pad64 prototype, heterogeneous-neighbor sweep (1 through 7 EP peers)	all 0
Native MCore pad64 prototype, forward + backward check	forward 0, target input grad 0, parameter grad 0
Full GPT-OSS 20B model, baseline final logits	0.25
Full GPT-OSS 20B model, native dropless pad64	0

The direct real-expert cross-shape probe also explains why 64 works for this kernel/model combination:

reference M=16, candidate M=32:   delta=0
reference M=16, candidate M=64+:  delta=0.0009765625

reference M=64, candidate M=128:  delta=0
reference M=64, candidate M=192:  delta=0
reference M=64, candidate M=256:  delta=0
reference M=64, candidate M=384:  delta=0
reference M=64, candidate M=512:  delta=0
reference M=64, candidate M=768:  delta=0
reference M=64, candidate M=1024: delta=0

Thus the observed discontinuity is a GEMM shape/kernel boundary at M=64. Once every expert M is a multiple of 64, all tested shapes use a numerically identical reduction path for the matching rows. This is consistent with TE's recent NVFP4 work also moving its required M alignment from 32 to 64 (TransformerEngine #2351), although that PR is not itself a BF16 batch-invariance guarantee.

Recommended solution

I think a practical MCore fix is to add post-dispatch expert padding as an explicit setting, separate from both router quantization padding and moe_expert_capacity_factor:

moe_expert_capacity_factor = None   # still dropless
moe_expert_padding_multiple = 64    # compute-only padding

The prototype is small (two MCore files, 43 insertions / 8 deletions) and works for both TEGroupedMLP and SequentialMLP. The patch and updated full-model/single-module reproducers are attached to the gist:

https://gist.github.com/yaoyu-33/c8847fa364d3bb2406b608c863cd465b

For a stronger cross-hardware contract, fixed M=64 chunks are safer than one rounded variable-M grouped GEMM: every GEMM launch then has exactly the same M, at the cost of more launches. However, the rounded-M=64 implementation is the lower-overhead option and passed the real module sweep, backward check, direct M=64..1024 probe, and full-model test here. It should still be validated on other GPU generations, TE versions, dtypes, and expert hidden sizes before treating 64 as universal.

Related batch-invariant-kernel work

Megatron-LM #4871 is directionally relevant and adds DeepGEMM-backed BF16 batch-invariant grouped GEMMs. I ported its current head onto our MCore ToT and built its exact pinned DeepGEMM revision. In this exact DeepEP training reproducer it still produced a maximum module delta of 0.001953125 for neighbor counts 1–6 (the tokenwise expert control remained exactly zero). So that PR is promising for the long-term kernel contract, but its current training-path integration is not yet a drop-in fix for this issue.

This also reinforces the corrected diagnosis from the previous comment: zeroing/replacing the DeepEP communication buffer is not the fix. The failure is deterministic cross-shape expert GEMM numerics; compute-only expert padding removes the shape transition while preserving dropless routing.

[yaoyu-33]

Follow-up with orthogonal experiments for the two batch-invariance hazards suggested in review: expert GEMM and expert-contribution combine order.

Test environment: H100 80GB, BF16, EP=8, MCore 8b3d8a5d9719c201efa7add2a8a112b0d095e846, the same exported GPT-OSS 20B layer-2 weights and real pre-MLP activations.

1. Expert GEMM: confirmed, with the first difference in FC2

I added a direct real-expert probe that completely bypasses the router, dispatch, communication, unpermute, and combine. It forwards byte-identical first rows while changing only GEMM M, with hooks at FC1 output, activated FC2 input, FC2 output, and final expert output.

SequentialMLP, reference M=16:

Candidate M	FC1 delta	Activation delta	FC2 delta	Expert-output delta
16	0	0	0	0
32	0	0	0	0
64	0	0	0.00048828125	0.0009765625
128..1024	0	0	0.00048828125	0.0009765625

Thus the first changed tensor is the output of FC2 GEMM itself. A GroupedLinear probe shows the same raw FC2 boundary, so this is not introduced by MoE combine and is not exclusive to either GroupedMLP or SequentialMLP.

I then enabled CUBLASLT_LOG_LEVEL=5. The trace directly shows the suspected split-K transition:

FC1 M=16: algoId=66 tile=64x16                                  # no split-K
FC2 M=16: algoId=66 tile=64x16 reduction=COMPUTE_TYPE splitK=2
FC1 M=32: algoId=66 tile=64x32                                  # no split-K
FC2 M=32: algoId=66 tile=64x32 reduction=COMPUTE_TYPE splitK=2
FC1 M=64: algoId=66 tile=64x64                                  # no split-K
FC2 M=64: algoId=66 tile=64x32 cluster=4x2x1                    # no split-K

The numerical families exactly match the algorithm families: M=16/32 use FC2 split-K=2 and are bit-identical to each other; M>=64 does not use split-K and the tested matching rows are bit-identical within that family. This explains why compute-only pad-to-64 works: it moves short FC2 calls out of the split-K regime.

The evidence confirms FC2 GEMM as the source and strongly supports the split-K selection boundary as the mechanism. Strictly isolating split-K from every other coupled cuBLASLt algorithm attribute would require pinning the algorithm and toggling only CUBLASLT_ALGO_CONFIG_SPLITK_NUM, which TE does not expose on this path. cuBLASLt documents split-K and its reduction schemes here: https://docs.nvidia.com/cuda/cublas/index.html#cublasltreductionscheme-t

2. Combine/unpermute: independently confirmed when permute fusion is off

For this control I removed all GEMMs. Each global expert emits a different signed power-of-two pattern, explicitly multiplied by the real routing probability. This makes top-k accumulation order-sensitive while retaining the real router, dispatch, communication, and combine. Each result sweeps 1..7 heterogeneous neighbor ranks, three repeats each.

Dispatcher	Top-k	Permute fusion	torch.use_deterministic_algorithms	Max composition delta	Homogeneous repeat delta
DeepEP	4	on	off	0	0
DeepEP	8	on	off	0	0
alltoall	4	on	off	0	0
alltoall	8	on	off	0	0
DeepEP	8	off	off	0.001953125	0.001953125
alltoall	8	off	off	0.0625	0.0625
DeepEP	8	off	on	0	0
alltoall	8	off	on	0	0

The nonzero homogeneous-repeat delta proves ordinary nondeterminism in the unfused accumulation path, not only a deterministic dependency on batch composition. Current MCore uses scatter_add_ for unfused unpermute by default and switches to index_add_ when PyTorch deterministic algorithms are enabled: https://github.com/NVIDIA/Megatron-LM/blob/8b3d8a5d9719c201efa7add2a8a112b0d095e846/megatron/core/transformer/moe/moe_utils.py#L518-L529

For the original issue configuration (moe_permute_fusion=True), the adversarial GEMM-free combine control remains exact for both DeepEP and alltoall. Therefore combine is a real second hazard, but it is not the source of the original fused GPT-OSS delta; FC2 split-K explains that path.

Implication for the fix

Both protections are needed for a general batch-invariant MoE mode:

1. A batch-invariant FC2/grouped-GEMM implementation that does not change reduction semantics with M. Pinning a suitable non-split-K algorithm or using a batch-invariant GEMM kernel is the principled fix. Pad-to-64 is a validated H100 workaround, not a universal kernel contract.
2. A deterministic combine/unpermute implementation. Keep the invariant fused path or use deterministic accumulation for the unfused path; fixing GEMM alone does not fix the scatter_add_ channel.

Updated reproducer and a compact result report with the trace excerpt:
https://gist.github.com/yaoyu-33/c8847fa364d3bb2406b608c863cd465b

[yaoyu-33]

Follow-up: a validated no-padding workaround for the split-K source identified above.

Workaround

Restrict Transformer Engine's cuBLAS workspace to 4 bytes before the first TE GEMM/workspace allocation. In the TE version in the current Megatron Bridge container, the helper lives in transformer_engine.pytorch.cpp_extensions.gemm:

from transformer_engine.pytorch.cpp_extensions import gemm as te_gemm

# Must run before constructing/running the model. Clear the cached workspace in
# case this is used in a process that imported TE earlier.
te_gemm.get_cublas_workspace.cache_clear()
te_gemm.get_cublas_workspace_size_bytes = lambda: 4

Equivalently, for a source-level WAR, make get_cublas_workspace_size_bytes() return 4 in:

transformer_engine/pytorch/cpp_extensions/gemm.py

Older TE versions placed the same helper in transformer_engine.pytorch.module.base; the location is version-dependent. The important requirement is that the workspace passed to cublasLtMatmulAlgoGetHeuristic is actually 4 bytes.

I added this to the reproducer as:

python -m torch.distributed.run --nproc_per_node=8 \
  gpt_oss_deepep_batch_invariance_reproducer.py \
  --hf-model-path /path/to/openai-gpt-oss-20b \
  --checkpoint-path /path/to/megatron-gpt-oss-20b \
  --seq-length 128 --repeats 3 \
  --te-cublas-workspace-bytes 4 --assert-invariant

For PyTorch GEMMs rather than TE-managed GEMMs, the corresponding process-level WAR is:

export CUBLAS_WORKSPACE_CONFIG=:0:0
export CUBLASLT_WORKSPACE_SIZE=4

Full GPT-OSS 20B E2E result

Environment: H100 80GB, BF16, TP=1/PP=1/CP=1/EP=8, flex + DeepEP + dropless, sequence length 128, MCore 8b3d8a5d9719c201efa7add2a8a112b0d095e846, Bridge b4fc0887eb6ef76a808f6c5f830d6fada9b9c96a, and the same real checkpoint/input ordering used throughout this issue.

The only A/B change was the TE workspace size:

Configuration	Repeat 0	Repeat 1	Repeat 2	Max logit delta
TE default workspace (33,555,456 bytes)	0.25	0.25	0.25	0.25
TE workspace = 4 bytes	0	0	0	0

Relevant output:

effective_config dispatcher=flex backend=deepep capacity_factor=None pad_to_capacity=False expert_padding_multiple=None batch_invariant_mode=False te_cublas_workspace_bytes=33555456 ep=8
repeat=0 target_max_abs_delta=0.25
repeat=1 target_max_abs_delta=0.25
repeat=2 target_max_abs_delta=0.25
status=INVARIANT_VIOLATED mode=dropless max_abs_delta=0.25 homogeneous_repeat_delta=0

effective_config dispatcher=flex backend=deepep capacity_factor=None pad_to_capacity=False expert_padding_multiple=None batch_invariant_mode=False te_cublas_workspace_bytes=4 ep=8
repeat=0 target_max_abs_delta=0
repeat=1 target_max_abs_delta=0
repeat=2 target_max_abs_delta=0
status=INVARIANT_HOLDS mode=dropless max_abs_delta=0 homogeneous_repeat_delta=0

The isolated real layer-2 MoELayer gives the same conclusion:

Configuration	Layer output delta
TE default workspace	0.00048828125
TE workspace = 4 bytes	0

Confirmation that the WAR removes split-K

With the default TE workspace, the earlier cuBLASLt trace showed FC2 at M=16/32 using numSplitsK=2 with REDUCTION_SCHEME_COMPUTE_TYPE, while M>=64 used a non-split-K algorithm.

With the 4-byte WAR:

• cuBLASLt reported maxWorkspaceSizeinBytes=4 and workSpaceSizeInBytes=4;
• across the traces from all eight ranks, there were zero matmuls containing numSplitsK or reductionScheme;
• therefore these calls use the default effective split-K=1/no separate reduction path.

A standalone PyTorch BF16 shape probe also changed from a cross-M max delta of 1.0 to exact 0 for every tested M=16..1024 with CUBLAS_WORKSPACE_CONFIG=:0:0 CUBLASLT_WORKSPACE_SIZE=4.

Scope and limitations

This is a useful correctness WAR for the reported fused-permute, BF16 GPT-OSS configuration, and unlike capacity padding it remains genuinely dropless and adds no dummy-token compute. It is not yet a production fix:

• it affects every TE cuBLASLt GEMM in the process, not only MoE FC2;
• performance impact has not yet been benchmarked;
• this E2E check is forward-only; training backward/optimizer parity under the WAR has not yet been tested;
• it does not replace deterministic combine protection for the separate unfused scatter_add_ hazard described above;
• other GPU architectures, dtypes, hidden sizes, FP8, and CUDA-graph configurations still need validation.

The preferable production solution is for TE/MCore batch-invariant mode to select/pin an appropriate non-split-K GEMM algorithm (and separately guarantee deterministic combine), rather than globally starving all cuBLASLt calls of workspace.