Functions — github.com/FlashSampling/FlashSampling

hub / github.com/FlashSampling/FlashSampling / functions

Functions315 in github.com/FlashSampling/FlashSampling

⨍ Functions 315◇ Types & classes 25 ↳ Endpoints 9

↓ 29 callersMethodrank0_print

(self, *args, **kwargs)

src/fused_mm_sampling/tp_info.py:23

↓ 12 callersFunctionmake_app

()

src/fused_mm_sampling/modal_lib/utils.py:31

↓ 8 callersFunctionbench_fn

(label, reps, warmup_reps, fn, *args)

src/fused_mm_sampling/persistent_matmul.py:722

↓ 8 callersFunctionfn

()

src/fused_mm_sampling/bench/speed_test.py:105

↓ 8 callersMethodprepare

(self)

src/fused_mm_sampling/core.py:645

↓ 8 callersMethodsample

(self, **kwargs)

src/fused_mm_sampling/core.py:648

↓ 7 callersFunctionfused_mm_sample_triton

( weights: torch.Tensor, # [V_local, D] (may be a TP shard) hidden_states: torch.Tensor, # [n_hidden

src/fused_mm_sampling/core.py:220

↓ 7 callersFunctionget_sampler

(provider: str, weights: torch.Tensor)

src/fused_mm_sampling/core.py:730

↓ 7 callersFunctionset_volume_caches

Point cache env vars to the Modal volume and enable Triton autotune logging. XDG_CACHE_HOME: used by flashinfer, torch.compile, etc. TRITON_C

src/fused_mm_sampling/modal_lib/utils.py:122

↓ 6 callersFunction_compute_pid

(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS)

src/fused_mm_sampling/persistent_matmul.py:286

↓ 6 callersMethodmake_tp

(self)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:86

↓ 5 callersFunctionis_cuda

()

src/fused_mm_sampling/persistent_matmul.py:13

↓ 5 callersMethodis_rank0

(self)

src/fused_mm_sampling/tp_info.py:20

↓ 5 callersFunctionmake_synthetic_inputs

Build weights and hidden_states that produce known logits. Creates up to two hidden states: one with ascending logits (favors high token indi

src/fused_mm_sampling/testing.py:27

↓ 5 callersFunctionrun_maybe_distributed

Run fn(*args) in a single process or spawn n_procs distributed workers. fn and args must be picklable (top-level functions, dataclasses, etc.)

src/fused_mm_sampling/tp_info.py:35

↓ 5 callersFunctionrun_test

(expect, fn, a, b, label, enabled=True)

src/fused_mm_sampling/persistent_matmul.py:771

↓ 5 callersFunctionsynchronize

(is_distributed: bool)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:252

↓ 4 callersFunction_allgather_logits

All-gather local logits along the vocab dimension to reconstruct [H, V_global].

src/fused_mm_sampling/core.py:77

↓ 4 callersFunction_provider_palette

Return a color mapping for the providers present in the data.

benchmarking/plot-triton-bench.py:61

↓ 4 callersFunctionargsort

Sort `x` along `dim` and apply the same permutation to `ids`. Args: x: values tensor (e.g. logits) ids: index tensor, same shape

src/fused_mm_sampling/tl_argsort.py:15

↓ 4 callersMethodfrom_world

(cls)

src/fused_mm_sampling/tp_info.py:28

↓ 4 callersFunctionmake_image

()

src/fused_mm_sampling/modal_lib/utils.py:35

↓ 4 callersMethodproviders

(self)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:91

↓ 4 callersFunctionshard_weights

Shard weights along vocab dim (same as vLLM's VocabParallelEmbedding.weight_loader).

src/fused_mm_sampling/testing.py:94

↓ 3 callersFunction_find_first

(lines: list[str], pattern: str)

benchmarking/insert_proton_records.py:181

↓ 3 callersFunction_find_first_after

(lines: list[str], pattern: str, after: int | None)

benchmarking/insert_proton_records.py:188

↓ 3 callersFunction_prepare_case

Set up sampler and fn for a benchmark case.

src/fused_mm_sampling/bench/speed_test.py:97

↓ 3 callersFunctionas_case

(args: Args, name: str)

src/fused_mm_sampling/bench/speed_test.py:35

↓ 3 callersFunctionclear_l2_cache

(cache: torch.Tensor)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:248

↓ 3 callersFunctiondevice_blas_name

()

src/fused_mm_sampling/persistent_matmul.py:35

↓ 3 callersFunctionget_gpu_name

()

src/fused_mm_sampling/bench/sys_metadata.py:18

↓ 3 callersFunctionis_hopper

()

src/fused_mm_sampling/persistent_matmul.py:43

↓ 3 callersFunctionload_variant

(model_dir: Path, variant_key: str)

benchmarking/vllm/plot_tpot.py:62

↓ 3 callersFunctionparse_ncu_csv

Parse a single NCU CSV file into a DataFrame. NCU CSV files contain non-CSV lines (==PROF== messages, Python stdout). Only lines starting wit

benchmarking/parse_ncu_sweep.py:75

↓ 3 callersFunctionread_triton_bench_csv

Read a Triton benchmark CSV, stripping the ' (Time (ms))' column suffix.

benchmarking/plot_lib.py:9

↓ 3 callersFunctionrun_speed_test

Run a speed test for a given set of arguments.

src/fused_mm_sampling/bench/speed_test.py:279

↓ 3 callersFunctiontp_from_dir

Parse the TP value out of the results dir name (e.g. ...-tp2). Defaults to 1.

benchmarking/vllm/plot_tpot.py:43

↓ 2 callersFunction_compute_tile_pid

Compute pid_v, pid_h from tile_id using grouped ordering for L2 cache efficiency.

src/fused_mm_sampling/core.py:391

↓ 2 callersFunction_find_last

(lines: list[str], pattern: str)

benchmarking/insert_proton_records.py:196

↓ 2 callersFunction_indicator

(n_dims: core.constexpr, j: core.constexpr)

src/fused_mm_sampling/tl_argsort.py:107

↓ 2 callersFunction_local_reduce

Reduce across V-tiles (dim=1) on this rank and adjust to global vocab indices.

src/fused_mm_sampling/core.py:344

↓ 2 callersFunction_log2

(i)

src/fused_mm_sampling/tl_argsort.py:94

↓ 2 callersFunction_plot_scatter_line

Plot median lines with scatter dots for one or more series. Each entry in *series* is (label, dataframe). Colors are looked up from VARIANT_C

benchmarking/vllm/plot_tpot.py:132

↓ 2 callersFunction_print_gpu_topology

()

src/fused_mm_sampling/tp_info.py:110

↓ 2 callersFunction_print_versions

(label: str)

src/fused_mm_sampling/modal_lib/modal_versions.py:10

↓ 2 callersFunction_provider_markers

Return a marker mapping for the providers present in the data.

benchmarking/plot-triton-bench.py:67

↓ 2 callersFunction_sm_version

Return the SM version string (e.g. '86' for RTX 3090).

src/fused_mm_sampling/cuda_impl.py:35

↓ 2 callersFunctionadd_library_code

Add the fused_mm_sampling source and benchmarking scripts to the image, then pip-install the package so subprocess-based tools (e.g. ncu) can impo

src/fused_mm_sampling/modal_lib/utils.py:136

↓ 2 callersFunctionassert_sampling_distribution

Verify that a sampler produces the correct distribution. Uses synthetic inputs with known logit vectors (ascending and/or descending), draws

src/fused_mm_sampling/testing.py:105

↓ 2 callersFunctionbench_cuda_events

Time a callable using CUDA events with L2 cache flushing. Returns per-iteration times in milliseconds. Uses fixed iteration counts (not adap

src/fused_mm_sampling/bench/triton_benchmark_lib.py:205

↓ 2 callersFunctionbench_cupti

Time a callable using FlashInfer's CUPTI-based bench_gpu_time. Returns per-iteration times in milliseconds.

src/fused_mm_sampling/bench/triton_benchmark_lib.py:188

↓ 2 callersMethodcompute_logits

( self, hidden_states: torch.Tensor, # [n_hidden_states, D] )

src/fused_mm_sampling/core.py:711

↓ 2 callersFunctioncreate_l2_cache

Allocate a 256 MB buffer for L2 cache flushing. Follows the same strategy as triton.testing.do_bench: the buffer is zeroed before each timed

src/fused_mm_sampling/bench/triton_benchmark_lib.py:237

↓ 2 callersFunctionflashinfer_create_logits_and_indices

( weights: torch.Tensor, # [V, D] hidden_states: torch.Tensor, # [n_hidden_states, D] num_sample

src/fused_mm_sampling/core.py:820

↓ 2 callersFunctionfn

()

src/fused_mm_sampling/bench/triton_benchmark_lib.py:174

↓ 2 callersMethodfrom_weights

( cls, weights: torch.Tensor, # [V, D] epsilon: float = 0.2, )

src/fused_mm_sampling/core.py:670

↓ 2 callersMethodmake_fn_kwargs

This function can be slow because it allocates tensors.

src/fused_mm_sampling/bench/speed_test.py:75

↓ 2 callersFunctionmatmul

(a, b, activation="")

src/fused_mm_sampling/tl_matmul.py:135

↓ 2 callersFunctionmodel_bytes

Minimum bytes transferred for fused matmul+sampling (read W + read X + write indices).

benchmarking/plot-triton-bench.py:236

↓ 2 callersFunctionparse_chrome_trace

Parse a chrome trace and return aggregate scope durations (in cycles).

benchmarking/parse_proton_intrakernel.py:19

↓ 2 callersFunctionplot

( df: pd.DataFrame, out_path: Path, fmt: str, y_col: str, y_label: str, y_cap: float |

benchmarking/plot_bsz_sweep_runtime.py:109

↓ 2 callersFunctionread_summary

(results_dir: Path, variant_key: str)

benchmarking/vllm/collect_results.py:32

↓ 2 callersFunctionresolve_model_dirs

Find all model directories matching the model name, including -trialN suffixes.

benchmarking/vllm/plot_tpot.py:49

↓ 2 callersFunctionrun_triton_bechmark

(args: Args)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:267

↓ 2 callersFunctionsample

( weights: torch.Tensor, # [V, D] (may be a TP shard over dim V) hidden_states: torch.Tensor, # [n_h

src/fused_mm_sampling/core.py:29

↓ 2 callersFunctionset_torch_allocator_for_tma_descriptors_cached

From https://triton-lang.org/main/python-api/generated/triton.language.make_tensor_descriptor.html

src/fused_mm_sampling/core.py:634

↓ 2 callersFunctionsupports_tma

()

src/fused_mm_sampling/persistent_matmul.py:39

↓ 2 callersFunctionsupports_warp_specialization_cached

()

src/fused_mm_sampling/core.py:336

↓ 2 callersFunctionunpack_grid

(grid)

src/fused_mm_sampling/core.py:417

↓ 1 callersFunction_apply_hatches

(ax, methods)

benchmarking/plot_bsz_sweep_runtime.py:101

↓ 1 callersFunction_as_torch_stream

(cs: "nvbench.CudaStream")

src/fused_mm_sampling/bench/speed_test.py:203

↓ 1 callersFunction_bitonic_merge

(x, ids, stage: core.constexpr, order: core.constexpr, n_dims: core.constexpr)

src/fused_mm_sampling/tl_argsort.py:55

↓ 1 callersFunction_check_nvcc_supports_gpu

Raise if the nvcc in cuda_home is too old for the current GPU.

tests/conftest.py:34

↓ 1 callersFunction_compare_and_swap

(x, ids, flip, i: core.constexpr, n_dims: core.constexpr)

src/fused_mm_sampling/tl_argsort.py:66

↓ 1 callersFunction_default_top_k_top_p

Set flashinfer defaults for top_k/top_p when not provided by the caller.

src/fused_mm_sampling/core.py:784

↓ 1 callersFunction_discover_rerun_folders

Find sibling rerun folders (e.g. b200-rerun1/tp1) for a primary folder. ``folder`` is typically ``.../<gpu>/<tp>``. The parent's siblings matchin

benchmarking/plot-triton-bench.py:473

↓ 1 callersFunction_fast_multinomial

Sample from a categorical distribution using the exponential race method. Avoids torch.multinomial's 10-kernel validation overhead (~2/3 of its r

src/fused_mm_sampling/core.py:64

↓ 1 callersFunction_find_cuda_home

Find a CUDA toolkit installation with nvcc that supports the current GPU.

tests/conftest.py:6

↓ 1 callersFunction_find_epilogue_if

Find the scf.if that guards the epilogue (masking + sample + store). This is the scf.if whose body contains a tt.reduce (the argmax).

benchmarking/insert_proton_records.py:149

↓ 1 callersFunction_find_free_port

()

src/fused_mm_sampling/tp_info.py:146

↓ 1 callersFunction_find_insertion_points

Scan the TTGIR and return (line_index, record_text) pairs. Each pair means "insert this proton.record line BEFORE lines[line_index]".

benchmarking/insert_proton_records.py:77

↓ 1 callersFunction_find_reduce_close

Find the closing line of the last tt.reduce block. Returns the index of the `}) : ...` line after tt.reduce.return.

benchmarking/insert_proton_records.py:167

↓ 1 callersFunction_get_module

()

src/fused_mm_sampling/cuda_impl.py:18

↓ 1 callersFunction_gpu_numa_node

Look up the NUMA node for a GPU from sysfs.

benchmarking/nsys_wrapper.py:21

↓ 1 callersFunction_gumbel_noise

(seed, pid_v, pid_h, sample_idx, noise_offsets)

src/fused_mm_sampling/core.py:620

↓ 1 callersFunction_is_torchrun

Check if the current process was launched by torchrun.

src/fused_mm_sampling/tp_info.py:60

↓ 1 callersFunction_label_bar

Place a centered label inside a bar segment if it's tall enough.

benchmarking/plot_ncu_kernel_breakdown.py:187

↓ 1 callersFunction_launch_kernel

(weights, hidden_states, out_vals, out_ids, V, D, H, top_k, temperature)

src/fused_mm_sampling/tl_fused_mm_topk.py:119

↓ 1 callersFunction_load_sweep

Load all NCU CSV files into a single DataFrame. Columns: bsz, method, kernel_name, duration_us. TP>1 durations are averaged across ranks.

benchmarking/parse_ncu_sweep.py:55

↓ 1 callersFunction_numa_bind

Pin this process to the CPU cores on the same NUMA node as the GPU. A binding failure (e.g. Modal's gVisor returning an empty CPU set) leaves

src/fused_mm_sampling/tp_info.py:118

↓ 1 callersFunction_numa_bind

Pin this process to the NUMA node of its GPU (CPU + memory).

benchmarking/nsys_wrapper.py:49

↓ 1 callersFunction_print_and_dump_cupti_results

(rows: list[dict], args: Args)

src/fused_mm_sampling/bench/speed_test.py:238

↓ 1 callersFunction_print_and_dump_own_results

(df: pd.DataFrame, args: Args)

src/fused_mm_sampling/bench/speed_test.py:262

↓ 1 callersFunction_print_numa_info

Print NUMA/CPU topology before nsys captures stdout.

benchmarking/nsys_wrapper.py:79

↓ 1 callersFunction_print_summary

Print all summary sections from the combined sweep DataFrame.

benchmarking/parse_ncu_sweep.py:106

↓ 1 callersFunction_python_version

()

src/fused_mm_sampling/bench/sys_metadata.py:22

↓ 1 callersFunction_read_gpu_name_from_logs

Parse the GPU name from the 'GPU: ...' line in logs.txt.

benchmarking/plot-triton-bench.py:90

↓ 1 callersFunction_resolve_cases

(case: str)

src/fused_mm_sampling/bench/triton_benchmark_lib.py:259

↓ 1 callersFunction_run_benchmark

Common benchmark logic for all modes.

src/fused_mm_sampling/bench/triton_benchmark_lib.py:153

↓ 1 callersFunction_system_cuda_version

()

src/fused_mm_sampling/bench/sys_metadata.py:28

next →1–100 of 315, ranked by callers