[DSv4][P5-5] Shared Expert MLP: strict CUDA + Triton kernels - #387
[DSv4][P5-5] Shared Expert MLP: strict CUDA + Triton kernels#387jyizheng wants to merge 4 commits into
Conversation
…n#64) Implements shared_expert_mlp_fwd/bwd per the P5-S0 contract: every valid token runs fc1 -> one-round SwiGLU -> fc2 on BF16 frozen weights, backward returns dX only (FP32 accumulator), and the shared output stays independent of the routed path. Both backends reproduce the FP32 oracle's numeric profile oracle-fp32-serial-v1 byte-for-byte on the same device: one lane owns one output element and reduces serially in ascending k, multiply and add rounded separately (__fmul_rn/__fadd_rn on CUDA, uncontracted IEEE fp32 in Triton), sigmoid computed as 1/(1+expf(-x)) to match torch.sigmoid on FP32 CUDA tensors. No cross-lane floating-point reduction exists anywhere, so results are batch/padding invariant by construction (fwd(x)[t] == fwd(x[t:t+1]) byte-equal). The one-round SwiGLU core runs in shared mode (p_s=None, no clamp, per S0 decision D6) and is the reuse point for P5-2 (RL-Align#63). Providers subclass ReferenceProvider and override only the two shared-expert methods, so the full acceptance command runs unchanged; unsupported input (non-CUDA device, missing extension or triton, foreign numeric profile) raises instead of falling back (fail-closed). Provenance records split_k=1 / serial-ascending-k / no-FMA per the P5-5 provenance requirement. Acceptance: python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider --device cuda python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider --device cuda pytest tests/test_shared_expert_mlp.py python benchmarks/benchmark_shared_expert_mlp.py
tl.exp is the fast exp2-based path and does not bit-match torch.sigmoid; libdevice __nv_expf does (0/4M mismatches on the device probe).
Neither tl.exp (exp2-based) nor libdevice __nv_expf bit-matches the nvcc expf inside torch.sigmoid (~45% / ~10% of fp32 values differ by 1 ulp); the tiny fixtures passed only because the BF16 round absorbed the difference, and the T=256 benchmark cross-check caught the divergence. The Triton path now takes torch.sigmoid(gate) as a kernel input and fuses the remaining SwiGLU math; a (256, 1024, 512) cross-backend byte-equality test locks the regression in.
… _rn ops The compiler may contract a * b + c into an FMA; at T=256 that rounded dsilu = sig * (1 + g * (1 - sig)) differently on 2/262144 dgate elements (1 ulp after the BF16 round). All mul/add/sub in the Triton strict GEMM and SwiGLU kernels now go through libdevice add_rn/mul_rn/sub_rn, the exact Triton spelling of the CUDA kernel's __fadd_rn/__fmul_rn/__fsub_rn.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
Flink-ddd
added
DSv4
deepseek-P5
platform: cuda
There was a problem hiding this comment.
Thanks for the clean and clear implementations. I pulled the branch onto dsv4-p5-dev and verified both backends on an H100 (torch 2.14+cu130, Triton 3.8). CUDA is byte-equal to the oracle on both fixtures, the T=256 cross-backend check, and a 2M-element SwiGLU sweep, and backward is batch-invariant too. Good Work!
Some issues about performance:
The performance is 1000x slower than cublas. I think there is still room for optimize.
Can you reference to Det GEMM to see if we can reuse this kernel or reference this design?
Note that reusing it may not align with oracle I provided, but we can ignore it for now.
Our goal is ensure batch-invariance and also high performance.
3 participants
Implements
shared_expert_mlp_fwd/bwd(#64) per theP5-S0contract: every valid token runs fc1 -> one-round SwiGLU -> fc2 on BF16 frozen weights; backward returns dX only (FP32 accumulator), and the shared output stays independent of the routed path (combine belongs to P6).Both backends reproduce
oracle-fp32-serial-v1byte-for-byte on the same device: one lane owns one output element and reduces serially in ascending k, multiply and add rounded separately (__fmul_rn/__fadd_rnon CUDA, libdevicemul_rn/add_rnon Triton). No cross-lane floating-point reduction exists, so results are batch/padding invariant by construction. The one-round SwiGLU core ships in shared mode (p_s = None, no clamp, per S0 decision D6) as the reuse point for P5-2 (#63).Acceptance (B300 / sm_103, torch 2.14+cu130)
Tests cover: fixture byte-equality (
shared_t1/shared_t16), batch/padding invariance (fwd(x)[t] == fwd(x[t:t+1])per row), nodW(weight hashes stable, no grad), shared/routed boundary independence, CUDA/Triton cross-backend byte-equality at (T=256, H=1024, F=512), and fail-closed behavior (CPU input / missing backend / foreign numeric profile raise; no silent oracle fallback).Benchmark (torch-native vs Triton vs CUDA, H=4096 F=2048, fwd+bwd ms)
The strict backends are asserted byte-equal to each other on every shape. The gap vs cuBLAS is the cost of the serial one-lane-per-element reduction; epilogue/perf work is the P5 step-10 follow-up, out of scope here.
Notes
expfbit-matchestorch.sigmoid(0/4M mismatches on a device probe) whiletl.exp(exp2-based) and libdevice__nv_expfdo not (~45% / ~10% of values off by 1 ulp). The CUDA kernel computes1/(1+expf(-x))inline; the Triton path takestorch.sigmoid(gate)as a kernel input and fuses the remaining SwiGLU math.*_rn: the compiler may contracta*b+cinto an FMA, which showed up as a 1-ulp dgate drift on 2/262144 elements at T=256. A cross-backend byte-equality test at that size locks the regression in.provenance()recordssplit_k=1,reduction=serial-ascending-k, no-FMA rounding, per the P5-5 provenance requirement.