[https://nvbugs/6633931][fix] Size the SM90 fp8 block-scale MoE workspace for the buffers it actually uses

[pranav-nvidia]

@coderabbitai summary

Description

On SM90, getWorkspaceDeviceBufferSizes over-sizes the two overlapped MoE buffers for the fp8 block-scale path, and #16849 made it worse.

Under fp8 activation fusion (useFp8BlockScaleActFusion(), SM90 only) the overlapped inputs buffer holds pre-quantized fp8 activations plus their padded 1x128 scales — never the bf16 activations that permuted_data_size and fc1_result_size are sized for. #16849 already computes the correct size with fp8BlockScaleRegionBytes(...), but takes a max() with the stale bf16 sizes instead of replacing them. Since sizeof(T) == sizeof(UnfusedGemmOutputType) == 2, permuted_data_size alone pins the buffer at permuted_elems * 2, roughly twice what the fused path needs. This replaces rather than maxes.

The FC2 GEMM output is now reserved in exactly one of the two buffers, selected by the same predicate the forward pass uses to route it (isActivationPrequantized() / use_fused_block_scale_quant): the fused path writes it into the outputs buffer (glu_inter_result_, because its input fc1_result_ is aliased onto fc2_result_), the unfused path writes fc2_result_ in the inputs buffer. It was previously counted in both. This also restores outputs = fc1_output on the unfused path — #16849 widened it to max(fc1, fc2) unconditionally, so non-SM90 architectures were reserving an FC2-output-sized buffer they never write to.

Measured on H100 (SM90) at DeepSeek-V3.2 dimensions (hidden 7168, inter 2048, 256 experts, top_k 8), sweeping MoE rows, reading the workspace size the op logs at DEBUG:

MoE rows	before	after
8192	1.759 GiB	1.336 GiB
16384	3.517 GiB	2.671 GiB
32768	7.034 GiB	5.340 GiB
65536	14.068 GiB	10.679 GiB

At 65536 rows that is 15,105,925,632 -> 11,467,047,424 bytes, a 24.1% reduction. This is sizing only; no kernel or launch parameter changes.

Context: the 14.02 GiB allocation is what OOMs in NVBug 6633931 on 8xH200. That test also runs MoE unchunked, which is fixed separately in #18322; the two changes are independent and both are needed to give the test real margin.

Test Coverage

tests/unittest/_torch/modules/moe/test_moe_backend.py -k "FP8_BLOCK_SCALES and CUTLASS" — 10 passed on H100, run on the unpatched build as a control.

Because that suite only covers small shapes (seq 1-8, hidden 256-2048) where the workspace is negligible, an under-sized buffer that only overflows at large num_moe_inputs would corrupt FC2's input silently rather than crash. So the change was additionally checked at full DeepSeek-V3.2 dimensions up to 65536 MoE rows (524288 expanded rows) with identical seeds: output fingerprints are bit-identical before and after at every row count (out_sum 194420032929792.0, out_absmax 14483456.0, out_mean 413869.1875, 0 NaN, 0 Inf), and getWorkspaceSize no longer needs the removed capacity.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.