ByteDance-Seed / ByteDance-Seed/Triton-distributed
Kernel API reference signatures are out of sync with implementations
- 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
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