Skip to content

Commit 889c161

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 b2928ac commit 889c161

6 files changed

Lines changed: 38 additions & 55 deletions

File tree

‎tests/jax/test_fused_router.py‎

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -615,14 +615,14 @@ def test_topk_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_func
615615
assert jnp.array_equal(probs_byte, probs_bit), "Probs must be identical across formats"
616616

617617
packed_expected = _bytemap_to_bitmap_u8(routing_map_byte)
618-
assert routing_map_bit.shape == (num_tokens, (num_experts + 7) // 8), (
619-
f"Bitmap shape {routing_map_bit.shape} != "
620-
f"({num_tokens}, {(num_experts + 7) // 8})"
621-
)
618+
assert routing_map_bit.shape == (
619+
num_tokens,
620+
(num_experts + 7) // 8,
621+
), f"Bitmap shape {routing_map_bit.shape} != ({num_tokens}, {(num_experts + 7) // 8})"
622622
assert routing_map_bit.dtype == jnp.uint8
623-
assert jnp.array_equal(routing_map_bit, packed_expected), (
624-
"Bitmap routing_map disagrees with np.packbits(bytemap, bitorder='little')"
625-
)
623+
assert jnp.array_equal(
624+
routing_map_bit, packed_expected
625+
), "Bitmap routing_map disagrees with np.packbits(bytemap, bitorder='little')"
626626

627627
# Backward parity: grad of probs.sum() must be identical for both formats.
628628
def loss_byte(logits_):
@@ -660,9 +660,7 @@ def loss_bit(logits_):
660660
)
661661
@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS)
662662
@pytest.mark.triton
663-
def test_score_for_aux_loss_bitmap_vs_bytemap(
664-
dtype, num_tokens, num_experts, topk, score_function
665-
):
663+
def test_score_for_aux_loss_bitmap_vs_bytemap(dtype, num_tokens, num_experts, topk, score_function):
666664
"""compute_aux_scores=True path: bitmap routing_map must equal LSB-packed
667665
bytemap; scores must be bitwise identical across formats."""
668666
from transformer_engine.jax.router import RoutingMapFormat
@@ -694,6 +692,6 @@ def test_score_for_aux_loss_bitmap_vs_bytemap(
694692
packed_expected = _bytemap_to_bitmap_u8(routing_map_byte)
695693
assert routing_map_bit.shape == (num_tokens, (num_experts + 7) // 8)
696694
assert routing_map_bit.dtype == jnp.uint8
697-
assert jnp.array_equal(routing_map_bit, packed_expected), (
698-
"Bitmap routing_map (aux-loss path) disagrees with packed bytemap"
699-
)
695+
assert jnp.array_equal(
696+
routing_map_bit, packed_expected
697+
), "Bitmap routing_map (aux-loss path) disagrees with packed bytemap"

‎transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,9 @@ 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(const DataType *logits, int num_tokens,
21-
int num_experts, int topk,
22-
int score_function, float *scores,
23-
uint8_t *routing_map,
24-
int routing_map_format,
25-
CompType *intermediate_output) {
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, int routing_map_format, CompType *intermediate_output) {
2623
/***
2724
* Section: Global Variables/Addresses init
2825
* - Each warp is responsible for one token, and has own shared memory buffer.
@@ -43,16 +40,14 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi
4340
const int bitmap_row_bytes = (num_experts + 7) / 8;
4441
uint32_t *bitmap_words_buf = nullptr;
4542
if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) {
46-
bitmap_words_buf =
47-
reinterpret_cast<uint32_t *>(topk_indices_buf + topk * num_token_per_block);
43+
bitmap_words_buf = reinterpret_cast<uint32_t *>(topk_indices_buf + topk * num_token_per_block);
4844
}
4945
// The address of buffers on the current warp
5046
CompType *local_logits = logits_buf + warp_id * num_experts;
5147
CompType *topk_logits = topk_logits_buf + warp_id * topk;
5248
int *topk_indices = topk_indices_buf + warp_id * topk;
53-
uint32_t *local_bitmap_words = (bitmap_words_buf != nullptr)
54-
? bitmap_words_buf + warp_id * bitmap_words_per_warp
55-
: nullptr;
49+
uint32_t *local_bitmap_words =
50+
(bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr;
5651

5752
/***
5853
* Section: Main Loop

‎transformer_engine/common/fused_router/fused_topk_with_score_function.cu‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -50,18 +50,16 @@ __global__ void fused_topk_with_score_function_forward_kernel(
5050
const int bitmap_row_bytes = (num_experts + 7) / 8;
5151
uint32_t *bitmap_words_buf = nullptr;
5252
if (routing_map_format == NVTE_ROUTING_MAP_FORMAT_BITMAP_U8) {
53-
bitmap_words_buf =
54-
reinterpret_cast<uint32_t *>(topk_indices_buf + topk * num_token_per_block);
53+
bitmap_words_buf = reinterpret_cast<uint32_t *>(topk_indices_buf + topk * num_token_per_block);
5554
}
5655
// The address of buffers on the current warp
5756
CompType *scores = scores_buf + warp_id * num_experts;
5857
CompType *topk_scores = topk_scores_buf + warp_id * topk;
5958
CompType *masked_scores = masked_scores_buf + warp_id * num_experts;
6059
CompType *group_scores = group_scores_buf + warp_id * num_groups;
6160
int *topk_indices = topk_indices_buf + warp_id * topk;
62-
uint32_t *local_bitmap_words = (bitmap_words_buf != nullptr)
63-
? bitmap_words_buf + warp_id * bitmap_words_per_warp
64-
: nullptr;
61+
uint32_t *local_bitmap_words =
62+
(bitmap_words_buf != nullptr) ? bitmap_words_buf + warp_id * bitmap_words_per_warp : nullptr;
6563

6664
/***
6765
* Section: Main Loop
@@ -574,8 +572,8 @@ void fused_topk_with_score_function_backward(const Tensor &routing_map, int rout
574572
void nvte_fused_topk_with_score_function_forward(
575573
const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax,
576574
int num_groups, int group_topk, float scaling_factor, int score_function,
577-
const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map,
578-
int routing_map_format, NVTETensor intermediate_output, cudaStream_t stream) {
575+
const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, int routing_map_format,
576+
NVTETensor intermediate_output, cudaStream_t stream) {
579577
NVTE_API_CALL(nvte_fused_topk_with_score_function_forward);
580578
using namespace transformer_engine;
581579
fused_router::fused_topk_with_score_function_forward(
@@ -586,13 +584,10 @@ void nvte_fused_topk_with_score_function_forward(
586584
*convertNVTETensorCheck(intermediate_output), stream);
587585
}
588586

589-
void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map,
590-
int 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) {
587+
void nvte_fused_topk_with_score_function_backward(
588+
const NVTETensor routing_map, int routing_map_format, const NVTETensor intermediate_output,
589+
const NVTETensor grad_probs, int num_tokens, int num_experts, int topk, int use_pre_softmax,
590+
float scaling_factor, int score_function, NVTETensor grad_logits, cudaStream_t stream) {
596591
NVTE_API_CALL(nvte_fused_topk_with_score_function_backward);
597592
using namespace transformer_engine;
598593
fused_router::fused_topk_with_score_function_backward(

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,8 @@ typedef enum {
5252
void nvte_fused_topk_with_score_function_forward(
5353
const NVTETensor logits, int num_tokens, int num_experts, int topk, int use_pre_softmax,
5454
int num_groups, int group_topk, float scaling_factor, int score_function,
55-
const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map,
56-
int routing_map_format, NVTETensor intermediate_output, cudaStream_t stream);
55+
const NVTETensor expert_bias, NVTETensor probs, NVTETensor routing_map, int routing_map_format,
56+
NVTETensor intermediate_output, cudaStream_t stream);
5757

5858
/*! \brief Backward pass for fused topk + softmax/sigmoid.
5959
*
@@ -70,13 +70,10 @@ 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(const NVTETensor routing_map,
74-
int 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);
73+
void nvte_fused_topk_with_score_function_backward(
74+
const NVTETensor routing_map, int routing_map_format, const NVTETensor intermediate_output,
75+
const NVTETensor grad_probs, int num_tokens, int num_experts, int topk, int use_pre_softmax,
76+
float scaling_factor, int score_function, NVTETensor grad_logits, cudaStream_t stream);
8077

8178
/*! \brief Forward pass for computing scores/routing map for auxiliary loss.
8279
*

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -139,17 +139,15 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI(
139139
grad_logits_tensor.data(), stream);
140140
} else {
141141
auto routing_map_dims = routing_map_buf.dimensions();
142-
auto routing_map_shape =
143-
std::vector<size_t>(routing_map_dims.begin(), routing_map_dims.end());
142+
auto routing_map_shape = std::vector<size_t>(routing_map_dims.begin(), routing_map_dims.end());
144143
auto routing_map_tensor =
145144
TensorWrapper(routing_map_buf.untyped_data(), routing_map_shape, DType::kByte);
146145

147146
nvte_fused_topk_with_score_function_backward(
148-
routing_map_tensor.data(), static_cast<int>(routing_map_format),
149-
intermediate_tensor.data(), grad_probs_tensor.data(), num_tokens, num_experts,
150-
static_cast<int>(topk), static_cast<int>(use_pre_softmax),
151-
static_cast<float>(scaling_factor), static_cast<int>(score_function),
152-
grad_logits_tensor.data(), stream);
147+
routing_map_tensor.data(), static_cast<int>(routing_map_format), intermediate_tensor.data(),
148+
grad_probs_tensor.data(), num_tokens, num_experts, static_cast<int>(topk),
149+
static_cast<int>(use_pre_softmax), static_cast<float>(scaling_factor),
150+
static_cast<int>(score_function), grad_logits_tensor.data(), stream);
153151
}
154152

155153
return ffi_with_cuda_error_check();

‎transformer_engine/jax/router.py‎

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

4444

4545
def _validate_routing_map_format(
46-
routing_map_format: Union[str, RoutingMapFormat, int]
46+
routing_map_format: Union[str, RoutingMapFormat, int],
4747
) -> RoutingMapFormat:
4848
"""Validate and convert routing_map_format to a RoutingMapFormat enum."""
4949
if isinstance(routing_map_format, RoutingMapFormat):

0 commit comments

Comments
 (0)