Skip to content

Commit 63409b8

Browse files
kawhil-amdkawhil-amd
committed
fix mem alloc within intranode ll
1 parent eec25c2 commit 63409b8

8 files changed

Lines changed: 21 additions & 14 deletions

File tree

‎.github/workflows/ci.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ jobs:
119119
- name: MORI-EP (intranode)
120120
run: |
121121
$CT exec -e PYTHONPATH=$GITHUB_WORKSPACE $CONTAINER bash -c "
122-
cd $GITHUB_WORKSPACE && timeout 420 pytest tests/python/ops/test_dispatch_combine_intranode.py -v
122+
cd $GITHUB_WORKSPACE && timeout 360 pytest tests/python/ops/test_dispatch_combine_intranode.py -v
123123
"
124124
125125
- name: MORI-EP (internode_v1)

‎.github/workflows/nightly.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,7 @@ jobs:
306306
- name: MORI-EP (intranode)
307307
run: |
308308
$CT exec -e PYTHONPATH=$GITHUB_WORKSPACE $CONTAINER bash -c "
309-
cd $GITHUB_WORKSPACE && timeout 420 pytest tests/python/ops/test_dispatch_combine_intranode.py -v
309+
cd $GITHUB_WORKSPACE && timeout 360 pytest tests/python/ops/test_dispatch_combine_intranode.py -v
310310
"
311311
312312
- name: MORI-EP (internode_v1)

‎include/mori/ops/dispatch_combine/dispatch_combine.hpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -245,23 +245,23 @@ class EpDispatchCombineHandle {
245245
int Fp8BlockwiseCombineScaleTypeSize() const { return fp8BlockwiseCombineScaleTypeSize; }
246246

247247
mori::application::SymmMemObjPtr GetShmemDispatchOutTokMemObj() const {
248-
if (config.kernelType == KernelType::IntraNode)
248+
if (config.kernelType == KernelType::IntraNode || config.kernelType == KernelType::IntraNodeLL)
249249
return std::get<ShmemBufsIntraNode>(shmemTokBufs).dispatchOut;
250250
if (config.kernelType == KernelType::InterNodeV1 ||
251251
config.kernelType == KernelType::InterNodeV1LL)
252252
return std::get<ShmemBufsInterNodeV1>(shmemTokBufs).dispatchOut;
253253
return std::get<ShmemBufsInterNode>(shmemTokBufs).dispatchOut;
254254
}
255255
mori::application::SymmMemObjPtr GetShmemCombineOutTokMemObj() const {
256-
if (config.kernelType == KernelType::IntraNode)
256+
if (config.kernelType == KernelType::IntraNode || config.kernelType == KernelType::IntraNodeLL)
257257
return std::get<ShmemBufsIntraNode>(shmemTokBufs).combineOut;
258258
if (config.kernelType == KernelType::InterNodeV1 ||
259259
config.kernelType == KernelType::InterNodeV1LL)
260260
return std::get<ShmemBufsInterNodeV1>(shmemTokBufs).combineOut;
261261
return std::get<ShmemBufsInterNode>(shmemTokBufs).combineOut;
262262
}
263263
mori::application::SymmMemObjPtr GetShmemCombineInpTokMemObj() const {
264-
if (config.kernelType == KernelType::IntraNode)
264+
if (config.kernelType == KernelType::IntraNode || config.kernelType == KernelType::IntraNodeLL)
265265
return std::get<ShmemBufsIntraNode>(shmemTokBufs).combineInp;
266266
if (config.kernelType == KernelType::InterNodeV1 ||
267267
config.kernelType == KernelType::InterNodeV1LL)

‎python/mori/ops/dispatch_combine.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -763,9 +763,12 @@ def combine(
763763
shared_mem = self._combine_shared_mem(actual_wpb)
764764

765765
if quant_type == EpDispatchCombineQuantType.Fp8BlockwiseQuant:
766-
if kt != EpDispatchCombineKernelType.IntraNode.value:
766+
if kt not in (
767+
EpDispatchCombineKernelType.IntraNode.value,
768+
EpDispatchCombineKernelType.IntraNodeLL.value,
769+
):
767770
raise ValueError(
768-
"Fp8BlockwiseQuant currently only supports IntraNode combine"
771+
"Fp8BlockwiseQuant currently only supports IntraNode/IntraNodeLL combine"
769772
)
770773
if sfx != "bf16":
771774
raise ValueError(f"Fp8BlockwiseQuant only supports bf16, got {sfx}")
@@ -1379,6 +1382,7 @@ def get_dispatch_src_token_pos(self):
13791382

13801383
if self.config.kernel_type.value in (
13811384
EpDispatchCombineKernelType.IntraNode.value,
1385+
EpDispatchCombineKernelType.IntraNodeLL.value,
13821386
EpDispatchCombineKernelType.InterNodeV1.value,
13831387
EpDispatchCombineKernelType.InterNodeV1LL.value,
13841388
EpDispatchCombineKernelType.AsyncLL.value,

‎python/mori/ops/tuning_configs/gfx950_mi355x_IntraNodeLL_ep8_dispatch.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
"version": "1.0",
33
"gpu_arch": "gfx950",
44
"gpu_model": "mi355x",
5-
"kernel_type": "IntraNode",
5+
"kernel_type": "IntraNodeLL",
66
"ep_size": 8,
77
"phase": "dispatch",
88
"rules": [

‎src/ops/dispatch_combine/convert.hpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,8 @@ template <typename T>
184184
__device__ inline void InvokeConvertDispatchOutput(const EpDispatchCombineArgs<T>& args, int myPe) {
185185
ConvertDispatchOutputArgs convArgs{};
186186
convArgs.config = args.config;
187-
if (args.config.kernelType == KernelType::IntraNode) {
187+
if (args.config.kernelType == KernelType::IntraNode ||
188+
args.config.kernelType == KernelType::IntraNodeLL) {
188189
convArgs.dispatchOutX = args.intraNodeTokBufs.dispatchOut->template GetAs<T*>(myPe);
189190
} else if (args.config.kernelType == KernelType::InterNodeV1 ||
190191
args.config.kernelType == KernelType::InterNodeV1LL) {
@@ -218,7 +219,8 @@ __device__ inline void InvokeConvertCombineInput(const EpDispatchCombineArgs<T>&
218219
convArgs.combineInput = nullptr;
219220
convArgs.dispTokToEpSlotMap = args.dispTokToEpSlotMap;
220221
convArgs.packedRecvCount = args.standardPackedRecvCount;
221-
if (args.config.kernelType == KernelType::IntraNode) {
222+
if (args.config.kernelType == KernelType::IntraNode ||
223+
args.config.kernelType == KernelType::IntraNodeLL) {
222224
convArgs.shmemCombineInpTokMemObj = args.intraNodeTokBufs.combineInp;
223225
} else if (args.config.kernelType == KernelType::InterNodeV1 ||
224226
args.config.kernelType == KernelType::InterNodeV1LL) {

‎src/ops/dispatch_combine/dispatch_combine.cpp‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ void EpDispatchCombineHandle::InitializeShmemBuf() {
198198
config.WeightBytes() + config.SrcTokenIdBytes() + blockwiseScaleBytes);
199199
}
200200

201-
if (config.kernelType == KernelType::IntraNode) {
201+
if (config.kernelType == KernelType::IntraNode || config.kernelType == KernelType::IntraNodeLL) {
202202
auto& bufs = shmemTokBufs.emplace<ShmemBufsIntraNode>();
203203
bufs.combineInp = ShmemMallocAndReturnMemObjPtr(maxStagingSize, hipDeviceMallocUncached);
204204
bufs.dispatchOut = ShmemMallocAndReturnMemObjPtr(dispatchOutSize, hipDeviceMallocUncached);
@@ -271,7 +271,7 @@ void EpDispatchCombineHandle::InitializeShmemBuf() {
271271
}
272272

273273
void EpDispatchCombineHandle::FinalizeShmemBuf() {
274-
if (config.kernelType == KernelType::IntraNode) {
274+
if (config.kernelType == KernelType::IntraNode || config.kernelType == KernelType::IntraNodeLL) {
275275
auto& bufs = std::get<ShmemBufsIntraNode>(shmemTokBufs);
276276
ShmemFree(bufs.dispatchOut->localPtr);
277277
ShmemFree(bufs.combineInp->localPtr);
@@ -463,7 +463,8 @@ EpDispatchCombineArgsRaw GetEpDispatchCombineArgsRaw(const EpDispatchCombineHand
463463
args.scalesBuf = handle.scalesBuf;
464464
args.destPeTokenCounter = handle.destPeTokenCounter;
465465
args.localPeTokenCounter = handle.localPeTokenCounter;
466-
if (handle.config.kernelType == KernelType::IntraNode) {
466+
if (handle.config.kernelType == KernelType::IntraNode ||
467+
handle.config.kernelType == KernelType::IntraNodeLL) {
467468
args.intraNodeTokBufs = std::get<ShmemBufsIntraNode>(handle.shmemTokBufs);
468469
} else if (handle.config.kernelType == KernelType::InterNodeV1 ||
469470
handle.config.kernelType == KernelType::InterNodeV1LL) {

‎src/ops/kernels/ep_intranode.hip‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ WRAP_ALL_TYPES_BOOL(EpDispatchIntraNodeKernel, _stdmoe, true)
88

99
// Dispatch Low Lantency (SMALL TOKEN)
1010
WRAP_ALL_TYPES_BOOL(EpDispatchIntraNodeLLKernel, , false)
11-
WRAP_ALL_TYPES_BOOL(EpDispatchIntraNodeLLKernel, _stdmoe, false)
11+
WRAP_ALL_TYPES_BOOL(EpDispatchIntraNodeLLKernel, _stdmoe, true)
1212

1313
// Combine
1414
WRAP_ALL_TYPES_BOOL3(EpCombineIntraNodeKernel, _p2p, true, false, false)

0 commit comments

Comments
 (0)