deepseek-ai / deepseek-ai/TileKernels

dtype_to_str: unsupported dtype raises ValueError for fp16/float8_e5m2

Open
#4 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.8k
Forks
162
PR merge metrics
No merged PRs in 30d

Description

## Problem

In `tile_kernels/testing/bench.py`, the `dtype_to_str()` function is incomplete:

```python
def dtype_to_str(dtype: torch.dtype) -> str:
mapping = {
torch.float32: 'fp32',
torch.bfloat16: 'bf16',
torch.float8_e4m3fn: 'e4m3',
torch.int8: 'e2m1', # int8 represents FP4 e2m1 format
}
if dtype not in mapping:
raise ValueError(f'Unsupported dtype: {dtype}. Only fp32, bf16, e4m3, and int8(e2m1) are supported')
return mapping[dtype]
```

**Missing mappings:**
- `torch.float16` → `'fp16'`
- `torch.float8_e5m2` → `'e5m2'`

When `torch.float16` or `torch.float8_e5m2` is passed, it raises a `ValueError`.

## Context

TileKernels includes quantization kernels (`tile_kernels/quant/`) that use `float16` and `float8_e5m2` dtypes. The `dtype_to_str` function is used in benchmark output formatting (via `_format_value` → `make_param_id`). If these dtypes are used in a benchmark, the function will crash instead of producing a human-readable string.

## Expected Fix

Add the two missing mappings to the `mapping` dict and update the error message accordingly:

```python
mapping = {
torch.float32: 'fp32',
torch.float16: 'fp16',
torch.bfloat16: 'bf16',
torch.float8_e4m3fn: 'e4m3',
torch.float8_e5m2: 'e5m2',
torch.int8: 'e2m1',
}
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.