[PR #408] [CLOSED] Support arbitrary hidden sizes in intranode combine kernel #552

Closed
opened 2026-10-27 21:03:58 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/408
Author: @varun-sundar-rabindranath
Created: 9/16/2025
Status: ❌ Closed

Base: main ← Head: intranode/fix-hidden-size


📝 Commits (4)

📊 Changes

3 files changed (+22 additions, -9 deletions)

View changed files

📝 csrc/kernels/intranode.cu (+15 -5)
📝 csrc/kernels/utils.cuh (+2 -2)
📝 tests/test_intranode.py (+5 -2)

📄 Description

Intranode combine kernels fail for arbitrary hidden sizes.

The kernel has all the required support to handle arbitrary hidden size. The PR simply fixes a spot that seemed to assume hidden_int4 % 32 == 0.

A for-loop inside the kernel invokes __syncwarp(). This is fine, when all the threads in a warp executes the __syncwarp() (the good case). But this is not always the case with arbitrary hidden size. There can be residuals where only a few threads in a warp would execute the for-loop. To facilitate this, the PR makes appropriate warp_masks that the __syncwarp() could be called with.

Test:
Tested with hidden-size 2880 (gpt-oss hidden-dim)

This PR

python3 test_intranode.py --num-processes=2 --hidden 2880 
[config] num_tokens=4096, hidden=2880, num_topk=8
[layout] Kernel performance: 0.033 ms

[testing] Running with BF16, without top-k (async=False, previous=False) ... passed
[testing] Running with BF16, with top-k (async=False, previous=False) ... passed
[testing] Running with BF16, without top-k (async=False, previous=False) ... passed
[testing] Running with BF16, with top-k (async=False, previous=False) ... passed
[testing] Running with BF16, without top-k (async=True, previous=False) ... passed
[testing] Running with BF16, with top-k (async=True, previous=False) ... passed
[testing] Running with BF16, without top-k (async=True, previous=False) ... passed
[testing] Running with BF16, with top-k (async=True, previous=False) ... passed
[testing] Running with BF16, without top-k (async=False, previous=True) ... passed
[testing] Running with BF16, with top-k (async=False, previous=True) ... passed
[testing] Running with BF16, without top-k (async=False, previous=True) ... passed
[testing] Running with BF16, with top-k (async=False, previous=True) ... passed
[testing] Running with BF16, without top-k (async=True, previous=True) ... passed
[testing] Running with BF16, with top-k (async=True, previous=True) ... passed
[testing] Running with BF16, without top-k (async=True, previous=True) ... passed
[testing] Running with BF16, with top-k (async=True, previous=True) ... passed

[tuning] SMs 24, NVL chunk 4: 97.48 GB/s (NVL), avg_t: 482.75 us
[tuning] SMs 24, NVL chunk 6: 124.54 GB/s (NVL), avg_t: 377.87 us
[tuning] SMs 24, NVL chunk 8: 139.62 GB/s (NVL), avg_t: 337.05 us
[tuning] SMs 24, NVL chunk 10: 150.22 GB/s (NVL), avg_t: 313.26 us
...

main

python3 test_intranode.py --num-processes=2 --hidden 2880  

[config] num_tokens=4096, hidden=2880, num_topk=8
[layout] Kernel performance: 0.034 ms

[testing] Running with BF16, without top-k (async=False, previous=False) ...WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources.
WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources.
[rank0]:[W916 00:21:13.202446982 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
W0916 00:21:14.463000 1862405 torch/multiprocessing/spawn.py:169] Terminating process 1862487 via signal SIGTERM
Traceback (most recent call last):
  File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 277, in <module>
    torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes)
  File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 340, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 296, in start_processes
    while not context.join():
  File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 215, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 1 terminated with the following error:
Traceback (most recent call last):
  File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 90, in _wrap
    fn(i, *args)
  File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 247, in test_loop
    test_main(args, i, local_rank, num_ranks, rank, buffer, group)
  File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 160, in test_main
    assert calc_diff(check_x, ref_x) < 5e-6
AssertionError

🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/408 **Author:** [@varun-sundar-rabindranath](https://github.com/varun-sundar-rabindranath) **Created:** 9/16/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `intranode/fix-hidden-size` --- ### 📝 Commits (4) - [`73ec503`](https://github.com/deepseek-ai/DeepEP/commit/73ec50348f0362f0709841990f11338a0b82a8aa) add warpmask for __syncwarp - [`f7fb36f`](https://github.com/deepseek-ai/DeepEP/commit/f7fb36f407252e26be743be7fd50a8eb07fd1266) pass membermask to elect_one_sync - [`592aa14`](https://github.com/deepseek-ai/DeepEP/commit/592aa1449145d6c475b22013dc2e97b5e45f3efc) compute mask - [`fd4e7e4`](https://github.com/deepseek-ai/DeepEP/commit/fd4e7e41c0d1dbadf4aac23f5343b3d669638b1a) fix whitespace ### 📊 Changes **3 files changed** (+22 additions, -9 deletions) <details> <summary>View changed files</summary> 📝 `csrc/kernels/intranode.cu` (+15 -5) 📝 `csrc/kernels/utils.cuh` (+2 -2) 📝 `tests/test_intranode.py` (+5 -2) </details> ### 📄 Description Intranode combine kernels fail for arbitrary hidden sizes. The kernel has all the required support to handle arbitrary hidden size. The PR simply fixes a spot that seemed to assume `hidden_int4 % 32 == 0`. A for-loop inside the kernel invokes `__syncwarp()`. This is fine, when all the threads in a warp executes the `__syncwarp()` (the good case). But this is not always the case with arbitrary hidden size. There can be residuals where only a few threads in a warp would execute the for-loop. To facilitate this, the PR makes appropriate `warp_masks` that the `__syncwarp()` could be called with. Test: Tested with hidden-size 2880 (gpt-oss hidden-dim) `This PR` ``` python3 test_intranode.py --num-processes=2 --hidden 2880 [config] num_tokens=4096, hidden=2880, num_topk=8 [layout] Kernel performance: 0.033 ms [testing] Running with BF16, without top-k (async=False, previous=False) ... passed [testing] Running with BF16, with top-k (async=False, previous=False) ... passed [testing] Running with BF16, without top-k (async=False, previous=False) ... passed [testing] Running with BF16, with top-k (async=False, previous=False) ... passed [testing] Running with BF16, without top-k (async=True, previous=False) ... passed [testing] Running with BF16, with top-k (async=True, previous=False) ... passed [testing] Running with BF16, without top-k (async=True, previous=False) ... passed [testing] Running with BF16, with top-k (async=True, previous=False) ... passed [testing] Running with BF16, without top-k (async=False, previous=True) ... passed [testing] Running with BF16, with top-k (async=False, previous=True) ... passed [testing] Running with BF16, without top-k (async=False, previous=True) ... passed [testing] Running with BF16, with top-k (async=False, previous=True) ... passed [testing] Running with BF16, without top-k (async=True, previous=True) ... passed [testing] Running with BF16, with top-k (async=True, previous=True) ... passed [testing] Running with BF16, without top-k (async=True, previous=True) ... passed [testing] Running with BF16, with top-k (async=True, previous=True) ... passed [tuning] SMs 24, NVL chunk 4: 97.48 GB/s (NVL), avg_t: 482.75 us [tuning] SMs 24, NVL chunk 6: 124.54 GB/s (NVL), avg_t: 377.87 us [tuning] SMs 24, NVL chunk 8: 139.62 GB/s (NVL), avg_t: 337.05 us [tuning] SMs 24, NVL chunk 10: 150.22 GB/s (NVL), avg_t: 313.26 us ... ``` `main` ``` python3 test_intranode.py --num-processes=2 --hidden 2880 [config] num_tokens=4096, hidden=2880, num_topk=8 [layout] Kernel performance: 0.034 ms [testing] Running with BF16, without top-k (async=False, previous=False) ...WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources. WARNING: destroy() was not called before DeepEP buffer destruction, which can leak resources. [rank0]:[W916 00:21:13.202446982 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator()) W0916 00:21:14.463000 1862405 torch/multiprocessing/spawn.py:169] Terminating process 1862487 via signal SIGTERM Traceback (most recent call last): File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 277, in <module> torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes) File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 340, in spawn return start_processes(fn, args, nprocs, join, daemon, start_method="spawn") File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 296, in start_processes while not context.join(): File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 215, in join raise ProcessRaisedException(msg, error_index, failed_process.pid) torch.multiprocessing.spawn.ProcessRaisedException: -- Process 1 terminated with the following error: Traceback (most recent call last): File "/home/varun/code/deps/deps-test/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 90, in _wrap fn(i, *args) File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 247, in test_loop test_main(args, i, local_rank, num_ranks, rank, buffer, group) File "/home/varun/code/deps/DeepEP/tests/test_intranode.py", line 160, in test_main assert calc_diff(check_x, ref_x) < 5e-6 AssertionError ``` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-27 21:03:58 +00:00
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#552
No description provided.