Add Pythonic API wrapper with automatic workspace management

[haix716]

Problem

The C++ function requires an externally-managed workspace tensor, but the Python flash_kda/__init__.py is essentially empty — it just imports the C extension with no Pythonic wrapper.

Users currently need to:

size = flash_kda.get_workspace_size(B, H, T, CHUNK)
workspace = torch.empty(size, dtype=torch.uint8, device='cuda')
flash_kda.fwd(q, k, v, g, beta, scale, out, workspace=workspace, ...)

Recommendation

Add a Pythonic wrapper that handles workspace allocation automatically:

def fwd(q, k, v, g, beta, scale, out=None, **kwargs):
    """Forward pass for Kimi Delta Attention.
    
    Args:
        q: [B, H, T, D] query tensor (bf16/fp16)
        k: [B, H, T, D] key tensor
        v: [B, H, T, D] value tensor
        g: [B, H, T] gate logits (fp32)
        beta: [B, H, T] beta modulation (fp32)
        scale: float, typically 1/sqrt(D)
    
    Returns:
        out: [B, H, T, D] output tensor
    """
    if out is None:
        out = torch.empty_like(q)
    size = get_workspace_size(q.shape[0], q.shape[1], q.shape[2], 16)
    workspace = torch.empty(size, dtype=torch.uint8, device=q.device)
    _fwd_impl(q, k, v, g, beta, scale, out, workspace, **kwargs)
    return out

Also add docstrings with expected shapes and dtypes, and define __all__ for clean from flash_kda import * behavior.

Impact

Usability. Makes the API accessible to users who don't want to manage CUDA workspace allocation manually. Reduces boilerplate and potential for OOM errors from incorrect size calculations.