Skip to content

Conversation

@yingxudeng
Copy link
Collaborator

No description provided.

#elif defined(USE_CUDA)
cuda::act_and_mul(params.output, params.input, params.act_mode);
#else
LOG(FATAL) << "active not implemented";
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove torch::Tensor active_tensor(ActivationParams& params) and add params.output = npu::active(params.input, params.act_mode) here for npu device.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

auto output = torch::empty(
    {batch_size,
     intermediate_size_ / parallel_args_.tp_group_->world_size()},
    gate_up.options());

This is a good modification. However, as described, the current code's output still allocates space preemptively. For NPU operators, they typically allocate their own space and return the result. This unavoidable difference still forces the external calling code to use an #if block to skip space allocation specifically for the NPU case.

To standardize the external calling code, I personally recommend aligning with the NPU's behavior: allocate the space within the operator wrapper/layer and then return it. This approach allows for a unified code structure for all external calls.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so don't add active_tensor and fused_layernorm_tensor these two func in ops_api.h, because no other platform will use such api.
put they in npu_ops_api.h and call them directly in npu layer.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image Regarding the code snippet above: if we implement the changes as suggested, we would need to introduce #if directives here to skip memory allocation, since the NPU operator handles this internally.

Could we instead consider moving the memory allocation logic for MLU and CUDA into their respective kernel wrappers? This would make the behavior more similar to PyTorch and allow us to unify the calling code here.

(PS: I haven't modified the CUDA or MLU code yet.)

#endif
}

torch::Tensor fused_layernorm_tensor(FusedLayerNormParams& params) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as above

Copy link
Collaborator Author

@yingxudeng yingxudeng Dec 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similar to the previous comment.

// Must be less than or equal to rope_seqlen if not using discrete
// position_ids.
int64_t max_query_len;
torch::Tensor positions;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

std::optional<torch::Tensor> position_ids already exists.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

During the implementation, I noticed that position_ids are set to empty during the prefill stage, so I initially added position. However, I see that the latest CUDA code addresses the same issue using a different approach. To ensure consistency, I plan to align my implementation with the CUDA method.

@yingxudeng yingxudeng force-pushed the feat/npu_backend_torch_2_kernels branch from 7485463 to 9711f41 Compare December 2, 2025 15:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants