Skip to content

Conversation

@chdhr-harshal
Copy link

@chdhr-harshal chdhr-harshal commented Nov 16, 2025

Note: This implementation creates mapping per request. Do not continue if output key is not found, raise warning or error.

TODO: Validate output mapping at the model loading time/use warmup and cache.

UPDATE:

  • Now the validation happens at first request, and caches the mapping.
  • Raises warnings if model outputs extra keys taht are not present in the config. Raises error if triton config keys are not present in model outputs.
  • Raises error if model outputs something that cannot be converted into a tensor (e.g., strings)

TODO:

  • Add unit tests for new functions.
  • Better errors and warnings with exact keys that are extra or missing.

if (${TRITON_PYTORCH_NVSHMEM})
set(PT_LIBS
${PT_LIBS}
"libtorch_nvshmem.so"
Copy link
Author

@chdhr-harshal chdhr-harshal Nov 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing this temporarily because I don't have GPU on my personal machine.

std::lock_guard<std::mutex> lock(dict_validation_mutex_);
if (dict_output_validated_.load(std::memory_order_acquire)) {
return nullptr;
}
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need to use double checking with mutex to make validation threadsafe.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

1 participant