Refine documentation

Signed-off-by: Barry Kang <[email protected]>

Author: Barry-Delaney

tensorrt_llm/quantization/utils/fp8_utils.py

@@ -555,7 +555,18 @@ def _transpose_kernel(input_ptr, output_ptr, M, N, stride_in_m, stride_in_n,
 
 
 def masked_transpose(input: torch.Tensor, n_available: int) -> torch.Tensor:
+    """
+    Perform a masked transpose operation on a 2D tensor.
+
+    Args:
+        input: Input tensor of shape (M, N)
+        n_available: Number of columns to transpose (must be <= N)
+
+    Returns:
+        Transposed tensor of shape (n_available, M)
+    """
     M, N = input.shape
+    assert n_available <= N, "n_available must be less than or equal to N"
     BLOCK_SIZE = 32
     output = torch.empty((n_available, M),
                          dtype=input.dtype,