DeepSpeed

bb5c23e0 - Use int64 program ids in the SwiGLU and AutoEP fused restore Triton kernels (#8591)

Use int64 program ids in the SwiGLU and AutoEP fused restore Triton kernels (#8591) # Use int64 program ids in the SwiGLU and AutoEP fused restore Triton kernels Fixes https://github.com/deepspeedai/DeepSpeed/issues/8590 ## The problem Four Triton kernels compute memory offsets from `tl.program_id`, which is int32. When the tensor has more than 2^31 elements, the offset wraps to a negative number. The mask looks like a bounds check, but a negative offset passes it. The kernels then load and store memory before the start of the tensor, and CUDA reports `an illegal memory access was encountered`. | kernel | int32 product | overflows when | |---|---|---| | `_swiglu_fwd_kernel`, `_swiglu_bwd_kernel` (`swiglu_triton.py`, #8244) | `pid * BLOCK_SIZE`, with `BLOCK_SIZE = 2048` | elements > 2^31 | | `_weighted_restore_forward_kernel` (`autoep_fused_token_ops.py`, #8326) | `token * out_stride`, where the stride is the hidden size | tokens x hidden > 2^31 | | `_weighted_restore_backward_kernel` (same file) | `token * grad_out_stride` | tokens x hidden > 2^31 | For example, in the SwiGLU kernels: ```python pid = tl.program_id(axis=0) # int32 offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) # wraps negative at pid = 1,048,576 mask = offsets < n_elements # a negative offset passes this check ``` ## When it is reached **SwiGLU** runs in `GroupedExperts` on the gate projection, which has shape `[rows received by this rank, moe_intermediate_size]`. With balanced routing, it faults when `tokens per rank x top_k x moe_intermediate_size > 2^31`: | model | top_k | moe_intermediate_size | faults above, tokens per rank | |---|---|---|---| | Qwen3.5-397B-A17B | 10 | 1024 | 209,715 | | DeepSeek-V3 | 8 | 2048 | 131,072 | These rows are arithmetic, not runs. Unbalanced routing reaches the limit sooner, because one rank can receive up to `expert-parallel size x tokens per rank x top_k` rows. We hit it on Qwen3.5-397B-A17B with 8,192-token sequences on 4 nodes x 8 H200. Every token had been routed to one expert-parallel rank of 32, so SwiGLU received 2,621,568 x 1,024 elements (1.25 x 2^31). **The fused restore** is opt-in (`combine_impl="fused_weighted_sum"`). It faults above 524,288 tokens per rank at hidden size 4096, or above 299,593 at hidden size 7168. ## The change Cast the program id to int64 in all four kernels, so every offset product that uses it is computed in int64: ```python # swiglu_triton.py, both kernels pid = tl.program_id(axis=0).to(tl.int64) # autoep_fused_token_ops.py, both _weighted_restore_*_kernel token = tl.program_id(0).to(tl.int64) ``` Each line has a short comment explaining why, because the mask makes the code look safe without it. Two other `program_id` products in `autoep_fused_token_ops.py` are left as they are: - `_invert_index_kernel` multiplies by a block of 256 over `tokens x top_k` indices. - The hidden-dimension tile in the forward kernel is at most the hidden size. ## Testing One H200, torch 2.8.0+cu126, Triton 3.4.0, `CUDA_LAUNCH_BLOCKING=1`, each size in a fresh process. **Existing unit tests:** `tests/unit/v1/ops/triton_ops/test_swiglu_triton.py` and `tests/unit/v1/ops/triton_ops/test_autoep_fused_token_ops.py`, 86 passed. **Above the limit.** Error is the largest absolute difference from a float32 reference: | kernel | size | `master` | this PR | |---|---|---|---| | SwiGLU, bf16 `[rows, 1024]` | 2,149,580,800 elements (1.001 x 2^31) | illegal memory access | forward 0.0132, backward 0.0077 / 0.0077 | | SwiGLU, bf16 `[rows, 1024]` | 2,684,485,632 elements (1.25 x 2^31) | illegal memory access | forward 0.0147, backward 0.0078 / 0.0075 | | fused restore, top_k 2, hidden 4096 | 524,800 tokens (tokens x hidden = 1.001 x 2^31) | illegal memory access | forward 0.0077, row gradient 0.0075, score gradient 7.6e-6 | For comparison, `master`'s SwiGLU errors below the limit are the same size: forward 0.0144 and backward 0.0076 / 0.0075. They come from rounding the float32 result to bf16. **Below the limit, results are bit-identical to `master`.** The sha256 of every output and gradient matches: | kernel | size | tensors compared | |---|---|---| | SwiGLU | 671,612,928 elements (0.31 x 2^31) | output, gate gradient, up gradient | | fused restore | 524,000 tokens (0.9995 x 2^31) | output, row gradient, score gradient | **No time cost.** SwiGLU, median of 50 calls at 671,612,928 elements: | | forward | backward | |---|---|---| | `master` | 0.962 ms | 1.634 ms | | this PR | 0.962 ms | 1.634 ms | Both run at about 4.2 TB/s forward and 4.1 TB/s backward, so memory traffic sets the time. **No new unit test.** A test needs more than 2^31 elements: 12.9 GB for SwiGLU in bf16, and 25.8 GB for the fused restore forward and backward. I can add one with a skip for GPUs that have less free memory, if you want it in CI. ## Reproducer ```python import torch import torch.nn.functional as F from deepspeed.ops.triton_ops.swiglu_triton import swiglu gate = torch.randn(2_099_200, 1024, dtype=torch.bfloat16, device="cuda") # 1.001 x 2**31 elements up = torch.randn(2_099_200, 1024, dtype=torch.bfloat16, device="cuda") out = swiglu(gate, up) torch.cuda.synchronize() # master: "an illegal memory access was encountered" print((out[-1].float() - F.silu(gate[-1].float()) * up[-1].float()).abs().max().item()) ``` On this PR it prints a last-row error of about 0.01. Signed-off-by: pengdurice <pengduhit@gmail.com>