deepseek-ai / deepseek-ai/FlashMLA

[CUDA/CUTLASS] Improvements for varlen option derivation, input validation, can_implement errors, and workspace handling

Open
#161 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
12.9k
Forks
1.2k
Avg merge
4h 20m
Merged PRs (30d)
2

Description

## Summary

In the current FMHA / attention-related CUTLASS wrapper (and MLA phase1 path), the following areas appear to have potential risks in terms of correctness, robustness, and debuggability:

1. Varlen option derivation is logically unsafe
2. Input tensor validation is too weak
3. `op.can_implement()` failure messages are not very informative
4. Workspace handling is left as TODO and may cause runtime issues

---

## (1) Varlen option derivation is logically unsafe (important)

**Observed in:**
`get_options()` where:

```cpp
options.q = q.size(0) / options.b;
options.k = k.size(0) / options.b;
```

This logic only makes sense for **dense batched (fixed sequence length)** inputs.

In **varlen** mode, `q.size(0)` and `k.size(0)` often represent the **total token counts** (e.g. last element of prefix-sum).
Dividing by `b` produces an average-like value that may not match the actual maximum sequence length.

### Why this is risky

Varlen kernels usually rely on:

* `max_seqlen_q`
* `max_seqlen_kv`

Using `total_tokens / b` may:

* Underestimate the required shape
* Break stride or shape assumptions
* Lead to `can_implement` failures or undefined behavior

### Example

```
b = 2
seqlen = [1, 4096]
total_tokens = 4097

options.q = total_tokens / b = 2048
actual max_seqlen = 4096
```

This creates an inconsistent problem shape.

### Suggested fix

When `kIsVarlen` (or equivalent flag) is set:

* Derive `options.q` and `options.k` from `max_seqlen_q` / `max_seqlen_kv`, or
* Leave them as placeholders and finalize the shape only in `initialize_varlen()`

---

## (2) Input tensor validation is too weak

**Observed near kernel argument preparation / launch path.**

Current checks only validate some stride conditions.
To avoid silent miscompilation or memory issues, stronger validation is recommended.

### Suggested checks

* All tensors are CUDA tensors
* All tensors use expected dtype (bf16/fp16/fp8 as required)
* All tensors are on the same device
* Layout/stride matches what the kernel supports

Example:

```cpp
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda() && o.is_cuda(),
"All tensors must be CUDA tensors");

TORCH_CHECK(q.scalar_type() == at::kBFloat16,
"Unsupported dtype for q");

TORCH_CHECK(q.device() == k.device() &&
k.device() == v.device() &&
v.device() == o.device(),
"All tensors must be on the same device");

TORCH_CHECK(q.is_contiguous() || is_supported_strided(q),
"Unsupported layout/stride for q");
```

---

## (3) Improve `can_implement` error handling

**Observed in:**

```cpp
CUTLASS_CHECK(op.can_implement(arguments));
```

If this fails, the user often gets an unhelpful error without context.

### Suggested improvement

Capture the status and provide a descriptive error:

```cpp
auto status = op.can_implement(arguments);

TORCH_CHECK(status == cutlass::Status::kSuccess,
"can_implement failed: ",
cutlass_status_to_string(status),
" | q shape=", ...,
" | q stride=", ...,
" | k shape=", ...,
" | dtype=", ...);
```

This significantly improves debuggability.

---

## (4) Workspace TODO may cause runtime or portability issues

**Observed where workspace logic is currently commented out or marked as TODO.**

Depending on the CUTLASS configuration or architecture, workspace may be required.
Ignoring it can cause:

* Runtime launch failures
* Architecture-dependent behavior
* Hard-to-diagnose issues

### Suggested fix

Query and validate workspace size:

```cpp
size_t ws = Operation::get_workspace_size(arguments);

TORCH_CHECK(workspace_nbytes >= ws,
"workspace too small: need ", ws, " bytes");

CUTLASS_CHECK(op.initialize(arguments, workspace_ptr));
CUTLASS_CHECK(op.run(stream));
```

---

## Expected benefits

* More stable varlen behavior
* Early detection of invalid inputs
* Much clearer debugging when `can_implement` fails
* Better portability across architectures and configurations

---

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.