Write a custom CUDA kernel to optimize `torch.logsumexp`.

The operation computes `log(sum(exp(input)))` along a given dimension. A naive implementation is numerically unstable and inefficient.

**Problem Analysis:**
1.  **Numerical Instability**: The `exp(input)` term can easily overflow standard floating-point types for moderately large inputs, leading to `inf` and incorrect results. This is a critical correctness issue.
2.  **Memory Bandwidth Bottleneck**: A sequential implementation (`exp` -> `sum` -> `log`) would materialize a large intermediate tensor for `exp(input)`, causing a severe performance penalty due to the unnecessary write and read from global memory.

**Optimization Strategy: Fused Two-Pass Parallel Reduction with the Log-Sum-Exp Trick**

The strategy is to implement the entire operation within a single, fully-fused CUDA kernel that addresses both numerical stability and performance by combining the "Log-Sum-Exp trick" with an efficient parallel reduction algorithm.

1.  **Numerical Stability (Log-Sum-Exp Trick)**: The kernel avoids overflow by implementing the stable formula: `m + log(sum(exp(x_i - m)))`, where `m` is the maximum value of the input vector `x`. This ensures the argument to `exp` is always non-positive.

2.  **Fused Two-Pass Reduction**: The kernel performs the stable computation in two reduction passes within a single launch, all orchestrated using fast **shared memory**:
    *   **Pass 1 (Find Max)**: Threads in a block collaboratively find the maximum value `m` of their assigned input slice using a parallel reduction.
    *   **Pass 2 (Sum Exp Diff)**: After syncing and reading the shared `m`, the threads perform a second parallel reduction to compute the sum of `exp(x_i - m)`.

3.  **Block-per-Slice Parallelism**: The kernel is launched with one thread block for each 1D slice of the input tensor along the specified `dim`.

4.  **Finalization**: The first thread of each block computes the final `m + log(sum)` and writes the single scalar result directly to the output tensor. This fuses all steps into a single memory pass over the input data, maximizing both performance and numerical robustness.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn

BATCH_SIZE = 4096
FEATURES = 2048
SHAPE = (BATCH_SIZE, FEATURES)
DIM = 1

class Model(nn.Module):
    """
    使用 PyTorch 内置的 torch.logsumexp 作为基准模型。
    """
    def __init__(self, dim, keepdim=False):
        super(Model, self).__init__()
        self.dim = dim
        self.keepdim = keepdim
    
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return torch.logsumexp(input_tensor, dim=self.dim, keepdim=self.keepdim)

def get_inputs():
    # 扩大值的范围以测试数值稳定性
    input_tensor = torch.randn(SHAPE, dtype=torch.float32) * 10.0
    return [input_tensor.contiguous()]

def get_init_inputs():
    return [DIM, False] # keepdim=False