Skip to content

Commit 600a30f

Browse files
pre-commit-ci[bot]pre-commit-ci[bot]
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent 2769576 commit 600a30f

6 files changed

Lines changed: 44 additions & 38 deletions

File tree

‎tests/pytorch/test_fused_router.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -537,9 +537,9 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func
537537
torch.testing.assert_close(probs_byte, probs_bit, atol=0.0, rtol=0.0)
538538

539539
expected_shape = (num_tokens, (num_experts + 7) // 8)
540-
assert routing_map_bit.shape == expected_shape, (
541-
f"Bitmap shape {tuple(routing_map_bit.shape)} != {expected_shape}"
542-
)
540+
assert (
541+
routing_map_bit.shape == expected_shape
542+
), f"Bitmap shape {tuple(routing_map_bit.shape)} != {expected_shape}"
543543
assert routing_map_bit.dtype == torch.uint8
544544
assert routing_map_byte.dtype == torch.bool
545545

@@ -558,9 +558,7 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func
558558
[(128, 32, 4), (256, 128, 8), (256, 130, 8)],
559559
)
560560
@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"])
561-
def test_score_for_aux_loss_bitmap_vs_bytemap(
562-
dtype, num_tokens, num_experts, topk, score_function
563-
):
561+
def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function):
564562
"""fused_compute_score_for_moe_aux_loss: bitmap routing_map must equal
565563
LSB-packed bytemap; scores must be bit-identical across formats."""
566564
if topk >= num_experts:

‎transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,12 @@ namespace transformer_engine {
1717
namespace fused_router {
1818

1919
template <typename DataType, TopkFuncType TopkFunc = TopkFuncType::Naive>
20-
__global__ void fused_score_for_moe_aux_loss_forward_kernel(
21-
const DataType *logits, int num_tokens, int num_experts, int topk, int score_function,
22-
float *scores, uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
23-
CompType *intermediate_output) {
20+
__global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens,
21+
int num_experts, int topk,
22+
int score_function, float *scores,
23+
uint8_t *routing_map,
24+
NVTERoutingMapFormat routing_map_format,
25+
CompType *intermediate_output) {
2426
/***
2527
* Section: Global Variables/Addresses init
2628
* - Each warp is responsible for one token, and has own shared memory buffer.
@@ -175,8 +177,8 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(
175177
template <typename DataType>
176178
void fused_score_for_moe_aux_loss_forward_kernel_launcher(
177179
const DataType *logits, int num_tokens, int num_experts, int topk, int score_function,
178-
float *scores, uint8_t *routing_map, NVTERoutingMapFormat routing_map_format, CompType *intermediate_output,
179-
cudaStream_t stream) {
180+
float *scores, uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
181+
CompType *intermediate_output, cudaStream_t stream) {
180182
// Meta data for the kernel
181183
size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;
182184
size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;
@@ -213,7 +215,8 @@ void fused_score_for_moe_aux_loss_forward_kernel_launcher(
213215

214216
void fused_score_for_moe_aux_loss_forward(const Tensor &logits, int num_tokens, int num_experts,
215217
int topk, int score_function, Tensor &scores,
216-
Tensor &routing_map, NVTERoutingMapFormat routing_map_format,
218+
Tensor &routing_map,
219+
NVTERoutingMapFormat routing_map_format,
217220
Tensor &intermediate_output, cudaStream_t stream) {
218221
TE_ROUTER_PROBS_TYPE_SWITCH_ALL(
219222
logits.data.dtype, DataType,

‎transformer_engine/common/fused_router/fused_topk_with_score_function.cu‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ template <typename DataType, typename BiasType, TopkFuncType TopkFunc = TopkFunc
1919
__global__ void fused_topk_with_score_function_forward_kernel(
2020
const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,
2121
int num_groups, int group_topk, float scaling_factor, int score_function,
22-
const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
23-
CompType *intermediate_output) {
22+
const BiasType *expert_bias, DataType *probs, uint8_t *routing_map,
23+
NVTERoutingMapFormat routing_map_format, CompType *intermediate_output) {
2424
/***
2525
* Section: Global Variables/Addresses init
2626
* - Each warp is responsible for one token, and has own shared memory buffer.
@@ -282,8 +282,8 @@ template <typename DataType, typename BiasType>
282282
void fused_topk_with_score_function_forward_kernel_launcher(
283283
const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax,
284284
int num_groups, int group_topk, float scaling_factor, int score_function,
285-
const BiasType *expert_bias, DataType *probs, uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
286-
CompType *intermediate_output, cudaStream_t stream) {
285+
const BiasType *expert_bias, DataType *probs, uint8_t *routing_map,
286+
NVTERoutingMapFormat routing_map_format, CompType *intermediate_output, cudaStream_t stream) {
287287
size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;
288288
size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;
289289
size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // scores
@@ -328,8 +328,8 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens,
328328
int group_topk, float scaling_factor,
329329
int score_function, const Tensor expert_bias,
330330
Tensor probs, Tensor routing_map,
331-
NVTERoutingMapFormat routing_map_format, Tensor intermediate_output,
332-
cudaStream_t stream) {
331+
NVTERoutingMapFormat routing_map_format,
332+
Tensor intermediate_output, cudaStream_t stream) {
333333
TE_ROUTER_PROBS_TYPE_SWITCH_ALL(
334334
logits.data.dtype, DataType,
335335
TE_ROUTER_PROBS_TYPE_SWITCH_ALL(
@@ -346,8 +346,8 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens,
346346
template <typename DataType>
347347
__global__ void fused_topk_with_score_function_backward_kernel(
348348
// Inputs tensor
349-
const uint8_t *routing_map, NVTERoutingMapFormat routing_map_format, const CompType *intermediate_output,
350-
const DataType *grad_probs,
349+
const uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
350+
const CompType *intermediate_output, const DataType *grad_probs,
351351
// Other parameters
352352
int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor,
353353
int score_function,
@@ -528,9 +528,10 @@ __global__ void fused_topk_with_score_function_backward_kernel(
528528

529529
template <typename DataType>
530530
void fused_topk_with_score_function_backward_kernel_launcher(
531-
const uint8_t *routing_map, NVTERoutingMapFormat routing_map_format, const CompType *intermediate_output,
532-
const DataType *grad_probs, int num_tokens, int num_experts, int topk, bool use_pre_softmax,
533-
float scaling_factor, int score_function, DataType *grad_logits, cudaStream_t stream) {
531+
const uint8_t *routing_map, NVTERoutingMapFormat routing_map_format,
532+
const CompType *intermediate_output, const DataType *grad_probs, int num_tokens,
533+
int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function,
534+
DataType *grad_logits, cudaStream_t stream) {
534535
// Meta data for the kernel
535536
size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp;
536537
size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block;
@@ -550,7 +551,8 @@ void fused_topk_with_score_function_backward_kernel_launcher(
550551
NVTE_CHECK_CUDA(cudaGetLastError());
551552
}
552553

553-
void fused_topk_with_score_function_backward(const Tensor &routing_map, NVTERoutingMapFormat routing_map_format,
554+
void fused_topk_with_score_function_backward(const Tensor &routing_map,
555+
NVTERoutingMapFormat routing_map_format,
554556
const Tensor &intermediate_output,
555557
const Tensor &grad_probs, int num_tokens,
556558
int num_experts, int topk, bool use_pre_softmax,
@@ -584,11 +586,13 @@ void nvte_fused_topk_with_score_function_forward(
584586
*convertNVTETensorCheck(intermediate_output), stream);
585587
}
586588

587-
void nvte_fused_topk_with_score_function_backward(
588-
const NVTETensor routing_map, NVTERoutingMapFormat routing_map_format,
589-
const NVTETensor intermediate_output, const NVTETensor grad_probs, int num_tokens,
590-
int num_experts, int topk, int use_pre_softmax, float scaling_factor, int score_function,
591-
NVTETensor grad_logits, cudaStream_t stream) {
589+
void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map,
590+
NVTERoutingMapFormat routing_map_format,
591+
const NVTETensor intermediate_output,
592+
const NVTETensor grad_probs, int num_tokens,
593+
int num_experts, int topk, int use_pre_softmax,
594+
float scaling_factor, int score_function,
595+
NVTETensor grad_logits, cudaStream_t stream) {
592596
NVTE_API_CALL(nvte_fused_topk_with_score_function_backward);
593597
using namespace transformer_engine;
594598
fused_router::fused_topk_with_score_function_backward(

‎transformer_engine/common/include/transformer_engine/fused_router.h‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,11 +70,13 @@ void nvte_fused_topk_with_score_function_forward(
7070
* \param[out] grad_logits Gradient of logits.
7171
* \param[in] stream CUDA stream used for the operation.
7272
*/
73-
void nvte_fused_topk_with_score_function_backward(
74-
const NVTETensor routing_map, NVTERoutingMapFormat routing_map_format,
75-
const NVTETensor intermediate_output, const NVTETensor grad_probs, int num_tokens,
76-
int num_experts, int topk, int use_pre_softmax, float scaling_factor, int score_function,
77-
NVTETensor grad_logits, cudaStream_t stream);
73+
void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map,
74+
NVTERoutingMapFormat routing_map_format,
75+
const NVTETensor intermediate_output,
76+
const NVTETensor grad_probs, int num_tokens,
77+
int num_experts, int topk, int use_pre_softmax,
78+
float scaling_factor, int score_function,
79+
NVTETensor grad_logits, cudaStream_t stream);
7880

7981
/*! \brief Forward pass for computing scores/routing map for auxiliary loss.
8082
*

‎transformer_engine/pytorch/csrc/extensions/router.cpp‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@ static std::map<std::string, int> score_function_map = {
2020
static at::Tensor allocate_routing_map(int num_tokens, int num_experts,
2121
NVTERoutingMapFormat format) {
2222
if (format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) {
23-
return at::empty({num_tokens, (num_experts + 7) / 8},
24-
at::dtype(at::kByte).device(at::kCUDA));
23+
return at::empty({num_tokens, (num_experts + 7) / 8}, at::dtype(at::kByte).device(at::kCUDA));
2524
}
2625
return at::empty({num_tokens, num_experts}, at::dtype(at::kBool).device(at::kCUDA));
2726
}

‎transformer_engine/pytorch/router.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@
3535

3636

3737
def _validate_routing_map_format(
38-
routing_map_format: Union[str, "RoutingMapFormat", int]
38+
routing_map_format: Union[str, "RoutingMapFormat", int],
3939
) -> "RoutingMapFormat":
4040
"""Coerce user-supplied routing_map_format into the NVTERoutingMapFormat enum.
4141

0 commit comments

Comments
 (0)