Fastllm Triton Ops

FastLLM Triton Ops

Overview

Use the existing Linear Triton prototype as the reference architecture: C++ decides whether an op is eligible, starts a local Python compiler server only on demand, asks it to emit a cached cubin plus metadata, then launches that cubin through the CUDA Driver API. Every Triton path must be environment-gated and must fall back to the original CUDA implementation on unsupported inputs or compile/launch failure.

Workflow

1. Inspect the existing CUDA op path in src/devices/cuda/cudadevice.cpp.

   • Find the op's Reshape, CanRun, Run, and lower-level helper functions.
   • Identify the exact tensor layout, dtype combinations, shape variables, optional bias/scale tensors, and output aliasing behavior.
   • Keep the original implementation as the fallback path.

2. Add or extend the Python compile path in tools/fastllm_triton_server.py.

   • Add a @triton.jit kernel for the new op.
   • Add <op>_cache_paths(payload) with a deterministic filename that includes op name, dtype/layout variants, SM arch, compile-time tile sizes, and feature flags.
   • Add compile_<op>(payload) that validates payload fields, compiles with ASTSource, writes .cubin, writes .json metadata, and returns the metadata.
   • Extend handle_compile(payload) by dispatching on payload["op"].
   • Keep /health and /compile stable; do not break existing "op": "linear" requests.

3. Add a CUDA launch wrapper in src/devices/cuda/fastllm-triton-cuda.cu.

   • Reuse LoadTritonKernel for cubin/module/function caching.
   • Add one extern "C" wrapper per op, for example FastllmCudaTriton<Op>(...).
   • Use FastllmCudaPrepareInput, FastllmCudaPrepareOutput, FastllmCudaFinishInput, and FastllmCudaFinishOutput when passing Data buffers.
   • Match the Triton kernel argument order exactly, including Triton's hidden global_scratch and profile_scratch pointer arguments when needed by AOT metadata.
   • Return false on load or launch failure so the C++ caller can fall back.

4. Declare the wrapper in include/devices/cuda/fastllm-cuda.cuh.

   • Keep the signature close to the CUDA fallback helper's shape arguments.
   • Pass metadata fields needed for launch, such as kernelName, shared, numWarps, and tile sizes.

5. Wire the op in src/devices/cuda/cudadevice.cpp.

   • Add a small metadata struct for the op's .json fields.
   • Reuse common helpers: CudaTritonCacheDir, CudaTritonDataTypeName, CudaTritonHttpRequest, and CudaTritonEnsureServer.
   • Add CudaTriton<Op>BaseName, CudaTritonRead<Op>Meta, CudaTritonRequest<Op>Kernel, and TryCudaTriton<Op>.
   • Gate with global FASTLLM_CUDA_TRITON; add an op-specific override such as FASTLLM_CUDA_TRITON_<OP>=0.
   • Validate device, pointer presence, dtype, layout, shape, arch, and feature constraints before requesting a kernel.
   • In the original op helper, call TryCudaTriton<Op>(...) immediately before the original CUDA implementation.

6. Update build wiring only when needed.

   • src/devices/cuda/fastllm-triton-cuda.cu is already in CMakeLists.txt.
   • If adding new files, update CMakeLists.txt and link dependencies without changing unrelated targets.