Initial commit to pass scale as Tensor for multi_tensor_scale op

[vasunvidia]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Overview

Greptile Summary

added tensor-based variant of multi_tensor_scale operation (multi_tensor_scale_tensor) to accept scale factor as a device tensor instead of host scalar

Key changes:

• implemented new ScalePtrFunctor CUDA kernel that dereferences scale from device memory
• added C API function nvte_multi_tensor_scale_tensor_cuda and PyTorch bindings
• exposed new function through Python API in optimizers module

Issues found:

• debug print statement left in scale.cpp:28 must be removed
• missing validation for scale tensor dtype (should be kFloat32) and size (should be 1 element) in scale.cu:219
• minor performance consideration: scale value read from device memory by all threads could use __ldg() or shared memory for optimization

Confidence Score: 3/5

• safe to merge after removing debug statement and adding validation checks
• implementation follows existing patterns and mirrors the scalar version correctly, but contains a debug print statement that must be removed and lacks input validation for the scale tensor that could cause runtime errors if invalid tensors are passed
• transformer_engine/common/multi_tensor/scale.cu and transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp need attention for the validation and debug statement issues

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp	added multi_tensor_scale_tensor_cuda function to accept scale as tensor; contains debug print statement that should be removed
transformer_engine/common/multi_tensor/scale.cu	implemented ScalePtrFunctor and multi_tensor_scale_tensor_cuda for tensor-based scaling; missing validation for scale tensor dtype and size

Sequence Diagram

sequenceDiagram
    participant Python as Python Code
    participant PyBind as PyTorch Binding
    participant Wrapper as C++ Wrapper
    participant CUDA as CUDA Kernel
    
    Python->>PyBind: multi_tensor_scale_tensor(chunk_size, noop_flag, tensor_lists, scale)
    PyBind->>Wrapper: multi_tensor_scale_tensor_cuda(...)
    Note over Wrapper: Convert PyTorch tensors<br/>to TE tensors
    Wrapper->>CUDA: nvte_multi_tensor_scale_tensor_cuda(scale as NVTETensor)
    Note over CUDA: Extract scale_tensor->data.dptr<br/>Cast to float*
    CUDA->>CUDA: multi_tensor_scale_tensor_cuda(scale as float*)
    Note over CUDA: Launch ScalePtrFunctor kernel<br/>with scale pointer
    CUDA->>CUDA: ScalePtrFunctor reads *scale_ptr<br/>Applies scaling to tensors
    CUDA-->>Python: Return scaled tensors

Loading

[greptile-apps]

_{6 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

debug print statement left in production code

Suggested change

std::cout << "multi_tensor_scale_cuda TENSOR\n";

[greptile-apps]

validate scale tensor dtype and size before casting. should check that:

• dtype is kFloat32
• numel() == 1 (single element tensor)

  Tensor *scale_tensor = convertNVTETensorCheck(scale);
  NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32, 
             "scale tensor must be float32, got ", to_string(scale_tensor->dtype()));
  NVTE_CHECK(scale_tensor->numel() == 1,
             "scale tensor must have exactly 1 element, got ", scale_tensor->numel());
  multi_tensor_scale::multi_tensor_scale_tensor_cuda(

[greptile-apps]

all threads in a block read from the same device memory location without synchronization - could cause redundant memory traffic but functionally correct. consider caching in shared memory or using __ldg() for read-only cache optimization

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}