ByteDance-Seed / ByteDance-Seed/Triton-distributed

Kernel API reference signatures are out of sync with implementations

Open
#194 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.5k
Forks
172
PR merge metrics
No merged PRs in 30d

Description

### Description

Several generated-looking kernel API reference signatures no longer match the functions at current main commit `8260bc34398c2b8f36dc840fd22f741ca9294584`. The stale signatures are also visible on the current ReadTheDocs site.

| Documentation | Documented parameters | Current implementation |
| --- | --- | --- |
| `docs/kernels/nvidia/allgather.rst` — `fast_allgather` | `ctx, buffer` | `symm_buffer, ctx, rank, node, num_ranks, num_nodes, mode` |
| `docs/kernels/nvidia/allgather_gemm.rst` — `create_ag_gemm_context` | `local_tensor, weight, rank, num_ranks, max_M, BLOCK_M, BLOCK_N, BLOCK_K, stages` | `max_M, N, K, dtype, rank, num_ranks, num_local_ranks, ag_intranode_stream, ag_internode_stream` |
| `docs/kernels/nvidia/gemm_reduce_scatter.rst` — `create_gemm_rs_context` | ends at `rs_stream` | additionally accepts `reduce_st` and `tp_group` |
| `docs/kernels/amd/all_gather_gemm.rst` — `create_ag_gemm_intra_node_context` | uses `world_size`, `M_PER_CHUNK=256` | uses `num_ranks`, adds `ag_streams`, `use_copy_kernel`, `comm_sms`, and defaults `M_PER_CHUNK=1024` |
| `docs/kernels/amd/gemm_reduce_scatter.rst` — `create_gemm_rs_intra_node_context` | uses `world_size` | uses `num_ranks` |

For example, the live NVIDIA AG-GEMM page currently recommends a `create_ag_gemm_context(A, B, rank, world_size, ...)` call that cannot bind to the current function: https://triton-distributed.readthedocs.io/en/latest/kernels/nvidia/allgather_gemm.html

### Expected behavior

Each API reference signature and example should match the current callable, including parameter names, ordering, and defaults. The prose descriptions also need to be updated where the meaning changed; this is why this report does not propose an automated signature-only patch.

Contributor guide

Open the contributing guide

Research direction

Start with the five listed .rst files under docs/kernels/nvidia and docs/kernels/amd, then compare each documented callable with its current implementation at commit 8260bc34398c2b8f36dc840fd22f741ca9294584. Update signatures, defaults, examples, and prose to match the callable behavior, and build the documentation to verify the API references render correctly.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, documentation
Issue type
Documentation
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.